refactor: review of sim state recording and replay

Author: pkrack

python/rcs/envs/sim.py

@@ -46,7 +46,6 @@ def reset(
 class SimStateObservationWrapper(ActObsInfoWrapper):
     STATE_KEY = "sim_state"
     STATE_SPEC_KEY = "sim_state_spec"
-    STATE_SIZE_KEY = "sim_state_size"
 
     def __init__(self, env):
         super().__init__(env)
@@ -58,7 +57,6 @@ def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tupl
         sim_state = self.sim.get_state()
         observation[self.STATE_KEY] = sim_state
         observation[self.STATE_SPEC_KEY] = self.sim.get_state_spec()
-        observation[self.STATE_SIZE_KEY] = sim_state.shape[0]
         return observation, info
 
 

python/rcs/sim_state_replay.py

@@ -19,16 +19,17 @@
 app = typer.Typer(help="Replay recorded MuJoCo trajectories from a parquet dataset.")
 
 DATASET_ARGUMENT = typer.Argument(..., exists=True, file_okay=False, dir_okay=True)
-ENV_ID_OPTION = typer.Option("rcs/FR3SimplePickUpSim-v0", help="Gymnasium env id used for replay.")
-TRAJECTORY_UUID_OPTION = typer.Option(None, help="UUID of the recorded trajectory to replay.")
-CAMERA_OPTION = typer.Option([], "--camera", help="Camera names to enable on the replay env.")
-RESOLUTION_OPTION = typer.Option((256, 256), help="Replay camera resolution as WIDTH HEIGHT.")
-FRAME_RATE_OPTION = typer.Option(0, help="Replay camera frame rate.")
-RENDER_MODE_OPTION = typer.Option("human", help="Gym render mode for the replay env.")
-CONTROL_MODE_OPTION = typer.Option(ControlMode.CARTESIAN_TRPY.name, help="Control mode name for env creation.")
-SLEEP_OPTION = typer.Option(0.0, help="Optional delay between restored states.")
-OUTPUT_DIR_OPTION = typer.Option(None, help="Optional directory for re-rendered RGB frames.")
-PREFER_DUCKDB_OPTION = typer.Option(True, help="Use duckdb for parquet loading when it is available.")
+
+ENV_ID_OPTION = typer.Option(help="Gymnasium env id used for replay.")
+TRAJECTORY_UUID_OPTION = typer.Option(help="UUID of the recorded trajectory to replay.")
+CAMERA_OPTION = typer.Option("--camera", help="Camera names to enable on the replay env.")
+RESOLUTION_OPTION = typer.Option(help="Replay camera resolution as WIDTH HEIGHT.")
+FRAME_RATE_OPTION = typer.Option(help="Replay camera frame rate.")
+RENDER_MODE_OPTION = typer.Option(help="Gym render mode for the replay env.")
+CONTROL_MODE_OPTION = typer.Option(help="Control mode name for env creation.")
+SLEEP_OPTION = typer.Option(help="Optional delay between restored states.")
+OUTPUT_DIR_OPTION = typer.Option(help="Optional directory for re-rendered RGB frames.")
+PREFER_DUCKDB_OPTION = typer.Option(help="Use duckdb for parquet loading when it is available.")
 
 
 @dataclass(frozen=True)
@@ -193,6 +194,7 @@ def replay_trajectory(
     env.reset()
     for recorded_step in recorded_steps:
         restore_sim_step(env, recorded_step)
+        env.get_wrapper_attr("sim").step(1)
         if output_dir is not None:
             save_rgb_frames(output_dir, recorded_step, collect_rgb_frames(env))
         if sleep_s > 0:
@@ -202,17 +204,19 @@ def replay_trajectory(
 @app.command()
 def replay(
     dataset: Annotated[Path, DATASET_ARGUMENT],
-    env_id: Annotated[str, ENV_ID_OPTION],
-    trajectory_uuid: Annotated[str | None, TRAJECTORY_UUID_OPTION],
-    camera: Annotated[list[str], CAMERA_OPTION],
-    resolution: Annotated[tuple[int, int], RESOLUTION_OPTION],
-    frame_rate: Annotated[int, FRAME_RATE_OPTION],
-    render_mode: Annotated[str, RENDER_MODE_OPTION],
-    control_mode: Annotated[str, CONTROL_MODE_OPTION],
-    sleep_s: Annotated[float, SLEEP_OPTION],
-    output_dir: Annotated[Path | None, OUTPUT_DIR_OPTION],
-    prefer_duckdb: Annotated[bool, PREFER_DUCKDB_OPTION],
+    env_id: Annotated[str, ENV_ID_OPTION] = "rcs/FR3SimplePickUpSim-v0",
+    trajectory_uuid: Annotated[str | None, TRAJECTORY_UUID_OPTION] = None,
+    camera: Annotated[list[str] | None, CAMERA_OPTION] = None,
+    resolution: Annotated[tuple[int, int], RESOLUTION_OPTION] = (256, 256),
+    frame_rate: Annotated[int, FRAME_RATE_OPTION] = 0,
+    render_mode: Annotated[str, RENDER_MODE_OPTION] = "human",
+    control_mode: Annotated[str, CONTROL_MODE_OPTION] = ControlMode.CARTESIAN_TRPY.name,
+    sleep_s: Annotated[float, SLEEP_OPTION] = 0.0,
+    output_dir: Annotated[Path | None, OUTPUT_DIR_OPTION] = None,
+    prefer_duckdb: Annotated[bool, PREFER_DUCKDB_OPTION] = True,
 ):
+    if camera is None:
+        camera = []
     resolved_uuid = resolve_trajectory_uuid(dataset, trajectory_uuid, prefer_duckdb=prefer_duckdb)
     env = gym.make(
         env_id,

python/tests/test_sim_state_record_replay.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-import importlib.util
-import sys
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -11,41 +9,10 @@
 import pyarrow.dataset as ds
 from rcs._core.common import RobotPlatform
 from rcs.camera.interface import CameraFrame, DataFrame, Frame, FrameSet
+from rcs.envs.sim import SimStateObservationWrapper
 from rcs.envs.storage_wrapper import StorageWrapper
-
-import rcs
-
-REPO_ROOT = Path(__file__).resolve().parents[2]
-
-
-def _load_local_module(module_name: str, relative_path: str):
-    module_path = REPO_ROOT / relative_path
-    spec = importlib.util.spec_from_file_location(module_name, module_path)
-    if spec is None or spec.loader is None:
-        msg = f"Could not create an import spec for {module_name} from {module_path}."
-        raise ImportError(msg)
-    module = importlib.util.module_from_spec(spec)
-    sys.modules[module_name] = module
-    parent_name, _, child_name = module_name.rpartition(".")
-    if parent_name:
-        parent_module = sys.modules[parent_name]
-        setattr(parent_module, child_name, module)
-    spec.loader.exec_module(module)
-    return module
-
-
-local_sim_module = _load_local_module("rcs.sim.sim", "python/rcs/sim/sim.py")
-rcs.sim.__dict__["Sim"] = local_sim_module.Sim
-_load_local_module("rcs.envs.sim", "python/rcs/envs/sim.py")
-_load_local_module("rcs.sim_state_replay", "python/rcs/sim_state_replay.py")
-
-from rcs.envs.sim import SimStateObservationWrapper  # noqa: E402
-from rcs.sim.sim import Sim  # noqa: E402
-from rcs.sim_state_replay import (  # noqa: E402
-    load_trajectory,
-    replay_trajectory,
-    restore_sim_step,
-)
+from rcs.sim.sim import Sim
+from rcs.sim_state_replay import load_trajectory, replay_trajectory, restore_sim_step
 
 XML = """
 <mujoco>