World Models From Scratch, Part 2: Frozen Versus Joint

That is the end goal of this post, running live: the trained weights, executing in your browser. The white ball is steered by the policy this post builds, from those 32x32 pixels and nothing else. Put the goal wherever you like; it was trained on goals that cruise and bounce, so one that teleports under your cursor is a situation it never saw, and it copes.

The strip underneath is the strange part. Nothing renders those frames: from what the model currently believes, it imagines the next 16 steps with no further pixels arriving (training used 15). That imagination was the policy's only training environment: the policy was trained entirely inside this world model, all 5,000 actor updates, and never saw a real transition until it was evaluated.

Part 1 ended with a zero: velocity $R^2 = 0.00$ from the VAE's latent, and no training could have fixed it, because a single frame contains no velocity. This post gives a model 33 frames at once, then does it again a second way, because there are two ways to attach memory to an encoder and the difference between them is the whole point of the series.

1. The frozen recipe: a GRU on a frozen encoder

Result of this block: a GRU trained on the frozen part-1 latents reads velocity at 0.81, and every variant of it has lost the real ball by 20 open-loop steps. The setup is the first six sections; the numbers start at "Results: four GRUs".

The recipe, and the seam

Ha and Schmidhuber's World Models (2018) trains a VAE first, freezes it, and trains a recurrent model on its latents. The defining property of the recipe is not the GRU. It is the seam.

The paper's parts are V, a VAE that compresses frames to latents, M, a recurrent network that predicts the next latent, and C, a tiny controller evolved with CMA-ES. I build V and M and skip C; block 3 replaces it with something better anyway.

Whatever reconstruction happened to preserve is all the dynamics model will ever get, and however the encoder arranged it is the coordinate system the dynamics model has to work in. Part 1 measured both costs in advance: about 5.4 nats per frame, in a code curled enough that a linear probe reads position at 0.75 rather than the 0.9+ the design doc expected.

Why split at all? Because in 2018 it was the thing that worked: training a big conv encoder jointly with a recurrent model was unstable, and decoupling them turns one hard optimization into two easy ones. What it costs is what block 2 measures.

The model: one GRU cell and a Gaussian head

One GRU cell with 128 hidden units, a small MLP head that emits a Gaussian over the next latent, and a negative log-likelihood (NLL) loss.

The hidden state is 128 and stays 128 for the rest of the post, including the RSSM, so that probes on different models read states of the same dimension. The code and the smaller details are in the box.

class GRUDynamics(nn.Module):
    latent_dim: int = 16
    action_dim: int = 2
    hidden: int = 128
    residual: bool = False

    @nn.compact
    def __call__(self, h, z_prev, action):
        x = jnp.concatenate([z_prev, action], axis=-1)
        h, _ = nn.GRUCell(features=self.hidden)(h, x)
        y = nn.silu(nn.Dense(128)(h))
        mu = nn.Dense(self.latent_dim)(y)
        sigma = jax.nn.softplus(nn.Dense(self.latent_dim)(y)) + 1e-3
        if self.residual:
            mu = z_prev + mu
        return h, mu, sigma
Copy

Why a GRU can recover velocity when the encoder cannot

The GRU state is the first object in the pipeline that ever sees a sequence, and the NLL objective forces a velocity estimate into it.

Nothing in the architecture is smarter than the VAE; the GRU cannot see pixels, only the same frozen 16 numbers per frame that scored zero on velocity. But velocity is $v_t \approx (p_t - p_{t-1})/\Delta t$, a function of two frames, absent from any single one, and $h_t$ is a function of everything up to $t$, so it can hold one.

Whether it will is decided by the loss. The NLL is minimized by predicting the conditional expectation,

and for a ball with momentum that expectation depends on position and velocity. A state that dropped velocity would predict the same next frame for a ball moving left as for one moving right, and pay for it in likelihood at every step of every sequence.

Training setup, and the two switches

Subsequences of 33 frames, teacher forced, with the first four transitions masked out of the loss. Then a 2x2 grid over two design forks that papers usually pick without comment.

Subsequences are sampled uniformly from the training episodes; batch 64, Adam at $10^{-3}$, 20k steps, one lax.scan over time. 32 transitions contain a couple of bounces (arithmetic in the box), so linear extrapolation will not do.

• $h_0 = 0$, and the first four transitions are excluded from the loss. With no history the velocity is unknowable, so those steps would only teach the model to hedge. The same warm-up of 4 is used at evaluation time, for both models in this post.
• Training is teacher forced. At every step the GRU receives the true $z_{t-1}$, never its own prediction.

That last one is standard, and it sets up the most useful failure in this post:

Now the switches:

1. Direct or residual? Should the head predict $z_t$, or the change $z_t - z_{t-1}$? Residual is standard wisdom for smooth dynamics, since the change is small and easy to fit.
2. Means or samples? The frozen encoder emits $q(z \mid o) = \mathcal{N}(\mu, \sigma)$. Encode the dataset to posterior means, or to samples?

Instead of choosing, I ran the 2x2, four GRUs identical except for those switches, about 3 GPU-minutes each:

uv run train-gru --predict direct   --latent-source mean
uv run train-gru --predict residual --latent-source mean
uv run train-gru --predict direct   --latent-source sample
uv run train-gru --predict residual --latent-source sample
Copy

Why a small error becomes a large one

In a world with walls, open-loop error does not grow smoothly. A bounce turns a small disagreement about where into a total disagreement about when.

Give the two balls below a nudge apart and watch what the wall does to it:

Nothing in that demo is a neural network. It is the environment from part 1, run twice from almost the same state, and it puts a floor under every drift curve in this post: before the first mispredicted bounce, errors grow slowly, and after it the two timelines are unrelated.

How to read a pixel-error curve

The evaluation is an open-loop rollout: warm up on 4 real transitions, run the model forward on its own mean predictions for 30 steps, decode each predicted latent, and measure pixel MSE against the real continuation. That number is drift in every table below. rollout.py implements it once and both models in this post call it.

A pixel MSE of 0.018 means nothing until you know the scale. The reference levels that make every number in this post readable:

The third level is specific to each pipeline rather than to the data, because every GRU prediction is decoded by the frozen decoder from latents produced by the frozen encoder. The RSSM's floor will turn out to be completely different, and the comparison in block 2 is meaningless without both.

Results: four GRUs

Bold marks the best value in each column, and one number in that table is bad enough to spot without help.

The zero moves: velocity 0.74 to 0.81

Velocity is linearly readable from every variant, against 0.00 from the encoder alone: the recurrence integrates motion, as the objective argument said it had to. The one-step number is startling for a model that has never seen a pixel, 0.00041 per pixel for direct-mean against a pipeline floor of 0.00027, so one step into the future costs about 50% more than repainting the present. Position reads 0.979 on $h$ where the raw latent managed 0.75; the box has the two mechanisms and the confound.

The switch that detonates: residual-mean

Residual prediction, trained on means, posts the best one-step error in the grid, and then explodes: 0.139 by horizon 30, more than six times above the hallucination ceiling.

The mechanism is the oldest one in numerical integration. Open loop, the residual model computes $\hat{z}{t+1} = \hat{z}t + \mu\Delta$, a running sum; every step's small bias accumulates, nothing pulls the sum back toward latents the decoder has seen, and the state random-walks out of the data distribution. The direct model's update is $\hat{z}{t+1} = \mu(h_t)$, a fresh prediction of a plausible latent at every step, so errors saturate at "plausible but wrong" instead of compounding to "impossible".

Same objective, nearly identical teacher-forced scores (-20.7 against -21.6 NLL), completely different failure geometry. If you are going to roll a model out for 30 steps, measure it at 30 steps.

The switch that rescues it: train on samples

Training on samples costs you the short game and wins the long one.

The sample-trained runs are about 9 times worse at horizon 1, because the encoder's posterior noise $z = \mu + \sigma\epsilon$ is an error floor no model can beat. Both plateau below every mean-trained run by horizon 30, and the residual variant goes from 0.139 trained on means to 0.0177 trained on samples, best in the grid.

Figure 2 is the explanation. Teacher forcing trains the model on inputs from the real data distribution; open loop feeds it its own slightly-wrong outputs, which are inputs it has never seen, unless it trained on noisy latents, in which case "slightly wrong" has been in-distribution all along. Noise injection at training time buys robustness where open-loop rollouts spend it.

Dreamer (Hafner et al., 2019), the architecture in block 2, feeds its dynamics model samples rather than means. I had read that as an implementation detail for years; the grid says it is load-bearing.

Where they all end up: one bounce of validity

Three of the four curves bend up toward the 0.022 line by $k \approx 15$ to $20$ and flatten there (residual-mean keeps climbing). That plateau is the frozen recipe's ceiling, and it means about one bounce.

As the two-ball demo showed, a small error changes when and whether the next bounce happens, so the drift curve measures how long the model's timeline stays synchronized with ours. Block 2 asks how much of that ceiling is the frozenness itself.

2. Removing the seam: the RSSM

Result of this block: train the encoder and the dynamics against one loss and velocity goes to 0.96, reconstruction gets eight times better than the VAE's, and the model is still tracking a bounce at step 22. The architecture is the first six sections; the numbers start at "Results: frozen against joint".

Why delete the seam

Everything in block 1's results that could explain the ceiling is a property of the seam.

• The GRU spends capacity untangling and denoising a code that was shaped by reconstruction alone, which part 1 called the untangling tax.
• It cannot push information back into the encoder, so anything reconstruction discarded is gone for good.

Dreamer's answer is to delete the seam with a recurrent state-space model, the RSSM: one model, one loss, encoder trained by the prediction objective from the first gradient step.

Predict, then correct

Every filtering algorithm ever written has the same two beats: predict where the world should be, then correct that prediction with what you actually observe. A Kalman filter does it with matrices; an RSSM does it with a GRU and two small MLPs, and learns all of them.

The state comes in two pieces:

• $h_t$, deterministic, 128 dimensions, carried by a GRU. This is memory. Given the past, it is a known quantity.
• $z_t$, stochastic, 16 dimensions, resampled every step. This is everything the frame told you that memory could not have predicted.

Neither piece works alone. A purely deterministic state cannot represent uncertainty, and the random nudges make this world uncertain by construction; a purely stochastic state has to re-derive the whole past at every step through a sampling bottleneck. Splitting them routes predictable structure through $h$, where it is cheap and exact, and reserves the noisy channel for what needs it.

The other half of the design is that the model gets asked the same question twice, once with the answer hidden and once with it visible. The prior $p(z_t \mid h_t)$ must guess the stochastic state from memory alone. The posterior $q(z_t \mid h_t, o_t)$ gets to look at the frame. Forcing the first to agree with the second is how dynamics get learned, and that forcing is a KL term.

The cell

A GRU core carries $h_t$; a prior head guesses $z_t$ from $h_t$ alone, a posterior head guesses it with the frame in hand, and a KL term between the two is the only thing connecting prediction to perception.

One step, given frame $o_t$ and action $a_t$:

Block 1 paid for several of those choices:

• The decoder sees $[h_t, z_t]$, the whole state. Reconstruction can lean on the deterministic path, which frees $z$ to carry only what $h$ could not predict; this is most of the explanation for the reconstruction result further down.
• $z_t$ is always a sample. The 2x2 grid showed that noisy inputs are what make a dynamics model robust to its own errors in open loop; what was a training trick there is how the architecture works here.
• $\sigma$ has a floor of 0.1 on both heads, 100x the GRU's $10^{-3}$ and the Dreamer default; like the GRU's floor it stops the model from driving $\sigma \to 0$ and earning unbounded likelihood.

The encoder and decoder are the same conv stacks as part 1's VAE, $h$ is the same 128 as the GRU's, and the parameter counts are within 7% (box), so whatever is about to happen to the drift curve is not because I brought a bigger model.

The loss, and why this one did not collapse

This is part 1's ELBO with the sequence wrapped around it, except that the KL is against a learned prior $p(z_t \mid h_t)$ where part 1 had $\mathcal{N}(0, I)$, and that single difference is why the collapse that ate part 1's first run never happened here.

I trained with $\beta = 1$ from step zero, warm-up knob untouched, and the KL settled at about 1.4 nats and stayed alive for 20k steps. What changed is the cost of "ignore the latent". Against a fixed prior the model can drive the KL to zero by matching $\mathcal{N}(0,I)$ and pay nothing; against a prior that chases the posterior, the KL is only large when the frame carried something the dynamics failed to predict, so carrying information is cheap from the first gradient step and there is no collapsed configuration to fall into.

The reward term is switched off until block 3; the plain ball has no reward.

The one trick: KL balancing

$\mathrm{KL}(q \Vert p)$ has two arguments and one gradient. Dreamer splits the gradient with stop-gradients while leaving the value alone, and no stabilization trick in the build carries more weight.

Left alone, the KL pulls both ways at once: the prior moves toward the posterior and the posterior gets dragged toward whatever the prior currently believes, which early in training is nothing useful, so that second pull is a regularizer toward ignorance.

def kl_balanced(mu_q, sig_q, mu_p, sig_p, alpha: float = 0.8):
    sg = jax.lax.stop_gradient
    toward_prior = kl_gauss(sg(mu_q), sg(sig_q), mu_p, sig_p)
    toward_post = kl_gauss(mu_q, sig_q, sg(mu_p), sg(sig_p))
    return alpha * toward_prior + (1.0 - alpha) * toward_post
Copy

Four lines, easy to get backwards, so there is a test: at $\alpha = 1$ the posterior parameters receive zero KL gradient, checked with jax.grad. Write it, because a sign error here produces a model that trains, converges, and is wrong in a way no loss curve shows.

Training

One lax.scan with carry $(h, z)$, pixels in, encoder inside the graph. Batches of 32 sequences (the GRU used 64); everything else, data, splits, 33 frames, warm-up 4, horizon 30, ridge probe, is pinned to the frozen run. One protocol difference: no loss mask on the first four transitions, because the posterior sees $o_t$ at every step, so early states are corrected rather than blind.

No learning-rate drop, no $\beta$ warm-up, no instability. I had lined up both levers in the design doc and used neither. RSSMs have a reputation for being touchy, and I record the non-event so that the next person's expectations are calibrated by something other than my anxiety.

def step(carry, xs):
    h, z = carry
    e_t, o_t, a_t, r_t, n_t = xs[:5]
    h = model.apply(params, h, z, a_t, method=RSSM.core_step)
    mu_p, sig_p = model.apply(params, h, method=RSSM.prior_dist)
    mu_q, sig_q = model.apply(params, h, e_t, method=RSSM.post_dist)
    z = mu_q + sig_q * n_t
    o_hat = model.apply(params, h, z, method=RSSM.decode)
    recon = ((o_hat - o_t) ** 2).sum(axis=(1, 2, 3)).mean()
    klb = kl_balanced(mu_q, sig_q, mu_p, sig_p, config.alpha).mean()
    ...
Copy

Results: frozen against joint

The middle column is generous on purpose: it takes the best of the four GRU variants per row, so the frozen recipe is represented by residual-mean on the short horizons, residual-sample on the long ones and direct-mean on the probes. RSSM probes are on $h$ alone, 128 dimensions, capacity-matched against the GRU's $h$; adding $z$ to the probe input buys about 0.003 more.

Subtract the decoders before reading the short-horizon gap

The joint curve enters 11 times better at one step and 13 times at five, and most of that is the decoder.

Each pipeline's drift is bounded below by its own reconstruction error, and the two floors are nothing alike: 0.00027 per pixel for the frozen stack, 0.000032 for the joint one. The frozen GRU's best one-step prediction is about 30% above its floor; the RSSM's is within rounding of its own, 0.000031 against 0.000032 (the floor decodes posterior samples and the drift rolls prior means, so the two are not strictly ordered). The faithful one-step statement is that predicting one step ahead costs the RSSM nothing measurable beyond redrawing the frame it already has, the frozen model pays about 30% extra, and the "11 times better" is that fact multiplied by a better decoder.

The long horizon is where the dynamics show

By $k = 30$ the frozen model is 65 times above its floor and the joint one 400 times above its own, so the decoders have dropped out of the comparison.

There the RSSM reads 0.0129, just above the mean-image line at 0.011, and nowhere near the 0.022 ceiling where every frozen variant had parked. The joint model still retains information about the true trajectory at a horizon where the frozen recipe had decorrelated into "a perfect ball in an uncorrelated place".

Why it wins

The encoder stopped being an adversary.

In the frozen pipeline the GRU spent capacity untangling and denoising a code shaped by reconstruction alone. Here the encoder is shaped by the prediction loss from the start, so the latent arrives prediction-ready.

The cleanest evidence is the reconstruction row, the result I would have bet against: 0.033 summed error against the standalone VAE's 0.28, and the comparison is if anything unfair to the RSSM, whose figure decodes a posterior sample where the VAE's decodes the mean. I expected a reconstruction-versus-dynamics trade and got the opposite, because $h$ hands the decoder temporal context no single-frame encoder ever had: knowing where the ball was for the last thirty frames is extremely useful for drawing where it is now.

The KL is a task-complexity meter. Its equilibrium value is how many nats per frame the dynamics could not predict, and it costs nothing to log; the box has the numbers across the three tasks.

Block 3 asks whether the model is good enough to train an agent inside.

3. Dreaming: an agent trained in imagination

Result of this block: an actor that never saw a real transition scores 0.88 on the follow task in the real environment, against baselines near 0.05, and its imagination underestimated the reward rather than gaming it.

The setup

A new version of the environment with a red goal ball rendered into the pixels, and a policy trained on nothing but the world model's own imagined rollouts.

The goal has to be in the pixels because the reward depends on it. The observation goes RGB, the reward is a Gaussian kernel on the agent-to-goal distance, and one knob covers two tasks: the goal either sits still (hover) or cruises and bounces at half the agent's speed cap (follow).

The recipe is offline Dreamer. Freeze the world model, then train an actor and a critic only on imagined rollouts:

• Start from real posterior states filtered out of the replay data.
• Roll the prior forward 15 steps under the actor, $a_i = a_{\max} \tanh\big(\mu_\pi(s_i) + \sigma_\pi(s_i) \epsilon\big)$ with $s = [h, z]$.
• Read rewards from the reward head and score with TD($\lambda$): $R_i = r_{i+1} + \gamma\big[(1-\lambda) V(s_{i+1}) + \lambda R_{i+1}\big]$, with $\gamma = \lambda = 0.95$.
• The critic regresses onto $\mathrm{sg}(R)$; the actor maximizes $R$ directly.

Because every piece is reparameterized, Gaussian latents, tanh-Gaussian actions and a differentiable reward head, the value gradient flows through the dynamics themselves: no REINFORCE, no target network, no environment interaction. The policy improves by backpropagating through its own imagination.

Results

Evaluated in the real environment over 100 fresh episodes, the actor scores 0.65 on hover and 0.88 on follow, against baselines near 0.05.

What matters more than the scores

The imagination was pessimistic, not exploitative. I built the left panel of Figure 11 expecting to catch the actor conning its own reward head, the classic model-based failure. The opposite happened: imagined reward sits near 0.2 while real reward climbs past 0.8. The cause is horizon bias. Imagination only sees 15-step windows starting from random replay states, most of them far from the goal, so it mostly measures transit time; real 200-step episodes let the agent park on the goal and collect. Imagined return is a local-improvement signal, and episode return is a different quantity.

The harder task produced the better policy, because it produced the better world model. Follow beats hover 0.879 to 0.652, and the standard deviations are the tell: 0.025 against 0.228. The hover world model overfit its 400 static goal positions, the first train/validation gap of the project, so on some episodes it mislocates the target and the actor inherits the error. The follow model never had that option: 400 goal trajectories forced it to track, and it probes goal position at 0.99. Its actor locks on every single episode.

The policy is a mirror held up to the world model. Every flaw transfers.

Where this breaks

• One seed, one environment, 32x32 pixels. Every number is a single run. I would defend the 0.96-against-0.81 velocity gap and the factor-of-six divergence; the 0.998-against-0.979 position gap and the 0.811-to-0.766 spread across the GRUs, I would not.
• None of the four GRUs was tuned individually. A residual model with its own gradient clipping or a lower learning rate might not detonate; what I showed is that the default recipe does.
• The objective is reconstruction. On a ball in an empty box, what you need to draw is what you need to predict; add a distractor background, waving trees or TV static, and reconstruction spends its capacity on the irrelevant pixels because they occupy most of the image. That is the case for reconstruction-free objectives, and I plan a distractor-background version to make it visible.
• Continuous Gaussian latents. DreamerV2 (2021) and V3 (2023) use categorical latents with straight-through gradients, which are better at multimodal futures. For this world the Gaussian is fine, but "fine on a bouncing ball" is a weak claim about anything else.
• Capacity-matched rather than compute-matched, and the open-loop wall moved rather than vanishing. The RSSM does far more work per step (an encoder and decoder inside a 33-step scan, against a GRU over precomputed latents), and what that bought is the mean-image crossing moving from $k \approx 10$ to $k \approx 28$ plus an order of magnitude in error short of it: large practically, the same kind of thing. Both models eventually lose the timeline, which is what a stochastic environment does to open-loop prediction, and if wall-clock is your constraint the frozen recipe is cheaper and debugs in two independent phases.

uv run make-data
uv run train-gru  --predict direct --latent-source mean
uv run train-rssm --run-name base

uv run make-data --task follow
uv run train-rssm --data data/ball_follow.npz --run-name follow
uv run train-ac --wm-run runs/rssm/follow --data data/ball_follow.npz \
                --goal-speed 1.0 --run-name follow
Copy

What moved the zero

A representation is shaped by the loss that reaches it. Train an encoder on reconstruction and you get a code that knows what a frame looks like. Train it on prediction and you get a code that knows what the world does next, and, as it turns out, one that also knows better what the frame looks like.

Next stop is the environment this was all rehearsal for: the same recipe, pointed at DOOM's take_cover, where nobody hands you the true state and every measurement in this series has to be earned a different way.