fix(robomimic): end episodes on success and at the task's horizon; make the extra installable

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/[email protected]",
+    "robosuite==1.4.1",
+    "mujoco==2.3.7",
     "PyOpenGL==3.1.4",
 ]
 atari = [

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,

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()