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) <noreply@anthropic.com>
2026-07-26 11:16:00 +00:00 · 2026-04-16 18:32:24 +02:00
parent cf1e4833fa
commit be70315583
7 changed files with 79 additions and 60 deletions
@@ -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 \
@@ -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=<your_policy_repo> \
    --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=<your_policy_repo> \
+    --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                              |
@@ -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 `<task>_<task_id>` 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}'.",
@@ -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 {}
@@ -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:
@@ -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


@@ -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()