What Is a World Model? › Building and Stress-Testing a Tiny World Model
World Models HandbookAll chapters What Is a World Model? › Building and Stress-Testing a Tiny World ModelOne section · Sections in this chapter7 / 10 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Everything in this chapter so far has been conceptual. Time to make it executable. We will proceed in five stages: simulate the true pendulum and collect experience, fit a linear world model to that experience, measure how prediction error compounds over recursive rollouts, contrast point prediction with stochastic simulation, and finally close the loop by designing a controller from the learned model and seeing what model bias does to real behavior. Keep a running tally as you go: each stage instantiates one term from the definitions above, and by the end you should be able to point at a specific line of code for each clause of the operational definition.
Simulating the true environment
First, the ground-truth pendulum. We integrate the nonlinear dynamics with a classical fourth-order Runge-Kutta scheme at a fixed timestep, and we gather training experience by injecting random torques from randomly sampled nearby-initial conditions. Random excitation is the system identification analogue of exploration: it keeps the dataset informative across the state region we care about. The choice of excitation range is itself a modeling decision: we sample angles within about half a radian of the hanging equilibrium because that is where our downstream local-regulation task will operate, and collecting data elsewhere would spend effort on states the controller will not visit. This is decision sufficiency shaping data collection before any model exists.
import matplotlib.pyplot as plt
import numpy as np
from scipy.linalg import expm, solve_discrete_are
## Physical constants of the true pendulum
g, l, m, c = 9.81, 1.0, 1.0, 0.15
dt = 0.05
def rk4(theta, omega, u, dt):
"""One RK4 step of the true nonlinear pendulum."""
def f(th, om):
return om, -(g / l) * np.sin(th) - c * om + u / (m * l**2)
k1t, k1o = f(theta, omega)
k2t, k2o = f(theta + 0.5 * dt * k1t, omega + 0.5 * dt * k1o)
k3t, k3o = f(theta + 0.5 * dt * k2t, omega + 0.5 * dt * k2o)
k4t, k4o = f(theta + dt * k3t, omega + dt * k3o)
th2 = theta + dt / 6 * (k1t + 2 * k2t + 2 * k3t + k4t)
om2 = omega + dt / 6 * (k1o + 2 * k2o + 2 * k3o + k4o)
return th2, om2
def true_step(x, u):
"""State-interface wrapper: array in, array out."""
th, om = x
th2, om2 = rk4(th, om, u, dt)
return np.array([th2, om2])
## Collect random-excitation experience near the hanging equilibrium
rng = np.random.default_rng(0)
n_episodes, n_steps = 25, 200
states, actions, next_states = [], [], []
for _ in range(n_episodes):
th = rng.uniform(-0.6, 0.6)
om = rng.uniform(-1.0, 1.0)
for _ in range(n_steps):
u = rng.uniform(-2.0, 2.0)
nxt = true_step(np.array([th, om]), u)
states.append([th, om])
actions.append([u])
next_states.append(nxt)
th, om = nxt
states = np.array(states)
actions = np.array(actions)
next_states = np.array(next_states)Transitions collected : 5000 Angle range : [-0.65, +0.64] rad Velocity range : [-1.93, +2.11] rad/s
Five thousand transitions span angles within about half a radian of the hanging equilibrium. This is our agent's entire lived experience of the world. Note what the agent does not have: the equations of motion, knowledge of or , or access to clean states at decision time. It must reverse-engineer the world from these pairs. Five thousand transitions sounds like a lot until you recall that each one is a single snapshot of a two-dimensional state, three-dimensional input space; the coverage is thin in absolute terms, and everything the model gets wrong below will trace back to what these five thousand snapshots never showed it.
Fitting the world model
We now fit the simplest honest world model: a linear one-step predictor , estimated by least squares on the collected transitions. Linearization is exactly the right complexity here because our data lives near the hanging equilibrium, where . This mirrors how classical control engineers identify linear models around operating points, a practice Part II, Chapter 3 formalizes as system identification. The regularization term added to the normal equations is tiny, but it is not cosmetic: it guards against the degenerate directions that arise when inputs are correlated, and it is the smallest possible acknowledgment of the statistical humility that larger learned models must show in far greater measure.
## Least-squares fit of s' = A s + B u on collected transitions
X = np.hstack([states, actions]) # features: [theta, omega, u]
Y = next_states # targets: [theta', omega']
W = np.linalg.solve(X.T @ X + 1e-9 * np.eye(3), X.T @ Y)
A_fit = W[:2].T # shape (2, 2)
B_fit = W[2:].T # shape (2, 1)Learned A (per 50 ms step): [[ 0.9881 0.0496] [-0.476 0.9806]] Learned B: [[0.0012] [0.0496]]
Read the learned matrices against the physics. The diagonal entry near 1 reflects the slow decay of angle over 50 milliseconds; the off-diagonal entries encode the coupling between velocity and angle; and captures how strongly torque kicks velocity. Nothing told the model these roles; least squares discovered them from data. This is the smallest possible demonstration that world models can be learned rather than authored. It is also worth noticing what the linear model structurally cannot capture: the sine nonlinearity, and therefore any behavior far from the hanging equilibrium. The model is not wrong so much as locally right, and the boundary of its rightness is the boundary of our data.
The heatmap makes the "locally right" claim visible. The dark valley of low error hugs the origin, exactly where the random excitation concentrated, and the bright ridges at large angles mark where the model's sine-free structure stops describing reality. Every failure mode discussed later in this chapter, compounding rollout drift and optimizer exploitation included, is a story about trajectories wandering from the dark valley into the bright ridges.
Stress test one: compounding rollout error
Now we interrogate the model the way this chapter insists you must: not by one-step accuracy alone, but by recursive rollout. We generate fresh test episodes unseen during training, then compare three things at every horizon: the model's one-step prediction made always from the true previous state, its recursive rollout in which each prediction feeds the next step, and a rival observation-only autoregressive model that, like our frosted-glass observer, sees only noisy angles and must limp along without velocity. The third competitor is included deliberately: it turns the abstract claim "partial observability taxes everything" into a measurable curve, and it previews the filtering machinery of Part II by showing what a model pays when it lacks an explicit state estimate.
## Fit an observation-only AR model: o_{t+1} = a1 o_t + a2 o_{t-1} + b u_t
obs_noise = 0.05
obs = states[:, 0] + rng.normal(0, obs_noise, size=len(states))
ep_ids = np.repeat(np.arange(n_episodes), n_steps)
phi, y_ar = [], []
for e in range(n_episodes):
idx = np.where(ep_ids == e)[0]
for j in range(1, len(idx) - 1):
phi.append([obs[idx[j]], obs[idx[j - 1]], actions[idx[j], 0]])
y_ar.append(obs[idx[j + 1]])
phi, y_ar = np.array(phi), np.array(y_ar)
w_ar = np.linalg.solve(phi.T @ phi + 1e-9 * np.eye(3), phi.T @ y_ar)
a1, a2, b_ar = w_ar
## Measure one-step vs recursive error on fresh test episodes
H = 120
n_test = 12
err_one = np.zeros(H) # state-space model, one-step
err_rec = np.zeros(H) # state-space model, recursive rollout
err_ar = np.zeros(H) # observation-only AR, recursive rollout
for _ in range(n_test):
x = np.array([rng.uniform(-0.6, 0.6), rng.uniform(-1.0, 1.0)])
us = rng.uniform(-2.0, 2.0, size=H + 2)
traj = [x.copy()]
for h in range(H + 1):
traj.append(true_step(traj[-1], us[h]))
traj = np.array(traj)
# One-step: always re-anchor on truth
for h in range(H):
pred = A_fit @ traj[h] + B_fit @ np.array([us[h]])
err_one[h] += abs(pred[0] - traj[h + 1, 0])
# Recursive: feed predictions back in
pred = traj[0].copy()
for h in range(H):
pred = A_fit @ pred + B_fit @ np.array([us[h]])
err_rec[h] += abs(pred[0] - traj[h + 1, 0])
# Observation-only AR: two noisy angle readings, then roll out blindly
o_prev = traj[0, 0] + rng.normal(0, obs_noise)
o_curr = traj[1, 0] + rng.normal(0, obs_noise)
for h in range(H):
o_pred = a1 * o_curr + a2 * o_prev + b_ar * us[h + 1]
err_ar[h] += abs(o_pred - traj[h + 2, 0])
o_prev, o_curr = o_curr, o_pred
err_one /= n_test
err_rec /= n_test
err_ar /= n_test
horizons = np.arange(1, H + 1)
The picture confirms the chapter's central warning quantitatively. The one-step curve stays essentially flat: anchored on truth each time, the linear model is a fine local predictor. The recursive curve climbs without bound: each small error becomes the next input, and the model wanders off the manifold where its least-squares fit is valid. Meanwhile the observation-only model trails at every horizon, paying a permanent tax for lacking velocity. If you ever evaluate a world model solely by one-step loss, this figure is the rebuttal. Read the log scale carefully: the vertical axis spans orders of magnitude, so the widening gap between the curves represents not a modest degradation but a qualitative collapse of usefulness at long horizons. Also notice the shape of the recursive curve: it accelerates before it climbs steadily, which is the signature of errors leaving the linear regime and entering the region where the model's local approximation no longer describes the true dynamics at all.
Stress test two: prediction versus simulation
Next we illustrate the prediction-simulation distinction directly. We fix an initial state and an action sequence, then do two things with the same fitted dynamics. First, a point prediction: propagate the deterministic model forward, one crisp line. Second, a simulation: add small process noise at every step and sample many stochastic rollouts, producing a fan of imagined futures. We also run one realization of the true nonlinear system under the same noise, to see whether reality stays inside the fan. The comparison is fair by construction: both the crisp line and the fan come from the same fitted matrices, so any difference in their behavior is attributable purely to how uncertainty is represented, not to which model was used.
## Point prediction vs stochastic simulation from the same fitted model
sigma_w = 0.03
T_fan = 80
x0 = np.array([0.4, 0.0])
us_fan = 1.2 * np.sin(2 * np.pi * np.arange(T_fan) / 30)
## Deterministic point prediction
det = [x0.copy()]
for u in us_fan:
det.append(A_fit @ det[-1] + B_fit @ np.array([u]))
det = np.array(det)
## Ensemble of stochastic simulated rollouts
n_sims = 200
sims = np.zeros((n_sims, T_fan + 1, 2))
for k in range(n_sims):
x = x0.copy()
sims[k, 0] = x
for t, u in enumerate(us_fan):
x = A_fit @ x + B_fit @ np.array([u]) + rng.normal(0, sigma_w, size=2)
sims[k, t + 1] = x
## One realized trajectory of the TRUE nonlinear system, same noise
real = [x0.copy()]
x = x0.copy()
for u in us_fan:
th, om = x
th2, om2 = rk4(th, om, u, dt)
x = np.array([th2, om2]) + rng.normal(0, sigma_w, size=2)
real.append(x.copy())
real = np.array(real)
lo = np.percentile(sims[:, :, 0], 5, axis=0)
hi = np.percentile(sims[:, :, 0], 95, axis=0)
mean_angle = sims[:, :, 0].mean(axis=0)
band_half_width = (hi - lo) / 2.0
t_axis = np.arange(T_fan + 1)
This single figure carries the whole conceptual payload of the prediction-versus-simulation section. The deterministic prediction is confident and specific, but it answers only "what is the most typical future?" The simulation fan answers a richer question: "what range of futures is plausible, and with what spread?" Notice two things. The fan widens with horizon, making uncertainty grow visibly with time rather than hiding it. And the true system's realization, the dotted line, generally lives inside the fan, which is exactly what a well-calibrated stochastic world model buys you: reality should look like a typical sample, not like an outlier. A planner that reasons only with the crisp center line will be repeatedly surprised; a planner that reasons with the fan can hedge. One caveat keeps this figure honest: the fan's width reflects the noise level we injected, , which we chose ourselves. In a real application the noise level must be estimated from data, and a miscalibrated noise estimate produces a fan that is either falsely reassuring or paralyzingly wide. Calibration of uncertainty is a first-class evaluation criterion, and the next chapter's qualification test includes it explicitly.
The widening band is the quantitative shadow of the fan chart above. Early in the rollout the ensemble members barely differ, so the model is effectively making a point prediction; by the end of the horizon the spread has grown several-fold, and the honest answer to "where will the pendulum be?" is a range, not a number. When you later meet planners that penalize model uncertainty, this curve is the quantity they consume: the further ahead a plan reaches, the wider the band, and the less any single imagined outcome should be trusted.
Stress test three: closing the loop under model bias
Our final test addresses decision sufficiency. We design a linear-quadratic regulator (LQR) around the stable hanging equilibrium, using the linearized pendulum as the design model, and we execute the resulting controller on the true nonlinear plant. Then we repeat the exercise with a deliberately biased model, one that believes torque is only 40 percent as effective as it really is, which is precisely the kind of error real identification pipelines produce when actuator gains are miscalibrated. The discretization here uses the exact matrix exponential of the augmented continuous-time system rather than a crude Euler approximation, a detail that matters: we want the only discrepancy between design model and plant to be the deliberate bias plus the nonlinearity, not an artifact of sloppy discretization.
## Exact zero-order-hold discretization of the hanging-equilibrium linearization
A_c = np.array([[0.0, 1.0], [-g / l, -c]])
B_c = np.array([[0.0], [1.0 / (m * l**2)]])
M_aug = np.zeros((3, 3))
M_aug[:2, :2] = A_c
M_aug[:2, 2:] = B_c
M_d = expm(M_aug * dt)
A_lin, B_lin = M_d[:2, :2], M_d[:2, 2:]
def lqr_gain(A, B, Q, R):
P = solve_discrete_are(A, B, Q, R)
return np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)
Q = np.diag([10.0, 1.0])
R = np.array([[1.0]])
K_true = lqr_gain(A_lin, B_lin, Q, R) # designed from the correct model
K_bias = lqr_gain(A_lin, 0.4 * B_lin, Q, R) # designed from a biased model
def closed_loop(K, x0, n_steps=150, u_max=3.0):
xs = np.zeros((n_steps + 1, 2))
us = np.zeros(n_steps)
x = x0.copy()
xs[0] = x
for t in range(n_steps):
u = float(np.clip(-(K @ x)[0], -u_max, u_max))
us[t] = u
x = true_step(x, u)
xs[t + 1] = x
return xs, us
x0_cl = np.array([0.35, 0.0])
xs_true, us_true = closed_loop(K_true, x0_cl)
xs_bias, us_bias = closed_loop(K_bias, x0_cl)
cost_true = np.sum(
xs_true[:-1, 0] ** 2 + 0.1 * xs_true[:-1, 1] ** 2 + 0.01 * us_true**2
)
cost_bias = np.sum(
xs_bias[:-1, 0] ** 2 + 0.1 * xs_bias[:-1, 1] ** 2 + 0.01 * us_bias**2
)Closed-loop cost, correct model : 2.107 Closed-loop cost, biased model : 2.267 Cost inflation from model bias : 1.08x
Here is the punchline of the entire chapter, delivered in one plot. Both controllers regulate the pendulum, because LQR is robust and the bias is not extreme. But they are not equally good. The controller designed from the biased model produces slightly larger oscillations and incurs about eight percent more accumulated cost on the true plant. The model was never scored on its prediction error here; it was scored on what its errors did to a real closed loop. That is decision sufficiency made concrete: the quantity that matters is not how closely matches in the abstract, but how much the mismatch costs the agent that trusts it. The small but repeatable difference also illustrates an important point: model bias can alter behavior even when both controllers eventually succeed, and those behavioral differences are often easier to diagnose than an aggregate prediction loss.
The bar chart condenses the trajectory comparison into the single number a decision-maker cares about: the biased model's controller pays measurably more for the same task, despite both controllers eventually succeeding. This is the operational definition's final clause rendered as arithmetic. A model can pass every open-loop accuracy check and still fail here, and conversely a modest improvement in model honesty converts directly into closed-loop savings.
You now hold a complete, if miniature, world model pipeline: experience collection, model fitting, open-loop stress tests, stochastic imagination, and closed-loop deployment with an eye on model bias. Every chapter that follows scales up one or more of these five stages. When a later chapter introduces a recurrent latent dynamics network or a foundation-scale video predictor, resist the temptation to treat it as a new kind of thing: it is still this pipeline, with each stage grown more powerful and each failure mode grown more expensive.