Teach a humanoid to run

Code available on GitHub.

A humanoid gait takes hundreds of millions of simulation steps to learn. The classic way to get them is to fan simulators out across many CPU workers and ship their rollouts to a GPU learner. This tutorial does the opposite: MuJoCo MJX compiles the physics into the same XLA program as the learner, so 8,192 copies of the Unitree G1 step in parallel on one GPU and there is no Python between simulation, scoring, and the gradient step. Flyte’s job is everything around that loop: the GPU, the image, a live report you can watch the robot learn in, CPU replays that prove the result, and a versioned artifact for the trained policy.

By the end you will have:

a running policy for the G1, trained with Brax PPO on DeepMind’s tuned MuJoCo Playground environment;
a training report that streams the reward curve and films the policy every few evaluations;
three replays filmed side on, a random policy and the trained policy at walking and running speed, each measuring the speed it reached and how often both feet left the ground;
the policy published as the model artifact g1-run-policy, with a model card.

Why one GPU, and where Flyte fits

Approach	Where the simulator runs	What limits it
CPU simulators, GPU learner	Many CPU workers	Moving rollouts to the learner, and Python in every step
MJX on the GPU (this tutorial)	Inside the learner’s XLA program	GPU memory for the parallel worlds

The single-GPU design keeps the loop simple, but it does not make the infrastructure go away. The image needs JAX with CUDA, MuJoCo, Brax, Playground, an EGL renderer, and the robot models. Rendering needs a different backend than training. And a policy is only worth something if you can see what it does and load it later. Each of those is one Flyte construct:

Need	Construct
A reproducible JAX, MuJoCo, and EGL image with the robot models baked in	`flyte.Image` layers, built remotely
One GPU for training, CPUs for filming	One `flyte.TaskEnvironment` per role
Watch the gait change while it trains	`report=True` and `flyte.report.replace` from a Brax callback
Prove the result next to a baseline	Three replay tasks run in parallel with `asyncio.gather`
Keep and find the trained policy	`flyte.artifacts` with a model card

            flowchart LR
    D["learn_to_run<br/>CPU driver"] --> T["train_policy<br/>1 GPU, 8,192 MJX worlds<br/>live report"]
    T -- "checkpoint (File)" --> R1["render_replay<br/>random policy"]
    T --> R2["render_replay<br/>walk command"]
    T --> R3["render_replay<br/>run command"]
    R1 & R2 & R3 --> P["publish_policy<br/>g1-run-policy artifact"]

    classDef gpu fill:#3b82f633,stroke:#3b82f6,stroke-width:1.5px
    classDef cpu fill:#94a3b826,stroke:#94a3b8,stroke-width:1px
    class T gpu
    class D,R1,R2,R3,P cpu

Getting started

You need a Union cluster with an NVIDIA L40S, or another GPU with at least 40 GB of memory. To use a different card, change gpu_env in learn_to_run.py.

        
    
git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/humanoid_locomotion
python learn_to_run.py

learn_to_run.py
                
if __name__ == "__main__":
    flyte.init_from_config()
    if "--checkpoint" in sys.argv:
        uri = sys.argv[sys.argv.index("--checkpoint") + 1]
        run = flyte.run(learn_to_run, checkpoint=File.from_existing_remote(uri))
    elif "--smoke" in sys.argv:
        run = flyte.run(learn_to_run, num_timesteps=20_000_000, publish=False)
    else:
        run = flyte.run(learn_to_run)
    print(run.url)

The first run builds the image, which takes a while because of the CUDA build of JAX. Later runs reuse it. --smoke trains 20 million steps in a few minutes, which is enough to check the plumbing but not enough for the robot to stay up. --checkpoint skips training and only replays and reports an existing policy.

The image and environments

The image pins JAX to 0.9 because Brax 0.14 still calls a JAX API that 0.10 removed. The last layer downloads the robot models into Playground’s package directory at build time. Without it, every pod would clone the mujoco_menagerie repository at startup.

learn_to_run.py
                
# brax 0.14.2 still calls jax.device_put_replicated, which jax removed in 0.10, so jax stays
# on 0.9.x until brax moves off the old pmap API.
image = (
    flyte.Image.from_debian_base(name="g1-mjx")
    .with_apt_packages("git", "libegl1", "libegl-mesa0", "libgl1", "libgl1-mesa-dri", "libgles2", "libglx-mesa0")
    .with_pip_packages(
        "jax[cuda12]==0.9.2", "mujoco==3.17.0", "mujoco-mjx==3.17.0", "brax==0.14.2", "playground==0.2.0",
        "numpy", "Pillow", "av",
    )
    # Bake the robot models into the image. Otherwise every pod clones mujoco_menagerie at startup.
    .with_commands(["python -c 'from mujoco_playground._src import mjx_env; mjx_env.ensure_menagerie_exists()'"])
)

Three environments share the image. Training gets the GPU, with JAX told not to preallocate most of the card, because MJX’s Warp kernels allocate memory outside JAX’s pool. Replays run MJX’s JAX backend on CPU and render with MuJoCo over headless EGL, since MJX has no renderer. The driver only waits on the other tasks, so it is small.

learn_to_run.py
                
                    
                
            
                
            

                
                    
                
            
EGL = {"MUJOCO_GL": "egl", "PYOPENGL_PLATFORM": "egl"}

# Physics and gradients, both on the GPU. MJX's Warp kernels allocate outside JAX's memory
# pool, so JAX must not preallocate most of the card.
gpu_env = flyte.TaskEnvironment(
    name="g1-run-train",
    image=image,
    resources=flyte.Resources(cpu=8, memory="48Gi", gpu="L40s:1", disk="50Gi"),
    env_vars={"XLA_PYTHON_CLIENT_PREALLOCATE": "false", "XLA_PYTHON_CLIENT_MEM_FRACTION": "0.75", **EGL},
)

# Replays. MJX has no renderer, so the policy runs on MJX's JAX backend on CPU and MuJoCo's
# renderer draws it over headless EGL.
render_env = flyte.TaskEnvironment(
    name="g1-run-render",
    image=image,
    resources=flyte.Resources(cpu=4, memory="16Gi"),
    env_vars={"JAX_PLATFORMS": "cpu", **EGL},
)

# The driver holds its pod while the other tasks run, so it stays small and CPU-only.
driver_env = flyte.TaskEnvironment(
    name="g1-run-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="8Gi"),
    env_vars={"JAX_PLATFORMS": "cpu"},
    depends_on=[gpu_env, render_env],
)

The running environment

The environment is DeepMind’s G1JoystickFlatTerrain: the policy gets a velocity command and proprioception, and it is rewarded for tracking the command while staying upright, with feet in the air on schedule. Upstream it learns to walk at up to 1 m/s. Two changes make it run:

forward commands go up to max_speed, 2.5 m/s by default;
the gait clock speeds up with the command, so a fast command asks for a faster cadence instead of ever-longer strides.

It also raises njmax, the constraint budget: upstream’s value is a little too small for a G1 in ground contact, and when MJX runs out it drops constraints silently and the feet sink into the floor. And it turns off a Warp solver warning that is printed for every world on every step, which otherwise dominates the wall-clock time.

g1.py
                
                    
                
            
                
            

                
                    
                
            
BASE_ENV = "G1JoystickFlatTerrain"
WALK_SPEED = 1.0
# The gait clock: 1.25 Hz at a standstill (upstream's lower bound), rising with the command.
# At 2.5 m/s that is ~2.1 Hz, a 0.6 m step, which is a run for a 1.3 m robot.
CADENCE_HZ = (1.25, 2.25)
CADENCE_PER_MPS = 0.35

def env_config(max_speed: float, impl: str = "warp"):
    """Upstream's G1 config, with run-speed forward commands and the njmax fix."""
    from mujoco_playground import registry

    cfg = registry.get_default_config(BASE_ENV)
    cfg.impl = impl
    # Upstream ships njmax=90 and a G1 in ground contact needs ~93; on overflow MJX drops
    # constraints silently and the feet sink into the floor.
    cfg.njmax = 128
    cfg.lin_vel_x = [0.0, max_speed]   # forward only: walk through run
    cfg.lin_vel_y = [-0.3, 0.3]
    cfg.ang_vel_yaw = [-0.5, 0.5]
    # At run speed the tracking term is the one worth chasing, and air time pays for flight.
    cfg.reward_config.scales.tracking_lin_vel = 1.5
    cfg.reward_config.scales.feet_air_time = 2.5
    return cfg

def make_env(max_speed: float, impl: str = "warp"):
    import jax.numpy as jp
    from mujoco_playground._src.locomotion.g1 import joystick

    class G1Run(joystick.Joystick):
        """G1 joystick env whose gait clock speeds up with the commanded forward speed."""

        def __init__(self, **kw):
            super().__init__(**kw)
            if impl == "warp":
                # MuJoCo-Warp printf's a solver-iteration warning from the kernel for every
                # world on every step: at 8,192 worlds that is most of the wall clock.
                self._mjx_model = self._mjx_model.tree_replace({"opt._impl.warn_overflow": 0})

        def _phase_dt(self, command):
            hz = jp.clip(CADENCE_HZ[0] + CADENCE_PER_MPS * jp.abs(command[0]), *CADENCE_HZ)
            return (2 * jp.pi * self.dt * hz).reshape(1)

        def reset(self, rng):
            state = super().reset(rng)
            state.info["phase_dt"] = self._phase_dt(state.info["command"])
            return state

        def step(self, state, action):
            state = super().step(state, action)
            state.info["phase_dt"] = self._phase_dt(state.info["command"])
            return state

    return G1Run(task="flat_terrain", config=env_config(max_speed, impl))

Training with a live report

train_policy is one Brax ppo.train call. Brax runs the whole loop on the GPU and calls back into Python twice: progress_fn after every evaluation, with the metrics, and policy_params_fn, with the current weights. The task uses the first to redraw its report, and the second to film the policy on the run command every few evaluations. A failed snapshot is logged and skipped, so filming can never kill a long run.

learn_to_run.py
                
                    
                
            
                
            

                
                    
                
            
@gpu_env.task(report=True, retries=1)
async def train_policy(
    num_timesteps: int = 150_000_000,
    num_envs: int = 8192,
    num_evals: int = 20,
    max_speed: float = 2.5,
    seed: int = 0,
    snapshot_every: int = 4,
    snapshot_steps: int = 200,
) -> File:
    """Train one G1 running policy on one GPU and return its checkpoint.

    The reward curve streams into this task's report at every eval. Every ``snapshot_every``
    evals, the current policy is filmed on the run command.
    """
    import functools

    import jax
    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

    log.info(f"devices: {jax.devices()}")
    env = make_env(max_speed)
    params = ppo_config(num_timesteps, num_envs, num_evals)
    net_cfg = dict(params.network_factory)
    history: list[dict] = []
    snaps: list[dict] = []
    latest: dict = {}
    started = time.time()
    evals = {"n": 0}

    def progress(step: int, metrics: dict) -> None:
        row = {"step": int(step), "reward": float(metrics.get("eval/episode_reward", 0.0)),
               "elapsed_s": round(time.time() - started, 1)}
        for k in ("tracking_lin_vel", "feet_air_time"):
            if f"eval/episode_reward/{k}" in metrics:
                row[k] = float(metrics[f"eval/episode_reward/{k}"])
        history.append(row)
        log.info(f"  step {row['step']:>12,} | reward {row['reward']:8.2f} | {row['step'] / max(row['elapsed_s'], 1):,.0f} steps/s")
        flyte.report.replace(live_html(history, snaps, latest, started, num_envs, num_timesteps, max_speed))
        flyte.report.flush()

    def policy_params_fn(step: int, make_policy, p) -> None:
        # The only Brax hook that hands over live weights. Filming must never kill the run.
        evals["n"] += 1
        if not snapshot_every or step == 0 or (evals["n"] - 1) % snapshot_every:
            return
        try:
            policy = jax.jit(make_policy(p, deterministic=True))
            r = rollout(env, lambda obs, key: policy(obs, key)[0], (max_speed, 0.0, 0.0), snapshot_steps)
            frames = render_side(env, r["traj"], 240, 320)
            latest.update(step=step, speed=r["speed"], flight=r["flight"], mp4=encode_mp4(frames, int(round(1 / env.dt))))
            snaps.append({"step": step, "speed": r["speed"], "flight": r["flight"], "still": still(frames[len(frames) * 2 // 3])})
        except Exception as exc:
            log.warning(f"  snapshot at {step:,} failed ({exc}); training continues")

    _, trained, _ = ppo.train(
        environment=env,
        progress_fn=progress,
        policy_params_fn=policy_params_fn,
        network_factory=functools.partial(ppo_networks.make_ppo_networks, **net_cfg),
        randomization_fn=registry.get_domain_randomizer(BASE_ENV),
        # Playground's wrapper applies the domain randomization. Brax's default drops it.
        wrap_env_fn=wrapper.wrap_for_brax_training,
        seed=seed,
        **{k: v for k, v in params.items() if k != "network_factory"},
    )
    train_s = time.time() - started
    log.info(f"trained {num_timesteps:,} steps in {train_s / 60:.1f} min")

    path = "/tmp/g1_run_policy.pkl"
    with open(path, "wb") as f:
        pickle.dump({
            "params": jax.device_get(trained), "network_factory": net_cfg, "history": history,
            "snapshots": snaps, "max_speed": max_speed, "num_timesteps": num_timesteps,
            "num_envs": num_envs, "seed": seed, "train_s": round(train_s, 1),
        }, f)
    return await File.from_local(path)

The policy is asymmetric: the actor sees sensor-shaped observations, and the critic also sees privileged simulator state, which it only needs during training. The checkpoint is a pickle of the parameters, the network configuration, the training history, and the snapshots, returned as a flyte.io.File.

Replays: proof, not a number

An eval reward says little on its own. The driver films three replays in parallel on CPU: a random policy, and the trained policy holding a walking and a running command. Each replay measures what the robot actually did.

g1.py
                
                    
                
            
                
            

                
                    
                
            
def rollout(env, act, command, steps: int, seed: int = 0) -> dict:
    """Hold one command for ``steps`` and measure what the robot actually did."""
    import jax
    import jax.numpy as jp
    import numpy as np

    reset, step = jax.jit(env.reset), jax.jit(env.step)
    feet = [env.mj_model.sensor_adr[s] for s in env._feet_floor_found_sensor]
    cmd = jp.array(command, dtype=jp.float32)
    rng = jax.random.PRNGKey(seed)
    state = reset(rng)
    state.info["command"] = cmd
    traj, speed, airborne, reward = [state], [], [], 0.0
    for _ in range(steps):
        rng, key = jax.random.split(rng)
        state = step(state, act(state.obs, key))
        state.info["command"] = cmd  # the env would resample it mid-episode
        traj.append(state)
        reward += float(state.reward)
        speed.append(float(env.get_local_linvel(state.data, "pelvis")[0]))
        airborne.append(not any(float(state.data.sensordata[a]) > 0 for a in feet))
        if float(state.done) > 0.5:
            break
    start, end = np.asarray(traj[0].data.qpos[:2]), np.asarray(traj[-1].data.qpos[:2])
    settled = speed[len(speed) // 4:] or speed  # skip the first second of getting up to speed
    return {
        "traj": traj, "steps": len(traj) - 1, "reward": reward, "fell": float(state.done) > 0.5,
        "speed": float(np.mean(settled)) if settled else 0.0,
        "flight": float(np.mean(airborne[len(airborne) // 4:] or airborne)) if airborne else 0.0,
        "distance": float(np.linalg.norm(end - start)),
    }

The flight phase is the share of steps with both feet off the ground. It is what separates a run from a fast walk.

learn_to_run.py
                
                    
                
            
                
            

                
                    
                
            
@render_env.task(retries=2)
async def render_replay(checkpoint: File, speed: float, policy: str = "trained", steps: int = 500) -> str:
    """Roll one policy out on one forward command, and return the clip and what it measured."""
    import jax

    with open(await checkpoint.download(), "rb") as f:
        ckpt = pickle.load(f)
    env = make_env(ckpt["max_speed"], impl="jax")
    if policy == "trained":
        fn = inference_fn(ckpt["network_factory"], ckpt["params"], env)

        def act(obs, key):
            return fn(obs, key)[0]
    else:
        def act(obs, key):
            return jax.random.uniform(key, (env.action_size,), minval=-1.0, maxval=1.0)

    r = rollout(env, act, (speed, 0.0, 0.0), steps)
    frames = render_side(env, r.pop("traj"), 360, 480)
    label = f"{policy} policy, {speed:.1f} m/s command"
    log.info(f"{label}: {r['steps']}/{steps} steps, {r['speed']:.2f} m/s, flight {r['flight']:.0%}, {r['distance']:.1f} m")
    return json.dumps({**r, "policy": policy, "command": [speed, 0.0, 0.0],
                       "html": video(encode_mp4(frames, int(round(1 / env.dt))), label),
                       "still": still(frames[len(frames) * 2 // 3], 360)})

The driver and the artifact

The driver trains, replays, writes the final report, and publishes the checkpoint as a model artifact. A later run, an evaluation job, or an app can load it with flyte.remote.Artifact.get("g1-run-policy").

learn_to_run.py
                
                    
                
            
                
            

                
                    
                
            
@driver_env.task(report=True)
async def learn_to_run(
    num_timesteps: int = 150_000_000,
    num_envs: int = 8192,
    num_evals: int = 20,
    max_speed: float = 2.5,
    seed: int = 0,
    replay_steps: int = 500,
    publish: bool = True,
    checkpoint: File | None = None,
) -> str:
    """Train a G1 running policy, replay it next to a random policy, and publish it.

    Pass ``checkpoint`` (a ``train_policy`` output) to skip training and only replay and report.
    """
    if checkpoint is None:
        await flyte.report.replace.aio(viz.page(
            "Unitree G1 learns to run",
            f"Training {viz.compact(num_timesteps)} steps across {num_envs:,} parallel G1s on one GPU. "
            "The live curve and snapshots are in the train_policy action's report."), do_flush=True)
        checkpoint = await train_policy(num_timesteps=num_timesteps, num_envs=num_envs, num_evals=num_evals,
                                        max_speed=max_speed, seed=seed)
    with open(await checkpoint.download(), "rb") as f:
        ckpt = pickle.load(f)
    max_speed = ckpt["max_speed"]

    # Three independent CPU replays. A trained clip means little without the random one next to it.
    rand, walk, run = await asyncio.gather(
        render_replay(checkpoint, max_speed, "random", replay_steps),
        render_replay(checkpoint, WALK_SPEED, "trained", replay_steps),
        render_replay(checkpoint, max_speed, "trained", replay_steps),
    )
    clips = {"random": json.loads(rand), "walk": json.loads(walk), "run": json.loads(run)}
    ckpt["replay_steps"] = replay_steps
    await flyte.report.replace.aio(final_html(ckpt, clips, flyte.ctx().action.run_name), do_flush=True)

    if publish:
        await publish_policy(checkpoint, json.dumps({k: {kk: vv for kk, vv in c.items() if kk != "html"} for k, c in clips.items()}))

    summary = {k: {kk: c[kk] for kk in ("command", "speed", "flight", "distance", "steps", "fell")} for k, c in clips.items()}
    summary["best_reward"] = max(h["reward"] for h in ckpt["history"])
    summary["train_min"] = round(ckpt["train_s"] / 60, 1)
    return json.dumps(summary)

learn_to_run.py
                
                    
                
            
                
            

                
                    
                
            
@driver_env.task(produces_artifacts=True)
async def publish_policy(checkpoint: File, clips_json: str) -> File:
    """Register the checkpoint as the model artifact ``g1-run-policy``, with a model card."""
    clips = json.loads(clips_json)
    with open(await checkpoint.download(), "rb") as f:
        ckpt = pickle.load(f)
    run_name = flyte.ctx().action.run_name
    card = await artifacts.Card.create_from.aio(content=policy_card(ckpt, clips, run_name), format="html", card_type="model")
    metadata = artifacts.Metadata.create_model_metadata(
        name="g1-run-policy",
        description=f"Unitree G1 running policy: {clips['run']['speed']:.2f} m/s at a {clips['run']['command'][0]:.1f} m/s command",
        framework="JAX",
        model_type="actor-critic MLP",
        architecture="Brax PPO, asymmetric actor-critic",
        task="humanoid locomotion",
        modality=("proprioception",),
        serial_format="pkl",
        card=card,
        attrs={"run_speed": f"{clips['run']['speed']:.2f}", "flight": f"{clips['run']['flight']:.2f}",
               "env_steps": str(ckpt["num_timesteps"])},
    )
    return artifacts.new(checkpoint, metadata)

Results

A run with the defaults trained 150 million steps in 19 minutes on one L40S, and took 36 minutes end to end. Each replay holds one command for 10 seconds (500 control steps):

Replay	Speed reached	Flight phase	Distance	Upright
Random policy, 2.5 m/s command	0.05 m/s	n/a	0.1 m	fell at step 8
Trained, 1.0 m/s command	0.97 m/s	7%	8.9 m	500 of 500 steps
Trained, 2.5 m/s command	2.42 m/s	15%	22.3 m	500 of 500 steps

At the run command the robot reaches 97% of the commanded speed, and both feet are off the ground in 15% of steps, twice as often as at walking speed. For a robot with legs about 0.7 m long, 2.4 m/s is a jog (a Froude number of about 0.85; people switch from walking to running at about 0.5), not a sprint.

The live report of train_policy shows progress, the latest filmed snapshot, and a filmstrip of the gait as it forms:

The training report: progress tiles, the latest filmed snapshot, and the gait filmstrip

The driver’s final report puts the three replays side by side:

The final report: running speed and flight phase tiles, and the random, walk, and run replays

The g1-run-policy artifact carries the same results in its model card:

The g1-run-policy model card

Going further

Train longer. The speed and the flight phase keep improving past 150 million steps. A billion steps takes under 2 hours on one L40S.
Rough terrain. Playground has a G1JoystickRoughTerrain environment. Change BASE_ENV, and train from the flat-terrain checkpoint.
Scale out instead. To run a simulator that cannot be compiled to the GPU, fan rollouts out to a warm pool of CPU workers, as the GRPO tutorial does for LLM generation.
Sweep it. Every train_policy call is independent, so a hyperparameter sweep is an asyncio.gather over several calls, each on its own GPU.

Reinforcement learning in the user guide
Reports
Artifacts