You have been training one robot at a time
In chapter 2.1 you built PPO and it worked: the cartpole learned to balance by
acting and watching the consequences. Look at how it collected its experience,
though — eight envs in a plain Python for loop, stepped one after another. That
is how you learn the algorithm. It is not how the field runs it.
Real reinforcement learning scales the other way: not a bigger network, but more robots. Thousands of simulated robots stepping in parallel on one accelerator, so the policy sees millions of transitions per second and a run that would take hours finishes in minutes. The question this chapter answers is simply: how? And what does it cost?
Meet the other framework the field uses
Here is the honest reason this chapter looks different from the ~22 before it. The
tool that runs MuJoCo massively in parallel is MJX — MuJoCo-XLA — and MJX is
JAX, not PyTorch. There is no parallel MuJoCo in torch; teaching "4096 robots
at once" in torch would be teaching a fiction (see infra/decisions/015). So this
is the curriculum's one deliberate excursion into JAX — the same way chapter 1.9
graduated you from hand-rolled training loops to LeRobot. The field genuinely runs
both: torch for policies (ACT, diffusion, VLAs), jax for parallel sim and RL
(Brax, MJX). You should meet the second one once, on purpose.
You are not leaving the algorithm behind. ppo_mjx.py is chapter 2.1's PPO — the
same Gaussian policy, the same GAE with the same termination-vs-truncation
bootstrap, the same clipped surrogate and minibatch epochs. What changes is the
idiom: from imperative torch to functional JAX.
The mental-model shift (read this before the code)
If you are torch-native, four things will feel wrong at first, and each is the point:
- Functions are pure. Params are data. A torch module owns its weights and
you mutate them in place (
opt.step()). In JAX the network is a pure function and its parameters are a plain pytree you pass in and get out:model.apply(params, obs). Every update returns a new params pytree; nothing is mutated. - Randomness is explicit. No hidden global seed. You carry a
PRNGKeyandsplitit every time you sample. That is why CPU JAX replays bit-for-bit: the randomness is threaded, not ambient. vmapis the parallelism. You write the env and the rollout for one world.jax.vmapturns it into N worlds by adding a leading batch axis — no env loop. This one transform is how a single device runs 4096 cartpoles.jitcompiles the whole step.jax.jittraces the entire rollout + GAE + update into one XLA program and compiles it. The first call is slow (compilation); every call after is the payoff. That compile-once / run-fast shape is exactly why throughput cliffs, as you are about to measure.
The file keeps this primer inline so you can refer to it beside the code:
# JAX vs TORCH — the mental-model shift (you spent ~22 chapters in torch):# torch: a model OWNS its weights (self.linear.weight), you mutate tensors in# place (opt.step()), randomness is a hidden global (manual_seed).# jax: FUNCTIONS ARE PURE. Params are plain data (a pytree of arrays) passed# IN and returned OUT — model.apply(params, obs). Nothing mutates: every# "update" returns a NEW params pytree. Randomness is EXPLICIT — thread a# PRNGKey and split it (key, sub = jax.random.split(key)) at every sample,# so the same seed replays bit-for-bit.# vmap: write the env/rollout for ONE world; jax.vmap makes it N worlds for free# (a leading batch axis). That is how one device runs 4096 cartpoles.# jit: jax.jit traces the whole rollout+GAE+update into one XLA program and# compiles it. First call is slow (compile); every call after is the# payoff — that compile-once/run-fast shape is why throughput cliffs.# pytree: params, optimizer state, and a batch of MjData are all pytrees; tree_map# applies an op to every array leaf at once (we use it for autoreset).The env, re-expressed as pure functions
The cartpole is the same MJCF you have used since chapter 2.1
(common/envs/cartpole/cartpole.xml), now loaded into MJX with mjx.put_model.
We do not reuse the CartpoleEnv class — that one is imperative C-MuJoCo plus
numpy. Instead reset, step, and obs become pure functions on an mjx.Data, so
vmap and jit can batch them over thousands of worlds:
# The env is curriculum/common/envs/cartpole/cartpole.xml loaded into MJX. We do# NOT reuse CartpoleEnv (that's C-MuJoCo + numpy, imperative); the reset/step/obs# logic is re-expressed as PURE functions on an mjx.Data so vmap+jit can batch# them. Constants mirror the env (cartpole_env.py): slider is qpos[0], hinge is# qpos[1]; reward is +1 per surviving step.OBS_DIM, ACT_DIM = 5, 1MAX_STEPS, FRAME_SKIP = 500, 2ANGLE_LIMIT, CART_LIMIT, RESET_BOUND = 0.2095, 2.4, 0.05 _MJ = mujoco.MjModel.from_xml_path( str(Path(__file__).resolve().parents[3] / "curriculum/common/envs/cartpole/cartpole.xml"))_MX = mjx.put_model(_MJ) # the XLA-side model; shared (closed over) by every env fn def env_reset(key): """One fresh episode: draw qpos/qvel uniform[-RESET_BOUND, RESET_BOUND] from `key`, then mjx.forward to fill derived quantities. Returns an mjx.Data.""" kq, kv = jax.random.split(key) data = mjx.make_data(_MJ).replace( qpos=jax.random.uniform(kq, (2,), minval=-RESET_BOUND, maxval=RESET_BOUND), qvel=jax.random.uniform(kv, (2,), minval=-RESET_BOUND, maxval=RESET_BOUND), ctrl=jnp.zeros(ACT_DIM)) return mjx.forward(_MX, data) def env_step(data, action): """Apply the clipped force and advance FRAME_SKIP physics steps (50 Hz control over 100 Hz physics). lax.scan keeps the inner loop inside XLA.""" data = data.replace(ctrl=jnp.clip(action, -1.0, 1.0)) data, _ = jax.lax.scan(lambda d, _: (mjx.step(_MX, d), None), data, None, length=FRAME_SKIP) return data def env_obs(data): theta = data.qpos[1] # hinge angle; 0 = upright. (cos, sin) is seam-free at the top. return jnp.array([data.qpos[0], data.qvel[0], jnp.cos(theta), jnp.sin(theta), data.qvel[1]]) def env_terminated(data): # a REAL failure: pole fell or cart ran off the rail return (jnp.abs(data.qpos[1]) > ANGLE_LIMIT) | (jnp.abs(data.qpos[0]) > CART_LIMIT) # vmap each single-world fn over the leading env axis: THIS is the parallelism.reset_batch, step_batch = jax.vmap(env_reset), jax.vmap(env_step)obs_batch, term_batch = jax.vmap(env_obs), jax.vmap(env_terminated) def tree_select(mask, a, b): """Per-env choose whole pytrees a or b (mask True -> a). Autoreset: where an episode ended, swap in its fresh-reset Data leaf-by-leaf.""" return jax.tree_util.tree_map( lambda x, y: jnp.where(mask.reshape((-1,) + (1,) * (x.ndim - 1)), x, y), a, b)The last two lines are the whole trick: jax.vmap(env_step) is a function that
steps all N envs at once. tree_select is the parallel autoreset — where an
episode ended, it swaps in that env's fresh reset, leaf by leaf across the whole
mjx.Data pytree.
The same PPO, now functional
The rollout is a lax.scan over timesteps instead of a Python loop, but every
line maps to something you wrote in 2.1: sample an action, step the envs, record
the transition, and store the bootstrap value of the real next state for the GAE.
The termination-vs-truncation distinction you fought for in 2.1 is here unchanged
— truncated (time ran out) still bootstraps; terminated (the pole fell) still
does not:
def make_rollout(num_steps, num_envs): """A jittable rollout: lax.scan over num_steps timesteps, all num_envs stepping together. At each step: sample actions from the policy, step every MJX env, compute reward/terminated/truncated, autoreset finished envs, record the transition. Returns the trajectory (T, N, ...) plus the final carry.""" def body(carry, _): ts, datas, obs, step_count, ep_return, key = carry key, ak, rk = jax.random.split(key, 3) mean, logstd, value = ts.apply_fn({"params": ts.params}, obs) action = mean + jnp.exp(logstd) * jax.random.normal(ak, mean.shape) # reparameterized sample logprob = gaussian_logprob(action, mean, logstd) ndatas = step_batch(datas, action) nobs = obs_batch(ndatas) terminated = term_batch(ndatas) step_count = step_count + 1 truncated = step_count >= MAX_STEPS # a time limit, NOT a failure -> bootstrap done = terminated | truncated _, _, next_value = ts.apply_fn({"params": ts.params}, nobs) bootstrap = jnp.where(done, next_value, 0.0) # V(real next state) for GAE on episode ends ep_return = ep_return + 1.0 # reward is +1 per surviving step finished = jnp.where(done, ep_return, jnp.nan) # this env's return, if it ended this step rdatas = reset_batch(jax.random.split(rk, num_envs)) # autoreset (Brax pattern) datas = tree_select(done, rdatas, ndatas) obs2 = jnp.where(done[:, None], obs_batch(rdatas), nobs) step_count = jnp.where(done, 0, step_count) ep_return = jnp.where(done, 0.0, ep_return) transition = (obs, action, logprob, value, terminated, done, bootstrap, finished) return (ts, datas, obs2, step_count, ep_return, key), transition return lambda carry: jax.lax.scan(body, carry, None, length=num_steps) def compute_gae(traj, last_value, gamma, gae_lambda): """GAE over the rollout, walking backward (lax.scan reverse=True). Mirrors ch2.1: `1 - terminated` masks the bootstrap so a FALLEN pole contributes no future value, while a TRUNCATED episode keeps its stored bootstrap value.""" _, _, _, values, terminated, done, bootstrap, _ = traj def step(carry, t): gae, next_value = carry next_v = jnp.where(done[t], bootstrap[t], next_value) # ended step bootstraps from real next state delta = 1.0 + gamma * next_v * (1.0 - terminated[t]) - values[t] # reward is a constant +1 gae = delta + gamma * gae_lambda * (1.0 - done[t]) * gae return (gae, values[t]), gae init = (jnp.zeros_like(values[0]), last_value) _, advantages = jax.lax.scan(step, init, jnp.arange(values.shape[0]), reverse=True) return advantages, advantages + values # (advantages, returns)The update is PPO's clipped surrogate, now written as a pure loss function that
jax.value_and_grad differentiates and optax applies — the same objective, the
same advantage normalization / value clipping / entropy terms:
def make_update(args, batch_size, minibatch_size): """The PPO update: flatten the rollout, then take update_epochs passes of minibatch SGD on the clipped surrogate — the SAME objective as ch2.1, now a pure loss function jax differentiates. Returns fn(train_state, flat, key).""" def ppo_loss(params, mb, apply_fn): obs, action, old_logprob, adv, ret, old_value = mb mean, logstd, value = apply_fn({"params": params}, obs) ratio = jnp.exp(gaussian_logprob(action, mean, logstd) - old_logprob) # pi_new/pi_old adv = (adv - adv.mean()) / (adv.std() + 1e-8) # normalize advantages per minibatch (trick 1) # clipped surrogate: pessimistic (max) of the two -> a too-helpful update clips like a harmful one pg = jnp.maximum(-adv * ratio, -adv * jnp.clip(ratio, 1 - args.clip_coef, 1 + args.clip_coef)).mean() v_clipped = old_value + jnp.clip(value - old_value, -args.clip_coef, args.clip_coef) # trick 2 v_loss = 0.5 * jnp.maximum((value - ret) ** 2, (v_clipped - ret) ** 2).mean() entropy = gaussian_entropy(logstd).mean() # trick 3: reward being uncertain, keep exploring return pg + args.vf_coef * v_loss - args.ent_coef * entropy, (pg, v_loss, entropy) grad_fn = jax.value_and_grad(ppo_loss, has_aux=True) def minibatch(ts, mb): (loss, aux), grads = grad_fn(ts.params, mb, ts.apply_fn) return ts.apply_gradients(grads=grads), (loss, *aux) def epoch(carry, _): ts, flat, key = carry key, pk = jax.random.split(key) perm = jax.random.permutation(pk, batch_size) shuffled = jax.tree_util.tree_map(lambda x: x[perm], flat) minibatches = jax.tree_util.tree_map( lambda x: x.reshape((-1, minibatch_size) + x.shape[1:]), shuffled) ts, stats = jax.lax.scan(minibatch, ts, minibatches) return (ts, flat, key), stats def update(ts, flat, key): (ts, _, _), stats = jax.lax.scan(epoch, (ts, flat, key), None, length=args.update_epochs) return ts, jax.tree_util.tree_map(lambda s: s.mean(), stats) # (loss, pg, v, entropy) return updateAnd the payoff of the JAX idiom: the entire rollout → GAE → K epochs of minibatch SGD is one jitted function. Call it once per iteration; XLA runs the whole thing as a single compiled program:
def build(args, key): """Assemble the train state + parallel env state + one jitted update_step(runner) doing rollout -> GAE -> PPO epochs. Rebuilt per num_envs so --sweep can time the cliff. The whole step is ONE XLA program.""" batch_size = args.num_envs * args.num_steps minibatch_size = batch_size // args.num_minibatches model = ActorCritic(args.hidden_dim) key, mk, rk = jax.random.split(key, 3) tx = optax.chain(optax.clip_by_global_norm(args.max_grad_norm), # trick 4: cap the step size optax.adam(args.lr, eps=1e-5)) train_state = TrainState.create( apply_fn=model.apply, params=model.init(mk, jnp.zeros((1, OBS_DIM)))["params"], tx=tx) datas = reset_batch(jax.random.split(rk, args.num_envs)) rollout = make_rollout(args.num_steps, args.num_envs) update = make_update(args, batch_size, minibatch_size) @jax.jit def update_step(runner): ts, datas, obs, step_count, ep_return, key = runner (ts, datas, obs, step_count, ep_return, key), traj = rollout(runner) _, _, last_value = ts.apply_fn({"params": ts.params}, obs) advantages, returns = compute_gae(traj, last_value, args.gamma, args.gae_lambda) t_obs, t_action, t_logprob, t_value, _, _, _, t_finished = traj flat = jax.tree_util.tree_map( lambda x: x.reshape((-1,) + x.shape[2:]), (t_obs, t_action, t_logprob, advantages, returns, t_value)) key, uk = jax.random.split(key) ts, stats = update(ts, flat, uk) return (ts, datas, obs, step_count, ep_return, key), (jnp.nanmean(t_finished), stats) runner = (train_state, datas, obs_batch(datas), jnp.zeros(args.num_envs, dtype=jnp.int32), jnp.zeros(args.num_envs), key) return runner, update_stepRun it
python curriculum/phase2_reinforcement/ch2.3_mjx/ppo_mjx.py --seed 0
The default is a small, CPU-jax, free-tier config — 64 parallel envs — and it learns cartpole in about 0.49 min (measured, cpu-laptop):
iter 10/36 return 54.4 80,445 env-steps/s
iter 20/36 return 216.6 81,794 env-steps/s
iter 35/36 return 312.5 81,078 env-steps/s
eval: mean return 407.2 over 64 envs (random ~10-30, cap 500) # seed 0
(Seed 1 solves outright at 499.7 — RL is graded on the seeded band, not one run; chapter 1.6's determinism tiering is why.) It is honestly small: 64 envs is not the 4096 in the title. Which is the whole next lesson.
The wall-clock cliff
MJX's headline is throughput. Measure it directly — --sweep times one update
step at each env count and reports env-steps per second:
python curriculum/phase2_reinforcement/ch2.3_mjx/ppo_mjx.py --sweep 16,64,256,1024
num_envs env-steps/sec sec/step
16 50,154 0.041
64 83,638 0.098
256 104,544 0.313 <- CPU peak
1024 76,387 1.716 <- past the plateau, throughput FALLS
There is the cliff. Throughput climbs steeply as you add parallel envs — the
parallelism win is real — but on a CPU laptop it plateaus around 256 envs and
then reverses: a handful of cores cannot actually run 1024 worlds at once, so
adding more only adds queueing. This is the honest free-tier result. On a GPU the
same curve keeps climbing to 4096 and beyond, which is exactly why the 4096-robot
regime is this chapter's Scale Lab (--platform gpu --num_envs 4096, needs a
jax[cuda] install), not its free-tier default.
Throughput is not learning
Here is the trap the throughput number sets, and the deeper lesson. More envs give you more env-steps per second — so more envs is strictly better, right?
Not for learning at a fixed data budget. --total_steps is a budget of
env-steps, and each PPO iteration spends num_envs * num_steps of it. So doubling
num_envs at a fixed budget halves the number of gradient updates. Measured, at
the default 300k-step budget:
| num_envs | env-steps/s | gradient updates | eval return |
|---|---|---|---|
| 64 | ~84k | 36 | 407 |
| 256 | ~104k+ | 9 | 90 |
The 256-env run is faster and learns worse — same data, a quarter of the
updates, nowhere near solved. That is the throughput-vs-gradient-quality tradeoff:
parallel envs buy you data per second, but each gradient is one big, correlated
batch, and past some point more parallelism stops buying you learning. On a GPU
you would raise total_steps to give the 4096-env run enough updates — the
tradeoff never vanishes, it just moves. (Investigate it yourself in exercise 2.)
Determinism, honestly (same rule as chapter 1.6)
With --platform cpu (the default) and a fixed --seed, this file is
bitwise-reproducible: two runs produce byte-identical metrics, because every
source of randomness flows from one PRNGKey. That is CI-enforced. On a GPU,
JAX/XLA kernels are only statistically reproducible — same seed, same
qualitative result within the seeded band, not the same bytes. This is the exact
tiering chapter 1.6 taught; the paradigm changed but the honesty rule did not.
Read the real thing
ppo_mjx.py is the from-scratch version; the field's parallel-PPO code is the
same shape, hardened. Three implementations are worth reading, all JAX:
- PureJaxRL (
luchris429/purejaxrl) — the closest match to this file: a singlejax.jit-compiled PPO where the entire rollout-plus-update is one XLA program, exactly the idiom you just built. - Brax PPO (
google/brax,brax/training/agents/ppo) — the canonical parallel-env PPO, with the vectorized envs and autoreset generalized well past one cartpole. - MJX (
google-deepmind/mujoco,mjx/) — the sim underneath everything here: readmjx.stepandmjx.make_datato see the physics you vmapped.
The guided-reading segment pins the exact upstream commit to study against.
Scale Lab
The 4096-robot headline is a GPU study, not a free-tier run. With a
jax[cuda] install on a 4090 or L40S (infra/decisions/014), re-run the sweep out
to 4096 envs — --platform gpu --num_envs 4096 — and chart where the throughput
curve keeps climbing past the CPU plateau you measured near 256. Then repeat the
throughput-vs-gradient-quality table at a GPU-scale --total_steps, large enough
to give the 4096-env run the gradient updates it needs: the tradeoff does not
vanish on a GPU, it moves. These numbers stay PENDING in wallclock.csv until
measured on a real GPU runner — this chapter does not estimate them.