admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-26 19:26:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(robomme): realign unit tests with current env API

Browse Source 
The tests were written against an earlier env layout and never updated when
the wrapper was refactored, so CI's fast-test job was failing with:

- KeyError: 'front_rgb' / 'wrist_rgb' — these were renamed to the
  lerobot-canonical 'image' / 'wrist_image' keys (matching the dataset
  columns and preprocess_observation's built-in fallbacks).
- AssertionError: 'robomme' not in result — create_robomme_envs now
  returns {task_name: {task_id: env}}, not {'robomme': {...}}, so
  comma-separated task lists work.
- ModuleNotFoundError: lerobot.envs.lazy_vec_env — LazyVectorEnv was
  removed; create_robomme_envs is straightforward synchronous now.

Rewrite the 7 failing cases against the current API, drop the three
LazyVectorEnv tests, and add a multi-task test so the new comma-separated
task parsing is covered. Stub install/teardown is moved into helpers
(`_install_robomme_stub` / `_uninstall_robomme_stub`) so individual tests
stop repeating six boilerplate lines.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
2026-04-15 10:50:05 +02:00
parent ee37fb9928
commit 4c00e39639
1 changed files with 100 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_robomme_env.py
+100 -153
View File
@@ -13,10 +13,9 @@
# limitations under the License. # limitations under the License.
"""Unit tests for the RoboMME env wrapper and config. """Unit tests for the RoboMME env wrapper and config.
RoboMME requires Linux + ManiSkill (Vulkan/SAPIEN), so all tests that RoboMME requires Linux + ManiSkill (Vulkan/SAPIEN), so tests that touch the
instantiate the real env mock the ``robomme`` package. Tests that only env wrapper mock the ``robomme`` package. Tests that only exercise the
exercise pure-Python logic (config defaults, obs conversion, lazy env) dataclass config run without any mocking.
run without any mocking.
""" """
from __future__ import annotations from __future__ import annotations
@@ -26,15 +25,10 @@ from types import ModuleType
from unittest.mock import MagicMock from unittest.mock import MagicMock
import numpy as np import numpy as np
import pytest
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _make_robomme_stub(): def _install_robomme_stub():
"""Return a minimal stub for the ``robomme`` package.""" """Register a minimal stub for the ``robomme`` package on sys.modules."""
stub = ModuleType("robomme") stub = ModuleType("robomme")
wrapper_stub = ModuleType("robomme.env_record_wrapper") wrapper_stub = ModuleType("robomme.env_record_wrapper")
@@ -56,7 +50,13 @@ def _make_robomme_stub():
wrapper_stub.BenchmarkEnvBuilder = FakeBuilder wrapper_stub.BenchmarkEnvBuilder = FakeBuilder
stub.env_record_wrapper = wrapper_stub stub.env_record_wrapper = wrapper_stub
return stub, wrapper_stub sys.modules["robomme"] = stub
sys.modules["robomme.env_record_wrapper"] = wrapper_stub
def _uninstall_robomme_stub():
sys.modules.pop("robomme", None)
sys.modules.pop("robomme.env_record_wrapper", None)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -89,8 +89,8 @@ def test_robomme_features_map():
cfg = RoboMMEEnv() cfg = RoboMMEEnv()
assert cfg.features_map[ACTION] == ACTION assert cfg.features_map[ACTION] == ACTION
assert cfg.features_map["front_rgb"] == f"{OBS_IMAGES}.front" assert cfg.features_map["image"] == f"{OBS_IMAGES}.image"
assert cfg.features_map["wrist_rgb"] == f"{OBS_IMAGES}.wrist" assert cfg.features_map["wrist_image"] == f"{OBS_IMAGES}.wrist_image"
assert cfg.features_map[OBS_STATE] == OBS_STATE assert cfg.features_map[OBS_STATE] == OBS_STATE
@@ -103,17 +103,14 @@ def test_robomme_features_action_dim_joint_angle():
def test_robomme_features_action_dim_ee_pose(): def test_robomme_features_action_dim_ee_pose():
"""ee_pose action space uses 7-D; config declares 8-D default (joint_angle). """`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
Users switching to ee_pose must override the features dict manually or override `features[ACTION]`. This test documents that.
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.envs.configs import RoboMMEEnv
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
cfg = RoboMMEEnv(action_space="ee_pose") 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,) assert cfg.features[ACTION].shape == (8,)
@@ -123,62 +120,57 @@ def test_robomme_features_action_dim_ee_pose():
def test_convert_obs_list_format(): def test_convert_obs_list_format():
"""_convert_obs must take the last element from list-format obs fields.""" """_convert_obs takes the last element from list-format obs fields and
stub, wrapper_stub = _make_robomme_stub() emits the lerobot-canonical keys (image, wrist_image, state)."""
sys.modules["robomme"] = stub _install_robomme_stub()
sys.modules["robomme.env_record_wrapper"] = wrapper_stub try:
from lerobot.envs.robomme import RoboMMEGymEnv
from lerobot.envs.robomme import RoboMMEGymEnv env = RoboMMEGymEnv.__new__(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)
front = np.full((256, 256, 3), 42, dtype=np.uint8) obs_raw = {
wrist = np.full((256, 256, 3), 7, dtype=np.uint8) "front_rgb_list": [np.zeros_like(front), front],
joints = np.arange(7, dtype=np.float32) "wrist_rgb_list": [np.zeros_like(wrist), wrist],
gripper = np.array([0.5, 0.5], dtype=np.float32) "joint_state_list": [np.zeros(7, dtype=np.float32), joints],
"gripper_state_list": [np.zeros(2, dtype=np.float32), gripper],
}
obs_raw = { result = env._convert_obs(obs_raw)
"front_rgb_list": [np.zeros_like(front), front], np.testing.assert_array_equal(result["image"], front)
"wrist_rgb_list": [np.zeros_like(wrist), wrist], np.testing.assert_array_equal(result["wrist_image"], wrist)
"joint_state_list": [np.zeros(7, dtype=np.float32), joints], assert result["state"].shape == (8,)
"gripper_state_list": [np.zeros(2, dtype=np.float32), gripper], np.testing.assert_array_almost_equal(result["state"][:7], joints)
} assert result["state"][7] == gripper[0]
finally:
result = env._convert_obs(obs_raw) _uninstall_robomme_stub()
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(): def test_convert_obs_array_format():
"""_convert_obs must also handle non-list (direct array) obs.""" """_convert_obs also handles non-list (direct array) obs."""
stub, wrapper_stub = _make_robomme_stub() _install_robomme_stub()
sys.modules["robomme"] = stub try:
sys.modules["robomme.env_record_wrapper"] = wrapper_stub from lerobot.envs.robomme import RoboMMEGymEnv
from lerobot.envs.robomme import RoboMMEGymEnv env = RoboMMEGymEnv.__new__(RoboMMEGymEnv)
env = RoboMMEGymEnv.__new__(RoboMMEGymEnv) front = np.zeros((256, 256, 3), dtype=np.uint8)
obs_raw = {
front = np.zeros((256, 256, 3), dtype=np.uint8) "front_rgb_list": front,
obs_raw = { "wrist_rgb_list": front,
"front_rgb_list": front, "joint_state_list": np.zeros(7, dtype=np.float32),
"wrist_rgb_list": front, "gripper_state_list": np.zeros(2, dtype=np.float32),
"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["image"].shape == (256, 256, 3)
result = env._convert_obs(obs_raw) assert result["wrist_image"].shape == (256, 256, 3)
assert result["front_rgb"].shape == (256, 256, 3) assert result["state"].shape == (8,)
finally:
del sys.modules["robomme"] _uninstall_robomme_stub()
del sys.modules["robomme.env_record_wrapper"]
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -187,98 +179,53 @@ def test_convert_obs_array_format():
def test_create_robomme_envs_returns_correct_structure(): def test_create_robomme_envs_returns_correct_structure():
stub, wrapper_stub = _make_robomme_stub() """Single task -> {task_name: {task_id: VectorEnv}} with one entry per task_id."""
sys.modules["robomme"] = stub _install_robomme_stub()
sys.modules["robomme.env_record_wrapper"] = wrapper_stub try:
from lerobot.envs.robomme import create_robomme_envs
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,
)
env_cls = MagicMock(return_value=MagicMock()) assert "PickXtimes" in result
result = create_robomme_envs( assert 0 in result["PickXtimes"]
task="PickXtimes", assert 1 in result["PickXtimes"]
n_envs=1, assert env_cls.call_count == 2
task_ids=[0, 1], finally:
env_cls=env_cls, _uninstall_robomme_stub()
)
assert "robomme" in result
assert 0 in result["robomme"]
assert 1 in result["robomme"]
assert env_cls.call_count == 2
del sys.modules["robomme"] def test_create_robomme_envs_multi_task():
del sys.modules["robomme.env_record_wrapper"] """Comma-separated task list produces one suite per task."""
_install_robomme_stub()
try:
from lerobot.envs.robomme import create_robomme_envs
env_cls = MagicMock(return_value=MagicMock())
result = create_robomme_envs(
task="PickXtimes,BinFill,StopCube",
n_envs=1,
env_cls=env_cls,
)
assert set(result.keys()) == {"PickXtimes", "BinFill", "StopCube"}
finally:
_uninstall_robomme_stub()
def test_create_robomme_envs_raises_on_invalid_env_cls(): def test_create_robomme_envs_raises_on_invalid_env_cls():
stub, wrapper_stub = _make_robomme_stub() _install_robomme_stub()
sys.modules["robomme"] = stub try:
sys.modules["robomme.env_record_wrapper"] = wrapper_stub import pytest
from lerobot.envs.robomme import create_robomme_envs from lerobot.envs.robomme import create_robomme_envs
with pytest.raises(ValueError, match="env_cls must be a callable"): with pytest.raises(ValueError, match="env_cls must be a callable"):
create_robomme_envs(task="PickXtimes", n_envs=1, env_cls=None) create_robomme_envs(task="PickXtimes", n_envs=1, env_cls=None)
finally:
del sys.modules["robomme"] _uninstall_robomme_stub()
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