Policies Learned in Imagination › Implementing Imagination-Based Learning
World Models HandbookAll chapters Policies Learned in Imagination › Implementing Imagination-Based LearningOne section · Sections in this chapter7 / 10 ↓
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
Let us build the experiment. We start with the environment and a data collection routine, then train the two models, then run actor-critic learning inside each.
The environment is small enough to express directly, and we will use the analytic reward for evaluation while the models learn their own reward head.
import numpy as np
import torch
import torch.nn as nn
DT, DRAG = 0.1, 0.99
GAMMA, LAM = 0.99, 0.95
TARGET, CTRL = 2.0, 0.05
EVAL_HORIZON = 40
torch.set_num_threads(1)
def true_reward(x, a):
"""Analytic reward: used to score the real environment during evaluation."""
return -((x - TARGET) ** 2) - CTRL * a**2The analytic reward supplies labels for the real training transitions and scores real evaluation rollouts. During imagined actor updates, the learner sees only the fitted reward head. Keeping those roles distinct lets us compare model-predicted return with the actual task reward; it does not mean the reward head was trained without true-reward labels.
Data collection draws episodes with a chosen initial-state distribution and a chosen action scale. Both dataset parameters are what distinguish Model A from Model B.
def collect(init_fn, action_scale, n_episodes=300, ep_len=40, seed=0):
"""Roll out the true environment under a uniform random behaviour policy."""
rng = np.random.default_rng(seed)
s_buf, a_buf, r_buf, sn_buf = [], [], [], []
for _ in range(n_episodes):
x, v = init_fn(rng)
for _ in range(ep_len):
a = rng.uniform(-action_scale, action_scale)
r = true_reward(x, a)
x_next = x + DT * v + rng.normal(0.0, 0.01)
v_next = DRAG * v + DT * a + rng.normal(0.0, 0.01)
s_buf.append((x, v))
a_buf.append((a,))
r_buf.append(r)
sn_buf.append((x_next, v_next))
x, v = x_next, v_next
return (
np.asarray(s_buf, dtype=np.float32),
np.asarray(a_buf, dtype=np.float32),
np.asarray(r_buf, dtype=np.float32)[:, None],
np.asarray(sn_buf, dtype=np.float32),
)
narrow = collect(
lambda rng: (rng.normal(0.0, 0.25), rng.normal(0.0, 0.25)),
action_scale=0.4,
seed=1,
)
broad = collect(
lambda rng: (rng.uniform(-3.0, 3.0), rng.uniform(-0.8, 0.8)),
action_scale=1.0,
seed=2,
)narrow: 12000 transitions | |x| max 2.64 | |v| max 0.68 broad: 12000 transitions | |x| max 7.17 | |v| max 1.47
The narrow dataset has only 9 of 12,000 positions within 0.2 of and none at or beyond it; the broad dataset has 641 within that band. These counts measure one aspect of the wider collection distribution. The histograms show the positional difference, not the full action or velocity coverage.
The world model predicts the state delta and the normalized reward. Predicting a delta rather than an absolute next state is a small but meaningful inductive bias: it makes the identity dynamics the default, so the network only has to learn the correction.
The delta parameterization is worth a sentence of justification, because it is the kind of choice that looks cosmetic but is not. The true dynamics are close to the identity over a short step, since the position changes by and the velocity changes by a small amount. A network that predicts the next state directly has to relearn the identity mapping from scratch, spending capacity on reproducing something it already knows. A network that predicts the delta starts from the identity and only has to learn the small correction. This is the same reason residual connections help deep networks: the identity path is a strong prior, and it costs nothing to build it in.
class WorldModel(nn.Module):
"""Predicts the next state and reward for (state, action) pairs."""
def __init__(self, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(3, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, 3),
)
self.register_buffer("r_mu", torch.zeros(1))
self.register_buffer("r_sigma", torch.ones(1))
def forward(self, s, a):
out = self.net(torch.cat([s, a], dim=-1))
return s + out[:, :2], out[:, 2]
def step(self, s, a):
s_next, r_norm = self.forward(s, a)
return s_next, self.r_mu + self.r_sigma * r_normdef train_world_model(
data, steps=3000, batch=256, lr=3e-3, seed=0, bootstrap=False
):
"""Fit a dynamics-and-reward model by minibatch least squares."""
torch.manual_seed(seed)
s, a, r, sn = (torch.as_tensor(x) for x in data)
r_norm = (r - r.mean()) / r.std()
n = len(s)
if bootstrap:
idx = torch.randint(0, n, (n,))
s, a, r_norm, sn = s[idx], a[idx], r_norm[idx], sn[idx]
model = WorldModel()
model.r_mu = r.mean().reshape(1)
model.r_sigma = r.std().reshape(1)
optimiser = torch.optim.Adam(model.parameters(), lr=lr)
for _ in range(steps):
idx = torch.randint(0, n, (batch,))
s_hat, r_hat = model(s[idx], a[idx])
loss = ((s_hat - sn[idx]) ** 2).mean() + (
(r_hat - r_norm[idx, 0]) ** 2
).mean()
optimiser.zero_grad()
loss.backward()
optimiser.step()
model.eval()
return model
model_narrow = train_world_model(narrow, seed=10)
model_broad = train_world_model(broad, seed=10)Both models fit their training data well. The interesting question is what happens away from it.
The mechanism behind Model A's failure is visible in its reward head. Plotting each model's predicted reward along a slice of position reveals how far each one is willing to extrapolate, and where the optimizer will be pulled.
probe_x = np.linspace(-3.0, 4.0, 200)
probe_v = np.zeros_like(probe_x)
probe_a = np.zeros_like(probe_x)
probe_s = torch.tensor(
np.stack([probe_x, probe_v], axis=1), dtype=torch.float32
)
probe_act = torch.tensor(probe_a[:, None], dtype=torch.float32)
with torch.no_grad():
probe_reward_true = true_reward(probe_s[:, 0], probe_act[:, 0]).numpy()
probe_reward_narrow = (
model_narrow.r_mu
+ model_narrow.r_sigma * model_narrow(probe_s, probe_act)[1]
).numpy()
probe_reward_broad = (
model_broad.r_mu
+ model_broad.r_sigma * model_broad(probe_s, probe_act)[1]
).numpy()
Figure 2. Both heads peak near the goal, but Model A penalizes overshoot too weakly: at and zero action it predicts about instead of . Model B predicts about there. The slice isolates reward extrapolation, not full closed-loop behavior.
The training loss does not report this out-of-distribution reward error. The next cell probes dynamics separately by rolling both models open loop from starts near the goal. Reward accuracy and transition accuracy are distinct; neither by itself certifies a useful policy.
We test them on a probe set of trajectories that start near the target at and roll out open loop under the true dynamics. Open-loop means the model receives only the initial state and the action sequence; every predicted state becomes the input for the next prediction, so errors compound.
Open-loop probing measures how state error grows when predictions become inputs. The narrow model has sparse goal-region data rather than none, and the probe compares its error with the broad model's under the same initial states and actions.
def open_loop_error(model, s0, actions, true_states):
"""Per-step mean squared state error of an open-loop model rollout."""
s = s0.clone()
errors = []
for t in range(actions.shape[1]):
with torch.no_grad():
s = model.step(s, actions[:, t])[0]
errors.append(((s - true_states[t + 1]) ** 2).sum(-1).mean().item())
return np.array(errors)
probe_rng = np.random.default_rng(7)
n_probe, probe_horizon = 32, 25
probe_s0 = torch.tensor(
np.stack(
[
2.0 + 0.2 * probe_rng.normal(size=n_probe),
0.5 * probe_rng.normal(size=n_probe),
],
axis=1,
),
dtype=torch.float32,
)
probe_actions = torch.tensor(
probe_rng.uniform(-0.5, 0.5, size=(n_probe, probe_horizon, 1)),
dtype=torch.float32,
)
true_probe_states = [probe_s0]
s = probe_s0
for t in range(probe_horizon):
a = probe_actions[:, t, 0]
s = torch.stack([s[:, 0] + DT * s[:, 1], DRAG * s[:, 1] + DT * a], dim=1)
true_probe_states.append(s)
error_narrow = open_loop_error(
model_narrow, probe_s0, probe_actions, true_probe_states
)
error_broad = open_loop_error(
model_broad, probe_s0, probe_actions, true_probe_states
)
Figure 3. Both models accumulate open-loop state error. At step one, the narrow and broad MSEs are approximately and ; at step 25 they are and . The narrow model is roughly 8.5 times worse at the end of this probe.
The first-step errors are small but differ by a factor of about 3.2. Both errors then grow, with Model A ending substantially higher. The models fit their respective training samples, but we have not shown that their training losses are equal or that Model B is accurate everywhere its actor will go.
The probe separates small one-step error from larger multi-step drift. It shows why reporting only in-distribution training loss would be insufficient: these rollouts begin near a goal that Model A's collection policy rarely visited. Closed-loop policy evaluation remains a separate test.
Now the policy. The actor samples a Gaussian raw action and applies so the environment action stays in . Its MLP mean and learned state-independent log-standard-deviation parameterize the raw Gaussian. The critic is a plain MLP. Bounding actions is part of the toy task specification, not a seed-specific rescue: without it, the earlier unbounded actor could produce exploding real trajectories even with the broad model.
class Actor(nn.Module):
"""Tanh-squashed Gaussian policy; environment actions lie in [-1, 1]."""
def __init__(self, hidden=64):
super().__init__()
self.body = nn.Sequential(
nn.Linear(2, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, 1),
)
self.log_std = nn.Parameter(torch.tensor([-1.0]))
def forward(self, s):
return torch.tanh(self.body(s))
def dist(self, s):
mean = self.body(s)
return torch.distributions.Normal(
mean, self.log_std.exp().expand_as(mean)
)
class Critic(nn.Module):
def __init__(self, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, hidden),
nn.SiLU(),
nn.Linear(hidden, hidden),
nn.SiLU(),
nn.Linear(hidden, 1),
)
def forward(self, s):
return self.net(s).squeeze(-1)The -return recursion is a short backward loop. Note that the values entering the target are detached, which keeps the critic's regression target fixed within a step.
def lambda_returns(rewards, values, gamma=GAMMA, lam=LAM):
"""Backward recursion for lambda-returns along an imagined rollout."""
vals = [v.detach() for v in values]
acc = torch.zeros_like(vals[0])
out = [None] * len(rewards)
for t in reversed(range(len(rewards))):
delta = rewards[t] + gamma * vals[t + 1] - vals[t]
acc = delta + gamma * lam * acc
out[t] = vals[t] + acc
return outThe imagination loop collects states, actions, rewards, values, and log-probabilities. Setting pathwise=True keeps the sampled action attached to the graph so gradients flow through the model; the default detaches it for the score-function estimate.
The pathwise flag is the switch between the two gradient estimators from earlier in the chapter, and it is worth tracing what it changes. When pathwise=False, the action is detached before it enters the model, so the model sees a constant input and the reward-objective gradient reaches the actor through the score-function term. If the optional entropy coefficient eta is positive, the actor also receives a direct entropy-regularization gradient; the plotted comparisons set eta=0. When pathwise=True, the action stays attached, so the reward and the next state both depend on the policy's parameters and the actor receives the analytic gradient through the model. The same rollout function supports both, which makes the comparison in the final cell a one-line change.
def imagine(model, actor, critic, s0, horizon, pathwise=False):
"""Roll the world model forward under the actor for `horizon` steps."""
s = s0
states, actions, rewards, values, logps = [], [], [], [], []
for _ in range(horizon):
states.append(s)
dist = actor.dist(s)
raw_a = dist.rsample()
a = torch.tanh(raw_a)
a_model = a if pathwise else a.detach()
actions.append(a_model)
s_next, r_hat = model.step(s, a_model)
rewards.append(r_hat)
# The critic learns from imagined states, but its update must not
# backpropagate through the actor or world-model rollout graph. The
# actor update below owns that graph in pathwise mode.
values.append(critic(s.detach()))
# The tanh transform is parameter-independent when its sampled raw
# action is held fixed; its score is the raw Normal score.
logps.append(dist.log_prob(raw_a.detach()).sum(-1))
s = s_next
values.append(critic(s.detach()))
return states, actions, rewards, values, logpsTwo evaluation helpers let us watch the gap. imagined_return scores the deterministic policy under the model; true_return scores the same policy in the real environment. Both use a 40-step horizon so the numbers are comparable even when the training horizon is shorter.
Both evaluations use the same deterministic mean policy, start states, horizon, and discount. The real evaluation uses the analytic, noise-free dynamics and reward; data collection alone adds small transition noise. The imagined evaluation uses the learned transition and reward heads. Their difference measures this complete model-versus-environment mismatch, not only transition error.
@torch.no_grad()
def true_return(actor, s0, horizon=EVAL_HORIZON):
"""Closed-loop return of the deterministic policy in the real environment."""
s = s0.clone()
total = torch.zeros(s.shape[0])
discount = 1.0
for _ in range(horizon):
a = actor(s)
total = total + discount * true_reward(s[:, 0], a[:, 0])
s = torch.stack(
[s[:, 0] + DT * s[:, 1], DRAG * s[:, 1] + DT * a[:, 0]], dim=1
)
discount *= GAMMA
return total.mean().item()
@torch.no_grad()
def imagined_return(actor, model, s0, horizon=EVAL_HORIZON):
"""Closed-loop return of the deterministic policy under the world model."""
s = s0.clone()
total = torch.zeros(s.shape[0])
discount = 1.0
for _ in range(horizon):
a = actor(s)
s, r_hat = model.step(s, a)
total = total + discount * r_hat
discount *= GAMMA
return total.mean().item()The training loop ties it together. The critic uses detached -targets, while the actor uses a sampled, finite-horizon reward-to-go and a state-dependent critic baseline. Both actor-gradient branches therefore target the same truncated discounted reward objective. The actor gradient is clipped, and the model parameters are excluded from its optimizer.
The score branch retains the outcome magnitude rather than normalizing advantages from the same samples, which would change the exact score estimate. Its sampled return is truncated at even though the critic target bootstraps there; this distinction is deliberate. Action squashing enforces the task's actuator limit and prevents the unbounded-policy divergence seen in the initial diagnostic run. Clipping controls update size but does not guarantee convergence or protection against model error.
def train_imagination(
model,
horizon=20,
iterations=250,
batch=64,
lr=1e-3,
gamma=GAMMA,
lam=LAM,
eta=1e-3,
seed=0,
pathwise=False,
uncertainty=None,
penalty=0.0,
eval_states=None,
record_every=25,
):
"""Actor-critic training entirely inside the learned model."""
if pathwise and penalty > 0.0:
raise ValueError(
"The toy pathwise branch does not differentiate uncertainty"
)
torch.manual_seed(seed)
for parameter in model.parameters():
parameter.requires_grad_(False)
actor, critic = Actor(), Critic()
opt_a = torch.optim.Adam(actor.parameters(), lr=lr)
opt_c = torch.optim.Adam(critic.parameters(), lr=lr)
history = {"iter": [], "imagined": [], "true": []}
for it in range(iterations):
s0 = torch.randn(batch, 2) * torch.tensor([0.25, 0.25])
states, actions, rewards, values, logps = imagine(
model, actor, critic, s0, horizon, pathwise
)
if uncertainty is not None and penalty > 0.0:
rewards = [
r - penalty * uncertainty(s.detach(), a.detach())
for r, s, a in zip(rewards, states, actions)
]
returns = lambda_returns(rewards, values, gamma, lam)
v_stack = torch.stack(values[:-1])
g_stack = torch.stack(returns).detach()
critic_loss = ((v_stack - g_stack) ** 2).mean()
opt_c.zero_grad()
critic_loss.backward()
opt_c.step()
if pathwise:
actor_loss = (
-sum((gamma**t) * rewards[t] for t in range(horizon)).mean()
/ horizon
)
else:
# A full sampled reward-to-go gives a score-function estimate of
# the same truncated H-step objective as the pathwise branch.
tail = torch.zeros_like(rewards[0])
reward_to_go = [None] * horizon
for t in reversed(range(horizon)):
tail = rewards[t].detach() + gamma * tail
reward_to_go[t] = tail
advantage = torch.stack(reward_to_go) - v_stack.detach()
discounts = torch.tensor(
[gamma**t for t in range(horizon)], dtype=advantage.dtype
)[:, None]
actor_loss = (
-(discounts * advantage * torch.stack(logps)).sum(dim=0).mean()
/ horizon
)
if eta > 0.0:
entropy = torch.stack(
[actor.dist(s).entropy().sum(-1) for s in states]
).mean()
actor_loss = actor_loss - eta * entropy
opt_a.zero_grad()
actor_loss.backward()
torch.nn.utils.clip_grad_norm_(actor.parameters(), 5.0)
opt_a.step()
if eval_states is not None and (
it % record_every == 0 or it == iterations - 1
):
history["iter"].append(it)
history["imagined"].append(
imagined_return(actor, model, eval_states)
)
history["true"].append(true_return(actor, eval_states))
return {"actor": actor, "critic": critic, "history": history}torch.manual_seed(123)
eval_states = 0.02 * torch.randn(64, 2)
eval_states[:, 0] = eval_states[:, 0] + 0.0
res_narrow = train_imagination(
model_narrow,
seed=0,
eval_states=eval_states,
horizon=30,
iterations=400,
eta=0.0,
)
res_broad = train_imagination(
model_broad,
seed=0,
eval_states=eval_states,
horizon=30,
iterations=400,
eta=0.0,
)
Figure 4. Both policies improve overall in this run, with visible fluctuations. Model A ends at approximately imagined versus real return; Model B ends at versus . A positive imagination gap is visible for Model A, but this chart alone does not establish its cause.
Model B's final predicted and real returns happen to agree closely here; Model A's do not. Neither curve rises smoothly, and neither comparison proves that coverage alone caused the gap. The reward probe and real trajectory help identify the plausible mechanism: Model A underprices overshoot outside its dense data region.
The plot does not show critic loss or training stability directly. It does show why an imagined learning curve is not enough: a model can predict more return than the same policy earns in the environment. Repeating the experiment across actor seeds is needed before treating the size of that gap as a stable property of the collection regime.
seed_check = []
for actor_seed in (0, 1, 2):
for name, model, first_run in (
("A/narrow", model_narrow, res_narrow),
("B/broad", model_broad, res_broad),
):
run = (
first_run
if actor_seed == 0
else train_imagination(
model,
seed=actor_seed,
eval_states=eval_states,
horizon=30,
iterations=400,
eta=0.0,
)
)
imagined = imagined_return(run["actor"], model, eval_states)
real = true_return(run["actor"], eval_states)
seed_check.append((actor_seed, name, imagined, real))
print(
f"seed {actor_seed} {name:>8}: imagined {imagined:7.2f} | real {real:7.2f}"
)seed 0 A/narrow: imagined -47.68 | real -51.26 seed 0 B/broad: imagined -47.24 | real -47.14
seed 1 A/narrow: imagined -60.33 | real -60.26
seed 1 B/broad: imagined -49.83 | real -50.71
seed 2 A/narrow: imagined -47.78 | real -51.26
seed 2 B/broad: imagined -52.70 | real -52.16
The three fixed seeds show variation in both real return and the sign or size of the gap. They are a robustness check for this toy protocol, not an uncertainty interval for a broader population of tasks.
The consequences for behavior are equally visible. Rolling each trained policy out in the real environment shows where it actually ends up, as opposed to where its imagination believes it is.
@torch.no_grad()
def policy_positions(actor, s0, horizon=EVAL_HORIZON):
"""Closed-loop mean position of the deterministic policy in the real environment."""
s = s0.clone()
xs = [s[:, 0].mean().item()]
for _ in range(horizon):
a = actor(s)
s = torch.stack(
[s[:, 0] + DT * s[:, 1], DRAG * s[:, 1] + DT * a[:, 0]], dim=1
)
xs.append(s[:, 0].mean().item())
return np.array(xs)
pos_narrow = policy_positions(res_narrow["actor"], eval_states)
pos_broad = policy_positions(res_broad["actor"], eval_states)
pos_steps = np.arange(len(pos_narrow))
Figure 5. Neither policy settles inside the shaded band at step 40. Model A ends near , Model B near . The legend maps each line to its trained model.
Next we sweep the imagination horizon while keeping the actor seed, iteration budget, optimizer, and 40-step evaluation protocol fixed. Horizon changes both the training objective and the trajectories seen by the optimizer. Short rollouts may omit useful delayed consequences; longer ones may accumulate more model error. Neither effect guarantees a monotone curve.
horizons = [5, 10, 20, 30, 40]
sweep = {"horizon": [], "imagined": [], "true": [], "gap": []}
for h in horizons:
out = train_imagination(
model_narrow, horizon=h, seed=0, iterations=400, eta=0.0
)
imag = imagined_return(out["actor"], model_narrow, eval_states)
real = true_return(out["actor"], eval_states)
sweep["horizon"].append(h)
sweep["imagined"].append(imag)
sweep["true"].append(real)
sweep["gap"].append(imag - real)
Figure 6. At , predicted return is about while real return is about ; at , they are about and . The gap can also be negative at shorter horizons. The comparison is a single-seed sensitivity check, not evidence that increasing must help or hurt.
Neither curve is monotone. The shortest training horizons perform poorly, plausibly because they provide little credit for reaching from near the origin. The run has a large positive gap, while the longer and runs perform better under this fixed iteration budget. Horizon is therefore a hyperparameter to validate in the closed loop, not a one-direction safety dial.
These policies are trained separately, with nonconvex function approximation and finite optimization. A longer rollout changes the objective rather than enlarging a nested feasible set. The gap at a particular horizon reports model-versus-environment return for that trained policy; it does not identify the cause without inspecting trajectories and errors.
Now to value expansion. The next diagnostic starts from real replay states but uses a fixed, generated action sequence in both model and true dynamics. This isolates prefix-return prediction error; it is not an implementation of MVE's policy-generated imagined actions. As grows, the model's mean partial return can drift from truth while the coefficient on an -step terminal bootstrap, , decreases. The panels have different units and do not locate an optimal horizon.
branch_rng = np.random.default_rng(11)
n_branch, branch_horizon = 64, 30
branch_idx = branch_rng.choice(len(narrow[0]), size=n_branch, replace=False)
branch_s0 = torch.tensor(narrow[0][branch_idx])
branch_actions = torch.tensor(
branch_rng.uniform(-0.8, 0.8, size=(n_branch, branch_horizon, 1)),
dtype=torch.float32,
)
true_branch_states = [branch_s0]
s = branch_s0.clone()
for t in range(branch_horizon):
a = branch_actions[:, t, 0]
s = torch.stack([s[:, 0] + DT * s[:, 1], DRAG * s[:, 1] + DT * a], dim=1)
true_branch_states.append(s)
model_rewards, true_rewards = [], []
s = branch_s0.clone()
with torch.no_grad():
for t in range(branch_horizon):
a = branch_actions[:, t]
s, r_hat = model_narrow.step(s, a)
model_rewards.append(r_hat)
true_rewards.append(true_reward(true_branch_states[t][:, 0], a[:, 0]))
def discounted_prefix(reward_list):
out, acc, discount = [], torch.zeros_like(reward_list[0]), 1.0
for r in reward_list:
acc = acc + discount * r
out.append(acc)
discount *= GAMMA
return out
model_prefix = discounted_prefix(model_rewards)
true_prefix = discounted_prefix(true_rewards)
bias_curve = np.array(
[
abs(m.mean().item() - t.mean().item())
for m, t in zip(model_prefix, true_prefix)
]
)
bootstrap_weight = np.array([GAMMA ** (h + 1) for h in range(branch_horizon)])
horizon_axis = np.arange(1, branch_horizon + 1)
Figure 7a. The plotted quantity is the absolute difference of two batch means, not the mean absolute per-trajectory error. It grows from about at one step to at 30 under this fixed generated action sequence.
Figure 7b. The terminal-value coefficient decreases from at one step to about at 30 steps. It is a coefficient, not a measurement of the critic's error.
Together, the panels show why model and critic errors must both be considered, but they cannot be crossed or compared numerically: one is in reward units and the other is dimensionless. We did not measure critic error, target variance, or complete MVE target MSE, so this diagnostic does not locate a sweet spot. STEVE estimates target variance across horizons to set mixture weights; bias still needs separate scrutiny.
Now the regularizer. We train a three-member bootstrapped ensemble on the narrow data. Its next-state prediction spread is a heuristic for epistemic uncertainty. We divide position and velocity spread by fixed toy scales of 2 and 1 before taking a root-mean-square, so the penalty is scalar and dimensionless. We do not assume that disagreement perfectly tracks error or that it must explode outside the data.
ensemble = [
train_world_model(narrow, seed=s, bootstrap=True) for s in (101, 102, 103)
]
def disagreement(models, s, a):
"""Dimensionless RMS of standardized next-state ensemble spread."""
preds = torch.stack([m(s, a)[0] for m in models])
spread = preds.std(dim=0, unbiased=False) / s.new_tensor([2.0, 1.0])
return torch.sqrt((spread**2).mean(dim=-1))
@torch.no_grad()
def trajectory_disagreement(actor, s0, models, horizon=EVAL_HORIZON):
"""Mean ensemble disagreement along the policy's closed-loop real trajectory."""
s = s0.clone()
readings = []
for _ in range(horizon):
a = actor(s)
readings.append(disagreement(models, s, a).mean().item())
s = torch.stack(
[s[:, 0] + DT * s[:, 1], DRAG * s[:, 1] + DT * a[:, 0]], dim=1
)
return float(np.mean(readings))penalties = [0.0, 1.0, 5.0, 20.0]
penalty_return, penalty_disagreement = [], []
for p in penalties:
out = train_imagination(
model_narrow,
seed=0,
horizon=30,
iterations=400,
eta=0.0,
uncertainty=lambda s, a: disagreement(ensemble, s, a),
penalty=p,
)
penalty_return.append(true_return(out["actor"], eval_states))
penalty_disagreement.append(
trajectory_disagreement(out["actor"], eval_states, ensemble)
)
Figure 8a. In this run, real return improves modestly from approximately with no penalty to at .
Figure 8b. Standardized mean ensemble disagreement falls slightly, from approximately to . This is not a direct measurement of real model error.
The left panel plots real return and the right panel plots disagreement. Both changes are modest for this seed: stronger penalty correlates with slightly better task return and slightly lower ensemble spread. The plot does not establish causality outside the specified training protocol or prove that the goal is unreachable under pessimism.
The lesson is narrower than either "pessimism always hurts" or "pessimism fixes exploitation." The penalty changes this actor's behavior enough to improve one real evaluation, but the ensemble score is not calibrated against transition error. A different seed, penalty scale, or misspecified ensemble could reverse the result. Closed-loop return and uncertainty diagnostics must be reported together.
Uncertainty penalties cannot create missing evidence. When interaction is available, targeted data collection can test the model in weakly covered regions. Choosing actions that improve both the task and what the agent knows is the subject of Active Perception, Dual Control, and Exploration.
For completeness, we compare score-function and pathwise actor updates on Model B under the same truncated discounted reward objective, seed, horizon, batch size, and iteration budget. Its open-loop probe was better than Model A's, but that does not guarantee accuracy on every visited state. One seed can illustrate behavior, not establish an estimator ranking.
score_fn = train_imagination(
model_broad,
seed=3,
horizon=30,
iterations=400,
eta=0.0,
eval_states=eval_states,
)
pathwise = train_imagination(
model_broad,
seed=3,
horizon=30,
iterations=400,
eta=0.0,
pathwise=True,
eval_states=eval_states,
)
comparison = {
"score-function": true_return(score_fn["actor"], eval_states),
"pathwise": true_return(pathwise["actor"], eval_states),
} score-function: real return -107.946
pathwise: real return -46.632
Imagination gap at the end of training
Model A: imagined -47.678 | real -51.259 | gap 3.581
Model B: imagined -47.239 | real -47.143 | gap -0.096With seed 3, the pathwise policy reaches about real return after 400 iterations while the score-function policy reaches about . This is a result of this configuration, not evidence that pathwise updates are always faster or safer. The actor losses are aligned on the same truncated objective; their gradient estimators and resulting optimization paths differ.
We have not run the corresponding estimator comparison on Model A, so we should not predict its outcome. To compare estimators more generally, vary actor seeds, inspect complete learning curves, and measure model error on the trajectories each estimator induces. Otherwise optimization and model mismatch remain entangled.