Week 12: World Models and Reinforcement Learning

Purpose of this lecture#

The generative models studied so far produce outputs in response to noise and conditioning signals. This lecture applies the same neural machinery to a fundamentally different goal: building a world model — a learned simulator that predicts how the environment transitions in response to agent actions, enabling an agent to plan by imagining future states rather than by acting in the real world. World models represent the deepest integration of generative modeling with decision-making, and they provide the conceptual foundation for understanding how foundation models will be used in physical AI.

World models: architecture and role#

World models (Ha & Schmidhuber, 2018) are a class of generative models that learn to simulate the environment in which an agent operates. The key insight is that the same neural machinery used to generate images, audio, or text can be applied to predict the next state of the environment given the current state and action: $p(s_{t+1} \mid s_t, a_t)$. This predictive capability enables model-based reinforcement learning, where the agent can plan by simulating many possible actions in its internal world model rather than interacting with the real environment.

The world model architecture typically consists of three components:

1. Encoder: maps observations $o_t$ to latent states $z_t$
2. Dynamics model: learns the transition function $p(z_{t+1} \mid z_t, a_t)$
3. Decoder: maps latent states back to observations $\hat{o}_t = p(o_t \mid z_t)$

This structure enables imagination: the agent can generate sequences of imagined states by sampling from the dynamics model, then decode them to produce simulated observations. This imagination capability is crucial for planning, as the agent can evaluate many possible action sequences without actually executing them in the real environment.

The RSSM architecture#

The recurrent state space model (RSSM; Hafner et al., 2019) is the workhorse world-model backbone behind PlaNet and Dreamer. It factors the latent state into a deterministic recurrent path $h_t$ and a stochastic path $z_t$:

At training time the posterior $q_\phi(z_t \mid h_t, o_t)$ sees the current observation (filter). At imagination time only the prior $p_\theta(z_t \mid h_t)$ is available — the model must predict the next latent without $o_t$:

This prior/posterior split is exactly the VAE structure from Week 2 applied to sequential data: $h_t$ plays the role of context, $z_t$ is the latent code, and the decoder reconstructs $o_t$ (and usually reward $r_t$).

RSSM ELBO#

For a trajectory $(o_{1:T}, a_{1:T})$, the world-model objective is an ELBO on $\log p(o_{1:T} \mid a_{1:T})$:

Intuition: reconstruction forces $z_t$ to carry information needed to explain $o_t$; the KL keeps the posterior close to the prior so that open-loop rollouts (no observations) stay on-distribution. If the KL is too weak, imagination diverges from filtered states; if too strong, $z_t$ collapses and the model becomes a deterministic RNN with limited multi-modality.

Dreamer: latent imagination for policy learning#

Dreamer (Hafner et al., 2020; DreamerV2/V3) freezes the RSSM as a differentiable simulator and trains policy components inside it:

1. World model — RSSM + reward head, trained on real environment data with the ELBO above
2. Actor $\pi_\psi(a_t \mid h_t, z_t)$ — proposes actions in latent space
3. Critic $V_\xi(h_t, z_t)$ — estimates expected return from a latent state

Latent imagination#

Starting from a real latent $(h_t, z_t)$ obtained by filtering observations, Dreamer rolls out $H$ imagined steps without new environment interaction:

for $\tau = t, \ldots, t+H-1$. Returns are bootstrapped with the critic (e.g. $\lambda$-returns / GAE-style targets). The critic fits these targets; the actor maximizes expected return by backpropagating through the differentiable dynamics (analytic gradients through $h,z$ when continuous, or reinforce-style estimators when discrete).

Why this is sample-efficient: each real transition can seed many imagined trajectories. Policy gradients use model-generated experience, so the agent needs far fewer environment steps than model-free methods — provided the world model is accurate on action-relevant directions.

DreamerV3 adds practical stabilizers (symlog observations/rewards, free-bits KL floors, percentile return normalization) that make the same recipe work across discrete and continuous control without per-domain retuning.

Model predictive control and latent-space planning#

Model predictive control (MPC) uses the world model at decision time rather than amortizing a policy:

1. Predict short-horizon trajectories under candidate action sequences
2. Optimize expected return or cost over the horizon
3. Execute only the first action, then replan (receding horizon)

In latent space this is cheaper than pixel-space planning: optimize over $z$-trajectories (dimension $d_z \ll H\times W\times C$), optionally decode only for visualization. Cross-entropy method (CEM), MPPI, or gradient-based planners all apply to $p_\theta(z_{t+1}\mid h_{t+1})$.

Closed-loop vs. open-loop: re-encoding observations (or resetting $z$ to the filtered posterior) every $k$ steps bounds error growth. Pure open-loop dreams of length 40 with a slightly wrong $A,B$ already diverge — the lab quantifies this.

Sample efficiency and model-based vs. model-free RL#

Sample efficiency measures real environment interactions needed to reach a performance level. Model-based methods reuse a learned simulator for many gradient steps; model-free methods require a real transition (or replay of one) for each update.

Model exploitation: the actor may discover adversarial action sequences that the world model scores highly but the real environment does not. Mitigations include short $H$, ensembles, uncertainty penalties, and continual fine-tuning of the model on on-policy data.

World models in physical AI#

In robotics and autonomous systems, world models support:

1. Lookahead — predict contact, occlusion, and object motion before acting
2. Uncertainty — stochastic $z_t$ expresses multi-modal futures (e.g. which way a door swings)
3. Sim-to-real — train policies in imagination; transfer with residual adaptation

Challenges remain: visual domain gap, long-horizon credit, and ensuring the latent is controllable (action-relevant) rather than merely reconstructive.

Cross-course context: world models across the curriculum#

The same ELBO and latent-dynamics machinery appears under different names: a VAE for images, an RSSM for pixels+actions, a diffusion policy for action chunks. Generative modeling supplies the likelihood and representation tools; RL supplies the decision objective.

Browser lab: latent dynamics rollout (RSSM / Dreamer-style)#

Linear latent world model $z_{t+1} = A z_t + B a_t + \xi$. One-step fit looks fine; multi-step open-loop imagination compounds error — why MBRL needs short horizons or closed-loop replan.

import numpy as np

rng = np.random.default_rng(0)
d_z, d_a = 4, 2
# True dynamics (stable spiral-ish)
A_true = np.array([
    [0.90, -0.10, 0.00, 0.00],
    [0.10,  0.90, 0.00, 0.00],
    [0.00,  0.00, 0.85, 0.05],
    [0.00,  0.00, -0.05, 0.85],
])
B_true = rng.normal(0, 0.3, size=(d_z, d_a))
noise = 0.05

def rollout(A, B, z0, actions, noise_std=0.0):
    z = z0.copy()
    traj = [z.copy()]
    for a in actions:
        z = A @ z + B @ a + rng.normal(0, noise_std, size=d_z)
        traj.append(z.copy())
    return np.array(traj)

# Collect data and fit least-squares ẑ' ≈ A z + B a
T_data = 800
z = rng.normal(size=d_z)
Zs, Zs_next, As = [], [], []
for _ in range(T_data):
    a = rng.normal(size=d_a)
    z_next = A_true @ z + B_true @ a + rng.normal(0, noise, size=d_z)
    Zs.append(z); Zs_next.append(z_next); As.append(a)
    z = z_next
Z, Zp, Act = map(np.array, (Zs, Zs_next, As))
# z' = A z + B a  ⇒  Zp = Z @ A.T + Act @ B.T
Phi = np.hstack([Z, Act])
Theta, *_ = np.linalg.lstsq(Phi, Zp, rcond=None)  # (d_z+d_a, d_z)
A_hat = Theta[:d_z, :].T
B_hat = Theta[d_z:, :].T

print(f"One-step train MSE: {np.mean((Zp - (Z @ A_hat.T + Act @ B_hat.T))**2):.4f}")
print(f"||A_hat - A_true||_F: {np.linalg.norm(A_hat - A_true):.3f}")

# Multi-step open-loop imagination error
z0 = rng.normal(size=d_z)
actions = rng.normal(size=(40, d_a))
true_tr = rollout(A_true, B_true, z0, actions, noise_std=0.0)
pred_tr = rollout(A_hat, B_hat, z0, actions, noise_std=0.0)
print(f"\n{'H':>4}  {'MSE@H':>10}  {'||z|| drift':>12}")
for H in [1, 5, 10, 20, 40]:
    mse = np.mean((true_tr[H] - pred_tr[H]) ** 2)
    drift = np.linalg.norm(pred_tr[H] - true_tr[H])
    print(f"{H:4d}  {mse:10.4f}  {drift:12.3f}")

# Closed-loop: replan every k steps (reset latent to true — like re-encoding obs)
def closed_loop_mse(k, H=40):
    z_t, z_p = z0.copy(), z0.copy()
    errs = []
    for t in range(H):
        a = actions[t]
        z_t = A_true @ z_t + B_true @ a
        z_p = A_hat @ z_p + B_hat @ a
        if (t + 1) % k == 0:
            z_p = z_t.copy()  # re-encode from observation
        errs.append(np.mean((z_t - z_p) ** 2))
    return np.mean(errs)

print("\nMean MSE over 40 steps with re-encode every k:")
for k in [40, 10, 5, 1]:
    print(f"  k={k:2d}  MSE={closed_loop_mse(k):.4f}")
print("Notice: one-step fit ≠ long dream accuracy; short imagination + frequent re-encoding stabilizes MBRL.")

What to try: raise process noise = 0.2 or shrink T_data — multi-step error blows up faster.

Key takeaways#

World models are action-conditioned generative models of environment dynamics. The RSSM splits latent state into a recurrent deterministic path $h_t$ and a stochastic code $z_t$, trained with an ELBO that balances reconstruction against prior–posterior KL so open-loop imagination stays on-distribution. Dreamer freezes that model and improves an actor–critic inside short latent rollouts, trading model bias for sample efficiency. Latent MPC plans online in the same space. One-step fit quality does not guarantee multi-step dream accuracy — short horizons, re-encoding, and uncertainty-aware latents are essential.

Conceptual questions#

1. A world model predicts the next observation given the current state and action. What are the advantages of this approach over direct policy learning? What are the potential disadvantages?

2. Write the RSSM ELBO terms for one timestep and explain the role of the KL between $q_\phi(z_t \mid h_t, o_t)$ and $p_\theta(z_t \mid h_t)$. What goes wrong if this KL is driven to zero? What goes wrong if it is ignored?

3. Dreamer uses latent imagination to improve an actor–critic. What are the key components of this process, and how does it differ from model-free learning and from online latent MPC?

4. Model-based methods can be more sample-efficient than model-free methods, but they also suffer from model error. How does Dreamer (and the browser lab) address compounding imagination error?

5. How do world models in generative modeling relate to world models in reinforcement learning? What are the key similarities and differences?

Knowledge Check#

Check world models and generative RL.

Looking ahead#

With world models linking generation to decision-making, the course turns to the risks that accompany powerful generative systems.

Week 13: Safety, Misuse, and Alignment. We examine misuse vectors (deepfakes, memorization, adversarial inputs), detection and differential-privacy defenses, and the RLHF/DPO alignment techniques that steer model behavior toward human preferences.

Further reading#

• Ha, D., & Schmidhuber, J. (2018). World Models. NeurIPS.
• Hafner, D., et al. (2019). Learning Latent Dynamics for Planning from Pixels (PlaNet, RSSM). ICML.
• Hafner, D., et al. (2020). Dream to Control: Learning Behaviors by Latent Imagination (Dreamer). ICML.
• Hafner, D., et al. (2023). Mastering Diverse Domains through World Models (DreamerV3). arXiv.
• Sutton, R. S. (1991). Dyna: An Integrated Architecture for Learning, Planning, and Reacting. ACM SIGART.