Refactor replay onto existing sim_state flow

Author: juelg

python/rcs/envs/sim.py

@@ -44,33 +44,32 @@ def reset(
 
 
 class SimStateObservationWrapper(ActObsInfoWrapper):
-    DYNAMIC_JOINT_SCHEMA_KEY = "dynamic_joint_schema"
-    DYNAMIC_JOINT_QPOS_KEY = "dynamic_joint_qpos"
-    DYNAMIC_JOINT_QVEL_KEY = "dynamic_joint_qvel"
+    STATE_KEY = "sim_state"
+    STATE_SPEC_KEY = "sim_state_spec"
+    STATE_SIZE_KEY = "sim_state_size"
 
     def __init__(self, env):
         super().__init__(env)
         assert self.env.get_wrapper_attr("PLATFORM") == RobotPlatform.SIMULATION, "Base environment must be simulation."
         self.sim = cast(sim.Sim, self.get_wrapper_attr("sim"))
-        self._dynamic_joint_schema = self.sim.get_dynamic_joint_schema()
-        self._include_schema_in_next_step = True
+        self._state_spec = self.sim.get_state_spec()
+        self._include_state_spec_in_next_step = True
 
     def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
         observation = dict(observation)
-        dynamic_joint_state = self.sim.get_dynamic_joint_state()
-        observation[self.DYNAMIC_JOINT_QPOS_KEY] = dynamic_joint_state["qpos"]
-        observation[self.DYNAMIC_JOINT_QVEL_KEY] = dynamic_joint_state["qvel"]
-        if self._include_schema_in_next_step:
-            observation[self.DYNAMIC_JOINT_SCHEMA_KEY] = self._dynamic_joint_schema
-            self._include_schema_in_next_step = False
+        sim_state = self.sim.get_state()
+        observation[self.STATE_KEY] = sim_state
+        observation[self.STATE_SIZE_KEY] = sim_state.shape[0]
+        if self._include_state_spec_in_next_step:
+            observation[self.STATE_SPEC_KEY] = self._state_spec
+            self._include_state_spec_in_next_step = False
         return observation, info
 
     def reset(
         self, *, seed: int | None = None, options: dict[str, Any] | None = None
     ) -> tuple[dict[str, Any], dict[str, Any]]:
         obs, info = super().reset(seed=seed, options=options)
-        # Re-emit the schema on the first recorded step after each reset.
-        self._include_schema_in_next_step = True
+        self._include_state_spec_in_next_step = True
         return obs, info
 
 

python/rcs/sim/sim.py

@@ -47,8 +47,6 @@ def gui_loop(gui_uuid: str, close_event):
 
 
 class Sim(_Sim):
-    STATE_SPEC = mj.mjtState.mjSTATE_INTEGRATION
-
     def __init__(self, mjmdl: str | PathLike | ModelComposer, cfg: SimConfig | None = None):
         if isinstance(mjmdl, ModelComposer):
             self.model = mjmdl.get_model()
@@ -73,31 +71,38 @@ def __init__(self, mjmdl: str | PathLike | ModelComposer, cfg: SimConfig | None
         if cfg is not None:
             self.set_config(cfg)
 
-    def get_state_spec(self) -> int:
-        return int(self.STATE_SPEC)
+    def get_state_spec(self) -> dict[str, list[str] | list[int]]:
+        return self.get_dynamic_joint_schema()
 
-    def get_state_size(self, spec: int | None = None) -> int:
-        state_spec = self.STATE_SPEC if spec is None else mj.mjtState(spec)
-        return mj.mj_stateSize(self.model, state_spec)
+    def get_state_size(self, spec: dict[str, list[str] | list[int]] | None = None) -> int:
+        state_spec = self.get_state_spec() if spec is None else spec
+        qpos_size = sum(int(value) for value in state_spec["qpos_sizes"])
+        qvel_size = sum(int(value) for value in state_spec["qvel_sizes"])
+        return qpos_size + qvel_size
 
-    def get_state(self, spec: int | None = None) -> np.ndarray:
-        state_spec = self.STATE_SPEC if spec is None else mj.mjtState(spec)
-        state = np.empty(self.get_state_size(int(state_spec)), dtype=np.float64)
-        mj.mj_getState(self.model, self.data, state, state_spec)
-        return state
+    def get_state(self, spec: dict[str, list[str] | list[int]] | None = None) -> np.ndarray:
+        del spec
+        dynamic_joint_state = self.get_dynamic_joint_state()
+        return np.concatenate((dynamic_joint_state["qpos"], dynamic_joint_state["qvel"]))
 
-    def set_state(self, state: np.ndarray, spec: int | None = None):
-        state_spec = self.STATE_SPEC if spec is None else mj.mjtState(spec)
+    def set_state(
+        self,
+        state: np.ndarray,
+        spec: dict[str, list[str] | list[int]] | None = None,
+    ):
+        state_spec = self.get_state_spec() if spec is None else spec
         state_array = np.asarray(state, dtype=np.float64)
-        expected_size = self.get_state_size(int(state_spec))
+        expected_size = self.get_state_size(state_spec)
         if state_array.shape != (expected_size,):
-            msg = (
-                f"Expected MuJoCo state with shape ({expected_size},), "
-                f"got {state_array.shape} for spec {int(state_spec)}."
-            )
+            msg = f"Expected state with shape ({expected_size},), got {state_array.shape}."
             raise ValueError(msg)
-        mj.mj_setState(self.model, self.data, state_array, state_spec)
-        mj.mj_forward(self.model, self.data)
+
+        qpos_size = sum(int(value) for value in state_spec["qpos_sizes"])
+        dynamic_joint_state = {
+            "qpos": state_array[:qpos_size],
+            "qvel": state_array[qpos_size:],
+        }
+        self.set_dynamic_joint_state(state_spec, dynamic_joint_state)
 
     def get_dynamic_joint_schema(self) -> dict[str, list[str] | list[int]]:
         schema = super().get_dynamic_joint_schema()

python/rcs/sim_state_replay.py

@@ -39,21 +39,13 @@ class RecordedSimStep:
     observation: dict[str, Any]
 
     @property
-    def dynamic_joint_schema(self) -> dict[str, Any] | None:
-        schema = self.observation.get(SimStateObservationWrapper.DYNAMIC_JOINT_SCHEMA_KEY)
-        return dict(schema) if schema is not None else None
+    def sim_state(self) -> np.ndarray:
+        return np.asarray(self.observation[SimStateObservationWrapper.STATE_KEY], dtype=np.float64)
 
     @property
-    def dynamic_joint_state(self) -> dict[str, np.ndarray] | None:
-        if (
-            SimStateObservationWrapper.DYNAMIC_JOINT_QPOS_KEY not in self.observation
-            or SimStateObservationWrapper.DYNAMIC_JOINT_QVEL_KEY not in self.observation
-        ):
-            return None
-        return {
-            "qpos": np.asarray(self.observation[SimStateObservationWrapper.DYNAMIC_JOINT_QPOS_KEY], dtype=np.float64),
-            "qvel": np.asarray(self.observation[SimStateObservationWrapper.DYNAMIC_JOINT_QVEL_KEY], dtype=np.float64),
-        }
+    def sim_state_spec(self) -> dict[str, Any] | None:
+        schema = self.observation.get(SimStateObservationWrapper.STATE_SPEC_KEY)
+        return dict(schema) if schema is not None else None
 
 
 class DuckDBUnavailableError(RuntimeError):
@@ -161,19 +153,15 @@ def resolve_trajectory_uuid(dataset_path: Path, trajectory_uuid: str | None, pre
 def restore_sim_step(
     env: gym.Env,
     recorded_step: RecordedSimStep,
-    dynamic_joint_schema: dict[str, Any] | None = None,
+    sim_state_spec: dict[str, Any] | None = None,
 ):
-    sim = env.get_wrapper_attr("sim")
-    dynamic_joint_state = recorded_step.dynamic_joint_state
-    if dynamic_joint_state is None:
-        msg = "Recorded step is missing dynamic joint state data."
+    resolved_spec = sim_state_spec or recorded_step.sim_state_spec
+    if resolved_spec is None:
+        msg = "Recorded sim state is missing its schema."
         raise ValueError(msg)
 
-    resolved_schema = dynamic_joint_schema or recorded_step.dynamic_joint_schema
-    if resolved_schema is None:
-        msg = "Recorded dynamic joint state is missing its schema."
-        raise ValueError(msg)
-    sim.set_dynamic_joint_state(resolved_schema, dynamic_joint_state)
+    sim = env.get_wrapper_attr("sim")
+    sim.set_state(recorded_step.sim_state, spec=resolved_spec)
 
 
 def collect_rgb_frames(env: gym.Env) -> dict[str, np.ndarray]:
@@ -209,17 +197,16 @@ def replay_trajectory(
     output_dir: Path | None = None,
 ):
     if not recorded_steps:
-        msg = "No recorded dynamic joint states found in the requested trajectory."
+        msg = "No recorded sim states found in the requested trajectory."
         raise ValueError(msg)
 
-    dynamic_joint_schema = next(
-        (recorded_step.dynamic_joint_schema for recorded_step in recorded_steps if recorded_step.dynamic_joint_schema),
-        None,
+    sim_state_spec = next(
+        (recorded_step.sim_state_spec for recorded_step in recorded_steps if recorded_step.sim_state_spec), None
     )
 
     env.reset()
     for recorded_step in recorded_steps:
-        restore_sim_step(env, recorded_step, dynamic_joint_schema=dynamic_joint_schema)
+        restore_sim_step(env, recorded_step, sim_state_spec=sim_state_spec)
         if output_dir is not None:
             save_rgb_frames(output_dir, recorded_step, collect_rgb_frames(env))
         if sleep_s > 0: