Learning in imagination

Real experience is slow and costly; imagined experience is cheap and endless. Dyna and Dreamer live off that gap — but every synthetic rehearsal is only as trustworthy as the model region it wanders into.

1 · Why learn from a model at all?

Here is the asymmetry that starts everything. One real robot grasp takes seconds and can shatter a mug; one imagined grasp takes a millisecond and breaks nothing. Computation is cheap and repeatable; interaction is slow, dangerous, or irreversible. A model-free learner can replay a stored real transition many times — but it can only ever learn from consequences that were actually collected. It cannot cheaply ask what several unexecuted action sequences would have done from the same grounded state. A world model extracts a reusable rule from many transitions and then applies it to new situations without paying the physical cost again.

Write a real transition as e_t = (o_t, a_t, r_t, o_t+1, d_t) — observation, action, the reward after it, next observation, and an episode-ending flag d_t ∈ {0,1}. A model-free update estimates a policy or value straight from stored e_t. A model-based learner also fits a transition distribution and outcome heads:

Term by term: θ are learned parameters; x_t is the compact belief inferred from history; p_θ is a distribution, not a single point; r_t is feedback after a_t; and c_t is continuation probability (near one mid-episode, zero at a true terminal). This is the same source→action→destination convention as Lesson 04. Once the distribution is good, sampling it yields imagined transitions that cost compute, not another collision.

One warning up front, because it prevents a common fantasy: the model creates no new evidence. It interpolates and recombines what the dataset already contains. Ten thousand imagined steps are not ten thousand independent real steps — they can cut variance and expose the optimizer to more action sequences, but they all inherit the same model assumptions and the same blind spots.

2 · Dyna is a schedule, not a network

Sutton’s Dyna gives the cleanest picture. A single real transition feeds three operations: improve the model, improve behavior from that real transition, and use the model for extra behavior updates.

K is the imagination ratio — model-generated updates per real step. At K = 0 Dyna is just a model-free learner; as K grows, one real experience is reused ever more aggressively. Where the model is accurate, reward information propagates far faster; where it’s wrong, the same error is injected over and over. Concretely:

Notice two things the loop is careful about. The model is trained on real replay, not recursively on its own hallucinations; and imagined rollouts start from anchored replay states so they don’t immediately leave data support. Modern methods vary how they represent state, propose actions, and compute the target — but this outer-loop / inner-loop split is a reliable debugging diagram.

3 · From pixels to a recurrent state-space model

Imagining future images is wasteful. A 256 × 256 RGB frame is 196,608 numbers, most of them appearance rather than decision-relevant change. So Dreamer-style agents imagine in a compact latent — a recurrent state-space model (RSSM) pairing deterministic memory h_t (predictable context) with a stochastic state z_t (what the past doesn’t determine):

Read in order: memory advances from the last memory, state, and action; when a real frame arrives the posterior combines it with memory to infer z_t; in imagination there is no future frame, so the agent samples the prior from h_t alone. Training aligns prior with posterior without letting the latent collapse to a constant. The full state handed to task heads is x_t = (h_t, z_t); from the destination state a decoder predicts o_t+1, a reward head r_t, and a continuation head c_t:

L_obs makes the latent explain the sensor; L_reward preserves task signal that may occupy few pixels; L_continue models termination; the KL aligns posterior and prior; β sets that pressure. Real variants add KL balancing, free bits, symlog targets, categorical latents — details that change optimization, not the division of labor: the posterior grounds state in evidence, and the prior is the simulator available between observations.

4 · Dreamer’s two loops, made explicit

“Dreamer learns in imagination” is true but hides that there are two coupled loops with different data. The world-model loop processes real replay sequences and updates θ. The behavior loop starts from posterior states inferred along those real sequences, then rolls the prior forward under the actor’s proposed actions and updates actor and critic from the predicted outcomes:

Start states matter. Begin each imagination at an arbitrary prior sample and you ask the model to simulate states it may never have learned; begin at replay posteriors and you get a grounded root. Horizon H matters too: longer rollouts propagate effects farther but let transition error compound. The actor is optimized by backpropagating through differentiable latent dynamics, by a policy-gradient estimator, or a mix; the critic summarizes consequences beyond the finite horizon. Pixel decoding stays in the outer loop as a representation signal, but no image is rendered per imagined step — that’s the compute win. If reconstruction is unnecessary or distracting, other self-supervised objectives can replace it, provided the latent still keeps the physical factors future tasks need.

5 · Return is the bridge from prediction to behavior

An actor needs one scalar reason to prefer an imagined future. Define discounted return from time t:

γ ∈ [0,1) discounts distant reward; each c_t gates terms after termination. A critic V_ψ(x_t) estimates expected future return under the current actor. Pure one-step bootstrapping has short model exposure but inherits critic bias; a full H-step imagined return uses more predicted rewards but more dynamics error. The λ-return blends them:

Every symbol earns its place: hats are model predictions; 1−λ weights stopping here and trusting the critic; λ weights following the next predicted transition; the recursion ends at the horizon with the critic. Small λ → lower-variance, more bootstrapped targets; large λ → more exposure to long model rollouts. Neither removes bias — it just chooses the source. The critic minimizes L_critic = E[(V_ψ(x_t) − stopgrad(G_t^λ))²]; the stop-gradient stops the target from chasing the critic. The actor maximizes J(φ) = E[Σ w_t(G_t^λ + η H[π_φ(·|x_t)])], where w_t carries cumulative continuation and discount, H is policy entropy, and η rewards exploration. For continuous actions, squashing or constraints keep outputs in the actuator’s range.

6 · Work the return: a robot near a fragile box

A robot’s latent encodes distance to a box and whether the gripper is aligned. It imagines three steps under fast approach then grasp, predicting rewards r̂₀ = 0, r̂₁ = 2, r̂₂ = 6 and continuations ĉ₀ = 1.0, ĉ₁ = 0.8, ĉ₂ = 1.0 — the middle value drops because fast motion might topple the box. Let γ = 0.9, λ = 0.5, terminal critic V(x₃) = 4, and V(x₁) = 5, V(x₂) = 5. Work backward. At the last imagined step, G₂^λ = 6 + 0.9 × 1 × 4 = 9.6. At step one:

The bracket is 2.5 + 4.8 = 7.3; times 0.72 is 5.256; plus reward two gives 7.256. At the root:

Now imagine a slow approach with lower immediate rewards but continuation near one, giving root return 5.9. The actor should prefer slow — and the continuation head is why: it prices failure into every downstream reward, not as bookkeeping but as a discount on the future.

But suppose the training data contained only slow approaches. The reward model may extrapolate that even higher speed yields a perfect grasp while underestimating breakage. Gradient ascent then pushes speed to the limit and reports return 12. The arithmetic is flawless; the premises are false; a real trial shatters the box. This is model exploitation, not a λ-return bug — and the fix is support-aware updates, uncertainty, more grounded data, or a hard constraint, never just a stronger optimizer.

7 · Model exploitation is a consequence of optimizing

Validation error averages over a dataset. Policy optimization does the opposite — it searches for actions with maximum predicted return. Write prediction error as ε(s,a) = Q̂(s,a) − Q(s,a); maximizing Q̂ = Q + ε favors both genuinely good actions and actions with positive error. So even zero-mean error becomes positively selected after a maximum. The actor is an adaptive adversary against its own simulator.

Separate three failure modes, since their fixes differ. Compounding error: small transition mistakes push later rollouts into unfamiliar states. Reward hacking: the actor finds a latent feature the reward head mistakes for success. Support extrapolation: the proposed state–action pairs were simply absent from real data. They co-occur but need different diagnostics. And picking one regularizer won’t save you, because the error can enter at different points — an ungrounded root, drift with depth, crossing action support, an exploited reward, a violated rule, or an uncorrected deployment. Derive defenses in that same order:

1. Ground the root with replay. Replay-rooted starts begin from posteriors that absorbed real observations — removing step-zero fantasy, though not training recovery from rare off-policy states. Diagnose by comparing error from replay roots vs perturbed roots.
2. Limit how long an unchecked chain runs. A short imagined horizon bootstraps before dynamics drift far — at the cost of trusting the critic more. Sweep horizon while logging transition error and terminal-value contribution.
3. Ask several plausible models where knowledge ends. Ensemble disagreement penalizes actions on which fitted models disagree — extra compute, and shared biases can fool them all, so validate disagreement against real error rather than assuming it is calibrated.
4. Keep the actor near observed actions. Behavior regularization limits divergence from replay actions, directly attacking support extrapolation — at the price of conservatism. Plot achieved improvement against distance from behavior instead of picking the penalty from offline return alone.
5. Price uncertainty inside the return. Optimize predicted return minus a risk coefficient times uncertainty, so a brilliant-but-unsupported future looks less attractive. Calibration is hard and excess pessimism stalls exploration, so tune the coefficient to consequence severity.
6. Refuse to trade safety rules for reward. A hard constraint or shield rejects actions violating trusted geometry, torque, or safety limits even under an optimistic reward — but the constraint itself must be reliable, so stress-test the shield separately.
7. Let reality answer the new policy’s questions. Online correction collects evidence exactly where the improved policy visits, turning exploitation failures into targeted updates — with a real cost and a need for safe deployment gates, so begin in shadow mode or with bounded action changes.

Together this shows the physical problem’s implementation shadow: because synthetic experience contains no new evidence, an imagination system needs a grounded root, a bounded interval, a measure of ignorance, a policy-support rule, independent constraints, and a controlled path back to reality. Drop any link and you get a recognizable diagnostic signature, not a mysterious drop in return. A practical pessimistic score is J_safe = E[Ĝ] − κ σ_ensemble(Ĝ), where κ converts epistemic uncertainty into a return penalty — not a universal guarantee (ensemble disagreement measures only variation inside the chosen model family), so known constraints stay explicit.

8 · Diagnostics that localize the failing loop

Don’t start from total episodic return — it’s the final symptom of a coupled model, critic, and actor. Instrument each boundary:

1. Posterior-vs-prior error by horizon. From a real posterior, roll the prior under recorded actions and compare predicted rewards, continuations, features, and decoded frames against the held-out trajectory — plotting error at each step, not an average.
2. Action sensitivity. Hold the start fixed and perturb actions. Predictions barely moving means the model ignores control; tiny changes causing huge rewards means gradients are exploitable.
3. Predicted-return calibration. Bucket real executions by imagined return and compare with achieved return. Optimistic top buckets are the dangerous ones, because policy improvement selects them.
4. Policy support. Measure actor likelihood under replay behavior, nearest-neighbor distance in latent state–action space, or ensemble disagreement — tracked through an imagined rollout.
5. Critic target decomposition. Log immediate rewards, continuation products, and terminal bootstrap separately. A value driven entirely by the terminal critic needs a different fix than a spurious reward spike.
6. Real-vs-imagined ablation. Compare the same learner at K = 0, short imagination, and long imagination. If more model usage monotonically hurts, inspect model bias before tuning the actor.
7. Counterfactual rollouts. From matched real roots, sweep safe actions against known monotonicity, conservation, and collision constraints — directly probing what the actor will optimize.

Also watch representation collapse, KL spikes, reward scale, gradient norms through time, continuation imbalance, and actor entropy. A continuation head trained on data with almost no terminals can predict “continue forever,” making imagined rewards effectively immortal; balanced sampling or calibrated losses may be required.

9 · Online and offline imagination have different contracts

Online, a newly improved policy eventually visits the states it imagines; surprises enter replay, the posterior grounds them, and the model repairs its policy-relevant errors — a corrective cycle (though unsafe proposals still need gates). Strictly offline, the fixed dataset never answers the actor’s new questions, so the inner loop can optimize past support with no reality check, and a high imagined return is just a hypothesis the data may be unable to confirm.

That changes design. Online systems can alternate conservative improvement with small real-data expansions. Offline systems usually need stronger behavior regularization, pessimistic values, support-aware constraints, or uncertainty penalties. One useful objective is J_offline(φ) = E[Ĝ] − α D(π_φ(·|s), π_b(·|s)), where π_b approximates the dataset behavior, D measures divergence, and α prices extrapolation. Large α hugs demonstrated actions; small α permits improvement but demands more of the model. And never call an offline model “safe because it never explores” — the environment is safe during training only because it’s absent; deployment risk can be higher if simulated optimization exploited untested regions. Use held-out policies, a trusted simulator, constraint stress tests, shadow mode, and bounded real trials before deployment.

10 · Choosing the imagination budget

The right horizon, batch size, and update ratio follow the bottleneck. If environment steps dominate cost and dynamics are smooth, spend more compute per real transition. If the world changes faster than the replay buffer, too much inner-loop optimization makes the learner stale. If latency only matters at deployment, heavy offline imagination amortizes into a fast actor. If goals change online, one amortized policy may be less flexible than the planner in Lesson 09.

Measure three currencies separately — environment steps, accelerator operations, and wall-clock time — because “sample efficient” can hide a hundred-fold compute increase. For a robot fleet, data diversity across machines can beat more gradient steps on one machine’s narrow replay. For an offline dataset, no online correction exists, so conservative support constraints become central. For safety-critical control, deploy the actor behind uncertainty thresholds, rate limits, a fallback controller, and shadow evaluation before its actions reach hardware. A useful staged recipe: first make one-step reward and continuation calibrated; then verify open-loop latent rollouts under recorded actions; then train a critic with the actor frozen; then enable short-horizon actor learning under action constraints; then compare imagined and real return rankings; only then expand horizon or ratio. It sacrifices some speed to make debugging causal.

11 · Dreamer, Dyna, and planning are relatives, not twins

Dyna is the architectural idea of interleaving real learning, model learning, and simulated updates. Dreamer is one modern realization: learn a latent RSSM and train an actor–critic on multi-step imagined latent trajectories, so that after training a single forward pass proposes an action. Model-predictive control instead optimizes actions at decision time. These hybridize freely — a learned actor can seed an MPC search, a planner can distill targets into a policy, and real planner failures can enrich replay. The central trade is where computation lives: actor learning pays at training time for fast deployment but compresses behavior for its training distribution; online planning pays repeatedly at deployment but adapts to a new goal, constraint, or observation. Both remain bounded by model validity.

Where this points next

Imagination has now trained a reusable decision rule. The alternative is to keep the current belief and goal online, generate candidate action sequences at decision time, and spend computation on the exact state in front of us. Lesson 09 derives that branch: random shooting and CEM for continuous trajectories, MPC for corrective replanning, MCTS for discrete branching, and MuZero for learning only the latent information search actually needs.

Interview prompts

• What is the exact difference between Dyna and ordinary experience replay? Replay reuses transitions that actually occurred. Dyna additionally fits a model and generates transitions that did not occur, then uses them for decision-learning updates. Its benefit and unique risk both come from that synthetic branch.
• Why does Dreamer need both posterior and prior latent distributions? The posterior uses the current observation to ground latent state during real training. The prior predicts state without a future observation, so it is the distribution available inside imagination. Their alignment teaches the prior to anticipate posterior states.
• Why can Dreamer learn behavior without decoding pixels inside every rollout? Reward, continuation, critic, and actor operate directly on the learned latent state. Pixel reconstruction can shape that state in the model-training loop, while imagined policy learning uses compact dynamics and task heads.
• What does λ control in a λ-return? It interpolates between earlier value bootstrapping and following more imagined rewards before bootstrapping. Larger λ depends more on multi-step model accuracy; smaller λ depends more on critic accuracy.
• Why is a policy an adversarial consumer of a learned model? It maximizes predicted return rather than sampling average data. Maximization preferentially selects actions with positive prediction error, so even unbiased validation error can become optimistic under the optimized policy.
• How would you tell reward-model exploitation from transition drift? Decompose predicted return and compare action-conditioned rollouts with held-out reality. An isolated reward spike in an otherwise plausible latent path implicates the reward head; growing state and outcome error with horizon implicates dynamics drift.
• Why do short rollouts help, and what new dependency do they create? They keep imagined states closer to grounded replay and limit compounding error. They terminate earlier with a critic bootstrap, so behavior becomes more sensitive to critic bias.
• When would online planning be preferable to a Dreamer actor? When goals or constraints change frequently, extra decision-time compute is available, and adapting search to the current belief matters more than a one-pass action. A hybrid can use the actor to seed the planner.