From 67916e6d1780b2b655f96b607f2940fd1f0c6e11 Mon Sep 17 00:00:00 2001 From: alberthli Date: Thu, 21 Aug 2025 10:50:37 -0700 Subject: [PATCH 1/2] add rollout_controls to pre_rollout function --- docs/source/interface/tasks.md | 3 ++- judo/controller/controller.py | 2 +- judo/tasks/base.py | 3 ++- judo/tasks/fr3_pick.py | 2 +- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/docs/source/interface/tasks.md b/docs/source/interface/tasks.md index 98e98127..5974b5a0 100644 --- a/docs/source/interface/tasks.md +++ b/docs/source/interface/tasks.md @@ -42,11 +42,12 @@ 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_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_controls: The intended controls to apply during the rollout - can modify in place. Shape=(T, nu). """ def post_rollout( diff --git a/judo/controller/controller.py b/judo/controller/controller.py index f261baf7..c3c7dad3 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, self.rollout_controls, self.task_cfg) self.states, self.sensors = self.rollout_backend.rollout( self.model_data_pairs, curr_state, diff --git a/judo/tasks/base.py b/judo/tasks/base.py index e21de260..6f4b91ac 100644 --- a/judo/tasks/base.py +++ b/judo/tasks/base.py @@ -89,11 +89,12 @@ 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_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_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). """ diff --git a/judo/tasks/fr3_pick.py b/judo/tasks/fr3_pick.py index 6a18e6fc..b6711daa 100644 --- a/judo/tasks/fr3_pick.py +++ b/judo/tasks/fr3_pick.py @@ -198,7 +198,7 @@ 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_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] From baed7e142a32ed1d3a40ea4ccf44c8e11be20fe0 Mon Sep 17 00:00:00 2001 From: alberthli Date: Thu, 21 Aug 2025 11:01:02 -0700 Subject: [PATCH 2/2] also add times to pre_rollout and post_rollout --- docs/source/interface/tasks.md | 13 +++++++++++-- judo/controller/controller.py | 3 ++- judo/tasks/base.py | 12 ++++++++++-- judo/tasks/fr3_pick.py | 8 +++++++- 4 files changed, 30 insertions(+), 6 deletions(-) diff --git a/docs/source/interface/tasks.md b/docs/source/interface/tasks.md index 5974b5a0..06751b87 100644 --- a/docs/source/interface/tasks.md +++ b/docs/source/interface/tasks.md @@ -42,25 +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, rollout_controls: 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 c3c7dad3..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.rollout_controls, 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 6f4b91ac..2e5b910a 100644 --- a/judo/tasks/base.py +++ b/judo/tasks/base.py @@ -89,11 +89,18 @@ 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, rollout_controls: 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). """ @@ -102,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 b6711daa..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, rollout_controls: 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]