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>

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