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>

.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 \

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=<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                              |

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 `<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}'.",

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(
+    features: dict[str, PolicyFeature] = field(default_factory=dict)
-        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_map: dict[str, str] = field(
         default_factory=lambda: {
             ACTION: ACTION,
-            "image": f"{OBS_IMAGES}.image",
+            "pixels/image": f"{OBS_IMAGES}.image",
-            "wrist_image": f"{OBS_IMAGES}.wrist_image",
+            "pixels/wrist_image": f"{OBS_IMAGES}.wrist_image",
-            OBS_STATE: OBS_STATE,
+            "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 {}

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),
+                "pixels/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),
+                "pixels/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),
+                "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,
+            "pixels/image": front_rgb,
-            "wrist_image": wrist_rgb,
+            "pixels/wrist_image": wrist_rgb,
-            "state": state,
+            "agent_pos": state,
         }
 
     def _convert_info(self, info: dict) -> dict:

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
 
 

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["pixels/image"] == f"{OBS_IMAGES}.image"
-    assert cfg.features_map["wrist_image"] == f"{OBS_IMAGES}.wrist_image"
+    assert cfg.features_map["pixels/wrist_image"] == f"{OBS_IMAGES}.wrist_image"
-    assert cfg.features_map[OBS_STATE] == OBS_STATE
+    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
+    """`ee_pose` uses a 7-D action; __post_init__ sets the correct shape."""
-    (the joint_angle default). Switching to ee_pose requires the user to
-    override `features[ACTION]`. This test documents that.
-    """
     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["pixels/image"], front)
-        np.testing.assert_array_equal(result["wrist_image"], wrist)
+        np.testing.assert_array_equal(result["pixels/wrist_image"], wrist)
-        assert result["state"].shape == (8,)
+        assert result["agent_pos"].shape == (8,)
-        np.testing.assert_array_almost_equal(result["state"][:7], joints)
+        np.testing.assert_array_almost_equal(result["agent_pos"][:7], joints)
-        assert result["state"][7] == gripper[0]
+        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["pixels/image"].shape == (256, 256, 3)
-        assert result["wrist_image"].shape == (256, 256, 3)
+        assert result["pixels/wrist_image"].shape == (256, 256, 3)
-        assert result["state"].shape == (8,)
+        assert result["agent_pos"].shape == (8,)
     finally:
         _uninstall_robomme_stub()