feat(replayer): added cli for dataset replaying

Author: juelg

python/rcs/__main__.py

@@ -3,6 +3,7 @@
 
 import typer
 from rcs.envs.storage_wrapper import StorageWrapper
+from rcs.sim.replayer import replay as replay_dataset
 
 app = typer.Typer()
 
@@ -32,5 +33,39 @@ def consolidate(
     typer.echo("Done.")
 
 
+@app.command("replay")
+def replay(
+    dataset: Annotated[
+        Path,
+        typer.Argument(
+            exists=True,
+            help="Parquet dataset directory to replay.",
+        ),
+    ],
+    headless: Annotated[bool, typer.Option(help="Whether to run without GUI.")] = True,
+    frequency: Annotated[int, typer.Option(help="Simulation frequency to use during replay.")] = 30,
+    relative_to: Annotated[
+        str,
+        typer.Option(help="RelativeTo enum name: CONFIGURED_ORIGIN, LAST_STEP, or NONE."),
+    ] = "CONFIGURED_ORIGIN",
+    scene: Annotated[
+        str,
+        typer.Option(help="Python expression that evaluates to a scene instance."),
+    ] = "env_configs.EmptyWorldFR3Duo()",
+    task_cfg: Annotated[
+        str,
+        typer.Option(help="Python expression that evaluates to a task config."),
+    ] = 'env_tasks.PickTaskConfig(robot_name="right")',
+):
+    replay_dataset(
+        dataset=dataset,
+        headless=headless,
+        frequency=frequency,
+        relative_to=relative_to,
+        scene=scene,
+        task_cfg=task_cfg,
+    )
+
+
 if __name__ == "__main__":
     app()

python/rcs/sim/replayer.py

@@ -1,15 +1,17 @@
+import typing
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 
 import duckdb
 import gymnasium as gym
 import numpy as np
+import rcs.envs.configs as env_configs
+import rcs.envs.tasks as env_tasks
 from rcs._core.sim import SimConfig
 from rcs.envs.base import RelativeTo, SimEnv
-from rcs.envs.configs import EmptyWorldFR3Duo
+from rcs.envs.scenes import SimEnvCreator
 from rcs.envs.storage_wrapper import StorageWrapper
-from rcs.envs.tasks import PickTaskConfig
 
 DATASET_PATH = "recorded_iris"
 
@@ -114,22 +116,36 @@ def replay_trajectory(env: gym.Env, recorded_steps: list[RecordedSimStep], headl
             env.get_wrapper_attr("success")()
 
 
-def replay():
-    dataset = "/home/tobi/coding/rcs_repos/robot-control-stack/test_iris"
-    headless = False
-    scene = EmptyWorldFR3Duo()
-    sim_cfg_data = scene.config()
-    sim_cfg_data.camera_cfgs = {}
-    sim_cfg_data.sim_cfg = SimConfig(async_control=True, realtime=not headless, frequency=30, max_convergence_steps=500)
+def replay(
+    dataset: Path | str,
+    headless: bool = True,
+    frequency: int = 30,
+    relative_to: str = RelativeTo.CONFIGURED_ORIGIN.name,
+    scene: str = "env_configs.EmptyWorldFR3Duo()",
+    task_cfg: str = 'env_tasks.PickTaskConfig(robot_name="right")',
+):
+    exec_scope = {**globals(), "__builtins__": __builtins__, "env_configs": env_configs, "env_tasks": env_tasks}
+    scene_locals: dict[str, Any] = {}
+    exec(f"_result = {scene}", exec_scope, scene_locals)
+    sc = typing.cast(SimEnvCreator, scene_locals["_result"])
+    sim_cfg_data = sc.config()
+    sim_cfg_data.sim_cfg = SimConfig(
+        async_control=True,
+        realtime=not headless,
+        frequency=frequency,
+        max_convergence_steps=500,
+    )
     sim_cfg_data.headless = headless
-    sim_cfg_data.relative_to = RelativeTo.CONFIGURED_ORIGIN
+    sim_cfg_data.relative_to = RelativeTo[relative_to.upper()]
     if sim_cfg_data.root_frame_objects is None:
         sim_cfg_data.root_frame_objects = {}
-    sim_cfg_data.task_cfg = PickTaskConfig(robot_name="right")
+    task_cfg_locals: dict[str, Any] = {}
+    exec(f"_result = {task_cfg}", exec_scope, task_cfg_locals)
+    sim_cfg_data.task_cfg = task_cfg_locals["_result"]
 
     uuids = load_distinct_uuids(dataset)
 
-    env_rel = scene.create_env(sim_cfg_data)
+    env_rel = sc.create_env(sim_cfg_data)
     env_rel = StorageWrapper(
         env_rel,
         DATASET_PATH,
@@ -148,7 +164,3 @@ def replay():
             replay_trajectory(env_rel, recorded_steps, headless)
     finally:
         env_rel.close()
-
-
-if __name__ == "__main__":
-    replay()