Rollout and Latent Imagination — Visual Experiments¶

Goal: See why autonomous multi-step prediction is harder than one-step prediction, and how uncertainty, termination, compute, and optimization change imagined planning.

What we will explore¶

#	Experiment	Question answered
1	Teacher forcing versus free rollout	Why can a good one-step model drift in open loop?
2	Observation refresh	How does receding-horizon correction limit accumulated error?
3	Particle imagination	What does a mean trajectory lose when futures branch?
4	Continuation masks	Why must imagined rewards stop after termination?
5	Latent-only compute	Where does the efficiency advantage over decoded rollouts come from?
6	Model exploitation	Why can more candidate search produce worse real actions?

Linked theory: Rollout and Latent Imagination

In [1]:

import warnings

import matplotlib.pyplot as plt
import numpy as np

warnings.filterwarnings('ignore')

plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})

def discounted_return(rewards, continuations, gamma=0.97):
    total = 0.0
    weight = 1.0
    for reward, continuation in zip(rewards, continuations):
        total += weight * reward
        weight *= gamma * continuation
    return total

import warnings import matplotlib.pyplot as plt import numpy as np warnings.filterwarnings('ignore') plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) def discounted_return(rewards, continuations, gamma=0.97): total = 0.0 weight = 1.0 for reward, continuation in zip(rewards, continuations): total += weight * reward weight *= gamma * continuation return total

---

Experiment 1: One-step accuracy can hide free-rollout drift¶

Question: Why does teacher-forced evaluation underestimate the error seen by planning?

The learned transition is only slightly biased. Teacher forcing applies it to the true state at every step, while free rollout repeatedly feeds the model its own previous prediction.

In [2]:

np.random.seed(1)
horizon = 85
time = np.arange(horizon)
actions = 0.42 * np.sin(0.17 * time) + 0.16 * np.cos(0.05 * time)
true_states = np.zeros(horizon)

def true_transition(state, action):
    return 0.94 * state + 0.16 * np.sin(1.8 * state) + 0.28 * action

def learned_transition(state, action):
    return 0.965 * state + 0.145 * np.sin(1.72 * state) + 0.265 * action + 0.006

for step in range(horizon - 1):
    true_states[step + 1] = true_transition(true_states[step], actions[step])

teacher_forced = np.zeros(horizon)
free_rollout = np.zeros(horizon)
teacher_forced[0] = true_states[0]
free_rollout[0] = true_states[0]
for step in range(horizon - 1):
    teacher_forced[step + 1] = learned_transition(true_states[step], actions[step])
    free_rollout[step + 1] = learned_transition(free_rollout[step], actions[step])

teacher_error = np.abs(teacher_forced - true_states)
free_error = np.abs(free_rollout - true_states)

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Small local bias becomes a different state distribution in free rollout', fontsize=14, fontweight='bold')

axes[0].plot(time, true_states, color='black', linewidth=2.4, label='true trajectory')
axes[0].plot(time, teacher_forced, color='#2a9d8f', linewidth=2.0, label='teacher-forced one-step')
axes[0].plot(time, free_rollout, color='#e76f51', linewidth=2.0, label='free rollout')
axes[0].set_title('State trajectories')
axes[0].set_xlabel('Rollout step')
axes[0].set_ylabel('Latent state')
axes[0].legend(fontsize=8)

axes[1].plot(time, teacher_error, color='#2a9d8f', linewidth=2.2, label='teacher-forced error')
axes[1].plot(time, free_error, color='#e76f51', linewidth=2.2, label='free-rollout error')
axes[1].set_title('Absolute prediction error')
axes[1].set_xlabel('Rollout step')
axes[1].set_ylabel('Absolute latent error')
axes[1].legend(fontsize=9)

window_horizons = np.arange(1, horizon)
teacher_rmse = [np.sqrt(np.mean(teacher_error[1:h + 1] ** 2)) for h in window_horizons]
free_rmse = [np.sqrt(np.mean(free_error[1:h + 1] ** 2)) for h in window_horizons]
axes[2].plot(window_horizons, teacher_rmse, color='#2a9d8f', linewidth=2.2, label='teacher-forced RMSE')
axes[2].plot(window_horizons, free_rmse, color='#e76f51', linewidth=2.2, label='free-rollout RMSE')
axes[2].set_title('Error as evaluation horizon grows')
axes[2].set_xlabel('Evaluation horizon')
axes[2].set_ylabel('Cumulative RMSE')
axes[2].legend(fontsize=9)

plt.tight_layout()
plt.show()

np.random.seed(1) horizon = 85 time = np.arange(horizon) actions = 0.42 * np.sin(0.17 * time) + 0.16 * np.cos(0.05 * time) true_states = np.zeros(horizon) def true_transition(state, action): return 0.94 * state + 0.16 * np.sin(1.8 * state) + 0.28 * action def learned_transition(state, action): return 0.965 * state + 0.145 * np.sin(1.72 * state) + 0.265 * action + 0.006 for step in range(horizon - 1): true_states[step + 1] = true_transition(true_states[step], actions[step]) teacher_forced = np.zeros(horizon) free_rollout = np.zeros(horizon) teacher_forced[0] = true_states[0] free_rollout[0] = true_states[0] for step in range(horizon - 1): teacher_forced[step + 1] = learned_transition(true_states[step], actions[step]) free_rollout[step + 1] = learned_transition(free_rollout[step], actions[step]) teacher_error = np.abs(teacher_forced - true_states) free_error = np.abs(free_rollout - true_states) fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) fig.suptitle('Small local bias becomes a different state distribution in free rollout', fontsize=14, fontweight='bold') axes[0].plot(time, true_states, color='black', linewidth=2.4, label='true trajectory') axes[0].plot(time, teacher_forced, color='#2a9d8f', linewidth=2.0, label='teacher-forced one-step') axes[0].plot(time, free_rollout, color='#e76f51', linewidth=2.0, label='free rollout') axes[0].set_title('State trajectories') axes[0].set_xlabel('Rollout step') axes[0].set_ylabel('Latent state') axes[0].legend(fontsize=8) axes[1].plot(time, teacher_error, color='#2a9d8f', linewidth=2.2, label='teacher-forced error') axes[1].plot(time, free_error, color='#e76f51', linewidth=2.2, label='free-rollout error') axes[1].set_title('Absolute prediction error') axes[1].set_xlabel('Rollout step') axes[1].set_ylabel('Absolute latent error') axes[1].legend(fontsize=9) window_horizons = np.arange(1, horizon) teacher_rmse = [np.sqrt(np.mean(teacher_error[1:h + 1] ** 2)) for h in window_horizons] free_rmse = [np.sqrt(np.mean(free_error[1:h + 1] ** 2)) for h in window_horizons] axes[2].plot(window_horizons, teacher_rmse, color='#2a9d8f', linewidth=2.2, label='teacher-forced RMSE') axes[2].plot(window_horizons, free_rmse, color='#e76f51', linewidth=2.2, label='free-rollout RMSE') axes[2].set_title('Error as evaluation horizon grows') axes[2].set_xlabel('Evaluation horizon') axes[2].set_ylabel('Cumulative RMSE') axes[2].legend(fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Teacher-forced predictions remain close because every input is corrected by the true trajectory. Free rollout gradually visits different states and accumulates much larger error.
2. Why it looks this way: Local transition error changes the next input. The learned dynamics then propagates old error and adds new error, so the rollout distribution diverges from the data distribution used for one-step evaluation.
3. Key takeaway: Always report prior-only multi-step curves. One-step loss primarily measures local fitting, not simulator reliability.

---

Experiment 2: Observation refresh limits error exposure¶

Question: Why does receding-horizon control usually execute only the first part of a model plan?

We use the same biased model and action sequence, but periodically replace the imagined state with the current true/posterior state. This isolates the benefit of fresh evidence before considering action re-optimization.

In [3]:

np.random.seed(2)
refresh_intervals = [horizon + 1, 15, 5, 1]
refresh_labels = ['no refresh', 'every 15 steps', 'every 5 steps', 'every step']
refresh_colors = ['#e76f51', '#f4a261', '#457b9d', '#2a9d8f']
refresh_rollouts = {}

for interval, label in zip(refresh_intervals, refresh_labels):
    prediction = np.zeros(horizon)
    prediction[0] = true_states[0]
    for step in range(horizon - 1):
        current = true_states[step] if step > 0 and step % interval == 0 else prediction[step]
        prediction[step] = current
        prediction[step + 1] = learned_transition(current, actions[step])
    refresh_rollouts[label] = prediction

fig, axes = plt.subplots(1, 2, figsize=(13.5, 4.8))
fig.suptitle('Fresh observations bound how long model error can compound', fontsize=14, fontweight='bold')

axes[0].plot(time, true_states, color='black', linewidth=2.5, label='true trajectory')
for label, color in zip(refresh_labels, refresh_colors):
    axes[0].plot(time, refresh_rollouts[label], color=color, linewidth=1.8, label=label)
axes[0].set_title('Rollouts with periodic state correction')
axes[0].set_xlabel('Time step')
axes[0].set_ylabel('Latent state')
axes[0].legend(fontsize=8)

rmse_values = [np.sqrt(np.mean((refresh_rollouts[label] - true_states) ** 2)) for label in refresh_labels]
max_errors = [np.max(np.abs(refresh_rollouts[label] - true_states)) for label in refresh_labels]
x_positions = np.arange(len(refresh_labels))
axes[1].bar(x_positions - 0.18, rmse_values, 0.36, color=refresh_colors, label='trajectory RMSE')
axes[1].bar(x_positions + 0.18, max_errors, 0.36, color=refresh_colors, alpha=0.45, label='maximum error')
axes[1].set_xticks(x_positions, refresh_labels, rotation=18)
axes[1].set_title('Error versus correction interval')
axes[1].set_xlabel('Observation refresh policy')
axes[1].set_ylabel('Latent error')
axes[1].legend(fontsize=9)

plt.tight_layout()
plt.show()

np.random.seed(2) refresh_intervals = [horizon + 1, 15, 5, 1] refresh_labels = ['no refresh', 'every 15 steps', 'every 5 steps', 'every step'] refresh_colors = ['#e76f51', '#f4a261', '#457b9d', '#2a9d8f'] refresh_rollouts = {} for interval, label in zip(refresh_intervals, refresh_labels): prediction = np.zeros(horizon) prediction[0] = true_states[0] for step in range(horizon - 1): current = true_states[step] if step > 0 and step % interval == 0 else prediction[step] prediction[step] = current prediction[step + 1] = learned_transition(current, actions[step]) refresh_rollouts[label] = prediction fig, axes = plt.subplots(1, 2, figsize=(13.5, 4.8)) fig.suptitle('Fresh observations bound how long model error can compound', fontsize=14, fontweight='bold') axes[0].plot(time, true_states, color='black', linewidth=2.5, label='true trajectory') for label, color in zip(refresh_labels, refresh_colors): axes[0].plot(time, refresh_rollouts[label], color=color, linewidth=1.8, label=label) axes[0].set_title('Rollouts with periodic state correction') axes[0].set_xlabel('Time step') axes[0].set_ylabel('Latent state') axes[0].legend(fontsize=8) rmse_values = [np.sqrt(np.mean((refresh_rollouts[label] - true_states) ** 2)) for label in refresh_labels] max_errors = [np.max(np.abs(refresh_rollouts[label] - true_states)) for label in refresh_labels] x_positions = np.arange(len(refresh_labels)) axes[1].bar(x_positions - 0.18, rmse_values, 0.36, color=refresh_colors, label='trajectory RMSE') axes[1].bar(x_positions + 0.18, max_errors, 0.36, color=refresh_colors, alpha=0.45, label='maximum error') axes[1].set_xticks(x_positions, refresh_labels, rotation=18) axes[1].set_title('Error versus correction interval') axes[1].set_xlabel('Observation refresh policy') axes[1].set_ylabel('Latent error') axes[1].legend(fontsize=9) plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Long autonomous intervals drift farther. More frequent state correction repeatedly resets the rollout to evidence and lowers both average and worst-case error.
2. Why it looks this way: Receding-horizon execution limits the number of consecutive model-generated inputs. It does not make the model more accurate, but it reduces exposure to accumulated bias before replanning.
3. Key takeaway: Planning horizon and execution horizon are different. A controller may inspect a long future while committing to only one short action prefix.

---

Experiment 3: A mean rollout can occupy an impossible future¶

Question: What information do particles preserve when stochastic dynamics branch into distinct modes?

Each trajectory commits to an upper or lower corridor. The mean state remains near the empty center, where a nonlinear safety reward assigns a misleadingly poor value.

In [4]:

np.random.seed(3)
particle_count = 1400
particle_horizon = 42
modes = np.random.choice([-1.0, 1.0], size=particle_count)
particle_x = np.zeros((particle_horizon, particle_count))
particle_y = np.zeros((particle_horizon, particle_count))

for step in range(1, particle_horizon):
    particle_x[step] = particle_x[step - 1] + 0.13
    particle_y[step] = particle_y[step - 1] + modes * (0.035 + 0.0015 * step) + 0.010 * np.random.randn(particle_count)

mean_x = np.mean(particle_x, axis=1)
mean_y = np.mean(particle_y, axis=1)
final_y = particle_y[-1]
safe_threshold = 0.8
particle_rewards = (np.abs(final_y) > safe_threshold).astype(float)
reward_of_mean = float(abs(np.mean(final_y)) > safe_threshold)
expected_reward = np.mean(particle_rewards)

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Particle imagination preserves branches that the mean path erases', fontsize=14, fontweight='bold')

for particle in range(0, particle_count, 28):
    axes[0].plot(particle_x[:, particle], particle_y[:, particle], color='#457b9d', alpha=0.18, linewidth=1.0)
axes[0].plot(mean_x, mean_y, color='#e76f51', linewidth=2.6, label='particle mean path')
axes[0].axhspan(-safe_threshold, safe_threshold, color='#f4a261', alpha=0.15, label='unsupported/unsafe center')
axes[0].set_title('Branching imagined trajectories')
axes[0].set_xlabel('Latent coordinate 1')
axes[0].set_ylabel('Latent coordinate 2')
axes[0].legend(fontsize=8)

axes[1].hist(final_y, bins=65, density=True, color='#457b9d', alpha=0.55, label='particle final states')
axes[1].axvline(np.mean(final_y), color='#e76f51', linewidth=2.2, linestyle='--', label='mean final state')
axes[1].axvspan(-safe_threshold, safe_threshold, color='#f4a261', alpha=0.16, label='unsupported/unsafe center')
axes[1].set_title('Multimodal final-state distribution')
axes[1].set_xlabel('Final latent coordinate 2')
axes[1].set_ylabel('Probability density')
axes[1].legend(fontsize=8)

axes[2].bar(['reward(mean state)', 'mean reward(particles)'], [reward_of_mean, expected_reward], color=['#e76f51', '#2a9d8f'])
axes[2].set_title('Nonlinear reward does not commute with averaging')
axes[2].set_xlabel('Evaluation rule')
axes[2].set_ylabel('Estimated safety reward')
axes[2].set_ylim(0.0, 1.05)

plt.tight_layout()
plt.show()

np.random.seed(3) particle_count = 1400 particle_horizon = 42 modes = np.random.choice([-1.0, 1.0], size=particle_count) particle_x = np.zeros((particle_horizon, particle_count)) particle_y = np.zeros((particle_horizon, particle_count)) for step in range(1, particle_horizon): particle_x[step] = particle_x[step - 1] + 0.13 particle_y[step] = particle_y[step - 1] + modes * (0.035 + 0.0015 * step) + 0.010 * np.random.randn(particle_count) mean_x = np.mean(particle_x, axis=1) mean_y = np.mean(particle_y, axis=1) final_y = particle_y[-1] safe_threshold = 0.8 particle_rewards = (np.abs(final_y) > safe_threshold).astype(float) reward_of_mean = float(abs(np.mean(final_y)) > safe_threshold) expected_reward = np.mean(particle_rewards) fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) fig.suptitle('Particle imagination preserves branches that the mean path erases', fontsize=14, fontweight='bold') for particle in range(0, particle_count, 28): axes[0].plot(particle_x[:, particle], particle_y[:, particle], color='#457b9d', alpha=0.18, linewidth=1.0) axes[0].plot(mean_x, mean_y, color='#e76f51', linewidth=2.6, label='particle mean path') axes[0].axhspan(-safe_threshold, safe_threshold, color='#f4a261', alpha=0.15, label='unsupported/unsafe center') axes[0].set_title('Branching imagined trajectories') axes[0].set_xlabel('Latent coordinate 1') axes[0].set_ylabel('Latent coordinate 2') axes[0].legend(fontsize=8) axes[1].hist(final_y, bins=65, density=True, color='#457b9d', alpha=0.55, label='particle final states') axes[1].axvline(np.mean(final_y), color='#e76f51', linewidth=2.2, linestyle='--', label='mean final state') axes[1].axvspan(-safe_threshold, safe_threshold, color='#f4a261', alpha=0.16, label='unsupported/unsafe center') axes[1].set_title('Multimodal final-state distribution') axes[1].set_xlabel('Final latent coordinate 2') axes[1].set_ylabel('Probability density') axes[1].legend(fontsize=8) axes[2].bar(['reward(mean state)', 'mean reward(particles)'], [reward_of_mean, expected_reward], color=['#e76f51', '#2a9d8f']) axes[2].set_title('Nonlinear reward does not commute with averaging') axes[2].set_xlabel('Evaluation rule') axes[2].set_ylabel('Estimated safety reward') axes[2].set_ylim(0.0, 1.05) plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Particles form two coherent corridors, while their mean travels through a region that almost no trajectory visits. The reward of the mean state differs sharply from the expected reward across particles.
2. Why it looks this way: Averaging collapses mode identity. For nonlinear rewards, constraints, and terminal events, evaluating the average state is not equivalent to averaging outcomes over possible futures.
3. Key takeaway: Use particles, mixtures, or explicit branches when decisions depend on event probabilities or tail risk—not only on the center of a predictive distribution.

---

Experiment 4: Rewards after termination are fictional¶

Question: How can ignoring continuation masks reverse the ranking of candidate plans?

A risky plan produces high predicted rewards but terminates early. A safe plan earns smaller rewards for the full horizon. An incorrect evaluator keeps accumulating risky rewards after the episode should have ended.

In [5]:

np.random.seed(4)
return_horizon = 14
return_time = np.arange(return_horizon)
risky_rewards = 1.30 - 0.025 * return_time
safe_rewards = np.full(return_horizon, 0.62)
risky_continuation = np.ones(return_horizon)
risky_continuation[4:] = 0.0
safe_continuation = np.ones(return_horizon)

risky_masked = discounted_return(risky_rewards, risky_continuation)
safe_masked = discounted_return(safe_rewards, safe_continuation)
risky_unmasked = discounted_return(risky_rewards, np.ones_like(risky_continuation))
safe_unmasked = discounted_return(safe_rewards, np.ones_like(safe_continuation))

def cumulative_return_curve(rewards, continuations, gamma=0.97):
    values = []
    total = 0.0
    weight = 1.0
    for reward, continuation in zip(rewards, continuations):
        total += weight * reward
        values.append(total)
        weight *= gamma * continuation
    return np.asarray(values)

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Continuation masks prevent imagined reward from leaking past terminal states', fontsize=14, fontweight='bold')

axes[0].plot(return_time, risky_rewards, color='#e76f51', linewidth=2.3, label='risky predicted reward')
axes[0].plot(return_time, safe_rewards, color='#2a9d8f', linewidth=2.3, label='safe predicted reward')
axes[0].axvline(4, color='black', linestyle='--', linewidth=1.2, label='risky terminal')
axes[0].set_title('Per-step model rewards')
axes[0].set_xlabel('Imagined step')
axes[0].set_ylabel('Predicted reward')
axes[0].legend(fontsize=8)

axes[1].plot(return_time, cumulative_return_curve(risky_rewards, risky_continuation), color='#e76f51', linewidth=2.3, label='risky, masked')
axes[1].plot(return_time, cumulative_return_curve(safe_rewards, safe_continuation), color='#2a9d8f', linewidth=2.3, label='safe, masked')
axes[1].plot(return_time, cumulative_return_curve(risky_rewards, np.ones_like(risky_continuation)), color='#e76f51', linestyle='--', linewidth=1.8, label='risky, unmasked')
axes[1].set_title('Accumulated imagined return')
axes[1].set_xlabel('Imagined step')
axes[1].set_ylabel('Discounted cumulative return')
axes[1].legend(fontsize=8)

x_positions = np.arange(2)
axes[2].bar(x_positions - 0.18, [risky_unmasked, safe_unmasked], 0.36, color=['#e76f51', '#2a9d8f'], label='ignore terminal')
axes[2].bar(x_positions + 0.18, [risky_masked, safe_masked], 0.36, color=['#e76f51', '#2a9d8f'], alpha=0.48, label='use continuation mask')
axes[2].set_xticks(x_positions, ['risky plan', 'safe plan'])
axes[2].set_title('Plan ranking depends on terminal semantics')
axes[2].set_xlabel('Candidate plan')
axes[2].set_ylabel('Predicted return')
axes[2].legend(fontsize=8)

plt.tight_layout()
plt.show()

np.random.seed(4) return_horizon = 14 return_time = np.arange(return_horizon) risky_rewards = 1.30 - 0.025 * return_time safe_rewards = np.full(return_horizon, 0.62) risky_continuation = np.ones(return_horizon) risky_continuation[4:] = 0.0 safe_continuation = np.ones(return_horizon) risky_masked = discounted_return(risky_rewards, risky_continuation) safe_masked = discounted_return(safe_rewards, safe_continuation) risky_unmasked = discounted_return(risky_rewards, np.ones_like(risky_continuation)) safe_unmasked = discounted_return(safe_rewards, np.ones_like(safe_continuation)) def cumulative_return_curve(rewards, continuations, gamma=0.97): values = [] total = 0.0 weight = 1.0 for reward, continuation in zip(rewards, continuations): total += weight * reward values.append(total) weight *= gamma * continuation return np.asarray(values) fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) fig.suptitle('Continuation masks prevent imagined reward from leaking past terminal states', fontsize=14, fontweight='bold') axes[0].plot(return_time, risky_rewards, color='#e76f51', linewidth=2.3, label='risky predicted reward') axes[0].plot(return_time, safe_rewards, color='#2a9d8f', linewidth=2.3, label='safe predicted reward') axes[0].axvline(4, color='black', linestyle='--', linewidth=1.2, label='risky terminal') axes[0].set_title('Per-step model rewards') axes[0].set_xlabel('Imagined step') axes[0].set_ylabel('Predicted reward') axes[0].legend(fontsize=8) axes[1].plot(return_time, cumulative_return_curve(risky_rewards, risky_continuation), color='#e76f51', linewidth=2.3, label='risky, masked') axes[1].plot(return_time, cumulative_return_curve(safe_rewards, safe_continuation), color='#2a9d8f', linewidth=2.3, label='safe, masked') axes[1].plot(return_time, cumulative_return_curve(risky_rewards, np.ones_like(risky_continuation)), color='#e76f51', linestyle='--', linewidth=1.8, label='risky, unmasked') axes[1].set_title('Accumulated imagined return') axes[1].set_xlabel('Imagined step') axes[1].set_ylabel('Discounted cumulative return') axes[1].legend(fontsize=8) x_positions = np.arange(2) axes[2].bar(x_positions - 0.18, [risky_unmasked, safe_unmasked], 0.36, color=['#e76f51', '#2a9d8f'], label='ignore terminal') axes[2].bar(x_positions + 0.18, [risky_masked, safe_masked], 0.36, color=['#e76f51', '#2a9d8f'], alpha=0.48, label='use continuation mask') axes[2].set_xticks(x_positions, ['risky plan', 'safe plan']) axes[2].set_title('Plan ranking depends on terminal semantics') axes[2].set_xlabel('Candidate plan') axes[2].set_ylabel('Predicted return') axes[2].legend(fontsize=8) plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Without masking, the evaluator keeps collecting high risky rewards and prefers the terminated plan. Correct continuation semantics stop accumulation and allow the longer safe plan to win.
2. Why it looks this way: A fixed-size rollout tensor still contains slots after terminal, but those slots are not future environment states. The alive mask must multiply rewards, values, and gradients after termination.
3. Key takeaway: Horizon length alone does not define return. Continuation probability, absorbing-state behavior, and mask alignment are part of the rollout contract.

---

Experiment 5: Latent-only imagination spends compute on decisions, not pixels¶

Question: How do candidate count, horizon, latent width, and output resolution affect the inner-loop workload?

This is an analytical element-count model, not a hardware benchmark. It compares touching a compact latent vector plus small heads with additionally materializing an RGB frame at every imagined step.

In [6]:

np.random.seed(5)
candidate_count = 1024
latent_dimension = 64
head_dimension = 4
horizons = np.arange(1, 65)
resolutions = [(64, 64), (128, 128), (256, 256)]

latent_elements = candidate_count * horizons * (latent_dimension + head_dimension)
decoded_elements = {
    f'{height}×{width} RGB': candidate_count * horizons * (latent_dimension + head_dimension + height * width * 3)
    for height, width in resolutions
}
ratios = {label: values / latent_elements for label, values in decoded_elements.items()}

candidate_grid = np.array([64, 256, 1024, 4096])
horizon_grid = np.array([5, 15, 30, 60])
memory_mb = candidate_grid[:, None] * horizon_grid[None, :] * latent_dimension * 4 / (1024**2)

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Decoder-free imagination avoids materializing high-dimensional observations', fontsize=14, fontweight='bold')

axes[0].plot(horizons, latent_elements, color='black', linewidth=2.5, label='latent + heads')
for label, values in decoded_elements.items():
    axes[0].plot(horizons, values, linewidth=2.0, label=label)
axes[0].set_title('Elements processed or materialized')
axes[0].set_yscale('log')
axes[0].set_ylabel('Element count (log scale)')
axes[0].set_xlabel('Rollout horizon')
axes[0].legend(fontsize=7)

for label, values in ratios.items():
    axes[1].plot(horizons, values, linewidth=2.2, label=label)
axes[1].set_yscale('log')
axes[1].set_title('Decoded-to-latent element ratio')
axes[1].set_xlabel('Rollout horizon')
axes[1].set_ylabel('Relative element count (log scale)')
axes[1].legend(fontsize=8)

image = axes[2].imshow(memory_mb, origin='lower', aspect='auto', cmap='viridis')
axes[2].set_xticks(np.arange(len(horizon_grid)), horizon_grid)
axes[2].set_yticks(np.arange(len(candidate_grid)), candidate_grid)
axes[2].set_title('Raw latent-state storage')
axes[2].set_xlabel('Horizon')
axes[2].set_ylabel('Candidate count')
for row in range(len(candidate_grid)):
    for column in range(len(horizon_grid)):
        axes[2].text(column, row, f'{memory_mb[row, column]:.1f} MB', ha='center', va='center', color='white' if memory_mb[row, column] > 2.0 else 'black', fontsize=8)
colorbar = fig.colorbar(image, ax=axes[2])
colorbar.set_label('Raw latent storage (MB)')

plt.tight_layout()
plt.show()

np.random.seed(5) candidate_count = 1024 latent_dimension = 64 head_dimension = 4 horizons = np.arange(1, 65) resolutions = [(64, 64), (128, 128), (256, 256)] latent_elements = candidate_count * horizons * (latent_dimension + head_dimension) decoded_elements = { f'{height}×{width} RGB': candidate_count * horizons * (latent_dimension + head_dimension + height * width * 3) for height, width in resolutions } ratios = {label: values / latent_elements for label, values in decoded_elements.items()} candidate_grid = np.array([64, 256, 1024, 4096]) horizon_grid = np.array([5, 15, 30, 60]) memory_mb = candidate_grid[:, None] * horizon_grid[None, :] * latent_dimension * 4 / (1024**2) fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) fig.suptitle('Decoder-free imagination avoids materializing high-dimensional observations', fontsize=14, fontweight='bold') axes[0].plot(horizons, latent_elements, color='black', linewidth=2.5, label='latent + heads') for label, values in decoded_elements.items(): axes[0].plot(horizons, values, linewidth=2.0, label=label) axes[0].set_title('Elements processed or materialized') axes[0].set_yscale('log') axes[0].set_ylabel('Element count (log scale)') axes[0].set_xlabel('Rollout horizon') axes[0].legend(fontsize=7) for label, values in ratios.items(): axes[1].plot(horizons, values, linewidth=2.2, label=label) axes[1].set_yscale('log') axes[1].set_title('Decoded-to-latent element ratio') axes[1].set_xlabel('Rollout horizon') axes[1].set_ylabel('Relative element count (log scale)') axes[1].legend(fontsize=8) image = axes[2].imshow(memory_mb, origin='lower', aspect='auto', cmap='viridis') axes[2].set_xticks(np.arange(len(horizon_grid)), horizon_grid) axes[2].set_yticks(np.arange(len(candidate_grid)), candidate_grid) axes[2].set_title('Raw latent-state storage') axes[2].set_xlabel('Horizon') axes[2].set_ylabel('Candidate count') for row in range(len(candidate_grid)): for column in range(len(horizon_grid)): axes[2].text(column, row, f'{memory_mb[row, column]:.1f} MB', ha='center', va='center', color='white' if memory_mb[row, column] > 2.0 else 'black', fontsize=8) colorbar = fig.colorbar(image, ax=axes[2]) colorbar.set_label('Raw latent storage (MB)') plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Both workloads grow linearly with candidates and horizon, but decoded rollouts add thousands to hundreds of thousands of output values per step. Even latent-state storage becomes significant when candidate and horizon axes multiply.
2. Why it looks this way: Latent imagination evaluates compact transition and task heads. Pixel imagination must also produce every RGB value, although those values may never enter the planning objective.
3. Key takeaway: Latent-only rollout is a systems advantage, not a universal complexity theorem. Network depth, attention, decoder architecture, batching, precision, and activation storage still determine real latency.

---

Experiment 6: Stronger search can exploit a biased model¶

Question: Why can evaluating more action candidates increase imagined return while decreasing real return?

The learned reward model is accurate near the data-supported action range but becomes optimistically biased for aggressive actions. Larger search budgets are more likely to discover those errors.

In [7]:

np.random.seed(6)
action_grid = np.linspace(-0.2, 2.7, 700)
true_return = 2.2 * action_grid - 1.18 * action_grid**2
optimistic_bias = 1.65 * np.maximum(action_grid - 1.05, 0.0) ** 2
predicted_return = true_return + optimistic_bias
uncertainty = 0.05 + 1.05 * np.maximum(action_grid - 0.9, 0.0) ** 2
pessimistic_return = predicted_return - 1.15 * uncertainty

search_budgets = np.array([4, 8, 16, 32, 64, 128, 256, 512])
repeat_count = 300
selected_predicted = []
selected_real = []
selected_real_pessimistic = []
selected_actions = []

for budget in search_budgets:
    budget_predicted = []
    budget_real = []
    budget_real_pessimistic = []
    budget_actions = []
    for _ in range(repeat_count):
        candidates = np.random.uniform(0.0, 2.7, size=budget)
        candidate_true = 2.2 * candidates - 1.18 * candidates**2
        candidate_bias = 1.65 * np.maximum(candidates - 1.05, 0.0) ** 2
        candidate_predicted = candidate_true + candidate_bias
        candidate_uncertainty = 0.05 + 1.05 * np.maximum(candidates - 0.9, 0.0) ** 2
        optimistic_index = np.argmax(candidate_predicted)
        pessimistic_index = np.argmax(candidate_predicted - 1.15 * candidate_uncertainty)
        budget_predicted.append(candidate_predicted[optimistic_index])
        budget_real.append(candidate_true[optimistic_index])
        budget_real_pessimistic.append(candidate_true[pessimistic_index])
        budget_actions.append(candidates[optimistic_index])
    selected_predicted.append(np.mean(budget_predicted))
    selected_real.append(np.mean(budget_real))
    selected_real_pessimistic.append(np.mean(budget_real_pessimistic))
    selected_actions.append(np.mean(budget_actions))

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8))
fig.suptitle('Optimization pressure turns model optimism into a selected action', fontsize=14, fontweight='bold')

axes[0].plot(action_grid, true_return, color='black', linewidth=2.4, label='real return')
axes[0].plot(action_grid, predicted_return, color='#e76f51', linewidth=2.3, label='model-predicted return')
axes[0].plot(action_grid, pessimistic_return, color='#2a9d8f', linewidth=2.3, label='prediction - uncertainty penalty')
axes[0].axvspan(1.05, action_grid[-1], color='#f4a261', alpha=0.13, label='weak data support')
axes[0].set_title('Action objective inside the model')
axes[0].set_xlabel('Action magnitude')
axes[0].set_ylabel('Return estimate')
axes[0].legend(fontsize=8)

axes[1].plot(search_budgets, selected_predicted, color='#e76f51', marker='o', linewidth=2.2, label='imagined return selected by model')
axes[1].plot(search_budgets, selected_real, color='black', marker='o', linewidth=2.2, label='real return of selected action')
axes[1].plot(search_budgets, selected_real_pessimistic, color='#2a9d8f', marker='o', linewidth=2.2, label='real return with penalty')
axes[1].set_xscale('log', base=2)
axes[1].set_title('Search budget changes selected outcomes')
axes[1].set_xlabel('Number of action candidates')
axes[1].set_ylabel('Mean selected return')
axes[1].legend(fontsize=8)

axes[2].plot(search_budgets, selected_actions, color='#6a4c93', marker='o', linewidth=2.3, label='mean selected action')
axes[2].axhline(action_grid[np.argmax(true_return)], color='black', linestyle='--', linewidth=1.2, label='true optimal action')
axes[2].axhline(1.05, color='#f4a261', linestyle='--', linewidth=1.2, label='support boundary')
axes[2].set_xscale('log', base=2)
axes[2].set_title('Stronger search moves toward model error')
axes[2].set_xlabel('Number of action candidates')
axes[2].set_ylabel('Mean selected action magnitude')
axes[2].legend(fontsize=8)

plt.tight_layout()
plt.show()

np.random.seed(6) action_grid = np.linspace(-0.2, 2.7, 700) true_return = 2.2 * action_grid - 1.18 * action_grid**2 optimistic_bias = 1.65 * np.maximum(action_grid - 1.05, 0.0) ** 2 predicted_return = true_return + optimistic_bias uncertainty = 0.05 + 1.05 * np.maximum(action_grid - 0.9, 0.0) ** 2 pessimistic_return = predicted_return - 1.15 * uncertainty search_budgets = np.array([4, 8, 16, 32, 64, 128, 256, 512]) repeat_count = 300 selected_predicted = [] selected_real = [] selected_real_pessimistic = [] selected_actions = [] for budget in search_budgets: budget_predicted = [] budget_real = [] budget_real_pessimistic = [] budget_actions = [] for _ in range(repeat_count): candidates = np.random.uniform(0.0, 2.7, size=budget) candidate_true = 2.2 * candidates - 1.18 * candidates**2 candidate_bias = 1.65 * np.maximum(candidates - 1.05, 0.0) ** 2 candidate_predicted = candidate_true + candidate_bias candidate_uncertainty = 0.05 + 1.05 * np.maximum(candidates - 0.9, 0.0) ** 2 optimistic_index = np.argmax(candidate_predicted) pessimistic_index = np.argmax(candidate_predicted - 1.15 * candidate_uncertainty) budget_predicted.append(candidate_predicted[optimistic_index]) budget_real.append(candidate_true[optimistic_index]) budget_real_pessimistic.append(candidate_true[pessimistic_index]) budget_actions.append(candidates[optimistic_index]) selected_predicted.append(np.mean(budget_predicted)) selected_real.append(np.mean(budget_real)) selected_real_pessimistic.append(np.mean(budget_real_pessimistic)) selected_actions.append(np.mean(budget_actions)) fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) fig.suptitle('Optimization pressure turns model optimism into a selected action', fontsize=14, fontweight='bold') axes[0].plot(action_grid, true_return, color='black', linewidth=2.4, label='real return') axes[0].plot(action_grid, predicted_return, color='#e76f51', linewidth=2.3, label='model-predicted return') axes[0].plot(action_grid, pessimistic_return, color='#2a9d8f', linewidth=2.3, label='prediction - uncertainty penalty') axes[0].axvspan(1.05, action_grid[-1], color='#f4a261', alpha=0.13, label='weak data support') axes[0].set_title('Action objective inside the model') axes[0].set_xlabel('Action magnitude') axes[0].set_ylabel('Return estimate') axes[0].legend(fontsize=8) axes[1].plot(search_budgets, selected_predicted, color='#e76f51', marker='o', linewidth=2.2, label='imagined return selected by model') axes[1].plot(search_budgets, selected_real, color='black', marker='o', linewidth=2.2, label='real return of selected action') axes[1].plot(search_budgets, selected_real_pessimistic, color='#2a9d8f', marker='o', linewidth=2.2, label='real return with penalty') axes[1].set_xscale('log', base=2) axes[1].set_title('Search budget changes selected outcomes') axes[1].set_xlabel('Number of action candidates') axes[1].set_ylabel('Mean selected return') axes[1].legend(fontsize=8) axes[2].plot(search_budgets, selected_actions, color='#6a4c93', marker='o', linewidth=2.3, label='mean selected action') axes[2].axhline(action_grid[np.argmax(true_return)], color='black', linestyle='--', linewidth=1.2, label='true optimal action') axes[2].axhline(1.05, color='#f4a261', linestyle='--', linewidth=1.2, label='support boundary') axes[2].set_xscale('log', base=2) axes[2].set_title('Stronger search moves toward model error') axes[2].set_xlabel('Number of action candidates') axes[2].set_ylabel('Mean selected action magnitude') axes[2].legend(fontsize=8) plt.tight_layout() plt.show()

What you see¶

1. What the plots show: Larger candidate sets improve the model's selected objective but push actions toward the weak-support region, where their real return deteriorates. An uncertainty penalty keeps selected actions closer to the reliable region.
2. Why it looks this way: Search selects the maximum of noisy or biased estimates. Increasing optimization pressure makes it more likely to find systematic optimism, not merely genuinely better actions.
3. Key takeaway: Planning quality depends on model calibration under the planner's action distribution. More samples, CEM iterations, or actor updates can hurt when the world model is exploitable.

---

Summary¶

Experiment	Key insight
1. Teacher forcing	Local accuracy does not measure the distribution shift created by autonomous rollout.
2. Observation refresh	Receding-horizon correction bounds consecutive exposure to model error.
3. Particle imagination	Branches and nonlinear event probabilities disappear in a mean trajectory.
4. Continuation masks	Terminal semantics determine which imagined rewards and values are valid.
5. Latent-only compute	Skipping pixel decoding can greatly reduce inner-loop data and activation work.
6. Model exploitation	Stronger optimization can amplify optimistic model bias outside data support.

Connection to Latent-Anything¶

A rollout API should preserve initial-state source, posterior-to-prior switch time, actions or policy, candidate/particle/ensemble axes, continuation masks, RNG identity, predicted heads, uncertainty, and model version. predict, fixed-action rollout, stochastic branch, and objective evaluate should remain separate so dynamics semantics are not silently coupled to one planner.