Learning Transition Models

Humans are world-modeling systems. We predict what will happen, imagine possible futures, simulate actions mentally before taking them, and plan using causal structure in the world. Hippocampal replay helps reuse experience for imagination and planning, while goal-directed control uses predicted outcomes to guide behavior.

This is why people can often act effectively in novel situations with relatively few samples. We do not need to try every bad option in the real world; we can ask "what if?" inside an internal model first.

Human world models

Brains predict before acting

Humans do not learn only by trial and error. We imagine, replay, and plan through causal models of the world.

Hippocampus

Replay and imagination

Past experience can be replayed and recombined into possible futures, including during sleep.

World model

Causal structure

The brain predicts how actions change the environment before we commit to them.

Prefrontal cortex

Goals guide action

Goals plus predicted outcomes help choose flexible behavior in novel situations.

Fewer real samples

Safer mental trials

Better transfer to new situations

Model-based RL borrows this instinct: predict what will happen, reason about alternatives, then choose an action.

Model-based reinforcement learning starts with a different question from model-free RL. Instead of only asking "Which action has the highest learned value?" the agent also asks "If I take this action, what will happen next?" A learned model gives the agent a small internal simulator it can query before committing to behavior in the real environment.

In a warehouse, that model might predict whether a loaded robot will clip a shelf after turning. In a power grid, it might predict how demand and voltage change after a battery dispatch decision. In a game, it might predict the next board position and reward. The common structure is the same: the model learns how actions transform states, and often how those transitions produce rewards.

Agent contrast

Habit or imagination?

Choose an agent style to see what it learns and how it acts.

Model-based agent

Deliberative, goal-directed, flexible

Learns transition and reward models.

Imagines futures, scores alternatives, then chooses.

More sample efficient, but vulnerable to model bias.

Model-based agents can predict, plan, reason counterfactually, learn safely, and generalize, but only if their learned model is reliable.

A one-step transition model can be written as:

$\left(\hat{s}_{t+1},\hat{r}_{t+1}\right)=f_\theta(s_t,a_t)$

The input is the current state $s_t$ and action $a_t$ . The output is a predicted next state $\hat{s}_{t+1}$ and, when useful, a predicted reward $\hat{r}_{t+1}$ . The hat is important: this is not the environment's true next state. It is the model's best guess, learned from data.

Training often minimizes prediction error on observed transitions:

$\mathcal{L}_{\text{model}}=\left\|s_{t+1}-\hat{s}_{t+1}\right\|^2+\left(r_{t+1}-\hat{r}_{t+1}\right)^2$

That objective is only a starting point. A model can have low average loss and still be bad for decision-making if it misses rare crashes, ignores delayed consequences, or learns shortcuts that do not hold after deployment.

Equation walkthrough

Learn a dynamics model

Click each term to see how a learned transition model predicts the next state.

=(,)

Predicted next state

where the model thinks the world goes

The hat means this is imagined by the model, not guaranteed by the real environment.

Model-based RL asks f_theta what would happen before taking the action for real.

Equation walkthrough

Train the model by matching reality

Click each term to see how the loss compares predicted transitions with observed transitions.

Model loss

what the dynamics model minimizes

This is the training objective for the learned model. Smaller loss means the model's predicted next state and reward better match the real logged transition.

Terms

Low prediction loss helps, but planning still needs multi-step checks because small one-step errors can compound.

⚠

Model Bias

Model bias happens when the policy exploits errors in the learned model. The agent is not trying to be wrong; it is optimizing exactly what the model says. If the model confidently predicts impossible shortcuts, unrealistic recoveries, or missing side effects, planning can amplify those mistakes into real actions.

ℹ

Deployment Check

Before using a learned transition model to choose real actions, ask what the model was validated against: one-step prediction loss, multi-step rollouts, rare-event detection, uncertainty calibration, simulator-to-real transfer, and performance on states created by the new policy rather than only by the old data-collection policy.

Checkpoint

Why can a transition model with good average prediction loss still be risky for planning?

←PreviousDecision TransformersOffline RL Next→Planning with Learned DynamicsModel-Based RL and World Models