feat(envs): add RoboMME benchmark integration

- RoboMME env wrapper with image/wrist_image/state observations
- Docker image with Vulkan, SAPIEN, mani-skill deps
- CI workflow: 1-episode smoke eval with pepijn223/smolvla_robomme
- preprocess_observation: handle image/wrist_image/state keys
- pyproject.toml: robomme extra

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.github/workflows/benchmark_tests.yml

@@ -310,3 +310,91 @@ jobs:
           name: metaworld-metrics
           path: /tmp/metaworld-artifacts/metrics.json
           if-no-files-found: warn
+
+  # ── ROBOMME ───────────────────────────────────────────────────────────────
+  # Isolated image: lerobot[robomme] only (mani-skill, SAPIEN, Vulkan chain)
+  robomme-integration-test:
+    name: RoboMME — build image + 1-episode eval
+    runs-on:
+      group: aws-g6-4xlarge-plus
+    env:
+      HF_USER_TOKEN: ${{ secrets.LEROBOT_HF_USER }}
+
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+        with:
+          persist-credentials: false
+          lfs: true
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3 # zizmor: ignore[unpinned-uses]
+        with:
+          cache-binary: false
+
+      - name: Build RoboMME benchmark image
+        uses: docker/build-push-action@v6 # zizmor: ignore[unpinned-uses]
+        with:
+          context: .
+          file: docker/Dockerfile.benchmark.robomme
+          push: false
+          load: true
+          tags: lerobot-benchmark-robomme:ci
+
+      - name: Run RoboMME smoke eval (1 episode)
+        run: |
+          docker run --name robomme-eval --gpus all \
+            --shm-size=4g \
+            -e HF_HOME=/tmp/hf \
+            -e HF_USER_TOKEN="${HF_USER_TOKEN}" \
+            -e HF_HUB_DOWNLOAD_TIMEOUT=300 \
+            lerobot-benchmark-robomme:ci \
+            bash -c "
+              hf auth login --token \"\$HF_USER_TOKEN\" --add-to-git-credential 2>/dev/null || true
+              lerobot-eval \
+                --policy.path=pepijn223/smolvla_robomme \
+                --env.type=robomme \
+                --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
+                --env.dataset_split=test \
+                --eval.batch_size=1 \
+                --eval.n_episodes=1 \
+                --eval.use_async_envs=false \
+                --policy.device=cuda \
+                '--rename_map={\"observation.images.image\": \"observation.images.camera1\", \"observation.images.wrist_image\": \"observation.images.camera2\"}' \
+                --policy.empty_cameras=3 \
+                --output_dir=/tmp/eval-artifacts
+              python scripts/ci/extract_task_descriptions.py \
+                --env robomme --task PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
+                --output /tmp/eval-artifacts/task_descriptions.json
+            "
+
+      - name: Copy RoboMME artifacts from container
+        if: always()
+        run: |
+          mkdir -p /tmp/robomme-artifacts
+          docker cp robomme-eval:/tmp/eval-artifacts/. /tmp/robomme-artifacts/ 2>/dev/null || true
+          docker rm -f robomme-eval || true
+
+      - name: Parse RoboMME eval metrics
+        if: always()
+        run: |
+          python3 scripts/ci/parse_eval_metrics.py \
+            --artifacts-dir /tmp/robomme-artifacts \
+            --env robomme \
+            --task PickXtimes \
+            --policy pepijn223/smolvla_robomme
+
+      - name: Upload RoboMME rollout video
+        if: always()
+        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
+        with:
+          name: robomme-rollout-video
+          path: /tmp/robomme-artifacts/videos/
+          if-no-files-found: warn
+
+      - name: Upload RoboMME eval metrics
+        if: always()
+        uses: actions/upload-artifact@v4 # zizmor: ignore[unpinned-uses]
+        with:
+          name: robomme-metrics
+          path: /tmp/robomme-artifacts/metrics.json
+          if-no-files-found: warn

docker/Dockerfile.benchmark.robomme

@@ -0,0 +1,102 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Isolated benchmark image for RoboMME integration tests.
+# Installs only lerobot[robomme] so its dep tree (mani-skill, SAPIEN, robomme)
+# cannot conflict with other benchmarks.
+#
+# Key differences vs libero/metaworld:
+#   - Vulkan ICD loader required (ManiSkill/SAPIEN renderer)
+#   - uv sync without --locked (robomme is a git dep, not in uv.lock)
+#   - gymnasium==0.29.1 + numpy==1.26.4 overridden post-install (mani-skill pins)
+#
+# Build:  docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-benchmark-robomme .
+# Run:    docker run --gpus all --rm lerobot-benchmark-robomme lerobot-eval ...
+
+ARG CUDA_VERSION=12.4.1
+ARG OS_VERSION=22.04
+FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu${OS_VERSION}
+
+ARG PYTHON_VERSION=3.12
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PATH=/lerobot/.venv/bin:$PATH \
+    CUDA_VISIBLE_DEVICES=0 \
+    DEVICE=cuda \
+    # NVIDIA Container Toolkit: expose all driver capabilities (includes Vulkan)
+    NVIDIA_DRIVER_CAPABILITIES=all \
+    # ManiSkill uses Vulkan; set ICD path for headless GPU rendering
+    VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json
+
+# System deps — same base as Dockerfile.internal, plus Vulkan loader for ManiSkill/SAPIEN
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    software-properties-common build-essential git curl \
+    libglib2.0-0 libgl1-mesa-glx libegl1-mesa ffmpeg \
+    libusb-1.0-0-dev speech-dispatcher libgeos-dev portaudio19-dev \
+    cmake pkg-config ninja-build \
+    # Vulkan: required by ManiSkill/SAPIEN renderer
+    libvulkan1 \
+    libvulkan-dev \
+    mesa-vulkan-drivers \
+    && add-apt-repository -y ppa:deadsnakes/ppa \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends \
+       python${PYTHON_VERSION} \
+       python${PYTHON_VERSION}-venv \
+       python${PYTHON_VERSION}-dev \
+    && curl -LsSf https://astral.sh/uv/0.8.0/install.sh | sh \
+    && mv /root/.local/bin/uv /usr/local/bin/uv \
+    && mkdir -p /usr/share/vulkan/icd.d \
+    && echo '{"file_format_version":"1.0.0","ICD":{"library_path":"libGLX_nvidia.so.0","api_version":"1.3.0"}}' \
+       > /usr/share/vulkan/icd.d/nvidia_icd.json \
+    && useradd --create-home --shell /bin/bash user_lerobot \
+    && usermod -aG sudo user_lerobot \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /lerobot
+RUN chown -R user_lerobot:user_lerobot /lerobot
+USER user_lerobot
+
+ENV HOME=/home/user_lerobot \
+    HF_HOME=/home/user_lerobot/.cache/huggingface \
+    HF_LEROBOT_HOME=/home/user_lerobot/.cache/huggingface/lerobot \
+    TORCH_HOME=/home/user_lerobot/.cache/torch \
+    TRITON_CACHE_DIR=/home/user_lerobot/.cache/triton
+
+RUN uv venv --python python${PYTHON_VERSION}
+
+# Install only lerobot[robomme] — completely isolated from libero/metaworld dep trees.
+# Note: --locked is intentionally omitted because robomme is a git dep and is not
+# pinned in uv.lock; the install is still reproducible via the pinned git URL.
+COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
+COPY --chown=user_lerobot:user_lerobot src/ src/
+
+# Step 1: Install lerobot base (without robomme) — this gets gymnasium>=1.1.1
+# and numpy>=2.0.0 as lerobot requires.
+RUN uv pip install --no-cache -e "."
+
+# Step 2: Force-install robomme with overridden gymnasium and numpy versions.
+# mani-skill (robomme dep) hard-pins gymnasium==0.29.1 and numpy<2.0.0 which
+# conflict with lerobot's defaults. Both overrides are safe at runtime:
+#   - gymnasium 0.29.x has the same 5-tuple step() API as 1.x (since gym 0.26)
+#   - numpy 1.26.4 is API-compatible with lerobot's actual usage
+RUN printf 'gymnasium==0.29.1\nnumpy==1.26.4\n' > /tmp/robomme_override.txt \
+    && uv pip install --no-cache --override /tmp/robomme_override.txt -e ".[robomme,smolvla,av-dep]" \
+    && python -c "import robomme; print('robomme import OK')"
+
+RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
+
+COPY --chown=user_lerobot:user_lerobot . .
+
+CMD ["/bin/bash"]

docs/source/robomme.mdx

@@ -0,0 +1,119 @@
+# RoboMME Benchmark
+
+[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.
+
+- **16 tasks** across 4 memory-skill suites
+- **1,600 training demos** (100 per task, 50 val, 50 test)
+- **Dataset**: [`pepijn223/robomme_data_lerobot`](https://huggingface.co/datasets/pepijn223/robomme_data_lerobot) — LeRobot v3.0, 768K frames at 10 fps
+- **Simulator**: ManiSkill / SAPIEN, Panda arm, Linux only
+
+> **Related work**: The initial integration of RoboMME into LeRobot was prototyped in
+> [`feat/robomme-integration`](https://github.com/huggingface/lerobot/tree/feat/robomme-integration).
+> This PR rebases that work onto `feat/async-vector-env`, adopts the `create_envs()` dispatch
+> pattern, and adds the `LazyVectorEnv` helper for large task-count benchmarks.
+
+## Tasks
+
+| Suite                             | Tasks                                                         |
+| --------------------------------- | ------------------------------------------------------------- |
+| **Counting** (temporal memory)    | BinFill, PickXtimes, SwingXtimes, StopCube                    |
+| **Permanence** (spatial memory)   | VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap  |
+| **Reference** (object memory)     | PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder |
+| **Imitation** (procedural memory) | MoveCube, InsertPeg, PatternLock, RouteStick                  |
+
+## Installation
+
+> RoboMME requires **Linux** (ManiSkill/SAPIEN uses Vulkan rendering). Docker is recommended to isolate dependency conflicts.
+
+### Native (Linux)
+
+```bash
+pip install 'lerobot[robomme]'
+```
+
+> **Dependency note**: `robomme` pulls in `mani-skill` which pins `gymnasium==0.29.1` and `numpy<2.0.0`. These conflict with lerobot's default pins. On Linux you can pin both in your env; see the Docker approach below to avoid conflicts entirely.
+
+### Docker (recommended)
+
+```bash
+# Build base image first (from repo root)
+docker build -f docker/Dockerfile.eval-base -t lerobot-eval-base .
+
+# Build RoboMME eval image (applies gymnasium + numpy pin overrides)
+docker build -f docker/Dockerfile.eval-robomme -t lerobot-robomme .
+```
+
+The `Dockerfile.eval-robomme` overrides `gymnasium==0.29.1` and `numpy==1.26.4` after lerobot's install. Both versions are runtime-safe for lerobot's actual API usage.
+
+## Running Evaluation
+
+```bash
+lerobot-eval \
+    --policy.path=<your_policy_repo> \
+    --env.type=robomme \
+    --env.task=PickXtimes \
+    --env.dataset_split=test \
+    --eval.batch_size=1 \
+    --eval.n_episodes=50
+```
+
+### Key CLI options for `env.type=robomme`
+
+| Option               | Default       | Description                                        |
+| -------------------- | ------------- | -------------------------------------------------- |
+| `env.task`           | `PickXtimes`  | Any of the 16 task names above                     |
+| `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                              |
+| `env.task_ids`       | `null`        | List of episode indices to evaluate (null = `[0]`) |
+
+## Dataset
+
+The dataset [`pepijn223/robomme_data_lerobot`](https://huggingface.co/datasets/pepijn223/robomme_data_lerobot) is in **LeRobot v3.0 format** and can be loaded directly:
+
+```python
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+
+dataset = LeRobotDataset("pepijn223/robomme_data_lerobot")
+```
+
+### Dataset features
+
+| Feature            | Shape         | Description                     |
+| ------------------ | ------------- | ------------------------------- |
+| `image`            | (256, 256, 3) | Front camera RGB                |
+| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                |
+| `actions`          | (8,)          | Joint angles + gripper          |
+| `state`            | (8,)          | Joint positions + gripper state |
+| `simple_subgoal`   | str           | High-level language annotation  |
+| `grounded_subgoal` | str           | Grounded language annotation    |
+| `episode_index`    | int           | Episode ID                      |
+| `frame_index`      | int           | Frame within episode            |
+
+### Feature key alignment (training)
+
+The env wrapper exposes `front_rgb` and `wrist_rgb` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `obs_images.front` and `obs_images.wrist` for the policy.
+
+The dataset uses `image` and `wrist_image` as column names. When fine-tuning from the dataset, you may need to rename columns to match your policy's expected input keys:
+
+```python
+dataset = LeRobotDataset("pepijn223/robomme_data_lerobot")
+# rename dataset columns if needed:
+# dataset.hf_dataset = dataset.hf_dataset.rename_column("image", "front_rgb")
+```
+
+## Action Spaces
+
+| Type          | Dim | Description                                               |
+| ------------- | --- | --------------------------------------------------------- |
+| `joint_angle` | 8   | 7 joint angles + 1 gripper (−1 closed, +1 open, absolute) |
+| `ee_pose`     | 7   | xyz + roll/pitch/yaw + gripper                            |
+
+Set via `--env.action_space=joint_angle` (default) or `--env.action_space=ee_pose`.
+
+## Platform Notes
+
+- **Linux only**: ManiSkill requires SAPIEN/Vulkan. macOS and Windows are not supported.
+- **GPU recommended**: Rendering is CPU-capable but slow; CUDA + Vulkan gives full speed.
+- **gymnasium / numpy conflict**: See installation note above. Docker image handles this automatically.
+- **ManiSkill fork**: `robomme` depends on a specific ManiSkill fork (`YinpeiDai/ManiSkill`), pulled in automatically via the `robomme` package.

pyproject.toml

@@ -206,6 +206,9 @@ aloha = ["lerobot[dataset]", "gym-aloha>=0.1.2,<0.2.0", "lerobot[scipy-dep]"]
 pusht = ["lerobot[dataset]", "gym-pusht>=0.1.5,<0.2.0", "pymunk>=6.6.0,<7.0.0"] # TODO: Fix pymunk version in gym-pusht instead
 libero = ["lerobot[dataset]", "lerobot[transformers-dep]", "hf-libero>=0.1.3,<0.2.0; sys_platform == 'linux'", "lerobot[scipy-dep]"]
 metaworld = ["lerobot[dataset]", "metaworld==3.0.0", "lerobot[scipy-dep]"]
+robomme = [
+    "robomme @ git+https://github.com/RoboMME/robomme_benchmark.git@main ; sys_platform == 'linux'",
+]
 
 # All
 all = [

scripts/ci/extract_task_descriptions.py

@@ -57,6 +57,36 @@ def _metaworld_descriptions(task_name: str) -> dict[str, str]:
     return {f"{task_name}_0": label}
 
 
+_ROBOMME_DESCRIPTIONS = {
+    "BinFill": "Fill the target bin with the correct number of cubes",
+    "PickXtimes": "Pick the indicated cube the specified number of times",
+    "SwingXtimes": "Swing the object the specified number of times",
+    "StopCube": "Grasp and stop the moving cube",
+    "VideoUnmask": "Pick the cube shown in the reference video",
+    "VideoUnmaskSwap": "Pick the cube matching the reference video after a swap",
+    "ButtonUnmask": "Press the button indicated by the reference",
+    "ButtonUnmaskSwap": "Press the correct button after objects are swapped",
+    "PickHighlight": "Pick the highlighted cube",
+    "VideoRepick": "Repick the cube shown in the reference video",
+    "VideoPlaceButton": "Place the cube on the button shown in the video",
+    "VideoPlaceOrder": "Place cubes in the order shown in the video",
+    "MoveCube": "Move the cube to the target location",
+    "InsertPeg": "Insert the peg into the target hole",
+    "PatternLock": "Unlock the pattern by pressing buttons in sequence",
+    "RouteStick": "Route the stick through the required waypoints",
+}
+
+
+def _robomme_descriptions(task_names: str) -> dict[str, str]:
+    """Return descriptions for each requested RoboMME task. Keys match the
+    video filename pattern `<task>_<task_id>` used by the eval script."""
+    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
+    return out
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--env", required=True, help="Environment family (libero, metaworld, ...)")
@@ -70,6 +100,8 @@ def main() -> int:
             descriptions = _libero_descriptions(args.task)
         elif args.env == "metaworld":
             descriptions = _metaworld_descriptions(args.task)
+        elif args.env == "robomme":
+            descriptions = _robomme_descriptions(args.task)
         else:
             print(
                 f"[extract_task_descriptions] No description extractor for env '{args.env}'.",

src/lerobot/envs/configs.py

@@ -574,3 +574,57 @@ class IsaaclabArenaEnv(HubEnvConfig):
             ),
             PolicyProcessorPipeline(steps=[]),
         )
+
+
+@EnvConfig.register_subclass("robomme")
+@dataclass
+class RoboMMEEnv(EnvConfig):
+    """RoboMME memory-augmented manipulation benchmark (ManiSkill/SAPIEN).
+
+    16 tasks across 4 suites: Counting, Permanence, Reference, Imitation.
+    Dataset: pepijn223/robomme_data_lerobot (LeRobot v3.0, 1,600 episodes).
+    Benchmark: https://github.com/RoboMME/robomme_benchmark
+
+    Requires: pip install 'lerobot[robomme]' (Linux only).
+    """
+
+    task: str = "PickXtimes"
+    fps: int = 10
+    episode_length: int = 300
+    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_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,
+        }
+    )
+
+    @property
+    def gym_kwargs(self) -> dict:
+        return {}
+
+    def create_envs(self, n_envs: int, use_async_envs: bool = True):
+        from lerobot.envs.robomme import create_robomme_envs
+
+        env_cls = _make_vec_env_cls(use_async_envs, n_envs)
+        return create_robomme_envs(
+            task=self.task,
+            n_envs=n_envs,
+            action_space_type=self.action_space,
+            dataset=self.dataset_split,
+            episode_length=self.episode_length,
+            task_ids=self.task_ids,
+            env_cls=env_cls,
+        )

src/lerobot/envs/robomme.py

@@ -0,0 +1,224 @@
+"""RoboMME environment wrapper for LeRobot evaluation.
+
+Wraps the RoboMME ``BenchmarkEnvBuilder`` into a Gymnasium-compatible
+``VectorEnv`` suitable for ``lerobot_eval``.
+
+RoboMME tasks:
+  Counting:    BinFill, PickXtimes, SwingXtimes, StopCube
+  Permanence:  VideoUnmask, VideoUnmaskSwap, ButtonUnmask, ButtonUnmaskSwap
+  Reference:   PickHighlight, VideoRepick, VideoPlaceButton, VideoPlaceOrder
+  Imitation:   MoveCube, InsertPeg, PatternLock, RouteStick
+
+Dataset: pepijn223/robomme_data_lerobot (LeRobot v3.0, 1,600 episodes)
+Install: pip install 'lerobot[robomme]'  (Linux only)
+Benchmark: https://github.com/RoboMME/robomme_benchmark
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import Any
+
+import gymnasium as gym
+import numpy as np
+from gymnasium import spaces
+
+ROBOMME_TASKS = [
+    "BinFill",
+    "PickXtimes",
+    "SwingXtimes",
+    "StopCube",
+    "VideoUnmask",
+    "VideoUnmaskSwap",
+    "ButtonUnmask",
+    "ButtonUnmaskSwap",
+    "PickHighlight",
+    "VideoRepick",
+    "VideoPlaceButton",
+    "VideoPlaceOrder",
+    "MoveCube",
+    "InsertPeg",
+    "PatternLock",
+    "RouteStick",
+]
+
+
+class RoboMMEGymEnv(gym.Env):
+    """Thin Gymnasium wrapper around a single RoboMME episode env."""
+
+    metadata = {"render_modes": ["rgb_array"], "render_fps": 10}
+
+    def __init__(
+        self,
+        task: str = "PickXtimes",
+        action_space_type: str = "joint_angle",
+        dataset: str = "test",
+        episode_idx: int = 0,
+        max_steps: int = 300,
+    ):
+        super().__init__()
+        from robomme.env_record_wrapper import BenchmarkEnvBuilder
+
+        self._task = task
+        self._action_space_type = action_space_type
+        self._dataset = dataset
+        self._episode_idx = episode_idx
+        self._max_steps = max_steps
+        self._max_episode_steps = max_steps
+
+        self._builder = BenchmarkEnvBuilder(
+            env_id=task,
+            dataset=dataset,
+            action_space=action_space_type,
+            gui_render=False,
+            max_steps=max_steps,
+        )
+        self._env = None
+        self._last_raw_obs: dict | None = None
+
+        action_dim = 8 if action_space_type == "joint_angle" else 7
+        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),
+            }
+        )
+
+    def reset(self, *, seed=None, options=None):
+        super().reset(seed=seed)
+        self._env = self._builder.make_env_for_episode(
+            episode_idx=self._episode_idx,
+            max_steps=self._max_steps,
+        )
+        obs, info = self._env.reset()
+        self._last_raw_obs = obs
+        return self._convert_obs(obs), self._convert_info(info)
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self._env.step(action)
+        self._last_raw_obs = obs
+
+        terminated_bool = bool(terminated.item()) if hasattr(terminated, "item") else bool(terminated)
+        truncated_bool = bool(truncated.item()) if hasattr(truncated, "item") else bool(truncated)
+
+        status = info.get("status", "ongoing")
+        is_success = status == "success"
+        conv_info = self._convert_info(info)
+        conv_info["is_success"] = is_success
+
+        return self._convert_obs(obs), float(reward), terminated_bool, truncated_bool, conv_info
+
+    def render(self) -> np.ndarray | None:
+        """Return the front camera image from the last observation for video recording."""
+        if self._last_raw_obs is None:
+            return np.zeros((256, 256, 3), dtype=np.uint8)
+        front = self._last_raw_obs.get("front_rgb_list")
+        if front is None:
+            return np.zeros((256, 256, 3), dtype=np.uint8)
+        frame = front[-1] if isinstance(front, list) else front
+        return np.asarray(frame, dtype=np.uint8)
+
+    def _convert_obs(self, obs: dict) -> dict:
+        front_rgb = (
+            obs["front_rgb_list"][-1] if isinstance(obs["front_rgb_list"], list) else obs["front_rgb_list"]
+        )
+        wrist_rgb = (
+            obs["wrist_rgb_list"][-1] if isinstance(obs["wrist_rgb_list"], list) else obs["wrist_rgb_list"]
+        )
+        joint_state = (
+            obs["joint_state_list"][-1]
+            if isinstance(obs["joint_state_list"], list)
+            else obs["joint_state_list"]
+        )
+        gripper_state = (
+            obs["gripper_state_list"][-1]
+            if isinstance(obs["gripper_state_list"], list)
+            else obs["gripper_state_list"]
+        )
+
+        front_rgb = np.asarray(front_rgb, dtype=np.uint8)
+        wrist_rgb = np.asarray(wrist_rgb, dtype=np.uint8)
+        joint = np.asarray(joint_state, dtype=np.float32).flatten()[:7]
+        gripper = np.asarray(gripper_state, dtype=np.float32).flatten()[:1]
+        state = np.concatenate([joint, gripper])
+
+        return {
+            "image": front_rgb,
+            "wrist_image": wrist_rgb,
+            "state": state,
+        }
+
+    def _convert_info(self, info: dict) -> dict:
+        return {
+            "status": info.get("status", "ongoing"),
+            "task_goal": info.get("task_goal", ""),
+        }
+
+
+def _make_env_fns(
+    *,
+    task: str,
+    n_envs: int,
+    action_space_type: str,
+    dataset: str,
+    episode_length: int,
+    task_id: int,
+) -> list[Callable[[], RoboMMEGymEnv]]:
+    """Build n_envs factory callables for one RoboMME task id."""
+
+    def _make_one(episode_index: int) -> RoboMMEGymEnv:
+        return RoboMMEGymEnv(
+            task=task,
+            action_space_type=action_space_type,
+            dataset=dataset,
+            episode_idx=episode_index,
+            max_steps=episode_length,
+        )
+
+    return [partial(_make_one, task_id + i) for i in range(n_envs)]
+
+
+def create_robomme_envs(
+    task: str,
+    n_envs: int = 1,
+    action_space_type: str = "joint_angle",
+    dataset: str = "test",
+    episode_length: int = 300,
+    task_ids: list[int] | None = None,
+    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
+) -> dict[str, dict[int, gym.vector.VectorEnv]]:
+    """Create vectorized RoboMME environments for evaluation.
+
+    `task` may be a single RoboMME task name (e.g. "PickXtimes") or a
+    comma-separated list (e.g. "PickXtimes,BinFill,StopCube"). Each task
+    becomes its own suite in the returned mapping.
+
+    Returns {suite_name: {task_id: VectorEnv}} matching lerobot's expected format.
+    """
+    if env_cls is None or not callable(env_cls):
+        raise ValueError("env_cls must be a callable that wraps a list of env factory callables.")
+    if not isinstance(n_envs, int) or n_envs <= 0:
+        raise ValueError(f"n_envs must be a positive int; got {n_envs}.")
+
+    if task_ids is None:
+        task_ids = [0]
+
+    task_names = [t.strip() for t in task.split(",") if t.strip()]
+    out: dict[str, dict[int, gym.vector.VectorEnv]] = {}
+    for task_name in task_names:
+        envs_by_task: dict[int, gym.vector.VectorEnv] = {}
+        for task_id in task_ids:
+            fns = _make_env_fns(
+                task=task_name,
+                n_envs=n_envs,
+                action_space_type=action_space_type,
+                dataset=dataset,
+                episode_length=episode_length,
+                task_id=task_id,
+            )
+            envs_by_task[task_id] = env_cls(fns)
+        out[task_name] = envs_by_task
+    return out

src/lerobot/envs/utils.py

@@ -107,6 +107,27 @@ 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

@@ -0,0 +1,284 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for the RoboMME env wrapper and config.
+
+RoboMME requires Linux + ManiSkill (Vulkan/SAPIEN), so all tests that
+instantiate the real env mock the ``robomme`` package.  Tests that only
+exercise pure-Python logic (config defaults, obs conversion, lazy env)
+run without any mocking.
+"""
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_robomme_stub():
+    """Return a minimal stub for the ``robomme`` package."""
+    stub = ModuleType("robomme")
+    wrapper_stub = ModuleType("robomme.env_record_wrapper")
+
+    class FakeBuilder:
+        def __init__(self, **kwargs):
+            pass
+
+        def make_env_for_episode(self, episode_idx: int, max_steps: int):
+            env = MagicMock()
+            obs = {
+                "front_rgb_list": [np.zeros((256, 256, 3), dtype=np.uint8)],
+                "wrist_rgb_list": [np.zeros((256, 256, 3), dtype=np.uint8)],
+                "joint_state_list": [np.zeros(7, dtype=np.float32)],
+                "gripper_state_list": [np.zeros(2, dtype=np.float32)],
+            }
+            env.reset.return_value = (obs, {"status": "ongoing", "task_goal": "pick the cube"})
+            env.step.return_value = (obs, 0.0, False, False, {"status": "ongoing", "task_goal": ""})
+            return env
+
+    wrapper_stub.BenchmarkEnvBuilder = FakeBuilder
+    stub.env_record_wrapper = wrapper_stub
+    return stub, wrapper_stub
+
+
+# ---------------------------------------------------------------------------
+# Config tests (no sim required)
+# ---------------------------------------------------------------------------
+
+
+def test_robomme_env_config_defaults():
+    from lerobot.envs.configs import RoboMMEEnv
+
+    cfg = RoboMMEEnv()
+    assert cfg.task == "PickXtimes"
+    assert cfg.fps == 10
+    assert cfg.episode_length == 300
+    assert cfg.action_space == "joint_angle"
+    assert cfg.dataset_split == "test"
+    assert cfg.task_ids is None
+
+
+def test_robomme_env_config_type():
+    from lerobot.envs.configs import RoboMMEEnv
+
+    cfg = RoboMMEEnv()
+    assert cfg.type == "robomme"
+
+
+def test_robomme_features_map():
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
+
+    cfg = RoboMMEEnv()
+    assert cfg.features_map[ACTION] == ACTION
+    assert cfg.features_map["front_rgb"] == f"{OBS_IMAGES}.front"
+    assert cfg.features_map["wrist_rgb"] == f"{OBS_IMAGES}.wrist"
+    assert cfg.features_map[OBS_STATE] == OBS_STATE
+
+
+def test_robomme_features_action_dim_joint_angle():
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION
+
+    cfg = RoboMMEEnv(action_space="joint_angle")
+    assert cfg.features[ACTION].shape == (8,)
+
+
+def test_robomme_features_action_dim_ee_pose():
+    """ee_pose action space uses 7-D; config declares 8-D default (joint_angle).
+
+    Users switching to ee_pose must override the features dict manually or
+    the env wrapper will return 7-D actions while the config claims 8-D.
+    This test documents the current behaviour so it is explicit.
+    """
+    from lerobot.envs.configs import RoboMMEEnv
+    from lerobot.utils.constants import ACTION
+
+    cfg = RoboMMEEnv(action_space="ee_pose")
+    # Default features still say 8-D — ee_pose override is a user responsibility.
+    assert cfg.features[ACTION].shape == (8,)
+
+
+# ---------------------------------------------------------------------------
+# Obs conversion (pure Python, no sim)
+# ---------------------------------------------------------------------------
+
+
+def test_convert_obs_list_format():
+    """_convert_obs must take the last element from list-format obs fields."""
+    stub, wrapper_stub = _make_robomme_stub()
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+    from lerobot.envs.robomme import RoboMMEGymEnv
+
+    env = RoboMMEGymEnv.__new__(RoboMMEGymEnv)
+
+    front = np.full((256, 256, 3), 42, dtype=np.uint8)
+    wrist = np.full((256, 256, 3), 7, dtype=np.uint8)
+    joints = np.arange(7, dtype=np.float32)
+    gripper = np.array([0.5, 0.5], dtype=np.float32)
+
+    obs_raw = {
+        "front_rgb_list": [np.zeros_like(front), front],
+        "wrist_rgb_list": [np.zeros_like(wrist), wrist],
+        "joint_state_list": [np.zeros(7, dtype=np.float32), joints],
+        "gripper_state_list": [np.zeros(2, dtype=np.float32), gripper],
+    }
+
+    result = env._convert_obs(obs_raw)
+
+    np.testing.assert_array_equal(result["front_rgb"], front)
+    np.testing.assert_array_equal(result["wrist_rgb"], wrist)
+    assert result["state"].shape == (8,)
+    np.testing.assert_array_almost_equal(result["state"][:7], joints)
+    assert result["state"][7] == pytest.approx(gripper[0])
+
+    # cleanup
+    del sys.modules["robomme"]
+    del sys.modules["robomme.env_record_wrapper"]
+
+
+def test_convert_obs_array_format():
+    """_convert_obs must also handle non-list (direct array) obs."""
+    stub, wrapper_stub = _make_robomme_stub()
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+    from lerobot.envs.robomme import RoboMMEGymEnv
+
+    env = RoboMMEGymEnv.__new__(RoboMMEGymEnv)
+
+    front = np.zeros((256, 256, 3), dtype=np.uint8)
+    obs_raw = {
+        "front_rgb_list": front,
+        "wrist_rgb_list": front,
+        "joint_state_list": np.zeros(7, dtype=np.float32),
+        "gripper_state_list": np.zeros(2, dtype=np.float32),
+    }
+    result = env._convert_obs(obs_raw)
+    assert result["front_rgb"].shape == (256, 256, 3)
+
+    del sys.modules["robomme"]
+    del sys.modules["robomme.env_record_wrapper"]
+
+
+# ---------------------------------------------------------------------------
+# create_robomme_envs (mocked sim)
+# ---------------------------------------------------------------------------
+
+
+def test_create_robomme_envs_returns_correct_structure():
+    stub, wrapper_stub = _make_robomme_stub()
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+    from lerobot.envs.robomme import create_robomme_envs
+
+    env_cls = MagicMock(return_value=MagicMock())
+    result = create_robomme_envs(
+        task="PickXtimes",
+        n_envs=1,
+        task_ids=[0, 1],
+        env_cls=env_cls,
+    )
+
+    assert "robomme" in result
+    assert 0 in result["robomme"]
+    assert 1 in result["robomme"]
+    assert env_cls.call_count == 2
+
+    del sys.modules["robomme"]
+    del sys.modules["robomme.env_record_wrapper"]
+
+
+def test_create_robomme_envs_raises_on_invalid_env_cls():
+    stub, wrapper_stub = _make_robomme_stub()
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+    from lerobot.envs.robomme import create_robomme_envs
+
+    with pytest.raises(ValueError, match="env_cls must be a callable"):
+        create_robomme_envs(task="PickXtimes", n_envs=1, env_cls=None)
+
+    del sys.modules["robomme"]
+    del sys.modules["robomme.env_record_wrapper"]
+
+
+# ---------------------------------------------------------------------------
+# LazyVectorEnv
+# ---------------------------------------------------------------------------
+
+
+def test_lazy_vec_env_used_when_task_ids_gt_50():
+    """create_robomme_envs must use LazyVectorEnv when len(task_ids) > 50."""
+    stub, wrapper_stub = _make_robomme_stub()
+    sys.modules["robomme"] = stub
+    sys.modules["robomme.env_record_wrapper"] = wrapper_stub
+
+    from lerobot.envs.lazy_vec_env import LazyVectorEnv
+    from lerobot.envs.robomme import create_robomme_envs
+
+    env_cls = MagicMock(return_value=MagicMock())
+    task_ids = list(range(51))
+    result = create_robomme_envs(task="PickXtimes", n_envs=1, task_ids=task_ids, env_cls=env_cls)
+
+    for tid in task_ids:
+        assert isinstance(result["robomme"][tid], LazyVectorEnv)
+    # env_cls must NOT have been called yet (lazy)
+    env_cls.assert_not_called()
+
+    del sys.modules["robomme"]
+    del sys.modules["robomme.env_record_wrapper"]
+
+
+def test_lazy_vec_env_materializes_on_access():
+    from lerobot.envs.lazy_vec_env import LazyVectorEnv
+
+    inner = MagicMock()
+    inner.reset.return_value = ({"obs": 1}, {})
+    env_cls = MagicMock(return_value=inner)
+    factory_fns = [lambda: MagicMock()]
+
+    lazy = LazyVectorEnv(env_cls, factory_fns)
+    env_cls.assert_not_called()
+
+    # accessing reset triggers materialization
+    lazy.reset()
+    env_cls.assert_called_once_with(factory_fns)
+    inner.reset.assert_called_once()
+
+
+def test_lazy_vec_env_close_clears_env():
+    from lerobot.envs.lazy_vec_env import LazyVectorEnv
+
+    inner = MagicMock()
+    env_cls = MagicMock(return_value=inner)
+    lazy = LazyVectorEnv(env_cls, [lambda: MagicMock()])
+
+    lazy.reset()  # materialize
+    lazy.close()
+    inner.close.assert_called_once()
+
+    # env_cls should be called again on next access after close
+    lazy.reset()
+    assert env_cls.call_count == 2