feat: gym make and replayer with other modules

python/rcs/__main__.py

@@ -25,6 +25,14 @@
 app = typer.Typer()
 
 
+def _exec_import_statements(imports: list[str] | None) -> None:
+    if imports is None:
+        return
+
+    for import_statement in imports:
+        exec(import_statement, {})
+
+
 @app.command()
 def consolidate(
     path: Annotated[
@@ -72,23 +80,26 @@ def replay(
         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[
+    env_id: Annotated[
         str,
-        typer.Option(help="Python expression that evaluates to a task config."),
-    ] = 'env_tasks.PickTaskConfig(robot_name="right")',
+        typer.Option(help="Environment id used in gym.make()."),
+    ] = "rcs/duo",
+    imports: Annotated[
+        list[str] | None,
+        typer.Option(
+            "--import",
+            help="Python import statement to execute before resolving the environment. Repeat for multiple imports. Example: --import 'from rcs_duobench.tasks import bin_sort'",
+        ),
+    ] = None,
 ):
+    _exec_import_statements(imports)
     replay_dataset(
         dataset=dataset,
         output=output,
         headless=headless,
         frequency=frequency,
         relative_to=relative_to,
-        scene=scene,
-        task_cfg=task_cfg,
+        env_id=env_id,
     )
 
 

python/rcs/envs/configs.py

@@ -2,6 +2,7 @@
 import time
 from typing import ClassVar, Literal
 
+import gymnasium as gym
 import numpy as np
 from rcs._core.common import FrankaHandTCPOffset, GripperType, RobotType
 from rcs._core.sim import (
@@ -482,11 +483,15 @@ def config(self) -> SimEnvCreatorConfig:
         return cfg
 
 
+gym.register(id="rcs/fr3", entry_point=EmptyWorldFR3())
+gym.register(id="rcs/duo", entry_point=EmptyWorldFR3Duo())
+gym.register(id="rcs/ur5e", entry_point=EmptyWorldUR5e())
+gym.register(id="rcs/xarm7", entry_point=EmptyWorldXArm7())
+gym.register(id="rcs/so101", entry_point=EmptyWorldSO101())
+
+
 if __name__ == "__main__":
-    scene = EmptyWorldFR3Duo()
-    # scene = EmptyWorldFR3()
-    # scene = EmptyWorldUR5e()
-    env = scene.create_env(scene.config())
+    env = gym.make("rcs/duo")
     obs, info = env.reset()
     print(obs)
     # Duo

python/rcs/sim/replayer.py

@@ -1,13 +1,10 @@
-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, SimStateSchema
 from rcs.envs.scenes import SimEnvCreator
@@ -132,14 +129,12 @@ def replay(
     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")',
+    env_id: str = "rcs/duo",
 ):
-    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()
+    env_creator = gym.envs.registry[env_id].entry_point
+    assert isinstance(env_creator, SimEnvCreator), "only rcs envs are supported for replay"
+
+    sim_cfg_data = env_creator.config()
     sim_cfg_data.sim_cfg = SimConfig(
         async_control=True,
         realtime=not headless,
@@ -150,22 +145,20 @@ def replay(
     sim_cfg_data.relative_to = RelativeTo[relative_to.upper()]
     if sim_cfg_data.root_frame_objects is None:
         sim_cfg_data.root_frame_objects = {}
-    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 = sc.create_env(sim_cfg_data)
+    env_rel = env_creator.create_env(sim_cfg_data)
     env_rel = StorageWrapper(
-        env_rel,
-        str(output),
-        "",
+        env=env_rel,
+        base_dir=str(output),
+        instruction=None,
         batch_size=32,
         max_rows_per_group=100,
         max_rows_per_file=1000,
         always_record=True,
-        allow_wrapper_instruction=False,
+        allow_wrapper_instruction=True,
+        success_from_env=False,
     )
     try:
         for uuid in uuids: