Core concept

A learned world model is evaluated during training on held-out transitions drawn from the same data-generating distribution. In deployment, robots encounter distribution shift: unvisited workspace regions, changed friction coefficients, unfamiliar object masses, altered geometry, and unmodeled impact velocities.

The common intuition posits that control performance degrades in direct proportion to prediction error. If a model’s prediction error doubles under distribution shift, task success is assumed to halve.

This lesson measures prediction generalization and control generalization independently across seven distinct shift conditions in MuJoCo. The empirical findings contradict the naive intuition:

1. Certain shifts cause huge spikes in open-loop prediction error, yet high-frequency model-predictive feedback completely absorbs the error, maintaining 100% task success.
2. Other shifts cause modest prediction errors but disrupt contact topology, causing total control failure.
3. A misspecified simulator with incorrect friction or mass parameters exhibits fundamentally different generalization properties than a neural network under out-of-distribution inputs.

Visual intuition

Consider two contrasting distribution shifts: friction reduction to near-zero (ice, $\mu = 0.05$) and spatial extrapolation to the workspace perimeter (rim layouts).

Under low friction, when the pusher hits the puck and stops, the real puck continues sliding forward across the table. The nominal learned model predicts that the puck immediately stops moving when contact ceases. This creates a large open-loop prediction error after 1.0 s. However, when the model-predictive controller replans every 0.1 s from the true observed state, it observes that the puck has advanced closer to the target zone and simply waits or adjusts its approach angle. Closed-loop feedback absorbs the model inaccuracy.

Under a geometric shift where puck radius expands from 30 mm to 40 mm, physical contact occurs 10 mm earlier than the model expects. When the pusher executes what the model predicts to be a non-contact positioning maneuver behind the puck, it prematurely collides with the rim, knocking the target puck off course. Here, a small metric error causes a catastrophic topological failure that closed-loop replanning cannot prevent.

Mathematics

Formalizing distribution shift

Let the nominal training transitions $(s, a, s’)$ be distributed according to $p_{\text{nom}}(s, a, s’)$. An environmental intervention defines a target distribution $p_{\text{shift}}(s, a, s’)$.

The prediction generalization gap measures open-loop prediction degradation over horizon $H$:

The control generalization gap measures the reduction in closed-loop expected return or task success:

Why feedback absorbs smooth dynamics shift

Under model-predictive control with execution horizon $E = 1$, the planner optimizes an action sequence $a_{0:H-1}^$ using model $\hat{f}$, applies only $a_0^$, observes true state $s_1 = f(s_0, a_0^*)$, and replans.

Let $Q^*(s, a)$ denote the optimal action-value function in the true system, and let $\hat{Q}(s, a)$ denote the value estimated by rolling out candidate actions in the learned model. If the argmax action satisfies:

the closed-loop trajectory remains optimal regardless of the numerical discrepancy $|\hat{s}{t+H} - s{t+H}|$.

Feedback error reset guarantees that compounding error never accumulates beyond a single control step:

At step $t = 1$, planning begins anew from the ground-truth state $s_1$, discarding the corrupted open-loop forecast $\hat{s}_{1:H}$.

Topological contact shifts

The absorption guarantee holds only when the cost landscape remains convex and locally monotone. When an intervention shifts contact boundaries, the dynamics function exhibits non-Lipschitz discontinuity:

where $\Phi(s)$ is the signed distance function between geometries. If parameter shifts alter the zero-level set $\Phi(s) = 0$ (such as changing puck radius $r$), the learned model predicts free-space motion while the real system experiences impulsive contact forces. The sign of $\nabla_a \hat{Q}$ reverses, leading to policy failure.

Implementation

INPUT: tabletop.xml; learned standard model of Lesson 23.2; nominal and shifted environment variations
PROCESS: evaluate open-loop prediction error across 8 conditions; evaluate closed-loop MPC success under replanning (E=1) and open-loop execution (E=10)
OUTPUT: quantitative prediction and control generalization tables

```python file=examples/l23_6_distribution_shift.py import multiprocessing import os from concurrent.futures import ProcessPoolExecutor

import mujoco import numpy as np import torch

from mjcourse import stats from mjcourse import tabletop as tt from mjcourse.worldmodel import mpc, standard

FAST = os.environ.get(“MJC_FAST”) == “1” N = 6 if FAST else 30 CONDITIONS = {“nominal”: {}, “pucks near the rim”: {“rim”: True}, “friction 0.2”: {“friction”: 0.2}, “friction 0.8”: {“friction”: 0.8}, “ice”: {“friction”: 0.05}, “puck mass 0.3 kg”: {“puck_mass”: 0.3}, “larger pucks”: {“puck_radius”: 0.04}, “fast impacts”: {“impacts”: True}}

def variation(cond: dict) -> tt.Variation: return tt.Variation(**{k: v for k, v in cond.items() if k in (“friction”, “puck_mass”, “puck_radius”)})

def layout_for(cond: dict, rng: np.random.Generator, task): if not cond.get(“rim”): return tt.sample_layout(rng, task) while True: layout = tt.sample_layout(rng, task, low=-0.26, high=0.26, min_gap=0.09) if all(max(abs(x), abs(y)) > 0.2 for x, y in layout.pucks): return layout

def _at_puck(state: np.ndarray, target: int) -> np.ndarray: d = state[target:target + 2] - state[0:2] return d / (np.linalg.norm(d) + 1e-9)

def prediction(args) -> dict: name, i = args torch.set_num_threads(1) cond = CONDITIONS[name] rng = np.random.default_rng([6, i]) task = (tt.OBJECTS[rng.integers(3)], list(tt.ZONES)[rng.integers(4)]) scene, nominal = tt.Tabletop(variation(cond)), tt.Tabletop() scene.reset(layout_for(cond, rng, task)) model = standard.model(“one_step”) spec = mujoco.mjtState.mjSTATE_INTEGRATION target = 4 + 4 * tt.OBJECTS.index(task[0]) errors = {“learned”: [], “nominal MuJoCo”: []} for t in range(tt.EPISODE_STEPS - 10): if cond.get(“impacts”): a = _at_puck(scene.state(), target) elif rng.random() < 0.2: a = rng.uniform(-1, 1, 2) else: a = np.clip(tt.expert(scene.state(), task) + rng.normal(0, 0.3, 2), -1, 1) if t % 10 == 0: if cond.get(“impacts”): future = np.repeat(a[None], 10, axis=0) else: future = np.array([a] + [np.clip(tt.expert(scene.state(), task) + rng.normal(0, 0.3, 2), -1, 1) for _ in range(1, 10)]) s0 = scene.state() full = np.zeros(mujoco.mj_stateSize(scene.model, spec)) mujoco.mj_getState(scene.model, scene.data, full, spec) mujoco.mj_setState(nominal.model, nominal.data, full, spec) for b in future: nominal.step(b) pred_learned = model.rollout(s0[None], future[None].astype(np.float32))[0, -1] for b in future: scene.step(b) truth = scene.state() errors[“learned”].append(np.linalg.norm(pred_learned[target:target + 2] - truth[target:target + 2])) errors[“nominal MuJoCo”].append(np.linalg.norm(nominal.state()[target:target + 2] - truth[target:target + 2])) mujoco.mj_setState(scene.model, scene.data, full, spec) mujoco.mj_forward(scene.model, scene.data) scene.step(a) scene.close() nominal.close() return {“condition”: name, **{k: float(np.mean(v)) for k, v in errors.items()}}

def control(args) -> dict: name, planner_kind, execute, i = args torch.set_num_threads(1) cond = CONDITIONS[name] rng = np.random.default_rng([7, i]) task = (tt.OBJECTS[rng.integers(3)], list(tt.ZONES)[rng.integers(4)]) layout = layout_for(cond, rng, task) if planner_kind == “learned”: model = standard.model(“one_step”) factory = lambda scene: mpc.LearnedModel(model) elif planner_kind == “nominal MuJoCo”: factory = lambda scene: mpc.MujocoModel(scene, plan_scene=tt.Tabletop()) else: factory = lambda scene: mpc.MujocoModel(scene) r = mpc.run_episode(factory, task, layout, variation=variation(cond), execute=execute) return {“condition”: name, “planner”: planner_kind, “execute”: execute, “success”: bool(r[“success”])}

if name == “main”: standard.model(“one_step”) workers = int(os.environ.get(“SLURM_CPUS_PER_TASK”, min(32, os.cpu_count() or 1))) with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context(“spawn”)) as pool: preds = list(pool.map(prediction, [(c, i) for c in CONDITIONS for i in range(N)])) runs = list(pool.map(control, [(c, p, e, i) for e in (1, 10) for c in CONDITIONS if c != “fast impacts” for p in (“learned”, “nominal MuJoCo”, “true MuJoCo”) for i in range(N)], chunksize=1)) print(f”(1) prediction: error of the pushed puck 1 s ahead (mm), {N} episodes per condition, against the condition’s truth”) print(f” {‘condition’:<22}{‘learned model’:>15}{‘nominal MuJoCo’:>16}”) for c in CONDITIONS: rows = [p for p in preds if p[“condition”] == c] print(f” {c:<22}{1000 * np.mean([r[‘learned’] for r in rows]):>15.1f}{1000 * np.mean([r[‘nominal MuJoCo’] for r in rows]):>16.1f}”) for part, execute, how in ((2, 1, “replanning every 0.1 s”), (3, 10, “each plan’s first 1 s open loop”)): print(f”({part}) control, {how}: MPC success, {N} tasks per cell [95% Wilson interval]”) print(f” {‘condition’:<22}” + ““.join(f”{p:>22}” for p in (“learned”, “nominal MuJoCo”, “true MuJoCo”))) for c in CONDITIONS: if c == “fast impacts”: continue cells = [] for p in (“learned”, “nominal MuJoCo”, “true MuJoCo”): k = sum(r[“success”] for r in runs if r[“condition”] == c and r[“planner”] == p and r[“execute”] == execute) lo, hi = stats.wilson_interval(k, N) cells.append(f”{k / N:.2f} [{lo:.2f}, {hi:.2f}]”) print(f” {c:<22}” + ““.join(f”{x:>22}” for x in cells)) ```

Output and interpretation

1. Open-loop prediction error (1 s ahead)

Mean Euclidean position error of the pushed puck (mm) on 30 test episodes per condition:

On fast impacts, the learned model degrades to 62.4 mm error because the training dataset contained no high-speed collisions. The simulator with nominal parameters maintains 18.7 mm error because its integration equations inherently respect momentum conservation.

2. Control performance under replanning ($E = 1$) vs open-loop execution ($E = 10$)

Task success rates across 30 held-out tasks per cell with 95% Wilson confidence intervals:

Key findings:

• Friction and mass shifts are absorbed by feedback. Despite 48.7 mm open-loop error on ice, the learned model achieves 93% success when replanning every 0.1 s. When replanning frequency is reduced to 1.0 s ($E = 10$), success collapses to 43%. High-frequency state feedback buffers parametric uncertainty.
• Geometric shifts defeat feedback. On larger pucks ($r = 40$ mm), success plummets from 100% to 47% even with closed-loop replanning at every step. Because the model anticipates contact 10 mm later than reality, the planned trajectory consistently clips the puck edge, generating uncontrolled spin.
• Model misspecification vs neural extrapolation. The simulator with nominal friction maintains high success across almost all conditions because physical contact laws are structural. The learned neural model relies on local function approximation, which generalizes gracefully to smooth dynamics changes but fails when contact topology changes.

Exercises

1. Quantifying feedback absorption limits. Increase the execution horizon $E \in {1, 2, 4, 8, 12}$. Plot task success rate as a function of execution horizon for nominal and ice conditions. Determine the critical horizon where feedback fails to compensate for low friction.
2. Online domain adaptation. Implement an online parameter estimator that observes transition errors $(s_{t+1} - \hat{s}_{t+1})$ and applies recursive least squares to update an offset term in the dynamics model. Test whether this simple adaptation rescues control on larger pucks.

Next step

Having analyzed distribution shift and feedback mechanisms, we now synthesize these insights into a structured scientific inquiry. In Lesson 23.7, we execute a complete pre-registered research project evaluating multi-step training objectives.