Monte Carlo tree search

Deep dive

How it works

Parts:

• tree — nodes with prior P, visits n, value sum W, children contiguous in an arena
• selection rule — UCT or PUCT with first-play urgency
• evaluator — random playout or a policy/value network
• backup — adds the value along the path, flipping sign per level for two-player alternation
• root policy — most-visited child, optional Dirichlet noise and temperature

Operations: select, expand and evaluate, back up, decide, reuse (steps below); the diagram shows the root after a few simulations.

Each edge (s, a) keeps visit count n, total value W, prior P. Q = W/n is stored from the view of the player who made the move, so every parent simply maximizes its children's Q.

1. Select. From the root, repeatedly take argmax Q(a) + c·P(a)·√N / (1 + n(a)) until reaching an unexpanded or terminal node. Unvisited children have no Q; use a first-play urgency (FPU): parent's Q minus a reduction (Lc0) or a fixed value (AlphaZero pseudocode used the loss value).
2. Expand + evaluate. Terminal: exact result. Otherwise call the network once: priors over legal moves (softmax restricted to legal moves) and value v ∈ [−1, 1] for the side to move. Create children with n = 0.
3. Back up. Walk the path upwards; add −v to the leaf's W (it is stored for the player who moved into it), negate at each level, increment n.
4. Decide. After the budget: argmax n at the root, or sample ∝ n^(1/τ) (self-play: τ = 1 for the first 30 plies in AlphaZero, then greedy).
5. Reuse. After the move, keep the chosen child's subtree as the new root.

root N=40
 ├─ e4  P=.40 n=25 Q=+.12   ← exploited
 ├─ d4  P=.35 n=13 Q=+.10
 └─ a3  P=.02 n=1  Q=-.30   ← explored once, prior keeps it low

Code

pub trait Env: Clone {
    type Move: Copy;
    fn legal(&self) -> Vec<Self::Move>;
    fn play(&mut self, m: Self::Move);
    /// Some(value for the side to move, in [-1, 1]) if the game is over.
    fn terminal(&self) -> Option<f32>;
}

pub trait Evaluator<E: Env> {
    /// Priors over `moves` (summing to 1) and value in [-1, 1] for the side to move.
    fn eval(&mut self, env: &E, moves: &[E::Move]) -> (Vec<f32>, f32);
}

struct Node<M> {
    mv: Option<M>,
    prior: f32,
    n: u32,
    w: f32,           // sum of values for the player who moved into this node
    first_child: u32, // children are contiguous in the arena
    n_children: u32,
    expanded: bool,
}

pub struct Mcts<E: Env> { nodes: Vec<Node<E::Move>>, pub c_puct: f32, pub fpu_reduction: f32 }

impl<E: Env> Mcts<E> {
    pub fn new(c_puct: f32, fpu_reduction: f32) -> Self { Mcts { nodes: Vec::new(), c_puct, fpu_reduction } }

    pub fn search(&mut self, root: &E, ev: &mut impl Evaluator<E>, sims: u32) -> Vec<(E::Move, u32)> {
        self.nodes.clear();
        self.nodes.push(Node { mv: None, prior: 1.0, n: 0, w: 0.0, first_child: 0, n_children: 0, expanded: false });
        let mut path = Vec::with_capacity(128);
        for _ in 0..sims {
            let (mut env, mut cur) = (root.clone(), 0usize);
            path.clear();
            path.push(0);
            while self.nodes[cur].expanded && self.nodes[cur].n_children > 0 {
                cur = self.select(cur);
                env.play(self.nodes[cur].mv.unwrap());
                path.push(cur);
            }
            let v = match env.terminal() {
                Some(v) => v,
                None => {
                    let moves = env.legal();
                    let (priors, v) = ev.eval(&env, &moves);
                    self.expand(cur, &moves, &priors);
                    v
                }
            };
            let mut v = -v; // leaf value is for its side to move; W is stored for the mover
            for &i in path.iter().rev() {
                self.nodes[i].n += 1;
                self.nodes[i].w += v;
                v = -v;
            }
        }
        self.children(0).map(|c| (c.mv.unwrap(), c.n)).collect()
    }

    fn children(&self, i: usize) -> impl Iterator<Item = &Node<E::Move>> {
        let (s, k) = (self.nodes[i].first_child as usize, self.nodes[i].n_children as usize);
        self.nodes[s..s + k].iter()
    }

    fn expand(&mut self, i: usize, moves: &[E::Move], priors: &[f32]) {
        let first = self.nodes.len() as u32;
        for (&m, &p) in moves.iter().zip(priors) {
            self.nodes.push(Node { mv: Some(m), prior: p, n: 0, w: 0.0, first_child: 0, n_children: 0, expanded: false });
        }
        let nd = &mut self.nodes[i];
        nd.first_child = first;
        nd.n_children = moves.len() as u32;
        nd.expanded = true;
    }

    fn select(&self, parent: usize) -> usize {
        let p = &self.nodes[parent];
        let sqrt_n = (p.n as f32).sqrt();
        let q_parent = if p.n > 0 { -p.w / p.n as f32 } else { 0.0 }; // parent's own view
        let fpu = q_parent - self.fpu_reduction;
        let start = p.first_child as usize;
        let (mut best, mut best_score) = (start, f32::NEG_INFINITY);
        for (k, c) in self.children(parent).enumerate() {
            let q = if c.n > 0 { c.w / c.n as f32 } else { fpu };
            let score = q + self.c_puct * c.prior * sqrt_n / (1.0 + c.n as f32);
            if score > best_score { best_score = score; best = start + k; }
        }
        best
    }

    /// Mix exploration noise into the root priors; `eta` ~ Dirichlet(alpha), i.e.
    /// independent Gamma(alpha, 1) samples normalized to sum 1. Call after the first simulation.
    pub fn add_root_noise(&mut self, eta: &[f32], eps: f32) {
        let s = self.nodes[0].first_child as usize;
        for (k, &e) in eta.iter().enumerate() {
            let c = &mut self.nodes[s + k];
            c.prior = (1.0 - eps) * c.prior + eps * e;
        }
    }
}

Complexity and constants

• Selection is O(depth × branching) per simulation: scanning ~35 children per level in chess is trivial next to a network call.
• Neural evaluation dominates: order of 10⁻⁵-10⁻² s per position depending on network size, hardware and batching. AlphaZero self-play used 800 simulations per move; Lc0 at tournament time controls searches order of 10⁴-10⁵ nodes/s on a modern GPU.
• Memory: one node per expanded child. The struct above is 28 B with a 16-bit move; Lc0 keeps edges (move + compressed prior) separate from visited nodes to shrink this further.
• Visit-count estimates converge like any Monte Carlo average: value noise shrinks as 1/√n per edge.

Efficiency

• Cost: per simulation, selection is O(depth × branching) on the CPU and one evaluation; the evaluator dominates, order of 10⁻⁵-10⁻² s per input depending on network, hardware and batching. Memory: one node per expanded child, 28 bytes in the struct above. Bound: evaluator throughput.
• Levers:
  • batch leaves with virtual loss, or one leaf per parallel search — fills the GPU (see batched MCTS inference) — learned evaluators
  • tree reuse after the move — keeps earlier visits — sequential decisions
  • evaluation cache by hash — skips repeated states — transpositions
  • solver propagation of proven results — stops spending visits on decided branches — problems with exact terminals
  • Gumbel root selection — policy improvement at very few simulations (Danihelka et al. 2022) — small budgets
• Naive → standard → state of the art:
  • UCT with random playouts · thousands of cheap simulations per second, no learned prior
  • PUCT with a network · AlphaZero: 800 simulations per move in self-play
  • batched, cached, Gumbel root, reused tree · Lc0 of order 10⁴-10⁵ nodes/s on a modern GPU

Implementation notes

• Arena + u32 indices, children contiguous: one allocation, cache-friendly scans, trivially freed. Tree reuse = copy the chosen subtree into a fresh arena (or keep a free list).
• Batching. GPUs need batches. Select many leaves per batch using virtual loss: during selection add a temporary visit and a loss to each node on the path so the next selection diverges; undo on backup. Also keep a cache keyed by Zobrist hash for repeated positions.
• Terminal handling. Store proven wins/losses and propagate them (MCTS-solver) so the search stops spending visits on decided moves.
• Value range. Keep Q in [−1, 1] (or [0, 1] consistently). Mixing conventions silently inverts FPU or the noise.
• Priors. Softmax over legal-move logits only, optionally with a policy temperature > 1 at the root.
• Transpositions. A tree treats transposed positions as different nodes. Sharing nodes via a hash table turns it into a DAG, but visit counts and backup rules then need care.
• Draws and repetitions. Detect game-rule draws during selection (path history), not only in the network.

Pitfalls

• Sign errors in backup: the search then plays the opponent's best move. Test on a mate-in-1 and on a position where every move but one loses.
• Picking the highest-Q child at the end: a child with n = 1 and a lucky value wins. Use visit counts.
• No root noise in self-play: the network's blind spots are never explored, training stalls.
• FPU too optimistic (unvisited = +1): every child is visited once before any is deepened, wasting simulations at wide nodes.
• Recursion or Vec allocation per simulation in the hot loop; reuse the path buffer.
• Virtual loss left in the tree after an aborted batch.

Variants and alternatives

In production

• Games: AlphaGo/AlphaGo Zero/AlphaZero and MuZero (DeepMind), Leela Chess Zero, KataGo, Leela Zero, ELF OpenGo.
• Chemistry: retrosynthesis route search with a neural policy guiding MCTS (Segler et al. 2018).
• Theorem proving: HyperTree Proof Search, an AlphaZero-inspired search over proof trees (Lample et al. 2022).
• Systems and compilers: MuZero rate control for VP9 video encoding (Mandhane et al. 2022); AlphaDev's search for faster sorting routines (Mankowitz et al. 2023).
• LLM reasoning: RAP uses MCTS over reasoning steps (Hao et al. 2023); a PUCT-style tree guides decoding (Feng et al. 2023).
• Partially observable planning: POMCP (Silver and Veness 2010).

Implementations

Sources

• Levente Kocsis, Csaba Szepesvári, Bandit Based Monte-Carlo Planning, ECML 2006. paper
• Rémi Coulom, Efficient Selectivity and Backup Operators in Monte-Carlo Tree Search, Computers and Games 2006.
• Peter Auer, Nicolò Cesa-Bianchi, Paul Fischer, Finite-time Analysis of the Multiarmed Bandit Problem, Machine Learning 47, 2002. paper
• Christopher D. Rosin, Multi-armed Bandits with Episode Context, Annals of Mathematics and Artificial Intelligence 61, 2011. paper
• David Silver et al., Mastering the Game of Go without Human Knowledge, Nature 550, 2017. paper
• David Silver et al., A General Reinforcement Learning Algorithm that Masters Chess, Shogi, and Go through Self-Play, Science 362, 2018. paper
• David J. Wu, Accelerating Self-Play Learning in Go, arXiv:1902.10565, 2019. paper
• Ivo Danihelka, Arthur Guez, Julian Schrittwieser, David Silver, Policy Improvement by Planning with Gumbel, ICLR 2022.
• Cameron Browne et al., A Survey of Monte Carlo Tree Search Methods, IEEE Trans. Computational Intelligence and AI in Games 4(1), 2012. paper
• Marc Segler, Mike Preuss, Mark P. Waller, Planning chemical syntheses with deep neural networks and symbolic AI, Nature 555, 2018. paper
• Guillaume Lample et al., HyperTree Proof Search for Neural Theorem Proving, NeurIPS 2022. paper
• Shibo Hao et al., Reasoning with Language Model is Planning with World Model, arXiv:2305.14992, 2023. paper
• Xidong Feng et al., Alphazero-like Tree-Search can Guide Large Language Model Decoding and Training, arXiv:2309.17179, 2023. paper
• Amol Mandhane et al., MuZero with Self-competition for Rate Control in VP9 Video Compression, arXiv:2202.06626, 2022. paper
• Daniel J. Mankowitz et al., Faster sorting algorithms discovered using deep reinforcement learning, Nature 618, 2023. paper
• David Silver, Joel Veness, Monte-Carlo Planning in Large POMDPs, NIPS 2010.