Model-Based Reinforcement Learning

Introduction and Context

Reinforcement learning (RL) studies how an agent learns to select actions through interaction with an environment. In model-free RL, the agent learns values or policies directly from experience, without explicitly learning how the environment changes. Model-based reinforcement learning (MBRL) takes a different approach: the agent learns, receives, or is given a model of the environment and uses that model to predict consequences, evaluate hypothetical actions, and plan.

A model-based agent separates two activities:

1. Learning: estimating how the environment behaves and, often, how rewards are generated.

2. Planning: using the estimated model to select actions or improve value estimates.

This separation is valuable when real-world interaction is expensive, dangerous, slow, or limited. A robot can test possible movements in simulation before moving its physical joints; an autonomous vehicle can predict the consequences of braking or changing lanes; and an industrial controller can optimize production without repeatedly experimenting on a live plant.

Model-based methods are not automatically superior. They depend on model quality. An inaccurate model can produce confident but harmful plans, a problem known as model bias. The central engineering challenge is therefore to exploit predictive models while managing uncertainty and computational cost.

Learning Objectives

By the end of this chapter, students should be able to:

• distinguish model-free and model-based reinforcement learning;
• formulate transition and reward models;
• explain planning through dynamic programming and Dyna-style methods;
• describe Monte Carlo tree search, trajectory optimization, and world models;
• analyze the effect of model error and uncertainty on planning;
• compare methods using sample efficiency, computational cost, robustness, and applicability;
• design a basic model-based RL workflow for a practical problem;
• identify safety, governance, and deployment concerns.

1. The Environment Model

Consider a Markov decision process (MDP) with state space \(\mathcal{S}\), action space \(\mathcal{A}\), transition dynamics, reward function, and discount factor \(\gamma\), where \(0 \leq \gamma < 1\). A stochastic environment can be represented by

\[

P(s'\mid s,a)=\Pr(S_{t+1}=s'\mid S_t=s,A_t=a)

\]

and

\[

R(r\mid s,a,s')=\Pr(R_{t+1}=r\mid S_t=s,A_t=a,S_{t+1}=s').

\]

A model may estimate both components. In a deterministic environment, these may reduce to functions \(s'=f(s,a)\) and \(r=r(s,a)\). A learned model is commonly written as

\[

\hat{s}_{t+1}=\hat{f}_\theta(s_t,a_t), \qquad \hat{r}_t=\hat{r}_\phi(s_t,a_t),

\]

where \(\theta\) and \(\phi\) are learned parameters. For stochastic or partially known systems, the model may output a probability distribution rather than a single prediction.

The value of a policy \(\pi\) is

\[

V^\pi(s)=\mathbb{E}_\pi\left[\sum_{k=0}^{\infty}\gamma^kR_{t+k+1}\mid S_t=s\right].

\]

If a model is available, the Bellman expectation equation is

\[

V^\pi(s)=\sum_a \pi(a\mid s)\sum_{s',r}p(s',r\mid s,a)\left[r+\gamma V^\pi(s')\right].

\]

The optimality equation is

\[

V^(s)=\max_a\sum_{s',r}p(s',r\mid s,a)\left[r+\gamma V^(s')\right].

\]

These equations assume the Markov property: the current state contains all information needed to predict future rewards and states. If observations omit important variables, a recurrent model, belief state, or state-estimation procedure may be required.

Explicit, Learned, and Hybrid Models

An explicit model is supplied using known system laws, such as rigid-body physics. A learned model is inferred from data. A hybrid model combines both, for example using known vehicle dynamics while learning tire friction or residual errors. Hybrid models are often attractive because structural knowledge improves data efficiency and constrains implausible predictions.

Models may operate at different levels:

• State-space models predict physical or symbolic states.
• Observation models predict raw observations such as images.
• Latent models predict a learned representation rather than pixels.
• Reward models estimate immediate utility or task progress.
• Termination models estimate failure, success, or episode-ending events.

A useful model must be judged not only by prediction accuracy but also by decision quality. Small errors in irrelevant variables may not matter, while a small error near a safety boundary can completely change the preferred action.

2. Planning with a Model

Planning means evaluating possible future action sequences using the model. For a fixed sequence \(a_t,a_{t+1},\ldots,a_{t+H-1}\), a deterministic model produces the predicted finite-horizon return

\[

\hat{G}_t=\sum_{k=0}^{H-1}\gamma^k\hat{r}_{t+k}+\gamma^H\hat{V}(\hat{s}_{t+H}),

\]

where \(H\) is the planning horizon and \(\hat{V}\) is a terminal value estimate. The action sequence may be selected by maximizing this quantity. The terminal value is important when the horizon is short; otherwise the planner may behave as if the world ends after \(H\) steps.

Planning methods differ in how they search:

• Dynamic programming systematically updates values across states.
• Tree search expands selected future branches.
• Trajectory optimization directly improves a sequence of continuous actions.
• Receding-horizon control executes only the first planned action, then replans using the new state.

The last strategy is called model predictive control (MPC). At each time step, the system plans over a finite horizon, applies the first action, observes the next state, and repeats. Replanning reduces the impact of prediction errors, although it does not remove them.

3. Dyna: Integrating Learning and Planning

Dyna, introduced by Sutton, combines direct reinforcement learning, model learning, and simulated planning. After each real transition \((s,a,r,s')\), the agent:

1. updates its value function from the real experience;

2. updates a model using the observed transition;

3. samples previously observed state-action pairs from the model;

4. performs additional simulated value updates.

For a tabular action-value function, a one-step Q-learning update from either real or simulated experience is

\[

Q(s,a)\leftarrow Q(s,a)+\alpha\left[r+\gamma\max_{a'}Q(s',a')-Q(s,a)\right],

\]

where \(\alpha\) is the learning rate. The formula applies when the target is based on the maximum next-action value and the transition is treated as a Q-learning target. For a policy-evaluation method, the target would instead use the policy’s action distribution.

Dyna improves sample efficiency because one real transition can generate many planning updates. However, if the model is wrong, simulated updates can reinforce incorrect beliefs. Dyna-Q+ addresses one form of insufficient exploration by adding a time-dependent exploration bonus to actions that have not been tried recently. The principle is optimism about potentially useful unexplored transitions, although the bonus must be designed carefully.

Illustrative Dyna Example

Suppose an agent observes that taking action \(a\) in state \(s\) gives reward 2 and leads to state \(s'\). Let \(Q(s,a)=1\), \(\alpha=0.5\), \(\gamma=0.9\), and \(\max_{a'}Q(s',a')=4\). The target is

\[

2+0.9(4)=5.6.

\]

Therefore,

\[

Q(s,a)\leftarrow 1+0.5(5.6-1)=3.3.

\]

Once the model stores this transition, the same update can be repeated as a simulated planning step. The simulated step does not create new environmental information; it reuses the learned transition to propagate value information more rapidly.

4. Monte Carlo Tree Search

Monte Carlo tree search (MCTS) builds a partial search tree from a current state. It is especially useful when the action space is discrete or can be sampled and when a complete model is available or learnable. A typical iteration contains four stages:

1. Selection: follow a tree policy from the root to a promising leaf.

2. Expansion: add one or more previously unrepresented actions or states.

3. Simulation or evaluation: estimate the value below the new node.

4. Backpropagation: update statistics along the path.

A common selection rule is the upper confidence bound for trees (UCT):

\[

\text{UCT}(i)=\frac{W_i}{N_i}+c\sqrt{\frac{\ln N}{N_i}},

\]

where \(W_i\) is accumulated return for child \(i\), \(N_i\) is the number of visits to that child, \(N\) is the number of visits to its parent, and \(c>0\) controls exploration. The first term favors actions with high estimated returns; the second favors less-visited actions. UCT assumes suitable reward scaling and is most naturally applied to a tree with repeated visits and stochastic uncertainty represented through sampling.

MCTS is successful in games because it can focus computation on promising branches rather than enumerate the whole tree. In physical control, however, branching can be enormous, actions may be continuous, and model errors can accumulate deeply. Learned value functions, progressive widening, action sampling, and short horizons are common adaptations.

5. Trajectory Optimization and Model Predictive Control

Trajectory optimization searches directly for an action sequence. Given a differentiable dynamics model, one may minimize a cost

\[

J(a_{0:H-1})=\sum_{t=0}^{H-1}\ell(s_t,a_t)+\ell_f(s_H),

\]

subject to

\[

s_{t+1}=f_\theta(s_t,a_t),

\]

where \(\ell\) is a stage cost and \(\ell_f\) is a terminal cost. Constraints may include actuator limits, collision avoidance, energy budgets, and state safety boundaries.

Methods such as iterative linear-quadratic regulation, differential dynamic programming, shooting methods, and the cross-entropy method optimize this problem in different ways. Gradient-based methods are efficient when the model is smooth and differentiable. Sampling-based methods are more tolerant of discontinuities and can work with black-box models, but may require many model evaluations.

In MPC, the optimized sequence is not executed in full. Only the first action is applied. After observing the next state, the optimization is repeated. This feedback loop is a major practical advantage over open-loop plans.

The main limitations are computational load, local optima, inaccurate long-term predictions, and difficulty representing multimodal futures. A learned model that predicts an average outcome can be dangerous: for example, averaging “left” and “right” collision outcomes may predict a physically impossible safe path.

6. World Models and Latent-Space Planning

A world model attempts to capture the aspects of the environment needed for prediction and control. In high-dimensional tasks, predicting raw images is often inefficient. Instead, an encoder maps observations to a latent state \(z_t\), and a latent dynamics model predicts

\[

z_{t+1}\sim p_\theta(z_{t+1}\mid z_t,a_t).

\]

A decoder may reconstruct observations, while a reward or value head predicts task-relevant quantities. The agent can then imagine trajectories in latent space and train a policy or value function from these imagined experiences.

Latent models offer computational savings, but a latent representation must preserve information relevant to future decisions. A representation that reconstructs visual detail well may still omit a hidden variable, such as whether a surface is slippery. Conversely, a control-oriented representation can ignore irrelevant appearance variation.

World models may be deterministic, probabilistic, recurrent, or transformer-based. Recurrent or sequence models are useful when observations are incomplete because they can maintain memory. Probabilistic models represent multiple possible futures and are important for risk-sensitive planning.

7. Uncertainty-Aware Planning

Uncertainty has at least two important sources:

• Aleatoric uncertainty: randomness inherent in the environment, such as sensor noise or an opponent’s unpredictable action.
• Epistemic uncertainty: lack of knowledge caused by limited data or unfamiliar states.

A model should not merely predict a mean; it should communicate confidence or a predictive distribution. An ensemble of models can approximate epistemic uncertainty by training several predictors on bootstrapped data or with different initializations. If ensemble predictions disagree strongly, the agent may avoid the region, gather information, or use a conservative fallback.

A risk-neutral planner maximizes expected return. A risk-sensitive planner may instead penalize variability or constrain the probability of failure. For example, a chance constraint can be expressed conceptually as

\[

\Pr(g(s_t,a_t)\leq 0\text{ for all relevant }t)\geq 1-\delta,

\]

where \(g\) represents a safety limit and \(\delta\) is an acceptable violation probability. Estimating this probability requires a calibrated uncertainty model; an arbitrary confidence score is not sufficient.

Uncertainty can also encourage exploration. An intrinsic information-gain objective rewards actions expected to reduce model uncertainty. This is useful in scientific discovery and robotics, but exploration must be constrained when failure is costly.

8. Model-Free and Model-Based RL: A Critical Comparison

Model-free methods learn a policy or value function without explicitly predicting transitions. They can be robust to certain modeling errors and may be simpler at execution time. Their disadvantages include high interaction requirements, weak generalization to changed dynamics, and limited ability to reason about unseen action sequences.

Model-based methods can be more sample-efficient, adapt to changed goals through replanning, and support constraint checking. Their disadvantages include model-learning cost, planning computation, model bias, distribution shift, and compounding prediction error.

The comparison should use the task’s actual requirements:

| Criterion | Model-free RL | Model-based RL |

|---|---|---|

| Real-world sample efficiency | Often lower | Often higher if the model is useful |

| Planning flexibility | Limited | Strong |

| Computational cost during deployment | Usually lower | Often higher |

| Sensitivity to model error | Not applicable in explicit form | Potentially high |

| Adaptation to changed objectives | May require retraining | Often easier through replanning |

| Safety analysis | Requires separate mechanisms | Model can support prediction and constraints |

Hybrid methods are frequently best. A learned value function can evaluate leaf states in tree search; a model can generate imagined data for a policy; and a model-free controller can serve as a fallback when planning confidence is low.

9. Worked Application: Warehouse Robot Navigation

Consider an illustrative autonomous warehouse robot that transports packages among shelves. Its state includes position, velocity, battery level, nearby obstacles, and package status. Actions are velocity commands. The robot receives positive reward for timely delivery and negative rewards for collisions, excessive energy use, and delay.

Background and Problem

A purely model-free controller may perform well in familiar aisles but fail when shelves are rearranged or a temporary obstacle appears. Real collision data are unacceptable, so collecting experience is expensive.

Model and Planning Approach

The engineering team builds a hybrid model. Known kinematic equations predict motion, while a learned residual model predicts wheel slip and braking distance. A perception system supplies obstacle estimates, and an ensemble predicts several possible motion outcomes. At each control cycle, MPC samples candidate velocity sequences over a two-second horizon. It scores each sequence using delivery progress, energy cost, and collision risk. Only the first command is executed before replanning.

The planner applies a conservative rule: if predicted uncertainty near an obstacle exceeds a threshold, it reduces speed and requests a wider clearance. A model-free emergency policy stops the robot when perception confidence is inadequate.

Analysis

The short horizon limits compounding error, while frequent replanning corrects deviations. The learned residual improves performance over a purely kinematic model. However, an obstacle detector that systematically misses low objects creates a model-input problem that planning alone cannot solve. The system therefore requires sensor validation, monitoring, and scenario testing.

Outcome

In an illustrative evaluation, the hybrid controller reduces simulation interactions by using imagined rollouts and adapts rapidly when aisle layouts change. Its principal benefit is not simply higher reward; it is the ability to test candidate actions against explicit motion and safety constraints before execution.

Learning Points

• Models should represent task-relevant uncertainty, not just average behavior.
• Receding-horizon control limits the damage from imperfect long-term predictions.
• Safety depends on sensing, model validation, constraints, and fallback policies together.
• Performance should be evaluated on changed layouts and rare hazards, not only average training conditions.

10. Implementation Workflow and Best Practices

A practical MBRL project can follow these steps:

1. Define the decision problem: specify state, action, reward, horizon, constraints, and termination conditions.

2. Audit prior knowledge: identify physical laws, invariants, conservation relationships, and known failure modes.

3. Collect representative data: include normal operation, boundary conditions, disturbances, and recovery behavior.

4. Choose model granularity: predict only what planning needs, at an appropriate temporal resolution.

5. Validate one-step and multi-step predictions: a model can have good one-step error but poor long-horizon behavior.

6. Measure calibration: predicted uncertainty should correspond to observed frequencies of error.

7. Use short horizons and replanning when appropriate: this often reduces compounding error.

8. Include hard constraints and fallback controllers: never rely solely on a learned planner for safety-critical actions.

9. Test distribution shift: evaluate new goals, layouts, users, weather, loads, and sensor conditions.

10. Monitor after deployment: log predictions, uncertainty, actions, near misses, and overrides.

Evaluation should report both task performance and system quality. Relevant metrics include return, success rate, collision or failure rate, sample efficiency, planning latency, energy use, calibration error, and robustness under perturbation. A high average reward is inadequate if rare failures have severe consequences.

11. Challenges, Ethics, and Governance

Model-based systems can create a misleading impression of explainability. A predicted trajectory is not necessarily a reliable explanation if the underlying model is biased. Organizations should document training data, modeling assumptions, uncertainty treatment, constraints, and known failure regions.

Safety-critical applications require hazard analysis, human override procedures, audit logs, and validation against regulatory requirements. In autonomous mobility, for example, responsibility cannot be delegated to an opaque planner without specifying operational limits and escalation policies.

Exploration raises ethical concerns when it affects people, property, health, or public services. Offline datasets may be safer, but they can encode historical bias and omit rare events. Privacy and data governance are also relevant when models learn from cameras, user behavior, or industrial records.

A responsible deployment process should include:

• explicit safety envelopes;
• conservative behavior under epistemic uncertainty;
• independent testing and red-team scenarios;
• clear ownership of decisions and incidents;
• mechanisms for model rollback and software updates;
• fairness and accessibility checks where human outcomes are involved.

12. Contemporary Developments

Current research increasingly combines differentiable simulators, sequence-based world models, offline RL, uncertainty ensembles, foundation models, and planning. Offline model-based RL is attractive when interaction is unavailable, but it must avoid extrapolating beyond the dataset. Decision transformers and sequence predictors can model behavior without traditional value iteration, yet their reliability still depends on coverage and evaluation.

Other developments include model-based multi-agent RL, where the environment includes strategic opponents; hierarchical planning, where high-level skills call lower-level controllers; and active system identification, where the agent deliberately gathers information about dynamics. These directions reinforce a central principle: planning quality depends on the relationship between model fidelity, uncertainty, computational resources, and decision consequences.

Chapter Summary

Model-based reinforcement learning uses an environment model to predict transitions and rewards, enabling planning as well as learning. Models may be explicit, learned, or hybrid, and may represent physical states, observations, or latent variables. Dyna combines real experience with simulated updates; Monte Carlo tree search selectively explores future branches; trajectory optimization and MPC optimize action sequences under dynamics and constraints; and world models support imagination in compact latent spaces.

Model error is the defining risk of MBRL. Prediction errors can compound over long horizons, and an inaccurate model can generate harmful plans. Uncertainty-aware ensembles, conservative constraints, short horizons, replanning, and fallback controllers help manage this risk. Compared with model-free RL, MBRL often improves sample efficiency and adaptability but requires additional modeling and computation. In professional systems, success must be assessed through safety, calibration, robustness, latency, and governance—not reward alone.

Key Terms

Model-based reinforcement learning: RL that uses an environment model for prediction or planning.

Transition model: A model of how states change after actions.

Reward model: A model estimating immediate rewards.

Model bias: Systematic error caused by an inaccurate learned or specified model.

Planning: Evaluating possible future actions using a model.

Dyna: An architecture combining direct learning, model learning, and simulated planning.

Model predictive control: Repeated finite-horizon planning in which only the first action is executed.

Monte Carlo tree search: Selective tree expansion using simulation and statistical action evaluation.

UCT: A tree-search rule balancing exploitation and exploration.

Trajectory optimization: Direct optimization of a sequence of future actions.

World model: A predictive model of environment dynamics, often in latent space.

Latent state: A learned internal representation used for prediction or control.

Epistemic uncertainty: Uncertainty caused by limited knowledge or data.

Aleatoric uncertainty: Inherent randomness in observations or transitions.

Compounding error: Growth of prediction inaccuracies over multi-step rollouts.

Receding horizon: Planning over a moving finite horizon and replanning after each action.

Hybrid model: A model combining known structure with learned components.

Chance constraint: A constraint requiring safety with at least a specified probability.

Sample efficiency: Performance gained per unit of real-world experience.

Fallback controller: A safer backup policy used when planning confidence is inadequate.

Review Questions

Short Answer Questions

1. What are the two principal activities separated by model-based reinforcement learning?

2. Distinguish aleatoric uncertainty from epistemic uncertainty.

3. What is the purpose of executing only the first action in MPC?

4. Name the four stages of Monte Carlo tree search.

5. Why can a model have good one-step accuracy but poor planning performance?

Long Answer Questions

1. Explain the Dyna architecture and analyze how simulated experience affects sample efficiency and model bias.

2. Compare model-free and model-based reinforcement learning using computational cost, adaptability, safety, and data requirements.

3. Describe trajectory optimization for a continuous-control problem, including the role of dynamics, costs, constraints, and terminal values.

4. Explain how world models use latent representations and discuss their advantages and risks.

5. Discuss how uncertainty-aware planning can improve the reliability of a deployed autonomous system.

Analytical / Application Questions

1. Design a model-based RL architecture for a robot operating in an unfamiliar environment. Specify its model, planner, uncertainty method, safety constraints, and fallback policy.

2. A planner performs well in simulation but fails after deployment because friction differs from the training environment. Diagnose the failure and propose a model-learning and validation strategy.

3. For a warehouse robot, compare MCTS, MPC, and Dyna. Recommend one primary method and one supporting method, justifying the choice using action-space structure, latency, uncertainty, and safety requirements.

Multiple Choice Questions with Answers

Question 1

Which component distinguishes model-based reinforcement learning from purely model-free reinforcement learning?

A. An explicit or learned model of environment dynamics

B. The use of discounted rewards

C. The existence of a policy

D. The use of states and actions

Correct Answer: A. An explicit or learned model of environment dynamics

Question 2

In Dyna, what is the main purpose of simulated planning updates?

A. To create observations that are guaranteed to be real

B. To reuse learned transitions and propagate value information

C. To eliminate the need for a reward function

D. To prevent the model from being updated

Correct Answer: B. To reuse learned transitions and propagate value information

Question 3

In the UCT formula, the exploration term primarily favors which actions?

A. Actions with the lowest immediate reward

B. Actions that have been visited most often

C. Actions that are relatively under-visited

D. Actions excluded from the search tree

Correct Answer: C. Actions that are relatively under-visited

Question 4

What does model predictive control typically do after optimizing a finite action sequence?

A. Executes the entire sequence without observation

B. Executes only the first action and replans

C. Discards the model permanently

D. Updates only the reward discount factor

Correct Answer: B. Executes only the first action and replans

Question 5

Which is an example of epistemic uncertainty?

A. Random sensor noise that cannot be eliminated

B. An opponent choosing randomly by design

C. Unknown braking behavior because few examples were observed

D. A fixed actuator limit

Correct Answer: C. Unknown braking behavior because few examples were observed

Question 6

Why are latent-space world models often used instead of models that predict raw images directly?

A. Latent models necessarily contain no uncertainty

B. Latent models can focus computation on decision-relevant representations

C. Raw images never contain useful information

D. Latent models remove the need for planning

Correct Answer: B. Latent models can focus computation on decision-relevant representations

Question 7

Given Q(s,a)=1, reward r=2, gamma=0.9, max Q(s',a')=4, and alpha=0.5, what is the one-step Q-learning update result?

A. 2.30

B. 3.30

C. 4.60

D. 5.60

Correct Answer: B. 3.30

Question 8

Which problem occurs when a learned model predicts an average of several qualitatively different possible futures?

A. The prediction may represent an unrealistic future

B. The model becomes automatically conservative

C. The discount factor becomes greater than one

D. The action space becomes discrete

Correct Answer: A. The prediction may represent an unrealistic future

Question 9

What is the main function of a fallback controller in a safety-conscious MBRL system?

A. To increase model error deliberately

B. To replace data collection in all circumstances

C. To provide safer behavior when planning confidence is inadequate

D. To maximize exploration near hazards

Correct Answer: C. To provide safer behavior when planning confidence is inadequate

Question 10

Which statement best describes a hybrid environment model?

A. It uses only random transitions

B. It combines known structure with learned components

C. It excludes physical constraints

D. It is a model-free policy with no predictions

Correct Answer: B. It combines known structure with learned components