Watch checkpoints walk: rollout video tab; NUM_EVALS sets checkpoint cadence

README.md

@@ -34,6 +34,13 @@ Trains `G1JoystickRoughTerrain` from [MuJoCo Playground](https://playground.mujo
 with Brax PPO on the Space's GPU. Thousands of MJX environments step in parallel,
 versus the 16 CPU processes a laptop manages.
 
+**Watch a checkpoint walk.** The reward curve says whether training is moving;
+the *Render rollout* section says whether the robot walks. Pick a checkpoint,
+a forward speed, snow friction and depth, and it films the policy on a Himalaya
+crop (`rollout.py`; MJX + `MUJOCO_GL=egl`) and reports how long it stayed up and
+how far it walked. `NUM_EVALS` (default 40) sets how many evals — and therefore
+log lines and checkpoint uploads — a run makes.
+
 Training starts automatically on boot. Checkpoints are written to `/data` when
 persistent storage is attached, and pushed to `HF_REPO` if that variable and a
 write-scoped `HF_TOKEN` secret are set.

app.py

@@ -17,6 +17,8 @@ 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))
@@ -35,6 +37,7 @@ SNOW = os.environ.get("SNOW", "1") == "1"                            # snow on t
 SNOW_FRICTION = tuple(float(v) for v in os.environ.get("SNOW_FRICTION", "0.3,0.7").split(","))
 SNOW_DEPTH = tuple(float(v) for v in os.environ.get("SNOW_DEPTH", "0.0,0.08").split(","))
 SURFACE = TERRAIN + ("-snow" if SNOW else "")                        # checkpoint prefix
+NUM_EVALS = int(os.environ.get("NUM_EVALS", 40))     # evals (and checkpoints) per run
 
 # /data exists only when persistent storage is attached; fall back to /tmp.
 OUT = Path("/data/ckpt") if Path("/data").is_dir() else Path("/tmp/ckpt")
@@ -144,6 +147,49 @@ def install_jax_pmap_shims() -> None:
             log(f"shimmed jax.{name} (removed in this JAX version)")
 
 
+ENVS = {}
+
+
+def build_envs():
+    """(env, eval_env, randomization_fn) for ENV_NAME with the Himalaya terrain
+    and snow applied. Built once and cached; the rollout tab reuses eval_env."""
+    if ENVS:
+        return ENVS["env"], ENVS["eval_env"], ENVS["randomization_fn"]
+    from mujoco_playground import registry
+
+    STATE["status"] = "building env"
+    log(f"loading {ENV_NAME} ...")
+    env = registry.load(ENV_NAME)
+    env_cfg = registry.get_default_config(ENV_NAME)
+    eval_env = registry.load(ENV_NAME, config=env_cfg)
+    randomization_fn = registry.get_domain_randomizer(ENV_NAME)
+    if TERRAIN == "himalaya":
+        from himalaya_terrain import apply_terrain, make_terrains, randomizer
+        log(f"building {NUM_TERRAINS} Himalaya crops: {HIMALAYA_PATCH:.0f} m of Khumbu "
+            f"-> 20 m arena, relief {HIMALAYA_RELIEF} m ...")
+        grids = make_terrains(NUM_TERRAINS, seed=SEED, patch_m=HIMALAYA_PATCH,
+                              relief=HIMALAYA_RELIEF)
+        apply_terrain(env, grids, HIMALAYA_RELIEF)
+        apply_terrain(eval_env, grids, HIMALAYA_RELIEF)
+        randomization_fn = randomizer(randomization_fn, grids)
+        log("terrain: himalaya (per-env crops via domain randomization)")
+    else:
+        log("terrain: playground stock rough terrain")
+    if SNOW:
+        from himalaya_terrain import snow_randomizer
+        randomization_fn = snow_randomizer(randomization_fn, SNOW_FRICTION, SNOW_DEPTH)
+        log(f"snow: foot-floor friction U{SNOW_FRICTION}, depth U{SNOW_DEPTH} m "
+            "(soft contact, per env)")
+    ENVS.update(env=env, eval_env=eval_env, randomization_fn=randomization_fn)
+    return env, eval_env, randomization_fn
+
+
+def net_config():
+    from mujoco_playground.config import locomotion_params
+    cfg = locomotion_params.brax_ppo_config(ENV_NAME).get("network_factory", None)
+    return dict(cfg) if cfg is not None else None
+
+
 def train_worker() -> None:
     try:
         STATE["status"] = "importing"
@@ -160,34 +206,13 @@ def train_worker() -> None:
 
         from brax.training.agents.ppo import networks as ppo_networks
         from brax.training.agents.ppo import train as ppo
-        from mujoco_playground import registry, wrapper
+        from mujoco_playground import wrapper
         from mujoco_playground.config import locomotion_params
 
-        STATE["status"] = "building env"
-        log(f"loading {ENV_NAME} ...")
-        env = registry.load(ENV_NAME)
-        env_cfg = registry.get_default_config(ENV_NAME)
-        eval_env = registry.load(ENV_NAME, config=env_cfg)
-        randomization_fn = registry.get_domain_randomizer(ENV_NAME)
-        if TERRAIN == "himalaya":
-            from himalaya_terrain import apply_terrain, make_terrains, randomizer
-            log(f"building {NUM_TERRAINS} Himalaya crops: {HIMALAYA_PATCH:.0f} m of Khumbu "
-                f"-> 20 m arena, relief {HIMALAYA_RELIEF} m ...")
-            grids = make_terrains(NUM_TERRAINS, seed=SEED, patch_m=HIMALAYA_PATCH,
-                                  relief=HIMALAYA_RELIEF)
-            apply_terrain(env, grids, HIMALAYA_RELIEF)
-            apply_terrain(eval_env, grids, HIMALAYA_RELIEF)
-            randomization_fn = randomizer(randomization_fn, grids)
-            log("terrain: himalaya (per-env crops via domain randomization)")
-        else:
-            log("terrain: playground stock rough terrain")
-        if SNOW:
-            from himalaya_terrain import snow_randomizer
-            randomization_fn = snow_randomizer(randomization_fn, SNOW_FRICTION, SNOW_DEPTH)
-            log(f"snow: foot-floor friction U{SNOW_FRICTION}, depth U{SNOW_DEPTH} m "
-                "(soft contact, per env)")
+        env, eval_env, randomization_fn = build_envs()
         ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
         ppo_params.num_timesteps = int(STATE["target"])
+        ppo_params.num_evals = NUM_EVALS       # each eval = one log line + checkpoint upload
         log(f"num_envs={ppo_params.get('num_envs')} "
             f"batch_size={ppo_params.get('batch_size')} "
             f"timesteps={int(STATE['target']):,}")
@@ -289,6 +314,45 @@ def logs() -> str:
     return "\n".join(LOG) or "(no output yet)"
 
 
+def list_ckpts():
+    files = sorted(OUT.glob("*.pkl"), key=lambda p: p.stat().st_mtime, reverse=True)
+    return [f.name for f in files]
+
+
+def refresh_ckpts():
+    names = list_ckpts()
+    return gr.Dropdown(choices=names, value=names[0] if names else None)
+
+
+def render(ckpt: str | None, vx: float, friction: float, depth: float, seconds: float):
+    """Load a checkpoint and film the policy walking on a snowy Himalaya crop."""
+    from rollout import apply_snow, load_params, make_inference_fn, rollout, write_video
+
+    if not ckpt:
+        return None, "no checkpoint yet -- train first (or wait for ckpt_0.pkl)"
+    path = OUT / ckpt
+    if not path.exists():
+        return None, f"{ckpt} not found"
+    try:
+        _, eval_env, _ = build_envs()
+        params = load_params(path)
+        policy = make_inference_fn(eval_env, net_config())(params, deterministic=True)
+        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)
+        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)"
+        log(msg)
+        return str(out), msg
+    except Exception as e:
+        log(f"rollout FAILED: {type(e).__name__}: {e}")
+        log(traceback.format_exc()[-1500:])
+        return None, f"rollout failed: {type(e).__name__}: {str(e)[:300]}"
+
+
 with gr.Blocks(title="G1 rough-terrain training") as demo:
     gr.Markdown("# Unitree G1 — Himalayan terrain locomotion training")
     st = gr.Markdown(status_md())
@@ -303,6 +367,20 @@ with gr.Blocks(title="G1 rough-terrain training") as demo:
     timer.tick(logs, outputs=out)
     timer.tick(status_md, outputs=st)
 
+    gr.Markdown("## Watch a checkpoint walk")
+    with gr.Row():
+        ckpt_dd = gr.Dropdown(choices=list_ckpts(), label="checkpoint",
+                              value=(list_ckpts() or [None])[0])
+        gr.Button("Refresh").click(refresh_ckpts, outputs=ckpt_dd)
+        vx_in = gr.Slider(0.0, 1.0, value=0.5, step=0.1, label="forward speed (m/s)")
+        mu_in = gr.Slider(0.2, 1.0, value=0.5, step=0.05, label="snow friction")
+        depth_in = gr.Slider(0.0, 0.10, value=0.05, step=0.01, label="snow depth (m)")
+        secs_in = gr.Slider(2, 20, value=8, step=1, label="seconds")
+    verdict = gr.Markdown("")
+    video = gr.Video(label="rollout", autoplay=True)
+    gr.Button("Render rollout", variant="primary").click(
+        render, inputs=[ckpt_dd, vx_in, mu_in, depth_in, secs_in], outputs=[video, verdict])
+
 if AUTO_START:
     start()
 

requirements.txt

@@ -5,3 +5,5 @@ playground==0.2.0
 brax>=0.12.1
 huggingface_hub>=0.35
 scipy>=1.10
+imageio>=2.31
+imageio-ffmpeg>=0.4.9

rollout.py

@@ -0,0 +1,108 @@
+"""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 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, 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."""
+    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)
+    cmd = jnp.asarray(command, dtype=jnp.float32)
+    state.info["command"] = cmd
+
+    n_steps = int(seconds / env.dt)
+    states = [state]
+    start = np.asarray(state.data.qpos[:3])
+    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)
+        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])),
+            "fell_at": fell_at}
+    return frames, info
+
+
+def write_video(frames, path: Path, fps: float) -> Path:
+    import imageio.v2 as imageio
+
+    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))
+    return path