World Models, SimPLe, and PlaNet › Worked example
World Models HandbookAll chapters World Models, SimPLe, and PlaNet › Worked exampleOne section · Sections in this chapter6 / 9 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
To make the ideas concrete without needing an Atari emulator, let us build a small variational latent dynamics model and compare two decision couplings: a tiny controller optimized through imagined rollouts and a CEM planner that searches the same learned prior online. SimPLe's alternating observation-space training loop is not reproduced by this example. The synthetic environment exposes its full two-dimensional state as the observation, so this is a dynamics example rather than a demonstration of partial observability. It is illustrative, not a benchmark.
We start with the shared plotting setup, then define the synthetic world.
The synthetic world is a controlled damped linear system. Its fully observed state contains position and velocity, and reward peaks only when both are near the origin. The system is already passively stable, so a no-action baseline is essential: a learned controller should be credited only for improvement beyond that contraction.
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
torch.manual_seed(7)
rng = np.random.default_rng(7)
## Fully observed state s = [position, velocity]; action a is a 1-D acceleration.
A = np.array([[0.95, 0.10], [0.00, 0.90]])
B = np.array([[0.0], [0.10]])
ACTION_LIMIT = 1.0
def step_world(s, a):
"""Advance and return the next fully observed state."""
s_next = A @ s + B.flatten() * a
return s_next
def reward_of(s):
"""A smooth reward peaking at the origin; asks the controller to stabilize."""
return float(np.exp(-0.5 * np.sum(s**2)))We now generate a training batch of short trajectories with exploratory actions and record transitions of the form (observation, action, next observation, reward). Exploration draws from an unbounded Gaussian. ACTION_LIMIT does not clip these training actions. It bounds only the tiny controller and CEM candidates later. The learned model therefore sees some action magnitudes outside their control range.
def collect_data(n_episodes=200, horizon=40):
obs, act, next_obs, rew = [], [], [], []
for _ in range(n_episodes):
s = rng.normal(scale=0.6, size=2)
for _ in range(horizon):
a = float(rng.normal(scale=0.8))
s_next = step_world(s, a)
obs.append(s.copy())
act.append([a])
next_obs.append(s_next.copy())
rew.append([reward_of(s_next)])
s = s_next
return (
np.asarray(obs, dtype=np.float32),
np.asarray(act, dtype=np.float32),
np.asarray(next_obs, dtype=np.float32),
np.asarray(rew, dtype=np.float32),
)
obs, act, next_obs, rew = collect_data()
print(
f"Collected {obs.shape[0]} transitions, observation dimension {obs.shape[1]}"
)We can inspect the collected transitions before training. Because reward depends jointly on next position and next velocity, color is a more faithful encoding than a curve obtained by sorting one coordinate.
With a two-dimensional observed state, we do not need a VAE for compression, but the same probabilistic interfaces make the temporal targets explicit. The encoder maps each state to a Gaussian latent. The prior maps to a distribution over . The decoder maps latents back to states. During training, the reward head receives samples from the observation-conditioned next-state posterior. During held-out one-step evaluation, imagined-controller optimization, and CEM planning, each reward-head query receives a transitioned prior mean instead. The deployed tiny controller selects actions from the encoded current state and does not itself query the reward head. The KL term encourages the next-posterior and transition-prior distributions to agree, but it does not make posterior samples equal to the prior mean used by these evaluation and planning paths. This toy setup therefore has a train-use shift in reward prediction. For image observations, SimPLe and PlaNet use convolutional encoders or decoders to exploit spatial structure. When the current observation omits decision-relevant information that can be recovered from history, a temporal state can help retain it. The small fully connected networks below are only for this toy system.
class Encoder(nn.Module):
def __init__(self, obs_dim, latent_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(obs_dim, 32), nn.ReLU(), nn.Linear(32, 2 * latent_dim)
)
def forward(self, x):
mu, log_var = self.net(x).chunk(2, dim=-1)
return mu, log_var.clamp(-6, 2)
class Prior(nn.Module):
def __init__(self, latent_dim, act_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim + act_dim, 32),
nn.ReLU(),
nn.Linear(32, 2 * latent_dim),
)
def forward(self, z, a):
mu, log_var = self.net(torch.cat([z, a], dim=-1)).chunk(2, dim=-1)
return mu, log_var.clamp(-6, 2)
class Decoder(nn.Module):
def __init__(self, latent_dim, obs_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim, 32), nn.ReLU(), nn.Linear(32, obs_dim)
)
def forward(self, z):
return self.net(z)
class RewardHead(nn.Module):
def __init__(self, latent_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim, 32), nn.ReLU(), nn.Linear(32, 1)
)
def forward(self, z):
return self.net(z)
def reparameterize(mu, log_var):
sigma = torch.exp(0.5 * log_var)
eps = torch.randn_like(sigma)
return mu + sigma * eps
def gaussian_kl(mu_q, log_var_q, mu_p, log_var_p):
var_q, var_p = torch.exp(log_var_q), torch.exp(log_var_p)
return 0.5 * torch.sum(
(var_q + (mu_q - mu_p) ** 2) / var_p - 1.0 + log_var_p - log_var_q,
dim=-1,
)
latent_dim = 3
encoder = Encoder(obs.shape[1], latent_dim)
prior = Prior(latent_dim, act.shape[1])
decoder = Decoder(latent_dim, obs.shape[1])
reward_head = RewardHead(latent_dim)
params = (
list(encoder.parameters())
+ list(prior.parameters())
+ list(decoder.parameters())
+ list(reward_head.parameters())
)
optimizer = torch.optim.Adam(params, lr=3e-3)
obs_t = torch.from_numpy(obs)
act_t = torch.from_numpy(act)
next_obs_t = torch.from_numpy(next_obs)
rew_t = torch.from_numpy(rew)
## Held-out split for honest evaluation.
n_train = int(0.9 * obs_t.shape[0])
train_idx = torch.arange(n_train)
test_idx = torch.arange(n_train, obs_t.shape[0])Each transition supplies both current and next posterior targets. We reconstruct from and from , match the transition prior to the next posterior, and predict from . An additional prior-decoding loss compares the transitioned prior mean with ; this is an auxiliary one-step prediction term rather than a separate ELBO likelihood.
def train_step(batch_size=256):
idx = train_idx[torch.randint(len(train_idx), (batch_size,))]
x = obs_t[idx]
a = act_t[idx]
x_next = next_obs_t[idx]
r_next = rew_t[idx]
mu_q, log_var_q = encoder(x)
mu_q_next, log_var_q_next = encoder(x_next)
z = reparameterize(mu_q, log_var_q)
z_next = reparameterize(mu_q_next, log_var_q_next)
mu_p_next, log_var_p_next = prior(z, a)
current_recon = F.mse_loss(decoder(z), x)
next_recon = F.mse_loss(decoder(z_next), x_next)
prior_next_recon = F.mse_loss(decoder(mu_p_next), x_next)
reward_loss = F.mse_loss(reward_head(z_next), r_next)
kl = gaussian_kl(
mu_q_next, log_var_q_next, mu_p_next, log_var_p_next
).mean()
loss = (
current_recon + next_recon + prior_next_recon + reward_loss + 0.1 * kl
)
optimizer.zero_grad()
loss.backward()
optimizer.step()
return (
float(loss.detach()),
float(current_recon.detach()),
float(prior_next_recon.detach()),
float(reward_loss.detach()),
float(kl.detach()),
)
history = []
for step in range(1500):
loss, recon, pred, rloss, kl = train_step()
if step % 250 == 0:
history.append((step, loss, recon, pred, rloss, kl))We can also evaluate the prior's one-step prediction error on held-out transitions, which is the first hint at whether planning in this model could work at all.
Training trace (step, loss, current recon, prior next, reward, kl): step 0: loss=1.9210 recon=0.1665 prior_next=0.1251 reward=1.4458 kl=0.2322 step 250: loss=0.1106 recon=0.0254 prior_next=0.0215 reward=0.0033 kl=0.3922 step 500: loss=0.0700 recon=0.0021 prior_next=0.0021 reward=0.0014 kl=0.6216 step 750: loss=0.0652 recon=0.0012 prior_next=0.0010 reward=0.0006 kl=0.6131 step 1000: loss=0.0731 recon=0.0011 prior_next=0.0010 reward=0.0003 kl=0.6970 step 1250: loss=0.0655 recon=0.0009 prior_next=0.0007 reward=0.0003 kl=0.6263 Held-out prior next-state MSE: 0.0003 Held-out prior next-reward MSE: 0.0005
After 1,500 optimizer updates, the held-out metrics above test one-step prior predictions rather than posterior reconstructions: the prior is invoked with , its next latent is decoded against , and the same next latent predicts . These checks do not establish long-horizon accuracy, but they provide a basic sanity check before control; multi-step and closed-loop evaluation are still needed to assess decision usefulness.
First, we optimize a tiny tanh-squashed affine controller through imagined rollouts generated by the trained prior. We backpropagate through deterministic prior-mean transitions and the reward head; no stochastic latent trajectory is sampled during this optimization. This is not the CMA-ES procedure from World Models. The analogy is limited to a small controller acting on learned latent features.
controller = nn.Linear(latent_dim, 1)
## Seed imagined rollouts from real encoded latents.
with torch.no_grad():
seed_idx = torch.randint(len(train_idx), (64,))
seed_mu, _ = encoder(obs_t[train_idx[seed_idx]])
seed_z = seed_mu
def imagine_return(z0, horizon=15, gamma=0.95):
"""Roll prior means forward and sum the predicted next-state rewards."""
z = z0
total = torch.zeros(z0.shape[0])
disc = 1.0
for _ in range(horizon):
a = ACTION_LIMIT * torch.tanh(controller(z))
mu_p, _ = prior(z, a)
z = mu_p
total = total + disc * reward_head(z).squeeze(-1)
disc *= gamma
return total
def train_controller(steps=400):
# Freeze the fitted world model; retain gradients through it to controller actions.
for module in (encoder, prior, decoder, reward_head):
module.requires_grad_(False)
opt = torch.optim.Adam(controller.parameters(), lr=1e-2)
for _ in range(steps):
returns = imagine_return(seed_z)
loss = -returns.mean()
opt.zero_grad()
loss.backward()
opt.step()
with torch.no_grad():
return float(imagine_return(seed_z).mean())
imagined_return = train_controller()
print(f"Imagined return of the trained tiny controller: {imagined_return:.4f}")Second, at each real step we infer the current latent and run a CEM-style search over action sequences using the same prior means and reward head. Prior means remove latent-transition sampling; CEM still samples action sequences. The fixed NumPy and PyTorch seeds support repeatable runs in the same software, platform, and device environment, but do not guarantee identical numbers across releases or hardware. PlaNet itself sampled one stochastic latent trajectory per candidate. As in the historical planner's bounded-action step, each iteration here samples an unbounded Gaussian parameterized in action coordinates, clips each sampled action once to the same range as the tiny controller, and refits the Gaussian moments from the clipped elites. The returned action is the first component of the final proposal mean, clamped to that range.
def cem_plan(z0, horizon=8, n_candidates=128, n_elites=16, n_iter=3):
"""Plan a single action for a single latent state using CEM."""
action_mean = torch.zeros(horizon, 1)
action_std = torch.full((horizon, 1), 0.7)
with torch.no_grad():
for _ in range(n_iter):
samples = action_mean + action_std * torch.randn(
n_candidates, horizon, 1
)
samples = samples.clamp(-ACTION_LIMIT, ACTION_LIMIT)
scores = torch.zeros(n_candidates)
z = z0.unsqueeze(0).repeat(n_candidates, 1)
disc = 1.0
for t in range(horizon):
mu_p, _ = prior(z, samples[:, t])
z = mu_p
scores += disc * reward_head(z).squeeze(-1)
disc *= 0.95
elite_idx = torch.topk(scores, n_elites).indices
elites = samples[elite_idx]
action_mean = elites.mean(dim=0)
action_std = elites.std(dim=0, correction=0).clamp_min(1e-2)
return action_mean[0].clamp(-ACTION_LIMIT, ACTION_LIMIT)
def evaluate(policy_fn, initial_states, horizon=25):
returns = []
for initial_state in initial_states:
s = initial_state.copy()
total = 0.0
for _ in range(horizon):
a = policy_fn(s)
s = step_world(s, a)
total += reward_of(s)
returns.append(total)
return np.asarray(returns)
def tiny_controller_policy(s):
with torch.no_grad():
mu_q, _ = encoder(torch.from_numpy(s).unsqueeze(0).float())
a = ACTION_LIMIT * torch.tanh(controller(mu_q)).item()
return a
def cem_policy(s):
with torch.no_grad():
mu_q, _ = encoder(torch.from_numpy(s).unsqueeze(0).float())
a = cem_plan(mu_q[0])
return float(a)
def no_action_policy(s):
return 0.0
evaluation_starts = rng.normal(scale=0.6, size=(24, 2)).astype(np.float32)
returns_no_action = evaluate(no_action_policy, evaluation_starts)
returns_tiny = evaluate(tiny_controller_policy, evaluation_starts)
returns_cem = evaluate(cem_policy, evaluation_starts)
for name, values in (
("No-action baseline", returns_no_action),
("Tiny imagined controller", returns_tiny),
("Latent CEM planner", returns_cem),
):
print(
f"{name:27s} mean 25-step return: {values.mean():.4f} "
f"(SE across starts {values.std(ddof=1) / np.sqrt(len(values)):.4f})"
)No-action baseline mean 25-step return: 22.3583 (SE across starts 0.4295) Tiny imagined controller mean 25-step return: 23.7577 (SE across starts 0.2144) Latent CEM planner mean 25-step return: 23.6934 (SE across starts 0.2070)
The printed results use the same 24 initial states for all three methods and include the passive baseline. They are a seeded illustration, not evidence that one coupling is generally better. After the common encoder pass, the tiny controller needs one controller pass per action, whereas CEM repeatedly evaluates its prior and reward head over a population of sequences online; the code therefore demonstrates a computation tradeoff, not a robustness guarantee.
Let's visualize the real-state trajectories of the passive baseline, tiny controller, and CEM planner, plus one imagined reward trajectory for the tiny controller.
The phase portrait shows three descriptive trajectories, one per method, from a shared start; the paired return summary above is the broader comparison. In the seeded prior-mean rollout, predicted reward declines slightly rather than rising monotonically. That curve does not, by itself, demonstrate model exploitation: doing so would require comparing imagined and true outcomes for the same action sequence and showing that the controller selected a systematic model error.