Policies Learned in Imagination › Regularizing Policies Against Model Errors

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.