World Models

10 min read

Compress pixels into a latent state, learn how that state evolves, then let a tiny controller act through the model's predictions or train entirely inside its imagined world.

the split

Put most of the intelligence in a learned world, not the policy

The agent is deliberately divided into vision, memory, and control. V and M learn reusable structure with gradient descent; C searches a much smaller reward-driven parameter space.

PublishedICML 2018Input64 x 64 RGBLatent32D racing / 64D DoomDynamics5-Gaussian MDN-RNNData10,000 random rolloutsRacing controller867 parameters

vision

VAE compresses each frame

A convolutional VAE turns a high-dimensional observation into a compact Gaussian latent. The decoder makes the compression visible, but the controller only needs z.

xₜ → q(zₜ | xₜ)

4,348,547 parameters in CarRacing

latent zₜ

memory

MDN-RNN predicts what comes next

An LSTM reads the current latent and action. Its mixture-density head represents several plausible next latents instead of blurring them into one deterministic guess.

p(zₜ₊₁ | aₜ, zₜ, hₜ)

422,368 parameters and 256 hidden units in CarRacing

zₜ + hₜ

controller

A linear policy chooses the action

The current visual code says what is here. The recurrent hidden state says what has happened and what is likely next. A single layer maps both directly to control.

aₜ = Wc [zₜ, hₜ] + bc

867 reward-trained parameters in CarRacing

Observe

pixels xₜ

Compress + remember

zₜ and hₜ

Act and update

aₜ changes the next observation

training recipe

Learn representation first, solve credit assignment second

The modules are trained separately. This avoids asking a sparse reward signal to shape millions of visual and recurrent parameters at the same time.

1
Collect experience
Run a random policy for 10,000 rollouts and store every RGB frame and action.
real environment
2
Learn spatial codes
Train V to reconstruct each frame while regularizing its latent toward a Gaussian prior.
framewise VAE loss
3
Learn latent dynamics
Encode the dataset, then train M with teacher forcing to predict the next latent distribution.
sequence likelihood
4
Evolve the policy
Freeze V and M. CMA-ES searches C's weights using average cumulative reward as fitness.
black-box optimization

The useful bargain: V and M absorb almost all model capacity without seeing reward. Only the tiny C model sees the task objective, so evolution searches hundreds of parameters instead of millions.

experiment one

Prediction features turn a snapshot into a driving state

CarRacing isolates the value of memory. All variants see the VAE latent; only the full model also receives the MDN-RNN hidden state that was trained to predict the future.

Average reward over 100 trials

solve at 900

V only

linear C reads zₜ

632 ± 251

V + hidden C

zₜ plus a 40-unit MLP

788 ± 141

V + M

linear C reads zₜ and hₜ

906 ± 21

Bar lengths encode the reported mean only. The printed values include the paper's standard deviations.

what hₜ adds

Predictive context

A frame can show where the road is, but not the car's velocity or which way it has been turning. The recurrent state summarizes the recent sequence and the distribution of likely next frames.

The full model reaches 906 ± 21, clearing the environment's 900-point solve threshold.

experiment two

Once M predicts termination, it becomes the environment

For VizDoom Take Cover, M predicts both the next latent and the probability of death. That is enough to expose a Gym-like interface and optimize C without rendering pixels or running the game engine.

Reality supplies the dataset

VizDoom frames

z sequences

A random policy collects 10,000 real rollouts once. The VAE and dynamics model learn from that fixed experience.

The controller trains inside a latent dream

C reads zₜ, hₜ

outputs aₜ

M samples next step

zₜ₊₁ + doneₜ₊₁

Repeat until done

reward = time alive

The VAE decoder is optional during training. It only renders the latent trajectory for humans; C and M operate directly on vectors.

train in dream

959

best average over 1,024 virtual rollouts

deploy in reality

1092 ± 556

average survival over 100 real trials

solve threshold

750

required average survival time

the reality gap

A harder dream can teach a more transferable policy

Controllers optimize whatever M permits. Temperature controls the diversity of MDN-RNN samples, trading an easy but exploitable simulator for a noisy but more robust training distribution.

VizDoom score by dream temperature

temperature τ	virtual score	actual score
0.10	2086 ± 140	193 ± 58
0.50	2060 ± 277	196 ± 50
1.00	1145 ± 690	868 ± 511
1.15used	918 ± 546	1092 ± 556
1.30	732 ± 269	753 ± 139

Low temperature rewards cheating

At τ = 0.10, mode collapse makes monsters stop firing. C nearly perfects the dream, then performs below the random-policy score in reality.

Moderate noise acts like domain randomization

At τ = 1.15, uncertain fireballs make the dream harder than the real game. A policy that survives that nightmare transfers well.

Too much noise erases the lesson

Temperature is still a hyperparameter. At τ = 1.30, the policy becomes less variable but its average real score falls.

the deeper limitation

The policy finds model errors, not just task solutions

Coverage

Random rollouts are sufficient here, but harder worlds need an iterative loop that collects data where the current model is weak.

Relevance

A reward-free VAE may discard tiny details that matter for the task. Joint training can recover them, but reduces reuse across tasks.

From the paper

The learned loop and the world it imagines

The original figures make three ideas concrete: how V, M, and C close the control loop; which visual details survive compression; and what a rollout sampled entirely from the learned dynamics looks like.

World Models paper diagram showing an observation entering the VAE, the VAE latent and MDN-RNN hidden state entering the controller, and the action returning to the environment and recurrent model — Figure 8. Agent model. V compresses the current observation, M carries predictive history, and C turns both signals into an action that updates the real or learned environment.

World Models paper comparison of a sharp CarRacing observation, its latent z controls, and the blurrier VAE reconstruction — Figure 10. VAE reconstruction. The 32-dimensional racing latent discards texture and sharp edges while retaining the road geometry and car position needed for control.

World Models paper frame of a blurry CarRacing track and red car generated inside the MDN-RNN dream environment at temperature 0.25 — Figure 13. CarRacing dream. M samples the next latent, and the VAE decoder renders it for a human viewer. The controller itself acts directly on the latent state and never needs this pixel reconstruction.

The mental model

World Models separates representation learning from control. The VAE compresses each image into z, an MDN-RNN turns the recent latent and action sequence into a predictive hidden state h, and a linear controller maps [z, h] to the next action. Most parameters learn from abundant observations; only the tiny controller must solve the reward-driven search problem.

The recurrent model has two jobs. Its hidden state is a compact feature for acting in the real environment, as the CarRacing result demonstrates. When it also predicts episode termination, it can replace the environment during training. The VizDoom controller is optimized against sampled latent trajectories, then transferred back to the real game without task training there.

The dream is useful precisely because it is imperfect, but its imperfections are also the main risk. An optimizer will exploit any shortcut the learned dynamics allow. Raising sampling temperature makes those shortcuts less reliable and produces a more robust policy, up to the point where the simulated world becomes too noisy to teach anything. The lasting idea is not that models must predict every pixel perfectly; they must preserve the dynamics needed for a policy to succeed outside the model.

Sources

Paper Interactive article Experiments