World Models

Why This Matters

Traditional model-free reinforcement learning (RL) agents learn by trial and error directly in their environments. While successful, these methods suffer from the credit assignment problem: when an agent receives a reward, it is difficult to determine which of its millions of parameters and past actions were responsible. Consequently, model-free RL algorithms are often bottlenecked, forcing researchers to use highly compact neural networks with few parameters to ensure training converges.

Yet, biological brains do not operate this way. When a baseball batter swings at a 100 mph fastball, the visual signal takes longer to reach their conscious brain than the actual duration of the pitch. They succeed because their brain continuously runs a predictive mental model of the environment, allowing them to reflexively act based on internal predictions of the future.

By separating the task of understanding the world (unsupervised representation learning) from the task of acting in the world (policy optimization), we can train massive, highly expressive models using standard backpropagation, while keeping the decision-making controller extremely simple.

The Big Idea

The core contribution of this paper is a modular, agentic architecture composed of three tightly integrated components:

Because V and M are trained in an unsupervised manner using backpropagation, they can scale to millions of parameters. The Controller (C), containing fewer than 1,000 parameters, can then be easily optimized using black-box algorithms like CMA-ES (Covariance Matrix Adaptation Evolution Strategy).

How It Works

Let's explore the execution loop. At each time step $t$, the raw environment frame is compressed by $V$ to produce $z_t$. This latent state, combined with the hidden recurrent state $h_t$ of $M$, is passed to $C$, which outputs action $a_t$. Finally, $M$ updates its internal state to $h_{t+1}$ using $z_t$ and $a_t$ to prepare for the next step.

Vision (V): Variational Autoencoder

The VAE acts as a spatial compressor. Raw observations from environments (like $64\times64\times3$ pixel screens) contain highly redundant spatial information. By training a convolutional autoencoder with a bottleneck layer, we compress each frame into a small latent vector $z_t \sim \mathcal{N}(\mu, \sigma^2 I)$.

Crucially, enforcing a Gaussian prior limits the information capacity of $z_t$, preventing overfitting and making the latent space robust to reconstructions generated by the memory model.

Memory (M): Mixture Density Network RNN

While $V$ compresses what the agent sees at a single instant, $M$ compresses what happens over time. Since complex environments are highly stochastic, predicting a single deterministic future frame is bound to fail.

Instead, $M$ is trained to output a probability density function $P(z_{t+1} \mid a_t, z_t, h_t)$ modeled as a Mixture of Gaussians (MDN). This allows the network to represent multi-modal outcomes—such as a fireball moving either left or right after being shot.

The state update is formulated as: $$h_{t+1} = \text{LSTM}(h_t, [z_t, a_t])$$ where the output layer models the parameters (means $\mu$, variances $\sigma$, and mixing coefficients $\pi$) of $K$ separate Gaussian distributions.

Controller (C): Action Selection

The Controller is designed to be as minimal as possible. Because the complex task of spatial representation and temporal forecasting is fully handled by $V$ and $M$, $C$ only needs to map the current latent state $z_t$ and memory state $h_t$ to an action $a_t$ via a simple linear projection: $$a_t = W_c [z_t, h_t] + b_c$$ This direct mapping contains a mere 867 parameters (in the Car Racing task), allowing it to be optimized efficiently using derivative-free algorithms like CMA-ES.

Results & Performance

The decoupled architecture achieves state-of-the-art results on challenging environments, being the first known model to solve the CarRacing-v0 task using raw pixel inputs.

Limitations & Open Questions

While incredibly powerful, training agents inside their own hallucinated dreams introduces unique challenges:

• Cheating the World Model: The agent can find "exploits" or adversarial policies that fool the world model. For example, in a low-temperature dream, an agent might discover a specific movement pattern that prevents the world model from generating fireballs, effectively "extinguishing" them.
• Task-Irrelevant Details: Because the VAE is trained in an unsupervised manner, it allocates capacity to compress all visual features, even those irrelevant to the task (like background wall textures), while potentially missing tiny, critical details (like distant fireballs).
• Capacity Limits & Forgetting: Recurrent models like LSTMs have finite storage capacity and suffer from catastrophic forgetting, limiting their ability to model highly complex, long-horizon environments.

Glossary

Citation

@article{ha2018worldmodels,
  title={World Models},
  author={Ha, David and Schmidhuber, J{\"u}rgen},
  journal={arXiv preprint arXiv:1803.10122},
  year={2018}
}