Policies Learned in Imagination

The planning methods in Part VII: Planning and Agency so far have used the world model at decision time. Sampling-based planning evaluates candidate action sequences, differentiable planning optimizes them through model gradients, and search expands imagined futures before committing to an action. Their model calls sit on the decision-time critical path.

That arrangement is demanding. A planner sits on the critical path of every decision, so it must fit within the control loop's latency budget. At 50 Hz, the entire sensing-to-action cycle lasts twenty milliseconds; state estimation, communication, inference, and actuation leave only part of that period for planning. The model also has to be accurate in a very specific way: not merely good on average, but trustworthy along the candidate trajectories the planner considers. Constraints may rule out large regions of state space, but a badly calibrated reachable candidate can still make a poor plan look attractive. In the language of classical control, this is the model predictive control regime, where an explicit model is queried online, repeatedly, and where closed-loop quality depends on the model along the plans actually considered and executed.

This chapter takes a different route: learn a policy inside the model, then act with that policy rather than running a planner for each decision. The agent rolls the model forward from states grounded in experience and uses synthetic trajectories to train a policy network. At action-selection time a deployed policy can map its current state estimate to an action without a planning rollout. Online variants may still collect data and update the model and policy between decisions.

The shift is a shift in when the model is used, and that single change ripples through everything else. At training time we can afford to run the model millions of times, to backpropagate through it, to compare it against real data, and to revise it. At deployment time we run a small network once per step. The model has become a teacher rather than an oracle, and the expensive, fragile, high-variance computation has been moved off the critical path entirely.

This is the amortization bargain: spend model compute during training to obtain a reactive controller with fixed per-step network inference cost. Dreamer illustrates this route; MuZero and TD-MPC instead retain learned-model planning at decision time. We will compare those distinct designs in Part VIII: Decision-Centric Research Lineages. First we need the mechanics of policy learning in imagination.

There is a classical analogy. In linear-quadratic control, a known or identified linear system and a specified quadratic objective yield an offline Riccati solution and a cheap feedback law. Here a learned transition and reward predictor supply imagined trajectories, while gradient descent fits a nonlinear policy. The task reward need not itself be learned as a preference; its predictor is learned because the actor trains without querying the environment on each imagined step. The analogy concerns moving computation out of the action-selection loop, not identical assumptions or guarantees.

This definition fixes the trajectory distribution over which the policy gradient is averaged: the one produced by the policy and learned model. A rollout may resemble real experience on the training distribution, but its transition law is still the model's. That difference underlies the failure modes we will study.

When real interaction is costly, model rollouts can provide additional training signal without another environment step, although those rollouts still consume compute and memory. The hazard is that the policy trains against a guess about the world. It can exploit optimistic errors that are reachable under its policy class and discoverable by its optimizer. A policy that scores well in imagination and fails in reality can therefore be optimizing its stated training objective faithfully.

The objective $\hat J(\omega)$ below is defined under the learned model. If a reachable model error makes a poor action look attractive, optimization may favor it over an action that works in the environment. The optimizer can be doing its job while optimizing the wrong objective.

We will build latent rollouts and the actor-critic losses that consume them, then examine value expansion, synthetic replay data, and ways to limit model exploitation.

The first design question is what the rollout operates on. A generative pixel-level model can roll forward in observation space if it predicts action-conditioned future images and rewards. But decoding images can be expensive, while the critic may need only a compact, task-relevant summary. Reconstruction can still be useful supervision for learning that summary even when decoding is unnecessary during behavior learning.

This is where notions of model quality separate. Predictive fidelity asks whether the next-step prediction matches reality. Rollout fidelity asks whether predictions remain useful when fed back over many steps. Representation quality asks whether the latent retains information needed for control. Uncertainty quality asks whether the model signals its limitations. Decision usefulness asks whether a policy trained against it works in the real environment. These properties can diverge. Latent imagination avoids repeated decoding during behavior learning; it does not make the other properties optional.

Latent imagination sidesteps this. Building on the recurrent and transformer state-space architectures from Part V: World-Model Architectures, the world model provides an encoder, a latent transition, and a reward head. The rollout happens entirely inside the latent space:

$$\hat z_t = q_\phi(o_{\le t}, a_{<t}), \qquad \hat z_{t+1} \sim p_\theta(z_{t+1} \mid \hat z_t, a_t), \qquad \hat r_t = r_\psi(\hat z_t, a_t).$$

where:

• $o_{\le t}$: the observation history up to and including time $t$, which the encoder consumes to produce a latent
• $a_{<t}$: the action history before time $t$, included so the encoder can account for how the current observation was reached
• $q_\phi$: the encoder, parameterized by $\phi$, that maps observations and past actions into the latent state $\hat z_t$
• $p_\theta$: the latent transition model, parameterized by $\theta$, which defines a distribution over the next latent state given the current latent and action
• $r_\psi$: the reward head, parameterized by $\psi$, that predicts the scalar reward from a latent state and action
• $\hat z_t$: the inferred latent state at time $t$, intended to summarize the task-relevant past; whether it succeeds is empirical
• $\hat r_t$: the predicted reward at time $t$, used as the synthetic training signal for the actor and critic

The equations resemble a classical control pipeline. The encoder $q_\phi$ is a learned state estimator, though it need not produce a calibrated sufficient belief. The transition model $p_\theta$ predicts how that state evolves under control. The reward head $r_\psi$ predicts a task reward supplied by the environment; it is not itself a learned preference specification. The actor and critic play the roles of controller and value estimate. Under partial observability, history or a belief estimate is generally useful, although a reactive observation policy can suffice in special tasks.

No observation needs to be decoded during these imagined actor and critic updates. The actor and critic read $\hat z_t$, while the reward head scores $(\hat z_t,a_t)$. A decoder may still be important when learning the representation.

There are three practical reasons this matters. First, skipping a high-dimensional decoder can reduce rollout compute, depending on architecture and observation resolution. Second, gradients need not pass through a pixel decoder to reach the policy. Third, reconstruction can pressure a model to preserve visual detail irrelevant to action selection; reward and transition objectives can additionally shape the latent toward control.

The representation question is which information the controller needs. A latent can reconstruct fine visual detail yet omit a small cue that changes the right action; another can reconstruct imperfectly but retain that cue. Neither ordering is guaranteed by the loss alone. Reward and value heads can help shape what the actor extracts from the latent, while the model's training objective still determines what information reaches those heads.

The rollout is then a single vectorized loop. Starting from a batch of $B$ latent states, at each of $H$ steps:

• Sample an action $a_t \sim \pi_\omega(\cdot \mid \hat z_t)$, where $\pi_\omega$ is the policy parameterized by $\omega$.
• Step the transition model to get $\hat z_{t+1}$.
• Score the reward $\hat r_t$ and query the critic for $\hat v_t = v_\xi(\hat z_t)$, where $v_\xi$ is the value function parameterized by $\xi$.

Because the $B$ rollouts are independent, each step can be batched on a GPU. There are still $H$ sequential transition steps and arithmetic scaling with $BH$; batching reduces wall-clock overhead, not the amount of work. This is the concrete form of the amortization bargain: training is expensive in aggregate but parallelizable across starts.

After the loop, query the critic once more at the final latent to get a bootstrap value $\hat v_H$. With $\gamma=0.99$ and $H=20$, that terminal value receives weight $\gamma^{20}\approx0.818$. For an infinite stream of constant rewards, about 82% of the normalized discounted mass lies after the first 20 rewards. In a continuing task, omitting the bootstrap can therefore undervalue delayed payoffs.

The scale $1/(1-\gamma)=100$ describes the geometric discount tail, not a literal cutoff. The first 20 rewards contain $1-0.99^{20}\approx18.2\%$ of the normalized mass for a constant stream. A terminal value estimate, as in finite-horizon dynamic programming or receding-horizon control, represents the continuing value beyond the sampled prefix. Its usefulness depends on its accuracy.

Whether the loop is differentiable matters. If the latent transition is a continuous distribution, we can sample with the reparameterization trick, $\hat z_{t+1} = \mu_\theta(\hat z_t, a_t) + \sigma_\theta(\hat z_t, a_t) \odot \epsilon$ with $\epsilon \sim \mathcal{N}(0, I)$, where $\mu_\theta$ and $\sigma_\theta$ are the mean and standard deviation predicted by the latent transition model and $\odot$ denotes elementwise multiplication, and gradients flow straight through the model. If the latent is discrete, as in many categorical state-space models, the sample is not differentiable. The usual fixes are a straight-through estimator, which passes the gradient of a relaxed sample through a hard sample, or a gradient-free actor update that treats the dynamics as a black box. Each choice changes what the actor gradient estimator can do, which we take up next.

Differentiability determines whether an exact pathwise derivative through sampled transitions is available. Continuous reparameterized models support it directly. Discrete or discontinuous models may use a relaxation, straight-through approximation, smoothing, or score-function estimator; none is automatically robust to model error. The choice affects both gradient bias or variance and the model derivatives the actor can exploit.

If actor gradients could update an unconstrained reward or transition head, the optimizer could improve the reported outcome without improving behavior. That shortcut is a risk, not an inevitable unbounded trajectory: predictive losses, regularization, and head ranges can counter it. In this chapter's code the model is fitted first and frozen during actor learning, so the two objectives remain separate.

Once we can generate rollouts, we need a learning signal. The natural choice is actor-critic, because it gives us both a critic that can reason about the tail of the horizon and an actor that can be optimized directly.

Actor-critic is itself a classical idea, and it is helpful to place it in that lineage. In approximate dynamic programming, the value function is the object of interest: policy iteration alternates between evaluating a policy and improving it greedily with respect to the value function. Actor-critic methods combine those two steps into a single loop with two function approximators, one for the value and one for the policy. Inside a learned model, the loop is unchanged in structure, but every state, reward, and transition it uses is synthetic. This is the sense in which imagination-based learning is classical control with a learned simulator swapped in: the algorithm is familiar, the data is not.

Recall from Chapter 6 of Part II that a critic $v_\xi$ estimates the expected discounted return from a state. Inside imagination, the critic is trained by regression onto multi-step return targets computed from the imaginary rewards. The standard target is the $\lambda$-return, a geometric blend of $n$-step returns that trades bias against variance. Define the temporal-difference residual along the rollout as

$$\delta_t = \hat r_t + \gamma\, v_\xi(\hat z_{t+1}) - v_\xi(\hat z_t),$$

where:

• $\delta_t$: the temporal-difference residual at step $t$, measuring the surprise of the imagined reward relative to the critic's prediction
• $\hat r_t$: the imagined reward at step $t$
• $\gamma$: the discount factor
• $v_\xi(\hat z_{t+1})$: the critic's value estimate at the next latent
• $v_\xi(\hat z_t)$: the critic's value estimate at the current latent

with the boundary condition that $\hat z_H$ is the last latent and $v_\xi(\hat z_H)$ is the bootstrap value. The $\lambda$-return is then

$$\hat G_t^{\lambda} = v_\xi(\hat z_t) + \sum_{k=t}^{H-1} (\gamma\lambda)^{k-t}\,\delta_k,$$

where:

• $\hat G_t^{\lambda}$: the $\lambda$-return target at step $t$, blending $n$-step returns to balance bias and variance
• $v_\xi(\hat z_t)$: the critic's baseline value estimate at the current latent
• $\lambda$: the blending parameter between one-step and Monte Carlo returns
• $\delta_k$: the temporal-difference residual at future step $k$
• $H$: the imagination horizon

which can be computed with a single backward pass over the rollout using the recursion $\hat G_t^{\lambda} = \delta_t + \gamma\lambda\,\hat G_{t+1}^{\lambda} + v_\xi(\hat z_t) - \gamma\lambda\, v_\xi(\hat z_{t+1})$ or, more simply, by accumulating an auxiliary variable $A_t = \delta_t + \gamma\lambda A_{t+1}$ and returning $v_\xi(\hat z_t) + A_t$. The critic loss is the mean squared error between $v_\xi(\hat z_t)$ and the detached target:

$$\mathcal{L}_{\text{critic}}(\xi) = \mathbb{E}\left[\frac{1}{H}\sum_{t=0}^{H-1}\left(\hat G_t^{\lambda} - v_\xi(\hat z_t)\right)^2\right].$$

where:

• $\mathcal{L}_{\text{critic}}(\xi)$: the critic loss as a function of the value-network parameters $\xi$
• $\mathbb{E}[\cdot]$: the expectation over imagined rollouts
• $H$: the number of rollout steps
• $\hat G_t^{\lambda}$: the detached $\lambda$-return target at step $t$
• $v_\xi(\hat z_t)$: the critic's predicted value at the latent state $\hat z_t$

In this separate critic regression, the target is detached during each update. Other algorithms deliberately differentiate through targets, so detachment is a design choice rather than a universal rule.

Why use a $\lambda$-return rather than only a one-step or full-rollout target? A one-step target relies heavily on the critic's bootstrap and can inherit its bias. A full sampled return avoids intermediate value bootstraps but may have higher sampling variance; in a deterministic model, that variance need not be large. If the rollout ends before the task does, its terminal bootstrap still carries value error. The $\lambda$-return interpolates among these targets. Its useful bias-variance balance depends on critic accuracy, model accuracy, and stochasticity rather than on a universal ordering.

For the actor there are two fundamentally different gradient estimators, and the distinction is one of the most important ideas in this chapter.

The score-function estimator. Also called REINFORCE or the likelihood-ratio estimator, this treats the sampled action as a fixed constant and reweights its log-probability by how good the outcome turned out to be:

$$\nabla_\omega \hat J_H(\omega) = \mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\left(R_t^H-b(\hat z_t)\right)\nabla_\omega \log \pi_\omega(a_t\mid\hat z_t)\right], \qquad R_t^H=\sum_{k=t}^{H-1}\gamma^{k-t}\hat r_k.$$

where:

• $\nabla_\omega\hat J_H$: the gradient of the finite, $H$-step discounted imagined-reward objective
• $R_t^H$: sampled discounted reward-to-go from step $t$ through the truncated rollout
• $b(\hat z_t)$: a state-dependent baseline, often the critic, whose expectation against the action score is zero
• $\gamma^t$: the time weight required by the stated discounted objective
• $\log \pi_\omega(a_t \mid \hat z_t)$: the log-probability of the sampled action under the current policy

The likelihood-ratio identity is unbiased for this finite imagined objective when $R_t^H$ is sampled from the stated rollout distribution and the baseline depends on state, not the sampled action. A practical update may replace $R_t^H$ with a bootstrapped $\lambda$-return and use $\operatorname{sg}(\hat G_t^\lambda-v_\xi(\hat z_t))$ as an advantage; with an approximate critic and $\lambda<1$, that surrogate is generally biased. The estimator works with discrete actions and non-differentiable transitions. Original Dreamer used analytic gradients through latent dynamics; DreamerV3 uses a REINFORCE-style actor update.

For one sampled action, the score is weighted by how far its return lies above or below the baseline. The magnitude matters, but the estimator does not observe the derivative of reward with respect to an untried nearby action. Sampling supplies that information indirectly. A critic baseline reduces variance by subtracting predictable state value without changing the exact score identity.

The pathwise estimator. Also called the analytic or reparameterized gradient, this differentiates directly through the reward and dynamics:

$$\nabla_\omega \hat J_H(\omega) = \mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\left(\frac{\partial \hat r_\psi}{\partial a_t}\frac{d a_t}{d\omega}+\frac{\partial \hat r_\psi}{\partial \hat z_t}\frac{d\hat z_t}{d\omega}\right)\right].$$

Here $d a_t/d\omega$ and $d\hat z_t/d\omega$ are total derivatives through all earlier sampled actions and transitions. If the actor objective includes a terminal bootstrap, its derivative must be included too. On a smooth reparameterized model, pathwise gradients often have lower variance than score-function gradients, though long Jacobian products can be unstable. They estimate the gradient of the imagined objective without requiring the model to be correct; model accuracy is needed to identify that gradient with the real-environment gradient. Discrete or discontinuous models may require a relaxation, straight-through estimator, or score-function route.

The score-function estimator uses sampled outcomes; the pathwise estimator additionally uses local model derivatives. Those derivatives can provide efficient credit assignment, but can also be misleading outside the training distribution. Neither estimator removes model bias. Their relative variance and learning speed depend on the model, horizon, policy, and optimizer.

The two can also be mixed, but their weights and bootstrap treatment define the actual objective. A REINFORCE-style bootstrapped actor surrogate, as used in later Dreamer variants, is:

$$\mathcal{L}_{\text{actor}}(\omega) = -\,\mathbb{E}\left[\sum_{t=0}^{H-1}\gamma^t\operatorname{sg}\!\left(\hat G_t^{\lambda} - v_\xi(\hat z_t)\right)\log \pi_\omega(a_t \mid \hat z_t)\right] - \eta\,\mathbb{E}\left[\sum_{t=0}^{H-1}\mathcal{H}\big[\pi_\omega(\cdot \mid \hat z_t)\big]\right].$$

The entropy coefficient $\eta$ can discourage premature concentration. Its effect depends on reward and gradient scales. The toy actor below has a learned global log-standard-deviation, so it cannot modulate exploration by state. Its optional actor.dist(s).entropy() term measures the pre-squash Normal's entropy, a surrogate regularizer rather than the entropy of bounded actions after $\tanh$. The plotted experiments set $\eta=0$ to isolate the two gradient routes.

At exactly zero variance, a point-mass policy has no ordinary log-density score; very small variance can make score estimates ill-conditioned and leave too little exploration. Entropy can prevent premature concentration, but stochasticity does not by itself protect against model exploitation.

Two possible stability tools are target critics and return normalization. A slowly updated target critic can reduce target movement; a return scale estimated independently of the current actor batch can moderate gradient magnitude across reward scales. They are not universal requirements. Our toy implementation instead uses a fixed reward scale, bounded actions, a conservative learning rate, and gradient clipping; its exact score-function branch does not normalize the sampled advantage.

Bootstrapping with function approximation can destabilize value learning, especially with off-policy data. Target networks can damp one feedback loop, while scaling changes optimizer step sizes. Neither proves convergence, and scaling the advantage using statistics of the same sampled actions can change the exact score-gradient estimator. The code below keeps that distinction visible.

Finally, note what the actor is not doing. It is not maximizing the true return. The objective is

$$\hat J(\omega) = \mathbb{E}_{\hat p_\theta,\, \pi_\omega}\left[\sum_{t=0}^{H-1}\gamma^t \hat r_\psi(\hat z_t, a_t)\right],$$

where:

• $\hat J(\omega)$: the imagined return objective as a function of policy parameters $\omega$
• $\mathbb{E}_{\hat p_\theta,\, \pi_\omega}$: the expectation taken under the learned model $\hat p_\theta$ and the current policy $\pi_\omega$
• $\gamma$: the discount factor
• $\hat r_\psi(\hat z_t, a_t)$: the model's reward prediction at latent state $\hat z_t$ and action $a_t$

an expectation under the model. Everything that follows is about the gap between $\hat J(\omega)$ and the real $J(\omega)$.

It is worth stating the mismatch in the sharpest possible terms. The world model is trained to predict observations, rewards, and transitions, which is a prediction objective. The actor is trained to maximize return under the model, which is a control objective. These two objectives are not the same, and nothing in the training procedure guarantees that the model's predictions are accurate on the states the actor will visit. This is sometimes called objective mismatch, and it is the structural reason that improving the model's average prediction error does not automatically improve the policy. A model can become better at predicting the data it was trained on while becoming worse as a training environment, if the improvement comes from fitting incidental structure that the policy then exploits.

Training a policy in imagination is one way to use a model. Another family uses it to improve value targets or augment replay data. Which approach performs better depends on model accuracy, data regime, task, and compute budget.

The distinction between the two families is a distinction in what the model is asked to do. In pure imagination, the model is asked to be an environment: it supplies entire trajectories, and the policy is trained on them. In value expansion and synthetic experience, the model is asked to be a data amplifier: it supplies a small amount of extra information that improves an otherwise model-free update. The model is still central, but its errors are more contained, because the update is anchored to real data rather than floating free in imagination.

An early ancestor is Sutton's Dyna architecture from 1990, elaborated in 1991. It interleaves updates from real transitions with planning updates from a learned model. The methods below refine how imagined transitions or returns enter a learner; they are not all interchangeable with arbitrary model-free updates.

Dyna separates two questions: where an update's transition comes from, and how the learner uses it. Real and simulated transitions can feed compatible updates, provided the algorithm handles their sampling distribution and model error. Generating extra experience is useful only when those conditions hold.

Model-based value expansion (MVE) uses short model rollouts to build value targets from real replay starts. For a state-value target, the model starts at a replay state and the current policy chooses the first and subsequent imagined actions. An action-value target can instead anchor its first action to the recorded replay action; later imagined actions still come from the current policy, not recorded future actions. The original formulation evaluates reward with the known task reward function; when that is unavailable, a learned reward head adds another error source. The following state-value equation depicts that learned-reward variant:

$$y_t^{H} = \sum_{i=0}^{H-1}\gamma^i \hat r_\psi(\hat s_{t+i}, a_{t+i}) + \gamma^{H} v_\xi(\hat s_{t+H}).$$

where:

• $y_t^{H}$: the $H$-step value expansion target at time $t$
• $\gamma$: the discount factor
• $\hat r_\psi(\hat s_{t+i}, a_{t+i})$: the model's predicted reward at the imaginary state $\hat s_{t+i}$ under action $a_{t+i}$
• $\hat s_{t+i}$: the $i$-th imaginary state along the rollout, initialized from a real state $s_t$
• $a_{t+i}$: an action sampled from the policy at the imagined state; in this state-value target, that includes $a_t$
• $\gamma^{H} v_\xi(\hat s_{t+H})$: the bootstrapped value term, weighted by the discount over the whole horizon

The parameter $H$ controls how many predicted rewards and transitions precede the bootstrap. A longer expansion downweights the terminal critic by $\gamma^H$ but exposes the target to more model error. STEVE mixes targets from multiple horizons, using inverse-variance weights estimated from ensembles of dynamics, reward, and value functions; those weights do not correct unestimated bias.

MVE can improve a value target when short-horizon model errors are smaller than the critic errors they replace. It can also make the target worse when the model is inaccurate. There is no universal ordering of model error, critic bias, and Monte Carlo variance. The useful horizon depends on those quantities and their covariance under the policy's state-action distribution.

Synthetic experience takes the Dyna idea literally. Rather than only changing the loss target, generated transitions are written into replay and sampled alongside real ones. The crucial design question is where those trajectories start. Model-based policy optimization (MBPO) branches from real replay states, takes a short number of model steps under the current policy, and adds those synthetic transitions to training. Its rollout length is tuned to model quality and training phase rather than fixed by a universal range.

Branching fixes the starting state to one observed in real replay. It limits the number of consecutive model steps but does not guarantee that subsequent policy actions or imagined states remain on the data distribution. MBPO's analysis makes the tradeoff between rollout length and model error explicit; the useful length depends on model quality and policy shift.

The containment strategy is to reduce, not eliminate, compounding error. Every synthetic transition is only a few model steps removed from a real start. Its distribution can still move far from the data if policy actions or one-step predictions do. Mixing in real data provides an anchor, but the mixture and resulting value estimates must be evaluated rather than assumed reliable.

Two further details matter: the synthetic-to-real sampling ratio, and which policy chooses imagined actions. Too much synthetic data can fit the critic to the model's distribution while concealing real-environment error in the training loss. Current-policy actions may improve relevance to the policy being evaluated but move rollouts away from the behavior data. This tension recurs in Offline and Conservative Model-Based RL.

Training loss on the synthetic mixture cannot establish accuracy under real rollouts. A held-out real transition set can test one-step value or model predictions; closed-loop environment rollouts test policy performance. Neither metric replaces the other. The action distribution deserves the same scrutiny because the model is used at states selected by the current policy, not merely at recorded replay states.

There is one more reason to prefer short rollouts, and it is the reason we return to in the next section: a long synthetic rollout gives the policy more opportunities to steer into a region where the model is wrong.

For two discounted MDPs with the same state and action spaces, $0\le\gamma<1$, and bounded true rewards $|r(s,a)|\le R_{\max}$, let $\epsilon_r=\sup_{s,a}|r(s,a)-\hat r(s,a)|$ and $\epsilon_m=\sup_{s,a}\|T(\cdot\mid s,a)-\hat T(\cdot\mid s,a)\|_1$. For any policy $\pi$, including one selected using the learned model, Bellman unrolling gives the uniform bound

$$\big\|V^\pi - \hat V^\pi\big\|_\infty \le \frac{\epsilon_r}{1-\gamma} + \frac{\gamma\,\epsilon_m\,R_{\max}}{(1-\gamma)^2}.$$

where:

• $V^\pi$: the true value function of policy $\pi$ in the real environment
• $\hat V^\pi$: the value function of the same policy under the learned model
• $\|\cdot\|_\infty$: the supremum norm, i.e. the worst-case absolute difference over all states
• $\epsilon_r$: a bound on the per-step reward error $|r(s,a) - \hat r(s,a)|$
• $\epsilon_m$: a bound on the per-step transition error $\|T(\cdot\mid s,a) - \hat T(\cdot\mid s,a)\|_1$
• $T$ and $\hat T$: the true and learned transition distributions
• $\gamma$: the discount factor
• $R_{\max}$: the maximum possible absolute reward
• $(1-\gamma)^{-1}$ and $(1-\gamma)^{-2}$: geometric accumulation factors from the Bellman unrolling

For one policy, write $V^\pi-\hat V^\pi=\delta_r+\gamma\hat T^\pi(V^\pi-\hat V^\pi)+\gamma(T^\pi-\hat T^\pi)V^\pi$. A Markov transition operator is non-expansive in the supremum norm: $\|\hat T^\pi f\|_\infty\le\|f\|_\infty$. Together with $\|V^\pi\|_\infty\le R_{\max}/(1-\gamma)$, this yields the two terms above. The $L^1$ norm in $\epsilon_m$ measures distance between transition distributions; it is not the induced matrix norm of $\hat T$.

Reward error carries a factor $(1-\gamma)^{-1}$; transition error carries $(1-\gamma)^{-2}$ because it changes downstream state values as well as accumulating over time. That is worse horizon scaling, not proof that the transition term dominates: dominance also depends on $\epsilon_r$, $\epsilon_m$, $R_{\max}$, and units. This is an infinite-horizon bound in $\gamma$, not a finite-imagination-horizon bound in $H$. The unbounded quadratic reward of the toy point mass below does not satisfy its global $R_{\max}$ assumption without restricting the state domain.

The bound can be too loose to certify useful performance because it uses global worst-case errors and large discount factors. It still applies to a policy selected after seeing the model: the errors are uniform over all state-action pairs. What fails under policy optimization is the comforting inference from a small average prediction error on a fixed dataset. The learned policy may select reachable regions where optimistic error is much larger than that average.

The practical distinction is therefore between uniform and distribution-weighted evidence. A uniform bound covers the selected policy but is often vacuous. A test-set average can be precise yet irrelevant to the states that policy visits. Closed-loop evaluation supplies the missing decision-focused check.

The imagination gap, $\hat J(\omega)-J(\omega)$, is a useful diagnostic when the two returns use the same starting-state distribution, policy, horizon, and discount. It can be positive or negative: models can be optimistic or pessimistic. A persistent positive gap alongside improving imagined and worsening real return is consistent with exploitation, but its cause still needs investigation.

In particular, the sign of the gap is not a theorem. The horizon sweep below contains negative as well as positive gaps. Track its magnitude and sign, then inspect the trajectories and model errors responsible rather than interpreting a single scalar as proof of intent or mechanism.

Four families of remedies are in common use.

Limit how far the policy can look. Shortening the imagination horizon or terminating a rollout outside estimated data support can reduce compounding model error, though even one-step errors can be exploited. Lowering the discount also changes the task objective and may suppress delayed rewards. The tradeoff prepares the hierarchical methods in Hierarchical, Symbolic, Language, and Multi-Agent Planning.

Shorter rollouts bound the number of consecutive model predictions; they do not guarantee that the visited states are trustworthy. When reward depends on delayed consequences, truncation can damage credit assignment unless a reliable terminal value supplies the tail. Hierarchical skills can reduce decision depth, but their option-level models also need validation.

Penalize uncertainty. Train an ensemble with different seeds or bootstrap resamples. For a $D$-coordinate next state, scale each coordinate by a declared task scale $d_j>0$, then form a scalar disagreement such as $u(s,a)=\sqrt{D^{-1}\sum_{j=1}^{D}(\operatorname{std}_k[\hat T_{k,j}(s,a)]/d_j)^2}$. This $u$ is dimensionless, so $\lambda$ has reward units in the shaped reward:

$$\tilde r(s,a) = \hat r_\psi(s,a) - \lambda\, u(s,a).$$

where:

• $\tilde r(s,a)$: the pessimism-adjusted reward used for actor learning
• $\hat r_\psi(s,a)$: the model's predicted reward at state $s$ and action $a$
• $\lambda$: a coefficient in reward units that sets the cost of one unit of standardized disagreement
• $u(s,a)$: a nonnegative scalar summary of standardized ensemble transition disagreement

The policy pays a price where the ensemble disagrees. MOPO uses a related model-uncertainty reward penalty. MOReL instead constructs a pessimistic MDP that sends unknown state-action pairs to a low-reward absorbing state. Ensemble spread is a heuristic for epistemic uncertainty, not a calibrated guarantee: members can disagree in familiar regions or agree on a shared misspecification.

Epistemic uncertainty reflects limited evidence; model misspecification reflects a representational or training limitation that may persist with more data. An ensemble can expose some epistemic uncertainty, but agreement alone cannot rule out common-mode error. The toy penalty sweep below shows one setting where a stronger penalty helps modestly; it does not certify safety or replace targeted data collection when interaction is possible.

Stay close to the data. A behavior-cloning or KL penalty added to the actor loss keeps the learned policy near the policy that generated the training data:

$$\mathcal{L}(\omega) = \mathcal{L}_{\text{actor}}(\omega) + \beta\,\mathbb{E}\left[D_{\mathrm{KL}}\!\left[\pi_\omega(\cdot\mid \hat z_t)\,\big\|\,\pi_{\text{data}}(\cdot\mid \hat z_t)\right]\right].$$

This is one family of offline-RL regularizers. It discourages departures from high-density behavior actions where the model may be poorly supported. A finite penalty does not guarantee that the learned policy stays in empirical support, nor does it impose an absolute performance ceiling. Offline learning can use a fixed dataset without an online exploration phase.

The tradeoff is that stronger conservatism may prevent valuable departures from the behavior policy. Its useful strength depends on data coverage and model error; neither the KL term nor behavior cloning by itself makes a policy safe. The point is to make unsupported actions costly while checking real-environment performance whenever interaction is available.

Keep one foot in reality. Online systems such as Dreamer and MBPO alternate real interaction with model and behavior updates. New transitions can reveal errors along the current policy's trajectory; whether retraining repairs them depends on model capacity and optimization. This is the subject of Online, Offline, and Continual Learning.

Interleaving real interaction attacks distribution shift directly, but it costs environment samples and cannot guarantee a corrected model under misspecification. It is complementary to short rollouts and policy constraints rather than a replacement for their validation.

A final, less obvious regularizer is gradient control. If the actor is optimized with pathwise gradients, the gradient magnitude scales with the Jacobian of the model, and a model with a spuriously large local slope will produce a spuriously large update. Clipping the actor gradient norm and scaling the pathwise term by a constant are cheap insurance. And the stop-gradient rule from earlier is itself a regularizer: if the actor loss could touch the model, the model would drift toward being optimistic, which is the opposite of what we want.

Gradient clipping limits the norm of a single parameter update direction before the optimizer step; it does not bound model error or closed-loop state excursions. Adjusting a pathwise term's weight changes the optimization surrogate and may help or hurt. The score-function term is not categorically more robust to a wrong model, since it still uses model-generated outcomes.

Abstract warnings about exploitation are easy to nod along to and hard to internalize. A concrete example makes the failure visible.

Our environment is a point mass with drag moving on a line. The state is position and velocity, $s = (x, v)$, and the action $a$ is an acceleration. The dynamics are linear:

$$x_{t+1} = x_t + \Delta t\, v_t, \qquad v_{t+1} = \rho\, v_t + \Delta t\, a_t,$$

with $\Delta t = 0.1$ and $\rho = 0.99$. The reward penalizes distance from a target at $x^\star = 2.0$ and the squared control effort:

$$r(s, a) = -(x - x^\star)^2 - 0.05\,a^2.$$

The maximum per-step reward is $0$ when $x=x^\star$ and $a=0$, regardless of velocity. Position, velocity, action, and reward are normalized toy quantities rather than physical SI measurements. Evaluation runs for 40 steps from states near rest at the origin.

The system is linear, low-dimensional, and fully observed. That removes representation learning and partial observability from this particular demonstration. The learned MLP is still an approximation, and optimizer behavior can interact with its errors. We will measure both prediction and closed-loop control rather than attribute every outcome to coverage alone.

The comparison uses two world models trained on different datasets.

• Model A is trained on data collected from initial states near the origin, with position and velocity independently drawn from zero-mean normals with standard deviation $0.25$, and with small actions. It has essentially never seen the region around $x = 2.0$.
• Model B is trained on data collected from initial states spread uniformly over $[-3, 3]$ with wider actions. It has seen the goal region many times.

Both models have the same architecture, transition count, loss, and optimizer. The treatment is the data collection distribution: Model A uses concentrated starts and smaller random actions; Model B uses wider starts and actions. This changes several state-action marginals together, not one isolated scalar.

The comparison can show how those two collection regimes affect learned predictions and downstream policies, but it cannot prove that goal-state coverage alone caused every behavioral difference. Training randomness, action coverage, and the actor optimizer also matter. A small seed check below keeps one attractive trajectory from standing in for a general result.

Model A's reward prediction near and beyond the goal relies much more on extrapolation. The learned policy may therefore see less cost to overshooting than the true reward assigns. Model B has direct samples around the goal, but broader data is not a guarantee of a good controller: the actor still needs a stable objective, bounded actions, and closed-loop evaluation.

In the reproducible run below, the narrow-model policy overshoots the target more than the broad-model policy and has a positive imagination gap. Both overshoot to some extent. That is evidence for this particular setup, not a claim that all narrow models fail or all broad models succeed.

We track the imagination gap: the return the learned model estimates for the policy minus its return in the real environment.

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.