World Models — Part 1: Inside the Latent (VAEs, RSSM, and Learning in a Dream)

The machinery, derived from scratch: why we compress pixels into a latent, how the VAE and the ELBO actually work, how a Recurrent State-Space Model turns a single-frame encoder into a simulator, and how Dreamer trains a policy entirely inside its own imagination — with the real systems that run on these ideas today.

TL;DR. A world model needs three things: a way to compress observations into a small state, a way to predict how that state evolves under actions, and a way to use those predictions. This part derives all three. We build the VAE and its ELBO from first principles, extend it across time into a Recurrent State-Space Model (RSSM), and then show how Dreamer trains an actor–critic entirely on imagined trajectories — never touching the environment. We close with the systems that run on this today, and a slide-by-slide skeleton for presenting it.


Where Part 0 left off

In Part 0 we established the vocabulary: a world model learns \(p(z_{t+1} \mid z_t, a_t)\) — the dynamics of an environment conditioned on actions — as opposed to a language model’s \(p(x_t \mid x_{<t})\). We sketched Ha & Schmidhuber’s V–M–C decomposition3 and said the magic was that a trained dynamics model lets you generate experience in imagination.

That was the what. This part is the how. Three questions drive everything below:

  1. Why compress? Why not predict pixels directly?
  2. How do we learn the compression when nobody hands us the latent variable?
  3. How do we turn a per-frame encoder into a simulator you can roll forward, and then plan in?

Each has a clean answer, and each answer is a piece of machinery you will meet in every modern system.


Why compress at all?

Take a modest \(64 \times 64\) RGB frame. That is \(64 \times 64 \times 3 = 12{,}288\) numbers per timestep. A one-minute episode at 20 Hz is roughly 15 million numbers. Predicting the next frame directly means modelling a distribution over a 12,288-dimensional space — and most of those dimensions are things you do not care about: the exact texture of gravel, the flicker of a shadow, sensor noise.

observation oₜ — 12,288 dims encoder latent state zₜ ~32 dims dynamics predicted zₜ₊₁ cheap to roll forward
Roughly a 380× reduction. Prediction happens in the small space, not the big one — which is what makes rolling the model forward hundreds of steps affordable.

The bet of latent world models is that a compact \(z_t\) retains everything decision-relevant and discards the rest. Two consequences follow immediately, and both matter in practice:

The cost is that the latent is unobserved. Nobody labels \(z_t\). We have to learn the encoder and the dynamics jointly, from observations alone. That is precisely the problem the VAE solves.


The variational autoencoder

Assume the data is generated by a latent-variable model: draw a latent from a simple prior, then decode it.

\[z \sim p(z) = \mathcal{N}(0, I), \qquad x \sim p_\theta(x \mid z).\]

To train \(\theta\) by maximum likelihood we would need the marginal

\[p_\theta(x) = \int p_\theta(x \mid z)\, p(z)\, dz,\]

which is intractable — it integrates over every possible latent. The variational trick1,2 is to introduce an inference network \(q_\phi(z \mid x)\) that guesses which latents could have produced \(x\), and then bound the likelihood.

Deriving the ELBO

Start from the log-likelihood and multiply inside by \(q_\phi(z\mid x)/q_\phi(z\mid x)\):

\[\log p_\theta(x) = \log \int q_\phi(z \mid x)\, \frac{p_\theta(x \mid z)\, p(z)}{q_\phi(z \mid x)}\, dz = \log \mathbb{E}_{q_\phi(z \mid x)}\!\left[\frac{p_\theta(x \mid z)\, p(z)}{q_\phi(z \mid x)}\right].\]

Since \(\log\) is concave, Jensen’s inequality moves it inside the expectation:

\[\log p_\theta(x) \;\ge\; \mathbb{E}_{q_\phi(z \mid x)}\big[\log p_\theta(x \mid z)\big] \; - \; D_{\mathrm{KL}}\!\big(q_\phi(z \mid x)\,\|\,p(z)\big) \;=\; \mathcal{L}_{\text{ELBO}}(\theta,\phi).\]

The gap between the two sides is exactly \(D_{\mathrm{KL}}(q_\phi(z\mid x) \,\|\, p_\theta(z \mid x))\) — how wrong our guessed posterior is. Maximising the ELBO therefore does two jobs at once: it fits the data and tightens the inference network.

Read the two terms plainly:

They pull against each other, and that tension is the representation.

x encoder q(z | x) z decoder p(x | z) μ, σ KL( q ‖ N(0, I) ) — keep the latent space smooth E[ log p(x | z) ] — keep enough to rebuild x reparameterise: z = μ(x) + σ(x) ⊙ ε — so gradients flow through the sample
The VAE in one picture. The two dashed annotations are precisely the two ELBO terms — one pulling the latent toward fidelity, the other toward a well-behaved prior.

The reparameterisation trick

One obstacle remains: we must backpropagate through a sample \(z \sim q_\phi(z\mid x)\), and sampling is not differentiable. The fix1 is to push the randomness into a parameter-free variable:

\[z = \mu_\phi(x) + \sigma_\phi(x) \odot \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, I).\]

Now \(z\) is a deterministic, differentiable function of \(\phi\) and an external noise draw. Gradients flow to \(\mu_\phi\) and \(\sigma_\phi\) cleanly. This single line is what made deep latent-variable models trainable at scale, and it is used unchanged inside every model below.


From frames to sequences — the RSSM

A VAE encodes one frame. A world model must answer: given this state and this action, what is the next state? The dominant answer is the Recurrent State-Space Model introduced with PlaNet4 and carried through the whole Dreamer line5,6,7.

Its key design decision is to split the state in two:

\[s_t = (h_t,\; z_t)\]

Why both? A purely deterministic model cannot represent uncertainty, so it blurs multiple futures into an average. A purely stochastic model tends to lose long-range information across sampling steps. Splitting gets memory and uncertainty. The four learned components are:

\[\begin{aligned} \text{recurrence:} &\quad h_t = f_\theta(h_{t-1},\, z_{t-1},\, a_{t-1}) \\[2pt] \text{prior (transition):} &\quad \hat{z}_t \sim p_\theta(\hat{z}_t \mid h_t) \\[2pt] \text{posterior (representation):} &\quad z_t \sim q_\phi(z_t \mid h_t,\, o_t) \\[2pt] \text{heads:} &\quad \hat{o}_t \sim p_\theta(o_t \mid h_t, z_t), \quad \hat{r}_t \sim p_\theta(r_t \mid h_t, z_t) \end{aligned}\]

The distinction between prior and posterior is the entire trick, and it is worth stating slowly:

Training pushes the prior toward the posterior. Once they agree, you can drop the observation entirely and run on the prior — and that is a simulator.

hₜ₋₁ hₜ hₜ₊₁ GRU · aₜ₋₁ GRU · aₜ zₜ₋₁ zₜ zₜ₊₁ oₜ₋₁ oₜ oₜ₊₁ deterministic (memory) stochastic (uncertainty) posterior uses oₜ · prior must guess without it KL(posterior ‖ prior) is what teaches the model to simulate
The RSSM unrolled. Solid path = memory that always flows; dashed boxes = observations, which are available during training but deliberately withheld when the model is used as a simulator.

The world-model objective

Everything above collapses into one loss — a sequential ELBO. Over a sampled trajectory,

\[\mathcal{L}(\theta,\phi) = \mathbb{E}_{q_\phi}\left[\sum_{t}\underbrace{-\log p_\theta(o_t \mid h_t, z_t)}_{\text{reconstruct observation}} \underbrace{-\log p_\theta(r_t \mid h_t, z_t)}_{\text{predict reward}} + \; \beta\underbrace{D_{\mathrm{KL}}\big(q_\phi(z_t \mid h_t, o_t)\,\big\|\,p_\theta(z_t \mid h_t)\big)}_{\text{make the prior match the posterior}}\right]\]

Three terms, three jobs:

  1. Reconstruction forces the latent to carry the scene.
  2. Reward prediction forces it to carry what matters for the task — the single line that turns a generic video model into a control model.
  3. KL is the simulator term. It penalises the gap between “state given the picture” and “state predicted from memory + action”. Drive it down and the model can dream.

That third term deserves emphasis, because it is the one people skim: the KL is not merely regularisation here — it is the training signal that makes open-loop rollout possible at all.


Learning in a dream

Now the payoff. With a trained RSSM you can generate trajectories without the environment: start from a real encoded state, then repeatedly sample the prior and feed back your own predictions.

Dreamer5 trains an actor and a critic purely on these imagined rollouts, typically ~15 steps deep:

\[\pi_\psi(a_t \mid s_t) \quad\text{(actor)}, \qquad v_\xi(s_t) \approx \mathbb{E}\Big[\textstyle\sum_{k\ge t} \gamma^{\,k-t} r_k\Big] \quad\text{(critic)}\]

Returns use a \(\lambda\)-weighted mixture8,15 that trades bias against variance:

\[V^\lambda_t = r_t + \gamma\Big[(1-\lambda)\, v_\xi(s_{t+1}) + \lambda\, V^\lambda_{t+1}\Big]\]

The critic regresses onto \(V^\lambda_t\); the actor maximises it. Crucially, because the whole rollout is differentiable, the actor can get gradients straight through the learned dynamics — a luxury no model-free method has.

real experience (posterior, uses oₜ) replayed from the buffer start state Vλ = 8.2 Vλ = 3.1 Vλ = −1.4 imagined futures (prior only — no environment) actor is pushed toward high Vλ
Training happens here. Each dashed branch costs a few matrix multiplies — no simulator step, no robot, no risk. The actor is updated by gradients flowing back through the learned dynamics.

The practical consequence is stark: Dreamer-style agents are often one to two orders of magnitude more sample-efficient than model-free baselines on pixel control, because most of their learning happens in a dream rather than in the world.


What actually made it work

The idea above dates to 2018–2019. What turned it from a promising result into a robust one was a set of unglamorous fixes in DreamerV37 — worth knowing, because they are the difference between a demo and a system:

The headline result: one fixed configuration across more than 150 tasks, and the first agent to collect diamonds in Minecraft from scratch, with no human data and no curriculum7. If you present one empirical fact about world models, that is a strong candidate — it is concrete, verifiable, and instantly legible to a non-specialist.


Real systems running on these ideas

The machinery above is not academic. Here is where it shows up, with what each system actually demonstrates:

DreamerV3 The RSSM + imagination recipe, hardened. One hyperparameter set, 150+ tasks, diamonds in Minecraft from scratch. Take-away: latent imagination is now a general method, not a per-task craft.7
Genie An 11B-parameter model trained on unlabelled internet gameplay video. It infers a latent action space with no action labels, then turns a single image — even a sketch — into a playable world. Take-away: actions can be discovered, not supplied. This is the bridge from passive video to interactive simulation.11
Sora Positioned explicitly as a step toward "world simulators": long-horizon video with emergent 3D consistency and object permanence — alongside well-documented physics failures. Take-away: scale buys visual coherence, but coherence is not physics. A useful cautionary example.12
GameNGen A diffusion model that runs DOOM interactively at ~20 FPS on a single TPU; human raters distinguishing short clips from the real engine perform barely above chance. Take-away: a neural network can be the game engine — real-time, interactive, learned.13
GAIA-1 A ~9B generative driving world model that produces plausible future driving video conditioned on past video, text, and actions. Take-away: the safety case — generate rare, dangerous scenarios instead of waiting to encounter them.14
The JEPA line Predicts in representation space rather than pixel space, deliberately refusing to model unpredictable detail. Take-away: the main live alternative to reconstruction — and the subject of Part 2.9,10

Notice the split running through this list. DreamerV3 and GAIA-1 reconstruct; Genie and GameNGen generate interactively; JEPA refuses to reconstruct at all. That disagreement — what should a world model actually predict? — is the live research question, and it is exactly where Part 2 picks up.


How to present this in ten minutes

Since a stated goal here is to be able to explain this from memory, here is a skeleton that has worked for me. Five slides, one idea each.

  1. The gap. “A language model predicts what a person would write next. A world model predicts what the world will do next — given what you do.” One sentence, one table (from Part 0). No maths.
  2. The compression. Show Figure 1. “12,288 numbers become 32. We predict in the small space.” The 380× number does the persuading.
  3. The two-part state. Show Figure 3. “Memory that never forgets, plus a random part for what genuinely can’t be predicted.” Then the one line that matters: the posterior sees the picture; the prior has to guess. Train them to agree, and you can throw the picture away.
  4. The dream. Show Figure 4. “Now training costs matrix multiplies instead of robot hours.” Land it with DreamerV3’s diamonds.
  5. The open question. “Reconstruct pixels, generate interactively, or predict representations?” Name Sora’s physics failures against Genie’s playability. Finish on the unresolved question — audiences remember tension, not summaries.

Two things to keep in reserve for questions: why the KL term is the simulator term (§ the objective), and why the reward head is what separates a world model from a video model. Those are the two points specialists probe first.


Where this goes next

Part 2 — Predicting representations: JEPA, energy-based learning, and the case against pixels. Why LeCun argues reconstruction is the wrong objective9, how joint-embedding predictive architectures avoid representation collapse10, and what the evidence actually shows when you stop asking a model to draw every pixel.

After that: generative interactive worlds (Genie / Sora / GameNGen) in depth, then evaluation — which, as the reading & resources hub argues, is where the field is currently rethinking itself hardest.

Comments are open below — corrections and pointers to papers I should cover are very welcome.


How to cite this post

If this is useful for your own work, please cite the primary sources in the References. To reference this post itself:

@misc{haque2026worldmodels1,
  author       = {Md Rezwanul Haque},
  title        = {World Models --- Part 1: Inside the Latent (VAEs, RSSM, and Learning in a Dream)},
  year         = {2026},
  howpublished = {\url{https://rezwan.xyz/blog/2026/world-models-latent-dynamics/}}
}

References

  1. D. P. Kingma and M. Welling. Auto-Encoding Variational Bayes. ICLR, 2014. arXiv:1312.6114
  2. D. J. Rezende, S. Mohamed, and D. Wierstra. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. ICML, 2014. arXiv:1401.4082
  3. D. Ha and J. Schmidhuber. World Models / Recurrent World Models Facilitate Policy Evolution. NeurIPS, 2018. arXiv:1803.10122
  4. D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning Latent Dynamics for Planning from Pixels (PlaNet). ICML, 2019. arXiv:1811.04551
  5. D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi. Dream to Control: Learning Behaviors by Latent Imagination. ICLR, 2020. arXiv:1912.01603
  6. D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mastering Atari with Discrete World Models (DreamerV2). ICLR, 2021. arXiv:2010.02193
  7. D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap. Mastering Diverse Domains through World Models (DreamerV3). 2023. arXiv:2301.04104
  8. J. Schulman, P. Moritz, S. Levine, M. Jordan, and P. Abbeel. High-Dimensional Continuous Control Using Generalized Advantage Estimation. ICLR, 2016. arXiv:1506.02438
  9. Y. LeCun. A Path Towards Autonomous Machine Intelligence. OpenReview, 2022. openreview.net
  10. M. Assran, Q. Duval, I. Misra, P. Bojanowski, P. Vincent, M. Rabbat, Y. LeCun, and N. Ballas. Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture (I-JEPA). CVPR, 2023. arXiv:2301.08243
  11. J. Bruce, M. Dennis, A. Edwards, J. Parker-Holder, et al. Genie: Generative Interactive Environments. ICML, 2024. arXiv:2402.15391
  12. T. Brooks, B. Peebles, et al. Video Generation Models as World Simulators (Sora). OpenAI technical report, 2024. openai.com
  13. D. Valevski, Y. Leviathan, M. Arar, and S. Fruchter. Diffusion Models Are Real-Time Game Engines (GameNGen). 2024. arXiv:2408.14837
  14. A. Hu, L. Russell, H. Yeo, Z. Murez, G. Fedoseev, A. Kendall, J. Shotton, and G. Corrado. GAIA-1: A Generative World Model for Autonomous Driving. Wayve, 2023. arXiv:2309.17080
  15. R. S. Sutton and A. G. Barto. Reinforcement Learning: An Introduction (2nd ed.). MIT Press, 2018.