refactor: build replay envs from scene configs

Author: juelg

python/rcs/sim_state_replay.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import copy
+import importlib
 import time
 from dataclasses import dataclass
 from pathlib import Path
@@ -12,6 +14,7 @@
 import typer
 from PIL import Image
 from rcs.envs.base import ControlMode
+from rcs.envs.scenes import SimEnvCreator
 from rcs.envs.sim import SimStateObservationWrapper
 
 import rcs  # noqa: F401
@@ -20,13 +23,17 @@
 
 DATASET_ARGUMENT = typer.Argument(..., exists=True, file_okay=False, dir_okay=True)
 
-ENV_ID_OPTION = typer.Option(help="Gymnasium env id used for replay.")
+SCENE_CONFIG_OPTION = typer.Option(
+    help="Scene creator class name or import path used to build the replay env, for example EmptyWorldFR3."
+)
 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.")
+CAMERA_OPTION = typer.Option("--camera", help="Optional camera names to enable on the replay env.")
+RESOLUTION_OPTION = typer.Option(help="Optional replay camera resolution as WIDTH HEIGHT.")
+FRAME_RATE_OPTION = typer.Option(help="Optional replay camera frame rate override.")
+RENDER_MODE_OPTION = typer.Option(
+    help="Replay viewer mode. Use 'human' to open the GUI, anything else keeps it headless."
+)
+CONTROL_MODE_OPTION = typer.Option(help="Optional control mode name override 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.")
@@ -202,30 +209,89 @@ def replay_trajectory(
             time.sleep(sleep_s)
 
 
+def resolve_scene_creator(scene_config: str) -> SimEnvCreator[Any]:
+    module_name, _, class_name = scene_config.rpartition(".")
+    if not module_name:
+        module_name = "rcs.envs.configs"
+        class_name = scene_config
+    module = importlib.import_module(module_name)
+    scene_creator_cls = getattr(module, class_name, None)
+    if scene_creator_cls is None:
+        msg = f"Could not find scene creator '{class_name}' in module '{module_name}'."
+        raise ValueError(msg)
+    scene_creator = scene_creator_cls()
+    if not isinstance(scene_creator, SimEnvCreator):
+        msg = f"{module_name}.{class_name} is not an RCS SimEnvCreator."
+        raise TypeError(msg)
+    return scene_creator
+
+
+def create_replay_env(
+    scene_config: str,
+    *,
+    camera: list[str] | None = None,
+    resolution: tuple[int, int] | None = None,
+    frame_rate: int | None = None,
+    render_mode: str = "human",
+    control_mode: str | None = None,
+) -> gym.Env:
+    scene_creator = resolve_scene_creator(scene_config)
+    cfg = scene_creator.config()
+
+    if control_mode is not None:
+        cfg.control_mode = ControlMode[control_mode]
+    cfg.sim_cfg.async_control = False
+    cfg.sim_cfg.realtime = False
+    cfg.headless = render_mode != "human"
+
+    if camera is not None:
+        if cfg.camera_cfgs is None:
+            msg = f"Scene config '{scene_config}' does not define any cameras to replay."
+            raise ValueError(msg)
+        missing_cameras = [name for name in camera if name not in cfg.camera_cfgs]
+        if missing_cameras:
+            available_cameras = sorted(cfg.camera_cfgs)
+            msg = (
+                f"Unknown replay cameras {missing_cameras} for scene config '{scene_config}'. "
+                f"Available cameras are {available_cameras}."
+            )
+            raise ValueError(msg)
+        cfg.camera_cfgs = {name: copy.deepcopy(cfg.camera_cfgs[name]) for name in camera}
+    elif cfg.camera_cfgs is not None:
+        cfg.camera_cfgs = {name: copy.deepcopy(camera_cfg) for name, camera_cfg in cfg.camera_cfgs.items()}
+
+    if cfg.camera_cfgs is not None:
+        for camera_cfg in cfg.camera_cfgs.values():
+            if resolution is not None:
+                camera_cfg.resolution_width, camera_cfg.resolution_height = resolution
+            if frame_rate is not None:
+                camera_cfg.frame_rate = frame_rate
+
+    return scene_creator.create_env(cfg)
+
+
 @app.command()
 def replay(
     dataset: Annotated[Path, DATASET_ARGUMENT],
-    env_id: Annotated[str, ENV_ID_OPTION] = "rcs/FR3SimplePickUpSim-v0",
+    scene_config: Annotated[str, SCENE_CONFIG_OPTION] = "EmptyWorldFR3",
     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,
+    resolution: Annotated[tuple[int, int] | None, RESOLUTION_OPTION] = None,
+    frame_rate: Annotated[int | None, FRAME_RATE_OPTION] = None,
     render_mode: Annotated[str, RENDER_MODE_OPTION] = "human",
-    control_mode: Annotated[str, CONTROL_MODE_OPTION] = ControlMode.CARTESIAN_TRPY.name,
+    control_mode: Annotated[str | None, CONTROL_MODE_OPTION] = None,
     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,
-        render_mode=render_mode,
-        control_mode=ControlMode[control_mode],
+    env = create_replay_env(
+        scene_config,
+        camera=camera,
         resolution=resolution,
         frame_rate=frame_rate,
-        cam_list=camera,
+        render_mode=render_mode,
+        control_mode=control_mode,
     )
     try:
         recorded_steps = load_trajectory(dataset, resolved_uuid, prefer_duckdb=prefer_duckdb)

python/tests/test_sim_state_record_replay.py

@@ -2,20 +2,25 @@
 
 from dataclasses import dataclass
 from pathlib import Path
+from types import SimpleNamespace
 
 import gymnasium as gym
 import mujoco as mj
 import numpy as np
 import pyarrow.dataset as ds
 from rcs._core.common import RobotPlatform
 from rcs.camera.interface import CameraFrame, DataFrame, Frame, FrameSet
+from rcs.envs.base import ControlMode
+from rcs.envs.configs import EmptyWorldFR3
 from rcs.envs.sim import SimStateObservationWrapper
 from rcs.envs.storage_wrapper import StorageWrapper
 from rcs.sim.sim import Sim
 from rcs.sim_state_replay import (
     RecordedSimStep,
+    create_replay_env,
     load_trajectory,
     replay_trajectory,
+    resolve_scene_creator,
     restore_sim_step,
 )
 
@@ -124,6 +129,62 @@ def reset(self, *, seed: int | None = None, options: dict | None = None):
         return {}, {}
 
 
+def test_resolve_scene_creator_defaults_to_rcs_env_configs():
+    assert isinstance(resolve_scene_creator("EmptyWorldFR3"), EmptyWorldFR3)
+
+
+@dataclass
+class DummyCameraConfig:
+    resolution_width: int
+    resolution_height: int
+    frame_rate: int
+
+
+class DummySceneCreator:
+    def __init__(self):
+        self.created_cfg = None
+
+    def config(self):
+        return SimpleNamespace(
+            control_mode=ControlMode.CARTESIAN_TQuat,
+            sim_cfg=SimpleNamespace(async_control=True, realtime=True),
+            headless=False,
+            camera_cfgs={
+                "bird_eye": DummyCameraConfig(1280, 720, 30),
+                "wrist": DummyCameraConfig(640, 480, 15),
+            },
+        )
+
+    def create_env(self, cfg):
+        self.created_cfg = cfg
+        return "dummy-env"
+
+
+def test_create_replay_env_uses_scene_config_and_overrides_camera_settings(monkeypatch):
+    scene_creator = DummySceneCreator()
+    monkeypatch.setattr("rcs.sim_state_replay.resolve_scene_creator", lambda _scene_config: scene_creator)
+
+    env = create_replay_env(
+        "EmptyWorldFR3",
+        camera=["wrist"],
+        resolution=(320, 240),
+        frame_rate=5,
+        render_mode="rgb_array",
+        control_mode=ControlMode.CARTESIAN_TRPY.name,
+    )
+
+    assert env == "dummy-env"
+    assert scene_creator.created_cfg is not None
+    assert scene_creator.created_cfg.control_mode == ControlMode.CARTESIAN_TRPY
+    assert scene_creator.created_cfg.sim_cfg.async_control is False
+    assert scene_creator.created_cfg.sim_cfg.realtime is False
+    assert scene_creator.created_cfg.headless is True
+    assert sorted(scene_creator.created_cfg.camera_cfgs) == ["wrist"]
+    assert scene_creator.created_cfg.camera_cfgs["wrist"].resolution_width == 320
+    assert scene_creator.created_cfg.camera_cfgs["wrist"].resolution_height == 240
+    assert scene_creator.created_cfg.camera_cfgs["wrist"].frame_rate == 5
+
+
 def test_replay_trajectory_syncs_gui_without_stepping():
     spy_sim = SpySim()
     env = SpyReplayEnv(spy_sim)