Spaces:
Sleeping
Sleeping
| """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 install_jax_pmap_shims() -> None: | |
| """Re-add jax.device_put_replicated / device_put_sharded if this JAX removed them. | |
| JAX deprecated both in 0.8.1 and removed them in 0.10.0 (April 2026), but the | |
| current brax *release* (0.14.2, which playground 0.2.0 requires) still calls | |
| device_put_replicated; only brax main has the sharding-based replacement. | |
| Shimming the two functions is a smaller intervention than pinning JAX down, | |
| which would risk the MJX/Warp stack that already compiles cleanly here. | |
| """ | |
| import jax | |
| import jax.numpy as jnp | |
| import numpy as np | |
| def _sharding(devices): | |
| mesh = jax.sharding.Mesh(np.array(list(devices)), axis_names=("i",)) | |
| return jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec("i")) | |
| def device_put_replicated(x, devices): | |
| sharding, n = _sharding(devices), len(devices) | |
| def rep(leaf): | |
| stack = jnp.stack if isinstance(leaf, jax.Array) else np.stack | |
| return jax.device_put(stack([leaf] * n), sharding) | |
| return jax.tree_util.tree_map(rep, x) | |
| def device_put_sharded(shards, devices): | |
| sharding = _sharding(devices) | |
| def put(*leaves): | |
| stack = jnp.stack if isinstance(leaves[0], jax.Array) else np.stack | |
| return jax.device_put(stack(list(leaves)), sharding) | |
| return jax.tree_util.tree_map(put, *shards) | |
| for name, fn in (("device_put_replicated", device_put_replicated), | |
| ("device_put_sharded", device_put_sharded)): | |
| try: | |
| getattr(jax, name) | |
| except AttributeError: | |
| setattr(jax, name, fn) | |
| log(f"shimmed jax.{name} (removed in this JAX version)") | |
| def train_worker() -> None: | |
| try: | |
| STATE["status"] = "importing" | |
| log("importing jax / playground / brax ...") | |
| import functools | |
| import jax | |
| log(f"jax {jax.__version__} devices: {jax.devices()}") | |
| install_jax_pmap_shims() | |
| import brax | |
| log(f"brax {brax.__version__}") | |
| 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() | |