World Models

An agent learns a model of its game, then trains inside that model.

Ha and Schmidhuber split an agent into a large world model, trained without any reward to compress each frame and predict the next one, and a controller of about a thousand parameters that is the only part to see the reward. In Doom, the controller trains only inside the learned model and then plays the real game.

On CarRacing-v0 the full agent averages 906 ± 21 over 100 trials, the first reported score above the 900 that counts as solving the task. The Doom controller transfers to the real game when its learned model is sampled with more randomness than it learned, and falls below a random policy when sampled with less.

A batter has a few milliseconds to decide how to swing at a 100 mph fastball, less time than it takes the visual signal to reach the brain. The paper opens with this example: the swing goes where the batter's internal model predicts the ball will be, not where the eye last saw it. World Models, by David Ha and Jürgen Schmidhuber, builds an artificial agent on that plan.

The agent has three networks. V (vision) compresses each 64×64 frame into a short vector $z$. M (memory) is a recurrent network that predicts a probability distribution for the next $z$. C (controller) reads what V and M produce and picks the action. V and M together are the world model. In the car-racing agent they hold about 4.8 million parameters and learn with no reward at all; C holds 867 parameters and is the only part trained on reward.

Because M predicts the next $z$ from the current one and the action, it can stand in for the game: give it an action, sample a next $z$, repeat. The paper trains a Doom controller entirely on M's samples, which the authors call training inside a dream, then runs that controller in the real game, where it survives $1092 \pm 556$ steps against a bar of 750. The page follows the paper's order: V, M and its sampling temperature, C, how C is trained, the car-racing result, the Doom experiment, and how a controller can exploit a flawed model.

A large world model and a small controller

Reinforcement learning has a credit assignment problem. An agent takes hundreds of actions and gets back rewards that may arrive long after the action that earned them, so it has to work out which actions to credit. Backpropagation trains large networks efficiently, but only given a differentiable loss, and a game's env.step is a black box with no gradient. The paper's footnote puts typical model-free RL agents of the time at $10^3$ to $10^6$ parameters. This agent has about $10^7$, arranged so that almost none of them learn from reward.

The split gives each network an objective suited to it:

• V is trained to reconstruct frames. Every pixel of every frame contributes to the loss, and the gradient flows by ordinary backpropagation.
• M is trained to predict the next compressed frame. Every time step contributes to the loss, again by backpropagation.
• C is trained on the episode's total reward, with a gradient-free search over its 867 parameters.

Neither V nor M needs a reward, so both train on episodes played by a policy that picks random actions. C's search space is then 867 numbers, small enough for an evolution strategy that needs only each episode's final score. In the paper's words, a small controller "lets the training algorithm focus on the credit assignment problem on a small search space, while not sacrificing capacity and expressiveness via the larger world model." The car-racing procedure (Section 3.2) has five steps:

1. Collect 10,000 rollouts from a random policy.
2. Train V to encode each frame into $z \in \mathbb{R}^{32}$.
3. Train M to model $P(z_{t+1} \mid a_t, z_t, h_t)$.
4. Define C as the linear map in equation (1) below.
5. Use CMA-ES to find C's weights and biases that maximize the expected total reward.

The three networks are trained separately, in that order. The paper notes they could be trained end to end, but separate training was more practical, and each of V and M took under an hour on a single GPU.

V: a frame compressed to 32 numbers

Each observation is resized to 64×64 RGB with every value in [0, 1], so a frame is 12,288 numbers. V is a convolutional variational autoencoder, or VAE (the site's VAE explainer has the full derivation). The encoder maps a frame to two vectors of length $N_z$, a mean $\mu$ and a log-variance $\log \sigma^2$, and draws the code $z$ as $\mu + \sigma \odot \varepsilon$, with $\varepsilon$ drawn from a standard normal. $N_z$ is 32 for car racing and 64 for Doom. The decoder maps $z$ back to a 64×64×3 image $\hat{\mathbf{x}}$. The layer shapes, from the authors' released code:

# V, the ConvVAE, car-racing sizes (N_z = 32). Every conv has stride 2.
x = frame                               # (64, 64, 3), RGB values in [0, 1]
conv   32 ch, 4x4, relu     -> (31, 31, 32)
conv   64 ch, 4x4, relu     -> (14, 14, 64)
conv  128 ch, 4x4, relu     -> (6, 6, 128)
conv  256 ch, 4x4, relu     -> (2, 2, 256)       # flatten: 1024 numbers
mu     = dense(1024 -> 32)
logvar = dense(1024 -> 32)
z      = mu + exp(logvar / 2) * randn(32)       # the code for this frame
dense(32 -> 1024), reshape  -> (1, 1, 1024)
deconv 128 ch, 5x5, relu    -> (5, 5, 128)
deconv  64 ch, 5x5, relu    -> (13, 13, 64)
deconv  32 ch, 6x6, relu    -> (30, 30, 32)
deconv   3 ch, 6x6, sigmoid -> (64, 64, 3)      # the reconstruction x_hat

Counting weights through those layers gives 4,348,547, the paper's exact figure for the car-racing V. Of those, 3,276,928 (75%) sit in the first deconvolution, which expands the 1024-number vector into a 5×5×128 map. The frame's 12,288 numbers pass through a 32-number code, a factor of 384 fewer.

V minimizes the sum of two terms:

The first term is the squared error summed over all 12,288 values. The second measures how far the encoder's Gaussian for this frame, $q(z \mid \mathbf{x})$, is from a standard normal. For a Gaussian with diagonal covariance it has a closed form:

It is zero only when every $\mu_j = 0$ and every $\sigma_j = 1$, a code that says nothing about the frame. The two terms pull against each other: lowering the reconstruction error takes precise, distinct codes, and lowering the KL term pushes every code toward standard noise around the origin. In the released code the KL term is floored at 0.5 nats times $N_z$, 16 nats in total for car racing: while a frame's KL is below that, the term contributes no gradient, so the encoder can keep up to 16 nats of information per frame without penalty.

Why a VAE and not a plain autoencoder, which has no KL term? The paper gives two reasons in one sentence of its appendix. The Gaussian prior limits how much information each code can carry, so V spends its capacity on large structure. And the prior makes the world model more robust to unrealistic $z$ vectors generated by M. The second reason comes into play once M runs as a simulator and produces codes that V never produced from a real frame: the KL term trained the decoder on codes spread across the region around the origin, with no gaps between them, so an invented code that lands in that region is more likely to decode to something frame-like.

V trains for one epoch over the frames of the 10,000 random rollouts and is then frozen. After that, the rest of the agent never sees pixels; it sees $z$. Figure 1 is a schematic of what a code this small keeps.

The reconstruction loss weights every pixel equally and carries no information about the task, so V keeps what is large and easy to rebuild. The paper reports the consequence in its discussion: the Doom V reproduced the brick pattern on the side walls, which does not matter for survival, while the car-racing V failed to reproduce task-relevant tiles on the road.

M: a distribution over the next code

V describes one frame; M describes how frames follow each other. At each step M reads the current code $z_t$ and the action $a_t$, updates a recurrent state, and outputs a probability distribution for the next code $z_{t+1}$. M is an LSTM (long short-term memory, a recurrent cell whose gates decide what to keep from step to step) with 256 hidden units for car racing and 512 for Doom. Its hidden state $h_t$ is a vector summarizing the codes and actions so far. M is trained to predict $z_{t+1}$ from it, so $h_t$ ends up holding the information that predicts the next frame, which is why C will read it.

Why a distribution and not a single predicted $z_{t+1}$? In the Doom task a monster may or may not launch a fireball on the next step. Say both outcomes have probability 0.5. A predictor trained with squared error on the next frame does best by outputting the average of the two outcomes, because squared error is smallest at the mean of the possible targets. That average is a fireball at half opacity, a frame the game never renders. A controller trained against it would practice for a half-fireball, while in the real game every fireball is either fully there or absent. A mixture keeps the two outcomes as separate modes with weight 0.5 each, and every sample is one of the two real frames. Figure 2 shows both predictors for one such event.

M's output layer is a mixture density network (Bishop, 1994). For each coordinate $i$ of the next code it outputs five weights $\pi_{ik}$, five means $\mu_{ik}$ and five standard deviations $\sigma_{ik}$. The paper does not model correlations between coordinates, so the density of the whole next code is a product over coordinates:

where each coordinate's five weights sum to one. The appendix calls this a factored mixture of Gaussians. In the released code the output layer has $N_z \times 5 \times 3$ units (for car racing, 32 × 5 × 3 = 480: a weight logit, a mean and a log standard deviation for each component of each coordinate), and sampling picks a component separately for every coordinate. An LSTM with a mixture density output is called an MDN-RNN. Ha had used the same combination in SketchRNN to draw sketches one pen stroke at a time.

The parameter count follows from the shapes. For car racing the LSTM reads 35 numbers per step (32 for $z_t$, 3 for the action), so its four gates hold 4 × 256 × (35 + 256 + 1) = 299,008 weights and biases. The output layer from 256 to 480 adds 256 × 480 + 480 = 123,360. The total, 422,368, is the paper's figure.

M trains for 20 epochs on the same random rollouts, with each frame converted to a code by the frozen V. The loss at each step is the negative log of the density above, evaluated at the code that actually came next:

Minimizing it raises the probability M assigns to what really happened, and a mixture can raise it for two different outcomes at once by putting a component on each. The released code averages the loss over coordinates and time steps. The appendix adds a guard against overfitting: V stores a $\mu$ and $\sigma$ for every recorded frame, and each time a training batch is built, a fresh $z$ is sampled from them, so M never fits one particular sample of a frame's code. One step of M, in the code's terms:

# M, one step, car-racing sizes. LSTM state (h, c): 256 numbers each.
h, c = lstm(concat(z_t, a_t), (h, c))          # input: 32 + 3 = 35 numbers
out  = dense(h)                                # 256 -> 480 = 32 x 5 x 3
logit, mean, logstd = split(out)               # each (32 coords, 5 comps)

# training: negative log-likelihood of the code that actually came next
log_pi = log_softmax(logit)
loss = -mean(logsumexp(log_pi + log_normal(z_next, mean, logstd)))

# sampling at temperature tau, each coordinate on its own
for i in range(32):
    k = choice(5, p=softmax(logit[i] / tau))   # pick a component
    z_next[i] = mean[i, k] + exp(logstd[i, k]) * sqrt(tau) * randn()

Sampling temperature

Sampling from M has one setting, a temperature $\tau$, which the paper takes from SketchRNN. It changes two things. Each component's variance is multiplied by $\tau$, so its standard deviation by $\sqrt{\tau}$. And the weight logits $\ell_{ik}$ are divided by $\tau$ before the softmax that turns them into weights:

At $\tau = 1$ samples come from the distribution M learned. Below 1, every component narrows and the weight moves toward the component with the largest logit; as $\tau \to 0$, M becomes deterministic. Above 1, the components widen and the rare ones gain weight.

Rare events disappear through the weight change. Take the mixture in Figure 3, where the most likely component has logit 2.4 and a rare "fireball" component has logit 0.4. At $\tau = 1$ the fireball weight is 7.4%. At $\tau = 0.5$ the gap between the two logits doubles from 2 to 4 and the weight falls to 1.5%. At $\tau = 0.1$ the gap is 20, and the weight is about $2 \times 10^{-9}$: the model essentially never samples a fireball. If temperature scaled only the variance, the fireball would keep its 7.4% at every $\tau$, only narrower.

The paper's Doom experiment runs into exactly this at $\tau = 0.1$: the monsters in the learned model stop firing, which the section on cheating below takes up.

C: one linear layer

C reads the current code $z_t$ and M's hidden state $h_t$, concatenated into one vector, and applies a single linear layer. This is the paper's only numbered equation:

For car racing, $[z_t\; h_t]$ has 32 + 256 = 288 numbers and the action has three (steering, gas, brake), so $W_c$ is 3 × 288 and $b_c$ has 3 entries: 288 × 3 + 3 = 867 parameters. The outputs pass through tanh to keep them in range. Steering stays in [−1, 1], gas is mapped to [0, 1], and brake is clipped to [0, 1]:

def action(z, h):                    # car-racing C: W is (3, 288), b is (3,)
    a = tanh(W @ concat(z, h) + b)   # 288 numbers in, 3 out
    steer = a[0]                     # in [-1, 1]
    gas   = (a[1] + 1) / 2           # in [0, 1]
    brake = clip(a[2], 0, 1)         # in [0, 1]
    return steer, gas, brake

The Doom controller differs in two ways. Take Cover's actions are discrete (move left, stay, move right), so C outputs one number in [−1, 1] and the range is divided into thirds. And C reads both of the LSTM's internal vectors, the hidden state $h$ and the cell state $c$, 512 numbers each. The paper's Doom procedure writes the controller without a bias, so the count is 64 + 512 + 512 = 1,088 weights, the paper's figure.

Figure 4 draws the loop. Each step, the game produces a frame, V encodes it, and C reads $z$ with M's state and acts. The action goes both to the game and to M, which advances its state for the next step. Compare the stacked bar under the loop with the parameter counts above: C is 0.018% of the car-racing agent.

The paper's pseudocode for one episode in the real environment:

def rollout(controller):              # the paper's pseudocode (Sec 2.4)
    obs = env.reset()
    h = rnn.initial_state()
    done, cumulative_reward = False, 0
    while not done:
        z = vae.encode(obs)           # V: 64x64x3 frame -> z
        a = controller.action([z, h]) # C: a = W_c [z h] + b_c
        obs, reward, done = env.step(a)
        cumulative_reward += reward
        h = rnn.forward([a, z, h])    # M: advance the recurrent state
    return cumulative_reward

Training C with CMA-ES

With under a thousand parameters, C can be trained without gradients. The paper uses CMA-ES (covariance matrix adaptation evolution strategy), which keeps a Gaussian distribution over parameter vectors: a mean $m$, an overall step size $\sigma$, and a covariance matrix $\Sigma$ that sets the shape of the search. Each generation:

1. Draw a population of candidate parameter vectors from $\mathcal{N}(m, \sigma^2 \Sigma)$. The paper uses 64.
2. Score every candidate. Here a score is the average total reward of 16 full episodes with different random seeds.
3. Rank the candidates and move the mean to a weighted average of the better half.
4. Stretch $\Sigma$ toward the directions in which the better candidates lay, and grow or shrink $\sigma$ depending on whether recent steps of the mean have kept going the same way.

Only step 2 touches the environment, and it uses nothing but each episode's total reward, so the reward can be sparse and env.step can stay a black box. The 64 candidates are independent, and the paper runs their episodes in parallel on the CPU cores of one machine. The paper chose CMA-ES because it works well up to a few thousand parameters; the full covariance for 867 parameters has 867² = 751,689 entries, and for a network the size of V it would have about $1.9 \times 10^{13}$.

Averaging 16 episodes per candidate is a response to noise. Car-racing tracks are generated at random, so one controller scores differently from one episode to the next, and ranking candidates by a single episode each would often put a lucky candidate above a better one. Averaging 16 episodes divides the standard deviation of that noise by 4, the square root of 16. The training loop, in outline:

es = CMA(n_params=867, popsize=64)
for generation in range(1800):
    candidates = es.ask()                      # 64 parameter vectors for C
    fitness = [mean(rollout(c, seed=s) for s in range(16))
               for c in candidates]            # 16 episodes each
    es.tell(candidates, fitness)               # move the mean, adapt covariance

Figure 5 runs the same algorithm on a controller with two weights and a known fitness ridge, small enough to draw the covariance. Watch the dashed ellipse stretch along the ridge over the first few generations and then shrink, and compare where the mean ends up with 1 and with 16 rollouts per candidate.

For car racing, C was trained in the real game. Every 25 generations the paper tested the current best candidate over 1024 episodes, and after 1800 generations that candidate averaged 900.46. 1800 generations of 64 candidates at 16 episodes each is 1,843,200 training episodes, each one a run of the 867-parameter controller through the game, V and M.

Car racing results

CarRacing-v0 is a top-down driving game from OpenAI Gym with a new randomly generated track every episode. The agent is rewarded for visiting as many track tiles as possible in the least time. The task counts as solved at an average reward of 900 over 100 consecutive trials, which, as the paper notes, leaves room for very few driving mistakes.

The paper tested three controllers on the same V and M:

• C reads only $z_t$ through the linear map $W_c z_t + b_c$, 32 × 3 + 3 = 99 parameters. It scores $632 \pm 251$. The car wobbles and misses the track on sharp corners.
• C reads only $z_t$, through a hidden layer of 40 tanh units: 1,443 parameters, more than the full controller. It scores $788 \pm 141$.
• C reads $z_t$ and $h_t$, the 867-parameter linear controller of (1). It scores $906 \pm 21$.

The hidden-layer controller has more parameters than the full one and still scores lower, so the gain comes from the input $h_t$, not from the size of C. The paper's explanation: $z_t$ "only captures a representation at a moment in time and does not have much predictive power," while $h_t$, the state M uses to predict $z_{t+1}$, gives C "what to expect in the future." The paper compares the full agent to the batter from the introduction: it does not plan or roll out futures before acting; it reads M's state and reacts.

The ± figures are standard deviations over the 100 trials. $906 \pm 21$ means the agent's score rarely strayed far from 900 on any track; $632 \pm 251$ means the z-only agent's score varied widely from track to track. The prior methods in the paper's Table 1: DQN scored $343 \pm 18$; A3C scored $591 \pm 45$ with continuous actions and $652 \pm 10$ with discrete ones; the best entry on the Gym leaderboard scored $838 \pm 11$. $906 \pm 21$ was the first reported score above 900. The earlier deep-RL entries also relied on hand-built preprocessing such as edge detection and stacks of recent frames, while this agent reads raw RGB frames.

Training C inside M

M maps a code, an action and a recurrent state to a distribution over the next code, the same interface as an environment (state and action in, next state out). The paper uses it in its second experiment, VizDoom's Take Cover scenario. Monsters on the far side of a room shoot fireballs at the agent, which can move left or right. There is no explicit reward; the score is the number of steps the agent stays alive, up to 2100 steps (about 60 seconds). The task counts as solved when the average survival over 100 consecutive episodes exceeds 750 steps (about 20 seconds).

To be a complete environment, M also has to end episodes, so the Doom M predicts whether the agent dies on the next step: its output layer gets one extra unit, a logit for that "done" event, and the episode ends when the predicted probability exceeds 50%. The paper chose this cutoff over sampling the event at random because death is rare at any single step and the cutoff was more stable. With $N_z = 64$ and 512 LSTM units, M's input per step is 66 numbers (the 64-number latent code, the one action, and a flag the released code uses to mark episode restarts), and its weights come to 4 × 512 × (66 + 512 + 1) + 512 × 961 + 961 = 1,678,785, the paper's Doom figure.

The authors wrapped M in the same Gym interface as the real game and trained C against it. A dream episode starts from the code of a real first frame (in the released code, sampled from stored encodings of recorded starting frames) and from there runs with no game and no V:

def dream_rollout(controller, tau):    # Doom, entirely inside M
    z = sample_initial_z()             # code of a recorded real first frame
    h = rnn.initial_state()
    done, steps = False, 0
    while not done and steps < 2100:
        a = controller.action([z, h])  # same C as in the real game
        z, done = rnn.sample(a, z, h, tau)  # M samples z_next and done
        h = rnn.forward([a, z, h])
        steps += 1                     # Take Cover: +1 per step alive
    return steps

The loop has the shape of the real rollout, with the next code sampled from M instead of produced by env.step and V, and a reward of 1 for each step survived. A step of the dream is one LSTM step and a sample, with no game engine and no rendering, so M has to reproduce what the engine did: move the agent when it steps left, stop it at the walls, move several fireballs at once, and detect a hit.

The paper states that it never trained the agent in the real VizDoom; the real game only supplied the 10,000 random-policy episodes. The controller trained inside M at $\tau = 1.15$ and then run in the real game survived $1092 \pm 556$ steps over 100 episodes. That clears the 750 bar and beats the best Gym leaderboard entry at the time, $820 \pm 58$. The same controller scored $918 \pm 546$ inside M: it survived longer in the real game than in the model it trained in. The paper's explanation is that M at this temperature is a harder game than the real one (it does not even keep the number of monsters right, and fireballs move less predictably), and an agent that survives the noisier version does better in the cleaner original.

Cheating the world model

A controller trained inside M maximizes the score M gives it, and M is wrong in places. In early Doom runs the paper saw C find an adversarial policy: it moved in a way that made the monsters in M never fire a single fireball during some rollouts, and when a fireball began to form, the agent moved so that it disappeared. The real game has no such behavior, so the policy was useless there.

The paper gives two reasons the search finds such policies. M learned from random play, and a policy that steers into unusual situations visits states far from that data, where M's predictions are least reliable. And in the dream C reads M's hidden state, which here plays the role of the game engine's internal memory, the part a human player never sees, so C can learn moves that push that state wherever the score is high. CMA-ES moves its search toward whichever candidates score higher, so if some sequence of moves drives M into a state where fireballs stop, the search converges on controllers that produce it.

Temperature controls how exploitable M is. Raising $\tau$ widens every sampled code and gives rare components more weight, so fireballs appear and move less predictably, and a move sequence that suppressed them in one rollout does not reliably do so in the next. The paper trained C inside M at five temperatures and measured each controller in M (virtual) and in the real game (actual):

At $\tau = 0.1$ the controller scores 2086 inside M (the paper says it reaches the perfect 2100 most of the time) and 193 in the real game, below the 210 of a random policy. At this temperature the fireball component's weight is near zero, as in the temperature section's example, so M cannot switch into the mixture component where fireballs form, so its monsters never shoot whatever the agent does, and the controller never learns to dodge. At $\tau = 0.5$ the scores barely move (2060 inside M, 196 in the real game). At $\tau = 1.0$ the virtual score drops to 1145 and the actual score rises to 868. At $\tau = 1.15$ the virtual score is 918 and the actual score 1092, the best transfer. At $\tau = 1.30$ the actual score falls to 753 with a standard deviation of 139 instead of 556; the paper describes this as a lower-scoring but less risky strategy.

The paper calls this a tradeoff between realism and exploitability and treats $\tau$ as a hyperparameter to tune: a higher temperature makes adversarial policies harder to find, and too high a temperature leaves a game too random to learn from. The setting exists because M outputs a mixture density, so the same trained weights can be sampled at any $\tau$. The earlier recurrent world models the paper builds on, Schmidhuber's from 1990 and 1991, were deterministic: one predicted future per state and nothing to turn up, so a controller could exploit them wherever they were wrong. The paper mentions Bayesian models such as PILCO as a partial fix.

Limits and proposed next steps

The paper lists its limits in its discussion:

• V learns without the task, so its code can spend capacity on irrelevant detail (the Doom wall bricks) and miss relevant detail (the car-racing road tiles). Training V together with an M that predicts rewards could focus it on what matters for the task, at the cost of a V that is harder to reuse for other tasks.
• The LSTM has limited capacity. Given more data it may not store everything, and networks trained by backpropagation suffer from catastrophic forgetting: learning new predictions degrades old ones. The paper suggests higher-capacity models or an external memory.
• M simulates the future one step at a time, with no hierarchical planning or abstraction over irrelevant detail.

It also proposes, without results, an iterative procedure for tasks where a random policy never reaches the parts of the world that matter:

1. Initialize M and C with random parameters.
2. Run the agent in the real environment N times, storing every action and observation.
3. Train M on the stored data to predict the next observation, reward, action and done flag, and train C inside M.
4. Go back to step 2 if the task is not yet solved.

One pass of this loop was enough for car racing and Take Cover. For harder tasks the paper suggests pushing C toward what M predicts badly, by flipping the sign of M's loss in the real environment so that high prediction error counts as reward. It ties this to Schmidhuber's work on artificial curiosity, which rewards improvement in compression. High error and improving prediction are different quantities: a source of pure noise keeps prediction error high forever while offering nothing to learn, so the sign flip is a loose version of that idea. The site's Curiosity-Critic explainer covers a method built on prediction-error improvement. The discussion also notes that, because V, M and C are all differentiable, C could in principle be trained by backpropagation through M instead of by evolution.

The car-racing agent in numbers: 10,000 random episodes; a 4,348,547-parameter V trained for one epoch; a 422,368-parameter M trained for 20 epochs; and an 867-parameter C evolved for 1800 generations of 64 candidates at 16 episodes each, scoring $906 \pm 21$ where 900 was the bar. The Doom agent: a 1,088-parameter C that trained only inside M at $\tau = 1.15$ and survived $1092 \pm 556$ steps in the real game, and that fell to 193 steps, below a random policy, when trained inside M at $\tau = 0.1$.