ai theory

Reinforcement Learning Basics (16) — Model-Based RL: Learning by Imagining the Environment

Junyoung Park · 2024-06-28 · 6 min

Introduction

The methods we have covered so far—Monte Carlo, TD, DQN, and PPO—do not explicitly learn the environment's transition rules. They update the value function or policy directly from experience. These methods are called model-free reinforcement learning.

Model-based RL learns one more thing:

If I take this action in this state, what will happen next?

With a model of the environment, an agent can imagine possible futures and plan without having to try every option in the real environment. The catch is that if it spends too long imagining with an inaccurate model, it may become confident in a future that is simply wrong.

In this article, we will examine what a model represents, how learning differs from planning, and how Dyna connects the two.

What Does an Environment Model Predict?

We represented the dynamics of an MDP with the following probability:

p(s',r\mid s,a)

It is the distribution of the next state $s'$ and reward $r$ , given the current state $s$ and action $a$ .

A model $\hat p_\eta$ approximates this distribution from data.

\hat p_\eta(s',r\mid s,a) \approx p(s',r\mid s,a)

In a deterministic environment, a neural network can predict the next state and reward directly.

(\hat s',\hat r) = f_\eta(s,a)

In a stochastic environment, predicting a single mean is not enough. The model needs to represent a distribution over possible outcomes.

Learning and Planning

Model-based RL becomes easier to understand once we distinguish these two terms.

Learning: update a value function, policy, or model using experience from the real environment.
Planning: update a value function or policy using simulated experience produced by a model.

A real transition and a model-generated transition can have exactly the same form in an update rule.

(s,a,r,s')

The difference is whether the tuple came from reality or from a simulation inside the model.

Dyna: Experience, Learn, and Imagine

The Dyna architecture places model-free learning and planning in a single loop.

Take an action in the real environment and obtain a transition.
Update the value function or policy with that real transition.
Train the model with the same transition.
Select a previously visited state-action pair and use the model to generate a simulated next state and reward.
Update the value function or policy several more times with the simulated transitions.

From one real experience, Dyna performs a direct learning update and trains the model, then adds planning updates using several simulated transitions.

With Q-Learning, both real and simulated experience can use the same update equation.

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

If the agent performs $n$ planning updates after every real environment step, it can extract more learning value from each expensive interaction.

A Simple Planning Example

Suppose the agent takes $(s_A,\text{RIGHT})$ in a maze and moves to $s_B$ . The model records this outcome.

\hat p(s_B,0\mid s_A,\text{RIGHT})=1

Later, even if the agent is somewhere else, it can retrieve $s_A$ from the model and simulate taking RIGHT again. It can repeatedly update the Q-value without physically returning to $s_A$ .

This resembles replaying previously collected experience, but there is an important difference.

A replay buffer reuses transitions that were actually observed.
A model uses learned dynamics to generate unobserved combinations and possible futures as well.

Why Model-Based RL Is Sample-Efficient

Moving a robot once or asking a person to evaluate one response can be expensive. Running a simulation inside a learned model may be relatively cheap.

From one real experience, we can obtain

1\text{ real update} + n\text{ planning updates}.

This provides more learning signal under the same interaction budget. It is one of the central advantages of model-based methods.

An explicit model can also answer counterfactual questions. The agent can compare “What if I had gone right instead of left?” without carrying out both actions in reality.

If the Model Is Wrong, the Planning Is Wrong Too

A model may be accurate in regions with abundant data and unreliable in regions the agent has barely visited. If the policy discovers an action that exploits a small modeling error, it may fail in the real environment.

Errors become especially problematic when the model is rolled out for many consecutive steps.

\hat S_{t+1} = f_\eta(S_t,A_t)

\hat S_{t+2} = f_\eta(\hat S_{t+1},A_{t+1})

A small error in the first prediction becomes the input to the second prediction, and the discrepancy can grow at every subsequent step.

An accurate model helps planning, but even a small one-step error can accumulate over a long rollout and produce an entirely different future.

Reducing Model Bias

Practical algorithms use several safeguards.

Short Model Rollouts

Start from a real state, use the model for only a few steps, and then return to real data. Short imagined rollouts accumulate less error than long ones.

Ensembles and Uncertainty

Train several models and treat regions where their predictions disagree as uncertain. The agent can reduce the planning weight of uncertain transitions or collect more real data there.

Mixing in Real Data

Do not train only on model-generated data. Continue mixing in real transitions so that the policy does not drift too far from reality.

Conservative Objectives

Apply a penalty so that state-action pairs the model is unsure about are not evaluated too optimistically.

Model-Free and Model-Based Are Not a Binary Choice

Practical methods lie along a continuum.

Method	Model use	Typical role
Model-free	None	Directly update a Q-function or policy from experience
Dyna-style	Learned model	Combine real updates with planning updates
Planning-focused	Known or learned model	Select actions through search
Hybrid	Partial model use	Combine short rollouts, value bootstrapping, and search

A model does not need to predict the entire future. It may predict only the reward, a latent state several steps ahead, or just the parts of a search tree where modeling is useful.

What to Remember

Model-based RL learns a model that predicts the environment's next state and reward, then plans with simulated experience generated by that model.

An environment model approximates $p(s',r\mid s,a)$ .
Learning updates from real experience; planning updates from model-generated experience.
Dyna performs several planning updates after each real step.
A model can improve sample efficiency, but inaccurate predictions can steer the policy in the wrong direction.
Short rollouts, uncertainty estimates, and mixtures of real data can reduce model bias.

The next article examines imitation learning and offline RL, where the agent cannot reconnect to the real environment and must learn only from an already collected dataset.