Faster resets

Author: daphne-cornelisse

Diff for: baselines/ppo/ppo_waypoint.py

@@ -117,6 +117,8 @@ def run(
     off_road_weight: Annotated[Optional[float], typer.Option(help="The weight for off-road penalty")] = None,
     goal_achieved_weight: Annotated[Optional[float], typer.Option(help="The weight for goal-achieved reward")] = None,
     waypoint_distance_scale: Annotated[Optional[float], typer.Option(help="Scale for realism rewards")] = None,
+    speed_distance_scale: Annotated[Optional[float], typer.Option(help="Scale for realism rewards")] = None,
+    jerk_smoothness_scale: Annotated[Optional[float], typer.Option(help="Scale for realism rewards")] = None,
     dist_to_goal_threshold: Annotated[Optional[float], typer.Option(help="The distance threshold for goal-achieved")] = None,
     randomize_rewards: Annotated[Optional[int], typer.Option(help="If reward_type == reward_conditioned, choose the condition_mode; 0 or 1")] = 0,
     sampling_seed: Annotated[Optional[int], typer.Option(help="The seed for sampling scenes")] = None,
@@ -174,6 +176,8 @@ def run(
         "off_road_weight": off_road_weight,
         "goal_achieved_weight": goal_achieved_weight,
         "waypoint_distance_scale": waypoint_distance_scale,
+        "jerk_smoothness_scale": jerk_smoothness_scale,
+        "speed_distance_scale": speed_distance_scale,
         "dist_to_goal_threshold": dist_to_goal_threshold,
         "sampling_seed": sampling_seed,
         "obs_radius": obs_radius,

Diff for: gpudrive/env/env_puffer.py

@@ -581,7 +581,10 @@ def step(self, action):
             self.last_obs = self.env.get_obs(self.controlled_agent_mask)
 
             # Asynchronously reset the done worlds and empty storage
-            self.env.reset(env_idx_list=done_worlds_cpu)
+            self.env.reset(
+                env_idx_list=done_worlds_cpu, 
+                mask=self.controlled_agent_mask
+            )
             self.episode_returns[done_worlds] = 0
             self.agent_episode_returns[done_worlds, :] = 0
             self.episode_lengths[done_worlds, :] = 0