Models, planning and imitation: Dyna, MCTS and AlphaZero

Before this lesson: lesson 02 (policy iteration), lesson 03 (UCB1), lesson 05 (Q-learning), lesson 07 (softmax policies); information theory 04 (cross-entropy) and supervised classification (logistic regression). Time: about 90 minutes, plus the notebook.

---

Why this matters #

• Games. AlphaGo, AlphaGo Zero and AlphaZero reached superhuman play in Go, chess and shogi with Monte Carlo tree search guided by a neural network. MuZero did the same without being told the rules.
• Discovery by search. AlphaZero-style agents (search plus a policy and value network) found new matrix multiplication algorithms (AlphaTensor, 2022) and faster sorting routines (AlphaDev, 2023). MCTS with neural networks plans chemical syntheses (Segler, Preuss and Waller, 2018).
• Sample efficiency. Real experience is expensive for robots, for plants in industry and for experiments in science. A learned model lets the agent reuse each sample many times (MBPO, Dreamer, model-predictive control).
• Language models. Supervised fine-tuning on demonstrations is behaviour cloning. On-policy distillation, where the teacher labels tokens that the student generated, is the DAgger idea.
• ML for science. A learned weather or fluid emulator that is rolled out step by step has the same compounding error as a learned dynamics model (ML for science 16).

1. Predict first #

2. The picture #

Learning and planning #

A model of the environment predicts what happens after an action: the next state $s'$ and the reward $r$. A distribution model gives the probabilities $P(s' \mid s, a)$ and the expected reward $R(s, a)$, as in lesson 02. A sample model only returns one possible $(r, s')$, like a simulator. Planning is any computation that uses a model to improve a policy or a value function, without new real experience.

• Model-free RL (lessons 04–09) learns $V$, $Q$ or $\pi$ directly from real transitions.
• Model-based RL learns a model from real transitions (or is given one), then plans with it.

There are two ways to plan. Background planning improves $Q$ or $\pi$ for all states, between real steps; Dyna-Q is the classic example. Decision-time planning searches from the current state only, just before choosing the next action; Monte Carlo tree search and model-predictive control do this.

Dyna-Q (Sutton, 1990) is Q-learning plus planning. After each real step it does one normal Q-learning update, stores the transition in a table (the model), and then makes $n$ more Q-learning updates on transitions replayed from the model. With $n = 0$ it is plain Q-learning.

Interactive figure: loading…

What to notice. The reward is 0 everywhere except at the goal. With $n = 0$ the value of the goal moves back about one cell per episode: after episode 2 only about 2 cells have a value. With $n = 50$ the planning updates replay the old transitions, so the value spreads back over almost the whole maze during episode 2. In our runs (30 seeds), episode 2 took on average 726 steps with $n = 0$, 156 with $n = 5$, and 41 with $n = 50$. Dyna-Q with $n = 50$ finds a near-shortest path in about 3 episodes; plain Q-learning needs about 25 to 30. The shortest path has 14 steps.

Rollouts drift: compounding model error #

A learned model is used again and again: the predicted next state is the input of the next prediction. A small error in one step changes the input of the next step, and the errors add up. If the true dynamics push nearby states apart, the errors also grow.

Interactive figure: loading…

Two lessons follow. First, a model must be judged by its multi-step error, not only by its one-step error. Second, a plan made with an imperfect model should not be run blindly: run its first action, look at the real state, and plan again. This is model-predictive control (section 3).

Search trees: Monte Carlo tree search in four phases #

In a game like Go or chess, the agent knows the rules, so it has a perfect model. The difficulty is the size of the tree of possible futures. Monte Carlo tree search (MCTS) grows a small part of this tree, from the current position, and spends its time on the promising branches. Each iteration has four phases:

1. Select. Start at the root. At each node, choose a child with a bandit rule that balances a high mean result (exploit) and few visits (explore). Stop at a node that is not yet in the tree.
2. Expand. Add that node to the tree.
3. Evaluate. Estimate how good the node is: play the game to the end with random moves (a playout or rollout), or ask a value network.
4. Backup. Add the result to the visit count $N$ and the total value $W$ of every node on the path. The mean $Q = W / N$ is the node's estimated value.

After many iterations, the agent plays the move at the root with the most visits.

Interactive figure: loading…

Two players. Here you move first and the opponent replies. The opponent wants your result to be low. We store each node's $Q$ for the player who made the move into it, so every player maximises at their own nodes. The backup therefore flips the sign: a result $z = +1$ (you win) adds $+1$ to your moves and $-1$ to the opponent's moves on the path. Forgetting the flip makes the opponent help you.

Why B wins. At first, A looks best, because random replies are mostly good for you. But once A is visited often, the selection at A finds the opponent's strong reply A3 and visits it most. The mean of A then moves towards $2 \cdot 0.2 - 1 = -0.6$. B's worst reply is still worth $2 \cdot 0.5 - 1 = 0$. MCTS approaches the minimax answer, not the average answer.

Copying an expert also drifts #

Imitation learning learns a policy from an expert's demonstrations instead of from rewards. The simplest method, behaviour cloning (BC), is supervised learning on the expert's state–action pairs. The problem is the data: the expert stays on the road, so the data contains only on-road states. After one mistake the learner is in a state it has never seen. There it has no reason to act well, so one mistake can cost many later steps. DAgger lets the learner drive, asks the expert what to do in the states the learner visited, and adds these labels to the data.

Interactive figure: loading…

BC's expected cost grows like $T^2$, DAgger's like $T$. Section 3 proves this.

3. Definitions and main results #

Models and Dyna-Q #

In a deterministic environment the simplest learned model is a table: $\text{Model}(s, a) = (r, s')$, the last outcome seen after $a$ in $s$. In a stochastic environment the table can store all outcomes and sample one of them.

Algorithm (tabular Dyna-Q, Sutton and Barto §8.2). Initialise $Q(s, a) = 0$ and an empty model. At every real step:

1. In state $s$, choose $a$ ε-greedily with respect to $Q$. Take it; observe $r$ and $s'$.
2. Direct RL: $Q(s, a) \leftarrow Q(s, a) + \alpha\,[r + \gamma \max_{a'} Q(s', a') - Q(s, a)]$ (the Q-learning update of lesson 05; $\max = 0$ after a terminal state).
3. Model learning: $\text{Model}(s, a) \leftarrow (r, s')$.
4. Planning: repeat $n$ times: pick a previously seen pair $(\tilde s, \tilde a)$ at random, read $(\tilde r, \tilde s') = \text{Model}(\tilde s, \tilde a)$, and apply the same Q-learning update to $(\tilde s, \tilde a, \tilde r, \tilde s')$.

Why it helps. Dyna-Q makes $n + 1$ Q-learning updates per real step, and each real transition is used again in many planning updates. The values then spread backwards through all known transitions, not only along the path of the current episode. The price is computation ($n$ updates per step) and memory for the model.

Two refinements from Sutton and Barto, chapter 8. Prioritized sweeping (§8.4) replays first the pairs whose values would change most, instead of random pairs. Dyna-Q+ (§8.3) adds a small bonus $\kappa\sqrt{\tau}$ to the reward of a pair not tried for $\tau$ steps, so the agent checks again whether the world has changed. A model of an old world can be worse than no model.

Planning with a learned model: MPC #

With continuous states and actions, a table is not enough. We fit a model $\hat f_\phi(s, a) \approx s'$ (and often a reward model $\hat r$) by supervised regression on real transitions, for example a neural network or an ensemble of networks.

Model-predictive control (MPC), also called receding-horizon control. At every real step:

1. In the model, search for an action sequence $a_0, \dots, a_{H-1}$ that maximises the predicted return $\sum_{k=0}^{H-1} \gamma^k \hat r(\hat s_k, a_k)$, where $\hat s_0 = s$ and $\hat s_{k+1} = \hat f(\hat s_k, a_k)$.
2. Run only the first action $a_0$ in the real environment.
3. Observe the real next state and start again from step 1.

The search in step 1 is often derivative-free:

• Random shooting: sample $N$ random action sequences, predict the return of each in the model, and keep the best one.
• Cross-entropy method (CEM): keep a Gaussian over action sequences with mean $\mu$ and standard deviation $\sigma$ (both of size $H \times \dim a$). Repeat a few times: sample $N$ sequences, keep the $M$ best ("elites"), and set $\mu$ and $\sigma$ to the mean and standard deviation of the elites. Return $\mu$. At the next real step, start from $\mu$ shifted by one step (a warm start).

MPC corrects model errors with feedback: every plan starts from the real state. PETS (Chua et al., 2018) is MPC with CEM on an ensemble of learned probabilistic models.

Model bias and compounding error #

Model bias. A planner or a policy that is optimised inside the model looks for the actions with the highest predicted return. If the model is wrong somewhere, the optimiser finds those places, because errors that look good are exactly what it searches for. So the plan is often better in the model than in reality.

Compounding error. Let the true dynamics be $s_{t+1} = f(s_t, a_t)$ and the model $\hat s_{t+1} = \hat f(\hat s_t, a_t)$, both run with the same actions from the same start $s_0$.

Theorem (open-loop error). Assume the one-step model error is at most $\varepsilon_m$: $\|\hat f(s, a) - f(s, a)\| \le \varepsilon_m$ for all $s, a$, and that $f$ is $L$-Lipschitz in $s$: $\|f(s, a) - f(s', a)\| \le L\|s - s'\|$. Then the error after $H$ steps satisfies

Proof.

1. Write $e_t = \|\hat s_t - s_t\|$, so $e_0 = 0$.
2. Add and subtract $f(\hat s_t, a_t)$: $e_{t+1} = \|\hat f(\hat s_t, a_t) - f(s_t, a_t)\| \le \|\hat f(\hat s_t, a_t) - f(\hat s_t, a_t)\| + \|f(\hat s_t, a_t) - f(s_t, a_t)\|$.
3. The first term is at most $\varepsilon_m$; the second is at most $L e_t$. So $e_{t+1} \le \varepsilon_m + L e_t$.
4. Unroll from $e_0 = 0$: $e_H \le \varepsilon_m (1 + L + \dots + L^{H-1})$. $\blacksquare$

When the dynamics expand ($L > 1$) the bound grows geometrically with $H$. In the figure, the true system is $x_{t+1} = 1.05\,x_t$ and the model is $\hat x_{t+1} = 1.06\,\hat x_t$. The relative error after $H$ steps is exactly $(1.06/1.05)^H - 1$: 0.0095 for $H = 1$, 0.099 for $H = 10$, and 0.61 for $H = 50$, more than the $50 \times 0.0095 = 0.48$ that a linear growth would give.

The trade-off. Model-based methods usually need far fewer real samples, because every sample is reused in planning. Model-free methods often reach a better final performance, because they optimise the real return and do not depend on the accuracy of a model. Practical methods take the middle road:

• MBPO (Janner et al., 2019) trains an ensemble of models, makes only short model rollouts ($k$ steps, often 1 to a few) that start from real states in the replay buffer, and trains SAC (lesson 08) on real and model data. Short rollouts keep the compounding error small; the ensemble shows where the models disagree.

The infinite-horizon version: the simulation lemma

Let $\hat P$ be a model with $\sum_{s'} |\hat P(s' \mid s, a) - P(s' \mid s, a)| \le \varepsilon$ for all $(s, a)$, with the true rewards and rewards in $[0, R_{\max}]$. For any policy $\pi$, the Bellman equations of lesson 01 give $\hat V^\pi - V^\pi = \gamma (I - \gamma \hat P_\pi)^{-1} (\hat P_\pi - P_\pi) V^\pi$. Each row of $(\hat P_\pi - P_\pi) V^\pi$ is at most $\varepsilon \|V^\pi\|_\infty \le \varepsilon R_{\max}/(1-\gamma)$ in absolute value, and $(I - \gamma \hat P_\pi)^{-1}$ multiplies the max-norm by at most $1/(1-\gamma)$. So

$$\|\hat V^\pi - V^\pi\|_\infty \le \frac{\gamma\,\varepsilon\,R_{\max}}{(1-\gamma)^2}.$$

The effective horizon is $1/(1-\gamma)$, so this is "horizon squared" again, the same order as the behaviour cloning bound below.

Monte Carlo tree search and UCT #

MCTS keeps, for each edge $(s, a)$ of the tree, a visit count $N(s, a)$, a total value $W(s, a)$ and the mean $Q(s, a) = W(s, a) / N(s, a)$. $N(s)$ is the visit count of node $s$. One iteration runs the four phases of section 2. The rule used in the selection phase is the tree policy; the rule used in random playouts is the default policy (or rollout policy).

UCT (Kocsis and Szepesvári, 2006) uses UCB1 from lesson 03 as the tree policy, at every node:

and tries every unvisited child first ($N(s, a) = 0$ counts as $+\infty$). With $c = \sqrt 2$ and results in $[0, 1]$ this is exactly UCB1: each node is a bandit whose arms are its children. Two differences from the bandit of lesson 03: the rewards of an arm are not stationary (the policy below it keeps changing), and the arms are nested. Kocsis and Szepesvári showed that, with a suitable $c$, the value estimate at the root converges to the optimal (minimax) value, and the probability of choosing a wrong move goes to 0 as the number of iterations grows.

Choosing the move. After the search, play the child of the root with the most visits (the "robust child"). A child with a high $Q$ but few visits may just be lucky; the visit count includes the search's confidence.

Two-player games (lesson 01 has one agent; here two agents alternate). Store $W$ and $Q$ for the player who moved into the node, and add $+z$ or $-z$ in the backup by whose move it was (the negamax convention). An equivalent convention stores all values for the first player and lets the second player choose $\arg\min$ of $Q - (\text{bonus})$.

Imitation learning: behaviour cloning and its error bound #

Behaviour cloning (BC): supervised learning of $\pi(a \mid s)$ on expert pairs $(s, a^*)$, for example by maximum likelihood (minimise the cross-entropy $-\log \pi(a^* \mid s)$, see information theory 04). Supervised fine-tuning of an LLM on demonstrations is BC.

Setting (Ross and Bagnell, 2010). Horizon $T$; a deterministic expert $\pi^*$; a cost per step in $[0, 1]$; $J(\pi)$ is the expected total cost of $\pi$ over $T$ steps. Note: here $J$ is a cost, so smaller is better (in lesson 07 $J$ was a return). Let $\varepsilon_t$ be the probability that $\pi$ picks a different action from $\pi^*$ at a state drawn from the expert's state distribution at step $t$, and $\varepsilon = \frac1T \sum_{t=1}^T \varepsilon_t$. So $\varepsilon$ is the error rate on the expert's own data: the test error of the supervised problem.

Theorem (compounding error of behaviour cloning). $J(\pi) \le J(\pi^*) + T^2 \varepsilon$.

Proof.

1. Run $\pi$ and $\pi^*$ with the same environment randomness. Until their first disagreement they visit the same states and pay the same costs.
2. Before step $t$ the two runs are identical, so the state at step $t$ has the expert's distribution. The first disagreement is at step $t$ only if $\pi$ disagrees with $\pi^*$ there: probability at most $\varepsilon_t$.
3. After a first disagreement at step $t$, the learner pays at most $T - t + 1 \le T$ more than the expert, because each cost is in $[0, 1]$.
4. So $J(\pi) - J(\pi^*) \le \sum_{t=1}^T \varepsilon_t \cdot T = T^2 \varepsilon$. $\blacksquare$

The bound is tight in order: there are problems where a policy with error $\varepsilon$ on the expert's states pays an extra cost of order $\varepsilon T^2$ (DAgger paper, §2). The road of section 2 is one. The reason is distribution shift: the learner is trained on the expert's states but tested on its own.

DAgger #

If $\varepsilon$ were measured on the learner's own state distribution, the extra cost would be only $O(T\varepsilon)$: each mistake costs a bounded amount, because the learner also knows how to recover. DAgger (Dataset Aggregation; Ross, Gordon and Bagnell, 2011) makes the training distribution the learner's own:

1. Train $\hat\pi_1$ on expert data $\mathcal D$.
2. For $i = 1, 2, \dots$: run the policy $\pi_i = \beta_i \pi^* + (1 - \beta_i)\hat\pi_i$ (the expert with probability $\beta_i$, often $\beta_1 = 1$ and then $\beta_i = 0$) and record the states it visits.
3. Ask the expert for the correct action $\pi^*(s)$ in each recorded state.
4. Add these pairs to $\mathcal D$ (aggregate), train $\hat\pi_{i+1}$ on all of $\mathcal D$, and repeat.

Guarantee (Theorem 3.2 of the paper, simplified). After $N$ rounds, one of the trained policies $\hat\pi$ satisfies $J(\hat\pi) \le J(\pi^*) + u T \varepsilon_N + O(1)$: linear in $T$. Here $\varepsilon_N$ is the error of the best policy in the class on the aggregated data, and $u$ bounds how much one wrong action can increase the expert's cost-to-go ($u \le 1$ when the cost is the 0–1 disagreement with the expert). The price: an expert that can be queried during training, not just a fixed dataset.

AlphaZero: search as the expert, the network as the student #

AlphaZero (Silver et al., 2018) learns Go, chess and shogi from self-play only. One neural network $f_\theta$ takes a position $s$ and returns

• a vector of move probabilities $p = p_\theta(\cdot \mid s)$ (the policy head), and
• a scalar value $v = v_\theta(s) \in [-1, 1]$, the predicted game result for the player to move (the value head).

Search with PUCT. AlphaZero runs MCTS from every position it plays (800 simulations per move in self-play). It changes MCTS in two ways:

1. No random playouts. A new leaf is evaluated by the network: its value $v$ is backed up, and its prior probabilities $P(s, a) = p_\theta(a \mid s)$ are stored on its new edges.
2. PUCT selection, which uses the prior:

The bonus is large for moves that the network likes and that the search has not visited much. Unlike UCT, an unvisited move is not tried automatically: a move with a tiny prior may never be tried. To keep some exploration, AlphaZero adds Dirichlet noise to the priors at the root. (The published pseudocode also lets $c_{\text{puct}}$ grow slowly with $N(s)$; a constant is the usual simplified form.)

Self-play targets. After the search from position $s_t$, the visit counts at the root give the search policy

with a temperature $\tau$ ($\tau = 1$: proportional to the visits; $\tau \to 0$: the most visited move). The move is sampled from $\pi_t$ in the first moves of a game, and later the most visited move is played. When the game ends with result $z \in \{-1, 0, +1\}$, each position gets the target $z_t = \pm z$, from the view of the player to move at $s_t$. The data is the set of triples $(s_t, \pi_t, z_t)$.

The loss. The network is trained by gradient descent on

the squared error of the value, the cross-entropy from the search policy $\pi$ to the network policy $p$, and weight decay.

Why the visit counts are the policy target. The search looks ahead with the network's own priors and values, so its move choice is usually stronger than the raw $p$: MCTS is a policy improvement operator. Training $p$ towards $\pi$ copies the improvement into the network; the next search starts from a better prior and improves again. The value $v$ is trained towards the real game result $z$, a Monte Carlo target for the current policy (as in lesson 04).

AlphaZero as approximate policy iteration. Compare with lesson 02:

Policy iteration	AlphaZero
policy improvement: $\pi' = \text{greedy}(Q^\pi)$	MCTS with the current network gives $\pi_t$, a better policy
policy evaluation: compute $V^\pi$	train $v$ towards the self-play results $z$
exact, all states	approximate, only the states that self-play visits

AlphaZero as imitation. The search is the expert and the network is the student; this is expert iteration (Anthony, Tian and Barber, 2017). The student is trained on the states that it reaches in its own games, with labels from the expert: exactly the DAgger recipe, so there is no distribution shift. And unlike a human expert, this expert gets better every time the student does.

MuZero and Dreamer: planning in a learned latent space

MuZero (Schrittwieser et al., 2020) does not get the rules. It learns three functions: a representation $h$ that maps the observations to a latent state, a dynamics function $g$ that maps a latent state and an action to the next latent state and a reward, and a prediction function $f$ that gives a policy and a value from a latent state. MCTS runs in the latent space with $g$ as the model. The model is trained to predict, $K$ steps ahead, the rewards, the values and the search policies, not the observations. So it only needs to model what matters for planning. MuZero matched AlphaZero in Go, chess and shogi and set a new state of the art on Atari at the time.

Dreamer (Hafner et al., 2020) learns a recurrent latent world model from pixels, then trains an actor and a critic only on imagined latent rollouts of the model, by backpropagating value estimates through the learned dynamics. Real data is used only to train the world model. DreamerV3 uses one set of hyperparameters across many domains and collected diamonds in Minecraft from scratch.

Answers to section 1. (1) About 3 episodes. The first episode is a long random walk for both, because all values are 0. After that, Dyna-Q's planning updates spread the goal's value over the whole maze, while Q-learning moves it back about one cell per episode (the figure in section 2). (2) Move 3: its UCT score is $0.2 + \sqrt2\sqrt{\ln 20 / 2} = 1.93$, larger than $0.6 + \sqrt2\sqrt{\ln 20 / 10} = 1.37$ for move 1 (section 4 has all the numbers). (3) Much more: about 2.47 steps, 25 times more. One mistake leaves the road for the rest of the episode, so the cost grows like $\varepsilon T^2/2$ (section 4).

4. Worked example #

(a) One MCTS node: UCT and PUCT. The node has $N(s) = 20$ visits and three children with $(N, Q, P)$ = (10, 0.6, 0.5), (8, 0.5, 0.3), (2, 0.2, 0.2). We need $\ln 20 = 2.996$ and $\sqrt{20} = 4.472$.

UCT with $c = \sqrt 2$:

child	$\sqrt{\ln 20 / N}$	bonus $\sqrt2\,\sqrt{\ln 20 / N}$	score $Q$ + bonus
1	$\sqrt{0.2996} = 0.547$	0.774	0.6 + 0.774 = 1.374
2	$\sqrt{0.3745} = 0.612$	0.865	0.5 + 0.865 = 1.365
3	$\sqrt{1.498} = 1.224$	1.731	0.2 + 1.731 = 1.931

UCT explores child 3, the one with few visits.

PUCT with $c_{\text{puct}} = 1.5$, bonus $1.5 \cdot P \cdot 4.472 / (1 + N)$:

child	bonus	score
1	$1.5 \cdot 0.5 \cdot 4.472 / 11 = 0.305$	0.6 + 0.305 = 0.905
2	$1.5 \cdot 0.3 \cdot 4.472 / 9 = 0.224$	0.5 + 0.224 = 0.724
3	$1.5 \cdot 0.2 \cdot 4.472 / 3 = 0.447$	0.2 + 0.447 = 0.647

PUCT picks child 1. The prior of child 3 is small, so its bonus is small too: the network's prior decides where the search explores.

Policy target. The visit shares are $\pi = (10, 8, 2)/20 = (0.5, 0.4, 0.1)$ with $\tau = 1$. With $\tau = 0.5$ we square the counts: $(100, 64, 4)/168 = (0.595, 0.381, 0.024)$. With $\tau \to 0$ we get $(1, 0, 0)$.

(b) A backup with the sign flip. Your move B has $N = 4$, $W = 0.4$ ($Q = 0.1$). The opponent's reply B2 has $N = 2$, $W = -0.2$ ($Q = -0.1$ for the opponent). An iteration selects root → B → B2, and the playout ends with $z = -1$ (you lose).

• B is your move: $N = 5$, $W = 0.4 + (-1) = -0.6$, $Q = -0.12$.
• B2 is the opponent's move: $N = 3$, $W = -0.2 - (-1) = 0.8$, $Q = 0.267$. B2 now looks better for the opponent, so the selection at B is more likely to pick it again.
• The root gets $N + 1$.

(c) Behaviour cloning vs DAgger on a narrow road. $T = 50$ steps, error rate $\varepsilon = 0.002$ per step. A wrong action drives the car off the road; the cloned policy has never seen off-road states and never comes back; every step off the road costs 1.

• The car is off the road after step $t$ if any of the first $t$ actions was wrong: probability $1 - (1 - \varepsilon)^t$.
• Expected cost: $\sum_{t=1}^{50} \big(1 - 0.998^t\big) = 2.47$. For small $\varepsilon T$, $1 - (1-\varepsilon)^t \approx \varepsilon t$, so the sum is about $\varepsilon T(T+1)/2 = 2.55$. Both are below the bound $T^2 \varepsilon = 5$.
• With DAgger the data also contains off-road states, labelled with the expert's recovery action. If one mistake costs one extra step ($u = 1$), the expected extra cost is at most $u T \varepsilon = 0.1$: 25 times less. The gap grows with $T$: for $T = 100$ BC costs 9.46 and DAgger about 0.2.

(d) Compounding model error. True dynamics $x_{t+1} = 1.05\,x_t$, model $\hat x_{t+1} = 1.06\,\hat x_t$, both from $x_0 = 1$. One step: $1.06/1.05 - 1 = 0.0095$, about 1%. After $H$ steps the relative error is $(1.06/1.05)^H - 1$:

• $H = 10$: $1.00952^{10} - 1 = 0.0994$;
• $H = 50$: $1.00952^{50} - 1 = 0.606$, while $50 \times 0.0095 = 0.476$.

(e) Two Dyna-Q updates by hand. In the maze, $\alpha = 0.1$, $\gamma = 0.95$, all $Q = 0$. Call $s_{13}$ the cell just below G and $s_{12}$ the cell below $s_{13}$. The first episode ends with the step "up" from $s_{13}$ into G, reward 1. Direct RL: $Q(s_{13}, \text{up}) = 0 + 0.1\,(1 + 0 - 0) = 0.1$. Now a planning step replays the stored transition $(s_{12}, \text{up}) \to s_{13}$ with reward 0: $Q(s_{12}, \text{up}) = 0 + 0.1\,(0 + 0.95 \cdot 0.1 - 0) = 0.0095$. Without the model, this second value would have to wait for the next real visit to $s_{12}$.

5. In code #

NumPy checks the numbers of the worked example. The notebook has these cells, the Dyna maze, CEM planning, UCT for tic-tac-toe, and the exercises.

import numpy as np

# (a) one MCTS node: N(s) = 20, children (N, Q, P)
N_s = 20
N = np.array([10, 8, 2]); Q = np.array([0.6, 0.5, 0.2]); P = np.array([0.5, 0.3, 0.2])
uct = Q + np.sqrt(2) * np.sqrt(np.log(N_s) / N)
puct = Q + 1.5 * P * np.sqrt(N_s) / (1 + N)
print(uct.round(3), uct.argmax())        # [1.374 1.365 1.931] 2
print(puct.round(3), puct.argmax())      # [0.905 0.724 0.647] 0
for tau in (1.0, 0.5):
    pi = N ** (1 / tau) / np.sum(N ** (1 / tau))
    print(tau, pi.round(3))              # 1.0 [0.5 0.4 0.1]   0.5 [0.595 0.381 0.024]

# (c) behaviour cloning vs DAgger on the road
T, eps = 50, 0.002
t = np.arange(1, T + 1)
print(np.sum(1 - (1 - eps) ** t), eps * T * (T + 1) / 2, eps * T**2, eps * T)
# 2.4687 2.55 5.0 0.1

# (d) compounding model error: true a = 1.05, model 1.06
for H in (1, 10, 50):
    print(H, (1.06 / 1.05) ** H - 1, H * (1.06 / 1.05 - 1))
# 1 0.0095 0.0095   10 0.0994 0.0952   50 0.6063 0.4762

6. Where it shows up #

• Board games. AlphaGo (2016) used MCTS with policy and value networks and rollouts, trained first on human games. AlphaGo Zero (2017) and AlphaZero (2018) removed the human data and the rollouts. MuZero (2020) also removed the rules. KataGo and Leela Chess Zero are open-source AlphaZero-style engines.
• Search for algorithms and molecules. AlphaTensor (matrix multiplication algorithms) and AlphaDev (sorting routines) use AlphaZero-style search with a policy and value network. Retrosynthesis planners search over chemical reactions with MCTS and neural networks.
• Robotics and control. MPC is standard in industry and robotics; with a learned model it becomes model-based RL (PETS, MBPO). Dreamer learns control tasks from pixels with few real samples. Robot policies learned from human demonstrations are behaviour cloning; their failures far from the demonstrated states are compounding error.
• LLMs. Supervised fine-tuning is behaviour cloning. On-policy distillation (the student generates, the teacher labels its tokens) applies the DAgger idea. Some reasoning systems search over partial solutions with tree search guided by a value or reward model.
• ML for science. Autoregressive emulators of weather, climate or fluids are rolled out for many steps, so they are often trained or fine-tuned with multi-step rollout losses. MPC with an emulator inherits model bias: check the multi-step error and whether the controller drives the system outside the training data.

7. Common mistakes #

• Running the whole MPC plan. Run only the first action, observe the real state, and plan again. The plan is only as good as the model, and noise adds up.
• Judging a model by its one-step error. Long open-loop rollouts can drift far; measure the error over the horizon that the planner uses.
• Trusting what the planner found. A plan that is much better in the model than in reality is a sign of model bias: the planner found the model's errors. Use ensembles, short rollouts or MPC.
• Forgetting the sign flip in a two-player backup. Then the opponent chooses the moves that are best for you, and the search plays as if the opponent helps.
• Playing the child with the highest $Q$. A child with two lucky visits can have a high $Q$; play the most visited child.
• UCT's $c$ with the wrong value range. $c = \sqrt 2$ fits results in $[0, 1]$. With results in $[-1, 1]$ the range is twice as large, so $c$ should be larger (or rescale the results).
• "BC error is $\varepsilon T$." Only on the expert's states. On the learner's own states the error can compound to $O(\varepsilon T^2)$.
• DAgger without an expert to ask. DAgger needs expert labels for the learner's states; a fixed dataset of demonstrations is not enough.
• An old model in a changed world. Dyna plans with what the model remembers. If the environment changes, planning spreads wrong values until real experience corrects the model (Dyna-Q+ adds an exploration bonus for this).

8. Check yourself #

1. Write the four steps of Dyna-Q at each real step.
   Answer

   Act ε-greedily and observe $(r, s')$; do a Q-learning update on the real transition; store $\text{Model}(s, a) = (r, s')$; do $n$ Q-learning updates on transitions $(\tilde s, \tilde a, \tilde r, \tilde s')$ replayed from the model at previously seen pairs.
2. Why does the open-loop error of a learned model grow faster than linearly in the horizon when the dynamics expand nearby states?
   Answer

   Each step adds the new one-step error $\varepsilon_m$ and also multiplies the old error by up to $L > 1$: $e_{t+1} \le \varepsilon_m + L e_t$. Unrolled, $e_H \le \varepsilon_m (L^H - 1)/(L - 1)$, which is geometric in $H$.
3. A node has $N(s) = 50$ and two children: (N = 40, Q = 0.55) and (N = 10, Q = 0.45). Compute the UCT scores with $c = \sqrt 2$. Which child is selected?
   Answer

   $\ln 50 = 3.912$. Child 1: $0.55 + \sqrt2\sqrt{3.912/40} = 0.55 + 0.442 = 0.992$. Child 2: $0.45 + \sqrt2\sqrt{3.912/10} = 0.45 + 0.885 = 1.335$. Child 2 is selected.
4. In AlphaZero, which quantity is the target of the value head, and which of the policy head?
   Answer

   Value head: the game result $z \in \{-1, 0, +1\}$ from the view of the player to move. Policy head: the MCTS visit-count distribution $\pi$ at the root (with temperature $\tau$).
5. Why is DAgger's extra cost linear in $T$ while BC's can be quadratic?
   Answer

   BC is trained only on the expert's states, so after a mistake the learner is in states without labels and may never recover: up to $T$ cost per mistake. DAgger adds expert labels for the states the learner visits, so its training and test distributions match and each mistake costs a bounded amount ($u$), giving $uT\varepsilon$.
6. Name one reason why model-based RL often reaches a worse final performance than model-free RL, and one method that reduces the problem.
   Answer

   Model bias: the planner exploits model errors, and compounding error limits long rollouts. MBPO uses an ensemble and short rollouts from real states; MPC replans from the real state at every step.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. What does Dyna-Q do after each real step?

   Show answer

2. What is model-predictive control (receding horizon)?

   Show answer

3. Bound the open-loop error after $H$ steps for one-step model error $\varepsilon_m$ and $L$-Lipschitz dynamics.

   Show answer

4. Write the UCT selection rule.

   Show answer

5. Write the PUCT selection rule of AlphaZero.

   Show answer

6. Give the four phases of one MCTS iteration in one sentence.

   Show answer

7. Write AlphaZero's loss.

   Show answer

8. What is the extra-cost bound of behaviour cloning, and why is it quadratic?

   Show answer

9. What is the idea of DAgger?

   Show answer

10. What does MuZero's model predict, and what does it not predict?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Model-based or model-free?" — Model-based is more sample-efficient because each sample is reused in planning; model-free often wins at the end because it does not suffer from model bias. Mention short rollouts (MBPO), ensembles and MPC.
• "Explain MCTS." — Four phases; UCT = UCB1 per node; play the most visited move; sign flip for two players. Know one UCT computation by heart.
• "How does AlphaZero train?" — Self-play with PUCT search; targets $(\pi, z)$; loss $(z - v)^2 - \pi^\top \log p + c\|\theta\|^2$; MCTS is policy improvement, so the loop is approximate policy iteration; no human data.
• "What is the problem with behaviour cloning?" — Distribution shift: compounding error $O(\varepsilon T^2)$. Fixes: DAgger (needs a queryable expert), noise injection and recovery data, or RL fine-tuning.
• "How is SFT related to RL?" — SFT is behaviour cloning; RLHF then optimises a reward with a KL constraint to the SFT policy (lesson 09).
• "Why would a planner do badly with a model that has a low test error?" — Low one-step error on the data distribution says nothing about long rollouts or about states the planner drives into.

10. Practice #

The notebook for this lesson runs in Colab. It has five exercises, each with a check cell and a folded solution:

1. Dyna-Q vs Q-learning on the Dyna maze: implement dyna_q and compare $n = 0, 5, 50$.
2. Plan with a learned model: fit a linear model by least squares and implement CEM planning with MPC; compare the true and the learned model.
3. UCT and PUCT scores: implement uct_scores and puct_scores and check them on the worked example.
4. UCT for tic-tac-toe: implement MCTS with random playouts and the sign flip; it must find a win in one move, block a threat, and beat a random player in almost every game.
5. Behaviour cloning vs DAgger on the toy road: implement the DAgger loop and compare the costs as $T$ grows.

11. Watch and read #

Monte Carlo Tree SearchJohn Levine · 16 min

The four phases with UCB1 on a small tree, computed by hand.

RL Course by David Silver - Lecture 8: Integrating Learning and PlanningGoogle DeepMind · 100 min

Models, Dyna, simulation-based search and MCTS.

RL Course by David Silver - Lecture 10: Classic GamesGoogle DeepMind · 111 min

Minimax, self-play and search in games, before AlphaZero.

DeepMind x UCL RL Lecture Series - Planning & models [8/13]Google DeepMind · 58 min

Dyna, model errors and planning at decision time.

L6 Model-based RL (Foundations of Deep RL Series)Pieter Abbeel · 18 min

Learned dynamics models, MPC, ensembles and MBPO in one short lecture.

CS 285: Lecture 2, Imitation Learning. Part 1RAIL · 25 min

Behaviour cloning and the distribution shift argument, by Sergey Levine.

MuZero: Mastering Atari, Go, Chess and Shogi by Planning with a Learned ModelYannic Kilcher · 19 min

How MuZero's latent model is trained to predict rewards, values and policies.

• Sutton and Barto, Reinforcement Learning: An Introduction, 2nd ed., free PDF: chapter 8 (§8.1–8.4 models, Dyna and prioritized sweeping; §8.11 MCTS).
• Kocsis and Szepesvári, Bandit based Monte-Carlo Planning (UCT), 2006.
• Silver et al., Mastering Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm (AlphaZero), 2017.
• Anthony, Tian and Barber, Thinking Fast and Slow with Deep Learning and Tree Search (expert iteration), 2017.
• Schrittwieser et al., Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model (MuZero), 2019.
• Ross, Gordon and Bagnell, A Reduction of Imitation Learning and Structured Prediction to No-Regret Online Learning (DAgger), §2–3.
• Ross and Bagnell, Efficient Reductions for Imitation Learning (the $T^2\varepsilon$ bound), 2010.
• Janner, Fu, Zhang and Levine, When to Trust Your Model: Model-Based Policy Optimization (MBPO).
• Chua, Calandra, McAllister and Levine, Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models (PETS).
• Hafner, Lillicrap, Ba and Norouzi, Dream to Control: Learning Behaviors by Latent Imagination (Dreamer).