Skip to content

36 · Legged locomotion and massively parallel training

Part8 · Reinforcement learning for robots
Run typeGPU · Colab GPU
Prerequisites34, 35
Learn firstIntroduction to JAX · Proximal Policy Optimization (PPO)
Time~90 min
Codeexperiments/36-legged-locomotion/experiment.py
Status✅ done

Act on this lesson

Tick each step when you have done it. Saved in this browser only.

A walking policy needs hundreds of millions of steps. The engineering question is where those steps come from.

Lesson 33 trained an arm reach with 300,000 policy steps. Unitree’s Go2 training configuration asks for a different order of magnitude: num_envs = 4096 simulated robots, num_steps_per_env = 24 per iteration and max_iterations = 1500 (legged_robot_config.py), which is 147,456,000 policy steps, or 589,824,000 physics steps at decimation = 4. Rudin et al. showed that running thousands of robots in parallel on one workstation GPU trains a flat-terrain quadruped policy in under four minutes (arXiv 2109.11978, CoRL 2021). This lesson explains what a legged robot adds to the RL problem, measures how far a CPU gets on the Go2 model, and shows how to run the GPU path with MuJoCo Playground.

Before you read · 2 questions. Guess first; the lesson tells you why.

1. A dog-like robot’s body is not bolted to anything. How does the policy make it move forward?

  • It commands a force on the body directly.
  • It sends joint targets; PD turns them into torques, and the feet, which can only push and only as hard as friction allows, move the body.
  • It sets the base velocity, which the simulator integrates.

Why: Core idea: the base is a floating body moved only by unilateral, friction-limited foot contacts. Standing, the four feet carry 149.2 N, exactly m g (section A).

2. Why does locomotion PPO run thousands of simulated robots in parallel instead of one robot for longer?

  • Each robot needs a different reward function.
  • A single simulated robot cannot run longer than 20 s.
  • PPO’s update is only valid for samples from (nearly) the current policy, and one robot’s consecutive steps are nearly the same data; many robots give many independent trajectories per update.

Why: Math details, on-policy sampling. The Go2 config collects 4096 × 24 = 98,304 samples per update, while each robot advances only 0.48 s.

An arm is bolted to the table; a dog-like robot is not. Its body (the base) floats: no motor holds it, and it moves only because the feet push on the floor. A foot can push but never pull, and it can push sideways only as hard as friction allows before it slips. Walking is a schedule of which feet are on the ground and which are swinging, called a gait. The policy never commands those foot forces. It sends joint targets, a PD loop turns them into torques, and the physics of the contacts turns the torques into forces on the body.

Learning that from trial and error takes on the order of a hundred million steps. PPO needs fresh attempts from the current policy for every update, and consecutive steps of one robot are nearly the same data. So the practical question becomes throughput, and the answer is to scale out like a stateless web service: run thousands of independent robot copies at once and step them all with one batched call.

flowchart TD
    P["policy, every 20 ms<br/>12 numbers a"] --> Q["joint targets<br/>q* = 0.25·a + q_default"]
    Q --> D["PD at every physics step<br/>τ = 20 (q* − q) − 0.5 q̇"]
    D --> J["12 leg joints move"]
    J --> C["foot contacts:<br/>push only, friction-limited"]
    C --> B["floating base moves"]
    B --> O["observation: base velocity,<br/>gravity direction, joints, command"]
    O --> P
num_envs, steps_per_env = 4096, 24 # Go2 training config
obs = envs.reset() # shape (4096, obs_dim): one row per simulated robot
for t in range(steps_per_env):
actions = policy(obs) # ONE batched network call for all robots
obs, rewards, dones = envs.step(actions) # ONE batched physics call (on a GPU with MJX)
samples = num_envs * steps_per_env # 98,304 samples for this PPO update
seconds_per_robot = steps_per_env * 0.02 # while each robot advances only 0.48 s
QuantityFormula
a contact only pushesfn≥0f_n \ge 0
friction limit∥ft∥≤μfn\lVert f_t \rVert \le \mu f_n
standing still∑fn=mg\sum f_n = m g
samples per updateNenv×TN_{\text{env}} \times T

fnf_n and ftf_t are the normal and tangential (sideways) parts of a foot’s contact force, μ\mu the friction coefficient, mm the robot’s mass, gg gravity, NenvN_{\text{env}} the number of parallel robots and TT the steps each one collects per iteration.

Math details: floating base, friction cones and why on-policy RL wants many copies

Floating base. The Go2 model has nq = 19, nv = 18, nu = 12: 12 actuated joints plus a 6-DoF base that no motor drives (nq is one larger than nv because the base orientation is stored as a 4-number quaternion). The base accelerates only through contact forces at the feet, which must lie inside the friction cone ∥ft∥≤μfn\lVert f_t \rVert \le \mu f_n (lesson 04). In a trot, diagonal pairs of feet alternate between stance (contact) and swing.

On-policy sampling. PPO’s update is only valid for samples from (nearly) the current policy, so old data cannot be reused across iterations. Samples from one robot are strongly correlated in time; 4096 robots give 4096 independent trajectories per update. The bottleneck becomes raw simulation throughput, which is why the physics, the policy and the PPO update all run batched on a GPU.

MJX and MuJoCo Playground. MJX runs MuJoCo in JAX (and a Warp version for NVIDIA GPUs); its documentation says it pays off when simulating thousands of scenes in parallel, while a single scene can be about 10× slower than CPU MuJoCo (MJX docs). MuJoCo Playground (Zakka et al., 2025, arXiv 2502.08844, playground.mujoco.org) builds environments and PPO training on MJX and reports training “in minutes on a single GPU”. Its default locomotion PPO config uses num_envs = 8192; for Go1JoystickFlatTerrain it sets num_timesteps = 200_000_000 (locomotion_params.py). The Playground locomotion registry has Go1, Spot and Barkour quadrupeds, but no Go2 environment (locomotion/__init__.py).

Go deeper: Introduction to JAX, Proximal Policy Optimization (PPO); Rudin et al. (arXiv 2109.11978); Zakka et al. (arXiv 2502.08844).

What the Go2 config rewards and randomizes

Section titled “What the Go2 config rewards and randomizes”

(legged_robot_config.py, go2_config.py)

ItemValue
velocity trackingtracking_lin_vel = 1.0, tracking_ang_vel = 0.5
gait shapingfeet_air_time = 1.0 (at touchdown, rewards swing time beyond 0.5 s, only when a velocity is commanded)
regularizersaction_rate = -0.01, torques = -0.0002 (Go2), dof_acc = -2.5e-7, collision = -1.
sign handlingonly_positive_rewards = True
friction randomizationrandomize_friction = True, friction_range = [0.5, 1.25]
pushespush_robots = True, every push_interval_s = 15, max_push_vel_xy = 1.
episodeepisode_length_s = 20; ends early on base contact (terminate_after_contacts_on = ["base"]) or when roll exceeds 0.8 rad or pitch 1.0 rad (check_termination in legged_robot.py)

Every row is a lesson-34 or lesson-35 decision. Randomizing friction over [0.5,1.25][0.5, 1.25] means each of the 4096 robots learns on a slightly different floor (Part 9).

You will now write the vectorized environment that the “as code” loop above assumes: one model, one MjData per robot, batched actions in, batched observations out, episodes that end and restart on their own, and counters that say how much was simulated. The robot is the Experiment’s PD-servo Go2 (load_go2). Nothing here measures time; the Experiment does that. The code is in physai/vecenv.py; uv run python experiments/36-legged-locomotion/build.py runs the steps (output in build_output.txt).

An MjModel holds what never changes (masses, geometry, gains); an MjData holds one robot’s state. N robots therefore need one model and N datas, and the batch step is a loop over them:

class VecEnv:
"""N robots, one model. reset_fn(m, d, rng) puts one MjData into a start state."""
def __init__(self, m, n, reset_fn, rng, decimation=4, max_steps=1000, max_roll=0.8, max_pitch=1.0,
min_height=0.15):
self.m, self.n, self.reset_fn, self.rng = m, n, reset_fn, rng
self.decimation, self.max_steps = decimation, max_steps
self.max_roll, self.max_pitch, self.min_height = max_roll, max_pitch, min_height
self.datas = [mujoco.MjData(m) for _ in range(n)]
self.t = np.zeros(n, dtype=int) # policy steps in the current episode, per robot
self.policy_steps = self.physics_steps = self.episodes = 0
def obs_one(self, d):
"""Base height, projected gravity, base velocity (6), joint angles and speeds."""
return np.r_[d.qpos[2], projected_gravity(d.qpos[3:7]), d.qvel[:6], d.qpos[7:], d.qvel[6:]]
def reset(self):
for d in self.datas:
self.reset_fn(self.m, d, self.rng)
self.t[:] = 0
return np.array([self.obs_one(d) for d in self.datas])
def step(self, actions):
"""actions: (n, nu) control targets. Returns obs, fell, timed_out (one row / entry per robot)."""
for d, a in zip(self.datas, actions):
d.ctrl[:] = a
mujoco.mj_step(self.m, d, nstep=self.decimation)
self.t += 1
self.policy_steps += self.n
self.physics_steps += self.n * self.decimation
fell, timed_out = np.zeros(self.n, bool), self.t >= self.max_steps
for i, d in enumerate(self.datas):
roll, pitch = roll_pitch(d.qpos[3:7])
fell[i] = abs(roll) > self.max_roll or abs(pitch) > self.max_pitch or d.qpos[2] < self.min_height
for i in np.flatnonzero(fell | timed_out): # auto-reset in the same step
self.reset_fn(self.m, self.datas[i], self.rng)
self.t[i] = 0
self.episodes += 1
return np.array([self.obs_one(d) for d in self.datas]), fell, timed_out
== Step 1. One model, 8 MjData: the batch vs 8 independent loops (100 policy steps) ==
max |qpos batch - qpos alone| over 8 robots: 0.0e+00
episodes ended: 0 (base heights 26.0, 26.1, 25.9, 26.3, 26.3, 26.2, 26.0, 26.5 cm)

The batch is bit for bit the same as eight separate simulations: robots in one batch do not interact, they only share the model. With small random targets nobody fell, and the base heights stay near the Experiment’s 26.2 cm.

The observation carries the gravity direction in the base frame (the policy cannot see the world frame) and the episode ends when roll or pitch passes legged_gym’s limits (0.8 and 1.0 rad). Both come from the base quaternion, through quat_to_mat from lesson 06:

def projected_gravity(quat):
"""Gravity direction (0, 0, -1) seen in the base frame: R^T g for the base quaternion (w, x, y, z)."""
return quat_to_mat(quat).T @ np.array([0.0, 0.0, -1.0])
def roll_pitch(quat):
"""Roll and pitch of the base (rad), from R = Rz(yaw) Ry(pitch) Rx(roll)."""
R = quat_to_mat(quat)
return np.arctan2(R[2, 1], R[2, 2]), np.arcsin(np.clip(-R[2, 0], -1.0, 1.0))
== Step 2. Projected gravity and roll / pitch from the base quaternion ==
projected gravity vs MuJoCo xmat^T g: max diff 3.9e-16
roll, pitch vs SciPy as_euler('ZYX'): max diff 1.6e-15
standing start: projected gravity [ 0. 0. -1.], base height 27.0 cm
== Step 3. Auto-reset: 16 robots standing, a roll / pitch kick of up to 18 rad/s every 25 steps, episodes of at most 150 steps ==
300 policy steps (2.4 s per robot): 22 falls, 18 time-outs, episodes ended 40
first fall after 13 policy steps; every ended row came back at base height 27.0-27.0 cm (the keyframe height) in the same step
steps since reset now, per robot: [0, 0, 0, 12, 135, 0, 143, 22, 12, 48, 65, 108, 49, 35, 0, 121]

Twenty-two falls, each reset in the step it happened, so the returned observation is already the fresh start (27.0 cm, the keyframe height); a robot that stays up times out after 150 steps. After the falls the episode clocks are scattered (12, 135, 143, …). That is the normal state of a vectorized environment: every robot is at a different point of its own episode, which is what makes 4096 of them 4096 different samples.

def samples_per_update(n_envs, steps_per_env):
"""On-policy batch size: every robot contributes steps_per_env transitions to one PPO update."""
return n_envs * steps_per_env
== Step 4. Accounting ==
Step 3 counters: policy steps 4,800 = 16 x 300, physics steps 19,200 = x4, simulated time 38.4 robot-seconds
Go2 training config: samples per update 98,304, policy steps 147,456,000, physics steps 589,824,000
each robot advances 24 x 0.02 s = 0.48 s per update

The counters are exact by construction and give the arithmetic of The problem: 98,304 samples per PPO update, 589,824,000 physics steps for one run of the config. Multiply by the Experiment’s 11–17 µs of CPU per Go2 step to price it.

physai.vecenv now has VecEnv, projected_gravity, roll_pitch and samples_per_update. Part 9 randomizes the model fields of such a batch (lesson 39). tests/test_36_legged_locomotion.py checks them against SciPy and mujoco.rollout: uv run pytest tests/test_36_legged_locomotion.py.

For open-loop batches MuJoCo’s C rollout does the same loop without Python, over a thread pool; Gymnasium wraps N environments with a vector API.

state, _ = rollout.rollout(m, [mujoco.MjData(m) for _ in range(4)], x0, U) # x0: (n, nstate), U: (n, steps, nu)
envs = gymnasium.vector.SyncVectorEnv([make_env] * n, autoreset_mode="SameStep") # default is "NextStep"
== Use it. mujoco.rollout: the same batch in C, open loop ==
rollout (4 threads) vs VecEnv after 100 policy steps: max |qpos diff| = 0.0e+00
ToolWhen an episode endsThe observation returned in that step
physai VecEnvreset in the same stepthe fresh start
unitree_rl_gymreset_idx in the same step, then compute_observations (legged_robot.py)the fresh start for joints; base velocity and projected gravity are computed before the reset
Gymnasium 1.4 vector, default NextStepreset on the next step() call, whose action is ignored and whose reward is 0the terminal observation
Gymnasium 1.4 vector, SameStepreset in the same stepthe fresh start; the terminal one is in info["final_obs"]

(The two Gymnasium rows were checked with CartPole-v1 and Gymnasium 1.4.0; mujoco.rollout is open loop, so a policy that reacts to observations needs the Python loop or MJX.)

Traps:

  • Which observation is terminal. With next-step reset, the value target of the last transition must use the terminal observation; with same-step reset it must come from final_obs, not from the returned one.
  • Time-outs are not failures. A robot that timed out could have kept going; PPO should bootstrap its value there, unlike after a fall. Keep the two flags apart, as step does.
  • nq vs nv. The Go2 has 19 positions but 18 velocities: the base quaternion has four numbers and three rotational speeds, so slice qpos[7:] but qvel[6:] for the joints.

Check after building · 2 questions.

1. In Build it, Step 1, eight robots stepped as one batch end bit for bit where eight separate simulations end (max difference 0.0e+00). What does that tell you about a batch?

  • Robots in a batch share one MjModel but never interact; each MjData is an independent simulation.
  • The robots collide with each other but the effect is too small to see.
  • MuJoCo averages the robots’ states.

Why: The model holds what never changes, the data one robot’s state. mujoco.rollout over the same batch also matches to 0.0e+00 (Use it).

2. Gymnasium 1.4’s vector environments default to autoreset_mode NextStep. What does the step() that ends an episode return?

  • The fresh start of the next episode, as physai’s VecEnv does.
  • The terminal observation; the reset happens on the next step() call, whose action is ignored and whose reward is 0.
  • Nothing; the environment has to be reset by hand.

Why: Use it table, checked with CartPole-v1. In SameStep mode the returned observation is the fresh start and the terminal one is in info[“final_obs”]; legged_gym also resets in the same step.

CPU part (measured): how many Go2 steps per second?

Section titled “CPU part (measured): how many Go2 steps per second?”

experiment.py loads the Menagerie Go2 (2 ms timestep, elliptic friction cones, impratio = 100), turns its 12 motors into PD servos with the training gains (Kp 20, Kd 0.5, torque capped at the motor limits) around the config’s default joint angles, and lets it stand. Because the PD lives inside MuJoCo, a constant control vector is a valid closed-loop stance, so batches of robots can be stepped by mujoco.rollout, MuJoCo’s C++ rollout module (Python docs): one MjData per thread, the GIL released, an internal thread pool.

for i in range(m.nu): # motor -> PD servo, keep the torque limit
m.actuator_gainprm[i, :3] = [KP, 0, 0]
m.actuator_biasprm[i, :3] = [0, -KP, -KD]
m.actuator_biastype[i] = mujoco.mjtBias.mjBIAS_AFFINE
...
datas = [mujoco.MjData(m) for _ in range(nthread)]
state, _ = rollout.rollout(m, datas, X0, U) # X0: (nbatch, nstate), U: (nbatch, nstep, nu)

MuJoCo 3.17.0, NumPy 2.5.4, Apple Silicon Mac (Apple M5, 4 performance + 6 efficiency cores). Other jobs were running, so all rates are approximate and change from run to run; sections A and B are deterministic, and the load average (how many processes compete for the cores) is printed. The script took 119 s; python experiment.py --quick is a smoke test whose rates do not reproduce the tables.

A. Contact basics. After 2 s of standing: base height 26.2 cm, exactly four floor contacts (FL, FR, RL, RR), and the sum of contact normal forces is 149.2 N against mg=15.21 kg×9.81=149.2m g = 15.21\,\text{kg} \times 9.81 = 149.2 N. The feet carry the whole weight; nothing else touches the floor.

B. Same physics. After 500 steps, rollout.rollout and a Python mj_step loop from the same state differ by exactly 0.0: batching changes the speed, not the result.

C. Throughput (Go2 standing from the settled stance, ~5·10⁴ steps per call, best of 5 round-robin rounds; load average 19 at the start, 24 at the end). Python loop over one MjData: 66,222 steps/s.

Wall-clock steps/s:

nbatch1 thread2 threads4 threads8 threads
153,68257,40854,81552,672
864,102151,081275,833488,908
6473,356154,889240,953490,195
51249,779100,212242,904428,027
409633,88185,478180,025272,937

Steps per CPU-second (CPU time summed over all threads):

nbatch1 thread2 threads4 threads8 threads
180,57881,86378,99078,729
882,57784,98878,20270,084
6488,72586,33573,87669,250
51272,78271,48069,63362,351
409661,34861,57558,36858,883

Two log-log panels against the number of robots per rollout call (1 to 4096), one line per thread count. Left, wall-clock steps per second: with one robot every thread count stays below the Python loop's 66,222; with 8 to 4096 robots, 8 threads reach up to 490,195. Right, steps per CPU second: all lines stay between 58,368 and 88,725.

Read the two tables together:

  • The cost of one Go2 step is roughly fixed: 58,368–88,725 steps per CPU-second in every cell, i.e. about 11–17 µs of CPU per step. Threads do not make a step cheaper (with 8 threads it gets slightly dearer: 62,351–78,729); they run more steps at once.
  • Threads need batch. With one robot, extra threads give nothing (53,682 → 52,672). With 8 to 4096 robots, 8 threads give 6.7–8.6× the single-thread wall rate. Ratios at or above 8× most likely mean the single-thread cell got a smaller share of a machine with ~20 other runnable processes; treat the wall-clock table as approximate.
  • C++ alone is not the win. One robot through rollout on one thread (53,682 steps/s) is no faster than the Python mj_step loop (66,222 steps/s); the best cell (7.4× the Python loop) comes from spreading a batch over threads.
  • Very large batches cost more per step (4096 robots: 58,368–61,575 steps per CPU-second); plausibly because rollout stores every state of every robot and the memory traffic grows with the batch; we did not profile it.
  • Arithmetic (section E): the Go2 training config needs 589,824,000 physics steps. At the best rate here that is 0.3 h of physics alone on this CPU, before any policy inference or learning, with our 2 ms step instead of the config’s 5 ms. And that best rate assumed 8 cores and robots that stand still.

The official Playground locomotion notebook trains a Go1 joystick policy:

  1. Open locomotion.ipynb in Colab.
  2. Runtime → Change runtime type → GPU. The notebook states that a GPU runtime is required, and a check cell raises an error if nvidia-smi fails.
  3. Run the install cells (pip install mujoco, mujoco_mjx, brax, playground, mediapy).
  4. Run the Quadrupedal → Joystick section (load, config and training cells): env_name = 'Go1JoystickFlatTerrain', ppo_params = locomotion_params.brax_ppo_config(env_name).
  5. Record, for your own report: GPU model (nvidia-smi), wall-clock time of the training cell, the final evaluation reward, and the rendered rollout. Then rerun with a different seed.

We report no numbers from this notebook: we did not run it, and the time depends on the GPU Colab assigns.

Section C started every robot from a settled stance. Start the same 512-robot batch from the Menagerie keyframe pose with joints perturbed by ±0.1 rad, so every robot first settles onto its feet:

start (512 robots × 25 steps, 8 threads)steps/ssolver iterations / stepcontacts
settled stance (+ small joint velocities)379,2091.74.0
released from keyframe pose, joints ±0.1 rad34,19410.34.4

The same model, batch and threads run about eleven times slower. The number of contacts barely changes (4.0 vs 4.4), but the constraint solver (the part of MuJoCo that computes contact forces) needs 10.3 iterations per step instead of 1.7 while feet land. Impacts, slipping and falling robots, most of early locomotion training, cost far more than standing: a benchmark on calm robots overstates training speed. Benchmark with the states your training will visit (Build it, Step 3, makes such falls on purpose).

  • The Go2 policies in unitree_rl_gym are trained in simulation with the gains above; Unitree’s low-level example for the real Go2 stands the robot with Kp 60, Kd 5 instead (go2_stand_example.py, see lesson 12). Training and deployment numbers are not interchangeable.
  • What transfers is shaped by what was randomized: friction range, pushes, observation noise (noise_scales in the config). Part 9 covers latency, system identification and domain randomization.
  • Nothing in this lesson was run on a real Go2. The standing test above is a simulation result; the real robot’s contact forces were not measured.

Quick check · 2 questions. Click an answer.

1. With one robot, 8 threads ran 52,672 steps/s; with 8 robots, 488,908. What explains the difference?

  • Threads make each physics step cheaper.
  • The cost of one step is roughly fixed; threads only help when there is a batch to split across them.
  • The C++ rollout module is much faster than a Python loop even for one robot.

Why: Section C: 58,368–88,725 steps per CPU-second in every cell. One robot through rollout on one thread (53,682 steps/s) is no faster than the Python loop (66,222 steps/s).

2. In Break it, robots released from the keyframe pose ran at 34,194 steps/s instead of 379,209. Why?

  • The robots have many more contacts (4.4 vs 4.0 is a large change).
  • The 8 threads stopped running in parallel.
  • While the feet land, the constraint solver needs 10.3 iterations per step instead of 1.7.

Why: Break it table. The number of contacts barely changes; the solver work does. Benchmark with the states training actually visits, impacts included.

  1. Easy. How many policy steps does one Go2 training iteration collect, and how many simulated seconds does each robot advance in it?
  2. Medium. Using your own row of the table in section C, estimate how long the physics of the Go2 training config would take on your CPU. Which two factors does this estimate ignore that make the real cost higher?
  3. Hard. Port the Go2: Menagerie ships go2_mjx.xml / scene_mjx.xml. Starting from Playground’s Go1JoystickFlatTerrain environment, replace the model and the default pose with the Go2’s (Unitree config: hip ±0.1, thigh 0.8 / 1.0, calf −1.5 rad), train on a Colab GPU, and report wall-clock and evaluation reward over 3 seeds.
  1. 4096×24=98,3044096 \times 24 = 98{,}304 policy steps; each robot advances 24×0.0224 \times 0.02 s =0.48= 0.48 s.
  2. On this Mac the best rate gives 0.3 h of physics alone (section E). Ignored: the policy forward pass and the PPO update, and the states visited: states with impacts, like the settling in section D, cost about eleven times more per step than standing. Also, 4096 robots per batch ran at a lower rate (272,937 steps/s) than the best cell.
  3. Open; we did not run it.
  • A legged robot is a floating base moved only by unilateral, friction-limited foot contacts; the policy commands joint targets and the gait emerges from contact physics and reward terms.
  • Locomotion PPO consumes ~10⁸ policy steps (Go2 config: 147,456,000); throughput, not the algorithm, decides whether training takes minutes or days.
  • On a CPU a Go2 step costs a roughly fixed 11–17 µs of CPU; threads help only when there is a batch to split, and impacts cost ~11× more per step. The next order of magnitude comes from thousands of environments on a GPU (MJX / Playground).

Next: lesson 37, the first sim-to-real gap: latency.

Sources