A planner needs a model that answers one question quickly and often: if I apply these actions from this state, what happens? MuJoCo answers it exactly, for the system MuJoCo was given. A learned world model answers it approximately, for whatever its training data covered. This lesson runs the same sampling-based planner with three kinds of model:
For the learned models it keeps three properties apart:
[!established] Predictive sampling with a model Every 40 ms the planner perturbs a 4-knot control spline over a 1 s horizon with Gaussian noise at three scales (0.5, 2 and 8 N), predicts each of 64 candidates’ trajectories with its model, scores them with a running cost (pole upright, cart centred, small velocities and force), and applies the best one until the next replan (Howell et al., arXiv:2212.00541; the same planner as the planner demo). Nothing in the controller depends on the model being MuJoCo: it needs a function from (state, force sequence) to predicted states.
The lab is the planner demo’s planner, live in this page. Every faint line is a future the model predicted; the orange one is the plan being executed. Make the real pole heavier with the slider while the planner keeps assuming 0.2 kg. It still balances a 0.4 kg pole; at 0.8 kg the script below measured 1 success in 10.
```lab mpc {“dock”: true, “title”: “A planner in the page”}
## Mathematics
### MuJoCo as the model, right and wrong
The task is to swing the cart-pole up from hanging and balance it, judged over 10 s on 10 starts. Success means within 0.2 rad of upright for the last 2 s, with the cart inside its rails. The planner's model is a 10 ms copy of the simulator with a 0.2 kg pole:
| Real pole | 0.2 kg | 0.4 kg | 0.8 kg |
|---|---|---|---|
| success | 10/10 | 10/10 | 1/10 |
Replanning every 40 ms from the measured state corrects the model's errors as they appear, so doubling the pole's mass costs nothing. Quadrupling it does: the predicted futures diverge from the real one faster than replanning can correct them (inference).
### A learned model
An MLP (two hidden layers of 128 units) predicts the 10 ms change of the state $(x, \theta, \dot x, \dot\theta)$ from $x$, $\sin\theta$, $\cos\theta$, the velocities and the force. It is rolled out step by step inside the same planner. Each model is evaluated on three axes:
- **Fidelity:** the one-step error on 5000 held-out transitions drawn uniformly over the state range.
- **Action-conditioned accuracy:** the median pole-angle error after 1 s of open-loop prediction under random forces, from 200 states.
- **Downstream utility:** closed-loop success, as above.
| Training data | Share near upright | One-step RMS error ($x$, $\theta$, $\dot x$, $\dot\theta$) | Angle error after 1 s | Success |
|---|---|---|---|---|
| uniform, 100 transitions | 0.140 | 2.6e-04, 1.0e-03, 5.2e-03, 3.1e-02 | 0.311 rad | 0/10 |
| uniform, 300 | 0.160 | 8.0e-05, 3.0e-04, 1.6e-03, 5.3e-03 | 0.017 rad | 10/10 |
| uniform, 1,000 | 0.167 | 4.0e-05, 1.8e-04, 9.2e-04, 2.2e-03 | 0.011 rad | 10/10 |
| uniform, 10,000 | 0.165 | 3.4e-05, 1.5e-04, 7.4e-04, 1.7e-03 | 0.010 rad | 10/10 |
| exploration, 100,000 | 0.0024 | 2.1e-03, 2.5e-03, 1.1e-01, 1.3e-01 | 0.209 rad | 0/10 |
With data spread over the task's states, 300 transitions are enough for the planner to succeed every time, and more data barely changes anything. With 100 the model's one-step errors look small (about a milliradian per step in angle), but they compound to 0.31 rad within the horizon, and the planner fails. The last row is the instructive one. Over three hundred times the data of the smallest successful set, collected by random forces from the hanging pose, covers the upright region in 0.24% of its states, so the model predicts velocities 50 to 120 times worse than the uniform models, and the planner never succeeds. Data volume did not buy coverage.
> [!recommendation] Evaluating a world model
> Report its fidelity, its action-conditioned accuracy and its downstream utility separately. A small one-step error does not imply a usable multi-step prediction, and neither implies that a planner succeeds. Measure accuracy on the states the planner will visit, and check the training data's coverage of them before training. When a simulator is available, compare against MuJoCo as the model: it tells you how much of the planner's failure is the model's fault.
### Where MuJoCo ends and a learned model begins
MuJoCo is a model of a system someone wrote down: exact for that system, wrong wherever the real one differs (Levels 17 to 19). A learned model can absorb what nobody wrote down, given data that covers it, and fails silently outside that data. In practice the two meet: simulated data trains learned models, learned models plan in place of simulators that are too slow or not available, and a simulator, as here, gives the ground truth needed to measure a learned model's errors (inference from this lesson's measurements).
## Implementation
```io
INPUT: cartpole.xml; the predictive-sampling planner (64 samples, 1 s, replanning every 40 ms)
PROCESS: plan with MuJoCo (right and wrong pole mass); train MLP dynamics on uniform or exploration transitions; measure one-step error, 1 s open-loop error and closed-loop success
OUTPUT: printed tables
```python file=examples/l21_4_world_models.py “"”Lesson 21.4: planning with MuJoCo, then with a learned world model (Level 21).
INPUT cartpole.xml; the predictive-sampling planner of the showcase (64 samples, 4-knot spline over 1 s, noise scales 0.5, 2 and 8 N, replanning every 40 ms); swing up from hanging and balance, 10 s, success when the pole stays within 0.2 rad of upright for the last 2 s with the cart inside its rails PROCESS (1) the planner with MuJoCo as its model (a 10 ms copy of the simulator, pole 0.2 kg), 10 starts, on the real cart-pole with a 0.2, 0.4 and 0.8 kg pole; (2) a learned world model: an MLP predicting the 10 ms state change from the state (pole angle as its sine and cosine) and the force; trained on 100 to 10,000 transitions drawn uniformly over the task’s state range, or on 100,000 transitions from random exploration (random forces held for 0.1 s, 1 s episodes from hanging); for each: the share of states within 0.5 rad of upright, the one-step error on uniform held-out transitions, the open-loop pole-angle error of 1 s predictions under random forces, and the closed-loop success of the same planner using the learned model OUTPUT printed tables
Run: python examples/l21_4_world_models.py (about 10 minutes) “””
import os
import mujoco import mujoco.rollout as rollout import numpy as np import torch from torch import nn
from mjcourse import model_path
FAST = os.environ.get(“MJC_FAST”) == “1” FULL = mujoco.mjtState.mjSTATE_FULLPHYSICS SCALES = np.array([0.5, 2.0, 8.0]) KNOTS, REPLAN, PLAN_DT, HORIZON, SAMPLES, LIMIT = 4, 0.04, 0.01, 1.0, 64, 20.0 LOW = np.array([-1.4, -np.pi, -3.0, -12.0]) # state ranges for the data: x, theta, xdot, thetadot HIGH = -LOW STEPS = round(HORIZON / PLAN_DT)
def cost(x: np.ndarray, u: np.ndarray) -> np.ndarray: “"”Summed running cost; x is (N, T, 4) states (x, theta, xdot, thetadot), u is (N, T) forces.””” return ((1 - np.cos(x[…, 1])) + 0.1 * x[…, 0] ** 2 + 0.01 * x[…, 2] ** 2 + 0.001 * x[…, 3] ** 2 + 1e-4 * u ** 2).sum(axis=1)
class MujocoModel: “"”The simulator at a 10 ms timestep, as a model: predict(states (N, 4), forces (N, T)) -> (N, T, 4).”””
def __init__(self):
self.model = mujoco.MjModel.from_xml_path(str(model_path("cartpole")))
self.model.opt.timestep = PLAN_DT
self.datas = [mujoco.MjData(self.model) for _ in range(8)]
self.template = np.zeros(mujoco.mj_stateSize(self.model, FULL))
def predict(self, states: np.ndarray, forces: np.ndarray) -> np.ndarray:
x0 = np.tile(self.template, (len(states), 1))
x0[:, 1:3], x0[:, 3:5] = states[:, :2], states[:, 2:] # full state: time, qpos, qvel
out, _ = rollout.rollout(self.model, self.datas, x0, forces[..., None])
return out[..., 1:5]
def features(states: np.ndarray, forces: np.ndarray) -> np.ndarray: “"”The network’s input: x, sin and cos of the pole angle (the dynamics are periodic in it), velocities, force.””” return np.c_[states[:, 0], np.sin(states[:, 1]), np.cos(states[:, 1]), states[:, 2:], forces]
class LearnedModel: “"”An MLP predicting the 10 ms state change, rolled out step by step.”””
def __init__(self, net: nn.Module, mean: np.ndarray, std: np.ndarray, dmean: np.ndarray, dstd: np.ndarray):
self.net, self.mean, self.std, self.dmean, self.dstd = net, mean, std, dmean, dstd
def step(self, states: np.ndarray, forces: np.ndarray) -> np.ndarray:
z = features(states, forces)
with torch.no_grad():
delta = self.net(torch.as_tensor((z - self.mean) / self.std, dtype=torch.float32)).numpy()
return states + delta * self.dstd + self.dmean
def predict(self, states: np.ndarray, forces: np.ndarray) -> np.ndarray:
out, x = np.empty((len(states), forces.shape[1], 4)), states.copy()
for t in range(forces.shape[1]):
x = self.step(x, forces[:, t])
out[:, t] = x
return out
def transitions(n: int, seed: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: rng = np.random.default_rng(seed) states, forces = rng.uniform(LOW, HIGH, (n, 4)), rng.uniform(-LIMIT, LIMIT, n) sim = MujocoModel() nxt = sim.predict(states, forces[:, None])[:, 0] return states, forces, nxt
def train(n: int, seed: int = 0, data: str = “uniform”) -> LearnedModel: s, u, s1 = transitions(n, seed) if data == “uniform” else explore(n, seed) z, d = features(s, u), s1 - s mean, std, dmean, dstd = z.mean(0), z.std(0) + 1e-8, d.mean(0), d.std(0) + 1e-8 torch.manual_seed(seed) net = nn.Sequential(nn.Linear(6, 128), nn.SiLU(), nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 4)) opt = torch.optim.Adam(net.parameters(), lr=2e-3) x = torch.as_tensor((z - mean) / std, dtype=torch.float32) y = torch.as_tensor((d - dmean) / dstd, dtype=torch.float32) updates = 300 if FAST else 6000 sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, updates) gen = torch.Generator().manual_seed(seed) for _ in range(updates): idx = torch.randint(0, n, (min(512, n),), generator=gen) loss = ((net(x[idx]) - y[idx]) ** 2).mean() opt.zero_grad() loss.backward() opt.step() sched.step() return LearnedModel(net, mean, std, dmean, dstd)
def explore(n: int, seed: int = 0) -> tuple[np.ndarray, np.ndarray, np.ndarray]: “"”Transitions from random exploration: 1 s episodes from hanging, a new random force every 0.1 s.””” rng, sim, chunks = np.random.default_rng(seed), MujocoModel(), [] while sum(len(c[0]) for c in chunks) < n: start = np.array([0.0, np.pi + rng.normal(0, 0.1), 0.0, 0.0]) forces = np.repeat(rng.uniform(-LIMIT, LIMIT, 10), 10)[None] traj = sim.predict(start[None], forces)[0] chunks.append((np.vstack([start, traj[:-1]]), forces[0], traj)) s, u, s1 = (np.concatenate([c[k] for c in chunks])[:n] for k in range(3)) return s, u, s1
def episode(world, seed: int, seconds: float = 10.0, pole_mass: float = 0.2) -> bool: rng = np.random.default_rng(seed) spec = mujoco.MjSpec.from_file(str(model_path(“cartpole”))) spec.geom(“pole”).mass = pole_mass model = spec.compile() data = mujoco.MjData(model) mujoco.mj_resetDataKeyframe(model, data, model.key(“hanging”).id) data.qpos[1] += rng.normal(0.0, 0.05) knot_t, step_t = np.linspace(0.0, HORIZON, KNOTS), np.arange(STEPS) * PLAN_DT nominal, t_nominal, angles = np.zeros(KNOTS), 0.0, [] for _ in range(round(seconds / REPLAN)): nominal = np.interp(knot_t + (data.time - t_nominal), knot_t, nominal) t_nominal = data.time noise = SCALES[np.arange(SAMPLES - 1) % 3][:, None] * rng.standard_normal((SAMPLES - 1, KNOTS)) candidates = np.vstack([nominal, np.clip(nominal + noise, -LIMIT, LIMIT)]) forces = np.stack([np.interp(step_t, knot_t, c) for c in candidates]) state = np.r_[data.qpos, data.qvel] predicted = world.predict(np.tile(state, (SAMPLES, 1)), forces) nominal = candidates[int(np.argmin(cost(predicted, forces)))] for _ in range(round(REPLAN / model.opt.timestep)): data.ctrl[0] = np.interp(data.time - t_nominal, knot_t, nominal) mujoco.mj_step(model, data) angles.append(abs((data.qpos[1] + np.pi) % (2 * np.pi) - np.pi)) return bool((np.array(angles[-round(2.0 / REPLAN):]) < 0.2).all() and abs(data.qpos[0]) < 1.45)
if name == “main”: starts = range(3) if FAST else range(10) print(f”MuJoCo {mujoco.version}, torch {torch.version}; planner: {SAMPLES} samples, {HORIZON} s horizon, “ f”replanning every {1000 * REPLAN:.0f} ms; {len(starts)} starts per row”) sim = MujocoModel() print(“(1) MuJoCo as the model (it assumes a 0.2 kg pole)”) for mass in (0.2, 0.4, 0.8): print(f” real pole {mass} kg: success {sum(episode(sim, s, pole_mass=mass) for s in starts)}/{len(starts)}”) print(“(2) a learned model (MLP, 2 x 128 units)”) test_s, test_u, test_s1 = transitions(5000, seed=99) rng = np.random.default_rng(7) starts_ol = rng.uniform(LOW * [0.5, 1, 0.3, 0.3], HIGH * [0.5, 1, 0.3, 0.3], (200, 4)) forces_ol = rng.uniform(-LIMIT, LIMIT, (200, STEPS)) truth = sim.predict(starts_ol, forces_ol) print(f” {‘data’:<28}{‘near upright’:>13}{‘one-step RMS error (x, th, xd, thd)’:>40}{‘pole angle error at 1 s’:>25}{‘success’:>9}”) sets = [(“uniform”, n) for n in ((300, 1000) if FAST else (100, 300, 1000, 10000))] + [(“exploration”, 10000 if FAST else 100000)] for data, n in sets: s_train = (transitions if data == “uniform” else explore)(n, 0)[0] near = np.mean(np.abs((s_train[:, 1] + np.pi) % (2 * np.pi) - np.pi) < 0.5) world = train(n, data=data) rms = np.sqrt(np.mean((world.step(test_s, test_u) - test_s1) ** 2, axis=0)) pred = world.predict(starts_ol, forces_ol) err = np.median(np.abs(np.angle(np.exp(1j * (pred[:, -1, 1] - truth[:, -1, 1]))))) wins = sum(episode(world, s) for s in starts) print(f” {data + ‘, ‘ + format(n, ‘,’) + ‘ transitions’:<28}{near:>13.4f} “ f”{rms[0]:.1e} {rms[1]:.1e} {rms[2]:.1e} {rms[3]:.1e}{err:>22.3f} rad{wins:>6}/{len(starts)}”)
Output:
```text
MuJoCo 3.14.0, torch 2.14.1+cpu; planner: 64 samples, 1.0 s horizon, replanning every 40 ms; 10 starts per row
(1) MuJoCo as the model (it assumes a 0.2 kg pole)
real pole 0.2 kg: success 10/10
real pole 0.4 kg: success 10/10
real pole 0.8 kg: success 1/10
(2) a learned model (MLP, 2 x 128 units)
data near upright one-step RMS error (x, th, xd, thd) pole angle error at 1 s success
uniform, 100 transitions 0.1400 2.6e-04 1.0e-03 5.2e-03 3.1e-02 0.311 rad 0/10
uniform, 300 transitions 0.1600 8.0e-05 3.0e-04 1.6e-03 5.3e-03 0.017 rad 10/10
uniform, 1,000 transitions 0.1670 4.0e-05 1.8e-04 9.2e-04 2.2e-03 0.011 rad 10/10
uniform, 10,000 transitions 0.1651 3.4e-05 1.5e-04 7.4e-04 1.7e-03 0.010 rad 10/10
exploration, 100,000 transitions 0.0024 2.1e-03 2.5e-03 1.1e-01 1.3e-01 0.209 rad 0/10
The learned model has a tiny training loss and the planner fails. Its errors compound over the horizon or its data misses the states the planner visits. Measure multi-step error and coverage.
More data did not help. Check what the data covers, not how much there is: exploration from one pose rarely visits the states that matter.
The planner with the true model fails too. The problem is the planner (samples, horizon, cost), not the model. The planner demo’s tables measure those.
A model trained on angles fails after a full swing. Represent angles as sine and cosine; the dynamics are periodic in them.
Train on exploration data collected by the planner itself, using the true model, and retrain on its trajectories. How few transitions does the learned model need when its data comes from the policy that will use it?
Learn the model from 64 x 64 images instead of states (a latent world model, as in PlaNet or DreamerV3 in the reading list), and evaluate it on this lesson’s three axes. Which axis degrades first?
[!research] World models in the reading list Ha and Schmidhuber (arXiv:1803.10122), PlaNet (arXiv:1811.04551), DreamerV3 (arXiv:2301.04104), TD-MPC2 (arXiv:2310.16828) and V-JEPA 2 (arXiv:2506.09985) learn dynamics from data and plan or train with them. For each, separate the three claims this lesson measured: how accurately it predicts, how well its predictions follow actions, and how well the policy or planner that uses it performs. Then ask what data covered the states the planner visits. Predictive sampling (Howell et al., arXiv:2212.00541) is the planner here.
{"id": "21.4-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "The planner assumes a 0.2 kg pole. Why does it still balance a 0.4 kg pole?",
"options": ["The cost function adapts", "It replans every 40 ms from the measured state, correcting the model's errors as they appear", "The pole's mass does not matter", "It learns online"],
"answer": 1,
"explain": "<p>At 0.8 kg the errors grow faster than replanning corrects them: 1/10.</p>"},
{"kind": "mcq", "q": "A model trained on 100,000 exploration transitions fails while one trained on 300 uniform transitions succeeds. Why?",
"options": ["Overfitting", "Coverage: only 0.24% of the exploration states are near upright, where the planner needs accurate predictions", "Too few epochs", "The network is too small"],
"answer": 1,
"explain": "<p>Its velocity predictions are a hundred times worse on uniform held-out states.</p>"},
{"kind": "mcq", "q": "The 100-transition model's one-step angle error is about 1 mrad, yet the planner fails. What does that show?",
"options": ["One-step error is enough to judge a model", "Errors compound over the horizon: 0.31 rad after 1 s", "The planner is broken", "Angles were wrapped wrongly"],
"answer": 1,
"explain": "<p>Report multi-step, action-conditioned error and downstream success separately.</p>"},
{"kind": "mcq", "q": "Which of the three axes did the planner's success measure?",
"options": ["Fidelity", "Action-conditioned accuracy", "Downstream utility"],
"answer": 2,
"explain": "<p>Fidelity is the one-step error, action-conditioned accuracy the open-loop error under random forces.</p>"}
]}
Lesson 21.5 builds evaluation for embodied AI on top of MuJoCo: generalization axes, ground truth for perception, and privileged state.