Every lesson in this course so far has had a perfect model of its world: MuJoCo, given an MJCF file. A planner that asked “what happens if I push this way?” could ask MuJoCo and get the right answer for the system described. A world model is what an agent uses when it does not have that: a model of how its world responds to its actions, learned from experience.
The phrase is used loosely, often for any network that predicts something. This level uses an operational definition, and every lesson tests one of its parts:
A world model is a learned, action-conditioned predictive model: given what the agent has observed and a sequence of actions it might take, it predicts future observations or task-relevant quantities (states, rewards, outcomes) over several steps, well enough to be used for a decision (planning, policy learning or evaluation).
Three requirements follow, and each one excludes something:
[!established] Where the term comes from Ha and Schmidhuber’s “World Models” (arXiv:1803.10122) learns “a compressed spatial and temporal representation” of an environment without supervision, and trains an agent “entirely inside of its own hallucinated dream” before transferring it back. PlaNet (arXiv:1811.04551) learns dynamics from images and plans in latent space; its model “must accurately predict the rewards ahead for multiple time steps”. Dreamer (arXiv:1912.01603) learns behaviours “purely by latent imagination”, propagating gradients of learned state values through trajectories imagined in a learned world model; DreamerV3 (arXiv:2301.04104) applies the approach across domains. TD-MPC (arXiv:2203.04955) optimizes trajectories over a short horizon with “a learned task-oriented latent dynamics model” and a learned terminal value, and TD-MPC2 (arXiv:2310.16828) does so in “the latent space of a learned implicit (decoder-free) world model”. None of these is reproduced in this level; its models demonstrate the concepts on one tabletop scene.
Output (1) below is the whole lesson in one table. For the same 200 one-second action sequences, four models predict where the pusher and the pushed puck will be. Read it column by column. Where nothing touches, a model written by a person from the scene’s description is good. Where contact happens, only MuJoCo and the learned model are good, and only MuJoCo is exact.
| Term | What it is | In this course |
|---|---|---|
| simulator | a program computing the next state from the current state, the action and a model written by people | MuJoCo with tabletop.xml |
| forward dynamics model | any function $s_{t+1} = f(s_t, a_t)$, analytic or learned | all four models of output (1) |
| system-identification model | a model whose structure comes from physics and whose few parameters are fitted to data | Level 18; MuJoCo with fitted friction below |
| predictive model | anything that predicts future observations, with or without actions | a video predictor |
| latent dynamics model | dynamics in a learned representation: $z = e(o)$, $z_{t+1} = g(z_t, a_t)$ | Lesson 23.3 |
| world model | the definition above: learned, action-conditioned, multi-step, decision-relevant | Lessons 23.2 to 23.6 |
| model-based reinforcement learning | using a model, learned or given, to plan or to train a policy | Lessons 21.4 and 23.4 |
The tabletop scene of Level 22 has a 16-number planar state: the pusher’s position and velocity, and each puck’s. The truth comes from MuJoCo. Two hundred test segments are made by running the data-collection policy of Lesson 23.2 for a random time, saving MuJoCo’s full state, then applying the policy’s next ten actions (1 s). Four models predict the final state from the saved state and the same actions:
mjSTATE_INTEGRATION): with mjSTATE_FULLPHYSICS the restart’s pushed-puck prediction was off by 0.06 mm on average after a second of contact.Output (1), mean error after 1 s (mm), split by whether the pusher touched a puck during the second (it did in 169 of 200):
| Model | pusher, no contact | pushed puck, no contact | pusher, contact | pushed puck, contact |
|---|---|---|---|---|
| MuJoCo | 0.00 | 0.00 | 0.00 | 0.00 |
| MuJoCo, wrong friction | 0.00 | 0.02 | 0.84 | 1.84 |
| hand-written | 6.33 | 0.28 | 4.84 | 36.14 |
| learned | 10.43 | 2.31 | 6.96 | 7.65 |
| nothing moves | 101.68 | 0.28 | 63.66 | 36.14 |
Four readings:
[!recommendation] Calling something a world model Say what it predicts (states, observations, rewards), from what (states or images), conditioned on which actions, over how many steps, and for which decision. Evaluate it against the simulator where one exists, and against the simple baselines: “nothing moves” and the model a person would write.
INPUT: tabletop.xml; the learned model of mjcourse.worldmodel.standard
PROCESS: 200 test segments; five predictions of the state 1 s later; errors split by contact
OUTPUT: the table above
```python file=examples/l23_1_kinds_of_model.py import mujoco import numpy as np
from mjcourse import tabletop as tt from mjcourse.worldmodel import data, standard
H = 10
def segment(seed: int, rng: np.random.Generator): task = (tt.OBJECTS[rng.integers(3)], list(tt.ZONES)[rng.integers(4)]) scene = tt.Tabletop() scene.reset(tt.sample_layout(rng, task)) policy = lambda s: rng.uniform(-1, 1, 2) if rng.random() < 0.2 else np.clip(tt.expert(s, task) + rng.normal(0, 0.3, 2), -1, 1) for _ in range(int(rng.integers(0, 100))): scene.step(policy(scene.state())) full = np.zeros(mujoco.mj_stateSize(scene.model, mujoco.mjtState.mjSTATE_INTEGRATION)) mujoco.mj_getState(scene.model, scene.data, full, mujoco.mjtState.mjSTATE_INTEGRATION) s0, actions, path = scene.state(), [], [scene.state()] for _ in range(H): actions.append(policy(path[-1])) scene.step(actions[-1]) path.append(scene.state()) scene.close() touched = data.contact_share([{“states”: np.array(path)}]) > 0 return s0, full, np.array(actions), path[-1], task, touched
def mujoco_prediction(full: np.ndarray, actions: np.ndarray, friction: float = 0.4) -> np.ndarray: scene = tt.Tabletop(tt.Variation(friction=friction)) mujoco.mj_setState(scene.model, scene.data, full, mujoco.mjtState.mjSTATE_INTEGRATION) mujoco.mj_forward(scene.model, scene.data) for a in actions: scene.step(a) out = scene.state() scene.close() return out
def hand_written(s0: np.ndarray, actions: np.ndarray) -> np.ndarray: s = s0.copy() tau, dt = 0.5 / 40.0, tt.CONTROL_DT for a in actions: v_cmd = np.clip(a, -1, 1) * tt.MAX_SPEED v0 = s[2:4].copy() decay = np.exp(-dt / tau) s[2:4] = v_cmd + (v0 - v_cmd) * decay s[0:2] += v_cmd * dt + (v0 - v_cmd) * tau * (1 - decay) s[4:] = np.r_[[s[4 + 4 * i:6 + 4 * i].tolist() + [0.0, 0.0] for i in range(3)]].ravel() return s
def main() -> None: mlp = standard.model(“one_step”) rng = np.random.default_rng(23) rows = {name: {“contact”: [], “free”: []} for name in (“MuJoCo”, “MuJoCo, wrong friction”, “hand-written”, “learned”, “nothing moves”)} for _ in range(200): s0, full, a, true, task, touched = segment(23, rng) target = 4 + 4 * tt.OBJECTS.index(task[0]) predictions = {“MuJoCo”: mujoco_prediction(full, a), “MuJoCo, wrong friction”: mujoco_prediction(full, a, 0.3), “hand-written”: hand_written(s0, a), “learned”: mlp.rollout(s0[None], a[None].astype(np.float32))[0, -1], “nothing moves”: s0} for name, p in predictions.items(): rows[name][“contact” if touched else “free”].append( (np.linalg.norm(p[:2] - true[:2]), np.linalg.norm(p[target:target + 2] - true[target:target + 2]))) n_contact = len(rows[“MuJoCo”][“contact”]) print(f”200 test segments of 1 s; the pusher touched a puck in {n_contact} of them”) print(f” {‘model’:<24}{‘pusher, no contact’:>20}{‘puck, no contact’:>18}{‘pusher, contact’:>18}{‘puck, contact’:>16}” “ (mean error, mm)”) for name, r in rows.items(): free, contact = np.array(r[“free”]) * 1000, np.array(r[“contact”]) * 1000 print(f” {name:<24}{free[:, 0].mean():>20.2f}{free[:, 1].mean():>18.2f}{contact[:, 0].mean():>18.2f}” f”{contact[:, 1].mean():>16.2f}”)
if name == “main”: main()
Output (MuJoCo 3.14.0; the learned model is built and cached by `mjcourse.worldmodel.standard` on first use, about 4 minutes):
```text
200 test segments of 1 s; the pusher touched a puck in 169 of them
model pusher, no contact puck, no contact pusher, contact puck, contact (mean error, mm)
MuJoCo 0.00 0.00 0.00 0.00
MuJoCo, wrong friction 0.00 0.02 0.84 1.84
hand-written 6.33 0.28 4.84 36.14
learned 10.43 2.31 6.96 7.65
nothing moves 101.68 0.28 63.66 36.14
The simulator disagrees with itself. Restarting MuJoCo from a saved state does not reproduce the original trajectory. Save mjSTATE_INTEGRATION, which includes the warm start, not only positions and velocities. Replaying recorded actions rounded to float32 also diverges in contact-rich scenes, by centimetres within seconds.
A “world model” cannot answer counterfactuals. If its predictions do not change when the actions do, it is not action-conditioned. Test with two different action sequences from the same state.
A learned model is evaluated only where it is accurate. Split the error by the situations that matter (contact against no contact here) before averaging.
Add a sixth model to output (1): MuJoCo with the right friction but the pucks’ mass doubled. Predict, before running it, which columns change and by roughly how much, then check. Then improve the hand-written model with the pusher’s damping (1 N s/m) and force limit (15 N), and measure what that buys in each column.
Make the hand-written model handle one contact: when the pusher reaches the target puck, move the puck along the line of centres at the pusher’s speed (the pushing rule the expert of Level 22 relies on). Measure its error in the contact columns against the learned model, and find the situations (sliding, two pucks touching, rim contact) where your rule is worst.
[!research] What counts as a world model The published systems differ in what they predict and why. PlaNet’s model must predict rewards several steps ahead, and adds a multi-step objective, “latent overshooting”, for it (arXiv:1811.04551). TD-MPC2’s model has no decoder at all: it is never asked to reproduce an observation (arXiv:2310.16828). Lambert et al. name the gap between a model’s training objective and its use “objective mismatch”: a model trained for one-step likelihood need not be the one that controls best (arXiv:2002.04523). Lessons 23.3 and 23.5 measure versions of that gap; Lesson 23.7 asks you to study it.
{"id": "23.1-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "A network predicts the next 10 video frames of a robot workspace from the last 4, without being given the robot's actions. Under this level's definition it is:",
"options": ["a world model", "a predictive model but not a world model, because it is not action-conditioned", "a simulator", "a system-identification model"],
"answer": 1,
"explain": "<p>Without actions it cannot answer what would happen under a different choice.</p>"},
{"kind": "mcq", "q": "In output (1), why does the hand-written model match 'nothing moves' exactly on the pushed puck in contact (36.14 mm)?",
"options": ["Both are learned", "The hand-written model has no contact, so it never moves the pucks either", "Coincidence", "Friction is wrong"],
"answer": 1,
"explain": "<p>Its author left contact out, so its puck predictions are the starting positions.</p>"},
{"kind": "mcq", "q": "Restarting MuJoCo from a saved state left a 0.06 mm difference after a second of contact. What was missing?",
"options": ["The actions", "The contact solver's warm start, which mjSTATE_INTEGRATION includes and mjSTATE_FULLPHYSICS does not", "The timestep", "The random seed"],
"answer": 1,
"explain": "<p>Exact replay needs the full integration state; without it the pushed puck was off by 0.06 mm on average.</p>"},
{"kind": "mcq", "q": "Where is the learned model better than the model a person wrote?",
"options": ["Everywhere", "On the free pusher", "On the pushed puck during contact, the physics the person left out", "Nowhere"],
"answer": 2,
"explain": "<p>7.65 against 36.14 mm in contact; worse on the free pusher, 10.43 against 6.33 mm.</p>"}
]}
Lesson 23.2 trains the learned model of this table, and follows its error from one step to fifty.