Offline and Conservative Model-Based RL › Out-of-Distribution Actions and Model Uncertainty

This section defines what "out of distribution" means for a learned model and why the action dimension is usually where the trouble starts. It builds the machinery for quantifying epistemic uncertainty, then constructs a gridworld where the failure is visible and unambiguous.

A dataset of continuous-control trajectories often covers a thin tube in joint state-action space. Even millions of transitions may leave reachable $(s,a)$ pairs poorly sampled. A robot moving through a hallway traces a narrow ribbon inside a higher-dimensional joint space. Three useful views of extrapolation arise; they overlap rather than partition the space.

• State OOD. The query state lies outside, or in a poorly covered part of, the behaviour state's distribution. This can happen when the policy drives into a new configuration; exact equality with a logged continuous state is not the test.
• Action OOD. The state is familiar but the proposed action is poorly supported conditional on that state. This is subtler: driving logs may cover a speed well but contain little evidence about full braking at that speed.
• Joint OOD. The pair $(s,a)$ is missing or rare in the logged joint distribution, even when the state and action have each appeared separately. State OOD or action OOD also makes the pair OOD; the distinctive case is a novel pairing of familiar marginals.

Action OOD is the one people underestimate. Consider a behaviour policy that is nearly deterministic, as some demonstration datasets are. In a well-covered state, the marginal $\rho_\beta(s)$ may be large, but conditional on $s$ the action distribution is a narrow spike. The model has evidence for only a narrow range of actions at that state, even if it has many transitions. If an improvement step proposes an action 2 standard deviations away from the spike, the transition model may have little evidence about the consequence, but the value function still returns a number, and the optimiser may still trust it.

These views interact. In a low-dimensional example, a state far from logged states may be visibly unusual; in high dimensions, that judgment is harder. A state-only density model cannot detect a poorly supported action at a familiar state. A familiar-looking state and a familiar-looking action can also form a rare pair, so checking only $p(s)$ cannot establish joint state-action coverage. The relevant question for a transition model is whether the pair has enough nearby evidence.

In reconstructive latent world models from Part III: Representing Agents and Worlds, an encoder maps observations into a latent space and a decoder reconstructs observations. Other architectures, including MuZero, need not have a decoder. If planning drives a latent state far beyond training support, its predicted transitions and values can become unreliable even though the vectors remain numerically valid. A model may not raise an exception or produce a NaN when this happens, so checking only numerical validity will miss the coverage problem.

We will compare five approaches to estimating or bounding uncertainty. They differ in cost, calibration, and how badly they can fail.

A schematic decomposition makes the distinction concrete before the individual estimator families are described.

No single family dominates: the choice depends on whether you need speed, a defensible statistical statement, or robustness to shared model bias.

Ensembles and disagreement. Train $K$ dynamics models $\{\hat P_{\theta_1}, \ldots, \hat P_{\theta_K}\}$ with different random initialisations and different bootstrapped data subsets. Each predicts a next-state mean $\mu_k(s,a)$. Then

$$u(s,a) = \max_k \Big\| \mu_{\theta_k}(s,a) - \bar\mu(s,a)\Big\|_2, \qquad \bar\mu(s,a) = \frac{1}{K}\sum_{k=1}^{K}\mu_{\theta_k}(s,a).$$

where:

• $u(s,a)$: the disagreement-based uncertainty score at $(s,a)$
• $\mu_{\theta_k}(s,a)$: the next-state mean predicted by ensemble member $k$ with parameters $\theta_k$
• $\bar\mu(s,a)$: the mean prediction averaged over the $K$ members
• $\|\cdot\|_2$: the Euclidean norm

This is one possible disagreement score, not MOPO's practical penalty. MOPO uses the maximum Frobenius norm of a member's predicted Gaussian scale matrix across its probabilistic model ensemble, which its authors describe as a maximum standard-deviation score. Ensemble disagreement can be useful, but it is not a calibrated posterior uncertainty; members with shared architecture and data may agree even when they are all wrong. Nor does predictive variance automatically isolate epistemic uncertainty from environmental randomness. The disagreement captures how much the members differ from each other, not how far the ensemble as a whole is from the truth. If all members are biased in the same way, the disagreement stays small while the error stays large.

Count-based uncertainty. In discrete or discretisable spaces, uncertainty is a function of the visit count $n(s,a)$. A common form is

$$u(s,a) = \frac{\alpha}{\alpha + n(s,a)}, \qquad \text{or} \qquad u(s,a) = \frac{1}{\sqrt{n(s,a)+1}}.$$

where:

• $u(s,a)$: the uncertainty score at $(s,a)$
• $n(s,a)$: the number of times $(s,a)$ was visited in the dataset
• $\alpha$: a smoothing prior (set to 1 in our implementation below)

The second has the $1/\sqrt n$ scaling of many concentration bounds. It is not, without a specified prior and observation-noise variance, an exact posterior standard deviation. Count-based scores are cheap and interpretable when counts are meaningful; we use one below to isolate the penalty mechanism from neural-network training noise. Their weakness is that "count" requires a definition of "the same event," and in continuous state spaces that definition is a discretisation, which introduces its own choices.

Density and likelihood models. Fit $\hat\rho_\beta(s,a)$ with a normalising flow or a VAE, then use a quantity such as $-\log \hat\rho_\beta(s,a)$ as a proxy for low data density. It is not automatically a calibrated bound on transition-model error or a reliable test of support. BEAR instead uses behavior-action samples and an MMD constraint in policy improvement. FisherBRC learns a behavior density for its critic decomposition $Q_\theta(s,a)=O_\theta(s,a)+\log\hat\mu(a\mid s)$ and regularizes action gradients of the offset; neither method defines this density score as a world-model reward penalty. Likelihood is not the same as support: a generative model can give high likelihood to out-of-distribution inputs, as likelihood-based OOD studies demonstrate.

Bayesian neural networks and variational inference. Place a prior over model parameters and use a posterior predictive distribution. These methods introduce inference and calibration choices that differ from those of a bootstrap ensemble; no single family wins across tasks and model classes. PILCO's Gaussian-process dynamics, discussed in Part VIII: Decision-Centric Research Lineages, provide posterior uncertainty under the GP model assumptions. Naive exact GP fitting has cubic time cost in the number of training points, so that particular implementation does not scale cheaply to large datasets.

Spectral and Lipschitz-based bounds. A smooth learned model alone does not guarantee a small error outside the data. If the error function is known to be Lipschitz, or both true and learned dynamics have compatible known Lipschitz bounds, an error measured near a query point can be extended with a distance term. Such a result is deterministic under its assumptions, though its constants may be too loose to guide planning. The needed quantity is a bound on how quickly error can change, not merely on how quickly the learned prediction changes.

Now we build something concrete. The goal is not to reproduce a benchmark. It is to construct the smallest system in which we can see the mechanism with our own eyes, compute everything exactly, and check the answer against ground truth. Minimality matters here. Offline RL failures on real benchmarks are opaque: you see a poor return, but you cannot easily separate model error from coverage. In a small grid we can isolate those two factors directly.

The environment is a $7 \times 60$ grid. Row 3 is a corridor that runs the full width; the goal sits at the far right end of that row. The grid edge acts as a wall. The goal is absorbing with zero reward, and reaching it pays 1.0. The discount factor is $\gamma = 0.99$.

The behaviour policy is a random walk along row 3, moving right with probability 0.7 and left with probability 0.3. It never leaves the corridor. This is the crucial design choice: the dataset covers one row out of seven, and covers it densely. Every other cell in the grid has never been seen.

The leftward moves matter. With a purely rightward policy, each of the 59 pre-goal corridor cells would be visited once per episode, while the absorbing goal would still receive 20 logged tail visits. The 30% leftward action instead makes the walk revisit earlier cells, producing a non-uniform pre-goal visit distribution.

The left edge clips a left action back to the same cell, while the rightward bias eventually carries most episodes to the goal. Because episodes reset at the start and stop after a short tail of goal self-transitions, the logged counts are not a stationary distribution of an endless reflecting walk. They need not be uniform along the corridor; the heatmap below shows the actual finite-episode coverage.

The optimal policy is trivially "walk right 59 times," so the return from the start state is $\sum_{t=0}^{58} \gamma^t \cdot r_{t+1} = \gamma^{58} \approx 0.99^{58} \approx 0.558$. (The reward is paid on arrival at the goal, after 59 steps, which is step index 58 from the start.) So we can compute every quantity exactly and ask whether each method finds it.

The printed counts show 120 of 1,680 state-action pairs observed, about 7.1%. The 60 visited cells are exactly the ones in row 3. Every other cell is unobserved, and every method below has to decide what to do there. The 7.1% figure is worth pausing on: it is not an extreme number, and yet it is enough to make a naive planner fail completely, as we are about to see.

The model is deliberately simple: a maximum-likelihood empirical transition model with a uniform fallback for unobserved pairs. This is the tabular analogue of a neural dynamics model with no regularisation, and it has the property we want to study: off-support, it returns something arbitrary, determined by our fallback choice rather than by evidence.

$$\hat P(s'\mid s,a) = \begin{cases} \dfrac{n(s,a,s')}{n(s,a)}, & n(s,a) > 0,\\[2mm] \dfrac{1}{|\mathcal{S}|}, & n(s,a) = 0, \end{cases} \qquad \hat r(s,a) = \begin{cases} \dfrac{1}{n(s,a)}\displaystyle\sum_{i:(s_i,a_i)=(s,a)} r_i, & n(s,a) > 0,\\[2mm] 0, & n(s,a) = 0. \end{cases}$$

where:

• $\hat P(s'\mid s,a)$: the estimated transition probability from $(s,a)$ to $s'$
• $\hat r(s,a)$: the estimated expected reward at $(s,a)$
• $n(s,a)$: the number of times $(s,a)$ appears in the dataset
• $n(s,a,s')$: the number of times $s'$ followed $(s,a)$
• $|\mathcal{S}|$: the number of states, used as the uniform fallback denominator

The uniform fallback is a stand-in for extrapolation, not a claim that other approximators behave identically. A neural network may continue a learned trend; a random forest prediction reflects its selected leaves; a Gaussian-process posterior mean approaches its prior mean when covariance with the training inputs vanishes, as can happen far away under a decaying kernel. None is automatically correct off-support. The uniform choice makes the failure mode vivid: an unvisited transition can land anywhere in the grid, including directly on the goal.

We use a count-based uncertainty proxy because the state-action visit counts are exact and interpretable in this tabular toy. In a larger continuous-state system, an ensemble is one common alternative; the count score isolates the penalty mechanism from neural-network training noise.

$$u(s,a) = \frac{\alpha}{\alpha + n(s,a)}, \qquad \alpha = 1.$$

where:

• $u(s,a)$: the count-based uncertainty score
• $n(s,a)$: the visit count from the dataset
• $\alpha$: the prior pseudo-count, fixed at 1 here

If a pair has never been observed, $u = 1$. If it has been observed 999 times, $u = 1/1000 = 0.001$. This is monotone, bounded in $(0,1]$, and it shrinks at roughly $1/n$ in the well-sampled regime. The $\sqrt{1/n}$ version has a different decay rate, closer to common concentration-bound scaling; here we use $1/n$ only as a simple bounded heuristic. Choosing a bounded score also keeps the penalty comparable across state-action pairs: no single pair can swamp the sum with an arbitrarily large uncertainty value.

Two things to notice. Right-action uncertainty is maximal off the corridor, because those pairs were never observed, and lower where rightward actions appear in the logs. The up and down actions remain maximally uncertain even on the corridor; that distinction is in the underlying state-action counts, though this heatmap shows only the right-action slice. A familiar state does not make every action at that state well supported.

A third thing to notice is that cell counts are not uniform even within the corridor. Fifty of the 59 nonterminal corridor cells have between 700 and 800 visits in this seeded run, but the two cells immediately before the goal have only 600 and 428. The absorbing goal has 6000 visits because every one of the 300 episodes adds a tail of 20 goal self-transitions. That bright goal cell should not be mistaken for stronger coverage of the approach states, and cell counts alone still say nothing about unsupported actions at those states.

Now we plan. The generic recipe is value iteration on the learned model:

$$\hat Q(s,a) = \hat r(s,a) + \gamma\sum_{s'}\hat P(s'\mid s,a)\hat V(s'), \qquad \hat V(s) = \max_a \hat Q(s,a),$$

where:

• $\hat Q(s,a)$: the estimated action-value at $(s,a)$ under the learned model
• $\hat V(s)$: the estimated state value, the maximum of $\hat Q$ over actions
• $\hat r(s,a)$: the estimated reward
• $\hat P(s'\mid s,a)$: the estimated transition probability
• $\gamma$: the discount factor

iterated to convergence. This is the same fixed-point computation that underlies every model-based planner in Part VII: Planning and Agency, with a learned model in place of the true one. The reason it is dangerous offline is precisely that value iteration trusts the model everywhere: the $\max_a$ operator steps through all four actions at every state, including states the model has never seen, and it takes whatever action looks best according to estimates that may be arbitrary.