Add a fall-recovery task: TASK=getup trains a G1 get-up policy on a full-collision model

README.md

@@ -41,6 +41,34 @@ 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.
 
+## Fall recovery (`TASK=getup`)
+
+Walking and getting up are different tasks, and the joystick task cannot learn
+the second: it ends the episode the moment the torso tips past horizontal (with
+a −100 reward), so the policy never sees a timestep while fallen. `TASK=getup`
+switches the Space to `g1_getup.py`, a humanoid port of Playground's `Go1Getup`:
+60% of episodes *start* fallen (dropped from 0.5 m with a random orientation and
+random joint angles, then left to settle), nothing terminates on being down, and
+the reward asks for an upright torso at standing height — then, only once both
+hold, for the nominal pose and for the policy to stop moving.
+
+Playground's G1 is a *feet-only* collision model, so a fallen torso would drop
+through the floor. `g1_getup_model.py` rewrites the robot XML: it enables the
+thigh, shin and hand collision geoms the model already carries but leaves off,
+and adds torso, pelvis and head capsules sized from the visual meshes. Body geoms
+get `contype=2 / conaffinity=1` so they collide with the ground but never with
+each other — without that the pelvis capsule overlaps the thighs and the standing
+robot is blown over by its own contacts.
+
+Observations (93) and actions (29) keep the joystick task's layout, so a fall
+detector can hand control between a walking policy and a get-up policy at run
+time. Checkpoints go under a `-getup` prefix. Flat ground only: `TERRAIN` and
+`SNOW` are ignored for this task.
+
+```
+TASK=getup      # then Start; ~100M steps
+```
+
 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

@@ -34,7 +34,11 @@ HIMALAYA_PATCH = float(os.environ.get("HIMALAYA_PATCH", 600.0))    # metres of r
 SNOW = os.environ.get("SNOW", "1") == "1"                            # snow on the terrain
 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
+# "walk" = Playground's joystick task (the terrain/snow work above).
+# "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 "")
 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.
@@ -156,6 +160,13 @@ def build_envs():
     from mujoco_playground import registry
 
     STATE["status"] = "building env"
+    if TASK == "getup":
+        import g1_getup
+        log("task: getup (fall recovery) -- flat ground, full-collision G1")
+        env, eval_env = g1_getup.G1Getup(), g1_getup.G1Getup()
+        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)
@@ -182,9 +193,17 @@ def build_envs():
     return env, eval_env, randomization_fn
 
 
-def net_config():
+def ppo_config():
+    """Brax PPO config for the active task."""
+    if TASK == "getup":
+        import g1_getup
+        return g1_getup.brax_ppo_config()
     from mujoco_playground.config import locomotion_params
-    cfg = locomotion_params.brax_ppo_config(ENV_NAME).get("network_factory", None)
+    return locomotion_params.brax_ppo_config(ENV_NAME)
+
+
+def net_config():
+    cfg = ppo_config().get("network_factory", None)
     return dict(cfg) if cfg is not None else None
 
 
@@ -205,10 +224,9 @@ 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 wrapper
-        from mujoco_playground.config import locomotion_params
 
         env, eval_env, randomization_fn = build_envs()
-        ppo_params = locomotion_params.brax_ppo_config(ENV_NAME)
+        ppo_params = ppo_config()
         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')} "
@@ -254,7 +272,7 @@ def train_worker() -> None:
             environment=env,
             eval_env=eval_env,
             wrap_env_fn=wrapper.wrap_for_brax_training,
-            randomization_fn=randomization_fn,
+            **({"randomization_fn": randomization_fn} if randomization_fn else {}),
             progress_fn=progress,
             policy_params_fn=policy_params_fn,
             seed=SEED,
@@ -335,14 +353,20 @@ def render(ckpt: str | None, vx: float, friction: float, depth: float, seconds:
         _, eval_env, _ = build_envs()
         params = load_params(path)
         policy = make_inference_fn(eval_env, net_config())(params, deterministic=True)
-        if SNOW:
+        if SNOW and TASK != "getup":
             apply_snow(eval_env, float(friction), float(depth))
-        log(f"rollout {ckpt}: vx={vx} friction={friction} depth={depth} m, {seconds}s ...")
+        log(f"rollout {ckpt}: task={TASK} vx={vx} friction={friction} depth={depth} m, {seconds}s ...")
         qpos, info = rollout(eval_env, policy, seconds=float(seconds), command=(float(vx), 0.0, 0.0))
         out = write_video(eval_env, qpos, 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)"
+        if TASK == "getup":
+            # Root height: ~0.76 m standing, ~0.1-0.3 m sprawled.
+            msg = (f"{ckpt}: started at {info['start_height_m']:.2f} m, ended at "
+                   f"{info['end_height_m']:.2f} m (peak {info['peak_height_m']:.2f}; "
+                   f"standing is ~0.76 m)")
+        else:
+            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:
@@ -352,7 +376,7 @@ def render(ckpt: str | None, vx: float, friction: float, depth: float, seconds:
 
 
 with gr.Blocks(title="G1 rough-terrain training") as demo:
-    gr.Markdown("# Unitree G1 — Himalayan terrain locomotion training")
+    gr.Markdown(f"# Unitree G1 — {'fall recovery' if TASK == 'getup' else 'Himalayan terrain locomotion'} training")
     st = gr.Markdown(status_md())
     with gr.Row():
         steps_in = gr.Number(value=NUM_TIMESTEPS, precision=0, label="timesteps",

g1_getup.py

@@ -0,0 +1,335 @@
+"""Fall recovery for the Unitree G1 -- a humanoid port of Playground's Go1Getup.
+
+Walking and getting up are different tasks, and the joystick task cannot learn
+the second one: `G1JoystickRoughTerrain._get_termination` ends the episode the
+moment the torso tips past horizontal (with a -100 reward), so the policy never
+sees a single timestep while fallen.
+
+This env inverts that. Most episodes *start* fallen -- the robot is dropped from
+0.5 m with a random orientation and random joint angles, then allowed to settle
+-- and nothing terminates on being down. The reward asks for an upright torso at
+standing height, and only once both hold does it ask for the nominal pose and
+for the policy to stop moving.
+
+It runs on the collision model from `g1_getup_model`, because Playground's G1 is
+feet-only and a fallen torso would otherwise fall through the floor.
+
+Observations and the action space match the joystick task's layout (29 joints,
+gyro + up-vector + joint state + last action), so the two policies can be
+switched between at run time by a fall detector.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Union
+
+import jax
+import jax.numpy as jp
+import mujoco
+import numpy as np
+from ml_collections import config_dict
+from mujoco import mjx
+
+from mujoco_playground._src import mjx_env
+from mujoco_playground._src.locomotion.g1 import base as g1_base
+
+from g1_getup_model import build_scene
+
+# Torso-site height standing in the "knees_bent" keyframe is 0.949 m; ask for a
+# little less so the target is comfortably reachable.
+STANDING_TORSO_HEIGHT = 0.90
+
+
+def default_config() -> config_dict.ConfigDict:
+    return config_dict.create(
+        ctrl_dt=0.02,
+        sim_dt=0.002,
+        episode_length=400,          # 8 s -- getting a humanoid up takes longer than a quadruped
+        drop_from_height_prob=0.6,
+        settle_time=0.5,
+        action_repeat=1,
+        action_scale=0.5,
+        restricted_joint_range=False,   # recovery needs the full range of motion
+        soft_joint_pos_limit_factor=0.95,
+        energy_termination_threshold=np.inf,
+        noise_config=config_dict.create(
+            level=1.0,
+            scales=config_dict.create(
+                joint_pos=0.03,
+                joint_vel=1.5,
+                gyro=0.2,
+                gravity=0.05,
+            ),
+        ),
+        reward_config=config_dict.create(
+            scales=config_dict.create(
+                orientation=1.0,
+                torso_height=1.0,
+                posture=1.0,
+                stand_still=1.0,
+                action_rate=-0.001,
+                dof_pos_limits=-0.1,
+                torques=-1e-5,
+                dof_acc=-2.5e-7,
+                dof_vel=-0.1,
+            ),
+        ),
+        impl="warp",
+        naconmax=80 * 8192,          # a fallen humanoid makes far more contacts than a walking one
+        njmax=500,
+    )
+
+
+class G1Getup(g1_base.G1Env):
+    """Recover from a fall and stand up.
+
+    Observation (state, 93): gyro (3), torso up-vector (3), joint angles minus
+    the nominal pose (29), joint velocities (29), last action (29).
+
+    Action (29): joint targets, scaled and added to the *current* joint angles
+    rather than to the nominal pose -- as in Go1Getup, this gives a policy that
+    starts in an arbitrary sprawl a much wider initial range of motion.
+
+    Reward: upright torso, torso at standing height, and -- gated on both of
+    those already holding -- the nominal posture and zero action, plus the usual
+    torque / action-rate / joint-limit regularizers.
+    """
+
+    def __init__(
+        self,
+        config: config_dict.ConfigDict = default_config(),
+        config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    ) -> None:
+        # G1Env.__init__ compiles from an XML path; this task needs the generated
+        # full-collision model instead, so set the model up here and skip it.
+        mjx_env.MjxEnv.__init__(self, config, config_overrides)
+        scene, assets = build_scene()
+        self._model_assets = assets
+        self._mj_model = mujoco.MjModel.from_xml_string(scene, assets=assets)
+        self._mj_model.opt.timestep = self.sim_dt
+        self._mj_model.vis.global_.offwidth = 3840
+        self._mj_model.vis.global_.offheight = 2160
+        self._mjx_model = mjx.put_model(self._mj_model, impl=self._config.impl)
+        self._xml_path = "<generated by g1_getup_model.build_scene>"
+        self._post_init()
+
+    def _post_init(self) -> None:
+        self._init_q = jp.array(self._mj_model.keyframe("knees_bent").qpos)
+        self._default_pose = jp.array(self._mj_model.keyframe("knees_bent").qpos[7:])
+
+        # First joint is the freejoint.
+        self._lowers, self._uppers = self.mj_model.jnt_range[1:].T
+        c = (self._lowers + self._uppers) / 2
+        r = self._uppers - self._lowers
+        self._soft_lowers = c - 0.5 * r * self._config.soft_joint_pos_limit_factor
+        self._soft_uppers = c + 0.5 * r * self._config.soft_joint_pos_limit_factor
+
+        self._settle_steps = int(self._config.settle_time / self.sim_dt)
+        self._z_des = STANDING_TORSO_HEIGHT
+        # The G1's "gravity" sensor is an up-vector: +z when upright (the joystick
+        # task calls it fallen at < 0), the opposite sign to Go1's convention.
+        self._up_vec = jp.array([0.0, 0.0, 1.0])
+        self._torso_site_id = self._mj_model.site("imu_in_torso").id
+
+    # -- episode ------------------------------------------------------------
+
+    def _get_random_qpos(self, rng: jax.Array) -> jax.Array:
+        """A sprawl: 0.5 m up, random orientation, random joint angles."""
+        rng, orientation_rng, qpos_rng = jax.random.split(rng, 3)
+
+        qpos = jp.zeros(self.mjx_model.nq)
+        qpos = qpos.at[2].set(0.5)
+        quat = jax.random.normal(orientation_rng, (4,))
+        quat /= jp.linalg.norm(quat) + 1e-6
+        qpos = qpos.at[3:7].set(quat)
+        qpos = qpos.at[7:].set(
+            jax.random.uniform(qpos_rng, (self.mjx_model.nu,),
+                               minval=self._lowers, maxval=self._uppers)
+        )
+        return qpos
+
+    def reset(self, rng: jax.Array) -> mjx_env.State:
+        rng, key1, key2 = jax.random.split(rng, 3)
+        qpos = jp.where(
+            jax.random.bernoulli(key1, self._config.drop_from_height_prob),
+            self._get_random_qpos(key2),
+            self._init_q,
+        )
+
+        rng, key = jax.random.split(rng)
+        qvel = jp.zeros(self.mjx_model.nv)
+        qvel = qvel.at[0:6].set(jax.random.uniform(key, (6,), minval=-0.5, maxval=0.5))
+
+        data = mjx_env.make_data(
+            self.mj_model, qpos=qpos, qvel=qvel, ctrl=qpos[7:],
+            impl=self.mjx_model.impl.value,
+            naconmax=self._config.naconmax, njmax=self._config.njmax,
+        )
+        data = mjx.forward(self.mjx_model, data)
+        # Let it flop and come to rest before the episode starts, so the policy
+        # is asked to recover from a settled pose rather than mid-air.
+        data = mjx_env.step(self.mjx_model, data, qpos[7:], self._settle_steps)
+        data = data.replace(time=0.0)
+
+        info = {
+            "rng": rng,
+            "last_act": jp.zeros(self.mjx_model.nu),
+            "last_last_act": jp.zeros(self.mjx_model.nu),
+        }
+        metrics = {f"reward/{k}": jp.zeros(())
+                   for k in self._config.reward_config.scales.keys()}
+        metrics["torso_height"] = jp.zeros(())
+
+        obs = self._get_obs(data, info)
+        reward, done = jp.zeros(2)
+        return mjx_env.State(data, obs, reward, done, metrics, info)
+
+    def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
+        motor_targets = state.data.qpos[7:] + action * self._config.action_scale
+        data = mjx_env.step(self.mjx_model, state.data, motor_targets, self.n_substeps)
+
+        obs = self._get_obs(data, state.info)
+        done = self._get_termination(data)
+
+        rewards = self._get_reward(data, action, state.info)
+        rewards = {k: v * self._config.reward_config.scales[k]
+                   for k, v in rewards.items()}
+        reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)
+
+        state.info["last_last_act"] = state.info["last_act"]
+        state.info["last_act"] = action
+        for k, v in rewards.items():
+            state.metrics[f"reward/{k}"] = v
+        state.metrics["torso_height"] = data.site_xpos[self._torso_site_id][2]
+
+        return state.replace(data=data, obs=obs, reward=reward, done=jp.float32(done))
+
+    def _get_termination(self, data: mjx.Data) -> jax.Array:
+        # Deliberately no fall termination: being down is the starting state.
+        energy = jp.sum(jp.abs(data.actuator_force * data.qvel[6:]))
+        return energy > self._config.energy_termination_threshold
+
+    # -- observation --------------------------------------------------------
+
+    def _noisy(self, info: dict[str, Any], value: jax.Array, scale: float) -> jax.Array:
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noise = (2 * jax.random.uniform(noise_rng, shape=value.shape) - 1)
+        return value + noise * self._config.noise_config.level * scale
+
+    def _get_obs(self, data: mjx.Data, info: dict[str, Any]) -> Dict[str, jax.Array]:
+        scales = self._config.noise_config.scales
+        gyro = self.get_gyro(data, "torso")
+        gravity = self.get_gravity(data, "torso")
+        joint_angles = data.qpos[7:]
+        joint_vel = data.qvel[6:]
+
+        state = jp.concatenate([
+            self._noisy(info, gyro, scales.gyro),                              # 3
+            self._noisy(info, gravity, scales.gravity),                        # 3
+            self._noisy(info, joint_angles, scales.joint_pos) - self._default_pose,  # 29
+            self._noisy(info, joint_vel, scales.joint_vel),                    # 29
+            info["last_act"],                                                  # 29
+        ])
+
+        privileged_state = jp.hstack([
+            state,
+            gyro,
+            self.get_accelerometer(data, "torso"),
+            self.get_local_linvel(data, "torso"),
+            self.get_global_angvel(data, "torso"),
+            joint_angles,
+            joint_vel,
+            data.actuator_force,
+            data.site_xpos[self._torso_site_id][2],
+        ])
+
+        return {"state": state, "privileged_state": privileged_state}
+
+    # -- reward -------------------------------------------------------------
+
+    def _get_reward(self, data: mjx.Data, action: jax.Array,
+                    info: dict[str, Any]) -> dict[str, jax.Array]:
+        torso_height = data.site_xpos[self._torso_site_id][2]
+        joint_angles = data.qpos[7:]
+        gravity = self.get_gravity(data, "torso")
+
+        gate = self._is_upright(gravity) * self._is_at_desired_height(torso_height)
+
+        return {
+            "orientation": self._reward_orientation(gravity),
+            "torso_height": self._reward_height(torso_height),
+            "posture": self._reward_posture(joint_angles, self._is_upright(gravity)),
+            "stand_still": self._reward_stand_still(action, gate),
+            "action_rate": self._cost_action_rate(action, info),
+            "torques": self._cost_torques(data.actuator_force),
+            "dof_pos_limits": self._cost_joint_pos_limits(joint_angles),
+            "dof_acc": self._cost_dof_acc(data.qacc[6:]),
+            "dof_vel": self._cost_dof_vel(data.qvel[6:]),
+        }
+
+    def _is_upright(self, up_vec: jax.Array, ori_tol: float = 0.05) -> jax.Array:
+        return jp.sum(jp.square(self._up_vec - up_vec)) < ori_tol
+
+    def _is_at_desired_height(self, torso_height: jax.Array,
+                              pos_tol: float = 0.02) -> jax.Array:
+        height = jp.minimum(torso_height, self._z_des)
+        return (self._z_des - height) < pos_tol
+
+    def _reward_orientation(self, up_vec: jax.Array) -> jax.Array:
+        return jp.exp(-2.0 * jp.sum(jp.square(self._up_vec - up_vec)))
+
+    def _reward_height(self, torso_height: jax.Array) -> jax.Array:
+        return jp.exp(jp.minimum(torso_height, self._z_des)) - 1.0
+
+    def _reward_posture(self, joint_angles: jax.Array, gate: jax.Array) -> jax.Array:
+        return gate * jp.exp(-0.5 * jp.sum(jp.square(joint_angles - self._default_pose)))
+
+    def _reward_stand_still(self, act: jax.Array, gate: jax.Array) -> jax.Array:
+        return gate * jp.exp(-0.5 * jp.sum(jp.square(act)))
+
+    def _cost_torques(self, torques: jax.Array) -> jax.Array:
+        return jp.sqrt(jp.sum(jp.square(torques))) + jp.sum(jp.abs(torques))
+
+    def _cost_action_rate(self, act: jax.Array, info: dict[str, Any]) -> jax.Array:
+        c1 = jp.sum(jp.square(act - info["last_act"]))
+        c2 = jp.sum(jp.square(act - 2 * info["last_act"] + info["last_last_act"]))
+        return c1 + c2
+
+    def _cost_joint_pos_limits(self, qpos: jax.Array) -> jax.Array:
+        out_of_limits = -jp.clip(qpos - self._soft_lowers, None, 0.0)
+        out_of_limits += jp.clip(qpos - self._soft_uppers, 0.0, None)
+        return jp.sum(out_of_limits)
+
+    def _cost_dof_vel(self, qvel: jax.Array) -> jax.Array:
+        return jp.sum(jp.square(jp.maximum(jp.abs(qvel) - 2.0 * jp.pi, 0.0)))
+
+    def _cost_dof_acc(self, qacc: jax.Array) -> jax.Array:
+        return jp.sum(jp.square(qacc))
+
+
+def brax_ppo_config(num_timesteps: int = 100_000_000) -> config_dict.ConfigDict:
+    """PPO config for the get-up task, following Playground's Go1Getup recipe
+    (short episodes, more parallel envs than the joystick task needs)."""
+    return config_dict.create(
+        num_timesteps=num_timesteps,
+        num_evals=40,
+        reward_scaling=1.0,
+        episode_length=default_config().episode_length,
+        normalize_observations=True,
+        action_repeat=1,
+        unroll_length=20,
+        num_minibatches=32,
+        num_updates_per_batch=4,
+        discounting=0.97,
+        learning_rate=3e-4,
+        entropy_cost=0.005,
+        num_envs=8192,
+        batch_size=256,
+        max_grad_norm=1.0,
+        network_factory=config_dict.create(
+            policy_hidden_layer_sizes=(512, 256, 128),
+            value_hidden_layer_sizes=(512, 256, 128),
+            policy_obs_key="state",
+            value_obs_key="privileged_state",
+        ),
+    )

g1_getup_model.py

@@ -0,0 +1,100 @@
+"""A G1 model that can lie on the ground.
+
+Playground's `G1JoystickRoughTerrain` is a *feet-only* collision model: the two
+foot boxes touch the floor through explicit `<pair>`s and every other geom has
+`contype=0` from the `g1` default class. That is right for walking (it is much
+cheaper) and useless for fall recovery -- a fallen torso passes straight through
+the floor, so there is nothing to push off.
+
+This module rewrites Playground's robot XML into a get-up variant:
+
+  * turns on the collision geoms the model already carries but leaves disabled
+    (thighs, shins, hands), with condim=3 so they can push rather than slide;
+  * adds the collision bodies it does not carry at all -- a torso capsule, a
+    pelvis capsule and a head sphere, sized from the visual meshes.
+
+Everything else -- actuators, sensors, keyframes, contact pairs -- is inherited
+unchanged, so observations and the action space stay identical to the joystick
+task and a walking policy remains loadable against this model.
+"""
+
+from __future__ import annotations
+
+import re
+
+from mujoco_playground._src.locomotion.g1 import base as g1_base
+from mujoco_playground._src.locomotion.g1 import g1_constants as consts
+
+ROBOT_XML = consts.ROOT_PATH / "xmls" / "g1_mjx_feetonly.xml"
+SCENE_XML = consts.FEET_ONLY_FLAT_TERRAIN_XML
+GETUP_ROBOT_NAME = "g1_mjx_getup.xml"
+
+# Existing collision geoms that the g1 default class disables.
+ENABLE_COLLISION = (
+    "left_thigh", "right_thigh",
+    "left_shin", "right_shin",
+    "left_hand_collision", "right_hand_collision",
+)
+# contype=2 / conaffinity=1 against a floor of contype=1 / conaffinity=1 means
+# "collide with the ground, never with each other": (2 & 1) == 0 for two body
+# geoms, while (floor.contype & body.conaffinity) == 1. Without this the pelvis
+# capsule overlaps the thigh capsules at the standing pose and the robot is
+# blown over by its own self-contacts. Self-collision is out of scope here, as
+# it is for the walking task.
+CONTACT_ATTRS = 'contype="2" conaffinity="1" condim="3" priority="1" friction="0.6"'
+
+# Collision bodies the model lacks, sized from the visual mesh AABBs:
+# pelvis mesh spans z in [-0.14, 0.0] and is ~0.14 wide; the torso mesh spans
+# z in [-0.01, 0.31] with the head above it at 0.385.
+PELVIS_GEOM = (
+    f'<geom name="pelvis_collision" type="capsule" group="3" '
+    f'fromto="0 -0.045 -0.075 0 0.045 -0.075" size="0.085" {CONTACT_ATTRS}/>'
+)
+TORSO_GEOM = (
+    f'<geom name="torso_collision" type="capsule" group="3" '
+    f'fromto="0 0 0.03 0 0 0.30" size="0.11" {CONTACT_ATTRS}/>'
+)
+HEAD_GEOM = (
+    f'<geom name="head_collision" type="sphere" group="3" '
+    f'pos="0.008 0 0.385" size="0.09" {CONTACT_ATTRS}/>'
+)
+
+# Anchor on the body open tags: the visual geoms are duplicated in the XML.
+PELVIS_ANCHOR = re.compile(r'(<body\s+name="pelvis"[^>]*>)')
+TORSO_ANCHOR = re.compile(r'(<body\s+name="torso_link"[^>]*>)')
+
+
+def _enable(xml: str, name: str) -> str:
+    """Give the named geom real contact parameters."""
+    pat = re.compile(rf'(<geom\s+name="{re.escape(name)}"\s+class="collision")')
+    xml, n = pat.subn(rf'\1 {CONTACT_ATTRS}', xml)
+    if n != 1:
+        raise RuntimeError(f"collision geom {name!r} not found in {ROBOT_XML.name} "
+                           f"({n} matches) -- Playground's G1 XML has changed")
+    return xml
+
+
+def build_robot_xml() -> str:
+    # Playground's XML keeps alternative collision shapes commented out, and a
+    # commented `<geom name="left_shin" ...>` is indistinguishable from a live
+    # one to a regex -- drop comments first so the anchors are unambiguous.
+    xml = re.sub(r"<!--.*?-->", "", ROBOT_XML.read_text(), flags=re.S)
+    for name in ENABLE_COLLISION:
+        xml = _enable(xml, name)
+    for anchor, added in ((PELVIS_ANCHOR, PELVIS_GEOM),
+                          (TORSO_ANCHOR, TORSO_GEOM + "\n      " + HEAD_GEOM)):
+        xml, n = anchor.subn(rf"\1\n      {added}", xml)
+        if n != 1:
+            raise RuntimeError(f"anchor {anchor.pattern} matched {n} times, expected 1 -- "
+                               "Playground's G1 XML has changed")
+    return xml
+
+
+def build_scene() -> tuple[str, dict[str, bytes]]:
+    """(scene XML string, assets) for a G1 that collides with the ground."""
+    scene = SCENE_XML.read_text().replace(ROBOT_XML.name, GETUP_ROBOT_NAME)
+    if GETUP_ROBOT_NAME not in scene:
+        raise RuntimeError("scene does not include the robot XML by name")
+    assets = g1_base.get_assets()
+    assets[GETUP_ROBOT_NAME] = build_robot_xml().encode()
+    return scene, assets

rollout.py

@@ -72,8 +72,13 @@ def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0), se
 
     rng = jax.random.PRNGKey(seed)
     state = jit_reset(rng)
+    # The joystick task steers by a "command" in info; the get-up task has no
+    # such key, and adding one would change the State pytree the step function
+    # was traced against.
+    steerable = "command" in state.info
     cmd = jnp.asarray(command, dtype=jnp.float32)
-    state.info["command"] = cmd
+    if steerable:
+        state.info["command"] = cmd
 
     n_steps = int(seconds / env.dt)
     qpos = [np.asarray(state.data.qpos)]
@@ -82,7 +87,8 @@ def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0), se
         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
+        if steerable:
+            state.info["command"] = cmd      # joystick envs resample on reset only
         qpos.append(np.asarray(state.data.qpos))
         if fell_at is None and float(state.done) > 0.5:
             fell_at = (i + 1) * env.dt
@@ -91,6 +97,9 @@ def rollout(env, inference_fn, seconds: float = 8.0, command=(0.5, 0.0, 0.0), se
     qpos = np.stack(qpos)
     info = {"seconds": len(qpos) * env.dt,
             "distance_m": float(np.linalg.norm((qpos[-1, :2] - qpos[0, :2]))),
+            "start_height_m": float(qpos[0, 2]),
+            "end_height_m": float(qpos[-1, 2]),
+            "peak_height_m": float(qpos[:, 2].max()),
             "fell_at": fell_at}
     return qpos, info