Fix: render rollouts in a subprocess with a probed GL backend (EGL import broke training)

app.py

@@ -17,8 +17,6 @@ from pathlib import Path
 
 import gradio as gr
 
-os.environ.setdefault("MUJOCO_GL", "egl")     # offscreen rendering for rollouts
-
 ENV_NAME = os.environ.get("ENV_NAME", "G1JoystickRoughTerrain")
 NUM_TIMESTEPS = int(os.environ.get("NUM_TIMESTEPS", 200_000_000))
 SEED = int(os.environ.get("SEED", 0))
@@ -340,8 +338,8 @@ def render(ckpt: str | None, vx: float, friction: float, depth: float, seconds:
         if SNOW:
             apply_snow(eval_env, float(friction), float(depth))
         log(f"rollout {ckpt}: vx={vx} friction={friction} depth={depth} m, {seconds}s ...")
-        frames, info = rollout(eval_env, policy, seconds=float(seconds), command=(float(vx), 0.0, 0.0))
-        out = write_video(frames, OUT / "rollouts" / f"{path.stem}.mp4", fps=1.0 / eval_env.dt)
+        qpos, info = rollout(eval_env, policy, seconds=float(seconds), command=(float(vx), 0.0, 0.0))
+        out = write_video(eval_env, qpos, OUT / "rollouts" / f"{path.stem}.mp4", fps=1.0 / eval_env.dt)
         verdict = (f"fell at {info['fell_at']:.1f}s" if info["fell_at"] is not None
                    else f"stayed up for {info['seconds']:.1f}s")
         msg = f"{ckpt}: {verdict}, walked {info['distance_m']:.2f} m (commanded {vx} m/s)"

packages.txt

@@ -0,0 +1,3 @@
+libegl1
+libgl1
+libosmesa6

render_worker.py

@@ -0,0 +1,30 @@
+"""Render a qpos trajectory to MP4. Run as a subprocess by rollout.write_video
+with MUJOCO_GL already set, because MuJoCo fixes its GL backend at import.
+
+    python render_worker.py model.mjb qpos.npy out.mp4 fps width height camera
+"""
+import sys
+
+import numpy as np
+
+
+def main(mjb, qpos_path, out, fps, width, height, camera):
+    import imageio.v2 as imageio
+    import mujoco
+
+    model = mujoco.MjModel.from_binary_path(mjb)
+    data = mujoco.MjData(model)
+    qpos = np.load(qpos_path)
+    cam = camera if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera) >= 0 else -1
+    with mujoco.Renderer(model, height=int(height), width=int(width)) as r, \
+            imageio.get_writer(out, fps=float(fps), codec="libx264", quality=7,
+                               macro_block_size=None) as w:
+        for q in qpos:
+            data.qpos[:] = q
+            mujoco.mj_forward(model, data)
+            r.update_scene(data, camera=cam)
+            w.append_data(r.render())
+
+
+if __name__ == "__main__":
+    main(*sys.argv[1:8])

rollout.py

@@ -11,6 +11,7 @@ same env, then `make_inference_fn(params, deterministic=True)`.
 
 from __future__ import annotations
 
+import os
 import pickle
 from pathlib import Path
 
@@ -58,11 +59,10 @@ def apply_snow(env, friction: float, depth: float) -> None:
     env._mjx_model = mjx.put_model(mj, impl=env._config.impl)
 
 
-def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0),
-            seed: int = 0, width: int = 640, height: int = 480, camera: str = "track"):
-    """Run the policy for `seconds` with a fixed joystick command and return
-    (frames, info): frames is a list of HxWx3 uint8 images at the env's control
-    rate, info has distance walked and whether it fell."""
+def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0), seed: int = 0):
+    """Run the policy for `seconds` with a fixed joystick command. Returns
+    (qpos, info): qpos is a (T, nq) trajectory at the env's control rate, info
+    has distance walked and when (if) it fell."""
     import jax
     import jax.numpy as jnp
 
@@ -76,33 +76,70 @@ def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0),
     state.info["command"] = cmd
 
     n_steps = int(seconds / env.dt)
-    states = [state]
-    start = np.asarray(state.data.qpos[:3])
+    qpos = [np.asarray(state.data.qpos)]
     fell_at = None
     for i in range(n_steps):
         rng, key = jax.random.split(rng)
         act, _ = jit_policy(state.obs, key)
         state = jit_step(state, act)
         state.info["command"] = cmd          # joystick envs resample on reset only
-        states.append(state)
+        qpos.append(np.asarray(state.data.qpos))
         if fell_at is None and float(state.done) > 0.5:
             fell_at = (i + 1) * env.dt
             break
 
-    end = np.asarray(states[-1].data.qpos[:3])
-    frames = env.render(states, width=width, height=height, camera=camera)
-    info = {"seconds": len(states) * env.dt,
-            "distance_m": float(np.linalg.norm((end - start)[:2])),
+    qpos = np.stack(qpos)
+    info = {"seconds": len(qpos) * env.dt,
+            "distance_m": float(np.linalg.norm((qpos[-1, :2] - qpos[0, :2]))),
             "fell_at": fell_at}
-    return frames, info
-
-
-def write_video(frames, path: Path, fps: float) -> Path:
-    import imageio.v2 as imageio
-
+    return qpos, info
+
+
+def offscreen_gl_backend() -> str | None:
+    """Which MuJoCo offscreen backend this machine can actually load. The
+    backend is fixed when `mujoco` is imported, so rendering happens in a
+    subprocess with MUJOCO_GL set to this."""
+    import ctypes.util
+    import sys
+
+    forced = os.environ.get("ROLLOUT_GL")
+    if forced:
+        return forced
+    if sys.platform == "darwin":
+        return "glfw"
+    if ctypes.util.find_library("EGL"):
+        return "egl"
+    if ctypes.util.find_library("OSMesa"):
+        return "osmesa"
+    return None
+
+
+def write_video(env, qpos: np.ndarray, path: Path, fps: float,
+                width: int = 640, height: int = 480, camera: str = "track") -> Path:
+    """Render a qpos trajectory of `env` to MP4 in a subprocess (see
+    render_worker.py), so a missing or different GL backend can never take the
+    training process down."""
+    import subprocess
+    import sys
+    import tempfile
+
+    backend = offscreen_gl_backend()
+    if backend is None:
+        raise RuntimeError("no offscreen GL backend (libEGL / libOSMesa) in this container -- "
+                           "add `libegl1 libosmesa6` to packages.txt")
     path.parent.mkdir(parents=True, exist_ok=True)
-    with imageio.get_writer(str(path), fps=fps, codec="libx264", quality=7,
-                            macro_block_size=None) as w:
-        for fr in frames:
-            w.append_data(np.asarray(fr))
+    with tempfile.TemporaryDirectory() as tmp:
+        import mujoco
+        mjb = Path(tmp) / "model.mjb"
+        mujoco.mj_saveModel(env._mj_model, str(mjb), None)
+        traj = Path(tmp) / "qpos.npy"
+        np.save(traj, qpos)
+        worker = Path(__file__).with_name("render_worker.py")
+        env_vars = dict(os.environ, MUJOCO_GL=backend, PYOPENGL_PLATFORM=backend)
+        r = subprocess.run([sys.executable, str(worker), str(mjb), str(traj), str(path),
+                            str(fps), str(width), str(height), camera],
+                           env=env_vars, capture_output=True, text=True, timeout=600)
+    if r.returncode != 0:
+        tail = (r.stderr or r.stdout).strip().splitlines()[-6:]
+        raise RuntimeError(f"render ({backend}) failed: " + " | ".join(tail))
     return path