diff --git a/pyproject.toml b/pyproject.toml index 285e3f6..729af48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,12 +27,17 @@ plugrl-run-env-client = "plugrl_env_client.cli:main" plugrl-run-worker = "plugrl_env_client.cli:main" [project.optional-dependencies] +# robomimic 0.3.0, the last on PyPI, imports mujoco_py in env_robosuite, and +# mujoco_py needs MuJoCo 2.1 binaries and system GL headers; v0.4.0 does not, +# and supports robosuite 1.2 onwards. robosuite 1.4.1 with mujoco 2.3.7 +# matches the environment metadata shipped in envs/robomimic/env_meta, and +# mujoco 3.x trips robosuite 1.4.1's joint-type assertion - so this extra and +# `mujoco` (gymnasium's MuJoCo 3) belong in separate environments. robomimic's +# dependency egl-probe builds with CMake: `pip install cmake` first. robomimic = [ - "cython<3", - "d4rl>=1.1,<2.0.0", - "patchelf>=0.17.2.4,<0.18.0", - "robomimic==0.3.0", - "robosuite<1.5.0", + "robomimic @ git+https://github.com/ARISE-Initiative/robomimic.git@v0.4.0", + "robosuite==1.4.1", + "mujoco==2.3.7", "PyOpenGL==3.1.4", ] atari = [ diff --git a/src/plugrl_env_client/envs/robomimic/robomimic_env.py b/src/plugrl_env_client/envs/robomimic/robomimic_env.py index 6255526..9380a76 100644 --- a/src/plugrl_env_client/envs/robomimic/robomimic_env.py +++ b/src/plugrl_env_client/envs/robomimic/robomimic_env.py @@ -37,9 +37,12 @@ import robomimic.envs.env_robosuite # type: ignore import robomimic.utils.env_utils as _env_utils # type: ignore import robomimic.utils.obs_utils as _obs_utils # type: ignore -except ImportError: +except ImportError as e: + # Say what failed. "Not installed" was the message when robomimic was + # installed and its own import of mujoco_py was what failed. raise ImportError( - "Robomimic is not installed. Please install it with the 'robomimic' extra, e.g. 'pip install plugrl-env-client[robomimic]'" + f"Could not import robomimic ({e}). Install it with the 'robomimic' " + "extra, e.g. 'pip install plugrl-env-client[robomimic]'" ) except Exception as e: raise ImportError(f"An error occurred while importing robomimic: {e}") @@ -50,6 +53,18 @@ UID = "Robomimic-v1" ENV_META_DIR = pathlib.Path(__file__).parent / "env_meta" +# robomimic's environments never end an episode themselves: the dataset +# metadata sets ignore_done and EnvRobosuite.is_done() is never true. +# robomimic's own rollouts stop on success or at these horizons, the ones its +# paper configs use per task. +_ROLLOUT_HORIZONS = { + "Lift": 400, + "PickPlaceCan": 400, + "NutAssemblySquare": 400, + "ToolHang": 700, + "TwoArmTransport": 700, +} + @register_env_config(UID) @dataclasses.dataclass @@ -64,6 +79,11 @@ class RobomimicConfig(BaseEnvConfig): ] ) agentview_image_size: tuple[int, int] = (720, 1280) + # Steps before an episode is truncated. None takes robomimic's rollout + # horizon for the task; a task without one must set it. + horizon: int | None = None + # End the episode, as terminated, on the step the task succeeds. + terminate_on_success: bool = True @register_env(UID, best_reward_threshold_for_success=1.0) @@ -97,6 +117,18 @@ def __init__( f"Environment metadata file not found for env name: {config.name}" ) + horizon = config.horizon + if horizon is None: + horizon = _ROLLOUT_HORIZONS.get(env_meta["env_name"]) + if horizon is None: + raise ValueError( + f"No rollout horizon is known for {env_meta['env_name']}; set " + "--env.horizon, or its episodes will never end" + ) + self.horizon = int(horizon) + self.terminate_on_success = config.terminate_on_success + self._steps = 0 + _obs_utils.initialize_obs_modality_mapping_from_dict( dict( low_dim=config.low_dim_keys, @@ -158,6 +190,7 @@ def reset( "simulation underneath is not seeded by this wrapper. Either " "run without --runner.seed, or seed the underlying env first." ) + self._steps = 0 obs = self.env.reset() agentview_image = self._render_agentview_image() return self.prepare_obs(obs, agentview_image), {} @@ -168,10 +201,15 @@ def step( assert actions.shape[0] == 1, "Batch size must be 1 for robomimic env" action = actions[0].tolist() obs, reward, done, info = self.env.step(action) + self._steps += 1 agentview_image = self._render_agentview_image() reward = np.array([float(reward)], dtype=np.float32) - terminated = np.array([bool(done)], dtype=np.bool_) - truncated = np.array([False], dtype=np.bool_) + succeeded = self.terminate_on_success and bool(self.env.is_success()["task"]) + ended = bool(done) or succeeded + terminated = np.array([ended], dtype=np.bool_) + truncated = np.array( + [not ended and self._steps >= self.horizon], dtype=np.bool_ + ) return ( self.prepare_obs(obs, agentview_image), reward, diff --git a/tests/test_robomimic_episodes.py b/tests/test_robomimic_episodes.py new file mode 100644 index 0000000..fa38d82 --- /dev/null +++ b/tests/test_robomimic_episodes.py @@ -0,0 +1,142 @@ +"""robomimic episodes have to end. + +robomimic's environments never end an episode on their own: the dataset +metadata the client ships sets `ignore_done`, `EnvRobosuite.is_done()` is +never true, and robomimic's own rollouts stop on success or at a fixed +horizon. The client did neither, so a robomimic run was one episode that +never ended - E27's pilot on NutAssemblySquare finished 0 episodes in 4,096 +steps. These tests pin the two stops robomimic uses: success, and the +task's rollout horizon. +""" + +import numpy as np +import pytest + +robomimic_env = pytest.importorskip( + "plugrl_env_client.envs.robomimic.robomimic_env", + exc_type=ImportError, + reason="requires the robomimic extra", +) +RobomimicConfig = robomimic_env.RobomimicConfig +RobomimicEnv = robomimic_env.RobomimicEnv + + +class _FakeRobosuite: + hard_reset = True + + +class _FakeEnvRobosuite: + """Stands in for robomimic's EnvRobosuite: never done, success on a step.""" + + action_dimension = 7 + success_at: int | None = None + + def __init__(self) -> None: + self.env = _FakeRobosuite() + self.t = 0 + + def _obs(self) -> dict: + return { + "object": np.zeros(14), + "robot0_eye_in_hand_image": np.zeros((96, 96, 3), dtype=np.uint8), + } + + def reset(self) -> dict: + self.t = 0 + return self._obs() + + def step(self, action): + self.t += 1 + return self._obs(), 0.0, False, {} + + def is_success(self) -> dict: + return {"task": self.success_at is not None and self.t >= self.success_at} + + def render(self, **kwargs) -> np.ndarray: + return np.zeros((kwargs["height"], kwargs["width"], 3), dtype=np.uint8) + + +@pytest.fixture +def make_env(monkeypatch: pytest.MonkeyPatch): + fake = _FakeEnvRobosuite() + monkeypatch.setattr( + robomimic_env._env_utils, "create_env_from_metadata", lambda **kw: fake + ) + monkeypatch.setattr( + robomimic_env._obs_utils, + "initialize_obs_modality_mapping_from_dict", + lambda mapping: None, + ) + monkeypatch.setattr( + robomimic_env.robomimic.envs.env_robosuite, "EnvRobosuite", _FakeEnvRobosuite + ) + + def make(success_at: int | None = None, **config): + fake.success_at = success_at + env = RobomimicEnv( + RobomimicConfig(name="square-img", agentview_image_size=(8, 8), **config) + ) + env.reset() + return env + + return make + + +def _step(env) -> tuple[bool, bool]: + _, _, terminated, truncated, _ = env.step(np.zeros((1, 7), dtype=np.float32)) + return bool(terminated[0]), bool(truncated[0]) + + +def test_an_episode_is_cut_at_the_task_s_rollout_horizon(make_env): + env = make_env() # NutAssemblySquare: robomimic rolls out 400 steps + + ends = [_step(env) for _ in range(400)] + + assert ends[:399] == [(False, False)] * 399 + assert ends[399] == (False, True) + + +def test_success_ends_the_episode(make_env): + env = make_env(success_at=5) + + ends = [_step(env) for _ in range(5)] + + assert ends[:4] == [(False, False)] * 4 + assert ends[4] == (True, False) + + +def test_success_can_be_left_running(make_env): + env = make_env(success_at=5, terminate_on_success=False) + + ends = [_step(env) for _ in range(6)] + + assert ends == [(False, False)] * 6 + + +def test_an_explicit_horizon_wins(make_env): + env = make_env(horizon=10) + + ends = [_step(env) for _ in range(10)] + + assert ends[9] == (False, True) + + +def test_reset_restarts_the_count(make_env): + env = make_env(horizon=3) + for _ in range(3): + _step(env) + + env.reset() + + assert [_step(env) for _ in range(3)] == [ + (False, False), + (False, False), + (False, True), + ] + + +def test_a_task_with_no_known_horizon_needs_one(make_env, monkeypatch): + monkeypatch.setattr(robomimic_env, "_ROLLOUT_HORIZONS", {}, raising=False) + + with pytest.raises(ValueError, match="horizon"): + make_env()