Latent Imagination Horizon¶
Companion to research/10-latent-imagination-horizon.md.
The imagination horizon $H$ is the one knob every Tier-7 planner must set. This capstone notebook makes its trade-off explicit, then demonstrates the three recurring fixes — value bootstrap, short branched rollouts, and adaptive truncation — that let latent planners dodge both ends of the curve.
What we will see:
- The U-curve — return-estimate error decomposes into a truncation term (myopia, shrinks with $H$) and a compounding model-error term (grows with $H$); their sum is U-shaped with an optimal $H^\star$.
- Value bootstrap — adding a terminal value cancels the truncation term, flattening the curve so short horizons are accurate and the planner is robust to $H$.
- MBPO — at a fixed model-step budget, many short rollouts incur far less compounding error than a few long ones, decoupling rollout length from task horizon.
- Adaptive truncation — stopping when model uncertainty crosses a threshold matches the best fixed horizon without having to tune $H$ per task.
import numpy as np
import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
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,
})
print('Setup complete.')
Setup complete.
Experiment 1: the U-curve of horizon error¶
Question. Why is there an optimal horizon rather than "longer is always better"?
A model with reward bias $b_k = d\cdot k$ that compounds with step $k$, true reward $r_k=1$, discount $\gamma=0.97$, and a truncated estimate without bootstrap $\hat G_H=\sum_{k<H}\gamma^k\hat r_k$. We plot the two error components — truncation (omitted tail) and compounding (accumulated bias) — and their total.
gamma, d = 0.97, 0.05
G_true = 1.0 / (1 - gamma) # sum of gamma^k * 1
Hs = np.arange(1, 80)
def errors(H, d=d, v_err=None):
k = np.arange(H)
disc = gamma ** k
partial_true = disc.sum() # sum_{k<H} gamma^k r_k
truncation = G_true - partial_true # omitted tail (>=0), the myopia error
compounding = np.sum(disc * d * k) # sum_{k<H} gamma^k b_k, grows with H
# the two are distinct error sources; total error budget is their sum (they do not cancel)
if v_err is None:
total = truncation + compounding
else:
total = gamma ** H * v_err + compounding # value bootstrap replaces truncation with v_err
return truncation, compounding, total
trunc = np.array([errors(H)[0] for H in Hs])
comp = np.array([errors(H)[1] for H in Hs])
tot = np.array([errors(H)[2] for H in Hs])
Hstar = Hs[np.argmin(tot)]
print('optimal horizon H* (min total error):', Hstar)
optimal horizon H* (min total error): 20
fig, ax = plt.subplots(figsize=(8.6, 4.8))
ax.plot(Hs, trunc, '--', color='#ff7f0e', label='truncation error (myopia, shrinks with H)')
ax.plot(Hs, comp, '--', color='#9467bd', label='compounding model error (grows with H)')
ax.plot(Hs, tot, '-o', ms=2, color='#1f77b4', lw=2, label='|total error|')
ax.axvline(Hstar, color='gray', ls=':', lw=1.5, label=f'H* = {Hstar}')
ax.set_xlabel('imagination horizon H'); ax.set_ylabel('error in return estimate')
ax.set_title('The horizon U-curve: myopia vs compounding model error')
ax.legend()
plt.tight_layout(); plt.show()
What you see¶
The truncation error (orange) is large at short horizons — the estimate ignores everything past $H$ — and decays as $\gamma^H$. The compounding model error (purple) is near zero at short horizons and climbs as the biased rollout extends. Their sum (blue) is U-shaped, dipping to a minimum at an intermediate $H^\star$.
Why. Extending the horizon trades one error for another: each extra step recovers some omitted reward (less myopia) but also feeds the model another biased prediction (more compounding). The sweet spot is where the marginal gain from seeing further equals the marginal cost of model drift.
Takeaway. "Longer is better" is false — there is a genuine optimal horizon. Every Tier-7 planner lives on this curve, and the rest of the notebook is about not having to find $H^\star$ by hand.
Experiment 2: value bootstrap flattens the curve¶
Question. Dreamer says a learned value makes it robust to horizon length. How?
We add a terminal value $\gamma^H\hat v(z_H)$ to the estimate. With a reasonable value head (small fixed error $v_{\text{err}}$) the truncation term is replaced by the value's estimate of the tail — so myopia at short horizons largely disappears. We compare the total error with and without the bootstrap.
v_err = 0.5
tot_nobs = np.array([errors(H)[2] for H in Hs])
tot_bs = np.array([errors(H, v_err=v_err)[2] for H in Hs])
print('no bootstrap : min error %.2f at H=%d ; error at H=5 is %.2f'
% (tot_nobs.min(), Hs[np.argmin(tot_nobs)], tot_nobs[4]))
print('with bootstrap: min error %.2f at H=%d ; error at H=5 is %.2f'
% (tot_bs.min(), Hs[np.argmin(tot_bs)], tot_bs[4]))
no bootstrap : min error 24.58 at H=20 ; error at H=5 is 29.08 with bootstrap: min error 0.48 at H=1 ; error at H=5 is 0.89
fig, ax = plt.subplots(figsize=(8.6, 4.8))
ax.plot(Hs, tot_nobs, '-o', ms=2, color='#d62728', label='no value bootstrap (U-curve)')
ax.plot(Hs, tot_bs, '-o', ms=2, color='#1f77b4', label='with terminal value bootstrap')
ax.set_xlabel('imagination horizon H'); ax.set_ylabel('|total error|')
ax.set_title('A value bootstrap removes truncation error -> robust to horizon')
ax.legend()
plt.tight_layout(); plt.show()
What you see¶
Without a bootstrap (red) the error is high at short horizons (myopia) and only dips at the U's minimum. With a terminal value (blue) the short-horizon error is already low and the curve is flat over a wide range of $H$ — the planner performs well almost regardless of horizon, only degrading at large $H$ where compounding finally dominates.
Why. The bootstrap $\gamma^H\hat v(z_H)$ supplies the discounted tail the rollout omitted, so the truncation term is replaced by the value head's small error instead of the full missing return. Short rollouts become far-sighted, exactly the property Dreamer relies on to keep $H$ small.
Takeaway. Value bootstrap decouples foresight from horizon: you get long-horizon reward accounting from a short, low-compounding rollout. The cost is that error moves from the dynamics to the critic — a short horizon cannot rescue a bad value function.
Experiment 3: MBPO — many short rollouts beat a few long ones¶
Question. Given a fixed budget of model steps, is it better to roll out far or to roll out often?
Model one-step error compounds geometrically within a rollout: a rollout of length $L$ has mean per-step error that grows with $L$. At a fixed total budget $B = (\text{number of rollouts})\times L$, we compare the average error of the synthetic transitions generated, as a function of $L$. MBPO's claim: short rollouts from real states keep error low.
B = 6000 # total model steps budget
L_lip = 1.06 # >1: one-step error compounds along the rollout
eps = 0.02 # base one-step model error (at a real start state)
def mean_transition_error(L):
k = np.arange(1, L + 1)
# error after k model steps grows geometrically (compounding from an on-support real start)
err_k = eps * (L_lip ** k - 1) / (L_lip - 1)
return err_k.mean() # averaged over the steps actually generated
Ls = np.array([1, 2, 3, 5, 8, 12, 20, 35, 60])
mean_err = np.array([mean_transition_error(L) for L in Ls])
n_rollouts = B // Ls
for L, e, nr in zip(Ls, mean_err, n_rollouts):
print(f'rollout length L={L:2d}: {nr:4d} rollouts, mean transition error={e:.3f}')
rollout length L= 1: 6000 rollouts, mean transition error=0.020 rollout length L= 2: 3000 rollouts, mean transition error=0.031 rollout length L= 3: 2000 rollouts, mean transition error=0.042 rollout length L= 5: 1200 rollouts, mean transition error=0.065 rollout length L= 8: 750 rollouts, mean transition error=0.104 rollout length L=12: 500 rollouts, mean transition error=0.163 rollout length L=20: 300 rollouts, mean transition error=0.317 rollout length L=35: 171 rollouts, mean transition error=0.792 rollout length L=60: 100 rollouts, mean transition error=2.806
fig, axL = plt.subplots(figsize=(8.6, 4.8))
axL.plot(Ls, mean_err, '-o', color='#1f77b4', label='mean model error per transition')
axL.set_xlabel('rollout length L (budget fixed: B = rollouts x L)')
axL.set_ylabel('mean model error of generated data', color='#1f77b4')
axL.tick_params(axis='y', labelcolor='#1f77b4')
axR = axL.twinx()
axR.plot(Ls, n_rollouts, '-s', color='#2ca02c', label='number of rollouts (coverage)')
axR.set_ylabel('number of rollouts at fixed budget', color='#2ca02c')
axR.tick_params(axis='y', labelcolor='#2ca02c'); axR.grid(False)
axL.set_title('MBPO: short rollouts -> low compounding error AND more state coverage')
axL.legend(loc='upper left'); axR.legend(loc='upper right')
plt.tight_layout(); plt.show()
What you see¶
As rollout length $L$ grows, the mean model error of the synthetic transitions climbs (blue) because later steps compound error away from the real start state. At the same fixed budget, longer rollouts also mean fewer rollouts (green), so worse state coverage. Short rollouts win on both axes.
Why. Compounding is geometric in rollout length, so the last steps of a long rollout are the least reliable data you could train on. Branching many short rollouts from real replay states keeps every generated transition close to on-support, where the model is accurate — MBPO's "disentangle rollout length from task horizon".
Takeaway. You do not need long rollouts to plan for a long task. Many short, on-support rollouts plus a value function (for the long-horizon credit) is the robust recipe — the same short-horizon lesson as MPC and Dreamer, here for generating model data.
Experiment 4: adaptive truncation without tuning H¶
Question. $H^\star$ depends on the task and model. Can we avoid tuning it by stopping when the model becomes unreliable?
Each episode the model stays reliable until a random breakdown step $b$ (when the rollout leaves the data support), after which its bias spikes. A fixed horizon either stops too early (myopic) or rolls past $b$ (compounding). Adaptive truncation watches an uncertainty signal and stops at $b$, then bootstraps a value. We compare average error of fixed-$H$ vs adaptive across episodes.
gamma2, v_err2 = 0.97, 0.5
rng = np.random.default_rng(3)
n_ep = 4000
breakdowns = rng.integers(4, 25, n_ep) # model becomes unreliable at a random step
def episode_error(b, H, adaptive):
# model is RELIABLE before breakdown b (bias ~0), then bias spikes; a value bootstrap covers the tail
stop = min(b, H) if adaptive else H # adaptive stops exactly at the breakdown
k = np.arange(stop)
bias = np.where(k < b, 0.0, 0.30) # 0 while reliable, large after breakdown
compounding = np.sum(gamma2 ** k * bias)
tail = gamma2 ** stop * v_err2 # value bootstrap covers reward beyond `stop`
return compounding + tail
Hs2 = np.arange(1, 40)
fixed_err = np.array([np.mean([episode_error(b, H, adaptive=False) for b in breakdowns]) for H in Hs2])
adapt_err = np.mean([episode_error(b, H=999, adaptive=True) for b in breakdowns])
print('best fixed-H error: %.3f at H=%d' % (fixed_err.min(), Hs2[np.argmin(fixed_err)]))
print('adaptive-truncation error: %.3f (no H tuning)' % adapt_err)
best fixed-H error: 0.443 at H=4 adaptive-truncation error: 0.334 (no H tuning)
fig, ax = plt.subplots(figsize=(8.6, 4.8))
ax.plot(Hs2, fixed_err, '-o', ms=3, color='#d62728', label='fixed horizon H')
ax.axhline(adapt_err, color='#1f77b4', lw=2.5, label='adaptive truncation (uncertainty-based)')
ax.axvline(Hs2[np.argmin(fixed_err)], color='gray', ls=':', lw=1.2, label='best fixed H*')
ax.set_xlabel('fixed imagination horizon H'); ax.set_ylabel('average error across episodes')
ax.set_title('Adaptive truncation matches the best fixed horizon without tuning it')
ax.legend()
plt.tight_layout(); plt.show()
What you see¶
The fixed-horizon error (red) is U-shaped again: short $H$ is myopic-ish (value bootstrap softens it) and long $H$ rolls past the breakdown into high-bias territory. The best fixed $H$ sits at the dip — but you would have to search for it. Adaptive truncation (blue line) stops each rollout exactly at its own breakdown step and lands at or below that best fixed error, with no horizon to tune.
Why. The reliable horizon varies episode to episode; a single fixed $H$ cannot be right for all of them. Truncating on an uncertainty/support signal adapts per rollout, cutting off precisely when the model stops being trustworthy, while the value bootstrap covers the remaining tail.
Takeaway. The cleanest answer to "what horizon?" is often "however far the model stays reliable" — measured by uncertainty, not fixed in advance. Combined with a value bootstrap, adaptive truncation is the most robust point on the horizon trade-off, and a natural job for Layer A introspection to drive.
Summary / Key Takeaways¶
- U-curve (Exp 1) — return-estimate error is truncation (myopia, $\downarrow$ with $H$) plus compounding model error ($\uparrow$ with $H$); their sum has a genuine optimal horizon $H^\star$. Longer is not better.
- Value bootstrap (Exp 2) — a terminal value replaces the truncation term, flattening the curve so short horizons are accurate and the planner is robust to $H$ — at the cost of moving error onto the critic.
- MBPO (Exp 3) — at a fixed budget, many short on-support rollouts beat a few long ones on both compounding error and state coverage, decoupling rollout length from task horizon.
- Adaptive truncation (Exp 4) — stopping when model uncertainty spikes matches the best fixed horizon without tuning it; "roll out as far as the model stays reliable" + value bootstrap is the robust recipe.
This closes Tier 7 — Planning in Latent Space: from model-based vs model-free, through reward/value heads and the planners (MPC, CEM, MPPI, Dreamer, MuZero, MCTS), to the horizon that bounds them all — using the latent world model to plan, not just to represent.