Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Interaction Data and Passive Observation › Worked Example: Coverage Determines Extrapolation Risk

Michael BrenndoerferJuly 3, 202610 min read

World Models HandbookAll chapters

Part 6 of 12Learning World Models at Scale

Chapter 33Interaction Data and Passive Observation › Worked Example: Coverage Determines Extrapolation Risk

Reading section 6 of 10 · Worked Example: Coverage Determines Extrapolation Risk · 10 minSection 6 of 10 · 10 min

Section 6 of 10 · 10 min read

Interaction Data and Passive Observation › Worked Example: Coverage Determines Extrapolation RiskOne section · Sections in this chapter6 / 10 ↓

Sections in this chapter10

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Let's make the coverage argument concrete with a tiny one-dimensional controlled system. We'll collect interaction data under a biased behavior policy and fit a deliberately misspecified linear dynamics model. The true system has a mild nonlinearity that is activated mainly where the behavior policy rarely goes. This makes the difference between interpolation and extrapolation visible. Then we'll demonstrate importance weighting as a partial remedy.

The environment has nonlinear controlled dynamics with additive Gaussian noise:

s_{t+1} = 0.9 \, s_t + 0.5 \, a_t + 0.35\,[a_t]_+^2 + \epsilon_t, \qquad \epsilon_t \sim \mathcal{N}(0, 0.05^2),

where:

$s_t \in \mathbb{R}$ : the scalar state at time $t$
$a_t \in [-1, 1]$ : the scalar action applied at time $t$
$[a_t]_+ = \max(a_t, 0)$ : the positive part of the action
$\epsilon_t \sim \mathcal{N}(0, 0.05^2)$ : independent zero-mean Gaussian process noise
$0.9$ : the state persistence coefficient
$0.5$ : the linear action coefficient
$0.35$ : the positive-action curvature coefficient

Negative actions follow a linear response; positive actions add a quadratic effect that a linear fitted model cannot represent exactly.

In[3]:

Code

import numpy as np
from scipy.stats import beta as beta_dist

rng = np.random.default_rng(0)

# Ground-truth dynamics
A_TRUE = 0.9
B_TRUE = 0.5
CURVATURE_TRUE = 0.35
NOISE_STD = 0.05

We define a rollout function and a behavior policy that prefers negative actions. The uniform target policy is sampled separately during evaluation.

In[4]:

Code

def behavior_biased(rng):
    # Beta distribution skewed toward negative actions, rescaled to [-1, 1].
    return rng.beta(1.5, 4.0) * 2.0 - 1.0


def rollout(policy, rng, n_episodes=200, horizon=20):
    states, actions, next_states = [], [], []
    for _ in range(n_episodes):
        s = rng.normal(scale=0.5)
        for _ in range(horizon):
            a = policy(rng)
            s_next = (
                A_TRUE * s
                + B_TRUE * a
                + CURVATURE_TRUE * max(a, 0.0) ** 2
                + rng.normal(scale=NOISE_STD)
            )
            states.append(s)
            actions.append(a)
            next_states.append(s_next)
            s = s_next
    return np.array(states), np.array(actions), np.array(next_states)

We collect $N=200\times20=4{,}000$ transitions under the biased behavior policy. The design matrix $X_b\in\mathbb{R}^{N\times3}$ has columns $s_t$ , $a_t$ , and $1$ ; the target vector $y_b\in\mathbb{R}^{N}$ contains $s_{t+1}$ ; and the fitted coefficient vector $w_b\in\mathbb{R}^{3}$ contains the state coefficient, action coefficient, and intercept. Ordinary least squares therefore computes the best linear approximation under the sampled behavior distribution.

In[5]:

Code

states_b, actions_b, next_states_b = rollout(behavior_biased, rng)

X_b = np.column_stack([states_b, actions_b, np.ones_like(states_b)])
w_b, *_ = np.linalg.lstsq(X_b, next_states_b, rcond=None)

Out[6]:

Console

Collected 4000 transitions under the biased behavior policy
Action range observed: [-1.00, +0.82]
Mean           action: -0.461
Fraction of actions below 0.0: 0.89

True dynamics:      s' = 0.90 * s + 0.50 * a + 0.35 * max(a, 0)^2
Fitted dynamics:    s' = 0.900 * s + 0.517 * a + 0.011

The fitted coefficients describe the best linear approximation in the region emphasized by the behavior data. Because positive actions are rare, the fit has little evidence for the omitted positive-action curvature. Let's bucket the realized error by action value. This seeded experiment demonstrates how a coverage gap interacts with model misspecification; coverage alone does not force a correctly specified model to have this exact error pattern.

In[7]:

Code

bins = np.linspace(-1.0, 1.0, 11)
centers = 0.5 * (bins[:-1] + bins[1:])
counts, _ = np.histogram(actions_b, bins=bins)
coverage_fractions = counts / counts.sum()

Now we evaluate on new one-step transitions. We resample states from the behavior occupancy but choose actions from a uniform target policy spanning the full action range. Holding the state distribution fixed isolates the action distribution shift used in the importance-weighting identity below.

In[8]:

Code

eval_rng = np.random.default_rng(11)
states_t = eval_rng.choice(states_b, size=len(states_b), replace=True)
actions_t = eval_rng.uniform(-1.0, 1.0, size=len(states_b))
next_states_t = (
    A_TRUE * states_t
    + B_TRUE * actions_t
    + CURVATURE_TRUE * np.maximum(actions_t, 0.0) ** 2
    + eval_rng.normal(scale=NOISE_STD, size=len(states_b))
)
X_t = np.column_stack([states_t, actions_t, np.ones_like(states_t)])
preds_t = X_t @ w_b
residuals_t = (preds_t - next_states_t) ** 2

mean_err_t = np.full(len(centers), np.nan)
for i, (lo, hi) in enumerate(zip(bins[:-1], bins[1:])):
    mask = (actions_t >= lo) & (actions_t < hi)
    if mask.sum() >= 5:
        mean_err_t[i] = residuals_t[mask].mean()

The visualization compares coverage (how many samples the behavior policy put in each action bucket) with prediction error under the target policy. The high-action region is the one the behavior policy avoided.

Out[9]:

Visualization

Bar chart showing training coverage concentrated on negative actions. — Behavior policy coverage by action value for the biased behavior policy. Coverage is concentrated on negative actions, with almost no samples above +0.5.

Out[10]:

Visualization

Line chart showing target-policy prediction error rising for positive actions. — One-step target-policy MSE by action value for the unweighted linear fit. Error rises sharply in the sparse positive-action region, where the omitted positive-action nonlinearity makes the linear model's extrapolation error visible.

In this seeded experiment, the largest realized errors occur in the poorly covered positive-action region, where the fitted linear model cannot represent the true curvature. The experiment does not prove that sparse coverage always causes high error. It shows that sparse coverage can hide misspecification and leave off-support accuracy weakly tested. The estimator can appear well-behaved on the observed data while remaining unreliable in the region the target policy will visit.

A partial remedy: importance weightingLink Copied

If the target policy is known and its conditional action distribution overlaps the behavior policy's support everywhere, action importance weighting can estimate a target-action loss while holding the behavior state occupancy fixed:

\mathbb{E}_{\substack{s\sim d^{\beta}(s),\,a\sim\pi(\cdot\mid s),\;s'\sim P(\cdot\mid s,a)}} \!\left[\ell(s,a,s')\right] = \mathbb{E}_{\substack{s\sim d^{\beta}(s),\,a\sim\beta(\cdot\mid s),\;s'\sim P(\cdot\mid s,a)}} \!\left[ \frac{\pi(a\mid s)}{\beta(a\mid s)}\,\ell(s,a,s') \right].

where:

$d^{\beta}(s)$ : the state-occupancy distribution under the behavior policy
$\pi(a\mid s)$ : the target policy's conditional action distribution
$\beta(a\mid s)$ : the behavior policy's conditional action distribution
$P(\cdot\mid s,a)$ : the environment transition kernel
$\ell(s,a,s')$ : the transition loss
$\pi(a\mid s)/\beta(a\mid s)$ : the action importance weight

The equality requires $\beta(a\mid s)>0$ wherever $\pi(a\mid s)>0$ , up to sets of measure zero, for states with positive $d^{\beta}(s)$ . It does not correct the separate shift from $d^{\beta}(s)$ to $d^{\pi}(s)$ . Correcting the full target-policy occupancy would require the joint ratio $d^{\pi}(s,a)/d^{\beta}(s,a)$ , which is generally unknown.

In[11]:

Code

def importance_weights(actions, target_pdf, behavior_pdf):
    # Ratio of target to behavior densities, clipped for numerical stability.
    num = target_pdf(actions)
    den = behavior_pdf(actions)
    return np.clip(num / np.maximum(den, 1e-6), 0.0, 20.0)


# Densities for our two policies (uniform target, shifted Beta behavior).
def target_pdf(a):
    return np.where((a >= -1.0) & (a <= 1.0), 0.5, 0.0)


def behavior_pdf(a):
    # Beta(1.5, 4.0) is supported on [0, 1]; our behavior maps it to [0, 1] -> [-1, 1].
    u = (a + 1.0) * 0.5
    pdf_u = beta_dist.pdf(u, 1.5, 4.0)
    return np.where((a >= -1.0) & (a <= 1.0), pdf_u * 0.5, 0.0)


rhos = importance_weights(actions_b, target_pdf, behavior_pdf)

The clipped sample weights form a vector $\rho\in\mathbb{R}^{N}$ . Multiplying each row of $X_b$ and the corresponding element of $y_b$ by $\sqrt{\rho_i}$ is equivalent to minimizing $\sum_i\rho_i(y_i-x_i^\top w)^2$ . Thus $X_w\in\mathbb{R}^{N\times3}$ , $y_w\in\mathbb{R}^{N}$ , and the weighted coefficient vector $w_{\mathrm{IW}}\in\mathbb{R}^{3}$ . Clipping stabilizes this finite-sample fit but means it no longer exactly optimizes the unclipped importance-sampling identity.

In[12]:

Code

W = np.sqrt(rhos)  # Weighted least squares uses sqrt of the sample weight.
X_w = X_b * W[:, None]
y_w = next_states_b * W
w_iw, *_ = np.linalg.lstsq(X_w, y_w, rcond=None)

We compare the target-policy errors of the unweighted and importance-weighted fits. Here $X_t\in\mathbb{R}^{N\times3}$ is a separate held-out design matrix whose states are resampled from the behavior occupancy and whose actions come from the uniform target policy.

In[13]:

Code

err_unweighted_t = np.mean((X_t @ w_b - next_states_t) ** 2)
err_weighted_t = np.mean((X_t @ w_iw - next_states_t) ** 2)

err_unweighted_b = np.mean((X_b @ w_b - next_states_b) ** 2)
err_weighted_b = np.mean((X_b @ w_iw - next_states_b) ** 2)

Out[14]:

Console

Mean squared error on held-out target-policy data
  Unweighted least squares:      0.01190
  Importance-weighted LS:        0.00695

Mean squared error on the biased behavior data
  Unweighted least squares:      0.00266
  Importance-weighted LS:        0.00305

For this fixed seed, action weighting trades behavior-distribution fit for lower error on the held-out target-policy data. That outcome is not guaranteed: clipping introduces bias, finite-sample variance can overwhelm the benefit, and the action ratio does not correct state-occupancy shift. If the target policy puts mass where the behavior policy has near-zero density, unclipped weights can become enormous and the estimator's variance can become unusable. This is why coverage is best treated as a design constraint on data collection, not as a problem that weighting is guaranteed to repair later.

The importance weights derived above trade fit on the behavior distribution for fit on the target distribution. To visualize that trade-off, we compute per-action errors for both fits on the same held-out target-policy sample; action buckets with fewer than five samples are omitted to reduce noise.

In[15]:

Code

# Compute one-step MSE for the unweighted and importance-weighted fits
# within each action bucket on the held-out target-policy sample.
weighted_preds_t = X_t @ w_iw
weighted_residuals_t = (weighted_preds_t - next_states_t) ** 2

mean_err_weighted_t = np.full(len(centers), np.nan)
for i, (lo, hi) in enumerate(zip(bins[:-1], bins[1:])):
    mask = (actions_t >= lo) & (actions_t < hi)
    if mask.sum() >= 5:
        mean_err_weighted_t[i] = weighted_residuals_t[mask].mean()

Out[16]:

Visualization

Unweighted and importance-weighted prediction error across action buckets. — One-step target-policy MSE by action bucket for the unweighted and importance-weighted linear fits. Weighting lowers error in the sparse positive-action region but does not eliminate the coverage gap.

The figure answers one narrow question: one-step squared error under a target action distribution while holding the behavior-state occupancy fixed. A real evaluation should keep four measurements separate:

held-out one-step error at target-relevant state-action pairs;
open-loop rollout error as predictions feed back over increasing horizons;
coverage and uncertainty diagnostics, including whether uncertainty rises on deliberately shifted inputs; and
closed-loop task performance and constraint violations when a planner uses the model.

Good one-step MSE does not imply good planning. A planner can repeatedly query a small biased region, exploit a localized model error, or amplify a small systematic error over a rollout. Conversely, a visually imperfect predictor can still preserve the task-relevant quantities needed for action selection.