Add rollout_controls to pre_rollout function

docs/source/interface/tasks.md

@@ -42,24 +42,34 @@ class MyTask(Task[MyTaskConfig]):
 If the system is our `SimulationNode` object, then there are two copies of the `Task` in the system and the controller respectively. The `SimulationNode` is responsible for stepping the `mujoco` simulation, while the `Controller` is responsible for rolling out the task. We expose functions for modifying the task before and after each of these steps. Additionally, we also allow a task-specific optimizer warm start, which is useful for tasks that require some initial setup before the optimization loop starts. The interface for these functions is as follows:
 ```python
 class MyTask(Task[MyTaskConfig]):
-    def pre_rollout(self, curr_state: np.ndarray, config: MyTaskConfig) -> None:
+    def pre_rollout(
+        self,
+        curr_state: np.ndarray,
+        rollout_times: np.ndarray,
+        rollout_controls: np.ndarray,
+        config: MyTaskConfig,
+    ) -> None:
         """Pre-rollout behavior for task (does nothing by default).
 
         Args:
             curr_state: Current state of the task. Shape=(nq + nv,).
+            rollout_times: The rollout times for the current rollout (global time). Shape=(T,).
+            rollout_controls: The intended controls to apply during the rollout - can modify in place. Shape=(T, nu).
+            config: The current task config (passed in from the top-level controller).
         """
 
     def post_rollout(
         self,
         states: np.ndarray,
         sensors: np.ndarray,
+        times: np.ndarray,
         controls: np.ndarray,
         config: MyTaskConfig,
         system_metadata: dict[str, Any] | None = None,
     ) -> None:
         """Post-rollout behavior for task (does nothing by default).
 
-        Same inputs as in reward function.
+        Same inputs as in reward function except times, which are the global rollout times associated with controls.
         """
 
     def pre_sim_step(self) -> None:

judo/controller/controller.py

@@ -183,7 +183,7 @@ def update_action(self, curr_state: np.ndarray, curr_time: float) -> None:
             self.rollout_controls = candidate_splines(curr_time + self.rollout_times)
 
             # Roll out dynamics with action sequences.
-            self.task.pre_rollout(curr_state, self.task_cfg)
+            self.task.pre_rollout(curr_state, curr_time + self.rollout_times, self.rollout_controls, self.task_cfg)
             self.states, self.sensors = self.rollout_backend.rollout(
                 self.model_data_pairs,
                 curr_state,
@@ -192,6 +192,7 @@ def update_action(self, curr_state: np.ndarray, curr_time: float) -> None:
             self.task.post_rollout(
                 self.states,
                 self.sensors,
+                curr_time + self.rollout_times,
                 self.rollout_controls,
                 self.task_cfg,
                 self.system_metadata,

judo/tasks/base.py

@@ -89,25 +89,34 @@ def dt(self) -> float:
         """Returns Mujoco physics timestep for default physics task."""
         return self.model.opt.timestep
 
-    def pre_rollout(self, curr_state: np.ndarray, config: ConfigT) -> None:
+    def pre_rollout(
+        self,
+        curr_state: np.ndarray,
+        rollout_times: np.ndarray,
+        rollout_controls: np.ndarray,
+        config: ConfigT,
+    ) -> None:
         """Pre-rollout behavior for task (does nothing by default).
 
         Args:
             curr_state: Current state of the task. Shape=(nq + nv,).
+            rollout_times: The rollout times for the current rollout (global time). Shape=(T,).
+            rollout_controls: The intended controls to apply during the rollout - can modify in place. Shape=(T, nu).
             config: The current task config (passed in from the top-level controller).
         """
 
     def post_rollout(
         self,
         states: np.ndarray,
         sensors: np.ndarray,
+        times: np.ndarray,
         controls: np.ndarray,
         config: ConfigT,
         system_metadata: dict[str, Any] | None = None,
     ) -> None:
         """Post-rollout behavior for task (does nothing by default).
 
-        Same inputs as in reward function.
+        Same inputs as in reward function except times, which are the global rollout times associated with controls.
         """
 
     def pre_sim_step(self) -> None:

judo/tasks/fr3_pick.py

@@ -198,7 +198,13 @@ def check_sensor_dists(
         dist = sensors[:, :, i]
         return dist
 
-    def pre_rollout(self, curr_state: np.ndarray, config: FR3PickConfig) -> None:
+    def pre_rollout(
+        self,
+        curr_state: np.ndarray,
+        rollout_times: np.ndarray,
+        rollout_controls: np.ndarray,
+        config: FR3PickConfig,
+    ) -> None:
         """Computes the current phase of the system."""
         # update the data object associated with the current state
         self._data.qpos[:] = curr_state[: self.model.nq]