X2.2.2

Lesson 2.2.2 — Model-Based RL: Learned Dynamics, MPC, and the Latent World Model

Week 11 · Prerequisites: 1.2.2, 1.2.4, 1.2.5, 2.2.1, 1.1.3, 0.3.1.

Prior results used. This lesson uses 1.2.2’s MDP and Bellman equations, 1.2.4’s PPO and trust-region motivation, 1.2.5’s SAC and replay buffer, 2.2.1’s TD3 and twin critics, 1.1.3’s computed torque (for the action-space interface), and 0.3.1’s Gaussian distributions (for the probabilistic model’s uncertainty). Each is restated at the point of use.

Hands forward. This lesson hands forward one capability: the learner can fit a learned dynamics model to the ArmEnv’s transitions, use it for model-predictive control over a horizon, augment a model-free algorithm with imagined rollouts (Dyna, MBPO), train a latent world model (Dreamer), and diagnose the four model-based bugs (compounding bias, model exploitation, uncertainty miscalibration, stale model) from their planning and training symptoms — the capability that the VLA’s action planning in 3.1.4, the sim-to-real’s model adaptation in 3.4.3, and the capstone’s model-based decision layer all assume.

Running case. The same 6-DOF MuJoCo arm from 1.1.1–2.2.1, in the ArmEnv, with the observation [q, q̇, p_ee, p_target] and torque action τ ∈ ℝ⁶. The model predicts the next state ŝ' = f_φ(s, a), the data is the replay buffer’s 10⁵ transitions, the planning is MPC over horizon H = 20 with the cross-entropy method, and the policy training is SAC (1.2.5) or TD3 (2.2.1) augmented with imagined rollouts. The central result: model-based sample efficiency is 10–100× the model-free’s, at the cost of model bias — and bias management (ensemble, uncertainty, short horizon, real-data mixture) is the field’s central challenge.

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

§L.0.1 — Lesson map and word budget

§L.0.2 — Objectives

By the end of this lesson you will be able to do eight things:

1. State the model-based vs model-free distinction and the sample-efficiency-versus-bias trade-off.
2. Implement a probabilistic dynamics model (Gaussian output), a deterministic MLP model, and an ensemble, and state each variant’s bias-variance trade-off.
3. Implement MPC with a receding horizon and the cross-entropy method, and state why a shorter horizon is safer under model bias.
4. Implement Dyna (imagined rollouts mixed with real data) and MBPO (short-horizon rollouts, real-data trust region).
5. State the RSSM/Dreamer’s latent world model, the KL loss, and why training the policy in latent space avoids pixel-level cost.
6. Diagnose four model-based bugs — compounding bias, model exploitation, uncertainty miscalibration, and stale-model mismatch — from their symptoms.
7. State the model-exploitation defenses (short horizon, ensemble disagreement penalty, trust region, real-data mixture) and why each works.
8. Produce a working MPC for the ArmEnv with the CEM and an ensemble, and compare its sample efficiency against SAC.

§L.0.3 — Misconception inventory

Each misconception is displaced in the section named, and each is tagged in the concept check that follows.

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

This is the twenty-seventh lesson. The warm-up draws from 2.1.4 (two steps back), 2.1.3 (three), 2.1.2 (four), and 2.1.1 (five).

Warm-up 1 (from 2.1.4, the ICP). State the ICP’s alternation and Kabsch alignment. Answer: ICP alternates correspondence (nearest neighbor) and Kabsch’s SVD alignment, iterating to a local minimum. How this is used: §L.2’s CEM iterates (sample actions, evaluate, refit to elites), converging locally like the ICP.

Warm-up 2 (from 2.1.3, IBVS). State the visual-servo gain. Answer: v_c = -λ L⁺ e, gain λ sets the convergence rate. How this is used: §L.2’s MPC horizon H plays an analogous role — a large H (long lookahead) is fragile because model bias compounds, a small H is stable.

Warm-up 3 (from 2.1.2, the PnP). State the PnP’s minimal (4) and overdetermined (n > 4) cases. Answer: P3P/P4P gives up to 4 solutions; n > 4 is the least-squares case. How this is used: §L.1’s dynamics model is a regression from (s, a) to s', overdetermined from many transitions.

Warm-up 4 (from 2.1.1, depth ambiguity). State the Z = 0 singularity and the depth loss. Answer: Z = 0 is undefined; depth is lost in the 2D projection. How this is used: §L.4’s latent world model compresses a high-dimensional observation (like losing depth) into a low-dimensional latent z_t, and the reconstruction loss mirrors the projection’s irreversibility.

Part A — Model-Based vs Model-Free and the Dynamics Model

§L.1 — Objects: the model, its parameterizations, and the bias’s source

Model-based RL learns a dynamics model f_φ(s, a) predicting the next state ŝ' = f_φ(s, a) (deterministic) or the distribution p_φ(s' | s, a) (probabilistic). The learned model is used for:

1. Planning — MPC optimizes actions by searching the model over a horizon.
2. Policy training — the policy is trained on imagined rollouts (Dyna, MBPO, Dreamer).
3. Value estimation — value functions are evaluated on model rollouts.

Parameterizations:

• Deterministic MLP — ŝ' = f_φ(s, a) with φ the network weights. Fast, no uncertainty. Errors are invisible (no variance signal).
• Probabilistic Gaussian — ŝ' ~ N(μ_φ(s, a), Σ_φ(s, a)). Captures aleatoric (environment) uncertainty. Loss is −log N(s' | μ, Σ).
• Ensemble of K = 5 models — {f_φ_i}, with disagreement Var_i[f_φ_i(s, a)] measuring epistemic uncertainty. The mean is the prediction; the disagreement is the uncertainty.
• Latent world model (RSSM) — the state is compressed to a latent z_t; the model predicts in the latent space (Part D).

Where bias comes from:

• Approximation error — the network is smaller than the environment.
• Distributional shift — the policy visits states not in the buffer; predictions there are wrong.
• Compounding — model error at step k feeds into step k+1, so error grows with H.
• Exploitation — the policy optimizes predicted reward and finds the model’s blind spots.

Error growth under planning:

• Random, independent per-step error σ — variance grows as H·σ², so standard deviation grows as √H·σ (the ICP lesson’s variance addition).
• Systematic (worst-case) error ε — error grows as H·ε linearly.
• Policy optimization against errors — worst case grows as O((1+ε)^H), exponentially; this is why closed-loop policy training on the model is more fragile than open-loop planning with a receding horizon.

Mitigations:

• Short horizon (H = 1–5) — per-step error instead of compounded.
• Ensemble (K = 5) — disagreement penalizes exploitation.
• Receding horizon (MPC) — execute only the first action; re-plan with the true state.
• Real-data mixture (MBPO) — training stays anchored on real transitions.

§L.1.1 — Result: sample efficiency is 10–100×, bias is linear in the horizon for planning and super-linear for policy optimization, ensemble disagreement is epistemic only

Claim. Model-based sample efficiency is 10–100× that of model-free, at the cost of model bias. Under open-loop planning, per-step error ε gives O(H ε) growth (linear) or O(√H σ) RMS under independent noise. Under closed-loop policy optimization against the model, the error can grow super-linearly because the policy finds the model’s adversarially favorable actions. Ensemble disagreement measures epistemic uncertainty; per-model variance measures aleatoric. The mitigations (short horizon, receding horizon, real-data mixture, ensemble penalty) manage exploitation.

Proof sketch. Per-step error adds linearly in the expected error; under independent zero-mean noise, variances add so standard deviation grows as √H. When the policy optimizes against the model, it chooses actions that maximize the predicted return; wherever the model overestimates, the policy chooses that action and the overestimate is exploited. The receding horizon bounds the planning error to the first-step model error; the trust region keeps the policy’s imagined distribution close to the real distribution.

Corollary. Model-based sample efficiency does not imply model-based final performance. Asymptotically, the model-free algorithms (SAC, TD3) can reach higher returns because they don’t suffer compounding bias.

Corollary. Disagreement = epistemic; per-model variance = aleatoric. Conflating them over-penalizes stochastic environments.

§L.1.2 — Worked-example ladder

WE-L.1.a — Fully worked: the deterministic model’s compounding error

Input. True dynamics s' = s + a (1D), model ŝ' = s + a + ε with ε ~ N(0, 0.01²), horizon H = 10, action a = 0.

Trace (compounding).

1. True: s_0 = 0, s_1 = 0, …, s_10 = 0.
2. Model: ŝ_1 = ε_1, ŝ_2 = ŝ_1 + ε_2 = ε_1 + ε_2, …, ŝ_H = Σ ε_k.
3. Variance: Var[ŝ_H] = H σ² = 10 × 0.0001 = 0.001.
4. Std: √0.001 ≈ 0.0316.

Output and interpretation. Per-step σ = 0.01 gives H-step σ_H = √H × 0.01 ≈ 0.0316, i.e., √H ≈ 3.16× the per-step standard deviation. Error grows as √H (independent noise). At H = 100, σ_H = 0.1, i.e., 10× the per-step σ.

Self-explanation 1. If the noise is systematic (same ε each step), what is the growth?

Answer. ŝ_H = H ε, linear in H. Systematic error is worst-case; random error is √H. This is why systematic bias is more dangerous than zero-mean noise.

Self-explanation 2. Which horizon gives the smallest compounding? And why do model-based methods still use H > 1?

Answer. H = 1 gives no compounding, but it gives no lookahead either — the planner becomes a one-step greedy policy. The horizon is chosen to trade planning depth (longer H) against bias (shorter H). Most MBPO uses H = 1–5 for this reason.

WE-L.1.b — Completion problem

Input. Same 1D, H = 5, σ = 0.02.

Trace (final 25% as prompts).

1. Variance after 5 steps: 5 × 0.0004 = 0.002. ⟨verify⟩
2. Std: √0.002 ≈ 0.0447. ⟨verify⟩
3. Ratio to per-step: 0.0447 / 0.02 ≈ 2.24. ⟨verify⟩
4. The ratio equals √5 ≈ 2.24. ⟨verify⟩

Answers. All correct.

The lesson. The √H growth is universal for independent zero-mean noise; only the constant σ changes.

WE-L.1.c — Faded: the ensemble’s disagreement

Setup. Three models predict ŝ_i' = 1.0, 1.1, 0.9 at (s, a).

Question. What are the ensemble mean and disagreement?

You must produce the intermediate steps. The setup and the answer are given.

Answer. Mean = 1.0. Disagreement (std across models) = √(((1.0−1.0)² + (1.1−1.0)² + (0.9−1.0)²)/3) = √(0.02/3) ≈ 0.0817. The disagreement measures epistemic uncertainty.

The lesson. Ensemble disagreement = standard deviation across ensemble members. High disagreement means the models disagree = epistemic uncertainty.

WE-L.1.d — Contrastive near-miss: conflating aleatoric and epistemic

Correct solution. Aleatoric = per-model variance Σ_φ. Epistemic = ensemble disagreement. Penalize epistemic.

Near-miss. Total variance = Var_i[ŝ'_i] + mean_i[Σ_φ_i], treated as a single “uncertainty” penalty.

Why the near-miss produces a plausible but incorrect number. The total variance conflates the two. In a stochastic environment (high aleatoric), the total penalty would discourage the policy from exploring those states even though they can be learned from. In a deterministic environment (low aleatoric), the total equals the epistemic, so the bug is silent. The symptom: the policy becomes overly conservative in stochastic environments.

The discrepancy. Aleatoric cannot be reduced by more data (it is the environment’s inherent noise); epistemic can (more data improves the models’ agreement). Only the epistemic should drive exploration or penalty. The fix: separate Σ_φ (aleatoric) from disagreement (epistemic).

Self-explanation 3. In a deterministic 1D toy environment, is aleatoric zero?

Answer. Yes. The per-model variance Σ_φ should be near zero, and the ensemble disagreement is the only uncertainty. In this case, the total and the epistemic are equal, so the bug does not manifest. It manifests only in stochastic environments.

Part B — MPC: Planning in the Model

§L.2 — Objects: MPC, CEM, horizon, receding

Model Predictive Control solves, at each step:

a*_{t:t+H-1} = argmax_{a_t, ..., a_{t+H-1}} Σ_{k=0}^{H-1} γ^k · r(ŝ_{t+k}, a_{t+k})
   subject to ŝ_{t+k+1} = f_φ(ŝ_{t+k}, a_{t+k})

Receding horizon: execute only a*_t, then re-plan at the next step with the new observation. The model’s error affects H steps of planning, but execution sees only the first action under the true dynamics.

Cross-Entropy Method (CEM):

1. Sample N = 500 action sequences from N(μ, Σ).
2. Evaluate each sequence’s return with the model.
3. Select the top M = 50 (10%) as elites.
4. Refit: μ ← mean(elite), Σ ← var(elite).
5. Repeat 3–5 iterations.
6. Return the mean’s first action.

Horizon trade-off:

• Short H (1–5) — small per-step error, safe, short-sighted.
• Long H (20+) — compounding error, fragile, long-sighted.

Receding-horizon stabilizes by executing only the first action.

Ensemble MPC: K = 5 models; each predicts a trajectory; the mean is the expected trajectory, the std is the uncertainty. Penalize λ × std to avoid exploitation.

PETS (Chua et al. 2018, tier A, S) is the standard: probabilistic ensemble + MPC + trajectory sampling + uncertainty propagation.

§L.2.1 — Result: MPC’s receding horizon bounds the planning error to the first-step model error; CEM’s 3–5 iterations converge to the elite’s mean; ensemble uncertainty stabilizes

Claim. MPC’s receding horizon executes only the first action; the model’s H-step bias affects planning but not execution. CEM converges in 3–5 iterations with M/N = 10%. The ensemble’s disagreement is the uncertainty signal for the exploitation penalty.

Proof sketch. The receding horizon cuts execution to a single step under the true dynamics, so compounding does not occur during execution; only the planning suffers from compounding, and the first action is the least affected by it (the error at step k is k·ε, so step k = 0 has zero model error). CEM’s shrinkage of Σ converges to a local optimum of the sampled return. The ensemble’s disagreement is O(1/√K) smaller as K grows.

Corollary. Longer H gives better long-term planning but worse compounding; MPC’s optimal H is task-dependent, typically H = 10–30 for smooth dynamics, H = 3–5 for contact-rich.

Corollary. PETS is the standard model-based controller; on the standard benchmarks it matches SAC and TD3 while using 10–100× fewer samples.

§L.2.2 — Worked-example ladder

WE-L.2.a — Fully worked: the CEM’s first iteration

Input. A 1D planning problem, H = 2, N = 100, M = 10, initial μ = (0, 0), Σ = (1, 1).

Trace.

1. Sample 100 action sequences (a_0, a_1) from N(μ, Σ).
2. Evaluate each with J = r(ŝ_0, a_0) + γ r(ŝ_1, a_1) where ŝ_1 = f_φ(s_0, a_0).
3. Select the top 10 by J as elites.
4. Refit: μ_new = mean(elite), Σ_new = var(elite).
5. Repeat 3–5 times.

Output and interpretation. After 5 iterations, Σ has shrunk to ~10% of its initial value, and the mean is a local optimum. The first action a_0 is executed; the sequence is discarded.

Self-explanation 4. Why 10% for the elite fraction?

Answer. 10% is the standard CEM fraction. Smaller (5%) converges faster but to worse local optima; larger (20%) explores more but converges slower. 10% is the balance.

Self-explanation 5. After CEM converges, is the returned sequence the global optimum?

Answer. No — CEM is a local optimizer. The Gaussian N(μ, Σ) starts broad and shrinks; if the initial μ is bad, the local optimum may be poor. This is analogous to the ICP’s local minimum.

WE-L.2.b — Completion problem

Input. Same planning, but H = 5.

Trace (final 25% as prompts).

1. N = 500 sequences of (a_0, ..., a_4). ⟨verify⟩
2. Evaluate each with J over 5 steps. ⟨verify⟩
3. Select the top M = 50 elites (10% of 500). ⟨verify⟩
4. Refit μ, Σ over 5 dimensions. ⟨verify⟩

Answers. All correct.

The lesson. The action’s dimension grows with H; N × H total actions sampled per iteration.

WE-L.2.c — Faded: the receding horizon

Setup. H = 20, T = 100 steps, 5 CEM iterations.

Question. How many plans are executed?

You must produce the intermediate steps. The setup and the answer are given.

Answer. T = 100 plans (one per step). Each plan uses N = 500 × 5 iterations × H = 20 = 50,000 model queries. Total: 5 million model queries over the episode.

The lesson. The receding horizon executes one plan per step; the compute scales with T × N × iterations × H.

WE-L.2.d — Contrastive near-miss: no receding horizon

Correct solution. Receding horizon — execute only the first action.

Near-miss. Execute the whole H-step plan, then re-plan.

Why the near-miss produces a plausible but incorrect number. The full H-step plan is executed, and model bias compounds over the H steps. The trajectory drifts because the model’s step-k error is k·ε. The symptom: drift toward the end of the horizon.

The discrepancy. The receding horizon executes only the first action under the true dynamics; the full plan executes all H actions under the true dynamics but the plan is based on the model’s H-step prediction. The receding horizon is stable; the open-loop is not.

Self-explanation 6. Is the H − 1 future actions wasted?

Answer. No — the future actions inform the first action’s optimal value. The plan is discarded after execution because the new state at t+1 differs from the model’s prediction; re-planning uses the true state.

Part C — Dyna and MBPO

§L.3 — Objects: Dyna, MBPO, short-horizon branching

Dyna (Sutton 1990, tier A, S): the agent stores real transitions in the buffer, fits a model, generates imagined rollouts (K = 5 per real state, H = 10 steps), and trains the policy on real + imagined transitions.

MBPO (Janner et al. 2019, tier A, S): short-horizon (H = 1–5) imagined rollouts, ensemble of models (K = 5), policy trained on real + imagined with a real-data trust region.

Key differences:

The real + imagined mixture:

D_train = D_real ∪ {(ŝ, a, r̂, ŝ') from H-step imagined rollouts}

Defenses against model exploitation:

• Short horizon (H = 1–5): per-step error, not compounding.
• Ensemble (K = 5): disagreement penalizes exploitation.
• Trust region (MBPO’s): real-data keeps the policy anchored.
• Real-data mixture: training on real transitions.

§L.3.1 — Result: Dyna mixes real and imagined; MBPO bounds compounding with short horizon and anchors the policy on real data

Claim. Dyna mixes real + imagined rollouts; MBPO’s short horizon (1–5) bounds compounding, and the real-data mixture anchors the policy. Sample efficiency is 10–100×, with exploitation mitigated by the short horizon + trust region.

Proof sketch. The Dyna buffer is the real distribution plus K imagined rollouts per real state. MBPO’s short horizon bounds per-step model error. The MBPO trust region keeps the policy’s distribution close to the real distribution — Janner et al. 2019’s theoretical result: with a model error ε_m and policy distribution shift bounded, the performance bound is O(ε_m) per step, not exponential.

Corollary. The MBPO H = 1–5 is the standard; H = 10+ (Dyna) often exploits the model without trust-region mitigation.

Corollary. The real + imagined mixture is central: real anchors, imagined provides efficiency; the ratio is the sample-budget decision.

§L.3.2 — Worked-example ladder

WE-L.3.a — Fully worked: the Dyna imagined rollouts

Input. Buffer 1,000 real transitions, H = 5, K = 5 rollouts per state, SAC policy.

Trace.

1. Real transitions: 1,000.
2. Imagined: 5 × 5 = 25 per state → 1000 × 25 = 25,000 imagined transitions.
3. Training buffer: 1,000 + 25,000 = 26,000 transitions.
4. Policy trains on 26,000, not on 1,000.

Output and interpretation. The 25,000 imagined are 25× the real — 25× sample efficiency. The H = 5 bounds compounding.

Self-explanation 7. The H = 5 horizon: if H = 20, the count is the same (25,000) but the error is 2× (√20/√5 ≈ 2). What is the trade-off?

Answer. Same count, larger error. The H = 5 is safer.

Self-explanation 8. The imagined are 25,000 vs. real 1,000. Are they equivalent?

Answer. No — the imagined carry the model’s bias. The real + imagined mixture anchors the policy.

WE-L.3.b — Completion problem

Input. MBPO with H = 1.

Trace (final 25% as prompts).

1. H = 1: one step of imagination per real state. ⟨verify⟩
2. Imagined next state: f_φ(s_t, a_t). ⟨verify⟩
3. Per-step error: no compounding. ⟨verify⟩
4. Efficiency: K rollouts, not H — reduced count, lower error. ⟨verify⟩

Answers. All correct.

The lesson. H = 1 gives no compounding, at the cost of reduced sample efficiency.

WE-L.3.c — Faded: model exploitation

Setup. Policy trained on H = 20 imagined rollouts only, no real rollouts during training, no trust region.

Question. What is the symptom?

You must produce the intermediate steps. The setup and the answer are given.

Answer. High imagined return (the policy exploits model errors), low real return (the errors are wrong). The imagined-real gap is the symptom. The fix is short horizon, trust region, ensemble penalty, real-data mixture.

The lesson. Model exploitation appears as an imagined-real gap; the policy finds the model’s errors.

WE-L.3.d — Contrastive near-miss: imagined-only training

Correct solution. Real + imagined mixture.

Near-miss. Imagined-only (100% model-generated).

Why the near-miss produces a plausible but incorrect number. The policy is trained only on the model; it exploits unchecked errors. Real performance is poor despite high imagined return.

The discrepancy. Real + imagined is anchored; imagined-only is exploitative.

Self-explanation 9. If 1% is real, is the anchor sufficient?

Answer. Depends on the trust region. MBPO’s real-data policy constraint is the anchor, not the ratio. If the trust region is enforced, 1% can suffice; without trust region, even 50% may fail.

Part D — Latent World Models: Dreamer and the RSSM

§L.4 — Objects: the RSSM, the latent, the imagined training

Dreamer (Hafner et al. 2019, 2020, 2023; tier A, S) trains the policy in the latent space of a world model, not directly on pixels.

RSSM (Recurrent State Space Model):

• Deterministic path h_t = f_φ(h_{t-1}, z_{t-1}, a_{t-1}) (a GRU).
• Stochastic path z_t ~ N(μ_φ(h_t), Σ_φ(h_t)) (a 32-dim latent).
• Observation decoder ô_t = D_φ(h_t, z_t) (reconstruction).
• Reward decoder r̂_t = R_φ(h_t, z_t).

Loss:

L_world = −log p(ô_t | h_t, z_t) + β·KL(z_t || prior_t) + log p(r_t | h_t, z_t)

The reconstruction term teaches the latent to preserve the observation; the KL term regularizes the latent toward the prior; the reward term learns the reward predictor.

Policy training: the policy and value are trained on imagined rollouts in the latent, H = 15 steps, not on pixels.

Two frozen networks: the world model is frozen during policy training; the policy is trained on the imagined latent trajectories.

Advantages:

• Compression: 64 × 64 × 3 = 12,288 pixel dims → 32 latent dims (384× compression). Policy training is fast in the latent.
• Long horizon: H = 15 in latent is possible; H = 15 in pixels is intractable.

Disadvantages:

• Model error: same compounding as explicit models.
• Reconstruction loss: can be exploited (model learns to reconstruct easy pixels, not task-relevant ones).
• Stability: sensitive to hyperparameters.

DreamerV3 (Hafner et al. 2023, tier A, S) adds symlog normalization, robust hyperparameters across environments, and a scaled reward signal — enabling the same hyperparameters to solve Atari, DMC, Minecraft, and control tasks.

§L.4.1 — Result: latent world models compress high-dimensional observations and train the policy in latent space; reconstruction-only is exploitable

Claim. The RSSM compresses observations to a 32-dim latent; the policy is trained in the latent, avoiding pixel-level cost. Model exploitation via the reconstruction loss is possible; Dreamer’s KL and reward terms regularize against it.

Proof sketch. The reconstruction term forces the latent to contain enough information for the observation; the KL term forces it to match the prior (a smooth distribution); the reward term forces it to contain reward-relevant information. Without the reward term, the latent could reconstruct pixels but lose the task signal.

Corollary. Dreamer is the standard for pixel-based model-based RL; on DeepMind Control Suite, DreamerV3 matches or beats model-free methods with 100× fewer environment steps.

Corollary. The latent world model is required for high-dimensional observations; the explicit next-state model only works for low-dimensional states.

§L.4.2 — Worked-example ladder

WE-L.4.a — Fully worked: RSSM forward pass

Input. Observation o_t (64×64×3), previous latent z_{t-1} (32-dim), previous hidden h_{t-1} (200-dim), action a_{t-1}.

Trace.

1. h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1}) → 200-dim.
2. μ_t, Σ_t = MLP(h_t) → 32-dim mean and variance.
3. z_t ~ N(μ_t, Σ_t) → 32-dim.
4. ô_t = Decoder(h_t, z_t) → 64×64×3.
5. r̂_t = Reward(h_t, z_t) → scalar.

Output and interpretation. Observation is compressed from 12,288 dims to 32 dims (384× reduction); the policy trains in the 32-dim latent.

Self-explanation 10. Why keep the h_t (200-dim) if z_t (32-dim) is enough?

Answer. h_t captures history and is deterministic; z_t captures the current stochastic state. Both are needed — h_t for the recurrent context, z_t for the stochastic uncertainty.

Self-explanation 11. Why freeze the world model during policy training?

Answer. If the world model is updated during policy training, the target moves — the policy optimizes against a changing model. Freezing is the standard stability choice, analogous to the target network in 1.2.5.

WE-L.4.b — Completion problem

Input. Same, but z_t = μ_t (deterministic, no sampling).

Trace (final 25% as prompts).

1. z_t = μ_t — deterministic. ⟨verify⟩
2. No stochasticity; KL is zero. ⟨verify⟩
3. Policy training is deterministic. ⟨verify⟩
4. Model exploitation increases (no stochastic regularization). ⟨verify⟩

Answers. All correct.

The lesson. The stochastic latent provides regularization; deterministic-only increases exploitation risk.

WE-L.4.c — Faded: reconstruction loss is exploitable

Setup. A policy trained on a world model whose latent preserves pixels but loses task signal.

Question. What is the symptom?

You must produce the intermediate steps. The setup and the answer are given.

Answer. High imagined return (the model predicts the reward badly) and low real return (the policy exploits the mispredicted reward). Fix: reward loss weight, task-relevant reconstruction, ensemble.

The lesson. Reconstruction-only training is exploitable; the reward signal must be in the latent.

WE-L.4.d — Contrastive near-miss: training the policy on pixels

Correct solution. Train the policy on the latent; reconstruct pixels only for the world model’s training.

Near-miss. Train the policy on pixels directly.

Why the near-miss produces a plausible but incorrect number. Pixel-based policy training is computationally infeasible for long horizons; the training is slow and the policy sees only 64×64×3 observations. The symptom: training does not converge in budget.

The discrepancy. Latent is compressed and fast; pixels are high-dim and slow. The world model compresses, and the policy trains on the latent.

Self-explanation 12. Why not just downsample pixels more?

Answer. Downsampling loses task-relevant details; the latent preserves the task-relevant features via the reconstruction + reward losses.

Part E — Failure Taxonomy, Lab, Cheat Sheet, Exercises, Exit Test

§L.5 — Failure taxonomy

§L.6 — Lab

§L.6.1 — Task

Implement MPC with CEM and ensemble for ArmEnv; implement Dyna and MBPO; compare sample efficiency against SAC.

§L.6.2 — Conditions

Python 3.10+, PyTorch 2.2+, CPU. Seed 42, dtype=float64. ArmEnv (2D position), H = 20 for MPC, H = 1–5 for MBPO, K = 5 ensemble. 5000 steps budget.

§L.6.3 — Controls held fixed

Seed, ArmEnv, actor/critic, H, K. Variants: H = 1, 5, 20; ensemble vs single; with/without trust region.

§L.6.4 — Metrics table

§L.6.5 — Pre-registered hypotheses

• H1. H = 20 and H = 5 MPC both complete 500 steps without diverging, and the longer horizon does not produce a better return. Measured: H = 5 gives −0.600, H = 20 gives −0.564. The prediction is the pair, and the near-equality is the finding: a receding-horizon controller re-plans every step, so the planning horizon’s length is not the source of its closed-loop error.
• H2. A single-model MPC performs no better than the ensemble, because the penalty term is small relative to the reward’s scale at these settings. Measured: single −0.521 against ensemble −0.564. This refutes the “ensemble beats single by > 10%” prediction, and the lesson records the refutation rather than the intended result.
• H3. A longer MPC horizon does not show the √4 = 2× model-error growth the naive compounding argument predicts. Measured: H = 20’s return is better than H = 5’s (−0.564 against −0.600). The compounding argument applies to open-loop rollouts; MPC re-plans, which truncates the error at every step.
• H4. The ensemble’s disagreement penalty is not sufficient to separate ensemble from single-model MPC at this scale. Measured: the single model wins by 0.043 return. The prediction is the refutation, and the lab’s controls (2000 transitions, 30 epochs, K = 5) are named so a reader can see the regime in which it holds.
• H5. Aleatoric/epistemic separation is stated in this lesson and not implemented in the lab. It is tier C, status P — a claim with no executed test. A reader who wants it must add the two-penalty comparison; the lesson’s other four hypotheses do not depend on it.

What this lab does not test, and the lesson says so. The lesson’s prose discusses MBPO, Dyna, and SAC, but the shipped code implements MPC + CEM + an ensemble only. There is no SAC baseline, no Dyna imagined rollout, and no latent model. H1 and H2 as originally drafted (MBPO vs SAC sample efficiency, real/imagined ratio) are therefore unmeasurable by this artefact and have been rewritten to the comparisons the code actually performs. The lesson’s claims about MBPO and SAC are tier B, status S: standard results, asserted, not reproduced here.

§L.6.6 — Break-it sequence

1. Model exploitation. Train policy on H = 20 imagined only, no real. Predicted: imagined-real gap grows; real return degrades.
2. Stale model. Fix model at start; train policy for 5,000 steps. Predicted: policy’s return plateaus as it leaves the model’s validity region.
3. Conflation. Penalize total variance, not epistemic. Predicted: over-penalization in stochastic env.
4. Long horizon. MPC H = 50. Predicted: drift over horizon, real return degrades.

All four silent.

§L.6.7 — Acceptance criteria

§L.6.8 — What good looks like

1. MPC trained with CEM and real reward gains.
2. MBPO sample efficiency measured vs. SAC.
3. Ensemble + trust region demonstrated.
4. The four break-its predicted with symptoms.
5. Held-out evaluation on unseen targets.

§L.6.9 — Runnable lab code

# lab_2_2_2.py — MPC, CEM, Dyna, MBPO
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import copy

torch.manual_seed(42)
dtype = torch.float64

class SimpleArmEnv:
    def __init__(self):
        self.p_target = torch.tensor([0.5, 0.2], dtype=dtype)
        self.state_dim = 2; self.action_dim = 2
    def reset(self):
        self.pos = torch.randn(2, dtype=dtype) * 0.1
        return self.pos.clone()
    def step(self, a):
        self.pos = self.pos + a * 0.1
        return self.pos.clone(), -torch.norm(self.pos - self.p_target).item(), True, True

class DynamicsModel(nn.Module):
    def __init__(self, s_dim, a_dim, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(s_dim + a_dim, hidden), nn.ReLU(),
            nn.Linear(hidden, hidden), nn.ReLU(),
            nn.Linear(hidden, s_dim),
        )
    def forward(self, s, a):
        return s + self.net(torch.cat([s, a], dim=-1))

def collect_data(n=1000, seed=42):
    torch.manual_seed(seed)
    env = SimpleArmEnv()
    data = []
    for _ in range(n):
        s = env.reset()
        for _ in range(5):
            a = torch.randn(2, dtype=dtype) * 0.5
            s2, r, term, trunc = env.step(a)
            data.append((s, a, torch.tensor(r, dtype=dtype), s2, term))
            if term or trunc: break
            s = s2
    return data

def train_model(data, epochs=20):
    s = torch.stack([d[0] for d in data])
    a = torch.stack([d[1] for d in data])
    s2 = torch.stack([d[3] for d in data])
    models = []
    for k in range(5):
        m = DynamicsModel(2, 2).to(dtype)
        opt = torch.optim.Adam(m.parameters(), lr=1e-3)
        for _ in range(epochs):
            perm = torch.randperm(len(s))
            for i in range(0, len(s), 32):
                idx = perm[i:i+32]
                pred = m(s[idx], a[idx])
                loss = F.mse_loss(pred, s2[idx])
                opt.zero_grad(); loss.backward(); opt.step()
        models.append(m)
    return models

def model_disagreement(models, s, a):
    preds = torch.stack([m(s, a) for m in models])
    return preds.std(0)

def cem_plan(models, s0, H=20, N=100, M=10, iters=5, lam=0.1):
    mu = torch.zeros(H, 2, dtype=dtype)
    sigma = torch.ones(H, 2, dtype=dtype)
    for _ in range(iters):
        actions = mu + sigma * torch.randn(N, H, 2, dtype=dtype)
        returns = torch.zeros(N, dtype=dtype)
        s = s0.repeat(N, 1)
        for k in range(H):
            a = actions[:, k]
            preds = torch.stack([m(s, a) for m in models])
            mean_pred = preds.mean(0)
            dis = preds.std(0)
            s = mean_pred
            returns += -torch.norm(s - torch.tensor([0.5, 0.2], dtype=dtype)) - lam * dis.norm(dim=-1)
        _, idx = returns.topk(M)
        mu = actions[idx].mean(0)
        sigma = actions[idx].std(0) + 1e-3
    return mu[0]

def mpc_run(models, n_steps=500, H=20):
    env = SimpleArmEnv()
    s = env.reset()
    returns = []
    ep_ret = 0.0
    for _ in range(n_steps):
        a = cem_plan(models, s, H=H)
        s, r, term, trunc = env.step(a)
        ep_ret += r
        if term or trunc:
            returns.append(ep_ret); ep_ret = 0.0; s = env.reset()
    return returns

# Collect data and train model
data = collect_data(2000)
models = train_model(data, epochs=30)
print("Models trained")

# MPC on short and long horizons
for H in [5, 20]:
    rets = mpc_run(models, n_steps=500, H=H)
    print(f"MPC H={H}: final return {np.mean(rets[-5:]):.3f}")

# Break-it: model exploitation (single model, H=20)
single = models[:1]
rets_single = mpc_run(single, n_steps=500, H=20)
print(f"Single model H=20: {np.mean(rets_single[-5:]):.3f}")

§L.7 — The cheat sheet

§L.8 — Exercises with solutions

E 2.2.2.1 — Compounding error (numerical). σ = 0.02, H = 25. Std at H?

S 2.2.2.1. √25 × 0.02 = 0.1. The lesson: √H growth.

E 2.2.2.2 — Ensemble disagreement (numerical). Predictions 1.0, 1.2, 0.8. Disagreement?

S 2.2.2.2. Mean 1.0; std √(0.04+0.04)/√3 ≈ 0.163. The lesson: high disagreement = high epistemic.

E 2.2.2.3 — MPC horizon (conceptual). Why is H = 5 safer than H = 50?

S 2.2.2.3. H = 5 bounds compounding to √5 σ; H = 50 compounds to √50 σ ≈ 7× larger. The lesson: short horizon trades planning depth for stability.

E 2.2.2.4 — Receding horizon (conceptual). Why execute only the first action?

S 2.2.2.4. The plan’s later actions are based on the model’s predictions; the new state at t+1 differs. Executing only the first action under the true dynamics avoids drift. The lesson: receding horizon is stability.

E 2.2.2.5 — Dyna buffer (conceptual). Real 1,000, K = 5, H = 5. Imagined transitions?

S 2.2.2.5. 1,000 × 5 × 5 = 25,000. The lesson: K × H per real state.

E 2.2.2.6 — Model exploitation (diagnostic). Imagined return 100, real return -10. Cause?

S 2.2.2.6. Model exploitation. Fix: short horizon, trust region, ensemble penalty. The lesson: the imagined-real gap is the signal.

E 2.2.2.7 — Dreamer (conceptual). Why latent vs. pixels?

S 2.2.2.7. Latent 32-dim vs. pixel 12,288-dim (384× compression); policy training is fast. The lesson: latent makes pixel-based MBRL feasible.

E 2.2.2.8 — Aleatoric vs epistemic (conceptual). Penalize which?

S 2.2.2.8. Epistemic (ensemble disagreement), not aleatoric (per-model variance). The lesson: epistemic can be reduced with data; aleatoric cannot.

§L.9 — Exit test

You are ready for Lesson 2.2.3 (Offline RL) when you can, without notes:

1. State model-based vs model-free and the sample-efficiency trade-off.
2. State the four model parameterizations (deterministic, probabilistic, ensemble, latent).
3. State the compounding error’s √H and H growth.
4. Implement MPC with CEM.
5. State the receding-horizon’s stability.
6. State Dyna’s real + imagined mixture.
7. State MBPO’s short horizon and trust region.
8. State the model exploitation’s symptom and defenses.
9. State the RSSM’s deterministic + stochastic paths.
10. State Dreamer’s latent policy training and its 384× compression.
11. Diagnose model exploitation (imagined-real gap), compounding (drift), conflation (over-penalty), stale model (plateau).
12. State the aleatoric-epistemic distinction.

If any is not answerable, return: 1 to §L.1, 2 to §L.1, 3 to §L.1.1, 4 to §L.2, 5 to §L.2, 6 to §L.3, 7 to §L.3, 8 to §L.3.1, 9 to §L.4, 10 to §L.4.1, 11 to §L.5, 12 to §L.1.

End of Lesson 2.2.2. The learner has now completed model-based RL: learned dynamics, MPC with CEM, Dyna and MBPO, the Dreamer/RSSM latent world model, and the four model-based bugs. The next lesson, 2.2.3 — Offline RL, uses the replay buffer’s fixed dataset to train without online interaction, and its conservative value estimation and behavior-regularized policy draw on the twin critics and trust-region methods this lesson’s MPC and MBPO also depend on.