Recurrent World Models Facilitate Policy Evolution

In plain words

Think of an agent that learns a rough picture of the world it lives in, the way a person forms a mental model from what they see. The agent first learns to squeeze each camera image into a short list of numbers that keep the important parts of the scene. A second network learns to guess how that list will change over time, including the uncertainty in that guess. Only then does the agent learn what to do, and it does this with a very small decision rule that is tuned by trial and error rather than by gradient descent.

The surprising part is how little the decision rule needs. Because the predictive network already summarises what has happened and what is likely to happen next, the decision rule can be a single linear layer. The authors show this on a top-down car racing game and on a game in which monsters shoot fireballs at the player. In the second game, the agent is trained entirely inside the model's own imagined version of the game, and the policy still works when it is moved back into the real game.

The imagined game is not perfect. The agent can learn to exploit mistakes in the model. The authors add noise to the model's predictions during training, which makes the imagined game harder and less easy to exploit. This knob is the temperature.

TL;DR

The paper compiles three parts. A convolutional VAE encodes each 64x64 frame into a latent vector z §2 app. A.2. An LSTM with a mixture density output (MDN-RNN) predicts the next z from the action and history §2 app. A.3. A single linear controller maps z and the LSTM hidden state to actions §2 app. A.4, and CMA-ES evolves its weights against the real reward app. A.5. The reported car racing score is 906 ± 21, which the authors say is the first reported solution of that task §3.1. A controller trained only inside the learned DoomTakeCover model scores 1092 ± 556 on the real task at temperature 1.15 §4.2 tab. 2. All numbers are author-reported.

Problem

Learning from raw pixels with a policy that is small enough to search by evolution is hard. The authors note that a good representation of the scene is the usual solution, and that hand-made features such as LIDAR and velocities have made small driving policies work before §3.1. They also want to test whether a learned model can replace the real environment when the controller is trained §4 §1. Earlier model-based work either still trained on the real environment or used models that were deterministic and easy to exploit §4.2 §5.

The method assumes frames can be modelled as a sequence of latent vectors, that a random policy yields enough data to train the frame encoder and the predictive model §3 alg. 1, and that the controller is small enough for evolution §2. It treats the reward as visible only to the controller §3.

Method

The agent has three parts §2 fig. 2:

1. V, the frame encoder. A convolutional VAE takes a 64x64 RGB frame scaled to [0, 1] and outputs the mean and log-variance of a diagonal Gaussian over a latent z. The latent size is 32 for car racing and 64 for Doom app. A.2. Training uses a reconstruction term plus a KL term app. A.2.
2. M, the predictive model. An LSTM with an MDN output predicts the density of the next z given the current z, the action and the hidden state §2. The density is a mixture of 5 Gaussians per latent dimension, with a diagonal covariance and no correlation terms app. A.3. Sampling from M uses a temperature τ §2. For Doom, M also predicts a binary death flag §4.1 app. A.3.
3. C, the controller. A single linear layer maps the sampled z and the LSTM hidden state h to the action, a_t = W_c [z_t h_t] + b_c §2. The car controller therefore has 32 + 256 inputs and 3 outputs, which gives 867 parameters §3 app. A.1. The Doom controller takes z, the LSTM cell state c and h, which is 1088 inputs to one output app. A.1 app. A.4.

Training runs in four steps alg. 1 §3:

• Collect 10,000 rollouts from a random policy §3 alg. 1.
• Train V on those frames to encode each frame into z §3 app. A.2.
• Train M on the sequences of encoded frames and actions, teacher-forced §3 app. A.3.
• Evolve C with CMA-ES to maximise the expected cumulative reward of a rollout §2 app. A.5.

Only step four uses the reward, and V and M never see it §3.

The paper's agent loop, in pseudocode, is:

# Pipeline  [alg. 1] [§3]
# Only the controller uses the reward. V and M see frames and actions only.

def build_agent(random_rollouts):
    # Fit the frame encoder V on frames from random rollouts  [app. A.2]
    vae = train_conv_vae(random_rollouts.frames, z_size=32)

    # Store mu and logvar per frame; z is drawn again at every batch  [app. A.3]
    mus, logvars, actions = encode_rollouts(vae, random_rollouts)

    # Fit the MDN-RNN M by teacher forcing on (z, action) sequences  [app. A.3]
    mdn_rnn = train_mdn_rnn(mus, logvars, actions)

    # Evolve the linear controller C against the real episode reward  [app. A.5]
    controller = evolve_cma_es(vae, mdn_rnn, num_generations=1800)

    return vae, mdn_rnn, controller

At each environment step the controller reads the encoder code and the RNN state §2 app. A.4:

# Agent step at time t  [§2] [app. A.4] [code: carracing/model.py#L98-L118]
# Only z_t and h_t reach the controller; c_t is not used for car racing.

def act(frame, rnn_state, controller):
    # Encode the frame and draw one latent sample  [app. A.2]
    z = sample_gaussian(vae_encode_mu_logvar(frame))    # (32,)

    # Read the LSTM output vector h_t from the recurrent state  [§2]
    features = concatenate(z, rnn_state.h)              # (288,)

    # Linear controller, then a bounded activation  [§2] [app. A.4]
    raw = tanh(features @ controller.W + controller.b)  # (3,)

    # Map to (steering, gas, brake)  [app. A.4] [code: model.py#L111-L114]
    steering = raw[0]
    gas = (raw[1] + 1.0) / 2.0
    brake = clip(raw[2], 0.0, 1.0)
    return (steering, gas, brake)

The MDN-RNN is sampled with temperature τ, which the paper uses to set how noisy the imagined environment is §2 §4.2:

# Sample the next latent from the MDN-RNN  [§2] [§4.2] [app. A.3]
# tau flattens the mixture weights and widens the noise.

def sample_next_z(mdn_out, tau):
    # Divide the mixture logits by tau, then normalise  [code: rnn/rnn.py#L81]
    weights = softmax(mdn_out.logmix / tau)              # (32, 5)

    # Choose one Gaussian per latent dimension, independently  [app. A.3]
    component = sample_categorical(weights)              # (32,)
    mean = mdn_out.mean[component]
    std = exp(mdn_out.logstd[component])

    # Scale the unit noise by sqrt(tau)  [code: rnn/rnn.py#L95-L96]
    noise = normal(0, 1, size=32) * sqrt(tau)
    return mean + std * noise

The VAE objective, as implemented, is shown below app. A.2 [code: carracing/vae/vae.py#L65-L79]:

# VAE loss for one batch  [app. A.2] [code: carracing/vae/vae.py#L65-L79]
# The KL term has a floor of kl_tolerance * z_size per sample (free bits).

def vae_loss(frames, recon, mu, logvar, kl_tolerance, z_size):
    # Squared error summed over pixels, averaged over the batch  [app. A.2]
    recon_loss = mean(sum(square(frames - recon), axis=(1, 2, 3)))

    # Gaussian KL to the unit prior, per sample  [app. A.2]
    kl = -0.5 * sum(1 + logvar - mu ** 2 - exp(logvar), axis=1)

    # Clamp each sample's KL from below  [code: carracing/vae/vae.py#L76]
    kl = maximum(kl, kl_tolerance * z_size)
    return recon_loss + mean(kl)

The controller is searched with CMA-ES over a population of 64, each candidate scored by the mean return of 16 rollouts with different seeds app. A.5:

# CMA-ES search over the controller weights  [app. A.5] [§2]
# Fitness is the mean return over 16 seeds; higher is better.

def evolve_cma_es(vae, mdn_rnn, num_generations, population=64, rollouts=16):
    # Start the search distribution at zero mean  [code: carracing/es.py#L100]
    search = cma_init(num_params=867, sigma=0.1, popsize=population)

    for generation in range(num_generations):
        # Sample candidate weight vectors  [app. A.5]
        candidates = search.ask()

        # Score each candidate on the real environment  [app. A.5]
        returns = [
            mean_return(candidate, rollouts, vae, mdn_rnn)
            for candidate in candidates
        ]

        # Update the distribution from the scores  [app. A.5]
        search.tell(returns)

    return search.current_mean()

The Doom experiment uses the same procedure with three changes §4.1:

• M also predicts a done flag, and the virtual episode ends when the probability of done exceeds 50% app. A.3.
• The virtual environment exposes z to C and steps M without rendering pixels, so the controller never sees a frame during training §4.1.
• The controller is trained at temperature τ = 1.15 and then tested on the real game app. A.6 §4.2.

Hyperparameters

Values are from the paper unless the Where column says code. "not stated" means the paper does not give it.

Results

Author-reported. All scores are mean ± standard deviation over 100 trials unless noted. Car racing setup: CarRacing-v0, 100 trials, the baselines are the cited ones §3.1 tab. 1.

Temperature sweep for the Doom experiment, with each controller trained in the virtual environment and then run on the real one tab. 2 §4.2:

Parameter counts tab. 3 tab. 4 app. A.1:

Paper vs code

Code: raw/code/1809.01999 at commit fd982b9691a941b52c6addbde29bc801ca6202c8, from the repository linked from the authors' project page (via the blog post it links). The repository has two folders, carracing and doomrnn. Its README states an MIT licence, but the clone has no LICENSE file, so code_license reflects the README. The code is not quoted in this card, only cited.

Discrepancies between paper and code:

• VAE epochs: the paper says 1 epoch of training on random rollouts app. A.2. Both carracing/vae_train.py and doomrnn/vae_train.py run 10 epochs vae_train.py vae_train.py.
• MDN-RNN length: the paper says 20 epochs app. A.3. The car script trains for 4000 steps of 100 sequences each, which is about 40 passes over the 10,000 series rnn_train.py rnn_train.py. The Doom script loops over 400 epochs rnn_train.py.
• Doom temperature: the paper trains the Doom controller at τ = 1.15 [app. A.6]. The code's virtual environment defaults to 1.25 [code: doomrnn/doomrnn.py#L24]. The code does not show which value produced the reported numbers.
• Data collection: the paper says the data comes from a random policy §3 alg. 1. The car script re-initialises the controller, the VAE and the MDN-RNN for each episode with Cauchy-distributed weights whose scale is a random value up to 0.01, then records the actions it takes [code: carracing/extract.py#L36] model.py. The actions therefore come from a random network, not from uniform random actions. The number of episodes depends on how many workers ran the script, which the code does not fix: each worker writes up to 200 episodes [code: carracing/extract.py#L13].
• Controller output mapping: the paper says the controller is bounded with tanh app. A.4. The code applies tanh to steering, rescales gas from [-1, 1] to [0, 1], and clips brake to [0, 1] model.py.
• Sampling temperature: the paper says τ controls uncertainty, citing SketchRNN §2. The code divides the mixture logits by τ and multiplies the unit noise by sqrt(τ) rnn.py rnn.py. The Doom virtual environment does the same doomrnn.py.
• Learning-rate schedule: the car script has a decay formula, but its decay rate is 1.0, so the rate stays at 1e-3 throughout rnn_train.py rnn_train.py.
• CMA-ES details not in the paper: zero initial mean, initial sigma 0.1 by default, and an L2 term of 0.01 added to the fitness es.py train.py es.py es.py.
• Doom restart handling: the MDN-RNN restart loss is weighted by 10 for positive examples doomrnn.py doomrnn.py. During training, the LSTM state is zeroed at each restart doomrnn.py. The paper mentions neither.
• Car racing dream environment: the code has a dream environment with a default temperature of 0.7 dream_env.py. The training scripts do not import it, which matches the paper's use of the real environment for car racing §3.

No discrepancies were found on the controller size, the population size of 64, the 16 rollouts per candidate, or the car latent size of 32 model.py train.py train.py vae_train.py.

Pitfalls

• Epoch counts are inconsistent. The paper's values (1 VAE epoch, 20 MDN epochs) do not match the released scripts (10 VAE epochs, 4000 car steps, 400 Doom epochs). Use the code's counts only if you need to match the released results, and state which you used.
• Temperature formula is not in the paper. The paper names τ but does not give the sampling rule. The code's rule (mixture logits divided by τ, noise scaled by sqrt(τ)) is the one you need to reproduce the Doom curve.
• Temperature is the exploit knob. At τ = 0.10 the virtual environment reaches 2086 but the real score is 193 tab. 2. Choose τ by the real score, not the virtual one.
• The model can be exploited. The paper says the controller can find policies that work only inside the model, because it sees the model's full hidden state §4.2. A deterministic model is easier to exploit than a stochastic one §4.2.
• Mode collapse at low τ. At τ = 0.10 the MDN-RNN stops producing fireballs, so a policy learned there does worse than random in the real game §4.2.
• KL floor. The code clamps each sample's KL from below at 0.5 * N_z [code: carracing/vae/vae.py#L76]. Below the floor the KL term has no gradient. The paper does not mention this.
• Restart weighting. The Doom restart loss is weighted 10 times [code: doomrnn/doomrnn.py#L21]. This changes how often the model predicts death. Not stated in the paper.
• Gradient clipping is per entry. The code clips each gradient element to ±1.0 rnn.py. This is not a global norm clip. The paper does not state a clipping rule.
• Random policy is not uniform. Data from randomly initialised networks covers a different set of states than uniform random actions. The paper does not say which was used extract.py.
• Cost of the search. The paper's numbers imply 64 candidates × 16 rollouts per generation, so 1024 episodes per generation. At 1800 generations that is about 1.8 million episodes (derived from app. A.5), and each episode runs the VAE and MDN-RNN on every frame.
• Parameter count of the model is large. The VAE and MDN-RNN together have about 4.8 million parameters for car racing tab. 3, far more than the 867 controller parameters. Only the controller is searched by evolution.
• The VAE reconstructs irrelevant detail. The paper reports that it reproduced wall tiles in Doom but not road tiles in car racing, and that it does not get the exact monster count §6 §4.1. The imagined environment inherits these errors.
• Limits stated by the authors. The per-frame latent is a single diagonal Gaussian §4.2. The authors name higher-capacity models and external memory as future work for harder tasks §6. A VAE trained without rewards may keep details that do not matter for the task §6.
• What the ± values measure. The reported standard deviations are over 100 test rollouts of the reported controller §3.1 tab. 2. The paper does not say whether training was repeated with other seeds.

Lineage

• Builds on: the VAE §2 app. A.2, the MDN-RNN as used for handwriting and sketches (Graves; SketchRNN) §2 app. A.3, CMA-ES §2 app. A.5, and the evolution-strategy method "Evolving Stable Strategies" app. A.5. Earlier RNN-based world models and controllers from 1990 to 2015 are cited as the basis of the controller-model split §1 §5.
• Compared against: DQN, A3C (continuous and discrete) and the public leaderboard tab. 1. For DoomTakeCover, the public leaderboard result tab. 2.
• Extended by: the paper lists an iterative training procedure, artificial curiosity, and a single network combining controller and model as future work §6. No later work is cited here as a successor, since none is in the wiki.