mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-23 12:40:08 +00:00
fix(robotwin): defer YAML lookup and realign tests with current API
__init__ was eagerly calling _load_robotwin_setup_kwargs just to read
head_camera_h/w from the YAML. That import (`from envs import CONFIGS_PATH`)
required a real RoboTwin install, so constructing the env — and thus every
test in tests/envs/test_robotwin.py — blew up with ModuleNotFoundError
on fast-tests where RoboTwin isn't installed.
Replace the eager lookup with DEFAULT_CAMERA_H/W constants (240×320, the
D435 dims baked into task_config/demo_clean.yml). reset() still resolves
the full setup_kwargs lazily — that's fine because reset() is only
called inside the benchmark Docker image where RoboTwin is present.
Also resync the test file with the current env API:
- mock get_obs() as the real nested {"observation": {cam: {"rgb": …}},
"joint_action": {"vector": …}} shape
- patch both _load_robotwin_task and _load_robotwin_setup_kwargs
(_patch_load → _patch_runtime)
- drop `front_camera` / `left_wrist` from assertions — aloha-agilex
exposes head_camera + left_camera + right_camera, not those
- black-frame test now uses left_camera as the missing camera
- setup_demo call check loosened to the caller-provided seed/is_test
bits (full kwargs include the YAML-derived blob)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -40,6 +40,9 @@ ACTION_DIM = 14 # 7 DOF × 2 arms
|
|||||||
ACTION_LOW = -1.0
|
ACTION_LOW = -1.0
|
||||||
ACTION_HIGH = 1.0
|
ACTION_HIGH = 1.0
|
||||||
DEFAULT_EPISODE_LENGTH = 300
|
DEFAULT_EPISODE_LENGTH = 300
|
||||||
|
# D435 dims from task_config/_camera_config.yml (what demo_clean.yml selects).
|
||||||
|
DEFAULT_CAMERA_H = 240
|
||||||
|
DEFAULT_CAMERA_W = 320
|
||||||
|
|
||||||
# Complete task list from RoboTwin 2.0 (60 tasks, as listed on the leaderboard).
|
# Complete task list from RoboTwin 2.0 (60 tasks, as listed on the leaderboard).
|
||||||
ROBOTWIN_TASKS: tuple[str, ...] = (
|
ROBOTWIN_TASKS: tuple[str, ...] = (
|
||||||
@@ -240,12 +243,11 @@ class RoboTwinEnv(gym.Env):
|
|||||||
self.episode_index = episode_index
|
self.episode_index = episode_index
|
||||||
self._reset_stride = n_envs
|
self._reset_stride = n_envs
|
||||||
self.camera_names = list(camera_names)
|
self.camera_names = list(camera_names)
|
||||||
# Match the declared observation_space to the camera resolution in
|
# Default to D435 dims (the camera type baked into task_config/demo_clean.yml).
|
||||||
# RoboTwin's YAML — otherwise gym.vector's concatenate pre-allocates
|
# The YAML-driven lookup is deferred to reset() so construction doesn't
|
||||||
# a buffer of the wrong shape and np.stack raises ValueError.
|
# import RoboTwin's `envs` module — fast-tests run without RoboTwin installed.
|
||||||
setup_kwargs = _load_robotwin_setup_kwargs(task_name)
|
self.observation_height = observation_height or DEFAULT_CAMERA_H
|
||||||
self.observation_height = observation_height or setup_kwargs["head_camera_h"]
|
self.observation_width = observation_width or DEFAULT_CAMERA_W
|
||||||
self.observation_width = observation_width or setup_kwargs["head_camera_w"]
|
|
||||||
self.episode_length = episode_length
|
self.episode_length = episode_length
|
||||||
self._max_episode_steps = episode_length # lerobot_eval.rollout reads this
|
self._max_episode_steps = episode_length # lerobot_eval.rollout reads this
|
||||||
self.render_mode = render_mode
|
self.render_mode = render_mode
|
||||||
|
|||||||
+68
-37
@@ -15,12 +15,14 @@
|
|||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
"""Unit tests for the RoboTwin 2.0 Gymnasium wrapper.
|
"""Unit tests for the RoboTwin 2.0 Gymnasium wrapper.
|
||||||
|
|
||||||
These tests mock out the SAPIEN-based RoboTwin runtime so they run without
|
These tests mock out the SAPIEN-based RoboTwin runtime (task modules +
|
||||||
the full RoboTwin installation (SAPIEN, CuRobo, mplib, etc.).
|
YAML config loader) so they run without the full RoboTwin installation
|
||||||
|
(SAPIEN, CuRobo, mplib, asset downloads, etc.).
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from contextlib import contextmanager
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import gymnasium as gym
|
import gymnasium as gym
|
||||||
@@ -42,13 +44,21 @@ from lerobot.envs.robotwin import (
|
|||||||
|
|
||||||
|
|
||||||
def _make_mock_task_env(
|
def _make_mock_task_env(
|
||||||
height: int = 480,
|
height: int = 240,
|
||||||
width: int = 640,
|
width: int = 320,
|
||||||
cameras: tuple[str, ...] = ROBOTWIN_CAMERA_NAMES,
|
cameras: tuple[str, ...] = ROBOTWIN_CAMERA_NAMES,
|
||||||
) -> MagicMock:
|
) -> MagicMock:
|
||||||
"""Return a mock that mimics the RoboTwin task class API."""
|
"""Return a mock that mimics the RoboTwin task class API.
|
||||||
obs_dict = {f"{cam}_rgb": np.zeros((height, width, 3), dtype=np.uint8) for cam in cameras}
|
|
||||||
obs_dict["joint_action"] = np.zeros(ACTION_DIM, dtype=np.float64)
|
RoboTwin's real get_obs returns
|
||||||
|
{"observation": {cam: {"rgb": img}}, "joint_action": {"vector": np.ndarray}, ...}
|
||||||
|
so the mock follows the same nested shape.
|
||||||
|
"""
|
||||||
|
obs_dict = {
|
||||||
|
"observation": {cam: {"rgb": np.zeros((height, width, 3), dtype=np.uint8)} for cam in cameras},
|
||||||
|
"joint_action": {"vector": np.zeros(ACTION_DIM, dtype=np.float64)},
|
||||||
|
"endpose": {},
|
||||||
|
}
|
||||||
|
|
||||||
mock = MagicMock()
|
mock = MagicMock()
|
||||||
mock.get_obs.return_value = obs_dict
|
mock.get_obs.return_value = obs_dict
|
||||||
@@ -60,10 +70,28 @@ def _make_mock_task_env(
|
|||||||
return mock
|
return mock
|
||||||
|
|
||||||
|
|
||||||
def _patch_load(mock_task_instance: MagicMock):
|
@contextmanager
|
||||||
"""Context manager that patches _load_robotwin_task to return a mock class."""
|
def _patch_runtime(mock_task_instance: MagicMock):
|
||||||
|
"""Patch both the task-class loader and the YAML config loader so the
|
||||||
|
env can construct + reset without a real RoboTwin install."""
|
||||||
task_cls = MagicMock(return_value=mock_task_instance)
|
task_cls = MagicMock(return_value=mock_task_instance)
|
||||||
return patch("lerobot.envs.robotwin._load_robotwin_task", return_value=task_cls)
|
fake_setup = {
|
||||||
|
"head_camera_h": 240,
|
||||||
|
"head_camera_w": 320,
|
||||||
|
"left_embodiment_config": {},
|
||||||
|
"right_embodiment_config": {},
|
||||||
|
"left_robot_file": "",
|
||||||
|
"right_robot_file": "",
|
||||||
|
"dual_arm_embodied": True,
|
||||||
|
"render_freq": 0,
|
||||||
|
"task_name": "beat_block_hammer",
|
||||||
|
"task_config": "demo_clean",
|
||||||
|
}
|
||||||
|
with (
|
||||||
|
patch("lerobot.envs.robotwin._load_robotwin_task", return_value=task_cls),
|
||||||
|
patch("lerobot.envs.robotwin._load_robotwin_setup_kwargs", return_value=fake_setup),
|
||||||
|
):
|
||||||
|
yield
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -79,12 +107,12 @@ class TestRoboTwinEnv:
|
|||||||
task_name="beat_block_hammer",
|
task_name="beat_block_hammer",
|
||||||
observation_height=h,
|
observation_height=h,
|
||||||
observation_width=w,
|
observation_width=w,
|
||||||
camera_names=["head_camera", "front_camera"],
|
camera_names=["head_camera", "left_camera"],
|
||||||
)
|
)
|
||||||
pixels_space = env.observation_space["pixels"]
|
pixels_space = env.observation_space["pixels"]
|
||||||
assert pixels_space["head_camera"].shape == (h, w, 3)
|
assert pixels_space["head_camera"].shape == (h, w, 3)
|
||||||
assert pixels_space["front_camera"].shape == (h, w, 3)
|
assert pixels_space["left_camera"].shape == (h, w, 3)
|
||||||
assert "left_wrist" not in pixels_space
|
assert "right_camera" not in pixels_space
|
||||||
|
|
||||||
def test_action_space(self):
|
def test_action_space(self):
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
@@ -94,7 +122,7 @@ class TestRoboTwinEnv:
|
|||||||
def test_reset_returns_correct_obs_keys(self):
|
def test_reset_returns_correct_obs_keys(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
obs, info = env.reset()
|
obs, info = env.reset()
|
||||||
|
|
||||||
assert "pixels" in obs
|
assert "pixels" in obs
|
||||||
@@ -107,15 +135,20 @@ class TestRoboTwinEnv:
|
|||||||
def test_reset_calls_setup_demo(self):
|
def test_reset_calls_setup_demo(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset(seed=42)
|
env.reset(seed=42)
|
||||||
mock_task.setup_demo.assert_called_once_with(seed=42, is_test=True)
|
# setup_demo receives the full YAML-derived kwargs plus seed + is_test;
|
||||||
|
# we only assert the caller-provided bits.
|
||||||
|
assert mock_task.setup_demo.call_count == 1
|
||||||
|
call_kwargs = mock_task.setup_demo.call_args.kwargs
|
||||||
|
assert call_kwargs["seed"] == 42
|
||||||
|
assert call_kwargs["is_test"] is True
|
||||||
|
|
||||||
def test_step_returns_correct_types(self):
|
def test_step_returns_correct_types(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset()
|
env.reset()
|
||||||
obs, reward, terminated, truncated, info = env.step(action)
|
obs, reward, terminated, truncated, info = env.step(action)
|
||||||
|
|
||||||
@@ -129,17 +162,17 @@ class TestRoboTwinEnv:
|
|||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
bad_action = np.zeros(7, dtype=np.float32) # wrong dim
|
bad_action = np.zeros(7, dtype=np.float32) # wrong dim
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset()
|
env.reset()
|
||||||
with pytest.raises(ValueError, match="Expected 1-D action"):
|
with pytest.raises(ValueError, match="Expected 1-D action"):
|
||||||
env.step(bad_action)
|
env.step(bad_action)
|
||||||
|
|
||||||
def test_success_terminates_episode(self):
|
def test_success_terminates_episode(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
mock_task.check_success.return_value = True # always succeeds
|
mock_task.check_success.return_value = True
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset()
|
env.reset()
|
||||||
_, _, terminated, _, info = env.step(action)
|
_, _, terminated, _, info = env.step(action)
|
||||||
assert terminated is True
|
assert terminated is True
|
||||||
@@ -149,7 +182,7 @@ class TestRoboTwinEnv:
|
|||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer", episode_length=2)
|
env = RoboTwinEnv(task_name="beat_block_hammer", episode_length=2)
|
||||||
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
action = np.zeros(ACTION_DIM, dtype=np.float32)
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset()
|
env.reset()
|
||||||
env.step(action) # step 1
|
env.step(action) # step 1
|
||||||
_, _, _, truncated, _ = env.step(action) # step 2 → truncated
|
_, _, _, truncated, _ = env.step(action) # step 2 → truncated
|
||||||
@@ -158,25 +191,25 @@ class TestRoboTwinEnv:
|
|||||||
def test_close_calls_close_env(self):
|
def test_close_calls_close_env(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
env.reset()
|
env.reset()
|
||||||
env.close()
|
env.close()
|
||||||
mock_task.close_env.assert_called_once_with(clearance=False)
|
mock_task.close_env.assert_called_once()
|
||||||
|
|
||||||
def test_black_frame_for_missing_camera(self):
|
def test_black_frame_for_missing_camera(self):
|
||||||
"""If a camera key is absent from get_obs(), a black frame is returned."""
|
"""If a camera key is absent from get_obs(), a black frame is returned."""
|
||||||
mock_task = _make_mock_task_env(cameras=("head_camera",)) # only head_camera
|
# Mock exposes only head_camera; we ask for both head_camera + left_camera.
|
||||||
|
mock_task = _make_mock_task_env(height=10, width=10, cameras=("head_camera",))
|
||||||
env = RoboTwinEnv(
|
env = RoboTwinEnv(
|
||||||
task_name="beat_block_hammer",
|
task_name="beat_block_hammer",
|
||||||
camera_names=["head_camera", "front_camera"],
|
camera_names=["head_camera", "left_camera"],
|
||||||
observation_height=10,
|
observation_height=10,
|
||||||
observation_width=10,
|
observation_width=10,
|
||||||
)
|
)
|
||||||
with _patch_load(mock_task):
|
with _patch_runtime(mock_task):
|
||||||
obs, _ = env.reset()
|
obs, _ = env.reset()
|
||||||
# front_camera is missing from mock → should be a black 10×10×3 array
|
assert obs["pixels"]["left_camera"].shape == (10, 10, 3)
|
||||||
assert obs["pixels"]["front_camera"].shape == (10, 10, 3)
|
assert obs["pixels"]["left_camera"].sum() == 0
|
||||||
assert obs["pixels"]["front_camera"].sum() == 0
|
|
||||||
|
|
||||||
def test_task_and_task_description_attributes(self):
|
def test_task_and_task_description_attributes(self):
|
||||||
env = RoboTwinEnv(task_name="beat_block_hammer")
|
env = RoboTwinEnv(task_name="beat_block_hammer")
|
||||||
@@ -196,12 +229,11 @@ class TestRoboTwinEnv:
|
|||||||
class TestCreateRoboTwinEnvs:
|
class TestCreateRoboTwinEnvs:
|
||||||
def test_returns_correct_structure(self):
|
def test_returns_correct_structure(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env_cls = gym.vector.SyncVectorEnv
|
with _patch_runtime(mock_task):
|
||||||
with _patch_load(mock_task):
|
|
||||||
envs = create_robotwin_envs(
|
envs = create_robotwin_envs(
|
||||||
task="beat_block_hammer",
|
task="beat_block_hammer",
|
||||||
n_envs=1,
|
n_envs=1,
|
||||||
env_cls=env_cls,
|
env_cls=gym.vector.SyncVectorEnv,
|
||||||
)
|
)
|
||||||
assert "beat_block_hammer" in envs
|
assert "beat_block_hammer" in envs
|
||||||
assert 0 in envs["beat_block_hammer"]
|
assert 0 in envs["beat_block_hammer"]
|
||||||
@@ -209,12 +241,11 @@ class TestCreateRoboTwinEnvs:
|
|||||||
|
|
||||||
def test_multi_task(self):
|
def test_multi_task(self):
|
||||||
mock_task = _make_mock_task_env()
|
mock_task = _make_mock_task_env()
|
||||||
env_cls = gym.vector.SyncVectorEnv
|
with _patch_runtime(mock_task):
|
||||||
with _patch_load(mock_task):
|
|
||||||
envs = create_robotwin_envs(
|
envs = create_robotwin_envs(
|
||||||
task="beat_block_hammer,click_bell",
|
task="beat_block_hammer,click_bell",
|
||||||
n_envs=1,
|
n_envs=1,
|
||||||
env_cls=env_cls,
|
env_cls=gym.vector.SyncVectorEnv,
|
||||||
)
|
)
|
||||||
assert set(envs.keys()) == {"beat_block_hammer", "click_bell"}
|
assert set(envs.keys()) == {"beat_block_hammer", "click_bell"}
|
||||||
|
|
||||||
@@ -265,12 +296,12 @@ class TestRoboTwinProcessorStep:
|
|||||||
step = RoboTwinProcessorStep()
|
step = RoboTwinProcessorStep()
|
||||||
obs = {
|
obs = {
|
||||||
f"{OBS_IMAGES}.head_camera": torch.zeros(1, 3, 4, 4),
|
f"{OBS_IMAGES}.head_camera": torch.zeros(1, 3, 4, 4),
|
||||||
f"{OBS_IMAGES}.front_camera": torch.zeros(1, 3, 4, 4),
|
f"{OBS_IMAGES}.left_camera": torch.zeros(1, 3, 4, 4),
|
||||||
OBS_STATE: torch.zeros(1, 14),
|
OBS_STATE: torch.zeros(1, 14),
|
||||||
}
|
}
|
||||||
result = step.observation(obs)
|
result = step.observation(obs)
|
||||||
assert f"{OBS_IMAGES}.head_camera" in result
|
assert f"{OBS_IMAGES}.head_camera" in result
|
||||||
assert f"{OBS_IMAGES}.front_camera" in result
|
assert f"{OBS_IMAGES}.left_camera" in result
|
||||||
assert result[OBS_STATE].dtype == torch.float32
|
assert result[OBS_STATE].dtype == torch.float32
|
||||||
|
|
||||||
def test_state_cast_to_float32(self):
|
def test_state_cast_to_float32(self):
|
||||||
|
|||||||
Reference in New Issue
Block a user