Writing · · 5 min read
Why World Models Matter
A world model learns the environment's dynamics so that planning can happen in imagination rather than in the world. That buys sample efficiency, transfer, and the ability to imagine better interventions.
#world-models#rl#generative-models
A world model converts the environment from a bottleneck into a function that can be evaluated offline. The minimal version is a table. The usable version is a learned latent simulator paired with a policy trained inside it.
The smallest world model
The model-free policy learns from interaction. The model-based alternative learns the transition first. In the tabular case, the dynamics are counts:
from collections import defaultdict, Counter
import numpy as np
class TabularWorldModel:
"""p(s'|s,a) estimated from observed transitions."""
def __init__(self):
self.counts = defaultdict(Counter)
self.reward_sum = defaultdict(float)
self.reward_count = defaultdict(int)
def observe(self, s, a, r, s_next, done):
self.counts[(s, a)][s_next] += 1
self.reward_sum[(s, a)] += r
self.reward_count[(s, a)] += 1
def mean_reward(self, s, a):
if self.reward_count[(s, a)] == 0:
return 0.0
return self.reward_sum[(s, a)] / self.reward_count[(s, a)]
def sample_next(self, s, a, rng):
dist = self.counts.get((s, a))
if dist is None:
return s # unknown transitions are treated as self-loops
states, weights = zip(*dist.items())
return rng.choice(states, p=[w / sum(weights) for w in weights])That table is a complete world model. Given a state and an action, sampling from sample_next produces a plausible next state. A planner or a value function can be trained on rollouts from the table without touching the real environment.
Why tables stop scaling
The table works because s, a, and s' are small integers. The environment breaks that assumption in two directions: the observation is high-dimensional, and the state is partially observed.
Camera frames, lidar scans, and proprioceptive vectors do not admit a small integer coding. A table over raw pixels has no entries. Even a discretized image space explodes combinatorially: every pixel value multiplies the state count.
The learning problem also changes. In model-free reinforcement learning, the objective is to maximize expected return:
where is the policy and is the discount factor. The agent queries the environment for data, applies a temporal-difference or policy-gradient update, and queries again. Data collection is the scarce resource. Compute is comparatively cheap.
A world model reverses the ratio. Dynamics are learned once from a modest dataset, then megabytes of imagined transitions are generated offline. The price is modeling error: the policy optimizes the model's dynamics, not the environment's dynamics.
Step 1: compress the observation
"World Models" by David Ha and Jürgen Schmidhuber (2018) separates the task. An encoder compresses each frame into a compact latent vector. A sequence model predicts the next latent vector and the reward. The policy operates on the latent space, where training is small and fast.
A minimal modern form uses a variational autoencoder followed by a recurrent dynamics model:
import torch
import torch.nn as nn
class Encoder(nn.Module):
def __init__(self, latent_dim=32):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(3, 32, 4, stride=2),
nn.ReLU(),
nn.Conv2d(32, 64, 4, stride=2),
nn.ReLU(),
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(64, latent_dim * 2), # mean and log-variance
)
def forward(self, frame):
params = self.net(frame)
mean, log_var = params.chunk(2, dim=-1)
std = (0.5 * log_var).exp()
eps = torch.randn_like(mean)
z = mean + std * eps
kl = -0.5 * (1 + log_var - mean.pow(2) - log_var.exp()).sum(-1).mean()
return z, kl
class Dynamics(nn.Module):
"""Predicts next latent, reward, and termination probability."""
def __init__(self, latent_dim=32, hidden=256):
super().__init__()
self.core = nn.Sequential(
nn.Linear(latent_dim + 6, hidden), # latent + one-hot action
nn.Tanh(),
nn.Linear(hidden, hidden),
nn.Tanh(),
)
self.next_latent = nn.Linear(hidden, latent_dim)
self.reward = nn.Linear(hidden, 1)
self.done = nn.Linear(hidden, 1)
def forward(self, z, action_onehot):
h = self.core(torch.cat([z, action_onehot], dim=-1))
return self.next_latent(h), self.reward(h).squeeze(-1), self.done(h).squeeze(-1)The encoder is trained for reconstruction. The dynamics model is trained for prediction:
def latent_dynamics_loss(encoder, dynamics, frames, actions, rewards, dones):
z, kl = encoder(frames[0])
losses = [kl * 1e-4]
for t in range(len(actions)):
z_pred, r_pred, done_pred = dynamics(z, actions[t])
# predict the encoding of the next frame, frozen as a target
with torch.no_grad():
z_target, _ = encoder(frames[t + 1])
losses.append(((z_pred - z_target) ** 2).mean())
losses.append(((r_pred - rewards[t]) ** 2).mean())
losses.append(torch.nn.functional.binary_cross_entropy_with_logits(done_pred, dones[t]))
z = z_pred # free-running rollout during training
return sum(losses)The bottleneck at this stage is the two-stage pipeline. The representation is chosen for reconstruction fidelity, not for the prediction task. A latent that preserves pixel detail may discard the small signal that determines the next reward. A latent that compresses aggressively may be easy to reconstruct and useless for control.
The failure mode is specific: the dynamics model learns to predict whatever is most predictable in the latent. The policy inherits those blind spots. That is compounded by compounding error. One-step prediction error of becomes multi-step error that grows with horizon, so short rollouts look accurate while long rollouts drift off the data manifold.
Step 2: train representation and dynamics together
The Dreamer line of work (Hafner et al., 2019) removes the split. Representation and dynamics are trained end to end, shaped by the prediction task. The recurrent state-space model keeps a deterministic hidden state for memory and a stochastic state for observation:
import torch
import torch.nn as nn
class RSSM(nn.Module):
def __init__(self, latent_dim=32, hidden=1024):
super().__init__()
self.hidden_dim = hidden
self.prior = nn.Sequential(
nn.Linear(hidden + 6, hidden // 2),
nn.ReLU(),
nn.Linear(hidden // 2, latent_dim * 2),
)
self.posterior = nn.Sequential(
nn.Linear(hidden + 6 + 64, hidden // 2), # + CNN embedding of frame
nn.ReLU(),
nn.Linear(hidden // 2, latent_dim * 2),
)
self.recurrent = nn.GRUCell(6 + latent_dim, hidden)
self.reward_head = nn.Linear(hidden, 1)
def prior_params(self, h, action):
out = self.prior(torch.cat([h, action], dim=-1))
return out.chunk(2, dim=-1)
def posterior_params(self, h, action, frame_features):
out = self.posterior(torch.cat([h, action, frame_features], dim=-1))
return out.chunk(2, dim=-1)
def observe(self, h, action, frame_features):
mean, log_var = self.posterior_params(h, action, frame_features)
std = (0.5 * log_var).exp()
z = mean + std * torch.randn_like(mean)
h_next = self.recurrent(torch.cat([action, z], dim=-1), h)
return h_next, z
def imagine_step(self, h, action):
mean, log_var = self.prior_params(h, action)
std = (0.5 * log_var).exp()
z = mean + std * torch.randn_like(mean)
h_next = self.recurrent(torch.cat([action, z], dim=-1), h)
return h_next, zThe encoder features are computed by a small convolutional network. The loss is a sum of reconstruction, reward prediction, and a KL term that balances prior and posterior. Training the representation on the prediction loss keeps the latent aligned with what the next decision requires.
Imagination as the training engine
Once the dynamics are stable, the policy is trained on imagined trajectories. The actor and critic never see the environment directly:
def imagine_train(rssm, actor, critic, replay, rng, imagine_horizon=15):
# sample a starting latent from real experience
batch = replay.sample()
h, z = rssm.observe(batch.initial_h, batch.initial_action, batch.initial_features)
# unroll the policy inside the model
hs, zs, rewards, log_probs = [h], [z], [], []
for _ in range(imagine_horizon):
action = actor(z)
h, z = rssm.imagine_step(h, action)
hs.append(h)
zs.append(z)
rewards.append(rssm.reward_head(h).squeeze(-1))
# bootstrapped value target: sum of rewards + value of final state
credits = torch.stack(rewards)
values = critic(torch.stack(zs))
# critic gradient flows only through values, not through imagined rewards
lambdas = torch.full_like(credits, 0.95)
targets = discounts_to_targets(credits, values, lambdas)
actor_loss = -(log_probs * targets.detach()).mean()
critic_loss = ((values - targets.detach()) ** 2).mean()
return actor_loss + critic_lossThe ledger is legible here. The critic is fitted to imagined returns. The actor is pushed up the imagined advantage. Both costs are computed without a single environment step. The real environment is consulted only to supply starting states and to refit the model.
Bottleneck: the policy exploits the model
A policy trained inside a single dynamics model learns the model's errors. Hallucinated transitions — shortcuts the model predicts but the real environment does not support — are systematically selected for. The symptom is a large gap between imagined return and real return.
The remedies are procedural, not magical. Each targets a different failure:
def ensemble_imagine_step(models, h, z, action):
"""Advance an ensemble; return mean prediction and disagreement."""
predictions = []
for m in models:
h_next, z_next = m.imagine_step(h, action)
predictions.append(z_next)
stacked = torch.stack(predictions)
mean = stacked.mean(dim=0)
disagreement = stacked.var(dim=0).sum(dim=-1) # scalar uncertainty per step
return h, mean, disagreement
def real_return_evaluation(policy, env, episodes=20):
returns = []
for _ in range(episodes):
obs, done, total = env.reset(), False, 0.0
while not done:
action = policy.act(obs)
obs, reward, done = env.step(action)
total += reward
returns.append(total)
return float(torch.tensor(returns).mean())
def training_protocol(model, policy, replay, env):
for iteration in range(1000):
model.update(replay)
policy = imagine_train(model, policy, replay)
if iteration % 100 == 0:
imagined = model.imagined_return(policy)
real = real_return_evaluation(policy, env)
print(f"step {iteration}: imagined={imagined:.2f} real={real:.2f}")
# a widening gap is the signal to refit the model, not to keep polishing the policyEnsembling raises the price of hallucination: a single model's shortcut scores low under disagreement penalty. Short horizons limit drift. Periodic real evaluation detects the gap before deployment. None of these fixes the underlying asymmetry — the model is always an approximation — but each makes the exploitation harder to sustain.
Transfer through the model
A trained dynamics model is a reusable simulator. Three transfers follow directly from the decoupling:
- A new reward head can be fitted on the same latent dynamics, producing a policy for a different objective without new interaction.
- The latent space can be adapted to a new environment with fine-tuning, while the policy architecture stays fixed, provided the interface is preserved.
- Disturbances never seen in the original environment can be injected as perturbations of the latent transition, producing robust policies at zero interaction cost.
Each transfer inherits the same caveat: compatibility is assumed, not earned. The new environment must admit the same latent interface, or the transfer is cosmetic.
What to measure
Four quantities separate a working world model from a decorative one:
- one-step prediction error on held-out real transitions
- multi-step prediction error across the planning horizon
- imagined return versus real return for the same policy
- transfer return when only the reward head or environment changes
A model with low one-step error and a wide imagined-to-real gap is predicting noise well and dynamics poorly. A model with exact multi-step error on held-out data has earned trust within the tested range.
Checklist
- Estimate dynamics from real transitions; validate on held-out real transitions.
- Train representation and dynamics together when the latent must serve control.
- Keep imagined horizons short relative to the drift curve; bootstrap values at the boundary.
- Use an ensemble or disagreement penalty when the policy is trained purely inside the model.
- Report imagined and real return side by side in every evaluation.
- Treat policy behavior as a proposal to be checked in the real environment, not as a conclusion.
- Refit the model when the gap widens; polishing the policy against a drifting model amplifies the error.
The bet is that dynamics are learnable and that the latent worth learning is the one that matters for the task. The bet pays in sample efficiency, in interventions that cost nothing, and in a simulator that can be re-used. The map is not the territory. A world model is a map drawn to be traveled.