Lesson 3.2.1 — Latent World Models: the RSSM and Dreamer

Week 20

Prerequisite lessons: 2.2.2 (model-based RL and MBPO), 2.2.5 (SAC), 1.2.2 (MDPs and value functions), 0.1.2 (autograd and the computational graph), 0.3.3 (filtering vs. smoothing), 3.1.4 (VLA fine-tuning).

Prior results used in this lesson. From 1.2.2: the MDP tuple and the value and action-value functions, with the Bellman expectation and optimality equations. From 2.2.2: the distinction between a model and a policy, and the MBPO result that short model rollouts interpolated with real data improve sample efficiency. From 0.3.3: a filter computes the posterior over a hidden state from past observations; the RSSM’s posterior is a learned filter. From 0.1.2: gradients flow through any differentiable composition; the RSSM’s training is one such composition. Each is restated at the point of use.

What this lesson hands forward. The latent-state formalism, the prior–posterior split, the KL-balancing objective, and the imagination loop. Lesson 3.2.2 (video prediction as world models) contrasts the RSSM’s compact latent with a pixel-space predictor. Lesson 3.2.3 (planning with world models) uses the imagination loop as the planner’s forward model. Module 4.2 uses the latent state as the policy layer’s working memory.

The running case in one sentence. The running arm, gripper, eye-in-hand camera, table, cube, bin, and divider are unchanged; this lesson learns a latent world model from the 100,000 transitions of the 2.2.4 dataset — collected on the reaching task — and asks what that model can and cannot predict when the arm is asked to grasp.

§L.0 — Lesson map, objectives, and misconceptions

Observable outcomes. By the end of this lesson you can, without notes: (1) state why a model in observation space is expensive and a model in latent space is not; (2) write the POMDP tuple and explain what the belief state adds to the MDP; (3) write the RSSM’s recurrence, prior, and posterior, and say what each is for; (4) state the ELBO for the RSSM and identify its four terms; (5) derive the KL-balancing objective and explain why naive KL collapses the latent; (6) describe the imagination loop and write the lambda-return used in it; (7) state why the actor gradient flows through the imagined dynamics and how the straight-through estimator makes it possible; (8) diagnose posterior collapse from a KL curve and name two fixes; (9) state the three failure modes of imagined rollouts (compounding error, reward extrapolation, terminal-state confusion); (10) evaluate the running case’s reaching-trained model on grasping and explain the extrapolation gap.

Misconception inventory.

Everything is self-contained. Where a prerequisite lesson is used, the result is restated at the point of use.

§L.0.5 — Retrieval warm-up

Q1. In 1.2.2, the MDP is $(S, A, P, R, \gamma)$ and the value function satisfies $V^\pi(s) = \mathbb{E}_{a \sim \pi}[R(s,a) + \gamma \mathbb{E}_{s' \sim P} V^\pi(s')]$. What does the model-based approach replace in this equation?

Answer. It replaces the expectation over the true transition $P$ with an expectation over a learned model $\hat{P}$, so the value can be estimated by rolling out the model rather than interacting with the environment. Used in §L.C.1 to define the imagination loop.

Q2. In 2.2.2, MBPO showed that short model rollouts branched from real states improve sample efficiency. Why short?

Answer. Because model error compounds with rollout length, so long rollouts drift off the data distribution and produce biased targets; short rollouts keep the model near its training distribution. Used in §L.C.3 to justify Dreamer’s finite horizon and bootstrapping.

Q3. In 0.3.3, a filter computes the posterior over a hidden state from observations up to the current time. What plays the role of the hidden state in a latent world model, and what plays the role of the observation?

Answer. The latent state $z_t$ is the hidden state; the image and proprioception are the observations. The RSSM’s posterior network is the learned filter. Used in §L.B.2 to define the posterior.

Q4. In 0.1.2, a computational graph records every operation so that gradients can be computed by the chain rule. What two properties must a world model’s dynamics have for the actor’s gradient to flow through it?

Answer. The dynamics must be differentiable with respect to the action, and the stochastic sampling must be reparameterizable or otherwise have a usable gradient estimator. Used in §L.C.2 to explain the straight-through estimator.

Part A — Why a latent model

§L.A.1 — The cost of modelling in observation space

The 2.2.2 lesson introduced model-based RL as learning $\hat{P}(s'|s,a)$ and rolling it out. The natural choice is to learn it in observation space: the image is the state, and the model predicts the next image.

That choice is expensive for three reasons, and each has a number.

Dimensionality. The eye-in-hand camera returns $640 \times 480 \times 3 = 921{,}600$ numbers per frame. A model that predicts this vector must have an output layer of that size, and the training signal is dominated by pixels that carry no task-relevant information — the table’s texture, the lighting gradient, the bin’s rim.

Irrelevance. The task depends on a handful of quantities: the cube’s pose, the gripper’s pose, the gripper’s opening, and the contact state. Everything else is nuisance. A pixel-space model spends its capacity on the nuisance because the loss weights every pixel equally.

Compounding error. Pixel prediction errors compound with rollout length in a way that is visually catastrophic: after ten steps the predicted image is a blur, after fifty it is unrecognizable. The 2.2.2 result that short rollouts are necessary is a statement about this compounding, and it is far more severe in pixel space than in a compact state.

A latent model addresses all three. It learns an encoder $e_\theta(o_t) \to z_t$ that maps the observation to a low-dimensional latent, and a dynamics model that predicts the next latent from the current one and the action. The latent is typically 30–200 dimensions. The model’s output layer is that size; the training signal is the reconstruction or the contrastive objective, not the raw pixel; and compounding error is bounded by the latent’s dimension.

Self-explanation 1. Why does a lower-dimensional latent reduce compounding error rather than merely moving it?

Answer. Because the latent is trained to discard the dimensions the task does not depend on, so the errors that compound are only in the task-relevant directions; a pixel model compounds errors in every pixel.

§L.A.2 — The POMDP and the belief state

The MDP of 1.2.2 assumed the state is fully observed. The running case violates that assumption: the cube’s pose is not in the image alone, because the cube can be occluded by the gripper, and the contact state is not visible at all. The correct formalism is the partially observed MDP (POMDP; Kaelbling, Littman, Cassandra 1998, tier A, S):

where $\mathcal{S}$ is the hidden state, $\mathcal{O}$ the observation space, $Z(o|s)$ the observation model, and the agent must act on a belief $b_t(s) = p(s_t | o_{1:t}, a_{1:t-1})$.

The belief is a distribution over hidden states, and its dimension grows with the state space unless it is summarized. The two classical summaries are the sufficient statistic — a finite vector that carries all the information the belief carries for the task — and the recurrent hidden state, which is what a learned filter produces. The RSSM uses the recurrent hidden state as its summary and adds a stochastic latent on top of it, and Part B explains why both are needed.

The key point for this lesson is that the latent $z_t$ is not the state; it is a sufficient statistic for control, learned jointly with the dynamics and the policy. This is why the latent can be 50-dimensional while the true state (cube pose, gripper pose, contact) is also small but unknown: the latent is trained to be whatever makes the prediction and control objectives small, not to recover the true state.

Concept check. A latent model is trained with a reconstruction loss and achieves near-zero reconstruction error. Does that guarantee the latent is a sufficient statistic for control?

Answer. No. Reconstruction only requires the latent to carry enough information to redraw the image, which includes the table texture; a latent that encodes texture and not contact state can reconstruct perfectly and control badly. The mechanism is that the reconstruction objective and the control objective are different, and the latent is only as good as the objective that shaped it. [displaces M1]

Worked-example ladder.

WE-L.A.a (fully worked). The eye-in-hand image is $640 \times 480 \times 3$ bytes. A latent is 64 floats at 4 bytes. Compute the compression ratio and the output-layer size ratio for a model that predicts the next observation versus one that predicts the next latent.

Solution. Image: 921,600 bytes. Latent: 256 bytes. Ratio: 3,600:1. Output layer for the image must produce 921,600 numbers (or a deconvolution producing them); for the latent, 64. The ratio of output sizes is the same 14,400:1 in numbers. The compression is not free — it is bought by the encoder and the training objective — but it is the reason the latent model is tractable.

WE-L.A.b (completion). The encoder is a CNN with 2 million parameters and the dynamics is an MLP with 200,000. Compute the parameter ratio of encoder to dynamics and state what that implies about where the training signal is concentrated.

Completion. 10:1 in favour of the encoder. The gradient signal is concentrated in the encoder during early training, because the dynamics cannot be learned until the latent is meaningful. This is why the RSSM trains the encoder and the dynamics jointly with a shared loss rather than in stages.

WE-L.A.c (faded). Derive the belief update for a POMDP with discrete states in two lines, and specialize it to a deterministic transition.

Faded step. $b_t(s') \propto Z(o_t|s') \sum_s P(s'|s,a_{t-1}) b_{t-1}(s)$. Completion. With a deterministic transition, the sum collapses to a single predecessor, so the belief update is a prediction followed by a correction — the same structure as the Kalman filter of 0.3.4, with the Gaussian replaced by a categorical or a learned distribution.

WE-L.A.d (contrastive near-miss). A learner trains a latent model with a reconstruction loss on the running case’s images and finds that the latent predicts the table’s texture pattern but not the cube’s position. Why it fails: the reconstruction loss is dominated by the large, high-contrast table region, and the cube occupies a small fraction of the pixels; the latent allocates capacity where the loss is largest. The fix: weight the loss by task-relevant regions, or use a contrastive or reward-predictive objective that does not require reconstruction, or add an auxiliary pose-prediction loss using the 2.1.2 estimator. The near-miss is the first sign that reconstruction is not the same as control.

Comparative table: model spaces.

Design principle P39.1. Model in the space the task depends on, not the space the sensor produces; the latent is a control statistic, not a compressed image.

Consolidation checkpoint 1

(a) Three capabilities. You can state the three costs of pixel-space modelling with numbers. You can write the POMDP tuple and the belief update. You can state why reconstruction does not guarantee a control-sufficient latent.

(b) Transfer probe. A depth camera returns $640 \times 480$ floats and the task depends on the cube’s 6D pose. Should the model predict depth or pose? Answer: pose, or a latent trained with a pose-auxiliary loss; depth prediction re-introduces the nuisance that the latent was meant to discard.

(c) Remediation rule. If you cannot state why a latent reduces compounding error, re-read §L.A.1 before Part B.

Part B — The RSSM

§L.B.1 — Two states, not one

The Recurrent State-Space Model (RSSM; Hafner et al. 2019, tier A, S) is the architecture behind Dreamer. Its defining choice is to carry two states per timestep: a deterministic recurrent state $h_t$ and a stochastic latent $z_t$.

The recurrence is

where $f_\theta$ is a GRU or an MLP with a recurrent connection. The stochastic latent has two distributions associated with it:

The prior predicts the latent from the recurrent state alone — that is, from the past. The posterior corrects the prior with the current observation. The observation model and the reward model are

and the continuation (episode-end) model is $p_\theta(c_t | h_t, z_t)$.

The two-state design is what distinguishes the RSSM from a plain recurrent model. The deterministic path $h_t$ carries the long-horizon information — the history of actions and the accumulated context — with no sampling noise. The stochastic path $z_t$ carries the uncertainty about the current state that the history cannot resolve. A deterministic-only model would have to represent a distribution over states as a single point, which either collapses the distribution or oscillates between modes; the stochastic latent lets the model represent several hypotheses simultaneously.

At training time, both the prior and the posterior are computed, because the posterior is what the ELBO compares against the prior. At inference and imagination time, only the prior is available, so the model rolls forward on its prior alone. The gap between the two is the KL term, and it is the central quantity of the training objective.

Self-explanation 2. Why is the posterior needed at training time if the model only uses the prior at imagination time?

Answer. Because the posterior is the target the prior is trained to match; without it, the prior has no supervision about what the latent should contain.

§L.B.2 — The ELBO

The RSSM is trained by maximizing the evidence lower bound (ELBO) on the log-likelihood of the observed sequence. For a sequence of length $T$,

The first three terms reward the model for explaining the observations, the reward, and the episode boundary. The fourth penalizes the posterior for departing from the prior, so that the prior remains a useful predictor at imagination time. The coefficient $\beta$ sets the trade-off.

The ELBO’s derivation is worth stating because it explains why the four terms appear together. The log-likelihood $\log p(o_{1:T}, r_{1:T})$ is intractable because it marginalizes over the latent sequence. Jensen’s inequality gives

and the four terms of Eq. (B.4) are the expansion of the right-hand side under the RSSM’s factorization. The bound is tight when the posterior equals the true posterior, which is why a well-trained posterior is a good filter.

§L.B.3 — KL balancing

The fourth term of Eq. (B.4) has a failure mode that is severe enough to have its own name. If $\beta$ is large, the cheapest way for the model to satisfy the KL penalty is to make the posterior equal to the prior — that is, to stop using the observation. The latent then carries only what the deterministic path carries, the stochastic path becomes decorative, and the model is effectively deterministic. This is posterior collapse, and its signature is a KL that falls to near zero while the reconstruction error plateaus.

The reverse failure is equally bad. If $\beta$ is small, the posterior is free to encode everything, but the prior is never trained to match it, so at imagination time the prior produces latents that the decoder cannot interpret. The model is accurate at training time and useless at rollout time.

KL balancing (Hafner et al. 2020, tier A, S) addresses both by splitting the KL into two terms with different gradient paths:

where $\mathrm{sg}(\cdot)$ is the stop-gradient operator. The first term trains the prior to match the posterior, with the posterior held fixed. The second term trains the posterior to match the prior, with the prior held fixed. With $\alpha = 0.8$, the balance favours training the prior, which is what the imagination loop depends on, while still allowing the posterior to use the observation.

This is the most important architectural detail in the lesson, and it is the one most often omitted in reimplementations.

Concept check. A reimplementation uses a single KL term with $\beta = 1$ and reports that the latent model’s rollouts are uninformative after three steps. Name the mechanism and the fix.

Answer. The posterior collapsed to the prior because the KL penalty had no balancing, so the latent carries no observation information and the prior’s predictions are uninformative. The fix is the split of Eq. (B.6) with $\alpha \approx 0.8$. [displaces M3]

Worked-example ladder.

WE-L.B.a (fully worked). A latent has 32 categorical variables with 32 classes each. Compute the number of possible latent configurations and the KL between the prior and a uniform posterior in the collapsed case.

Solution. $32^{32} \approx 10^{48}$ configurations. In the collapsed case the posterior equals the prior, so the KL is zero and the latent carries no information — the model is deterministic in disguise. The lesson: a large latent capacity does not help if the objective drives the posterior to the prior.

WE-L.B.b (completion). With KL balancing at $\alpha = 0.8$, the prior-matching term is weighted 0.8 and the posterior-matching term 0.2. If the observed KL is 2.0 nats, compute the contribution of each and state which gradient dominates.

Completion. Prior term $0.8 \times 2.0 = 1.6$, posterior term $0.2 \times 2.0 = 0.4$. The gradient through the prior dominates, so the prior is trained harder to match the posterior — which is the intent, because the prior is what imagination uses.

WE-L.B.c (faded). Derive the ELBO’s reconstruction term under a Gaussian observation model and show that it reduces to an MSE up to a constant.

Faded step. $\log p(o|z) = -\frac{1}{2}|o - \hat{o}|^2/\sigma^2 - \frac{d}{2}\log(2\pi\sigma^2)$, so maximizing it is minimizing the MSE when $\sigma$ is fixed. Completion. The constant is the log-normalizer; the gradient is the MSE. This is why a reconstruction loss is an ELBO term in disguise, and why a categorical observation model reduces to cross-entropy.

WE-L.B.d (contrastive near-miss). A learner trains the RSSM with a single KL term and $\beta = 0.1$ and finds that the model’s training-time reconstruction is excellent but its imagination rollouts diverge within five steps. Why it fails: $\beta = 0.1$ is small enough that the posterior is free to encode everything and the prior is never trained to match it; the KL between prior and posterior at test time is large, and the imagined latents are outside the decoder’s training distribution. The fix: the balanced KL of Eq. (B.6), which explicitly trains the prior.

Comparative table: latent model variants.

Design principle P39.2. A deterministic path carries memory, a stochastic path carries uncertainty, and the KL balance decides which one the model actually uses.

Consolidation checkpoint 2

(a) Three capabilities. You can write the RSSM recurrence, prior, posterior, and observation/reward models. You can write the ELBO and identify its four terms. You can derive the KL-balancing objective and explain both its terms.

(b) Transfer probe. The prior–posterior KL is 0.01 nats and reconstruction error is high. What has happened, and what is the fix? Answer: posterior collapse; the latent has stopped using the observation. The fix is to lower $\beta$ or rebalance with a lower $\alpha$ so the posterior has more freedom.

(c) Remediation rule. If you cannot state why a single KL term causes collapse, re-derive Eq. (B.6) before Part C.

Part C — Imagination

§L.C.1 — The imagination loop

Once the RSSM is trained, the policy can be trained inside it. The imagination loop starts from a state drawn from the training data, then rolls the model forward using the policy’s actions, producing a synthetic trajectory

of length $H$. The reward and continuation models produce $\hat{r}_t$ and $\hat{c}_t$ at each step, and the critic produces a value estimate at the horizon.

The starting states are sampled from the training distribution, not from a random initialization. This is the 2.2.2 MBPO insight restated: model rollouts must start where the model is accurate. Dreamer’s contribution is to make the starting states latents from the training sequences, which guarantees that the first step of the rollout is on the model’s manifold.

The horizon $H$ is finite, typically 15–50. Beyond that horizon, the model’s error dominates and the returns are meaningless. The 2.2.2 lesson’s short-rollout result is the same lesson.

§L.C.2 — Lambda-returns and the actor–critic

The actor’s objective is the expected return under the model. Directly maximizing the sum of imagined rewards has high variance because the rewards are themselves predictions; the standard solution is the lambda-return, a weighted average of $n$-step returns that trades bias against variance:

The critic $V_\psi$ is trained to regress $G_t^\lambda$ at every imagined step; the actor $\pi_\phi$ is trained to maximize $G_t^\lambda$ with respect to its parameters. The lambda parameter $\lambda \in [0,1]$ interpolates between one-step bootstrapping (low variance, high bias) and the Monte Carlo return (high variance, low bias). Dreamer uses $\lambda \approx 0.95$, which is close to Monte Carlo but still relies on the critic at the horizon.

Two implementation details make the loop work.

The actor gradient flows through the dynamics. The actor’s parameters enter through the action $a_t$ at each step, and the gradient of the return with respect to the actor must pass through the dynamics $f_\theta$, the prior $p_\theta(z|h)$, and the reward model. This is what makes the world model useful: it is differentiable, so the actor can be improved without interacting with the environment. The gradient path requires the dynamics to be differentiable in the action, which is a design requirement on $f_\theta$.

The stochastic latent’s gradient is estimated. The latent sample $z_t \sim p_\theta(z|h)$ is not differentiable under a categorical distribution. Dreamer uses a straight-through estimator: forward, $z_t$ is the sampled categorical or the sampled relaxed vector; backward, the gradient flows as if the sample were the mean of the distribution. For continuous latents, the reparameterization trick gives an exact gradient; for categorical latents, the straight-through is an approximation that works in practice because the categorical is one-hot and the downstream computation is smooth in the latent’s logits.

Self-explanation 3. Why does the actor’s gradient need to flow through the dynamics rather than being estimated by finite differences?

Answer. Because finite differences require one forward pass per action dimension and are noisy; the analytic gradient through a differentiable dynamics is exact and costs one backward pass.

§L.C.3 — Why imagination works and where it stops

Imagination works because the model’s error is bounded within its training distribution, and the rollouts start from that distribution and stay near it for a few steps. It stops working beyond the horizon where the error dominates, and it stops working outside the distribution where the model’s predictions are extrapolations.

Three failure modes bound the approach.

Compounding error. Model error grows with rollout length. The standard mitigation is the finite horizon and the learned value’s bootstrapping, which transfers the long-horizon estimate to the critic rather than to the model. This is why the horizon can be 15 while the task’s horizon is 200: the value supplies the tail.

Reward extrapolation. The reward model is trained on the rewards in the data, which may not cover the full reward range. Outside the training range, the reward model can predict values that are too high or too low, and the actor exploits the error. The mitigation is to keep the imagined states near the data, and to clip or regularize the reward model’s output.

Terminal-state confusion. The continuation model predicts whether the episode continues. If it predicts continuation where the real environment terminates, the imagined return overestimates; if it predicts termination early, the return underestimates. The mitigation is the same horizon discipline: terminations inside the horizon are used, and the horizon’s bootstrap value handles the tail.

Worked-example ladder.

WE-L.C.a (fully worked). An imagined rollout has rewards $r = [1, 1, 1]$ and continuation probabilities $\hat{c} = [1, 1, 0.5]$, with $\gamma = 0.99$, $\lambda = 0.95$, and a critic value $V(s_3) = 2$. Compute the lambda-return at $t=2$.

Solution. $G_3 = V(s_3) = 2$. $G_2 = \hat{r}_2 + \gamma \hat{c}_2 [(1-\lambda)V_3 + \lambda G_3] = 1 + 0.99 \times 0.5 \times [0.05 \times 2 + 0.95 \times 2] = 1 + 0.495 \times 2.0 = 1.99$. The return at the last imagined step is essentially the reward plus the discounted value, which is the intended behaviour.

WE-L.C.b (completion). With $H = 5$ and $\lambda = 0.95$, what fraction of the return’s weight comes from beyond the horizon when the critic is accurate? Completion. The weight on the critic’s horizon value is $\lambda^H \approx 0.77$, so about three quarters of the tail is the critic’s estimate and a quarter is the model’s; this is why an accurate critic matters as much as an accurate model.

WE-L.C.c (faded). Derive the straight-through gradient for a one-hot categorical sample and show that it equals the gradient through the logits as if the sample were the softmax probabilities.

Faded step. Forward: $z = \text{onehot}(\arg\max \ell)$. Backward: $\partial z / \partial \ell = \partial \text{softmax}(\ell)/\partial \ell$. Completion. The straight-through is biased but consistent in the sense that the expected gradient equals the score-function gradient plus a lower-variance term; in practice it trains stably because the actor’s loss is smooth in the latent’s expected value.

WE-L.C.d (contrastive near-miss). A learner rolls out the model for $H = 200$ steps and trains the actor on the Monte Carlo return. The actor learns to seek a predicted reward spike that does not exist in the environment. Why it fails: the model’s error at step 200 is large, and the reward model’s extrapolation at that horizon is unconstrained; the actor exploits the error. The fix: shorten the horizon to 15–50 and bootstrap with the critic, and regularize the reward model.

Comparative table: return estimators in imagination.

Design principle P39.3. Imagined returns are only as trustworthy as the horizon is short and the critic is accurate; bootstrap the tail, never imagine it.

Consolidation checkpoint 3

(a) Three capabilities. You can write the imagination loop and the lambda-return. You can state the two gradient-flow requirements. You can name the three failure modes and their mitigations.

(b) Transfer probe. The critic’s value at the horizon is systematically 10% too high. What happens to the actor’s policy? Answer: the actor is trained to reach states the critic overvalues, so it drifts toward the critic’s bias; the fix is critic recalibration or a longer horizon with more model steps, but the former is cheaper.

(c) Remediation rule. If you cannot write the lambda-return recursion, re-derive Eq. (C.2) before Part D.

Part D — Pitfalls

§L.D.1 — Posterior collapse and its diagnosis

The first pitfall is the one §L.B.3 introduced, and its diagnosis is a specific curve. Plot the KL between posterior and prior over training. A healthy model shows the KL rising and then plateauing at a small positive value — the latent is informative but not wildly divergent. A collapsing model shows the KL falling monotonically toward zero while reconstruction error stops improving. The signature is the joint pattern: a falling KL alone is fine if reconstruction is also falling; the collapse is the pair.

The second signature is the prior–posterior gap at rollout time. Measure the KL between the prior’s sample and the posterior’s mean on held-out sequences. A gap that grows with rollout length indicates that the prior is not tracking the posterior, which is the trained-up symptom of the same disease. The fix is KL balancing, and the operational test is to re-train with $\alpha = 0.8$ and compare the gap.

A third signature is the entropy of the latent. In a categorical latent, collapse shows as the posterior’s entropy falling to zero — the latent becomes a constant. Plot the entropy per latent variable; a collapsed variable has near-zero entropy across the dataset, which means it never distinguishes states.

§L.D.2 — Reward scaling and the robust objective

The reward scale interacts with everything. A reward of order 1 trains stably with a standard MSE; a reward of order 1000 dominates the gradient and makes the reconstruction term invisible. DreamerV3 (Hafner et al. 2023, tier A, S) addresses this with a symlog transform on the reward and the value targets: $\mathrm{symlog}(x) = \mathrm{sign}(x)\log(1+|x|)$, which compresses large magnitudes and is invertible. The transform makes a single set of hyperparameters work across reward scales, which is why DreamerV3 trains on tasks with rewards ranging over orders of magnitude without retuning.

The running case’s reward is the negative distance to the goal plus a sparse success term, of order 1 and 10 respectively. The scale is small enough that symlog is not required, but the sparse term’s magnitude relative to the dense term is a design choice that the world model inherits: if the success reward is 100 times the distance penalty, the reward model will fit the spike and miss the gradient.

§L.D.3 — Data coverage and the running case’s gap

The third pitfall is the one the running case will hit. The 2.2.4 dataset is 100,000 transitions from a mixed SAC + PID behaviour policy on the reaching task. The arm moves to positions, but it does not grasp. The world model trained on this data will be accurate for reaching and will extrapolate poorly for grasping, because the contact dynamics — the forces at the gripper’s fingers, the cube’s response — are absent from the data.

The signature is a latent prediction error that is low on the training distribution and high on the new task’s states. The diagnostic is to compute the model’s prediction error on a held-out set from the new task and compare it with the held-out set from the training task; a ratio above 2–3 is a coverage gap, not a model error.

The fix is not architectural. It is data: collect a small number of grasping episodes, add them to the training set, and retrain. The 3.1.3 §L.B.3 corrective-data argument applies here with a different target — the data is correcting the model’s distribution, not the policy’s — and the amount needed is set by the coverage gap, not by the task’s complexity.

Self-explanation 1. Why is a world model’s coverage gap more dangerous than a policy’s?

Answer. Because the actor is trained to exploit the model, so a model error in an uncovered region becomes a policy that seeks that region; a policy’s own error is at least confined to states it actually visits.

Worked-example ladder.

WE-L.D.a (fully worked). The reaching model’s latent prediction error is 0.02 on held-out reaching and 0.15 on held-out grasping. Compute the ratio and decide whether it is a coverage gap.

Solution. Ratio $0.15/0.02 = 7.5$, well above the 2–3 threshold, so it is a coverage gap. The model is not broken; it has not seen grasping.

WE-L.D.b (completion). Adding 2,000 grasping transitions drops the error to 0.06. Recompute the ratio and state whether more data is needed.

Completion. Ratio $3.0$, at the threshold. More data would help; a further 2,000 episodes is a reasonable next step, and the decision should be made against the imagination test of §L.E.2 rather than the raw error.

WE-L.D.c (faded). Derive the symlog transform and its inverse, and show that it is monotone and compresses large magnitudes.

Faded step. $\mathrm{symlog}(x) = \mathrm{sign}(x)\log(1+|x|)$; inverse is $\mathrm{sign}(y)(e^{|y|}-1)$. Monotone because the derivative is $1/(1+|x|) > 0$. Completion. The transform maps $[0, \infty)$ onto $[0, \infty)$ and compresses by a log factor, so a reward of 1000 becomes about 6.9 and a reward of 1 becomes 0.69 — a 10× compression ratio rather than 1000×.

WE-L.D.d (contrastive near-miss). A learner adds 10,000 grasping transitions to the reaching dataset and retrains the RSSM. The grasping prediction error improves but the reaching error rises. Why it fails: the added data now dominates the loss because grasping episodes are longer and more varied; the model has drifted toward the new distribution. The fix: balance the training batches by task, as in 3.1.3 §L.D.3, and monitor both errors.

Comparative table: pitfalls and signatures.

Design principle P39.4. Diagnose a latent model by comparing its error across distributions, not by its error on the training distribution.

Consolidation checkpoint 4

(a) Three capabilities. You can diagnose posterior collapse from two curves. You can apply symlog and state why it matters. You can compute and interpret a coverage-gap ratio.

(b) Transfer probe. The KL is healthy, reconstruction is good, but imagined rollouts diverge after eight steps. Which pitfall? Answer: prior–posterior gap or horizon overrun; check the gap on held-out sequences first, because a healthy KL on training does not imply a small gap at rollout.

(c) Remediation rule. If you cannot state why a coverage gap is a data problem rather than an architectural one, re-read §L.D.3 before Part E.

Part E — The running case

§L.E.1 — What the model sees and what it must predict

The running case’s dataset is the 2.2.4 100,000 transitions. Each transition is an observation (image and proprioception), an action (the mixed SAC and PID policy’s output, in the end-effector delta space), a reward (negative distance to the goal), and a continuation flag.

The observation encoder is a small CNN for the image plus an MLP for the proprioception, producing a 64-dimensional feature. The RSSM’s deterministic state $h_t$ is 200-dimensional; the stochastic latent $z_t$ is 32 categorical variables with 32 classes each. The reward model is an MLP from $(h_t, z_t)$ to a scalar; the continuation model is an MLP to a Bernoulli logit.

The training protocol follows §L.B.2 and §L.B.3: ELBO with reconstruction, reward, continuation, and balanced KL at $\alpha = 0.8$, $\beta = 1.0$. Batches are sequences of length 32, drawn from the training split at the episode level per 3.1.3 §L.D.2.

The model’s output at each step is a distribution over the next latent, a reward prediction, and a continuation probability. What it does not output is a 3D pose of the cube; that is the 2.1.2 estimator’s job, and the world model’s latent is a control statistic, not a pose.

§L.E.2 — The imagination test on grasping

The model is trained on reaching, so its imagination is meaningful for reaching and not for grasping. The running case’s test is to compute two quantities:

The reaching imagination error. Start from a held-out reaching state, roll the model forward for 15 steps under the recorded actions, and compare the predicted rewards and continuations with the recorded ones. A well-trained model matches the rewards within a few percent and the continuation within a few percentage points.

The grasping imagination error. Start from a held-out grasping state (which the model has not seen), roll forward, and compare. The error will be larger, and the ratio is the coverage gap of §L.D.3.

The decision rule is the ratio. If the ratio is below 2, the model is usable for grasping imagination with the caveat that its confidence is not calibrated; if it is above 3, the model requires more grasping data before it can be used for planning in 3.2.3.

§L.E.3 — What the model buys

Two capabilities, each with a cost.

Sample-efficient policy improvement. Once the model is accurate on a task, the actor can be trained on imagined rollouts without further environment interaction. For the running case, the reaching task’s policy can be improved from a smaller real dataset than a model-free method would need, and the 2.2.5 SAC comparison is the baseline.

A forward model for planning. The 3.2.3 lesson uses the model to evaluate candidate action sequences before executing them. This is the model-predictive control formulation, and it depends on the model’s accuracy and on its differentiable or sampleable form.

The cost is the training time and the coverage requirement. The model is an investment that pays off if the task is revisited many times; it does not pay off for a one-off task. This is the same marginal-cost argument as 3.1.1 §L.A.2, applied to the model rather than the policy.

Self-explanation 2. Why does the model’s value increase with the number of tasks that share its dynamics?

Answer. Because the dynamics of the arm and the table are shared across tasks, so a model trained on one task’s data partly transfers to another’s; the model is a shared asset, and the marginal cost of adding a task is the data to cover its new states.

Worked-example ladder.

WE-L.E.a (fully worked). The model’s reaching imagination error is 0.03 in normalized return units and the grasping error is 0.14. Compute the ratio and state the decision.

Solution. Ratio $0.14/0.03 = 4.7$, above the 3 threshold, so the model requires more grasping data before it can be used for grasping planning. It is usable for reaching planning immediately.

WE-L.E.b (completion). After collecting 3,000 grasping transitions, the grasping error falls to 0.06. Recompute and state whether the model is usable for planning with a 0.10 error budget. Completion. Ratio $2.0$, within the coverage threshold; the absolute error 0.06 is within the 0.10 budget, so the model is usable, with a caution that its confidence is not calibrated.

WE-L.E.c (faded). Design the retraining schedule: how many transitions per grasping task, and how the batches are balanced against the reaching data.

Faded step. Use the coverage ratio as the stopping criterion: collect until the ratio falls below 2, then balance batches by task so the reaching error does not rise.

Completion. The operational schedule is 2,000–3,000 transitions per task with a ratio check after each batch; balance by task, not by frame, for the reasons in 3.1.3 §L.D.3.

WE-L.E.d (contrastive near-miss). A learner trains the world model on reaching and then uses it to plan a grasping trajectory without collecting grasping data. The plan is to drive the gripper through the cube and the model assigns high reward because the reward model extrapolates. Why it fails: the reward model has never seen a collision, so its output there is an extrapolation, and the planner exploits it. The fix: never plan in an uncovered region, or add a safety layer that rejects plans outside the model’s training distribution.

Comparative table: the running case’s model uses.

Design principle P39.5. A world model is an investment in a shared dynamic; its value scales with the number of tasks whose states it covers, and its risk scales with the states it does not.

Consolidation checkpoint 5

(a) Three capabilities. You can state the running case’s model architecture and training protocol. You can compute and interpret the coverage-gap ratio. You can design the retraining schedule from the ratio.

(b) Transfer probe. The reaching error is 0.02 and the grasping error is 0.05. The grasping task’s error budget is 0.03. Is the model usable? Answer: no, because the absolute error exceeds the budget even though the ratio is within the coverage threshold; the ratio and the budget are independent criteria.

(c) Remediation rule. If you cannot state why a ratio and a budget are both needed, re-read §L.E.2 before the lab.

Failure taxonomy

The table covers the three lab corruptions plus three failures not exercised in the lab: critic bias, nuisance encoding, and planning-horizon error.

The lab

Task. Implement a minimal RSSM on a synthetic partially observed dataset and train it with the balanced KL of Eq. (B.6). Demonstrate the collapse of a single-term KL at high $\beta$, recover it with balancing, and roll the model forward in imagination to compare predicted and true returns.

Conditions. Python 3.10+, PyTorch 2.x, CPU sufficient. The synthetic task is a 1-D point mass observed through a noisy 4-D observation with a hidden velocity, so the observation is partially informative. Episode length 50, 500 episodes. Seed all RNGs.

Controls. (i) Single KL term at $\beta = 1$; (ii) single KL term at $\beta = 0.1$; (iii) balanced KL at $\alpha = 0.8$, $\beta = 1$; (iv) balanced KL with a horizon sweep at $H \in {5, 15, 50}$.

Metrics.

Pre-registered hypotheses.

H1. Control (i) collapses: its KL falls below 0.05 nats while its reconstruction error remains above control (iii)‘s, because the single penalty has no path to train the prior independently.

H2. Control (iii) maintains a KL between 0.1 and 2.0 nats with the lowest reconstruction error, because balancing trains both the prior and the posterior.

H3. Imagination error grows with $H$, and the growth is steeper for control (i) than for control (iii), because the collapsed prior cannot roll forward.

H4. Control (ii) has a high KL and good reconstruction but poor imagination, because the prior is untrained.

Break-it sequence.

1. Set $\beta$ very high silently. The KL collapses to zero and the model becomes deterministic. Silent: the loss curve is smooth.
2. Remove the continuation model silently. The imagined returns ignore termination and overestimate. Silent: reconstruction is unaffected.
3. Scale the reward by 100 silently. The reward term dominates the loss and reconstruction degrades. Silent: the total loss falls.

Acceptance criteria.

What good looks like. A good artefact has five properties. (1) The KL and reconstruction curves are plotted together, because the collapse signature is the pair. (2) The prior–posterior gap is measured on held-out sequences, not only on training. (3) Imagination error is reported against horizon, not as a single number. (4) The break-it results are a table. (5) The reward scale is logged and the symlog choice is justified.

A minimal RSSM core, seeded and annotated:

# lab_3_2_1.py — minimal RSSM with KL balancing
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
torch.manual_seed(0); np.random.seed(0)

$D_OBS$, $D_ACT$, $D_H$, $D_Z$, $N_CAT$, $N_CLS$ = 4, 1, 64, 8, 8, 8

class RSSM(nn.Module):
    def __init__(self):
        super().__init__()
        self.gru   = nn.$GRUCell(D_Z$ * $N_CLS$ + $D_ACT$, $D_H$)       # Eq. (B.1)
        self.prior = nn.$Linear(D_H$, $N_CAT$ * $N_CLS$)              # p(z|h)
        self.post  = nn.$Linear(D_H$ + $D_OBS$, $N_CAT$ * $N_CLS$)      # q(z|h,o)
        self.dec   = nn.$Linear(D_H$ + $D_Z$ * $N_CLS$, $D_OBS$)        # p(o|h,z)
        self.rew   = nn.$Linear(D_H$ + $D_Z$ * $N_CLS$, 1)            # p(r|h,z)
        self.con   = nn.$Linear(D_H$ + $D_Z$ * $N_CLS$, 1)            # p(c|h,z)

    def sample(self, logits):
        logits = logits.view(-1, $N_CAT$, $N_CLS$)
        probs  = F.softmax(logits, dim=-1)
        idx    = torch.distributions.Categorical(probs=probs).sample()
        onehot = F.one_hot(idx, $N_CLS$).float()
        # straight-through: forward uses onehot, backward uses probs
        z = (onehot - probs).detach() + probs
        return z.view(-1, $D_Z$ * $N_CLS$), probs

    def forward(self, obs, act, h):
        post_logits = self.post(torch.cat([h, obs], dim=1))
        z, $q_probs$   = self.sample(post_logits)                 # posterior sample
        prior_logits = self.prior(h)
        $p_probs$      = F.softmax(prior_logits.view(-1, $N_CAT$, $N_CLS$), dim=-1)
        $h_next$ = self.gru(torch.cat([z, act], dim=1), h)        # recurrence
        recon  = self.dec(torch.cat([h, z], dim=1))
        rew    = self.rew(torch.cat([h, z], dim=1)).squeeze(-1)
        con    = self.con(torch.cat([h, z], dim=1)).squeeze(-1)
        return $h_next$, z, $q_probs$, $p_probs$, recon, rew, con

def $kl_balanced(q_probs$, $p_probs$, alpha=0.8):
    q = $q_probs$.clamp_min(1e-6); p = $p_probs$.clamp_min(1e-6)
    kl_qp = (q * (q.log() - p.log())).sum(-1).mean()           # KL(q||p)
    # Eq. (B.6): train prior on sg(q) and posterior on sg(p)
    prior_loss = (q.detach() * (q.detach().log() - p.log())).sum(-1).mean()
    post_loss  = (q * (q.log() - p.detach().log())).sum(-1).mean()
    return alpha * prior_loss + (1 - alpha) * post_loss, kl_qp

def make_episode():
    v = np.random.randn() * 0.5                               # hidden velocity
    x = np.random.randn() * 0.5
    obs, act, rew, con = [], [], [], []
    for t in range(50):
        o = np.array([x, np.sin(x), v, np.random.randn() * 0.1])  # noisy obs
        a = np.array([np.random.randn() * 0.3])
        v = 0.9 * v + a[0] * 0.1                               # dynamics
        x = x + v * 0.1
        r = -abs(x)                                            # reward
        obs.append(o); act.append(a); rew.append(r)
        con.append(1.0 if t < 49 else 0.0)
    return (torch.tensor(obs, dtype=torch.float32),
            torch.tensor(act, dtype=torch.float32),
            torch.tensor(rew, dtype=torch.float32),
            torch.tensor(con, dtype=torch.float32))

model = RSSM()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
data = [make_episode() for _ in range(500)]

for step in range(3000):                                      # training loop
    ep = data[np.random.randint(len(data))]
    h = torch.zeros(1, $D_H$)
    total_kl = 0.0; total_rec = 0.0
    for t in range(50):
        o, a, r, c = ep[0][t:t+1], ep[1][t:t+1], ep[2][t:t+1], ep[3][t:t+1]
        h, z, q, p, recon, rew, con = model(o, a, h)
        rec_loss = F.mse_loss(recon, o)                        # reconstruction
        rew_loss = F.mse_loss(rew, r)                          # reward
        con_loss = F.binary_cross_entropy_with_logits(con, c)   # continuation
        kl, kl_raw = kl_balanced(q, p)                          # Eq. (B.6)
        loss = rec_loss + rew_loss + con_loss + kl
        total_kl += kl_raw.item(); total_rec += rec_loss.item()
    opt.zero_grad(); loss.backward(); opt.step()

print(f"KL={total_kl/50:.3f}  recon={total_rec/50:.4f}")

Line 9 defines the GRU recurrence of Eq. (B.1); lines 10–13 are the prior, posterior, decoder, reward, and continuation heads of Eqs. (B.2)–(B.3). Line 23 is the straight-through sample: the forward pass uses the one-hot, the backward pass uses the softmax. Lines 33–40 are the balanced KL of Eq. (B.6), with the stop-gradients in the right places. Lines 46–58 generate the partially observed episode. Lines 64–75 are the training loop; the printed KL and reconstruction are the pair to watch for collapse.

Cheat sheet

Exercises with solutions

E 3.2.1.1. Write the POMDP tuple and state what the belief adds to the MDP.

S 3.2.1.1. $(S, A, O, P, Z, R, \gamma)$; the belief $b_t(s)$ is a distribution over hidden states given the history, which the MDP does not need because the state is observed. The lesson: partial observability replaces the state with a belief, and the latent is the learned summary of the belief.

E 3.2.1.2. Write the RSSM’s recurrence, prior, and posterior, and state which is available at imagination time.

S 3.2.1.2. $h_t = f(h_{t-1}, z_{t-1}, a_{t-1})$; prior $p(z|h)$; posterior $q(z|h,o)$. At imagination time only the prior and recurrence are available, because the observation is not. The lesson: the prior is the part that must be trained for the model to be usable without observations.

E 3.2.1.3 (transfer). A latent has 16 categorical variables with 16 classes each. Compute the number of configurations and the KL in the collapsed case.

S 3.2.1.3. $16^{16} \approx 1.8 \times 10^{19}$ configurations; in the collapsed case the posterior equals the prior so the KL is zero and the latent carries no information. The lesson: capacity is irrelevant if the objective collapses the posterior.

E 3.2.1.4. State the ELBO’s four terms and explain the role of each.

S 3.2.1.4. Reconstruction, reward, continuation, and KL. The first three reward explaining the observations, rewards, and terminations; the KL keeps the posterior close to the prior so the prior is usable at imagination time. The lesson: the ELBO is a sum of prediction objectives and a consistency constraint, not a single loss.

E 3.2.1.5 (diagnostic). The posterior–prior KL falls from 1.2 to 0.02 nats while reconstruction error rises from 0.05 to 0.09. Diagnose and fix.

S 3.2.1.5. Posterior collapse: the KL penalty is dominating and the posterior is being driven to the prior. Fix with KL balancing at $\alpha \approx 0.8$ and a check of $\beta$. The lesson: the collapse signature is the pair of curves, not either alone.

E 3.2.1.6. Compute the lambda-return at $t=1$ for rewards $[1, 0.5]$, continuations $[1, 1]$, $\gamma = 0.99$, $\lambda = 0.95$, and $V(s_2) = 3$.

S 3.2.1.6. $G_2 = 3$. $G_1 = 0.5 + 0.99 \times 1 \times [0.05 \times 3 + 0.95 \times 3] = 0.5 + 0.99 \times 3.0 = 3.47$. The lesson: with an accurate critic, the return is dominated by the critic’s horizon value, which is why critic accuracy matters as much as model accuracy.

E 3.2.1.7 (transfer). A world model trained on reaching has held-out error 0.02 on reaching and 0.11 on grasping. The budget is 0.05. Diagnose and state the action.

S 3.2.1.7. The ratio $0.11/0.02 = 5.5$ indicates a coverage gap, and the absolute error exceeds the budget. Action: collect grasping data, balance batches by task, and re-evaluate after each batch against both the ratio and the budget. The lesson: coverage and budget are independent criteria, and both must hold.

E 3.2.1.8 (diagnostic). An actor trained in imagination seeks a reward spike that does not exist in the environment. Name two causes and the discriminating test.

S 3.2.1.8. (1) Reward extrapolation outside the data: test by mapping the imagined states against the training distribution. (2) Horizon overrun: test by shortening the horizon and observing whether the spike disappears. The lesson: exploitation of model error is diagnosed by comparing the imagined trajectory to the data distribution.

E 3.2.1.9. Explain why the straight-through estimator is needed for a categorical latent and what it approximates.

S 3.2.1.9. Categorical sampling is not differentiable, so the gradient cannot flow through the sample; the straight-through uses the one-hot in the forward pass and the softmax probabilities in the backward pass, approximating the true gradient with a biased but low-variance estimate. The lesson: the estimator is a gradient device, not a modelling choice, and its bias is acceptable because the downstream loss is smooth in the expected latent.

Exit test

You are ready for Lesson 3.2.2 when you can, without notes:

1. State the three costs of pixel-space modelling with numbers.
2. Write the POMDP tuple and the belief update.
3. Explain why reconstruction does not guarantee a control-sufficient latent.
4. Write the RSSM recurrence, prior, posterior, and observation and reward models.
5. State which states are available at training and which at imagination time.
6. Write the ELBO and identify its four terms.
7. Derive the KL-balancing objective and explain both of its terms.
8. Diagnose posterior collapse from the KL and reconstruction curves together.
9. Write the imagination loop and the lambda-return recursion.
10. Explain the two gradient-flow requirements: differentiable dynamics and straight-through sampling.
11. State the three imagination failure modes and their mitigations.
12. Apply symlog and state why DreamerV3 uses it.
13. Compute a coverage-gap ratio and decide from it plus the error budget.
14. Design a retraining schedule with task-balanced batches.
15. State the running case’s model architecture and its two usable regimes.

Cross-references. 1.2.2 supplies the MDP and value functions restated in §L.C.2. 2.2.2 supplies the model-versus-policy distinction and the short-rollout result restated in §L.C.3. 0.3.3 supplies the filtering formulation restated in §L.A.2. 0.1.2 supplies the computational graph that makes the actor gradient possible, restated in §L.C.2. 3.1.4 supplies the fine-tuning pipeline and the held-out evaluation discipline restated in §L.E.2. Lesson 3.2.2 contrasts the RSSM’s compact latent with pixel-space video prediction. Lesson 3.2.3 uses the imagination loop as the planner’s forward model.

Ready. The next lesson to write is 3.2.2 — Video Prediction Models as World Models. Say “next” to generate it.