G1 rough-terrain training app (MuJoCo Playground + Brax PPO)

README.md

@@ -1,13 +1,29 @@
 ---
-title: Armins
-emoji: 🌍
+title: G1 Rough Terrain Training
+emoji: 🦿
 colorFrom: red
 colorTo: green
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.13'
+python_version: '3.11'
 app_file: app.py
 pinned: false
+short_description: GPU training for Unitree G1 rough-terrain locomotion
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Unitree G1 — rough-terrain locomotion training
+
+Trains `G1JoystickRoughTerrain` from [MuJoCo Playground](https://playground.mujoco.org/)
+with Brax PPO on the Space's GPU. Thousands of MJX environments step in parallel,
+versus the 16 CPU processes a laptop manages.
+
+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.
+
+**Settings that matter**
+- *Sleep time* → **never**. The Space idles on HTTP traffic, not CPU/GPU load, so
+  a background training thread will not keep it awake.
+- *Persistent storage* — without it, checkpoints are lost on restart. Set
+  `HF_REPO` so they are pushed to the Hub instead.
+- `NUM_TIMESTEPS` (default 200M), `ENV_NAME`, `SEED` are configurable variables.

app.py

@@ -0,0 +1,178 @@
+"""Unitree G1 rough-terrain locomotion training, on the Space's GPU.
+
+Training runs in a background thread so the Gradio server stays responsive; the
+UI is a log tail plus start/stop. Brax PPO cannot be interrupted from outside,
+so the stop button sets a flag that the progress callback checks and raises on.
+"""
+
+from __future__ import annotations
+
+import os
+import pickle
+import threading
+import traceback
+from collections import deque
+from datetime import datetime
+from pathlib import Path
+
+import gradio as gr
+
+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))
+HF_REPO = os.environ.get("HF_REPO", "").strip()
+AUTO_START = os.environ.get("AUTO_START", "1") == "1"
+
+# /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")
+
+LOG: deque[str] = deque(maxlen=800)
+STATE = {"thread": None, "status": "idle", "stop": False, "step": 0}
+
+
+def log(msg: str) -> None:
+    LOG.append(f"[{datetime.now():%H:%M:%S}] {msg}")
+    print(msg, flush=True)
+
+
+class Stopped(Exception):
+    pass
+
+
+def _upload(path: Path) -> None:
+    if not HF_REPO:
+        return
+    try:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        api.create_repo(HF_REPO, repo_type="model", exist_ok=True)
+        api.upload_file(path_or_fileobj=str(path), path_in_repo=path.name,
+                        repo_id=HF_REPO, repo_type="model")
+        log(f"uploaded {path.name} -> {HF_REPO}")
+    except Exception as e:                       # never let upload kill training
+        log(f"upload failed ({type(e).__name__}): {e}")
+
+
+def train_worker() -> None:
+    try:
+        STATE["status"] = "importing"
+        log("importing jax / playground / brax ...")
+        import functools
+        import jax
+
+        log(f"jax devices: {jax.devices()}")
+        if not any(d.platform == "gpu" for d in jax.devices()):
+            log("WARNING: no GPU visible to JAX -- this will be very slow")
+
+        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.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)
+        ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
+        ppo_params.num_timesteps = NUM_TIMESTEPS
+        log(f"num_envs={ppo_params.get('num_envs')} "
+            f"batch_size={ppo_params.get('batch_size')} "
+            f"timesteps={NUM_TIMESTEPS:,}")
+
+        OUT.mkdir(parents=True, exist_ok=True)
+
+        def progress(step, metrics):
+            if STATE["stop"]:
+                raise Stopped()
+            STATE["step"] = int(step)
+            rew = metrics.get("eval/episode_reward", float("nan"))
+            ln = metrics.get("eval/avg_episode_length", float("nan"))
+            log(f"step {int(step):>12,}  reward {rew:8.2f}  ep_len {ln:7.1f}")
+
+        def policy_params_fn(step, make_policy, params):
+            p = OUT / f"ckpt_{int(step)}.pkl"
+            with open(p, "wb") as f:
+                pickle.dump(params, f)
+            _upload(p)
+
+        ppo_kwargs = dict(ppo_params)
+        net_cfg = ppo_kwargs.pop("network_factory", None)
+        if net_cfg is not None:
+            ppo_kwargs["network_factory"] = functools.partial(
+                ppo_networks.make_ppo_networks, **dict(net_cfg))
+
+        STATE["status"] = "training"
+        log("compiling (first step takes several minutes) ...")
+        _, params, _ = ppo.train(
+            **ppo_kwargs,
+            environment=env,
+            eval_env=registry.load(ENV_NAME, config=env_cfg),
+            wrap_env_fn=wrapper.wrap_for_brax_training,
+            randomization_fn=registry.get_domain_randomizer(ENV_NAME),
+            progress_fn=progress,
+            policy_params_fn=policy_params_fn,
+            seed=SEED,
+        )
+        final = OUT / "final.pkl"
+        with open(final, "wb") as f:
+            pickle.dump(params, f)
+        _upload(final)
+        STATE["status"] = "done"
+        log("training complete")
+
+    except Stopped:
+        STATE["status"] = "stopped"
+        log("stopped by user")
+    except Exception as e:
+        STATE["status"] = f"error: {type(e).__name__}"
+        log(f"FAILED: {type(e).__name__}: {e}")
+        log(traceback.format_exc()[-2000:])
+
+
+def start():
+    t = STATE["thread"]
+    if t is not None and t.is_alive():
+        return status_md()
+    STATE["stop"] = False
+    STATE["thread"] = threading.Thread(target=train_worker, daemon=True)
+    STATE["thread"].start()
+    return status_md()
+
+
+def stop():
+    STATE["stop"] = True
+    log("stop requested -- will halt at the next eval boundary")
+    return status_md()
+
+
+def status_md() -> str:
+    alive = STATE["thread"] is not None and STATE["thread"].is_alive()
+    pct = 100.0 * STATE["step"] / max(NUM_TIMESTEPS, 1)
+    return (f"**{ENV_NAME}** — status: `{STATE['status']}`"
+            f"{' (running)' if alive else ''}  \n"
+            f"step {STATE['step']:,} / {NUM_TIMESTEPS:,}  ({pct:.1f}%)  \n"
+            f"checkpoints: `{OUT}`"
+            + (f" → pushing to `{HF_REPO}`" if HF_REPO else
+               "  \n⚠️ `HF_REPO` unset — checkpoints are not pushed to the Hub"))
+
+
+def logs() -> str:
+    return "\n".join(LOG) or "(no output yet)"
+
+
+with gr.Blocks(title="G1 rough-terrain training") as demo:
+    gr.Markdown("# Unitree G1 — rough-terrain locomotion training")
+    st = gr.Markdown(status_md())
+    with gr.Row():
+        gr.Button("Start", variant="primary").click(start, outputs=st)
+        gr.Button("Stop", variant="stop").click(stop, outputs=st)
+    out = gr.Textbox(label="training log", lines=26, max_lines=26,
+                     autoscroll=True, value=logs())
+    timer = gr.Timer(3.0)
+    timer.tick(logs, outputs=out)
+    timer.tick(status_md, outputs=st)
+
+if AUTO_START:
+    start()
+
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,6 @@
+jax[cuda12]>=0.4.35
+mujoco>=3.2.7
+mujoco-mjx>=3.2.7
+playground==0.2.0
+brax>=0.12.1
+huggingface_hub>=0.35