Model-predictive control turns any model of the dynamics into a controller. At every step it samples candidate action sequences, predicts what each would do, scores the predicted futures, applies the first action of the best one, observes the real result, and plans again. The planner does not care whether its predictions come from MuJoCo or from a network; it needs a function from (state, action sequence) to predicted states.
That makes MPC the cleanest way to measure what a learned world model is worth for control. Run the same planner with the same cost and the same random seeds, once with MuJoCo as its model and once with the learned model, on the same tasks. Any difference is the model’s. This lesson does that on the tabletop. Then it opens individual decisions, to see how the model’s errors change which plan gets chosen.
Each panel of the figure shows one decision taken while the pusher is touching the puck. In grey are the pusher paths of the 64 candidate plans of the planner’s last iteration; in orange, the pushed puck’s paths; in bold, the plan the learned model chose. On the left, as the learned model predicts them; on the right, what MuJoCo says the same action sequences would do, from the same state. The clouds have the same shape. The puck paths differ by millimetres to centimetres, and the chosen plan ends a few millimetres from where the model expected.
The cross-entropy method (CEM) searches over sequences of $H = 12$ actions (1.2 s). Each iteration samples $K = 64$ sequences from a Gaussian around the current mean, predicts their states with the model, and refits the mean and spread to the best 16; three iterations per step. The cost of a predicted trajectory $\hat s_{1..H}$ is
\[c = \frac1H\sum_{k=1}^{H} w_k\Big(\big\|\hat p_k - g\big\| + 0.5\,\big\|\hat u_k - b(\hat p_k)\big\|\Big) + 0.01\sum_k\|a_k\|^2,\]with $\hat p_k$ the target puck’s predicted position, $g$ the zone’s centre, $\hat u_k$ the pusher’s, $b(p) = p - r\,(g - p)/|g - p|$ the point behind the puck on the line to the zone, and weights $w_k$ rising from 0.5 to 1.5. The second term says where to push from; it says nothing about how pucks move, which only the model supplies. After applying the first action, the mean is shifted by one step to warm-start the next plan. The planner is a standard sampling-based MPC (see Howell et al., arXiv:2212.00541, for the predictive-sampling variant used in Lesson 21.4).
Output (1), 50 held-out tasks, each run once per model:
| Model | Success [95% Wilson] | Median time to success | Planning time per step |
|---|---|---|---|
| MuJoCo (ground truth, from the full simulator state) | 50/50 [0.93, 1.00] | 5.5 s | 719 ms |
| learned model of Lesson 23.2 (from the 16-number state) | 50/50 [0.93, 1.00] | 5.8 s | 10 ms |
(Planning time on one core of a compute node, with 32 episodes running side by side.) The learned model loses nothing in success and 0.3 s in median time, and plans 72 times faster: 64 batched forward passes of a small MLP against 64 MuJoCo rollouts of 240 physics steps each. Both solved exactly the same tasks.
The learned model’s open-loop error after 1.2 s is several millimetres on the pushed puck (Lesson 23.2). Why does it not matter? The planner never commits to a prediction for more than 0.1 s. It replans from the measured state, so a prediction only has to be right enough to choose the right first action.
Output (2) measures that directly. In 10 learned-model episodes, at every one of 555 decisions, the 64 candidates of the planner’s last iteration were also simulated in MuJoCo from the same state:
| Decisions taken with | Number | Rank correlation of predicted and true costs (median; 10th percentile) | Chosen candidate truly the best | Chosen candidate in the true top 5 |
|---|---|---|---|---|
| the pusher not touching a puck | 314 | 0.99; 0.94 | 1.00 | 1.00 |
| the pusher touching a puck | 241 | 0.91; 0.72 | 0.93 | 0.99 |
Away from contact the learned model orders the candidates almost exactly as MuJoCo does. During contact its ordering degrades: one decision in ten has a rank correlation below 0.72, and 7% of the time it picks a candidate that MuJoCo ranks below the best. But in 99% of contact decisions the pick is among MuJoCo’s top five. That is close enough, because the next replanning step, 0.1 s later, corrects the choice. A planner needs a model that ranks futures, not one that predicts them exactly.
[!recommendation] Comparing a learned model with a simulator in MPC Keep everything but the model fixed: planner, cost, samples, seeds, tasks. Report success with intervals and the time per step. Then look inside decisions: compare the model’s ranking of the candidates with the simulator’s, split by regime. Success alone hides how close to failure the model runs.
INPUT: tabletop.xml; the learned model of Lesson 23.2
PROCESS: the same CEM planner with MuJoCo and with the learned model on 50 tasks; candidate rankings at 555 decisions
OUTPUT: tables; runs/l23_4/candidates.png; runs/l23_4/candidates.json
```python file=examples/l23_4_planning.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 mpc, standard
OUT = Path(file).resolve().parents[1] / “runs” / “l23_4” FAST = os.environ.get(“MJC_FAST”) == “1” TASKS, PROBES = (6, 2) if FAST else (50, 10)
def task_and_layout(i: int): rng = np.random.default_rng([4, i]) task = (tt.OBJECTS[rng.integers(3)], list(tt.ZONES)[rng.integers(4)]) return task, tt.sample_layout(rng, task)
def episode(args) -> dict: kind, i = args torch.set_num_threads(1) task, layout = task_and_layout(i) if kind == “MuJoCo”: factory = lambda scene: mpc.MujocoModel(scene) else: model = standard.model(“one_step”) factory = lambda scene: mpc.LearnedModel(model) start = time.perf_counter() r = mpc.run_episode(factory, task, layout) return {“kind”: kind, “i”: i, “success”: bool(r[“success”]), “steps”: r[“steps”], “seconds_per_step”: (time.perf_counter() - start) / r[“steps”]}
def probe(i: int) -> dict: torch.set_num_threads(1) task, layout = task_and_layout(i) scene = tt.Tabletop() scene.reset(layout) planner = mpc.Planner(mpc.LearnedModel(standard.model(“one_step”)), task) truth = mpc.MujocoModel(scene) decisions, kept = [], None for _ in range(tt.EPISODE_STEPS): state = scene.state() a = planner.act(state) actions, cost_pred, chosen = planner.last[“actions”], planner.last[“costs”], planner.last[“best”] true_states = truth.predict(state, actions) cost_true = mpc.cost(true_states, task, actions) gap = min(np.linalg.norm(state[:2] - state[4 + 4 * k:6 + 4 * k]) for k in range(3)) contact = gap < tt.PUCK_RADIUS + tt.PUSHER_RADIUS + 0.003 ranks = (np.argsort(np.argsort(cost_pred)), np.argsort(np.argsort(cost_true))) decisions.append({“contact”: bool(contact), “spearman”: float(np.corrcoef(*ranks)[0, 1]), “true_rank”: int((cost_true < cost_true[chosen]).sum()), “regret”: float(cost_true[chosen] - cost_true.min())}) if contact and kept is None: kept = {“task”: list(task), “predicted”: planner.last[“states”].tolist(), “true”: true_states.tolist(), “cost_pred”: cost_pred.tolist(), “cost_true”: cost_true.tolist(), “chosen”: int(chosen)} scene.step(a) if scene.success(task): break scene.close() return {“i”: i, “decisions”: decisions, “example”: kept}
def figure(p: dict) -> None: import matplotlib matplotlib.use(“Agg”) import matplotlib.pyplot as plt
i = 4 + 4 * tt.OBJECTS.index(p["task"][0])
pred, true = np.array(p["predicted"]), np.array(p["true"])
fig, axes = plt.subplots(1, 2, figsize=(9, 4.2), constrained_layout=True, sharex=True, sharey=True)
for ax, states, title in ((axes[0], pred, "predicted by the learned model"), (axes[1], true, "what MuJoCo says happens")):
for k in range(len(states)):
ax.plot(states[k, :, 1], states[k, :, 0], color="#7a7a7a", alpha=0.25, lw=0.8)
ax.plot(states[k, :, i + 1], states[k, :, i], color="#E69F00", alpha=0.25, lw=0.8)
c = p["chosen"]
ax.plot(states[c, :, 1], states[c, :, 0], color="#0072B2", lw=2, label="chosen plan: pusher")
ax.plot(states[c, :, i + 1], states[c, :, i], color="#D55E00", lw=2, label="chosen plan: puck")
zone = tt.ZONES[p["task"][1]]
ax.add_patch(plt.Circle((zone[1], zone[0]), tt.SUCCESS_RADIUS, color="#009E73", alpha=0.3))
ax.set_title(title, fontsize=10)
ax.set_aspect("equal")
ax.set_xlabel("y (m)")
axes[0].invert_xaxis()
axes[0].set_ylabel("x (m)")
axes[0].legend(frameon=False, fontsize=8)
fig.savefig(OUT / "candidates.png", dpi=150)
plt.close(fig)
if name == “main”: OUT.mkdir(parents=True, exist_ok=True) standard.model(“one_step”) with ProcessPoolExecutor(min(32, os.cpu_count() or 1), mp_context=multiprocessing.get_context(“spawn”)) as pool: runs = list(pool.map(episode, [(k, i) for k in (“MuJoCo”, “learned model”) for i in range(TASKS)])) probes = list(pool.map(probe, range(PROBES))) print(f”(1) the same planner, two models, {TASKS} held-out tasks”) for kind in (“MuJoCo”, “learned model”): rows = [r for r in runs if r[“kind”] == kind] k = sum(r[“success”] for r in rows) lo, hi = stats.wilson_interval(k, len(rows)) steps = [r[“steps”] for r in rows if r[“success”]] print(f” {kind:<14} success {k}/{len(rows)} [{lo:.2f}, {hi:.2f}] median time to success “ f”{np.median(steps) / 10:.1f} s planning {1000 * np.median([r[‘seconds_per_step’] for r in rows]):.0f} ms per step”) both = {r[“i”]: r[“success”] for r in runs if r[“kind”] == “MuJoCo”} learned = {r[“i”]: r[“success”] for r in runs if r[“kind”] == “learned model”} print(f” tasks solved by MuJoCo only {sum(both[i] and not learned[i] for i in both)}, “ f”by the learned model only {sum(learned[i] and not both[i] for i in both)}”) decisions = [d for p in probes for d in p[“decisions”]] print(f”(2) inside the decisions of {PROBES} learned-model episodes ({len(decisions)} decisions, 64 candidates each)”) for label, contact in ((“pusher not touching a puck”, False), (“pusher touching a puck”, True)): rows = [d for d in decisions if d[“contact”] == contact] print(f” {label:<28} {len(rows):4d} decisions; rank correlation of predicted and true costs: median “ f”{np.median([d[‘spearman’] for d in rows]):.2f}, 10th percentile {np.percentile([d[‘spearman’] for d in rows], 10):.2f}; “ f”chosen candidate truly best in {np.mean([d[‘true_rank’] == 0 for d in rows]):.2f}, in the true top 5 in “ f”{np.mean([d[‘true_rank’] < 5 for d in rows]):.2f}”) example = next(p[“example”] for p in probes if p[“example”] is not None) figure(example) (OUT / “candidates.json”).write_text(json.dumps(example)) print(f”(3) wrote {OUT / ‘candidates.png’} and {OUT / ‘candidates.json’}”)
The planner, the cost and the two models:
```python file=src/mjcourse/worldmodel/mpc.py
from __future__ import annotations
from dataclasses import dataclass
import mujoco
import numpy as np
from mujoco import rollout as mj_rollout
from mjcourse import tabletop as tt
@dataclass
class PlannerConfig:
horizon: int = 12
samples: int = 64
elites: int = 16
iterations: int = 3
init_std: float = 0.6
seed: int = 0
def cost(states: np.ndarray, task: tuple[str, str], actions: np.ndarray) -> np.ndarray:
i = 4 + 4 * tt.OBJECTS.index(task[0])
puck = states[:, 1:, i:i + 2]
to_zone_vec = tt.ZONES[task[1]] - puck
to_zone = np.linalg.norm(to_zone_vec, axis=-1)
behind = puck - to_zone_vec / (to_zone[..., None] + 1e-6) * (tt.PUCK_RADIUS + tt.PUSHER_RADIUS + 0.005)
to_behind = np.linalg.norm(states[:, 1:, 0:2] - behind, axis=-1)
weights = np.linspace(0.5, 1.5, to_zone.shape[1])
return (weights * (to_zone + 0.5 * to_behind)).mean(1) + 0.01 * (actions ** 2).sum((1, 2))
class MujocoModel:
def __init__(self, scene: tt.Tabletop, threads: int = 1, plan_scene: tt.Tabletop | None = None):
self.scene = scene
self.plan_model = (plan_scene or scene).model
self.datas = [mujoco.MjData(self.plan_model) for _ in range(threads)]
self.spec = mujoco.mjtState.mjSTATE_FULLPHYSICS
self.nq, self.nv = scene.model.nq, scene.model.nv
self.qadr = [scene.model.jnt_qposadr[scene.model.joint(f"{o}_puck").id] for o in tt.OBJECTS]
self.vadr = [scene.model.jnt_dofadr[scene.model.joint(f"{o}_puck").id] for o in tt.OBJECTS]
def planar(self, full: np.ndarray) -> np.ndarray:
q, v = full[..., 1:1 + self.nq], full[..., 1 + self.nq:1 + self.nq + self.nv]
parts = [q[..., 0:2], v[..., 0:2]]
for a, b in zip(self.qadr, self.vadr):
parts += [q[..., a:a + 2], v[..., b:b + 2]]
return np.concatenate(parts, -1)
def predict(self, state: np.ndarray, actions: np.ndarray) -> np.ndarray:
m, d = self.scene.model, self.scene.data
x0 = np.zeros(mujoco.mj_stateSize(m, self.spec))
mujoco.mj_getState(m, d, x0, self.spec)
k = actions.shape[0]
skip = self.scene.frame_skip
control = np.repeat(actions * self.scene.variation.max_speed, skip, axis=1)
full, _ = mj_rollout.rollout(self.plan_model, self.datas, np.tile(x0, (k, 1)), control)
planar = self.planar(full[:, skip - 1::skip])
return np.concatenate([np.repeat(state[None, None], k, 0), planar], 1)
class LearnedModel:
def __init__(self, dynamics):
self.dynamics = dynamics
def predict(self, state: np.ndarray, actions: np.ndarray) -> np.ndarray:
return self.dynamics.rollout(np.repeat(state[None], len(actions), 0), actions)
class Planner:
def __init__(self, model, task: tuple[str, str], cfg: PlannerConfig | None = None):
self.model, self.task, self.cfg = model, task, cfg or PlannerConfig()
self.rng = np.random.default_rng(self.cfg.seed)
self.mean = np.zeros((self.cfg.horizon, 2))
self.last: dict = {}
def act(self, state: np.ndarray) -> np.ndarray:
return self.plan(state)[0]
def plan(self, state: np.ndarray, execute: int = 1) -> np.ndarray:
cfg = self.cfg
mean, std = self.mean, np.full_like(self.mean, cfg.init_std)
for _ in range(cfg.iterations):
actions = np.clip(mean + std * self.rng.standard_normal((cfg.samples, *mean.shape)), -1, 1)
actions[0] = np.clip(mean, -1, 1)
states = self.model.predict(state, actions)
costs = cost(states, self.task, actions)
elite = actions[np.argsort(costs)[:cfg.elites]]
mean, std = elite.mean(0), elite.std(0) + 0.05
best = int(np.argmin(costs))
self.last = {"actions": actions, "states": states, "costs": costs, "best": best}
self.mean = np.concatenate([actions[best][execute:], np.repeat(actions[best][-1:], execute, 0)])
return actions[best]
def run_episode(planner_model_factory, task, layout, variation: tt.Variation | None = None, cfg: PlannerConfig | None = None,
steps: int = tt.EPISODE_STEPS, record: bool = False, execute: int = 1) -> dict:
scene = tt.Tabletop(variation)
scene.reset(layout)
planner = Planner(planner_model_factory(scene), task, cfg)
states, predicted_cost, queue = [scene.state()], [], []
for _ in range(steps):
if not queue:
queue = list(planner.plan(states[-1], execute)[:execute])
if record:
predicted_cost.append(float(planner.last["costs"][planner.last["best"]]))
scene.step(queue.pop(0))
states.append(scene.state())
if scene.success(task) and np.linalg.norm(states[-1][4 + 4 * tt.OBJECTS.index(task[0]) + 2:][:2]) < 0.02:
break
out = {"success": scene.success(task), "steps": len(states) - 1, "states": np.array(states)}
if record:
out["predicted_cost"] = predicted_cost
scene.close()
return out
Output (MuJoCo 3.14.0, PyTorch 2.14 CPU, 32 worker processes):
(1) the same planner, two models, 50 held-out tasks
MuJoCo success 50/50 [0.93, 1.00] median time to success 5.5 s planning 719 ms per step
learned model success 50/50 [0.93, 1.00] median time to success 5.8 s planning 10 ms per step
tasks solved by MuJoCo only 0, by the learned model only 0
(2) inside the decisions of 10 learned-model episodes (555 decisions, 64 candidates each)
pusher not touching a puck 314 decisions; rank correlation of predicted and true costs: median 0.99, 10th percentile 0.94; chosen candidate truly best in 1.00, in the true top 5 in 1.00
pusher touching a puck 241 decisions; rank correlation of predicted and true costs: median 0.91, 10th percentile 0.72; chosen candidate truly best in 0.93, in the true top 5 in 0.99
MPC with MuJoCo fails. Before blaming any learned model, make the planner work with the true model: check the cost (does the best predicted future actually solve the task?), the horizon (does it reach the event that matters, here contact?) and the samples. The first version of this lesson’s cost pulled the pusher toward the puck instead of toward the point behind it, over a 1 s horizon, and MuJoCo-as-model solved only 2 of 6 tasks.
The learned planner looks perfect. Check whether it is being tested on decisions that matter. The first version of output (2) probed only the first decision of each episode, where the pusher is far from every puck; all candidates differ only in free motion, and the learned model ranked them perfectly (correlation 1.00).
Planning is too slow. A learned model rolls out all candidates as one batched tensor operation; a simulator rolls them out one by one (in parallel threads with mujoco.rollout). That difference, 10 against 719 ms here, is often the practical reason to learn a model of a system one already has.
Halve the planner’s samples (32) and iterations (2) and repeat output (1) for both models. Predict which model loses more success, and why, before running it.
Replace the learned model by an ensemble of five, and plan against the mean prediction plus a penalty on the ensemble’s spread (as in PETS, arXiv:1805.12114). Measure the contact rows of output (2) again.
[!research] Learned models for planning TD-MPC (arXiv:2203.04955) plans over a short horizon with a learned latent model and adds a learned terminal value for what lies beyond the horizon. That is a different answer to the same problem this lesson solved with shaping (the pusher-placement term) and fast replanning. The rank-correlation measurement of output (2) is a way to evaluate a model for planning directly, without running full episodes (inference: a proposal, not a standard benchmark).
{"id": "23.4-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "Why does the learned model reach the same success as MuJoCo despite centimetre errors after 1.2 s?",
"options": ["The errors cancel", "The planner applies only the first 0.1 s of each plan and replans from the measured state, so the model only has to choose good first actions", "The tasks are easy", "MuJoCo is wrong"],
"answer": 1,
"explain": "<p>Replanning corrects the model's errors before they accumulate.</p>"},
{"kind": "mcq", "q": "During contact, how often was the learned model's chosen candidate among MuJoCo's five best?",
"options": ["0.72", "0.91", "0.93", "0.99"],
"answer": 3,
"explain": "<p>Truly best in 0.93 of contact decisions, in the top five in 0.99.</p>"},
{"kind": "mcq", "q": "Which quantity measures what a planner needs from a model, decision by decision?",
"options": ["One-step error", "The rank correlation between predicted and true costs of the candidates", "Training loss", "Parameter count"],
"answer": 1,
"explain": "<p>A planner chooses by ranking; exact prediction is not required.</p>"}
]}
Lesson 23.5 breaks this agreement on purpose: the same network trained on different data, planning horizons from 0.4 to 2 s, and a measurement of how far the chosen plans fall short of what the model promised.