From be70315583dbe862e306ac1da06a8f9da3c4cca1 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 16 Apr 2026 18:32:24 +0200 Subject: [PATCH] fix: integrate PR #3311 review feedback - envs: rename obs keys to pixels/image, pixels/wrist_image, agent_pos - envs: add __post_init__ for dynamic action_dim in RoboMMEEnv config - envs: remove special-case obs conversion in utils.py (no longer needed) - ci: add Docker Hub login, HF_USER_TOKEN guard, --env.task_ids=[0] - scripts: extract_task_descriptions supports multiple task_ids - docs: title to # RoboMME, add image, restructure eval section - tests: update all key assertions to match new obs naming Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/workflows/benchmark_tests.yml | 8 +++++++ docs/source/robomme.mdx | 24 +++++++++++++++++-- scripts/ci/extract_task_descriptions.py | 19 ++++++++++++--- src/lerobot/envs/configs.py | 24 ++++++++++--------- src/lerobot/envs/robomme.py | 12 +++++----- src/lerobot/envs/utils.py | 21 ----------------- tests/test_robomme_env.py | 31 +++++++++++-------------- 7 files changed, 79 insertions(+), 60 deletions(-) diff --git a/.github/workflows/benchmark_tests.yml b/.github/workflows/benchmark_tests.yml index e112c5842..a0a5c99a7 100644 --- a/.github/workflows/benchmark_tests.yml +++ b/.github/workflows/benchmark_tests.yml @@ -332,6 +332,12 @@ jobs: with: cache-binary: false + - name: Login to Docker Hub + uses: docker/login-action@v3 # zizmor: ignore[unpinned-uses] + with: + username: ${{ secrets.DOCKERHUB_LEROBOT_USERNAME }} + password: ${{ secrets.DOCKERHUB_LEROBOT_PASSWORD }} + - name: Build RoboMME benchmark image uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses] with: @@ -342,6 +348,7 @@ jobs: tags: lerobot-benchmark-robomme:ci - name: Run RoboMME smoke eval (1 episode) + if: env.HF_USER_TOKEN != '' run: | docker run --name robomme-eval --gpus all \ --shm-size=4g \ @@ -356,6 +363,7 @@ jobs: --env.type=robomme \ --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \ --env.dataset_split=test \ + --env.task_ids=[0] \ --eval.batch_size=1 \ --eval.n_episodes=1 \ --eval.use_async_envs=false \ diff --git a/docs/source/robomme.mdx b/docs/source/robomme.mdx index 219395215..426462254 100644 --- a/docs/source/robomme.mdx +++ b/docs/source/robomme.mdx @@ -1,4 +1,4 @@ -# RoboMME Benchmark +# RoboMME [RoboMME](https://robomme.github.io) is a memory-augmented manipulation benchmark built on ManiSkill (SAPIEN). It evaluates a robot's ability to retain and use information across an episode — counting, object permanence, reference, and imitation. @@ -7,6 +7,8 @@ - **Dataset**: [`lerobot/robomme`](https://huggingface.co/datasets/lerobot/robomme) — LeRobot v3.0, 768K frames at 10 fps - **Simulator**: ManiSkill / SAPIEN, Panda arm, Linux only +![RoboMME benchmark tasks overview](https://raw.githubusercontent.com/RoboMME/robomme_benchmark/main/assets/teaser.png) + ## Tasks | Suite | Tasks | @@ -44,12 +46,30 @@ The `Dockerfile.eval-robomme` overrides `gymnasium==0.29.1` and `numpy==1.26.4` ## Running Evaluation +### Default (single task, single episode) + ```bash lerobot-eval \ --policy.path= \ --env.type=robomme \ --env.task=PickXtimes \ --env.dataset_split=test \ + --env.task_ids=[0] \ + --eval.batch_size=1 \ + --eval.n_episodes=1 +``` + +### Multi-task evaluation + +Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to control which episodes are evaluated per task. Recommended: 50 episodes per task for the test split. + +```bash +lerobot-eval \ + --policy.path= \ + --env.type=robomme \ + --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \ + --env.dataset_split=test \ + --env.task_ids=[0,1,2,3,4,5,6,7,8,9] \ --eval.batch_size=1 \ --eval.n_episodes=50 ``` @@ -58,7 +78,7 @@ lerobot-eval \ | Option | Default | Description | | -------------------- | ------------- | -------------------------------------------------- | -| `env.task` | `PickXtimes` | Any of the 16 task names above | +| `env.task` | `PickXtimes` | Any of the 16 task names above (comma-separated) | | `env.dataset_split` | `test` | `train`, `val`, or `test` | | `env.action_space` | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D) | | `env.episode_length` | `300` | Max steps per episode | diff --git a/scripts/ci/extract_task_descriptions.py b/scripts/ci/extract_task_descriptions.py index 217fd38da..70a068a5d 100644 --- a/scripts/ci/extract_task_descriptions.py +++ b/scripts/ci/extract_task_descriptions.py @@ -77,13 +77,16 @@ _ROBOMME_DESCRIPTIONS = { } -def _robomme_descriptions(task_names: str) -> dict[str, str]: +def _robomme_descriptions(task_names: str, task_ids: list[int] | None = None) -> dict[str, str]: """Return descriptions for each requested RoboMME task. Keys match the video filename pattern `_` used by the eval script.""" + if task_ids is None: + task_ids = [0] out: dict[str, str] = {} for name in (t.strip() for t in task_names.split(",") if t.strip()): desc = _ROBOMME_DESCRIPTIONS.get(name, name) - out[f"{name}_0"] = desc + for tid in task_ids: + out[f"{name}_{tid}"] = desc return out @@ -91,9 +94,19 @@ def main() -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)") parser.add_argument("--task", required=True, help="Task/suite name (e.g. libero_spatial)") + parser.add_argument( + "--task-ids", + type=str, + default=None, + help="Comma-separated task IDs (e.g. '0,1,2'). Default: [0]", + ) parser.add_argument("--output", required=True, help="Path to write task_descriptions.json") args = parser.parse_args() + task_ids: list[int] | None = None + if args.task_ids: + task_ids = [int(x.strip()) for x in args.task_ids.split(",")] + descriptions: dict[str, str] = {} try: if args.env == "libero": @@ -101,7 +114,7 @@ def main() -> int: elif args.env == "metaworld": descriptions = _metaworld_descriptions(args.task) elif args.env == "robomme": - descriptions = _robomme_descriptions(args.task) + descriptions = _robomme_descriptions(args.task, task_ids=task_ids) else: print( f"[extract_task_descriptions] No description extractor for env '{args.env}'.", diff --git a/src/lerobot/envs/configs.py b/src/lerobot/envs/configs.py index a7731a7bb..21030ce63 100644 --- a/src/lerobot/envs/configs.py +++ b/src/lerobot/envs/configs.py @@ -595,23 +595,25 @@ class RoboMMEEnv(EnvConfig): action_space: str = "joint_angle" # or "ee_pose" (7-D) dataset_split: str = "test" # "train" | "val" | "test" task_ids: list[int] | None = None - features: dict[str, PolicyFeature] = field( - default_factory=lambda: { - ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(8,)), - "image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)), - "wrist_image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)), - OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,)), - } - ) + features: dict[str, PolicyFeature] = field(default_factory=dict) features_map: dict[str, str] = field( default_factory=lambda: { ACTION: ACTION, - "image": f"{OBS_IMAGES}.image", - "wrist_image": f"{OBS_IMAGES}.wrist_image", - OBS_STATE: OBS_STATE, + "pixels/image": f"{OBS_IMAGES}.image", + "pixels/wrist_image": f"{OBS_IMAGES}.wrist_image", + "agent_pos": OBS_STATE, } ) + def __post_init__(self): + action_dim = 8 if self.action_space == "joint_angle" else 7 + self.features = { + ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)), + "pixels/image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)), + "pixels/wrist_image": PolicyFeature(type=FeatureType.VISUAL, shape=(256, 256, 3)), + "agent_pos": PolicyFeature(type=FeatureType.STATE, shape=(8,)), + } + @property def gym_kwargs(self) -> dict: return {} diff --git a/src/lerobot/envs/robomme.py b/src/lerobot/envs/robomme.py index 0f155710e..7a0a76ed1 100644 --- a/src/lerobot/envs/robomme.py +++ b/src/lerobot/envs/robomme.py @@ -81,9 +81,9 @@ class RoboMMEGymEnv(gym.Env): self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(action_dim,), dtype=np.float32) self.observation_space = spaces.Dict( { - "image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), - "wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), - "state": spaces.Box(-np.inf, np.inf, shape=(8,), dtype=np.float32), + "pixels/image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), + "pixels/wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), + "agent_pos": spaces.Box(-np.inf, np.inf, shape=(8,), dtype=np.float32), } ) @@ -146,9 +146,9 @@ class RoboMMEGymEnv(gym.Env): state = np.concatenate([joint, gripper]) return { - "image": front_rgb, - "wrist_image": wrist_rgb, - "state": state, + "pixels/image": front_rgb, + "pixels/wrist_image": wrist_rgb, + "agent_pos": state, } def _convert_info(self, info: dict) -> dict: diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index f57df43ca..b0d834a05 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -107,27 +107,6 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if "camera_obs" in observations: return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"] - # Handle flat image keys (e.g., RoboMME: "image", "wrist_image", "state") - for key in ("image", "wrist_image"): - if key in observations: - img = observations[key] - img_tensor = torch.from_numpy(img) if isinstance(img, np.ndarray) else img - if img_tensor.ndim == 3: - img_tensor = img_tensor.unsqueeze(0) - img_tensor = einops.rearrange(img_tensor, "b h w c -> b c h w").contiguous() - img_tensor = img_tensor.float() / 255 if img_tensor.dtype == torch.uint8 else img_tensor.float() - return_observations[f"{OBS_IMAGES}.{key}"] = img_tensor - - if "state" in observations and OBS_STATE not in return_observations: - state = ( - torch.from_numpy(observations["state"]).float() - if isinstance(observations["state"], np.ndarray) - else observations["state"].float() - ) - if state.dim() == 1: - state = state.unsqueeze(0) - return_observations[OBS_STATE] = state - return return_observations diff --git a/tests/test_robomme_env.py b/tests/test_robomme_env.py index 6802368bb..840a8bdcd 100644 --- a/tests/test_robomme_env.py +++ b/tests/test_robomme_env.py @@ -89,9 +89,9 @@ def test_robomme_features_map(): cfg = RoboMMEEnv() assert cfg.features_map[ACTION] == ACTION - assert cfg.features_map["image"] == f"{OBS_IMAGES}.image" - assert cfg.features_map["wrist_image"] == f"{OBS_IMAGES}.wrist_image" - assert cfg.features_map[OBS_STATE] == OBS_STATE + assert cfg.features_map["pixels/image"] == f"{OBS_IMAGES}.image" + assert cfg.features_map["pixels/wrist_image"] == f"{OBS_IMAGES}.wrist_image" + assert cfg.features_map["agent_pos"] == OBS_STATE def test_robomme_features_action_dim_joint_angle(): @@ -103,15 +103,12 @@ def test_robomme_features_action_dim_joint_angle(): def test_robomme_features_action_dim_ee_pose(): - """`ee_pose` uses a 7-D action; the features dict still declares 8-D - (the joint_angle default). Switching to ee_pose requires the user to - override `features[ACTION]`. This test documents that. - """ + """`ee_pose` uses a 7-D action; __post_init__ sets the correct shape.""" from lerobot.envs.configs import RoboMMEEnv from lerobot.utils.constants import ACTION cfg = RoboMMEEnv(action_space="ee_pose") - assert cfg.features[ACTION].shape == (8,) + assert cfg.features[ACTION].shape == (7,) # --------------------------------------------------------------------------- @@ -121,7 +118,7 @@ def test_robomme_features_action_dim_ee_pose(): def test_convert_obs_list_format(): """_convert_obs takes the last element from list-format obs fields and - emits the lerobot-canonical keys (image, wrist_image, state).""" + emits the lerobot-canonical keys (pixels/image, pixels/wrist_image, agent_pos).""" _install_robomme_stub() try: from lerobot.envs.robomme import RoboMMEGymEnv @@ -141,11 +138,11 @@ def test_convert_obs_list_format(): } result = env._convert_obs(obs_raw) - np.testing.assert_array_equal(result["image"], front) - np.testing.assert_array_equal(result["wrist_image"], wrist) - assert result["state"].shape == (8,) - np.testing.assert_array_almost_equal(result["state"][:7], joints) - assert result["state"][7] == gripper[0] + np.testing.assert_array_equal(result["pixels/image"], front) + np.testing.assert_array_equal(result["pixels/wrist_image"], wrist) + assert result["agent_pos"].shape == (8,) + np.testing.assert_array_almost_equal(result["agent_pos"][:7], joints) + assert result["agent_pos"][7] == gripper[0] finally: _uninstall_robomme_stub() @@ -166,9 +163,9 @@ def test_convert_obs_array_format(): "gripper_state_list": np.zeros(2, dtype=np.float32), } result = env._convert_obs(obs_raw) - assert result["image"].shape == (256, 256, 3) - assert result["wrist_image"].shape == (256, 256, 3) - assert result["state"].shape == (8,) + assert result["pixels/image"].shape == (256, 256, 3) + assert result["pixels/wrist_image"].shape == (256, 256, 3) + assert result["agent_pos"].shape == (8,) finally: _uninstall_robomme_stub()