Model-Based Reinforcement Learning

Michael BrenndoerferJune 12, 202642 min read

Part of World Models Handbook

Explains how learned world models power planning, policy optimization, and value learning in RL, plus model bias, rollout error.

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

Model-Based Reinforcement Learning

Imagine a robot arm that must learn to insert a peg into a hole. A model-free reinforcement learning agent treats the arm as a black box: it tries actions, observes outcomes, and slowly accumulates value estimates through trial and error. On physical hardware this is painful. Every bit of learning costs real time, real wear on the motors, and real risk of damage. If a policy requires a hundred thousand training interactions, and each interaction takes a second of wall-clock time and a little life off the gearboxes, the arithmetic of deployment becomes unforgiving. A model-based agent takes a different route: it first learns how the world responds to its actions, then uses that learned model to rehearse. It can imagine thousands of peg insertions in the privacy of its own simulator, extract what it needs, and only occasionally return to reality to check its work. The expensive resource, real interaction, is spent sparingly; the cheap resource, computation inside the learned model, is spent liberally.

This is the promise of model-based reinforcement learning (MBRL): use a learned model to obtain more decision-making signal from each real interaction. Whether that promise is realized depends on the task, the data, and how the model is consumed. Model errors do not enter every decision uniformly; optimization can favor trajectories where the errors look advantageous. This chapter follows both sides of that bargain: how dynamics models support planning, policy optimization, and value learning, and how model bias, compounding rollout error, and model exploitation constrain those uses.

Throughout, keep two questions separate, because conflating them causes most MBRL confusion:

  • What does the model predict? States, observations, rewards, sometimes whole futures.
  • How does the agent use those predictions? To plan at decision time, to train a policy, or to train a value function.

The first question is about the model itself: its inputs, outputs, accuracy, and uncertainty. The second is about the decision machinery wrapped around it: whether predictions are consumed one step at a time, unrolled open-loop, or differentiated through. A given model can be paired with any of these consumers, and the same model can be excellent for one use and poor for another. Keeping the questions separate lets you diagnose failures precisely: when an MBRL system underperforms, the fault may lie in the model, in the consumer, or, most often and most subtly, in the interaction between the two.

Model-based reinforcement learning

An RL agent is model-based when it learns or is given a predictive model of the environment, a transition and reward model p^ϕ(st+1,rt∣st,at)\hat{p}_\phi(s_{t+1}, r_t \mid s_t, a_t), and uses that model as part of the decision-making process, rather than learning a policy or value function from real experience alone.

What Model-Based RL Adds to the Decision Loop

Building on the MDP formalism from Part II: Inference, Decisions, and Control, the agent faces a Markov decision process with transition kernel p(st+1∣st,at)p(s_{t+1} \mid s_t, a_t) and reward function r(st,at)r(s_t, a_t). A model-free decision method learns a policy πθ\pi_\theta or value function QθQ_\theta from experience without querying an explicit pp or rr during decision learning. An implementation may still train predictive auxiliary tasks; the distinction is whether a transition-and-reward model supplies the policy, value, or planning update. Model-based methods insert that intermediate object:

experience  ⟶  p^ϕ(s′,r∣s,a)⏟learned world model  ⟶  {π,Q,plan}⏟decision machinery\text{experience} \;\longrightarrow\; \underbrace{\hat{p}_\phi(s', r \mid s, a)}_{\text{learned world model}} \;\longrightarrow\; \underbrace{\{\pi, Q, \text{plan}\}}_{\text{decision machinery}}

Read the diagram left to right as a data flow, not a one-time pipeline: real experience refines the model, and the model feeds decision machinery whose behavior determines what data is collected next. A learned model is an explicit, reusable summary of experience. It can be queried for actions not taken in a recorded transition when its representation and data support such generalization, and it can generate additional optimization inputs. Those synthetic transitions are not new independent evidence. Model-free policies, values, representations, and replay data can also be reused; the practical distinction is that a model exposes a transition-and-reward query interface.

Chapters 3 and 4 of this part already showed you both ends of this pipeline: system identification taught you how to fit a dynamics model from data, and classical optimal control taught you how to act with a known model. Model-based RL couples these activities online. Adaptive and dual control also study this controller-model-data feedback; MBRL brings the same problem into reinforcement-learning systems with learned, often high-capacity models. As the controller changes, its state-action distribution can change, so validation must track the distribution on which the model will actually be used.

The components of a model-based RL system are:

  • Dataset D\mathcal{D} of real transitions (st,at,rt,st+1)(s_t, a_t, r_t, s_{t+1}) collected by earlier behavior. Its distribution often shifts as the policy changes, although it can be stationary under a fixed behavior policy.
  • World model p^ϕ\hat{p}_\phi and reward model r^ψ\hat{r}_\psi, fit to D\mathcal{D} by likelihood or regression. MBRL often augments ordinary system-identification losses with uncertainty models, multi-step objectives, online updates, or control-aware weighting.
  • Decision machinery, one of the three mechanisms described below: a planner, a policy, or a value function. This is the consumer whose needs the model must ultimately serve.
  • Data-collection strategy, which decides what real experience to gather next, often by exploring where the model is uncertain or where the current policy performs poorly. This component closes the loop between the agent and the dataset. Model uncertainty can support model-aware exploration, although model-based and model-free systems can also use the same behavior policy.

Learning a Model Is Not the Same as Using One

Here is the trap that separates people who have run MBRL experiments from people who have only read about them. You fit the model to minimize one-step prediction error on the data you collected, exactly as system identification recommends. The model gets impressively accurate. You then let the agent plan with it, and the agent behaves worse than a model-free baseline. What went wrong?

The first time you see this, it feels like a contradiction. The model is measurably good; the decisions built on it are measurably bad. The resolution is that "good" was measured under the wrong conditions.

Objective mismatch

A model trained purely for prediction accuracy can be suboptimal for control. A planner or learning rule queries the model under distributions shaped by its current policy, candidate actions, and optimizer, often through recursive rollouts, while prediction loss is commonly computed on one-step errors under the data-collection distribution. Even a one-step consumer can query unfamiliar state-action pairs. A model can therefore be excellent under its training distribution and useless, or actively misleading, under its decision-time query distribution. This is the objective mismatch identified by Lambert and colleagues, and it means the model's training objective should ideally reflect how the model will be consumed.

The mismatch has two root causes. First, distribution shift: as the policy improves, it visits states the dataset never contained, and the model extrapolates there. A neural network fit on random-walk data has no idea what the world looks like along the precise, purposeful trajectories a trained policy would take, and nothing in its training loss asked it to. Second, compounding: planners use the model recursively, feeding its own predictions back in as inputs, so one-step errors accumulate over the rollout horizon. Even a model whose one-step error is negligible on-distribution can drift badly over a twenty-step rollout, because step twenty's input is nineteen of its own guesses. Both deserve their own treatment, and we return to them in detail below.

There is also a subtler conceptual point, foreshadowed in the world-model qualification test from Part I: The World-Model Idea: open-loop prediction quality and closed-loop decision usefulness are different properties. A model can pass every prediction benchmark you throw at it and still produce bad policies, because policies live in a different part of the distribution than prediction benchmarks do. A benchmark scores the model where the data was; a policy uses the model where the optimizer steered. We will develop evaluation criteria for this in Part XI: Evaluation and Understanding of the handbook; for now, hold onto the distinction as a design principle. Whenever you train a model inside a decision system, ask not only "how accurate is it?" but "accurate where, and under what distribution of queries?"

Three Ways to Put a Model to Work

A useful taxonomy groups model use into three mechanisms: decision-time planning, policy optimization through imagined trajectories, and value learning from imagined experience. Many systems combine them, and some auxiliary-model or representation-learning methods do not fit the categories cleanly.

The most direct use: at every decision point, optimize an action sequence through the model, execute a prefix, and repeat. This is model predictive control (MPC), which you met in classical form in the optimal control chapter; the only change is that the dynamics come from p^ϕ\hat{p}_\phi instead of a hand-derived ODE:

at:t+H−1∗  =  arg⁡max⁡at,…,at+H−1  Ep^ϕ ⁣[∑k=0H−1γk r^ψ(st+k,at+k)],st+k+1∼p^ϕ(⋅∣st+k,at+k).a_{t:t+H-1}^* \;=\; \arg\max_{a_{t},\dots,a_{t+H-1}} \; \mathbb{E}_{\hat{p}_\phi}\!\left[\sum_{k=0}^{H-1} \gamma^k \, \hat{r}_\psi(s_{t+k}, a_{t+k})\right], \qquad s_{t+k+1} \sim \hat{p}_\phi(\cdot \mid s_{t+k}, a_{t+k}).

This objective is risk-neutral: it scores a candidate by expected return under the model. Risk-sensitive or pessimistic MPC replaces that expectation with another functional. After executing at∗a_t^*, the agent observes the next state and replans. This feedback usually limits exposure to open-loop drift, but it does not guarantee robustness to a biased model or make the inner optimization easy.

For continuous actions, sampling-based optimizers such as random shooting, CEM, and MPPI are common, alongside gradient methods and trajectory optimizers such as iLQR. Stochastic rollouts can still provide pathwise gradients through reparameterization or score-function estimates; sampling is attractive when such gradients are unavailable, unstable, or not worth their cost. PETS (Chua et al., 2018) paired CEM with a probabilistic neural-network ensemble and reported strong sample efficiency on its continuous-control benchmarks. Tree search is important for many discrete-action problems; MuZero searches in a learned latent model, as discussed in Part VIII: Decision-Centric Research Lineages.

Planning has a distinctive profile:

  • No separate policy training is required. Plain online MPC selects actions directly, although learned proposal policies and planner warm starts are common. With learned nonlinear dynamics, the inner optimization is often nonconvex; linear MPC with convex costs and constraints can instead be convex. Numerical behavior and closed-loop stability depend on the full formulation and solver.
  • Decision-time compute. Sampling budgets range widely with the optimizer, horizon, model, and hardware. This can be difficult under hard real-time constraints, although learned-model MPC can run at high rates in suitably engineered systems.
  • Short horizons are natural. Because reality corrects the state estimate after each executed action, MPC can succeed with horizons where open-loop planning would fail. The model is exposed to at most HH consecutive predicted transitions even when the task lasts much longer.

Optimizing a Policy Through the Model

The second mechanism treats the model as a differentiable simulator and backpropagates value through it to train a policy directly. For the expression below, assume a deterministic differentiable transition model, a deterministic differentiable policy ak=πθ(sk)a_k=\pi_\theta(s_k), and a differentiable modeled reward. Then s^t+k+1=fϕ(s^t+k,πθ(s^t+k))\hat{s}_{t+k+1} = f_\phi(\hat{s}_{t+k}, \pi_\theta(\hat{s}_{t+k})), the finite-horizon return is differentiable in θ\theta, and you can compute gradients through the imagined trajectory. A stochastic policy instead needs reparameterization or another valid gradient estimator.

dJdθ=∑k=0H−1γk[∂r^k∂skdskdθ+∂r^k∂ak(∂πθ∂θ+∂πθ∂skdskdθ)],\frac{dJ}{d\theta} = \sum_{k=0}^{H-1}\gamma^k\left[ \frac{\partial \hat r_k}{\partial s_k}\frac{d s_k}{d\theta} + \frac{\partial \hat r_k}{\partial a_k} \left(\frac{\partial \pi_\theta}{\partial\theta} + \frac{\partial \pi_\theta}{\partial s_k}\frac{d s_k}{d\theta}\right) \right],

where ak=πθ(sk)a_k=\pi_\theta(s_k) and the state sensitivity obeys the corresponding recursion through fϕf_\phi. These sensitivities multiply repeated model Jacobians, so they can vanish or explode with horizon. PILCO (Deisenroth and Rasmussen, 2011) used Gaussian-process dynamics, approximate analytic moment propagation, and analytic policy gradients, solving its reported control tasks with very little interaction. Its posterior uncertainty is meaningful under its modeling assumptions, but is not automatically calibrated. Gradient-based and sampling-based optimizers can both exploit model error; the result depends on their search pressure, regularization, and the geometry of the learned model.

The Dreamer family learns latent dynamics and trains actor and value functions from imagined latent rollouts. The original Dreamer propagates actor gradients through imagined dynamics; estimator details vary across versions. A latent state can omit task-irrelevant observation detail, but it is not necessarily lower-dimensional and does not shorten the temporal chain by itself. Part III: Representing Agents and Worlds develops the representation machinery.

Training Values on Imagined Experience

The third mechanism is the oldest and, in some ways, the most durable: use the model to generate synthetic transitions, and feed them to an otherwise ordinary model-free learner. The value backup

Q(s,a)←Q(s,a)+α[r+γmax⁡a′Q(s′,a′)−Q(s,a)]Q(s,a) \leftarrow Q(s,a) + \alpha\left[r + \gamma\max_{a'}Q(s',a') - Q(s,a)\right]

has the same algebra for a real or modeled transition. In practice, model bias and the synthetic-data distribution often require different rollout lengths, sampling ratios, weighting, or pessimism. The model supplies additional optimization inputs and can branch from real states into continuations not present in the replay buffer; it does not supply additional independent evidence.

This is the Dyna family, which gets its own section below. MBPO (Janner et al., 2019) is a modern relative that trains soft actor-critic on short model rollouts branched from real replay states and reported strong sample efficiency on its benchmarks.

A useful refinement is model-based value expansion (Feinberg et al., 2018): unroll HH modeled steps from a real state, sum predicted rewards, and bootstrap from a terminal value estimate. Increasing HH reduces reliance on the terminal bootstrap, but it can increase sampling variance and exposure to accumulated model bias; the balance depends on the estimator, model, and task. The rollout rewards are model predictions, not additional real reward observations.

The three mechanisms are compared in Table model-uses:

The three mechanisms for using a learned model, their decision-time compute cost, and their characteristic failure modes.
MechanismDecision-time computeHow it uses the modelCharacteristic failure
Planning (MPC, search)High, every stepRollout candidate actions, pick the bestExtrapolation errors at unfamiliar states
Policy gradients through the modelLow, after trainingDifferentiable rollouts for ∇θJ\nabla_\theta JExploiting smooth-but-wrong model regions
Value learning on imagined dataLow, after trainingSynthetic transitions for backupsModel bias baked into QQ

These are not rivals so much as a menu. PETS combines a learned ensemble model with planning. Dreamer combines value learning with gradients through the model. MBPO combines Dyna-style imagined data with off-policy value learning. The table's failure column is a diagnostic aid: ask how the model is consumed, then ask how the method limits harmful distribution shift and model exploitation.

Sample Efficiency, Model Bias, and Model Exploitation

This section covers the central tradeoff of the chapter. Models buy sample efficiency, and they charge for it in bias. The two subsections that follow give each side its due, first the mechanisms by which models reduce the demand for real data, then the quantitative ways in which model error enters and accumulates. By the end, you should be able to identify the controlling tradeoffs and design evaluations that reveal whether a model-based system helps on a particular problem.

Why Models Improve Sample Efficiency

It is worth being precise about where the efficiency comes from, because "the model lets you imagine" is only part of the story. There are at least four distinct sources, and they interact:

  • Compression. A suitable parametric model can share structure across states rather than learn each tabular state-action pair independently. This statistical advantage depends on whether the hypothesis class captures structure that actually transfers.
  • Re-use and re-weighting. Both model-free replay and model-based training can revisit real transitions many times. A model adds structured branching: it can propose continuations and reweight where computation is spent.
  • Off-policy imagination. A generalized model can evaluate actions absent from a particular trajectory, but reliable counterfactual prediction still requires action coverage or defensible structural assumptions. State coverage alone is not enough.
  • Auxiliary predictive signal. Predicting observations or dynamics provides dense self-supervised targets even when reward is sparse. A learned reward model does not by itself densify sparse reward labels; that requires shaping, auxiliary objectives, or other added structure.

These potential benefits do not require long-horizon rollouts. A system may use a model only for short branching, auxiliary prediction, or one-step backups.

Compounding Rollout Error

Suppose the true and learned transition maps start from the same state, so e0=0e_0=0, and are evaluated under the same actions. Assume the learned map has uniform one-step error at most ε\varepsilon on the relevant region and the true transition map is LL-Lipschitz in state. If ehe_h is the state error after hh modeled transitions, then

eh+1≤Leh+ε,eH≤ε∑k=0H−1Lk={εLH−1L−1,L≠1,Hε,L=1.e_{h+1} \le L e_h + \varepsilon, \qquad e_H \le \varepsilon\sum_{k=0}^{H-1}L^k = \begin{cases} \varepsilon\dfrac{L^H-1}{L-1}, & L\ne 1,\\[6pt] H\varepsilon, & L=1. \end{cases}

An error injected kk transitions before the end is multiplied by at most LkL^k. The regimes differ: the bound stays finite as HH grows when L<1L<1, grows linearly when L=1L=1, and grows exponentially only when L>1L>1.

The bound is easiest to internalize as a curve. Fix a one-step error of ε=0.01\varepsilon = 0.01, a deliberately small value that would look excellent on any one-step benchmark, and plot the worst-case rollout error against horizon for a few values of the Lipschitz constant LL:

In[3]:
Code
import numpy as np

EPSILON = 0.01
HORIZONS = np.arange(1, 51)
LIPSCHITZ_VALUES = [0.8, 1.0, 1.2]


def rollout_error_bound(L, horizons, epsilon=EPSILON):
    if np.isclose(L, 1.0):
        return epsilon * horizons
    return epsilon * (L**horizons - 1.0) / (L - 1.0)


bound_curves = {L: rollout_error_bound(L, HORIZONS) for L in LIPSCHITZ_VALUES}
Out[4]:
Visualization
Three state-error bounds: a saturating curve for L 0.8, a linear curve for L 1, and a steep exponential curve for L 1.2.
Worst-case open-loop state-error bound against horizon with one-step error 0.01. Contractive dynamics (L = 0.8) approach a finite bound, L = 1 grows linearly, and expansive dynamics (L = 1.2) grow exponentially.

The plot does not say that every rollout diverges. It isolates how transition sensitivity and repeated one-step error interact under explicit assumptions. Short rollouts, replanning, and value bootstrapping reduce exposure to the expansive regime.

Discounting supplies a separate amplification factor. If εB\varepsilon_B is a uniform one-step Bellman-operator error, the induced fixed-policy value error is bounded on the order of εB/(1−γ)\varepsilon_B/(1-\gamma). That statement does not follow from an unspecified transition prediction error. Simulation lemmas based on reward error and transition-kernel error usually contain a reward term proportional to (1−γ)−1(1-\gamma)^{-1} and a transition term proportional to (1−γ)−2(1-\gamma)^{-2}, with constants determined by reward bounds, the norm convention, and policy selection.

The amplification factor 1/(1−γ)1/(1-\gamma) is worth seeing as a curve rather than a formula, because its growth is not gentle:

In[5]:
Code
import numpy as np

GAMMA_GRID = np.linspace(0.50, 0.999, 300)
amplification = 1.0 / (1.0 - GAMMA_GRID)
gamma_marks = np.array([0.9, 0.99, 0.999])
mark_amplification = 1.0 / (1.0 - gamma_marks)
Out[6]:
Visualization
Rising Bellman-error amplification curve with markers at factors 10, 100, and 1000.
Amplification factor 1/(1-gamma) for a uniform Bellman-backup error, with markers at gamma = 0.9, 0.99, and 0.999. A fixed value-error budget therefore permits ten times less backup error at each successive marker.

The markers make the arithmetic concrete: moving from γ=0.9\gamma=0.9 to 0.990.99 multiplies a fixed Bellman error by ten, so a fixed value-error budget permits ten times less one-step backup error. Moving to 0.9990.999 tightens the budget by another factor of ten.

Let's make compounding concrete with a small experiment. We take a damped pendulum, fit two models to a few thousand random-policy transitions, a linear model that is structurally misspecified for the nonlinear dynamics and a more flexible small neural network, and compare their one-step error against their open-loop rollout error. The experiment measures the resulting errors directly; it does not separately identify approximation bias, estimation error, and recursive feedback.

In[7]:
Code
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.neural_network import MLPRegressor

rng = np.random.default_rng(0)

DT = 0.1


def pendulum_next(s, a):
    """Deterministic damped pendulum: state is (angle, angular velocity)."""
    theta, omega = s[0], s[1]
    return np.array(
        [
            theta + DT * omega,
            omega + DT * (-np.sin(theta) - 0.1 * omega + a),
        ]
    )


rng = np.random.default_rng(0)


def collect_transitions(n_traj, horizon):
    S, A, S2 = [], [], []
    for _ in range(n_traj):
        s = np.array([rng.uniform(-np.pi, np.pi), rng.uniform(-2.0, 2.0)])
        for _ in range(horizon):
            a = rng.uniform(-2.0, 2.0)
            S.append(s)
            A.append(a)
            S2.append(pendulum_next(s, a))
            s = S2[-1]
    return np.array(S), np.array(A), np.array(S2)


S, A, S2 = collect_transitions(100, 20)  # 2000 random-policy transitions
X = np.hstack([S, A[:, None]])
lin = LinearRegression().fit(X, S2)
mlp = MLPRegressor(
    hidden_layer_sizes=(64, 64), max_iter=5000, random_state=0
).fit(X, S2)

The raw linear model cannot represent the sin⁡(θ)\sin(\theta) term exactly. The neural network is more flexible, but capacity alone does not guarantee lower held-out error. We therefore evaluate both models on exactly the same held-out initial states and action sequences.

In[8]:
Code
eval_rng = np.random.default_rng(2026)


def predict_linear(s, a):
    return lin.predict(np.append(s, a)[None, :])[0]


def predict_mlp(s, a):
    return mlp.predict(np.append(s, a)[None, :])[0]


N_TEST_TRAJ, TEST_HORIZON = 50, 30
test_initial_states = np.column_stack(
    [
        eval_rng.uniform(-np.pi, np.pi, size=N_TEST_TRAJ),
        eval_rng.uniform(-2.0, 2.0, size=N_TEST_TRAJ),
    ]
)
test_actions = eval_rng.uniform(-2.0, 2.0, size=(N_TEST_TRAJ, TEST_HORIZON))


def evaluate_model(predict):
    n_traj, horizon = test_actions.shape
    onestep_err = np.zeros(horizon)
    rollout_err = np.zeros(horizon)
    for initial_state, actions in zip(test_initial_states, test_actions):
        s_true = initial_state.copy()
        s_model = s_true.copy()
        for t in range(horizon):
            a = actions[t]
            s_true_next = pendulum_next(s_true, a)
            onestep_err[t] += np.mean((s_true_next - predict(s_true, a)) ** 2)
            rollout_err[t] += np.mean((s_true_next - predict(s_model, a)) ** 2)
            s_true, s_model = s_true_next, predict(s_model, a)
    return onestep_err / n_traj, rollout_err / n_traj


err_lin_onestep, err_lin_rollout = evaluate_model(predict_linear)
err_mlp_onestep, err_mlp_rollout = evaluate_model(predict_mlp)
horizons = np.arange(1, TEST_HORIZON + 1)

Both models see the same held-out trajectories. One-step error always starts from the true state; rollout error feeds each model's prediction back as its next input.

Out[9]:
Console
Linear model  one-step MSE (steps 1-5): [0.0014 0.0013 0.0013 0.0012 0.0011]
MLP model     one-step MSE (steps 1-5): [0.0015 0.0014 0.0011 0.0013 0.0016]
Linear model  rollout MSE  at step 30:  1.234
MLP model     rollout MSE  at step 30:  1.016
Out[10]:
Visualization
Two similar-scale one-step error curves evaluated on identical trajectories.
Paired held-out one-step MSE for the linear and neural pendulum models. Their errors are of the same order, with no consistent winner over the first five steps.
Two rising rollout-error curves on a log scale, with the neural curve lower at long horizons.
Paired held-out open-loop rollout MSE. Error grows for both models, while the neural model has lower error by step 30 despite its weaker one-step result.

The panels support a narrower and more useful conclusion. On these paired trajectories, the two models have one-step errors of the same order, with no consistent winner over the first five steps. By step 30 the neural model has the lower rollout error. One-step ranking therefore does not determine long-horizon ranking, and recursive use must be evaluated directly. A trained policy would induce another test distribution, so these random-action results should not be generalized beyond the displayed experiment.

Model Exploitation and Distribution Shift

Compounding error degrades rollouts passively. A nastier failure mode is active: the decision machinery seeks out the model's errors.

Model exploitation

Model exploitation occurs when optimization favors trajectories that look good under the learned model because of model error, for example an overestimated reward or implausibly convenient dynamics. Greater search pressure can increase this failure, but optimizer choice alone does not determine it.

The interaction is between optimization and distribution shift. A planner may steer toward unfamiliar states or actions where the model is unreliable; those regions need not be the true high-reward regions. If optimistic errors dominate there, additional planning can worsen real performance even while imagined return improves.

One common defense is model uncertainty. Bootstrap ensembles often provide a useful proxy for epistemic uncertainty. Disagreement can expose sensitivity to finite data, but it can miss shared model-class bias and is not automatically calibrated.

To see both the value and the limitation, fit quadratic models to bootstrap resamples of noisy sine data. The members share the same misspecified feature class, so their spread measures resampling sensitivity within that class, not total predictive uncertainty.

In[11]:
Code
import numpy as np
from sklearn.linear_model import LinearRegression

rng = np.random.default_rng(7)
N_ENSEMBLE = 20


def make_features(x):
    return np.column_stack([np.ones_like(x), x, x**2])


x_train = rng.uniform(-1.0, 1.0, size=60)
y_train = np.sin(2.0 * np.pi * x_train) + rng.normal(0.0, 0.1, size=60)
F_train = make_features(x_train)

x_eval = np.linspace(-1.6, 1.6, 200)
F_eval = make_features(x_eval)

ensemble_preds = np.zeros((N_ENSEMBLE, len(x_eval)))
for m in range(N_ENSEMBLE):
    idx = rng.integers(0, len(x_train), size=len(x_train))
    coef, *_ = np.linalg.lstsq(F_train[idx], y_train[idx], rcond=None)
    ensemble_preds[m] = F_eval @ coef

pred_mean = ensemble_preds.mean(axis=0)
pred_std = ensemble_preds.std(axis=0)
true_curve = np.sin(2.0 * np.pi * x_eval)
covered = (x_eval >= -1.0) & (x_eval <= 1.0)
Out[12]:
Visualization
True sine curve, quadratic ensemble mean, and a widening bootstrap spread that does not cover all shared bias.
Quadratic bootstrap-ensemble mean and two-standard-deviation spread for noisy sine data observed only on [-1, 1]. Spread widens outside the data range, yet the ensemble can agree while sharing substantial misspecification inside it.

The widening band outside [−1,1][-1,1] is a useful warning about extrapolation. The dashed truth also exposes the limitation: a thin band is not a certificate of correctness when all members share a misspecified quadratic class. An uncertainty penalty or rollout truncation can still be useful, but it should be calibrated and checked against control-relevant held-out data.

PETS averages over ensemble rollouts in its planner; many other methods add an uncertainty penalty directly to the objective. Other common defenses include:

  • Short rollouts and branching. Generate imagined data only a few steps from real states, then bootstrap with the value function. This bounds compounding by construction: every imagined trajectory starts from ground truth and stays short. This is MBPO's core insight and, as we saw, it is also implicit in MPC's re-planning.
  • Pessimism. One common implementation subtracts an uncertainty estimate from imagined reward so that unfamiliar, possibly-wrong regions look bad rather than promising. Other implementations use lower-confidence values, robust transition sets, constraints, or worst-case objectives. These constructions connect to the same pessimism principle used in offline RL.
  • Model trust diagnostics. Before letting the agent act on a model, test it the way the qualification test from Part I: The World-Model Idea suggested: not only on held-out data, but on the distribution the policy will induce, and at the rollout lengths the decision machinery will use. A model validated only on its training distribution has not been validated at all, as far as the decision machinery is concerned.

Planner error and model error are distinct. A precise optimizer can expose optimistic model errors, while an approximate optimizer can simply be suboptimal. Robust objectives, constraints, uncertainty penalties, and real-environment validation are more principled controls than assuming that fewer optimizer iterations always make an uncertain model safer.

Dyna: Integrating Real and Imagined Experience

Richard Sutton presented an early version of Dyna at ICML in 1990 and developed the framework further in 1991. Dyna remains a clear template for algorithms that interleave real experience, model learning, and planning updates. It is an important lineage for modern MBRL, not an exhaustive description of every model-based method.

The Algorithm

Each cycle of interaction does three things:

  1. Act in the real environment: choose ata_t from the current policy, observe (rt,st+1)(r_t, s_{t+1}).
  2. Learn the model: update p^ϕ\hat{p}_\phi and r^ψ\hat{r}_\psi with the new transition.
  3. Learn the value, twice: apply one real Q-learning update from the observed transition, then apply nn additional planning updates, each drawn from the model: sample a previously visited (s,a)(s, a), query the model for (r^,s^′)(\hat{r}, \hat{s}'), and apply a backup as if it had really happened.
In[21]:
Code
## Dyna's inner loop, in pseudocode
for each real step:
    a, s2, r, done = act_and_observe(s)
    model[(s, a)] = (s2, r, done)       # learn the model
    q_update(s, a, r, s2, done)         # real backup
    for _ in range(n_planning):
        ps, pa = sample_visited_pair()
        ps2, pr, pdone = model[(ps, pa)]  # imagined transition
        q_update(ps, pa, pr, ps2, pdone)  # imagined backup
    s = s2                              # advance the real state

The nn planning updates are what distinguishes Dyna from plain Q-learning. Each one propagates value information backward through the state space without touching the environment. The imagined backups are interleaved with the real ones, so real updates continue to provide ground-truth corrections while modeled updates propagate information more deeply. Those real corrections do not guarantee protection from bias when inaccurate modeled backups dominate. Notice also the division of labor: the model absorbs the raw statistics of the environment, while the value function absorbs the long-horizon consequences. Neither has to do the other's job.

Why It Works and When It Fails

The tabular Q-learning backup is local: a transition moves value information one edge from s′s' to ss. Dyna repeats such backups from stored model entries between real interactions, so reward information can propagate through already visited state-action pairs without revisiting them in the environment. The size of the gain depends on exploration, planning samples, hyperparameters, and model accuracy.

Two assumptions carry this argument:

  • The Markov assumption from Part II: Inference, Decisions, and Control, Chapter 1 must hold for the model's inputs. If the environment is partially observable and you model raw observations, the model is predicting from an insufficient statistic, and the imagined backups will be systematically wrong. The model may fit the data well in the one-step sense while still being useless for planning, because no function of the observation alone can capture the hidden state the future depends on. Belief-state modeling from Chapter 2 of this part is the fix, and latent-state models in Part III: Representing Agents and Worlds are the learned version of the same idea.
  • The model must be right where the planning distribution samples it. In tabular Dyna with a deterministic environment, the model is exact after one visit per pair, so imagined data is as good as real data. The moment the model is approximate, imagined backups inject model bias directly into QQ, and the value function becomes a compounding of model error over the effective planning depth. Dyna amplifies a good model and a bad one. This is the sample-efficiency-versus-bias tradeoff in its purest form: the planning multiplier nn scales both the honest signal and the bias.

Dyna is closely related to prioritized sweeping (Moore and Atkeson, 1993), which spends planning updates where the value estimate changed most and is an ancestor of prioritized replay in model-free deep RL. And the modern branch-based methods are best read as Dyna with a smarter sampling rule: instead of replaying visited pairs, sample real states from a replay buffer, branch short imagined continuations from them, and train a deep value function on the mixture. Dreamer's actor and critic train on imagined latent trajectories, while its world model is still fit to real replay. Read that way, three decades of MBRL research look like a long refinement of a single question: from which states, and over what horizons, should imagined experience be generated?

Worked Example: One Trajectory in a Six-by-Six Room

Let's make the Dyna effect concrete before coding it. The room is a 6×66 \times 6 grid. The agent starts in the top-left corner, the goal is in the bottom-right, each nonterminal step costs −1-1 reward, and the transition into the goal pays +10+10 instead. The optimal path is ten moves, for a return of 10−9=1.010 - 9 = 1.0. The geometry matters: the reward sits at the far corner of the state space from the start, so the value information must travel the full diagonal, and any method that propagates value one edge per real visit will pay the full cost of that distance in real interactions.

Now suppose the agent's first episode is an undirected random walk that stumbles into the goal after 60 moves, collecting a return of about −49-49. What does each kind of agent extract from that single trajectory?

The pure Q-learning agent performs 60 update operations, although repeated state-action pairs mean fewer than 60 distinct entries may change. The sample does not guarantee that its second episode will improve.

With plan_steps=30, the Dyna agent performs 30 model backups after each real step, or 1,800 planning backups during a 60-step episode. This implementation samples only state-action pairs that were actually visited. It replays their learned outcomes in a new order; it does not infer untried actions or leave the visited-pair support. Repeated backups can nevertheless propagate reward information backward through the stored graph before another real episode begins.

The honest caveat: this gridworld has a deterministic environment, so the learned model is exact. That is the best case for Dyna, and we chose it deliberately so you can see the mechanism cleanly, with zero model bias contaminating the comparison. Whatever gap appears between the curves is attributable purely to the amplification of honest experience, not to any advantage the model gained by being wrong. The experiment below quantifies exactly how large the advantage is in real interactions.

Code Implementation: Dyna Against Model-Free Q-Learning

We now build the worked example in full. The plan: implement the gridworld and a tabular Dyna agent, run three variants (no planning, 10 planning steps, 30 planning steps) across many random seeds, and plot episode return against real environment interactions, which is the resource we care about. Plotting against episodes instead would flatter nothing and show nothing: an episode is a variable-length bundle of interactions, and the whole point of model-based methods is their exchange rate between computation and reality.

First, the environment. It is deterministic, with bounded grid movement and a terminal goal state. The clipping at the walls means attempted moves off the grid result in staying in place, which gives the agent a small but real amount of transition structure to learn about the boundary.

In[13]:
Code
from collections import defaultdict

GRID_SIZE = 6
START = (0, 0)
GOAL = (GRID_SIZE - 1, GRID_SIZE - 1)
ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]
STEP_REWARD = -1.0
GOAL_REWARD = 10.0


def env_step(state, action_idx):
    dr, dc = ACTIONS[action_idx]
    r = int(np.clip(state[0] + dr, 0, GRID_SIZE - 1))
    c = int(np.clip(state[1] + dc, 0, GRID_SIZE - 1))
    next_state = (r, c)
    done = next_state == GOAL
    reward = GOAL_REWARD if done else STEP_REWARD
    return next_state, reward, done

Next, the agent. Q maps each state to four action values, defaulting to zeros for unvisited states. The model dictionary maps visited (s,a)(s, a) pairs to outcomes; because the environment is deterministic, one observation per pair makes the model exact. The plan_steps parameter is Dyna's dial: set it to zero and this is plain Q-learning.

In[14]:
Code
def dyna_episode(
    Q,
    model,
    behavior_rng,
    planning_rng,
    alpha=0.5,
    gamma=0.95,
    eps=0.1,
    plan_steps=0,
):
    """Run one episode: epsilon-greedy control, tabular Q-learning,
    and plan_steps imagined backups per real step."""
    state = START
    total_reward, env_steps = 0.0, 0
    while state != GOAL and env_steps < 250:
        if behavior_rng.random() < eps:
            action_idx = int(behavior_rng.integers(len(ACTIONS)))
        else:
            greedy = np.flatnonzero(Q[state] == np.max(Q[state]))
            action_idx = int(behavior_rng.choice(greedy))
        next_state, reward, done = env_step(state, action_idx)

        target = reward if done else reward + gamma * np.max(Q[next_state])
        Q[state][action_idx] += alpha * (target - Q[state][action_idx])
        model[(state, action_idx)] = (next_state, reward, done)

        keys = list(model.keys())
        for _ in range(plan_steps):
            ps, pa = keys[planning_rng.integers(len(keys))]
            pnext, preward, pdone = model[(ps, pa)]
            ptarget = preward if pdone else preward + gamma * np.max(Q[pnext])
            Q[ps][pa] += alpha * (ptarget - Q[ps][pa])

        state = next_state
        total_reward += reward
        env_steps += 1
    return total_reward, env_steps

Planning pairs are sampled uniformly from visited pairs. Separate random-number streams keep exploratory actions paired across planning budgets until their learned values cause the policies to diverge; planning draws no longer advance the behavior stream. Random tie-breaking also avoids a fixed preference for action 0.

Now the experiment. We run each variant for 15 seeds and 120 episodes, recording the return of every episode and the cumulative count of real environment steps, then average across seeds. This cell is the expensive one; everything after it only rearranges these arrays.

In[15]:
Code
def run_experiment(plan_steps, n_seeds=15, n_episodes=120):
    """Return per-seed cumulative real steps and episode returns."""
    all_returns = np.zeros((n_seeds, n_episodes))
    all_cum_steps = np.zeros((n_seeds, n_episodes))
    for seed in range(n_seeds):
        behavior_rng = np.random.default_rng(seed)
        planning_rng = np.random.default_rng(10_000 + seed)
        Q = defaultdict(lambda: np.zeros(len(ACTIONS)))
        model = {}
        cum_steps = 0
        for ep in range(n_episodes):
            ret, steps = dyna_episode(
                Q, model, behavior_rng, planning_rng, plan_steps=plan_steps
            )
            cum_steps += steps
            all_returns[seed, ep] = ret
            all_cum_steps[seed, ep] = cum_steps
    return all_cum_steps, all_returns


steps_q, returns_q = run_experiment(plan_steps=0)
steps_d10, returns_d10 = run_experiment(plan_steps=10)
steps_d30, returns_d30 = run_experiment(plan_steps=30)

Before plotting, let's compute a crisp summary: for each variant, how many real environment interactions were needed before the smoothed return first crossed a fixed threshold. This converts the learning curves into a single interpretable number per method, one that can be compared and quoted, rather than eyeballed off a chart.

In[16]:
Code
def steps_to_threshold_per_seed(steps, returns, threshold=-2.0, window=10):
    """Per-seed real steps to the first qualifying moving-average return."""
    crossings = []
    kernel = np.ones(window) / window
    for seed_steps, seed_returns in zip(steps, returns):
        smooth = np.convolve(seed_returns, kernel, mode="valid")
        crossed = np.flatnonzero(smooth >= threshold)
        crossings.append(
            np.nan if len(crossed) == 0 else seed_steps[crossed[0] + window - 1]
        )
    return np.asarray(crossings)
Out[17]:
Console
Q-learning: 15/15 seeds reached -2; mean  846.8 +/-  21.9 real steps
   Dyna-10: 15/15 seeds reached -2; mean  273.1 +/-  14.5 real steps
   Dyna-30: 15/15 seeds reached -2; mean  237.0 +/-   9.6 real steps

The threshold is computed per seed before aggregation, so a crossing in one run cannot hide a failure in another. The plot below interpolates each seed's smoothed returns onto one common real-interaction grid before taking the mean and standard error. It does not extrapolate a seed before its first complete 10-episode window; a method's curve begins only where all 15 seeds contribute observations.

Out[18]:
Visualization
Learning curves with standard-error bands showing both Dyna variants improving earlier than Q-learning on a shared interaction axis.
Mean 10-episode-smoothed return on a common real-interaction grid for Q-learning and two Dyna variants across 15 seeds; shaded bands show one standard error. Each curve begins only once all 15 seeds for that method have completed a smoothing window. In this deterministic visited-pair model, both planning variants reach a -2 return threshold in about one-third of Q-learning's real interactions.

In this exact deterministic tabular setting, the two planning variants reach the displayed threshold in roughly one-third of Q-learning's real interactions. The difference between 10 and 30 planning updates is smaller than the difference between either Dyna variant and no planning. This is one experiment, not a monotonic law: with an approximate model, more planning can reinforce bias, while with an exact model and a compute-free objective, no theorem here requires a finite optimal planning budget.

You now have a complete Dyna-style MBRL system: model learning, imagined backups, and explicit accounting of real interactions. Other model-based methods may instead plan directly or differentiate through a model without this value-learning loop.

Key Parameters

The key parameters for the Dyna implementation are:

  • plan_steps: Number of imagined backups per real environment step. Set it to 0 for plain Q-learning; larger values spend more computation reusing the stored model. The displayed run shows a smaller gain from 10 to 30 than from 0 to 10, but other problems can improve, saturate, or degrade differently. This is the sample-efficiency-versus-compute dial, and, when the model is imperfect, the sample-efficiency-versus-bias dial as well.
  • alpha: Learning rate for real and imagined Q-learning updates. It interacts with plan_steps, but their product is only a rough update-budget heuristic because repeated sampling and Bellman contraction cause saturation.
  • gamma: Discount factor. Larger values reduce attenuation of distant rewards and weaken Bellman contraction; they do not change graph distance or guarantee that planning helps more.
  • eps: Epsilon-greedy exploration rate during data collection, which influences the model's visited-pair coverage. Greedy actions, random tie-breaking, dynamics, and the start state also shape that coverage. Only real interaction adds new pairs in this implementation; planning cannot leave the support already stored in the model.
  • n_seeds, n_episodes: Number of random seeds and episodes. Averaging helps, while paired random streams, dispersion bands, and per-seed threshold calculations make the comparison interpretable.

Limitations and Impact

The risk of model error is structural when predictive models are learned from finite data, even though a finite deterministic tabular model can become exact on each observed pair. Open-loop state error can remain bounded for contractive dynamics, grow linearly at L=1L=1, or grow exponentially for L>1L>1. Discounted value error has its own scaling and depends on whether error is measured in rewards, transitions, or Bellman backups. Common defenses include uncertainty estimates, short rollouts, pessimism, constraints, and validation on the decision distribution, but no one defense is universal. Model-based methods also trade data for computation. Under partial observation, the model needs a sufficient statistic such as a belief or recurrent state; one-step validation on aliased observations can otherwise look good while planning fails.

Model-based RL has produced influential sample-efficient systems. PILCO reported solving a physical cart-pole task with 17.5 seconds of interaction. PETS, MBPO, and Dreamer extended different model-use mechanisms to neural models and high-dimensional observations. DreamerV3 learned Minecraft diamond collection without human data by learning behavior from imagined rollouts while still collecting and acting on real environment experience. Dyna remains one important lineage among the decision-centric methods in Part VIII: Decision-Centric Research Lineages and the planning agents of Part VII: Planning and Agency.

For you as a practitioner, the chapter leaves you with a design checklist rather than a single recipe:

  • Decide how the model is used, using planning, policy optimization, and imagined value learning as a practical taxonomy rather than an exhaustive partition.
  • Keep imagined rollouts within the model's validated horizon. When the consumer uses a terminal value function, let that value carry the burden beyond the rollout.
  • Use an uncertainty-aware diagnostic or objective when model uncertainty is material and can be estimated; an exact known model does not need a proxy for uncertainty.
  • Evaluate the model on the decision distribution, not just held-out one-step accuracy, remembering the qualification test from Part I: The World-Model Idea.
  • Report and monitor the real-interaction budget, the number that model-based RL exists to shrink.

Summary

  • Model-based RL inserts a dynamics and reward model between experience and decisions. When data and representation support the query, that model can evaluate alternative actions and generate synthetic optimization inputs.
  • A useful taxonomy asks whether the model supports decision-time planning, policy optimization through imagination, or value learning from synthetic transitions; many algorithms combine these uses.
  • Under the stated same-start, uniform-error, and Lipschitz assumptions, the worst-case open-loop state-error bound is finite, linear, or exponential according to transition sensitivity. Discounted performance bounds must likewise state the error metric and assumptions.
  • Model exploitation occurs when optimization favors optimistic model errors. Ensembles, short branched rollouts, pessimism, constraints, and decision-distribution validation are common defenses, none of which certifies correctness alone.
  • Dyna interleaves real updates, model learning, and backups from visited model entries. Direct planners and differentiable policy-search methods need not follow this loop.
  • In the deterministic gridworld experiment, both Dyna variants reached a smoothed return of −2-2 in about one-third of Q-learning's mean real interactions across seeds. That result is specific to the implementation and hyperparameters shown.

The next part of the book turns to the representation question this chapter deliberately set aside: what should the model predict over? Raw observations are often an awkward target for control-relevant prediction, and the road to latent, memory-equipped, object-centric world models starts with a closer look at where observations come from.

Quiz

Ready to test your understanding? Take this quick quiz to reinforce what you've learned about model-based reinforcement learning.

Model-Based Reinforcement Learning

Question 1 of 80 of 8 completed
What does the 'objective mismatch' problem describe in MBRL?

Comments

No comments yet. Be the first to share your thoughts!

Reference

Citation details

Cite or share this article.

BIBTEXAcademic
@misc{brenndoerfer2026modelbased, author = {Michael Brenndoerfer}, title = {Model-Based Reinforcement Learning}, year = {2026}, url = {https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models}, organization = {mbrenndoerfer.com}, note = {Accessed: 2026-10-13} }
APAAcademic
Michael Brenndoerfer (2026). Model-Based Reinforcement Learning. Retrieved from https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models
MLAAcademic
Michael Brenndoerfer. "Model-Based Reinforcement Learning." 2026. Web. October 13, 2026. <https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models>.
CHICAGOAcademic
Michael Brenndoerfer. "Model-Based Reinforcement Learning." Accessed October 13, 2026. https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models.
HARVARDAcademic
Michael Brenndoerfer (2026) 'Model-Based Reinforcement Learning'. Available at: https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models (Accessed: October 13, 2026).
SimpleBasic
Michael Brenndoerfer (2026). Model-Based Reinforcement Learning. https://mbrenndoerfer.com/writing/model-based-reinforcement-learning-world-models

About the author

Continue with the full handbook

This chapter is part of World Models Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.

Explore World Models Handbook
Newsletter

Stay up to date

Get articles, book updates, and news delivered to your inbox.

No spam, unsubscribe anytime.

or

Join the community

Sign in to remove popups, track your reading progress, and join the discussion.