diff --git a/docs/source/interface/tasks.md b/docs/source/interface/tasks.md index 98e98127..06751b87 100644 --- a/docs/source/interface/tasks.md +++ b/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: diff --git a/judo/controller/controller.py b/judo/controller/controller.py index f261baf7..8bfdda5a 100644 --- a/judo/controller/controller.py +++ b/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, diff --git a/judo/tasks/base.py b/judo/tasks/base.py index e21de260..2e5b910a 100644 --- a/judo/tasks/base.py +++ b/judo/tasks/base.py @@ -89,11 +89,19 @@ 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). """ @@ -101,13 +109,14 @@ 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: diff --git a/judo/tasks/fr3_pick.py b/judo/tasks/fr3_pick.py index 6a18e6fc..a151556e 100644 --- a/judo/tasks/fr3_pick.py +++ b/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]