explainer
World models for robotics: learn dynamics, then plan
Fit a small action-conditioned world model, use it to choose a plan, and check what actually happens. See how perfect training predictions can hide a poor decision.
What you will learn
- Distinguish a model that predicts action outcomes from a policy that chooses actions.
- Fit a forward model from state, action, and next-state records.
- Trace how a planner selects a constant action using model predictions.
- Compare training error, rollout prediction error, and executed task error.
- Explain the limits of training ranges and hindsight comparisons.
Before you start
A model can predict every training example correctly and still guide a robot toward a poor action. Planning searches through possibilities. That search can favor the places where the model makes optimistic mistakes.
This lesson fits a small forward model, uses it to choose an action, and checks the resulting motion. Every training record and calculation is visible. The experiment uses synthetic dynamics so you can inspect the exact source of the error.
What does a world model predict?
A world model represents how an environment changes. For action selection, a useful model predicts what could happen after a specified action. A deterministic state model can take this form:
x̂ₜ₊₁ = f̂(xₜ, uₜ)
Here, x describes the current state, u describes an action, and the hat marks a prediction. A probabilistic version predicts a distribution over possible next states. Models can also predict observations, rewards, or episode termination when those outputs support the task.
The world model vs policy distinction concerns their jobs. A policy chooses an action from available information: u = π(observation, goal). A forward model evaluates the consequences of a supplied action; a planner can compare those predictions to choose one.
Chua and colleagues describe this sequence in PETS, their 2018 model-based reinforcement learning paper. They learn dynamics from transitions, predict trajectories under candidate actions, and use the predicted outcomes for decision making. Our example isolates that prediction-and-selection connection with two fitted coefficients.
Choose what the model represents
A robot model might work with joint positions, object poses, or a learned hidden representation. An image does not always reveal velocity, contact state, or an occluded object. A model that needs this missing information can use observation and action history to build its internal state.
Ha and Schmidhuber's World Models project from 2018 uses a recurrent model to predict a distribution over future latent representations, conditioned on actions and internal memory. Its controller then chooses actions from the available representation and memory. Their demonstrations use simulated environments.
Dreamer, published at ICLR 2020, learns latent dynamics and rewards, then learns behavior through imagined trajectories. It can perform that imagination in a compact hidden state without rendering an image at every predicted step. A world model therefore need not produce a watchable video.
A geometric world map answers questions about spatial structure. An action-conditioned predictor answers questions about changes over time. A system can contain both, and the robot foundation models lesson explains a separate question: how broadly a pretrained model can support reuse.
Fit a model and let it choose a plan
Start with An optimistic plan. The model fits all nine records, selects action 1.5, and predicts an endpoint of 2.43. Execution reaches 0.27 under the known synthetic dynamics.
The first plot compares one-step responses from x = 0. Three dots mark recorded actions at that state; the shaded interval includes unsampled values. The second plot follows the selected constant action through prediction and execution.
Change the target, horizon, or candidate action limit. These settings change the planning problem while leaving the training records and fit fixed. The rollout chart rescales its vertical axis; connecting lines only join discrete states.
Specify the dynamics and training records
The teaching environment has one fully observed state x and one action u. Both are dimensionless normalized quantities. Each step spans the same fixed interval, and actions remain between −1.5 and 1.5.
xₜ₊₁ = 0.8xₜ + uₜ − 0.4uₜ³
The cubic term makes the action effect bend downward at larger positive inputs. This is a deliberately simple nonlinear system. It supplies no claim about a particular motor, contact model, image generator, or robot experiment.
The dataset contains the nine combinations of x and u from −0.5, 0, and 0.5. Each record stores the current state, applied action, and resulting next state. The following table gives every target value:
| Current state x | Next state for u = −0.5 | Next state for u = 0 | Next state for u = 0.5 |
|---|---|---|---|
| −0.5 | −0.85 | −0.4 | 0.05 |
| 0 | −0.45 | 0 | 0.45 |
| 0.5 | −0.05 | 0.4 | 0.85 |
Training receives only these records. The planner receives only the fitted model. Known truth provides separate execution and retrospective checks; it never helps select the candidate action.
Understand a perfect but incomplete fit
Choose a model with no intercept and two coefficients:
f̂(x, u) = ax + bu
minimize Σᵢ₌₁⁹ (axᵢ + buᵢ − x′ᵢ)²
A least-squares solve gives a = 0.8 and b = 0.9. On the sampled actions, u³ = u/4, so the true action term becomes u − 0.4(u/4) = 0.9u. The linear and nonlinear rules agree on every training row.
Training next-state RMSE is therefore zero in exact arithmetic. Those records identify the two coefficients within the chosen linear model family. They do not establish that this family describes other actions.
At x = 0 and u = 1.5, the learned model predicts 1.35. The true next state is 1.5 − 0.4(1.5³) = 0.15. The parameter identification lesson develops the related distinction between estimating coefficients and checking the model that contains them.
Select actions using predicted outcomes
Every plan starts at x₀ = 0. The planner repeats one constant action for H steps, where H ranges from 1 to 4. It considers actions spaced by 0.25 inside the selected limit.
For each action, it feeds predicted states back into the fitted transition model. It then scores the predicted endpoint against target g:
x̂₀ = 0
x̂ₜ₊₁ = 0.8x̂ₜ + 0.9u
Ĵ(u) = (x̂H − g)²
The lowest predicted cost wins. Costs within 10⁻¹² count as numerical ties; the implementation chooses smaller absolute action, then smaller signed action. This makes the bounded grid search reproducible.
Execution repeats the selected action for the whole horizon. It never updates the plan from a new observation. PETS applies the first planned action, obtains new state information, and plans again using model predictive control.
The best audited grid action uses true dynamics after selection to score the same constant-action candidates. This hindsight comparison exposes a missed option within that grid. It makes no claim about the best varying action sequence or the globally optimal controller.
Compare the imagined plan with execution
Use target g = 2.5, horizon H = 2, and action limit 1.5. The fitted model selects u = 1.5. Follow both transitions:
| Step | Predicted state | Executed state | Prediction minus execution |
|---|---|---|---|
| 0 | 0 | 0 | 0 |
| 1 | 1.35 | 0.15 | 1.2 |
| 2 | 2.43 | 0.27 | 2.16 |
The predicted endpoint follows 0.8(1.35) + 0.9(1.5) = 2.43. Execution follows 0.8(0.15) + 1.5 − 0.4(1.5³) = 0.27. The small predicted terminal error hides a large executed error.
Use separate names for three measurements:
- Training RMSE: 0, measured on the nine fitted records.
- Rollout prediction RMSE: √((1.2² + 2.16²)/2) = 1.747226, measured on executed steps 1 and 2.
- Executed terminal error: 0.27 − 2.5 = −2.23, measured against the task target.
The shared reset contributes no sample to rollout RMSE. Terminal errors use endpoint minus target; the planner squares that error to form its cost.
The best true endpoint within this two-step constant-action grid comes from u = 1: x₂ = 1.08. It still falls short of 2.5. A hindsight improvement over the chosen plan does not establish that the stated target is reachable under these candidates.
Check action coverage and rollout length
Restrict the action range limits candidates to ±0.5. The planner selects u = 0.5, and both trajectories reach 0.81. That agreement follows from the two dynamics agreeing at this specific action; the target remains far away.
An unsampled interior action sets the target to 0.5. The planner selects u = 0.25 and predicts 0.405 after two steps. Execution reaches 0.43875.
That action lies inside the observed action range, but no training record uses it. At x = 0, its one-step prediction is 0.225 and its true outcome is 0.24375. An interval around training values describes a range; it does not supply observations at every point in that interval.
State coverage matters too. Reusing a sampled action can lead to a state-action pair absent from the nine records. This toy's error happens to depend only on u because it learns the state coefficient exactly.
Let eₜ = x̂ₜ − xₜ for a fixed action and shared reset. The exact error recurrence is:
eₜ₊₁ = 0.8eₜ + 0.4u³ − 0.1u, with e₀ = 0
Repeated prediction carries earlier error forward and adds a fresh action-model error. Changing H also changes the planner's chosen action, so the displayed error need not rise every time you increase H. A zero-action agreement provides a useful counterexample: from reset zero, both models remain at zero for every horizon.
Test predictions and decisions separately
A useful world-model evaluation answers several specific questions:
- On held-out transitions, how large is the next-state error for relevant states and actions?
- On held-out action sequences, how does prediction error change with horizon?
- When the planner chooses actions, what outcomes occur during separate execution?
- Which tasks, resets, action limits, and sensor conditions do those checks cover?
Keep evaluation data out of model fitting and tuning until the protocol allows it. Use grouped splits when nearby records share an episode or scene, as explained in robot data leakage. Report the scope of the held-out conditions alongside the score.
The default plan illustrates model exploitation: the optimizer favors an overly optimistic prediction and execution disappoints. No agent needs to understand or deliberately exploit the mistake. Action selection alone can concentrate decisions in a region where the model is wrong.
PETS uses probabilistic model ensembles and trajectory sampling to account for different sources of uncertainty during planning. Its Sections 4 and 5 describe those mechanisms. This two-coefficient lab produces point predictions only; it estimates no uncertainty or coverage probability.
A parameter bound or ensemble disagreement needs its own assumptions and checks. The model uncertainty lesson makes those conditions explicit. A low prediction error also leaves the task specification open: safety, energy, contact success, and other requirements need defined tests of their own.
Reproduce the model and plan in Python
This standard-library example fits the same nine records. It selects actions using the learned model, then audits the result using the separate true transition. All reported plans hold one action constant.
from math import sqrt
def truth(x, u):
return 0.8*x + u - 0.4*u**3
rows = [(x, u, truth(x, u))
for x in (-0.5, 0.0, 0.5)
for u in (-0.5, 0.0, 0.5)]
xx = sum(x*x for x, u, y in rows)
uu = sum(u*u for x, u, y in rows)
xu = sum(x*u for x, u, y in rows)
xy = sum(x*y for x, u, y in rows)
uy = sum(u*y for x, u, y in rows)
det = xx*uu - xu*xu
a = (xy*uu - uy*xu) / det
b = (uy*xx - xy*xu) / det
training_rmse = sqrt(sum((a*x + b*u - y)**2
for x, u, y in rows) / len(rows))
def model(x, u):
return a*x + b*u
def rollout(step, u, horizon):
states = [0.0]
for _ in range(horizon):
states.append(step(states[-1], u))
return states
def choose(step, actions, target, horizon):
best = None
for u in sorted(actions, key=lambda value: (abs(value), value)):
states = rollout(step, u, horizon)
cost = (states[-1] - target)**2
if best is None or cost < best[0] - 1e-12:
best = (cost, u, states)
return best
def experiment(target, horizon, limit):
actions = [k/4 for k in range(-int(4*limit), int(4*limit)+1)]
_, u, predicted = choose(model, actions, target, horizon)
executed = rollout(truth, u, horizon)
rmse = sqrt(sum((p-q)**2 for p, q in
zip(predicted[1:], executed[1:])) / horizon)
# Retrospective audit cannot affect the selection above.
_, best_u, best_states = choose(truth, actions, target, horizon)
return u, predicted[-1], executed[-1], rmse, best_u, best_states[-1]
print(f"fit: a={a:.6f}, b={b:.6f}, training RMSE={training_rmse:.6f}")
for name, target, limit in [("default", 2.5, 1.5),
("restricted", 2.5, 0.5),
("interior", 0.5, 1.5),
("zero", 0.0, 1.5)]:
u, predicted, actual, rmse, best_u, best_x = experiment(target, 2, limit)
print(f"{name}: u={u:.6f}, predicted={predicted:.6f}, actual={actual:.6f}")
if name == "default":
print(f"rollout RMSE={rmse:.6f}, audited u={best_u:.6f}, endpoint={best_x:.6f}")
Expected output:
fit: a=0.800000, b=0.900000, training RMSE=0.000000
default: u=1.500000, predicted=2.430000, actual=0.270000
rollout RMSE=1.747226, audited u=1.000000, endpoint=1.080000
restricted: u=0.500000, predicted=0.810000, actual=0.810000
interior: u=0.250000, predicted=0.405000, actual=0.438750
zero: u=0.000000, predicted=0.000000, actual=0.000000
Try it yourself
Exercise 1. At x = 0.5 and u = 0.25, calculate the learned and true next states. Does being inside both observed coordinate ranges make this a training example?
Show the prediction and coverage check
The learned prediction is 0.8(0.5) + 0.9(0.25) = 0.625. Truth gives 0.8(0.5) + 0.25 − 0.4(0.25³) = 0.64375. Prediction minus truth is −0.01875.
The dataset contains state 0.5 but never action 0.25. This pair falls inside the coordinate ranges and remains unobserved. Range membership alone cannot certify prediction accuracy.
Exercise 2. Start at zero, use target 2.5 and horizon 2, and compare the constant actions 1 and 1.5. Calculate both predicted costs and both executed costs. Which action wins each comparison?
Show the two action comparisons
For u = 1, the model predicts endpoint 1.62 and cost (1.62 − 2.5)² = 0.7744. Execution reaches 1.08, with cost 2.0164.
For u = 1.5, the model predicts endpoint 2.43 and cost 0.0049. Execution reaches 0.27, with cost 4.9729.
The model selects 1.5 in this comparison. The true outcomes favor 1. Both true endpoints miss the target; this comparison supplies no success claim or guarantee about other plans.
Sources and further study
- David Ha and Jürgen Schmidhuber, World Models, 2018: latent observation representations, action-conditioned recurrent prediction, and a separate controller.
- Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine, Deep Reinforcement Learning in a Handful of Trials Using Probabilistic Dynamics Models, NeurIPS 2018: learned transitions, uncertainty-aware prediction, trajectory sampling, and model predictive control.
- Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi, Dream to Control: Learning Behaviors by Latent Imagination, ICLR 2020: latent transition and reward models used to learn behavior through imagined trajectories.
Next, use robot policy evaluation to define what the executed plan must accomplish. A prediction becomes useful for robotics when the resulting decision survives the test you actually care about.