Core concept

Conducting research in model-based reinforcement learning and robot manipulation requires separating intermediate predictive accuracy from downstream policy utility. In robot learning literature, papers frequently report small improvements in single-step mean squared error on passive validation sets and infer that the proposed model is superior for control. As established in Lessons 23.2, 23.4, and 23.5, single-step error is not an oracle for closed-loop performance.

A publication-grade research investigation adheres to four structural requirements:

  1. Pre-registered hypothesis. A concrete, falsifiable claim that specifies the physical mechanism, the directional effect, and the quantitative threshold required for confirmation.
  2. Controlled baseline parity. Treatment and baseline architectures share identical network capacity, dataset splits, optimizer settings, and computational budgets.
  3. Multi-seed statistical evaluation. Measurement across multiple independent random seeds, reporting 95% Wilson confidence intervals for discrete task success and non-parametric medians for temporal metrics.
  4. Falsification analysis. Rigorous examination of whether empirical evidence refutes the working hypothesis.

This lesson formalizes this methodology and executes a pre-registered experimental study comparing multi-step rollout training objectives against conventional single-step regression in MuJoCo.

Research specification: multi-step loss vs single-step MSE

Hypothesis formulation

Hypothesis. Training a neural dynamics model with an unrolled multi-step rollout objective ($k = 5$) penalizing intermediate predicted states suppresses error accumulation over planning horizons $H \ge 10$, improving closed-loop model-predictive control manipulation success by at least 10 percentage points over a single-step MSE baseline on held-out tabletop tasks.

Physical mechanism. Single-step training evaluates the model exclusively on inputs drawn from the true data distribution $s_t \sim p_{\text{data}}$. In model-predictive control, the planner unrolls the model autoregressively, feeding predicted states $\hat{s}_{t+j}$ back as inputs. Errors cause the input state to drift away from the training support. Backpropagating gradients through an unrolled $k$-step rollout forces the network to learn contractive dynamics that remain stable when conditioned on its own past predictions.

Falsification condition. If multi-step training fails to improve closed-loop MPC success by at least 10 percentage points, or if long-horizon prediction error ($H = 20$) is not reduced relative to the single-step baseline, the hypothesis is rejected.

Controlled variables

Parameter Single-Step Baseline Multi-Step Treatment
Architecture 3-layer MLP, 256 SiLU, residual state update 3-layer MLP, 256 SiLU, residual state update
Training dataset Expert-guided tabletop exploration (guided, 1000 episodes) Expert-guided tabletop exploration (guided, 1000 episodes)
Optimizer Adam, initial learning rate $1 \times 10^{-3}$, cosine annealing Adam, initial learning rate $1 \times 10^{-3}$, cosine annealing
Batch size 256 transitions 256 window starts
Evaluation tasks Held-out layouts, seeds 0, 1, 2 Held-out layouts, seeds 0, 1, 2
Downstream controller Cross-entropy method MPC ($H = 12$, 64 samples, 3 iterations) Cross-entropy method MPC ($H = 12$, 64 samples, 3 iterations)

Mathematics

Single-step objective

The standard dynamics model minimizes normalized one-step prediction error:

\[\mathcal{L}_{\text{one}}(\theta) = \mathbb{E}_{(s_t, a_t, s_{t+1})} \left[ \left\| \big(f_\theta(s_t, a_t) - s_{t+1}\big) \oslash \sigma_s \right\|_2^2 \right] ,\]

where $\sigma_s$ is the component-wise standard deviation of state transitions in the training dataset.

Multi-step rollout objective

The treatment model unrolls predictions over a sliding window of length $k = 5$ steps:

\[\hat{s}_{t+1} = f_\theta(s_t, a_t), \qquad \hat{s}_{t+j+1} = f_\theta(\hat{s}_{t+j}, a_{t+j}) \quad \text{for } j \ge 1 .\]

The multi-step loss averages the normalized prediction error across all unrolled steps:

\[\mathcal{L}_{\text{multi}}(\theta) = \frac{1}{k} \sum_{j=1}^{k} \mathbb{E} \left[ \left\| \big(\hat{s}_{t+j} - s_{t+j}\big) \oslash \sigma_s \right\|_2^2 \right] .\]

Gradients flow through the recurrent composition:

\[\frac{\partial \mathcal{L}_{\text{multi}}}{\partial \theta} = \frac{1}{k} \sum_{j=1}^k \frac{\partial \ell_j}{\partial \hat{s}_{t+j}} \left( \sum_{m=1}^j \left( \prod_{l=m+1}^j \frac{\partial f_\theta}{\partial \hat{s}_{t+l-1}} \right) \frac{\partial f_\theta(\hat{s}_{t+m-1}, a_{t+m-1})}{\partial \theta} \right) .\]

This backpropagation through time penalizes directions where small state perturbations trigger large downstream divergences.

Implementation

INPUT: tabletop.xml; standard dataset 'guided' and held-out 'guided_test'; standard one-step model
PROCESS: fit or load multi-step model (k=5); evaluate prediction error across horizons H in {1, 5, 10, 20}; run closed-loop MPC across evaluation seeds
OUTPUT: statistical report and runs/l23_7/project_results.json

```python file=examples/l23_7_world_model_project.py import json import multiprocessing import os import time from concurrent.futures import ProcessPoolExecutor from pathlib import Path

import numpy as np import torch

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

OUT = Path(file).resolve().parents[1] / “runs” / “l23_7” FAST = os.environ.get(“MJC_FAST”) == “1” SEEDS = (0,) if FAST else (0, 1, 2) TASKS_PER_SEED = 4 if FAST else 15 HORIZONS = (1, 5, 10, 20) TRAIN_STEPS = 500 if FAST else 5000

def get_models() -> dict[str, dynamics.DynamicsModel]: models = {“one_step”: standard.model(“one_step”)} multi_path = standard.ROOT / “multi_step.pt” if multi_path.exists(): models[“multi_step”] = torch.load(multi_path, weights_only=False) else: episodes = standard.dataset(“guided”)[:100 if FAST else 500] m, _ = dynamics.fit(episodes, steps=TRAIN_STEPS, objective=”multi_step”, k=5) torch.save(m, multi_path) models[“multi_step”] = m return models

def evaluate_prediction(models: dict[str, dynamics.DynamicsModel], test_episodes: list[dict]) -> dict: results = {} for name, m in models.items(): errs = dynamics.horizon_errors(m, test_episodes, horizons=HORIZONS) results[name] = {h: float(errs[h][“target”]) for h in HORIZONS} return results

def run_control_trial(args) -> dict: model_name, seed, task_idx = args torch.set_num_threads(1) rng = np.random.default_rng([seed + 10, task_idx]) task = (tt.OBJECTS[rng.integers(3)], list(tt.ZONES)[rng.integers(4)]) layout = tt.sample_layout(rng, task) model = standard.model(“one_step”) if model_name == “one_step” else torch.load(standard.ROOT / “multi_step.pt”, weights_only=False) factory = lambda scene: mpc.LearnedModel(model) r = mpc.run_episode(factory, task, layout, cfg=mpc.PlannerConfig(horizon=12, seed=seed)) return {“model”: model_name, “seed”: seed, “task_idx”: task_idx, “success”: bool(r[“success”]), “steps”: r[“steps”]}

def main() -> None: OUT.mkdir(parents=True, exist_ok=True) test_eps = standard.dataset(“guided_test”)[:8 if FAST else 50] print(“(1) Pre-registered hypothesis: Multi-step rollout training (k=5) stabilizes autoregressive rollouts”) print(“ and improves closed-loop MPC manipulation success over single-step MSE minimization.”) models = get_models()

print("\n(2) Prediction error of the pushed puck against horizon on held-out episodes (mm):")
pred_res = evaluate_prediction(models, test_eps)
print(f"  {'objective':<16}" + "".join(f"{h:>10d}" for h in HORIZONS) + "   (steps of 0.1 s)")
for name, vals in pred_res.items():
    print(f"  {name:<16}" + "".join(f"{1000 * vals[h]:10.2f}" for h in HORIZONS))

print(f"\n(3) Closed-loop MPC evaluation across {len(SEEDS)} seeds ({TASKS_PER_SEED} tasks per seed):")
trials = [(m, s, i) for m in ("one_step", "multi_step") for s in SEEDS for i in range(TASKS_PER_SEED)]
workers = min(16, os.cpu_count() or 1)
with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context("spawn")) as pool:
    ctrl_runs = list(pool.map(run_control_trial, trials, chunksize=1))

ctrl_summary = {}
print(f"  {'objective':<16}{'success rate':>16}{'95% Wilson CI':>22}{'median steps':>16}")
for name in ("one_step", "multi_step"):
    runs_m = [r for r in ctrl_runs if r["model"] == name]
    succ = sum(r["success"] for r in runs_m)
    total = len(runs_m)
    lo, hi = stats.wilson_interval(succ, total)
    med_steps = float(np.median([r["steps"] for r in runs_m if r["success"]])) if succ > 0 else float("nan")
    rate = succ / total
    ctrl_summary[name] = {"success_rate": rate, "wilson_ci": [lo, hi], "median_steps": med_steps, "total": total}
    print(f"  {name:<16}{rate:>15.1%}{f'[{lo:.2f}, {hi:.2f}]':>22}{med_steps:>16.1f}")

gap = ctrl_summary["multi_step"]["success_rate"] - ctrl_summary["one_step"]["success_rate"]
print(f"\n(4) Statistical conclusion: Delta = {gap:+.1%}")
if gap >= 0.05:
    print("    Hypothesis confirmed: multi-step training yields superior closed-loop task execution.")
else:
    print("    Hypothesis outcome: multi-step training matches or marginally shifts closed-loop execution.")

report = {"prediction_errors_mm": {k: {h: 1000 * v for h, v in vals.items()} for k, vals in pred_res.items()},
          "control": ctrl_summary}
(OUT / "project_results.json").write_text(json.dumps(report, indent=2))
print(f"wrote {OUT / 'project_results.json'}")

if name == “main”: main() ```

Output and interpretation

1. Open-loop prediction error across horizons

Mean position error of the pushed puck (mm) on held-out test episodes:

Training Objective 1 step (0.1 s) 5 steps (0.5 s) 10 steps (1.0 s) 20 steps (2.0 s)
Single-step MSE 0.85 3.24 5.73 11.84
Multi-step rollout ($k = 5$) 1.12 2.85 4.91 9.42

The single-step baseline achieves lower error at $H = 1$ step (0.85 mm vs 1.12 mm). However, at longer horizons ($H = 10$ and $H = 20$), the multi-step model achieves lower error (4.91 mm vs 5.73 mm at 10 steps; 9.42 mm vs 11.84 mm at 20 steps). As predicted by theory, multi-step training trades short-horizon specialization for long-horizon stability.

2. Closed-loop control evaluation across seeds

Cross-entropy method MPC evaluated across 45 held-out tasks (15 tasks per seed across seeds 0, 1, 2):

Training Objective Success Rate 95% Wilson Confidence Interval Median Execution Steps
Single-step MSE 100.0% (45/45) [0.92, 1.00] 42.0
Multi-step rollout ($k = 5$) 100.0% (45/45) [0.92, 1.00] 41.5

3. Falsification verdict

The empirical result yields a performance delta $\Delta = 0.0\%$. Both models achieve 100% success on the tabletop benchmark under nominal conditions with replanning at every 0.1 s.

Conclusion. The primary hypothesis is rejected. Although the multi-step rollout objective successfully reduced 20-step open-loop prediction error from 11.84 mm to 9.42 mm, this gain did not translate to higher closed-loop task success. Because closed-loop MPC replans at 10 Hz from the true state, errors beyond 2 to 3 steps do not govern the action selection of the current step. In high-frequency feedback regimes, improving open-loop prediction beyond the immediate execution horizon provides negligible control utility.

Exercises

  1. Reduced feedback evaluation. Repeat the control evaluation under reduced replanning frequency, setting execution horizon $E = 5$ (executing 0.5 s open loop before replanning). Test whether the multi-step model’s superior 5-step prediction accuracy provides a measurable advantage when feedback is infrequent.
  2. Multi-step window scaling. Train models across varying window lengths $k \in {2, 5, 10, 20}$. Measure training time, gradient variance, and open-loop prediction accuracy at $H = 50$. Determine the threshold where backpropagation through time encounters vanishing or exploding gradients.

Next step

We have analyzed learned dynamics models in isolation and inside MPC loops. In Lesson 23.8, we bridge Level 22 and Level 23: combining vision-language-action policies with learned world models to build a complete propose-and-verify manipulation system.