Train on real Himalayan (Khumbu SRTM) terrain: per-env heightfield crops

README.md

@@ -1,5 +1,5 @@
 ---
-title: G1 Rough Terrain Training
+title: G1 Himalaya Terrain Training
 emoji: 🦿
 colorFrom: red
 colorTo: green
@@ -8,10 +8,19 @@ sdk_version: 6.26.0
 python_version: '3.11'
 app_file: app.py
 pinned: false
-short_description: GPU training for Unitree G1 rough-terrain locomotion
+short_description: GPU training for Unitree G1 on real Himalayan terrain
 ---
 
-# Unitree G1 — rough-terrain locomotion training
+# Unitree G1 — Himalayan terrain locomotion training
+
+Trains on **real elevation data**: `dem_khumbu.npy` is SRTM-derived terrain
+(AWS Terrain Tiles, 4.2 m/px, 3.2 km square) of the Khumbu glacier valley below
+Everest Base Camp. `himalaya_terrain.py` cuts random 600 m crops, squeezes them
+onto Playground's 20 m arena with 0.3 m of relief plus scree detail, and gives
+every environment its own crop through the domain randomizer. Set `TERRAIN=playground`
+to fall back to the stock 5 cm noise field; `HIMALAYA_RELIEF`, `HIMALAYA_PATCH`,
+`NUM_TERRAINS` tune the terrain.
+
 
 Trains `G1JoystickRoughTerrain` from [MuJoCo Playground](https://playground.mujoco.org/)
 with Brax PPO on the Space's GPU. Thousands of MJX environments step in parallel,

app.py

@@ -24,6 +24,14 @@ HF_REPO = os.environ.get("HF_REPO", "").strip()
 HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None
 AUTO_START = os.environ.get("AUTO_START", "1") == "1"
 
+# Terrain: "himalaya" swaps Playground's 5 cm noise heightfield for crops of real
+# SRTM elevation of the Khumbu valley (see himalaya_terrain.py); "playground"
+# keeps the stock terrain. Every env trains on its own crop.
+TERRAIN = os.environ.get("TERRAIN", "himalaya")
+NUM_TERRAINS = int(os.environ.get("NUM_TERRAINS", 64))
+HIMALAYA_RELIEF = float(os.environ.get("HIMALAYA_RELIEF", 0.3))    # metres over the 20 m arena
+HIMALAYA_PATCH = float(os.environ.get("HIMALAYA_PATCH", 600.0))    # metres of real ground per arena
+
 # /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")
 
@@ -153,6 +161,20 @@ def train_worker() -> None:
         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")
         ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
         ppo_params.num_timesteps = int(STATE["target"])
         log(f"num_envs={ppo_params.get('num_envs')} "
@@ -196,9 +218,9 @@ def train_worker() -> None:
         _, params, _ = ppo.train(
             **ppo_kwargs,
             environment=env,
-            eval_env=registry.load(ENV_NAME, config=env_cfg),
+            eval_env=eval_env,
             wrap_env_fn=wrapper.wrap_for_brax_training,
-            randomization_fn=registry.get_domain_randomizer(ENV_NAME),
+            randomization_fn=randomization_fn,
             progress_fn=progress,
             policy_params_fn=policy_params_fn,
             seed=SEED,
@@ -244,7 +266,7 @@ def status_md() -> str:
     alive = STATE["thread"] is not None and STATE["thread"].is_alive()
     target = int(STATE["target"])
     pct = 100.0 * STATE["step"] / max(target, 1)
-    return (f"**{ENV_NAME}** — status: `{STATE['status']}`"
+    return (f"**{ENV_NAME}** · terrain `{TERRAIN}` — status: `{STATE['status']}`"
             f"{' (running)' if alive else ''}  \n"
             f"step {STATE['step']:,} / {target:,}  ({pct:.1f}%)  \n"
             f"checkpoints: `{OUT}`"
@@ -257,7 +279,7 @@ def logs() -> str:
 
 
 with gr.Blocks(title="G1 rough-terrain training") as demo:
-    gr.Markdown("# Unitree G1 — rough-terrain locomotion training")
+    gr.Markdown("# Unitree G1 — Himalayan terrain locomotion training")
     st = gr.Markdown(status_md())
     with gr.Row():
         steps_in = gr.Number(value=NUM_TIMESTEPS, precision=0, label="timesteps",

dem_khumbu.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a9f73e91d7ee40949349f7ea4b6a4e035ebc4a97d257a70921e498bd5136548
+size 2359424

himalaya_terrain.py

@@ -0,0 +1,111 @@
+"""Himalayan heightfields for MuJoCo Playground's G1JoystickRoughTerrain.
+
+`dem_khumbu.npy` is real SRTM-derived elevation (AWS Terrain Tiles, zoom 15,
+4.2 m/px, 3.2 km square) of the Khumbu glacier valley below Everest Base Camp,
+generated by examples/g1_himalaya.py in the mujoco-mcp-server repo.
+
+Playground's rough-terrain scene is a 256x256 heightfield over a 20 m arena.
+A real valley is not walkable at 1:1 by a humanoid, so `patch_m` metres of real
+ground are squeezed onto the arena and the relief normalized to `relief` metres.
+Each crop is centred so the ground under the spawn sits at data value 0.5 --
+the floor geom is lowered by relief/2 to put that at z = 0 -- and a 1 m disc
+around the spawn is levelled (Playground jitters the spawn by +-0.5 m).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+
+DEM_PATH = Path(__file__).with_name("dem_khumbu.npy")
+DEM_METRES_PER_PX = 4.2
+
+
+def load_dem() -> np.ndarray:
+    return np.load(DEM_PATH).astype(np.float64)
+
+
+def sample_patch(dem, mpp, rng, patch_m, res, arena, relief,
+                 detail_m=0.06, detail_width_m=0.15):
+    """One res x res crop of `patch_m` metres of the DEM at a random position and
+    heading, in [0, 1]. A slope-weighted fine layer adds scree (SRTM has no
+    rock-scale detail)."""
+    from scipy.ndimage import gaussian_filter, map_coordinates
+
+    n = dem.shape[0]
+    px = patch_m / mpp
+    margin = px * 0.75
+    ci, cj = rng.uniform(margin, n - margin, size=2)
+    theta = rng.uniform(0, 2 * np.pi)
+
+    u = (np.arange(res) - (res - 1) / 2) / (res - 1) * px
+    U, V = np.meshgrid(u, u)
+    c, s = np.cos(theta), np.sin(theta)
+    h = map_coordinates(dem, [ci - s * U + c * V, cj + c * U + s * V], order=3, mode="nearest")
+    h = (h - h.min()) / (h.max() - h.min() + 1e-9)
+
+    if detail_m > 0:
+        gy, gx = np.gradient(h)
+        slope = np.hypot(gx, gy)
+        slope = slope / (np.percentile(slope, 90) + 1e-9)
+        fine = gaussian_filter(rng.normal(size=(res, res)),
+                               sigma=max(0.6, detail_width_m / (arena / res)))
+        fine /= np.abs(fine).max() + 1e-9
+        h = h + (detail_m / relief) * np.clip(slope, 0.15, 1.0) * fine
+        h = (h - h.min()) / (h.max() - h.min() + 1e-9)
+    return h
+
+
+def make_terrains(n_terrains: int, seed: int = 0, res: int = 256, arena: float = 20.0,
+                  patch_m: float = 600.0, relief: float = 0.3,
+                  spawn_clear_r: float = 1.0) -> np.ndarray:
+    """(n_terrains, res, res) heightfield grids in [0, 1], spawn-centred at 0.5."""
+    dem = load_dem()
+    rng = np.random.default_rng(seed)
+    u = np.linspace(-arena / 2, arena / 2, res)
+    X, Y = np.meshgrid(u, u)
+    blend = np.clip((np.hypot(X, Y) - spawn_clear_r) / 0.8, 0.0, 1.0)
+
+    out = np.empty((n_terrains, res, res), dtype=np.float32)
+    for k in range(n_terrains):
+        h = sample_patch(dem, DEM_METRES_PER_PX, rng, patch_m, res, arena, relief)
+        h0 = float(h[res // 2, res // 2])
+        h = h * blend + h0 * (1.0 - blend)
+        # centre the spawn height at 0.5 without leaving [0, 1]
+        h = 0.5 + 0.5 * (h - h0) / max(h0, 1.0 - h0, 1e-6)
+        out[k] = np.clip(h, 0.0, 1.0)
+    return out
+
+
+def apply_terrain(env, grids: np.ndarray, relief: float) -> None:
+    """Bake grids[0] into a Playground env's model and rebuild its MJX model.
+    Later per-env variation goes through `randomizer`."""
+    from mujoco import mjx
+
+    mj = env._mj_model
+    hf = mj.hfield("hfield").id
+    floor = mj.geom("floor").id
+    mj.hfield_size[hf, 2] = relief            # vertical range
+    mj.geom_pos[floor, 2] = -0.5 * relief     # data 0.5 -> z = 0 under the spawn
+    mj.hfield_data[:] = grids[0].ravel()
+    env._mjx_model = mjx.put_model(mj, impl=env._config.impl)
+
+
+def randomizer(base_fn, grids: np.ndarray):
+    """Wrap Playground's domain randomizer so every env also gets its own crop
+    of the Himalaya (env i uses grids[i % len(grids)])."""
+    import jax
+    import jax.numpy as jnp
+
+    flat = jnp.asarray(grids.reshape(len(grids), -1))
+
+    def fn(model, rng):
+        model, in_axes = base_fn(model, rng)
+        n = rng.shape[0]
+        idx = jnp.arange(n) % len(grids)
+        model = model.tree_replace({"hfield_data": flat[idx]})
+        in_axes = in_axes.tree_replace({"hfield_data": 0})
+        return model, in_axes
+
+    return fn

requirements.txt

@@ -4,3 +4,4 @@ mujoco-mjx>=3.2.7
 playground==0.2.0
 brax>=0.12.1
 huggingface_hub>=0.35
+scipy>=1.10