World Models

In plain words

Teaching a program to play a game from raw pixels is hard. Standard reinforcement learning works best with small networks, because it struggles to train large ones from reward alone. This paper splits the problem in two. A large model first learns, without any reward, to compress what the agent sees and to predict what will come next. That model is the agent's "world model". A very small decision-maker then chooses actions using only the compressed view and the world model's memory.

The small decision-maker is trained with an evolutionary search: many random variants are tried, the best ones are kept and their settings are perturbed again. The authors report that this approach solves a car-racing task with a policy of fewer than a thousand parameters.

The second idea is that a good enough world model can stand in for the real game. The agent practises inside the model's predictions, which the authors call a dream, and the policy learned there is then run in the real game. In a shooting-game test this transfer worked, but the agent also found ways to exploit mistakes in the model. The authors counter this by making the dream more random, which makes it harder to cheat.

TL;DR

World Models is a three-part agent: a VAE compresses each frame, an MDN-RNN predicts the next compressed frame from the current action and its hidden state, and a single-layer linear controller picks actions from both §2. The controller is evolved with CMA-ES, which needs only a final reward per rollout §2.3 §A.4. Use it when you want a compact policy on pixel inputs, or to see how training inside a learned model can work, and budget for training data from roughly 10,000 random rollouts and about 1,800 evolution generations §3.2 §A.4.

Problem

Reinforcement learning with a large model is hard: the credit-assignment problem makes it difficult for traditional RL algorithms to learn millions of weights, so practitioners use small networks that iterate quickly §1. Learning a dynamics model directly from high-dimensional frames is also difficult, which motivates first compressing the frames into a low-dimensional representation §6. The paper's goal is a large unsupervised world model plus a small controller, so the complexity sits in the model while the policy stays small §1 §2.3.

The authors name these limits of their own approach:

• The VAE keeps detail that is not relevant to the task, because it is trained without rewards §7.
• The memory of the MDN-RNN has limited capacity, and catastrophic forgetting is a concern for the iterative variant §7.
• A controller can exploit a dream model that is not accurate, and a deterministic dynamics model is especially easy to exploit §4.5.

Method

Overview. The agent has three components. V (Vision) is a VAE that maps each frame to a latent vector z_t. M (Memory) is an LSTM with a mixture-density output that models the distribution of the next z given the action, z_t and its hidden state h_t. C (Controller) is a single linear layer that maps [z_t, h_t] to an action §2 fig. 4.

Steps.

1. Collect data. Run a random policy for 10,000 rollouts on the environment and record frames and actions §3.2 step 1 §A.1.
2. Train V. Train a convolutional VAE to encode each frame into z. Frames are resized to 64×64 RGB with values in [0, 1]. The loss is a squared-error reconstruction term plus a KL term. Car Racing uses a 32-dimensional z, Doom a 64-dimensional z §3.2 step 2 §4.2 §A.1.
3. Train M. Encode the recorded frames with V and train an MDN-RNN to model P(z_{t+1} | a_t, z_t, h_t) as a mixture of Gaussians with a diagonal covariance §2.2 §3.2 step 3 §A.2.
4. Define C. The controller is one linear map from [z_t, h_t] to the three Car Racing actions, with an added bias §2.3 eq. 1 §3.2 step 4.
5. Evolve C. Use CMA-ES on the controller's weights, with the fitness of a candidate equal to its average reward over several rollouts §3.2 step 5 §A.4.
6. Dream (Doom only). Train C inside the M model. Each dream step samples the next z from M, and M also predicts whether the agent dies in the next frame. Death is declared when that predicted probability exceeds 50% §4.1 §4.2 §A.2.
7. Transfer. Run the policy trained in the dream on the real environment, using V to encode real frames §4.4.

The rollout loop below restates the paper's listing of how the three parts interact at run time §2.3.

# One episode of the agent  [§2.3 listing]
# Input:  env, V (vae), M (rnn), C (controller)
# Output: total reward of the episode

def rollout(env, vae, rnn, controller):
    # M starts from its initial hidden state h_0  [§2.3]
    h = rnn.initial_state()
    obs = env.reset()
    total_reward = 0.0
    done = False
    while not done:
        # V compresses the frame into a latent z_t  [§2.1]
        z = vae.encode(obs)                           # (32,)
        # C maps [z_t, h_t] to the action a_t  [§2.3] [eq. 1]
        a = controller.action(concat(z, h))
        obs, reward, done = env.step(a)
        total_reward += reward
        # M advances its hidden state with (z_t, a_t)  [§2.2]
        h = rnn.forward(a, z, h)
    return total_reward

The controller equation is the only numbered equation in the paper:

a_t = W_c \, [z_t \; h_t] + b_c eq. 1

Sampling from M during the dream uses a temperature τ. The paper says τ controls the uncertainty of M's predictions, as in SketchRNN, but does not define the sampling procedure precisely §2.2 §4.3. The released code divides the mixture logits by τ and scales the Gaussian noise by √τ. The procedure below follows that code and is not stated in the paper.

# One dream step of M at temperature tau  [§2.2] [§4.3]
# Code path: carracing/rnn/rnn.py, doomrnn/doomrnn.py
# Input:  logits (D, K), means (D, K), log_stds (D, K), tau
# Output: next latent z_next (D,), with D = 32 (Car) or 64 (Doom), K = 5

def dream_step(logits, means, log_stds, tau):
    # Flatten or sharpen the mixture weights by tau  [code: rnn.py#L280-L285]
    weights = softmax(logits / tau)               # (D, K)
    # Choose one Gaussian per latent dimension  [code: rnn.py#L286-L293]
    k = sample_categorical(weights)               # (D,)
    mu = take(means, k)
    sigma = exp(take(log_stds, k))
    # Add noise scaled by sqrt(tau)  [code: rnn.py#L294-L295]
    noise = normal(0, 1, shape=D) * sqrt(tau)
    return mu + sigma * noise

The evolution loop for the controller is summarised below. The ask and tell calls mirror the code's CMA-ES wrapper train.py; the paper does not name them.

# Evolve the controller with CMA-ES  [§3.2 step 5] [§A.4]
# Code path: carracing/train.py, carracing/es.py#L86-L100

def train_controller(generations):
    # Population of 64 candidates per generation  [§A.4]
    # sigma_init = 0.1 is a code default, not stated in the paper
    cma = CMAES(num_params, sigma_init=0.1, population=64)
    for generation in range(generations):
        candidates = cma.ask()
        for each candidate:
            # Fitness: average reward over 16 seeded rollouts  [§A.4]
            fitness = mean(rollout(candidate, seed) for seed in 16_seeds)
        cma.tell(fitness)
        # Every 25 generations, re-test the best candidate  [§A.4]
        if generation % 25 == 0:
            evaluate(cma.best(), rollouts=1024)

The paper reports the Car Racing agent reaching the target score after 1,800 generations §A.4; the loop above does not fix a generation count.

Hyperparameters

Parameter counts from the paper's tables. Car Racing tab. 1:

VizDoom Take Cover tab. 3:

The 867 Car Racing controller weights are consistent with an input of z (32) plus h (256), giving 288 × 3 + 3 = 867 model.py. The Doom controller input is z (64) plus c (512) plus h (512), giving 1,088 doomrnn.py config.py.

Results

Author-reported. Setup: CarRacing-v0, controller scores averaged over 100 random trials §3.3 tab. 2. Each model is trained as described in Method.

Car Racing tab. 2:

The text reports the full model as the first reported solution of the task §3.3. The V-only result is the controller seeing z alone, with no memory state §3.3 §A.4. V and M never see reward; only C does §3.1.

Take Cover, dream and real scores tab. 4. Setup: DoomTakeCover, 100 consecutive trials on the real environment; virtual scores are inside the M dream §4.3 §4.4:

The text states the best real score as 1,092 ± 556 at τ = 1.15 §4.5 tab. 4 and reports about 900 time steps in the dream §4.3. The paper also reports 959 over 1,024 rollouts in the dream at τ = 1.15 §A.5.

Paper vs code

Code: hardmaru/WorldModelsExperiments, commit fd982b9691a941b52c6addbde29bc801ca6202c8, owned by David Ha, a paper author. The paper's interactive article points to the author's blog for reproducing the experiments, and that post names this repository as the reference implementation. The repository's README gives the licence as MIT. There is no LICENSE file in the repository, so the licence rests on that README statement.

Discrepancies, each with both citations:

• VAE epochs. The paper trains the Car Racing VAE for one epoch §A.1. The code's training script runs 10 epochs vae_train.py, and its README says the same README.md. The Doom VAE script also uses 10 epochs vae_train.py.
• MDN-RNN training. The paper trains the MDN-RNN for 20 epochs §A.2. The code trains for 4,000 steps, each on a fresh random batch of 100 sequences of length 999 rnn_train.py rnn_train.py.
• Doom dream temperature. The paper uses τ = 1.15 for the reported Take Cover agent §A.5. The code's dream temperature constant is 1.25 doomrnn.py doomrnn.py.
• Default optimizer. The module-level default is PEPG, while the command-line default is CMA-ES train.py train.py. Running the code as documented uses CMA-ES.
• Antithetic seeds. The command-line default enables antithetic seed reuse, in which each seed is used for two candidates train.py train.py. The paper does not mention this.
• VAE KL term. The code floors the KL term at 0.5 × z_size per batch vae.py. The paper describes only a KL loss §A.1.

Consistent with the paper: controller input for Car Racing is [z, h] (MODE_ZH) model.py; Doom uses [z, c, h] doomrnn.py; population 64 and 16 rollouts train.py train.py; evaluation every 25 generations train.py; the Doom done rule at 50% doomrnn.py; Doom action thresholds at ±1/3 doomreal.py.

Pitfalls

• Latent sampling at inference. The code samples z = μ + σ·ε for every frame at run time model.py. The paper does not say whether the agent should use the mean instead. Test both and record which one you use.
• Temperature. The paper does not define how τ acts. The code divides the mixture logits by τ and scales noise by √τ rnn.py. A different reading of τ will change the Take Cover results in Table 4.
• Action mapping. The paper gives the ranges of the three Car Racing outputs §A.3. The code also remaps the gas output from [−1, 1] to [0, 1] with (x + 1) / 2 model.py. Skipping that remap gives a different action.
• Parameter rounding. The code rounds all parameters to four decimal places before evaluation train.py train.py, and it stores integers scaled by 10,000 in its JSON files vae.py. The paper does not mention this. Small differences in scores may come from it.
• Hidden state reset. Both the rollout and the dream reset M's hidden state at the start of each episode model.py doomrnn.py. A policy evaluated without this reset starts each episode with a stale memory.
• Cost hidden by the abstract. The evolution stage runs about 1,800 generations, each with 64 candidates and 16 rollouts §A.4. That is roughly 1.8 million episodes. The "fewer than an hour per model" figure refers only to the two world-model stages §3.1 footnote.
• The dream can be cheated. The paper reports that the controller found a policy that exploits the dream model and fails in the real game at low τ §4.5 tab. 4. Raising τ reduces this, but at high τ the dream becomes too hard to learn in.
• Missing optimizer details. The paper gives no learning rates, batch sizes or gradient clipping for V and M. The code uses Adam with learning rates 1e-4 (VAE) and 1e-3 (MDN-RNN) and clips MDN-RNN gradients at ±1 vae.py rnn.py.
• Scores are author-reported. The 906 ± 21 figure comes from 100 trials of one trained controller. The paper does not give variance across independent training runs.

Lineage

• Builds on: the paper's own earlier framework Learning to Think, which it cites for its C-M structure and its iterative training procedure §1 §5. No arXiv id for that work appears in the source, so it has no card here.
• Compared against: DQN, A3C and the Gym Leaderboard for Car Racing tab. 2; PILCO and related Gaussian-process or Bayesian dynamics models §6.
• Superseded by: not known from content in this wiki.
• Related card, not cited by this paper: DreamerV3, which also learns a recurrent world model and trains a policy on imagined rollouts.