| cpu-laptop | expected wall-clock on cpu-laptop: ~0.33 min (measured) | measured |
| mps | wall-clock on mps: not yet measured | pending |
| t4 | wall-clock on t4: not yet measured | pending |
| 4090 | wall-clock on 4090: not yet measured | pending |
The one function you never questioned
For thirty chapters there was a line you never looked twice at: obs, reward, done, info = env.step(action). You hand MuJoCo an action, it hands you back the next state. mj_step is the oracle at the center of everything — every rollout, every training loop, every eval. You trusted it completely, because it is real physics and it is not yours to learn.
This chapter takes it away. We are going to learn that function from data — build a model that, given where things are and what you did, dreams where things go next. No mj_step inside. Just a neural network that watched the simulator run and learned to imitate its stepping. This is a world model, and it is the idea underneath Dreamer, underneath MuZero, underneath the whole "learn to imagine, then plan in imagination" line of work you'll build on in 3.2.
The payoff is not that the learned simulator is better than MuJoCo — it never will be. The payoff is that it is differentiable, fast to roll out, and yours — and in 3.2 you'll train a policy entirely inside it, never touching the real env during learning.
Reconstruction is not prediction
Here is the distinction the entire field turns on, and the one this chapter is built to make you feel.
Give a model the current observation and ask it to spit the same observation back out. That is reconstruction. It is an autoencoder. It is easy — squeeze the state through a bottleneck and expand it again, and with enough capacity the error goes to zero. Reconstruction tells you the model can represent a state. It tells you nothing about whether it understands dynamics.
Now hide the future from the model. Give it a few frames to orient, then feed it only the actions and ask it to roll forward — to say what the state will be in one step, two steps, ten steps, without ever seeing those frames. That is prediction, and it is a completely different question. To predict, the model has to have learned how a state flows into the next one under an action. It has to have learned the simulator.
The honest test of prediction is a baseline so dumb it is almost insulting: copy-last. Just assume nothing moves — predict that every future state equals the last one you saw. On slow dynamics this is shockingly strong. At 10 Hz control the pusher moves maybe a centimetre per step; for one step, "nothing changed" is nearly right. A world model that cannot beat copy-last has not learned anything a constant couldn't tell you.
So that is the bar. Not "does reconstruction work" (it will). The bar is: does k-step prediction beat copy-last?
The architecture: prior and posterior
# A deterministic-latent RSSM-lite (the skeleton of Dreamer's world model, minus# the stochastic-KL machinery — see "What we cut"). Four learned pieces:# encoder obs_t -> embed_t (features from an observation)# gru (z,a), h_{t-1} -> h_t (the DETERMINISTIC recurrent state)# posterior [h_t, embed_t] -> z_t (latent that HAS SEEN obs_t: for reconstruction)# prior [h_t] -> zhat_t (latent PREDICTED from h_t alone: for prediction)# decoder [h_t, z_t] -> obs_t (reconstruct the observation)# The prior/posterior split is the whole lesson: at train time the posterior teaches# the decoder to reconstruct; the prior learns to MATCH the posterior WITHOUT the# observation, so at test time we can roll the prior forward on actions alone. class WorldModel(nn.Module): def __init__(self, obs_dim, act_dim, latent_dim, hidden_dim, embed_dim): super().__init__() self.latent_dim, self.hidden_dim = latent_dim, hidden_dim self.encoder = nn.Sequential(nn.Linear(obs_dim, embed_dim), nn.ReLU(), nn.Linear(embed_dim, embed_dim)) self.gru = nn.GRUCell(latent_dim + act_dim, hidden_dim) self.posterior = nn.Sequential(nn.Linear(hidden_dim + embed_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, latent_dim)) self.prior = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, latent_dim)) self.decoder = nn.Sequential(nn.Linear(hidden_dim + latent_dim, embed_dim), nn.ReLU(), nn.Linear(embed_dim, obs_dim)) def _init(self, batch): h = torch.zeros(batch, self.hidden_dim, device=next(self.parameters()).device) z = torch.zeros(batch, self.latent_dim, device=h.device) return h, z def observe(self, obs_seq, act_seq): """Teacher-forced pass over a whole sequence. At each step advance the GRU, form the posterior latent from (state, observation), and decode. Returns the reconstruction loss and the DYNAMICS loss (prior matching posterior) — the two terms that, minimized together, make the prior a one-step simulator.""" batch, steps = obs_seq.shape[0], obs_seq.shape[1] embed = self.encoder(obs_seq) # (B, T, embed) — encode every frame at once h, z = self._init(batch) recon_loss, dyn_loss = 0.0, 0.0 for t in range(steps): if t > 0: # advance the deterministic state with the PREVIOUS latent + action h = self.gru(torch.cat([z, act_seq[:, t - 1]], dim=-1), h) zhat = self.prior(h) # predicted latent (no obs) z = self.posterior(torch.cat([h, embed[:, t]], dim=-1)) # corrected latent (sees obs) recon = self.decoder(torch.cat([h, z], dim=-1)) recon_loss = recon_loss + F.mse_loss(recon, obs_seq[:, t]) if t > 0: # train the prior to hit the posterior it can't see (detached target) dyn_loss = dyn_loss + F.mse_loss(zhat, z.detach()) return recon_loss / steps, dyn_loss / max(steps - 1, 1) @torch.no_grad() def warm_state(self, obs_seq, act_seq): """Filter through `context` observed steps to get (h, z) at the last one — the state the model predicts FROM. Uses the posterior: warmup gets to see.""" batch = obs_seq.shape[0] embed = self.encoder(obs_seq) h, z = self._init(batch) for t in range(obs_seq.shape[1]): if t > 0: h = self.gru(torch.cat([z, act_seq[:, t - 1]], dim=-1), h) z = self.posterior(torch.cat([h, embed[:, t]], dim=-1)) return h, z @torch.no_grad() def imagine(self, h, z, future_actions, true_future=None): """Open-loop rollout: step the GRU on actions alone and take the PRIOR latent each step (no observation) — dreaming forward. Decode to predicted states. `true_future` is only read under --break peek, where we illegally re-filter the posterior on the real frame: 'prediction' that secretly reconstructs.""" preds = [] for k in range(future_actions.shape[1]): h = self.gru(torch.cat([z, future_actions[:, k]], dim=-1), h) if true_future is not None: # BROKEN: peek at the frame we are meant to predict z = self.posterior(torch.cat([h, self.encoder(true_future[:, k])], dim=-1)) else: z = self.prior(h) preds.append(self.decoder(torch.cat([h, z], dim=-1))) return torch.stack(preds, dim=1) # (B, K, obs_dim)Four learned pieces, and one idea holding them together:
- an encoder turns an observation into a feature vector;
- a GRU carries a deterministic recurrent state
h— the running summary of everything that has happened; - a decoder turns
(h, latent)back into an observation; - and two heads that both produce the latent
z, differing only in what they are allowed to look at.
That last part is the whole design. The posterior forms the latent from (h, encoded observation) — it has seen the current frame, so it is the accurate, corrected belief. The prior forms a latent from h alone — it must predict what the posterior will be, before the observation arrives. This is the RSSM prior/posterior split, stripped to its skeleton (we keep the deterministic latent and drop the stochastic-KL machinery — see "What we cut").
Training ties them together. The decoder learns to reconstruct from the posterior latent; the prior learns to match the posterior it can't see. Minimize both and you get a model that reconstructs accurately when it has an observation and predicts accurately when it doesn't — because the prior has been trained to stand in for the posterior.
Data: sequences, not steps
# The world model trains on SEQUENCES, not single steps: it must learn how one# state flows into the next under an action, so the data is (state_0..T, action_0..T-1)# rollouts. We reuse the ch1.1 PushT env and the scripted expert (with exploration# noise for coverage) — the SAME simulator we are about to learn to imitate. def collect(num, seed_base): """Roll `num` scripted-expert episodes, each cropped to seq_len+1 states and seq_len actions. Deterministic: episode i uses env+expert seed (seed_base + i), so the whole dataset is reproducible from the CLI args alone (like gen_demos).""" states, actions = [], [] env = PushTEnv() need = args.seq_len + 1 i = 0 while len(states) < num: seed = seed_base + i i += 1 obs = env.reset(seed) expert = ScriptedExpert(noise=args.noise, seed=seed) obs_seq, act_seq, done = [obs], [], False while not done and len(obs_seq) < need: action = expert.action(env) obs, _, done, _ = env.step(action) obs_seq.append(obs) act_seq.append(action) if len(obs_seq) == need: # keep only full-length sequences (short successes dropped) states.append(np.stack(obs_seq)) actions.append(np.stack(act_seq)) return np.stack(states).astype(np.float32), np.stack(actions).astype(np.float32) train_states, train_actions = collect(args.episodes, seed_base=0)val_states, val_actions = collect(args.val_episodes, seed_base=EVAL_SEED_BASE + args.seed) # Standardize states to zero-mean/unit-std so every dim carries comparable weight in# the MSE (constant dims — the fixed target pose — get std 1, never divided to inf).# The sampler and losses live in this normalized space; we only de-normalize to report.obs_mean = train_states.reshape(-1, OBS_DIM).mean(0)obs_std = train_states.reshape(-1, OBS_DIM).std(0)obs_std = np.where(obs_std < 1e-4, np.float32(1.0), obs_std) def to_tensor(states, actions): norm = (states - obs_mean) / obs_std return torch.from_numpy(norm).to(device), torch.from_numpy(actions).to(device) train_obs, train_act = to_tensor(train_states, train_actions)val_obs, val_act = to_tensor(val_states, val_actions)print(f"dataset: {len(train_obs)} train / {len(val_obs)} val sequences of {args.seq_len} steps " f"(latent_dim={args.latent_dim}, hidden_dim={args.hidden_dim})")A world model cannot learn dynamics from a pile of independent (state, action) pairs — it has to see how a state becomes the next state, so the unit of data is a rollout. We reuse the PushT env and its scripted expert (the same simulator we are learning to imitate), add exploration noise for coverage, and crop each episode to a fixed-length sequence of states and the actions between them. States are standardized so every coordinate carries comparable weight in the loss; the fixed target-pose dimensions get unit scale rather than a divide-by-zero.
Training: one loss, two terms
# One loss, two terms: reconstruction (the decoder must recreate the state the# posterior saw) + dynamics (the prior must predict that posterior latent from the# recurrent state alone). The dynamics term is what turns an autoencoder into a# simulator — without it the prior is untrained and prediction is noise.torch.manual_seed(args.seed) # model init reproducible, independent of data collection abovemodel = WorldModel(OBS_DIM, ACT_DIM, args.latent_dim, args.hidden_dim, args.embed_dim).to(device)optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)recon_hist, global_step = float("nan"), 0for epoch in range(args.epochs): epoch_recon, epoch_dyn, num_batches = 0.0, 0.0, 0 for batch in torch.randperm(len(train_obs), generator=shuffle_gen).split(args.batch_size): recon_loss, dyn_loss = model.observe(train_obs[batch], train_act[batch]) optimizer.zero_grad() (recon_loss + dyn_loss).backward() optimizer.step() epoch_recon += recon_loss.item() epoch_dyn += dyn_loss.item() num_batches += 1 if args.rerun: rr.set_time("step", sequence=global_step) rr.log("train/recon", rr.Scalars([recon_loss.item()])) rr.log("train/dyn", rr.Scalars([dyn_loss.item()])) global_step += 1 recon_hist = epoch_recon / num_batches if epoch % 15 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:3d} recon {recon_hist:.5f} dyn {epoch_dyn / num_batches:.5f}")The objective is reconstruction loss plus a dynamics loss (the prior matching the detached posterior). The dynamics term is what turns an autoencoder into a simulator — delete it and the prior is never trained, so prediction is noise no matter how good reconstruction looks. On the default config this trains in about 20 seconds on a CPU laptop. It is deliberately tiny — the lesson is the mechanism, not a competitive Dreamer.
The headline, measured
# The headline. RECONSTRUCTION: decode the posterior on val states — the model# recreating what it sees (its floor). PREDICTION: warm the state on `context`# frames, then imagine `horizon` steps on actions alone and score each k against# COPY-LAST (repeat the last observed state). Copy-last is a strong, honest baseline# on slow dynamics: at k=1 almost nothing moved, so it wins. The world model wins# once the horizon is long enough that integrating the actions beats assuming# stasis — the CROSSOVER, and the widening gap after it, is "it learned to step."model.eval()with torch.no_grad(): val_recon, val_dyn = model.observe(val_obs, val_act) val_recon = float(val_recon) ctx_obs = val_obs[:, : args.context + 1] # frames 0..context (observed)ctx_act = val_act[:, : args.context] # actions between themfuture_act = val_act[:, args.context : args.context + args.horizon]future_obs = val_obs[:, args.context + 1 : args.context + 1 + args.horizon] # ground truth to predictlast_obs = val_obs[:, args.context] # copy-last baseline state h0, z0 = model.warm_state(ctx_obs, ctx_act)peek = future_obs if args.break_mode == "peek" else None # --break peek: illegal re-filteringpred = model.imagine(h0, z0, future_act, true_future=peek) # (B, K, obs_dim) wm_err = ((pred - future_obs) ** 2).mean(dim=(0, 2)).cpu().numpy() # per-k world-model MSEcopy_err = ((last_obs[:, None] - future_obs) ** 2).mean(dim=(0, 2)).cpu().numpy() # per-k copy-last MSEcrossover_k = next((k + 1 for k in range(args.horizon) if wm_err[k] < copy_err[k]), 0)wm_mean, copy_mean = float(wm_err.mean()), float(copy_err.mean()) # WHERE the win lives — the honest breakdown. Split obs into the fast PUSHER dims# (whose next state is a trivial first-order integral of the commanded velocity) and# the OBJECT dims (the T-block pose — the actual contact/pushing dynamics PushT is# ABOUT). The aggregate win is carried by the pusher; the object dynamics are the hard# half this tiny free-tier model does NOT yet beat copy-last on. The chapter must say so.PUSHER_DIMS, OBJECT_DIMS = [0, 1], [2, 3, 4, 5] # pusht obs layout: pusher xy | tee xy + yaw(sin,cos)def _mse(dims): # (world-model, copy-last) mean MSE over the given obs dims + all k w = ((pred[..., dims] - future_obs[..., dims]) ** 2).mean().item() c = ((last_obs[:, None][..., dims] - future_obs[..., dims]) ** 2).mean().item() return w, cwm_push, copy_push = _mse(PUSHER_DIMS)wm_obj, copy_obj = _mse(OBJECT_DIMS) print(f"reconstruction (posterior, sees the frame): val recon {val_recon:.5f}")print("prediction (prior rollout, actions only) vs copy-last, per horizon k:")for k in range(args.horizon): flag = " <- world model wins" if wm_err[k] < copy_err[k] else "" print(f" k={k + 1:2d} world_model {wm_err[k]:.5f} copy_last {copy_err[k]:.5f}{flag}") if args.rerun: rr.set_time("horizon", sequence=k + 1) rr.log("eval/pred/world_model", rr.Scalars([float(wm_err[k])])) rr.log("eval/pred/copy_last", rr.Scalars([float(copy_err[k])]))print(f"crossover at k={crossover_k or '>horizon'} | mean world_model {wm_mean:.5f} vs " f"copy_last {copy_mean:.5f} ({copy_mean / max(wm_mean, 1e-9):.2f}x lower error)")print(f" where the win lives: PUSHER dims wm {wm_push:.5f} vs copy {copy_push:.5f} " f"({copy_push / max(wm_push, 1e-9):.1f}x lower — learned to integrate the commanded velocity) | " f"OBJECT dims wm {wm_obj:.5f} vs copy {copy_obj:.5f} " + ("(wm wins)" if wm_obj < copy_obj else "(COPY-LAST WINS — the block/contact dynamics are the hard " "half this tiny model does not yet beat; that gap is the honest lesson, not a failure)")) if args.rerun: # a readable rollout: predicted vs true pusher_x on one val sequence (de-normalized) px_pred = pred[0, :, 0].cpu().numpy() * obs_std[0] + obs_mean[0] px_true = future_obs[0, :, 0].cpu().numpy() * obs_std[0] + obs_mean[0] for k in range(args.horizon): rr.set_time("horizon", sequence=k + 1) rr.log("eval/rollout/pusher_x_predicted", rr.Scalars([float(px_pred[k])])) rr.log("eval/rollout/pusher_x_true", rr.Scalars([float(px_true[k])]))Warm the state on a few observed frames, then roll the prior forward on actions alone and score each horizon k against copy-last. Measured across seeds 0-2 at the default config:
- Reconstruction (posterior, sees the frame): val MSE ~0.035 — the model recreates what it sees. Easy half, done.
- Prediction vs copy-last: at k=1 copy-last wins — almost nothing moved in a tenth of a second. The world model overtakes it by k=2-3 (the crossover), and the gap widens with the horizon. Averaged over the 12-step horizon, the world model's error is ~2.3x lower than copy-last (range 2.18-2.45x across seeds).
That crossover is the entire chapter. Copy-last decays because the world keeps moving and "nothing changed" gets more wrong every step. The world model holds low because it knows what the action does and integrates it forward. The moment its curve dips under copy-last's is the moment you can say, with a measurement behind it, that it learned to step the simulator.
But read the eval's per-dim line before you celebrate, because it keeps you honest. Split the observation into the two things that move — the pusher and the object — and the aggregate win comes apart. On the pusher dimensions the model beats copy-last by roughly 8x, which sounds like triumph until you notice why: the pusher's next position is a near-trivial integral of the velocity you just commanded. On the object dimensions — the T-block's own pushed-around pose, the contact dynamics PushT is actually about — copy-last still wins, on every seed. The block barely moves per step, so "nothing changed" is nearly right, and the model's prediction only adds noise.
So state the headline honestly: this tiny model learned the easy half of the simulator — the kinematics of the thing you drive directly — and not yet the hard half, the dynamics of the thing you drive it into. The aggregate crossover is real and worth measuring, but it is not "it learned the physics." That gap between the halves is not a failure to hide; it is the honest reason chapter 3.2 and the pixel Scale Lab exist.
A note on the --break peek mode: it lets the prediction rollout illegally re-filter the posterior on each true frame. "Prediction" error collapses toward the reconstruction floor and the model looks fantastic — because it is secretly reconstructing, not predicting. It is the single most common way to fool yourself with a world model, and it is worth running once to see how good cheating looks.
Why state, not pixels (the honest part)
The curriculum map says "PushT pixels" for this chapter, and we ran a feasibility spike to hold that promise honestly. The spike measured three things:
- Rendering pixels is cheap. PushT renders 32x32 grayscale frames offscreen at ~0.7 ms each on CPU — a pixel dataset is entirely free-tier.
- Reconstruction on pixels is beautiful. A conv encoder/deconv decoder reconstructs frames to ~5e-4 MSE. The frames look right.
- Prediction on pixels cannot beat copy-last at free-tier scale. And this is the finding. A T and a pusher are a handful of pixels on a static background, so copy-last's pixel-MSE is tiny — it is pixel-perfect on the ~95% of the frame that never changes. A small model carries a reconstruction-blur floor across the whole frame, and that floor never dips below copy-last's motion error. We measured no crossover through a 22-step horizon. Beating copy-last on pixels needs a sharper (larger) decoder — more compute — which is exactly the point.
That is the honest reason world models "eat compute," and it is the map's own framing for this chapter: genuinely instructive, deliberately small. So the free-tier chapter teaches the mechanism on the low-dim state, where the same recipe wins cleanly and runs in seconds; the pixel world model is the Scale Lab and the substrate that 3.2 grows into. The physics of the lesson — encoder, prior/posterior, decoder, reconstruction-vs-prediction — is identical either way. Only the observation changed, and with it the honest reach of a free-tier model.
Report and what carries forward
metrics = { "break_mode": args.break_mode or "none", "context": args.context, "copy_last_pred_mean": round(copy_mean, 6), "crossover_k": crossover_k, "epochs": args.epochs, "final_train_recon": round(recon_hist, 6), "hidden_dim": args.hidden_dim, "horizon": args.horizon, "latent_dim": args.latent_dim, "num_train_sequences": len(train_obs), "pred_ratio_copy_over_wm": round(copy_mean / max(wm_mean, 1e-9), 6), "seed": args.seed, "seq_len": args.seq_len, "smoke": bool(args.smoke), "val_recon": round(val_recon, 6), "world_model_pred_mean": round(wm_mean, 6),}(args.out / "metrics.json").write_text(json.dumps(metrics, indent=2, sort_keys=True) + "\n")print(f"metrics: {args.out / 'metrics.json'} (world model predicts {metrics['pred_ratio_copy_over_wm']}x " f"lower than copy-last; crossover k={crossover_k or '>horizon'})")if args.rerun: print(f"recording: {args.out / 'world_models.rrd'} — open it with: rerun {args.out / 'world_models.rrd'}")The metrics record the reconstruction floor, the prediction-vs-copy-last means, the ratio, and the crossover horizon — the numbers the exercises reproduce. In 3.2 you'll stop evaluating the prior rollout and start learning inside it: a policy trained on imagined trajectories, the real env untouched during learning. The imagine() rollout and the prior/posterior split you built here are exactly what that needs.
What we cut
- The stochastic latent. Real RSSMs make
za distribution (Gaussian or categorical) and train the prior/posterior gap as a KL divergence; we use a single deterministic vector and an MSE between prior and posterior — the same shape of idea, none of the sampling. This costs prediction sharpness and rules out proper uncertainty, and it is the first thing to add when you read the real thing. - Pixels. The observation is the 10-D state, not an image (see above). The encoder/decoder are MLPs, not conv stacks.
- KL balancing, free bits, reward/continue heads, symlog — all the DreamerV3 stabilizers. We predict only the observation, not reward or episode-continuation, because 3.1 is about learning the simulator; acting comes in 3.2.
Read the real thing
Everything we cut is sitting, fully assembled, in one file of Danijar Hafner's DreamerV3 — dreamerv3/rssm.py, pinned at commit e3f02248. It is the production version of the exact skeleton you just built. Reading it is not "here is the better code"; it is watching each simplification we made get its real machinery bolted back on.
Start with the prior/posterior split — our prior() and posterior() in wm.py#model. One honesty note first: this chapter's meta.yaml names them obs_step/img_step, the historical DreamerV2 names, but at this pinned commit those methods no longer exist under those names. The posterior path is now RSSM.observe (and its inner _observe); the prior path is RSSM.imagine, which calls _prior(feat); the GRU lives in _core. Same two questions — correct the latent from an observation, or predict it from the recurrent state alone — different names. Map our posterior() onto _observe, our prior() onto _prior, and our imagine() onto RSSM.imagine, and the architectures line up beat for beat.
Now the latent itself. Our z is a single deterministic vector — one nn.Linear output, tied to the posterior by MSE. The real RSSM keeps both halves of the state: a deterministic recurrent part (deter: int = 4096, the _core GRU state, our h) and a stochastic part that is categorical — stoch: int = 32, classes: int = 32, so 32 discrete variables of 32 classes each, sampled with a unimix: float = 0.01 uniform floor. That stochastic latent is precisely what lets the model represent a distribution over next states instead of one averaged guess — the same multi-modality problem that made BC blur in ch1.1. We kept only the deterministic half, which is exactly why our object dynamics collapsed into noise: a T-block's contact outcome is multi-modal, and a deterministic latent can only average the modes.
Finally the loss. Ours is dyn_loss = F.mse_loss(zhat, z.detach()) in WorldModel.observe — a point-wise MSE with a stop-gradient. In RSSM.loss it becomes a KL between distributions, split two ways: dyn = self._dist(sg(post)).kl(self._dist(prior)) and rep = self._dist(post).kl(self._dist(sg(prior))), each floored by free_nats: float = 1.0. That is KL balancing (the sg stop-gradient on opposite sides sets how fast prior and posterior move toward each other) and free bits (the floor stops the KL collapsing to zero). Our .detach() on the posterior is the stripped cousin of dyn; we simply dropped rep and the entire distributional layer.
Read these next, in order: (1) dreamerv3/rssm.py observe/_observe and imagine/_prior — trace them against our posterior/prior; (2) RSSM.loss — the dyn/rep/free_nats block, the KL our dyn_loss approximates; (3) _core, _prior, _dist — the GRU and the categorical latent we cut. Then, when you reach 3.2, dreamerv3/agent.py — the actor-critic that learns inside this model by rolling imagine forward.