rollout.py

"""Watch a checkpoint walk: rebuild the Brax PPO policy from a saved params
pickle, run it on one Himalaya/snow environment, and encode the frames to MP4.

A reward curve says whether training is moving; only a rollout says whether the
robot walks. `ppo.train` returns `(make_inference_fn, params, metrics)`, and the
checkpoints written by `policy_params_fn` are the same `params` tuple
`(normalizer_params, policy_params, value_params)`, so the policy is rebuilt
exactly the way `ppo.train` builds it: the same `network_factory` config on the
same env, then `make_inference_fn(params, deterministic=True)`.
"""

from __future__ import annotations

import os
import pickle
from pathlib import Path

import numpy as np


def make_inference_fn(env, net_cfg: dict | None):
    """The `make_inference_fn` that `ppo.train(...)` returns, without training."""
    from brax.training.acme import running_statistics
    from brax.training.agents.ppo import networks as ppo_networks

    ppo_network = ppo_networks.make_ppo_networks(
        observation_size=env.observation_size,
        action_size=env.action_size,
        preprocess_observations_fn=running_statistics.normalize,
        **dict(net_cfg or {}),
    )
    return ppo_networks.make_inference_fn(ppo_network)


def load_params(path: Path):
    with open(path, "rb") as f:
        params = pickle.load(f)
    # policy_params_fn and the final return both give (normalizer, policy, value);
    # the inference fn wants the first two.
    return tuple(params)[:2] if len(params) == 3 else params


def apply_snow(env, friction: float, depth: float) -> None:
    """Fixed snow on a single (non-randomized) Playground env: same foot-floor
    contact-pair parameters `snow_randomizer` draws per env, then rebuild the
    MJX model like `apply_terrain` does."""
    from mujoco import mjx

    from himalaya_terrain import snow_params

    mj = env._mj_model
    solref, solimp = snow_params(float(depth))
    mj.pair_friction[0:2, 0:2] = friction
    mj.pair_solref[0:2] = np.asarray(solref)
    mj.pair_solimp[0:2] = np.asarray(solimp)
    floor = mj.geom("floor").id                      # make it look like snow too
    mj.geom_matid[floor] = -1
    mj.geom_rgba[floor] = (0.92, 0.94, 0.98, 1.0)
    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):
    """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

    jit_reset = jax.jit(env.reset)
    jit_step = jax.jit(env.step)
    jit_policy = jax.jit(inference_fn)

    rng = jax.random.PRNGKey(seed)
    state = jit_reset(rng)
    # The joystick task steers by a "command" in info; the get-up task has no
    # such key, and adding one would change the State pytree the step function
    # was traced against.
    steerable = "command" in state.info
    cmd = jnp.asarray(command, dtype=jnp.float32)
    if steerable:
        state.info["command"] = cmd

    n_steps = int(seconds / env.dt)
    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)
        if steerable:
            state.info["command"] = cmd      # joystick envs resample on reset only
        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

    qpos = np.stack(qpos)
    info = {"seconds": len(qpos) * env.dt,
            "distance_m": float(np.linalg.norm((qpos[-1, :2] - qpos[0, :2]))),
            "start_height_m": float(qpos[0, 2]),
            "end_height_m": float(qpos[-1, 2]),
            "peak_height_m": float(qpos[:, 2].max()),
            "fell_at": fell_at}
    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 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