A note from the author:

This article was first written for my personal notes when doing research for my project Py2048-ML v2. It is assumed that readers already have some experience with the fundamentals of reinforcement learning, such as states, environments, value functions, policies, etc. If you wish to learn about these ideas, I highly recommend this lecture series by David Silver.

What is MCTS?

Monte Carlo Tree Search (MCTS) is a scheme for approximating the action-value function of a game-tree. It aims to improve on Monte Carlo policy evaluation, which suffers from very slow convergence due to needing to trajectories to completion. It overcomes this by only evaluating partial trajectories, then estimating the value for the remainder of the trajectory.

Implementing MCTS requires the following components:

• A model of the environment (or a simulator) that when given a state-action pair, returns the next state (and/or resulting afterstate), and any transition rewards that were generated by taking that state-action pair.
• A scheme for selecting actions when generating partial trajectories. Usually this is done by assigning some “MCTS score” (eg. UCB, UCT, PUCT) to each state (or state-action pair), then acting greedily wrt. to that score.
• A scheme for quickly estimating the value of a state. Since the value of a state is the expected sum of future rewards, estimating the value of the last state in a partial trajectory is equivalent to estimating the total reward of the rest of the trajectory. These quick estimates are used to determine if a trajectory is worth exploring further, without having to do a full evaluation.

Common approximation techniques are rollouts and neural networks. With rollouts, we play out games according to some fast policy (random moves, greedy, etc.), and set the value of the leaf to be the mean reward acquired during rollout. With a neural network, and use that.

Note: I will use the term “node” to interchangeably refer to a state, state-action pair, or afterstate. Whether you consider a node to be a state, an afterstate, or the state-action pair which resulted in that state/afterstate is a matter of convention

We estimate the value of nodes by sampling trajectories through the game tree. The value of a node is simply average reward acquired after visiting that node in each trajectory. Each trajectory consists of 2 phases: the in-tree phase, and the after-tree phase.
The in-tree phase consists of nodes which have previously been visited in the search. In this phase, actions are selected according to MCTS score. The reward accumulated during this phase comes from the actual transition rewards given by the model (or simulator).
The after-tree phases begins when we reach a leaf node: either a state-action pair that hasn’t been taken, or a state that hasn’t been visited, depending on scheme. At this point, we use the value approximation scheme to assign an initial value to this leaf.

Exploring a trajectory consists of 4 steps:

1. Selection. We use the selection scheme to choose actions and navigate through the tree, keeping track of any transition rewards we accumulate. Selection stops when we encounter a leaf.
2. Expansion. We expand the leaf, adding its children to the tree. We use the value approximation scheme to assign initial values to the children. This might be rollouts or neural network evaluation.
3. Backup. For each node in the selected path, increase its total value by the (potentially discounted) value of the newly expanded leaf, and any transition rewards that were accumulated after that node.
4. Repeat. After exploring some number of trajectories (or exploring for some fixed amount of time), the search terminates. The selected action (from the root) is the one which was taken the most number of times. Final value estimates for each node are given by their total value, divided by the number of times that node was visited.

DeepMind’s MCTS, Summarized

Deep Mind uses the same general scheme for its MCTS in AlphaZero, MuZero etc.

Overview from AlphaZero Paper:

Each search (i.e. determination of the “best” move) consists of a series of simulated games of self-play that traverse a tree from root $s_{root}$ to leaf. Each simulation proceeds by selecting in each state $s$ a move $a$ with low visit count, high move probability and high value (averaged over the leaf states of simulations that selected $a$ from $s$) according to the current neural network $f_θ$. The search returns a vector $\pi$ representing a probability distribution over moves, either proportionally or greedily with respect to the visit counts at the root state.

In AlphaGo and AlphaZero, environment model is actually a simulator with perfect knowledge of the state transitions. These networks were only used on zero-sum games, so no transition rewards exist. In MuZero, the environment model is learned, and allows for transition rewards.

The value approximation and selection schemes both depend on the output of a neural network. This network takes in a state and outputs 1. the value of that state, and 2. the probability that it would take each action in that state (i.e. a policy) For value approximation of a leaf, we simply query the network and take the value of the leaf state. The selection scheme for the in-tree phase is PUCT, which mixes the sum of returns found during MCTS and the networks current policy. The MCTS score in PUCT is given as follows:

\[PUCT(s,a) = Q(s,a) + c*p(s,a) * \frac{\sqrt{parent.N}}{child.N + 1}\]

where

• $Q(s,a)$ is the current monte-carlo estimate of the node’s value (i.e. the avg return of trajectories in this search). Note while this includes value estimates generated by the neural network, it also includes “real” transition rewards generated by the environment model (or simulator).
• $c$ is a tunable parameter between 0 and 1 (higher = more exploration)
• $p(s,a)$ is the prior probability of the action being selected, according to the policy network
• parent.N is the number of times the current state has been visited
• child.N is the number of times the afterstate of the pair $(s,a)$ has been visited.

While many games such has chess have value functions which are bounded between 0 and 1, many games (such as 2048) can have theoretically unbounded values. In order for $Q(s,a)$ term to not dominate the PUCT rule, we substitute $Q(s,a)$ for normalized value estimates $\overline{Q} \in [0, 1]$ by using the max and min values observed in the search tree up to that point. Let $Q_{max}$ and $Q_{min}$ be the largest and smallest $Q$ among all nodes in the current tree. Then, for an edge $(s, a)$, the normalized value is $\overline{Q}(s,a) = \frac{Q(s,a)-Q_{min}}{Q_{max}}$.Since the max and min values will change as search progresses, the normalized value must be re-computed each time the edge is considered for selection (i.e. each time the PUCT rule is used).

To play a game, we repeatedly use MCTS to select actions, each time discarding the tree and stating fresh. Each completed MCTS returns a value estimate for the root node, and counts for how many times we took each possible action, from the root node. We can use the value estimate as a “value label” for the root state, and the counts as a “policy label” for the root state (since the counts can be normalized and interpreted as a probability distribution). These labels can then be used as targets to train the network.