World Models

The core insight of World Models is to decompose an RL agent into three specialized modules, each solving a different subproblem. This decomposition mirrors how cognitive scientists describe human perception: we compress sensory input, maintain a predictive model of the world, and make decisions using both.

Vision (V) — a Variational Autoencoder — compresses each raw observation (a 64×64 RGB image) into a compact latent vector $z_t$ of just 32 dimensions. This is like reducing a photograph to a quick sketch that captures the essential layout while discarding pixel-level noise.

Memory (M) — an MDN-RNN (a recurrent network with a Mixture Density Network output) — receives the current latent $z_t$ and action $a_t$, then predicts the distribution of the next latent $z_{t+1}$. It maintains a hidden state $h_t$ that accumulates temporal context. Think of it as an internal crystal ball that says "given what I've seen and what I just did, here is what might happen next."

Controller (C) — a single linear layer — takes the concatenation of $z_t$ (what I see now) and $h_t$ (what I remember and predict) and outputs the action $a_t$. With only 867 parameters, it is deliberately tiny so that the world model carries the complexity.

The Vision model is a convolutional Variational Autoencoder. The encoder takes a 64×64 RGB frame and maps it to the parameters of a Gaussian distribution in a 32-dimensional latent space: a mean $\mu$ and a log-variance $\log \sigma^2$. A sample $z_t$ is drawn from this distribution using the reparameterization trick. The decoder reconstructs the frame from $z_t$.

Why a VAE and not a plain autoencoder? Because the VAE learns a smooth, continuous latent space. Nearby points in this space correspond to visually similar scenes. This smoothness is critical for the Memory model — it needs to predict trajectories through latent space, and a jumpy, fragmented space would make prediction nearly impossible. The VAE's KL divergence penalty acts like a gentle force that prevents the latent space from collapsing into disconnected islands.

The result: a 32-number "sketch" that captures the essential structure of each frame — the position of the car, the shape of the road, the presence of obstacles — while discarding irrelevant pixel details.

With a compact representation of the present ($z_t$), the next challenge is predicting what happens next. The Memory model is an LSTM whose output is not a single prediction but a probability distribution over possible futures.

Specifically, the MDN-RNN outputs the parameters of a Mixture of Gaussians: for each of the 32 latent dimensions, it produces $K$ Gaussian components (means $\mu_k$, standard deviations $\sigma_k$, and mixing weights $\pi_k$). This means the model can express uncertainty and multimodality — if the car is approaching a fork, the model can assign probability to both branches rather than blurring between them.

At each step, the LSTM receives the concatenation of $z_t$ and $a_t$ and updates its hidden state $h_t$. The hidden state accumulates temporal context: it remembers the trajectory so far, the velocity, the rhythm of the track. The MDN head then maps $h_t$ to the mixture parameters.

The Controller is deliberately minimal: a single linear layer that maps the concatenated input $[z_t, h_t]$ to an action vector.

Why is the Controller so small? Because all the complexity lives in the world model. The VAE has already extracted the meaningful features. The LSTM has already compressed temporal patterns. All the Controller needs to do is draw a simple linear boundary in this well-organized space. Think of it as a CEO who receives a one-page executive summary from two brilliant analysts — the decision becomes straightforward.

This separation has a practical benefit: since the Controller has so few parameters, it can be trained with evolutionary strategies like CMA-ES instead of gradient descent. Evolution is gradient-free — it does not require backpropagation through the world model, so V, M, and C can be trained independently.

This is the paper's most striking contribution. Once the world model (V + M) is trained on observations from the real environment, it can generate synthetic experience — the agent can explore a hallucinated version of the environment entirely inside its own Memory model.

The process works like this: at each dream step, the Controller takes an action based on the current latent state and hidden state. The MDN-RNN then samples the next latent state from its predicted distribution. The Controller reacts to this hallucinated state, takes another action, and the cycle continues — all without touching the real environment.

A crucial detail: the temperature parameter $\tau$ controls the stochasticity of the dream. At $\tau=1$, dreams match the real environment's randomness. At $\tau > 1$, dreams become more chaotic and unpredictable — the authors found this acts as a form of regularization, making the Controller more robust because it trains against worst-case scenarios.

The entire system trains in three independent stages, each taking less than an hour on a single GPU.

Stage 1 — Collect rollouts. A random policy interacts with the environment, recording observations and actions. No reward signal is needed for this step — the world model is trained purely on sensory experience.

Stage 2 — Train V and M. The VAE learns to compress frames, then the MDN-RNN learns temporal dynamics in the latent space. Both are trained with unsupervised objectives (reconstruction + KL for the VAE, negative log-likelihood for the MDN-RNN).

Stage 3 — Train C. The Controller is evolved using CMA-ES, either by interacting with the real environment (V + M + C mode) or entirely inside dreams (Dream mode). Each candidate controller is evaluated across multiple rollouts, and CMA-ES updates the parameter distribution toward higher rewards.

The paper evaluates on two environments. CarRacing-v0 is a continuous-control driving task with procedurally generated tracks. The agent sees a top-down view and must visit as many tiles as possible. The full V+M+C model scored 906 (out of ~950 maximum), significantly outperforming the previous best methods. The Vision-only baseline (V+C, no memory) scored 632 — proving that memory and prediction are crucial for handling curves and planning ahead.

VizDoom Take Cover is a first-person survival game where the agent must dodge fireballs thrown by monsters. Here, the paper demonstrated dream training: the controller was trained entirely inside the world model's imagination, then transferred to the real environment, achieving over 1100 timesteps of survival — competitive with agents trained with direct environment access.

World Models introduced a blueprint that the entire field followed. The core recipe — learn a latent representation, predict dynamics in that space, plan using imagined trajectories — appears in virtually every modern model-based RL system.

The paper also introduced a profound conceptual shift: the agent does not directly perceive reality. It only sees what its world model shows it. This "Plato's Cave" perspective highlights both the power and the danger: a good world model enables efficient learning, but a flawed model can create systematic blind spots that the agent cannot detect.

The limitations were clear too. The MDN-RNN struggled with long horizons. The VAE discarded details that sometimes mattered. The CMA-ES controller did not scale to high-dimensional action spaces. But these were engineering limitations, not conceptual ones — and the follow-up work (PlaNet, Dreamer, MuZero) addressed each in turn.

The trajectory from World Models to Dreamer v3 is a story of engineering the same core idea: learn a latent model of the world, imagine rollouts, and optimize behavior in imagination. Each successor addressed a specific limitation — better state models, better policy learning, better scaling — but the conceptual foundation laid in 2018 remains the blueprint.