From 1d2b8f646a0d41428803fe7e61f584bc6b05da4f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20J=C3=BClg?= Date: Thu, 14 May 2026 21:05:59 +0200 Subject: [PATCH 1/2] feat(replayer): gym make for env creators - add gym.make registration in build in envs - add gym.make for replayer --- python/rcs/__main__.py | 13 ++++--------- python/rcs/envs/configs.py | 13 +++++++++---- python/rcs/sim/replayer.py | 29 +++++++++++------------------ 3 files changed, 24 insertions(+), 31 deletions(-) diff --git a/python/rcs/__main__.py b/python/rcs/__main__.py index 95d30a6d..62e7b784 100644 --- a/python/rcs/__main__.py +++ b/python/rcs/__main__.py @@ -72,14 +72,10 @@ def replay( str, typer.Option(help="RelativeTo enum name: CONFIGURED_ORIGIN, LAST_STEP, or NONE."), ] = "CONFIGURED_ORIGIN", - scene: Annotated[ + env_id: 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")', + typer.Option(help="Environment id used in gym.make()."), + ] = "rcs/duo", ): replay_dataset( dataset=dataset, @@ -87,8 +83,7 @@ def replay( headless=headless, frequency=frequency, relative_to=relative_to, - scene=scene, - task_cfg=task_cfg, + env_id=env_id, ) diff --git a/python/rcs/envs/configs.py b/python/rcs/envs/configs.py index c2a8f1ca..6601d35e 100644 --- a/python/rcs/envs/configs.py +++ b/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 diff --git a/python/rcs/sim/replayer.py b/python/rcs/sim/replayer.py index 05a091a1..1f2f3a36 100644 --- a/python/rcs/sim/replayer.py +++ b/python/rcs/sim/replayer.py @@ -1,4 +1,3 @@ -import typing from dataclasses import dataclass from pathlib import Path from typing import Any @@ -6,8 +5,6 @@ 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: From 0b8f58a01a2c3ecb7fbb9ba115923fdba929a776 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20J=C3=BClg?= Date: Thu, 14 May 2026 21:06:11 +0200 Subject: [PATCH 2/2] impr(replayer): allow imports from other modules --- python/rcs/__main__.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/python/rcs/__main__.py b/python/rcs/__main__.py index 62e7b784..32dd98a4 100644 --- a/python/rcs/__main__.py +++ b/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[ @@ -76,7 +84,15 @@ def replay( str, 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,