Add a snow gait: higher swing, wider stance, slower cadence, foot-separation cost

README.md

@@ -41,6 +41,23 @@ 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.
 
+## Snow gait (`GAIT=snow`)
+
+Playground's joystick task is tuned for pavement: a 15 cm swing-height reference
+at 1.25–1.5 Hz with the hips pinned to a narrow stance. `GAIT=snow` retunes it
+for deep snow via `snow_gait.py` — a 22 cm swing (`FOOT_HEIGHT`), a slower
+1.0–1.3 Hz cadence (`GAIT_FREQ`) with more air-time reward for longer strides, a
+relaxed hip-deviation cost so the stance can widen, and a new cost for letting
+the feet come within `FOOT_SEPARATION` (20 cm) of each other *laterally*.
+
+That last term is also a fix, not just styling. The joystick task ends an episode
+when the two feet touch, and the 200M-step run plateaued because the gait crossed
+its feet after ~45 steps and took the −100 termination every episode while still
+upright. The separation cost penalises exactly that, and is measured along the
+pelvis's y-axis so feet may still pass fore-aft as walking requires.
+
+Checkpoints go under a `-snowgait` suffix.
+
 ## Fall recovery (`TASK=getup`)
 
 Walking and getting up are different tasks, and the joystick task cannot learn

app.py

@@ -38,7 +38,15 @@ SNOW_DEPTH = tuple(float(v) for v in os.environ.get("SNOW_DEPTH", "0.0,0.08").sp
 # "getup" = fall recovery (g1_getup.py): most episodes start fallen, nothing
 # terminates on being down, and the model gains torso/pelvis collision geoms.
 TASK = os.environ.get("TASK", "walk").strip().lower()
-SURFACE = TERRAIN + ("-snow" if SNOW else "") + ("-getup" if TASK == "getup" else "")
+# "snow" retunes the walking gait for deep snow: higher swing, wider stance,
+# slower cadence, and a lateral foot-separation cost (see snow_gait.py).
+GAIT = os.environ.get("GAIT", "street").strip().lower()
+FOOT_HEIGHT = float(os.environ.get("FOOT_HEIGHT", 0.22))        # swing reference, m
+FOOT_SEPARATION = float(os.environ.get("FOOT_SEPARATION", 0.20))  # min lateral gap, m
+GAIT_FREQ = tuple(float(v) for v in os.environ.get("GAIT_FREQ", "1.0,1.3").split(","))
+SURFACE = (TERRAIN + ("-snow" if SNOW else "")
+           + ("-getup" if TASK == "getup" else "")
+           + ("-snowgait" if TASK == "walk" and GAIT == "snow" else ""))
 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.
@@ -167,10 +175,21 @@ def build_envs():
         randomization_fn = None            # no terrain or snow: this task is about the body
         ENVS.update(env=env, eval_env=eval_env, randomization_fn=randomization_fn)
         return env, eval_env, randomization_fn
-    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)
+    if GAIT == "snow":
+        import snow_gait
+        log(f"gait: snow -- swing {FOOT_HEIGHT * 100:.0f} cm (street 15), "
+            f"stance >= {FOOT_SEPARATION * 100:.0f} cm, {GAIT_FREQ[0]}-{GAIT_FREQ[1]} Hz "
+            "(street 1.25-1.5), relaxed hips")
+        cfg = snow_gait.snow_gait_config(foot_height=FOOT_HEIGHT,
+                                         foot_separation=FOOT_SEPARATION,
+                                         gait_freq=GAIT_FREQ)
+        env = snow_gait.G1SnowGait(config=cfg)
+        eval_env = snow_gait.G1SnowGait(config=cfg)
+    else:
+        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
@@ -376,7 +395,8 @@ def render(ckpt: str | None, vx: float, friction: float, depth: float, seconds:
 
 
 with gr.Blocks(title="G1 rough-terrain training") as demo:
-    gr.Markdown(f"# Unitree G1 — {'fall recovery' if TASK == 'getup' else 'Himalayan terrain locomotion'} training")
+    gr.Markdown(f"# Unitree G1 — {'fall recovery' if TASK == 'getup' else 'Himalayan terrain locomotion'} training"
+                + (" (snow gait)" if TASK == "walk" and GAIT == "snow" else ""))
     st = gr.Markdown(status_md())
     with gr.Row():
         steps_in = gr.Number(value=NUM_TIMESTEPS, precision=0, label="timesteps",

snow_gait.py

@@ -0,0 +1,109 @@
+"""A snow gait for the G1: high, wide, deliberate steps instead of a street walk.
+
+Playground's joystick task rewards tracking a swing-height reference of 15 cm at
+1.25-1.5 Hz with the hips held near the nominal (narrow) pose. That is a gait for
+pavement. Walking through snow, people lift their feet clear of the surface,
+plant them wider apart, and take slower, longer strides.
+
+Four changes, three of them just numbers the joystick task already exposes:
+
+  * `max_foot_height` 0.15 -> 0.22 m. This is the reference trajectory the
+    `feet_phase` reward tracks (`gait.get_rz`), so raising it literally asks for
+    a higher swing -- the direct "pick your feet up" knob.
+  * `joint_deviation_hip` -0.25 -> -0.05. That cost pins the hips (including
+    hip *roll*) to the nominal narrow stance; relaxing it permits a wide one.
+  * gait frequency 1.25-1.5 Hz -> 1.0-1.3 Hz, and `feet_air_time` 2.0 -> 3.0.
+    Slower cadence at the same commanded speed means a longer stride.
+  * a new `feet_separation` cost, which the joystick task has no equivalent of.
+
+That last one is not cosmetic. The joystick task ends an episode when the two
+feet touch each other, and the 200M-step run plateaued because the learned gait
+crossed its feet after ~45 steps and ate the -100 termination every episode
+while still upright. Penalising *lateral* separation below a threshold attacks
+that directly, and a wide stance is what you want in snow anyway. The cost is
+lateral-only, measured along the pelvis's y-axis, so the feet may still pass each
+other fore-aft as any walking gait requires.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Union
+
+import jax
+import jax.numpy as jp
+from ml_collections import config_dict
+from mujoco import mjx
+
+from mujoco_playground._src.locomotion.g1 import joystick as g1_joystick
+
+# Defaults, all overridable from the Space (see app.py).
+FOOT_HEIGHT = 0.22          # m, swing-height reference (street gait: 0.15)
+FOOT_SEPARATION = 0.20      # m, minimum lateral gap between the feet
+SEPARATION_COST = -2.0      # scale for falling below that gap
+HIP_DEVIATION_COST = -0.05  # street gait: -0.25
+AIR_TIME_REWARD = 3.0       # street gait: 2.0
+GAIT_FREQ = (1.0, 1.3)      # Hz, street gait: (1.25, 1.5)
+
+
+def snow_gait_config(
+    foot_height: float = FOOT_HEIGHT,
+    foot_separation: float = FOOT_SEPARATION,
+    separation_cost: float = SEPARATION_COST,
+    hip_deviation_cost: float = HIP_DEVIATION_COST,
+    air_time_reward: float = AIR_TIME_REWARD,
+    gait_freq: tuple[float, float] = GAIT_FREQ,
+) -> config_dict.ConfigDict:
+    cfg = g1_joystick.default_config()
+    rc = cfg.reward_config
+    rc.max_foot_height = foot_height
+    rc.min_foot_separation = foot_separation          # new
+    rc.scales.feet_separation = separation_cost       # new
+    rc.scales.joint_deviation_hip = hip_deviation_cost
+    rc.scales.feet_air_time = air_time_reward
+    cfg.gait_freq = list(gait_freq)                   # new
+    return cfg
+
+
+class G1SnowGait(g1_joystick.Joystick):
+    """The joystick task, retuned for deep snow."""
+
+    def __init__(
+        self,
+        task: str = "rough_terrain",
+        config: Optional[config_dict.ConfigDict] = None,
+        config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    ) -> None:
+        super().__init__(task=task, config=config or snow_gait_config(),
+                         config_overrides=config_overrides)
+        self._pelvis_body_id = self._mj_model.body("pelvis").id
+
+    def reset(self, rng: jax.Array) -> Any:
+        # The base class samples gait frequency internally from a hard-coded
+        # 1.25-1.5 Hz; phase_dt is the only thing it derives from it, so a slower
+        # cadence is a matter of rewriting that one entry after the fact.
+        rng, freq_rng = jax.random.split(rng)
+        state = super().reset(rng)
+        lo, hi = self._config.gait_freq
+        gait_freq = jax.random.uniform(freq_rng, (1,), minval=lo, maxval=hi)
+        state.info["phase_dt"] = 2 * jp.pi * self.dt * gait_freq
+        return state
+
+    def _get_reward(self, data: mjx.Data, action: jax.Array, info: dict[str, Any],
+                    metrics: dict[str, Any], done: jax.Array,
+                    first_contact: jax.Array, contact: jax.Array) -> dict[str, jax.Array]:
+        rewards = super()._get_reward(data, action, info, metrics, done,
+                                      first_contact, contact)
+        rewards["feet_separation"] = self._cost_feet_separation(data)
+        return rewards
+
+    def _cost_feet_separation(self, data: mjx.Data) -> jax.Array:
+        """How far the feet are inside the minimum lateral gap, in metres.
+
+        Measured along the pelvis's y-axis so that fore-aft passing during swing
+        costs nothing -- only the feet converging sideways does.
+        """
+        feet = data.site_xpos[self._feet_site_id]          # (2, 3), world frame
+        lateral_axis = data.xmat[self._pelvis_body_id].reshape(3, 3)[:, 1]
+        lateral_gap = jp.abs(jp.dot(feet[0] - feet[1], lateral_axis))
+        return jp.clip(self._config.reward_config.min_foot_separation - lateral_gap,
+                       0.0, None)