Learn how the world responds, then plan or practice inside that model.
On this page
A model-based agent learns, or is given, a model of the world. It plans or practices inside that model, and acts in the real world to collect steps that improve the model, the policy or both (Figure 1). The methods in Value-Based RL, Policy Gradients: REINFORCE, Actor-Critic and PPO and Continuous-Action Actor-Critic: TD3 and SAC learn from real steps only; they are model-free.
A chess player tries moves in their head before touching a piece; a model lets an agent do the same before spending real steps.
A model predicts the next state and reward from a state and a move; planning uses it to look ahead or practice without the real world. Dyna (Sutton 1991) runs the Q-learning update from Value-Based RL on imagined steps as well as real ones.
A real step into the goal teaches one square; the model lets the agent replay the square before it without moving (Figure 2).
Dyna-Q adds n imagined steps from a remembered model after each real one (Sutton and Barto’s RL textbook, 2018, section 8.2). The code runs one Dyna-Q step on the book’s 6 by 9 maze. env_step(s, a) is the maze’s rule, returning (r, s_next).
import random
import numpy as np
alpha, gamma, eps, n = 0.1, 0.95, 0.1, 5 # book's maze
Q = np.zeros((54, 4)) # 6 x 9 squares, 4 moves
model = {} # (s, a) -> (r, s_next)
GOAL = 8 # top-right square, episode ends
def q_update(s, a, r, s_next):
future = 0.0 if s_next == GOAL else Q[s_next].max()
Q[s, a] += alpha * (r + gamma * future - Q[s, a])
def dyna_step(env_step, s):
if random.random() < eps:
a = random.randrange(4)
else:
a = int(random.choice(np.flatnonzero(Q[s] == Q[s].max())))
r, s_next = env_step(s, a)
q_update(s, a, r, s_next)
model[(s, a)] = (r, s_next)
for _ in range(n):
(s_m, a_m), (r_m, s_m_next) = random.choice(list(model.items()))
q_update(s_m, a_m, r_m, s_m_next)
return s_next
Each imagined step can carry value, a square’s best Q, one square further back from the goal, so the reward found there spreads through the maze without new real steps (the RL textbook, Example 8.1). Where a square has no value yet, all four moves there score 0 and the agent picks one at random; once value reaches the start, following the best move from each square gives a route to the goal. Reaching the goal ends an episode, and the agent starts again from the start.
Until the goal is first found every value is 0, so imagined steps change nothing and every n takes the same real steps to find it, on average. After that, a larger n forms the route in fewer real steps on average, at the cost of n + 1 updates per real step, and n = 0 is plain Q-learning, where each real step updates only the square it left. One random run can occasionally flip that order.
At step 150, no route yet.
In the book’s averaged runs a good route took about 25, 5 and 3 episodes for n = 0, 5 and 50. Imagined steps matter most when real steps are slow or costly, as for a robot. DQN’s experience replay reuses stored real steps the same way, with a network in place of the table.
AlphaZero (Silver et al. 2018) uses a game’s rules for a tree search, a look-ahead over move sequences that a network steers, and learns from games against itself.
In one Go position the network’s first guesses are rough; 100 search visits sharpen them (Figure 4).
MuZero (Schrittwieser et al. 2020) learns the model instead. MuZero is not told the rules of Go; it must also predict what each move pays (Figure 5).
World Models (Ha and Schmidhuber 2018) train a small controller inside a learned model of the game, then run it for real. Making the model’s predictions more random kept the controller from exploiting the model’s errors. An agent that counts only the reward inside a fixed window of imagined steps turns shortsighted. Dreamer (Hafner et al. 2020) learns a compact hidden-state model from replayed experience and trains its actor and critic on imagined runs inside it; the critic estimates the reward past the window, so the actor still aims for rewards that come later.
Dreamer’s walker practices inside its own model of the body, a few steps at a time (Figure 6).
The reward 1 comes when it reaches walking speed; γ is 0.9, as in the return strip of RL Basics.
What is known about the rules and what a real step costs decide the pick; Choosing a method sets these methods beside the model-free ones.
| Method | When | How often it learns or plans | Carries | Cost |
|---|---|---|---|---|
| Dyna-Q | few states that fit in a table, rules unknown, real steps cost more than compute | one real update plus n imagined ones per real step; more n gets more from each real step at more compute, 0 is plain Q-learning; 5 or 50 in the book | the Q-table and the stored model of seen steps | n + 1 updates per real step; maze.py |
| AlphaZero | rules known and exactly simulated, as in board games | 800 search simulations per move; more simulations look further ahead at more compute per move, and AlphaZero played stronger with more thinking time per move (the AlphaZero paper’s thinking-time plot); network trained 700,000 steps on batches of 4,096 positions | the network and its stored self-play games | 5,000 first-generation TPUs (Google’s AI chips) playing, 16 second-generation TPUs training |
| MuZero | rules unknown, look-ahead matters | 800 simulations per move in board games, 50 in Atari, which has far fewer moves to pick from; more helped a lot in Go but leveled off near 100 in Atari, where even 1, the policy alone, played well (the MuZero paper’s search-scaling plots); the dynamics unroll 5 moves per replayed sample, always 5 in the paper, which does not test other values; a longer unroll trains the model to predict further ahead at more compute per sample; the playing copy refreshes every 1,000 training steps | the three networks and a replay buffer of past games | one learned-model step per imagined move |
| World Models | rules unknown, pixels in; the idea Dreamer builds on | 10,000 random runs first, then the vision model, which compresses each frame, the memory model, which predicts the next compressed frame, then the controller by evolution, trying many random weight changes and keeping the best | the vision, memory and controller models | controller 867 weights, vision model 4.3 million |
| Dreamer | rules unknown, real steps costly | batches of 50 replayed sequences of 50 steps; imagines 15 steps ahead, 10 in Atari; a longer window counts more imagined reward before the critic takes over, and with its critic Dreamer stayed robust to the window length and did well even with short ones, unlike a version without a critic (the Dreamer paper’s horizon plot) | the world model, actor, critic and a replay buffer | no real steps for the actor and critic |
The methods in this series ran on small examples such as a cliff walk, CartPole and a maze. What does it take to train one on a real agent, the banana collector from RL Basics that sees 37 sensor numbers? Deep Q-Networks (DQN) builds that agent with replay, a target network and the Double and Dueling changes from Value-Based RL.