fixed planning at end of trahectory

Author: FedeClaudi

proj/environment/environment.py

@@ -83,6 +83,8 @@ def plan(self, curr_x, g_traj, itern):
             Given the current state and the goal trajectory, 
             find the next N sates, based on planning
         """
+
+        traj_length = len(g_traj)
         n_ahead = self.model.planning["n_ahead"]
         pred_len = self.model.planning["prediction_length"] + 1
 
@@ -96,24 +98,35 @@ def plan(self, curr_x, g_traj, itern):
         self.model.curr_traj_waypoint_idx = self.curr_traj_waypoint_idx
         self.current_traj_waypoint = g_traj[min_idx, :]
 
-        start = min_idx + n_ahead
-        if start > len(g_traj):
-            start = len(g_traj)
+        start = min_idx + n_ahead  # don't overshoot
+        if start > traj_length:
+            start = traj_length
 
         end = min_idx + n_ahead + pred_len
 
-        if start + pred_len > len(g_traj):
-            end = len(g_traj) - 2
+        if start + pred_len > traj_length:
+            end = traj_length
 
-        if abs(start - end) != pred_len:
-            g_traj = g_traj[start:end]
-            len_diff = (end - start) - pred_len
+        # Make sure planned trajectory has the correct length
+        if (end - start) != pred_len:
+            planned = g_traj[start:end]
+            len_diff = len(planned) - pred_len
 
             if len_diff <= 0:
-                len_diff = 1
-                return np.pad(g_traj, ((0, len_diff), (0, 0)), mode="edge")
+                planned = np.pad(
+                    planned, ((0, abs(len_diff)), (0, 0)), mode="edge"
+                )
+            else:
+                raise ValueError("Something went wrong")
+        else:
+            planned = g_traj[start:end]
+
+        if len(planned) != pred_len:
+            raise ValueError(
+                f"Planned trajecotry length should be {pred_len} but it is {len(planned)} instead"
+            )
         else:
-            return g_traj[start:end]
+            return planned
 
     def isdone(self, curr_x, trajectory):
         """

proj/environment/trajectories.py

@@ -50,6 +50,10 @@ def compute_trajectory_stats(
     min_dist_travelled=150,
     mute=False,
 ):
+
+    # ! shorten traj
+    trajectory = trajectory[: planning_params["prediction_length"] + 1, :]
+
     # Compute stats
     n_points = len(trajectory)
     distance_travelled = np.sum(
@@ -136,11 +140,11 @@ def compute_trajectory_stats(
             print(
                 f"[bold red]Lookahead of {lookahead} is {perc_lookahead} of the # of waypoints, that might be too low. Values closer to 5% are advised."
             )
-        if distance_travelled < min_dist_travelled * params["px_to_cm"]:
-            logger.warning(
-                "Distance travelled below minimal requirement, erroring"
-            )
-            return None, None
+        # if distance_travelled < min_dist_travelled * params["px_to_cm"]:
+        #     logger.warning(
+        #         "Distance travelled below minimal requirement, erroring"
+        #     )
+        #     return None, None
 
     return trajectory, duration, metadata
 

proj/run/runner.py

@@ -117,6 +117,8 @@ def run_experiment(
                 logger.exception(
                     f"Failed to take next step in simulation.\nError: {e}\n\n"
                 )
+                break
+
         try:
             environment.conclude()
         except:

run_allocentric.py

@@ -25,7 +25,7 @@
 # plot_trajectory(env.reset())
 
 # %%
-run_experiment(env, control, model, n_secs=0.1)
+run_experiment(env, control, model, n_secs=0.5)
 
 
 # %%