20. MCTS Node Selection via UCB1

Problem

Monte Carlo Tree Search balances exploiting children that have looked good so far against exploring children that haven't been visited enough to trust yet. UCB1 (Auer, Cesa-Bianchi & Fischer, 2002) makes that tradeoff precise, and it's exactly the selection rule at the core of AlphaGo/AlphaZero's tree search (in a modified, policy-weighted form called PUCT). At a parent node with several children, each child i gets a score: UCB1(i) = Q(i) + c * sqrt(ln(N_parent) / N(i)), and search always descends into the highest-scoring child.

A child that has never been visited (N(i) == 0) needs a special case — the formula divides by zero, and the whole point of UCB1 is that unvisited children should be explored before their true value is trusted, so they must win selection unconditionally. This problem works on a batch of parent nodes at once, as you'd want when selecting simultaneously across many in-flight simulations.

Source: src/20_mcts_ucb1_selection.py

def ucb1_select(
    q_values: torch.Tensor, visit_counts: torch.Tensor,
    parent_visits: torch.Tensor, c: float = 1.4142135623730951,
) -> torch.Tensor: ...

>>> q = torch.tensor([[0.5, 0.1, 0.3]])
>>> visits = torch.tensor([[10.0, 10.0, 0.0]])
>>> parent_visits = torch.tensor([20.0])
>>> ucb1_select(q, visits, parent_visits)
tensor([2])

Step-by-Step Approach

1. Compute the exploration bonus for every child: c * sqrt(log(parent_visits) / visit_counts), broadcasting the per-parent parent_visits against the per-child visit_counts.
2. Add a tiny epsilon to the visit-count denominator so the divide stays finite even before the unvisited-child override is applied.
3. Add the exploration bonus to the mean Q-value to get each child's UCB1 score.
4. Override any child with visit_counts == 0 to a score of +inf — unconditionally beating every visited child regardless of Q-value, using torch.where rather than a Python-level branch so it stays vectorized across the batch.
5. Take argmax over the child dimension to get the selected child index per parent node.

UCB1's exploration term already goes to infinity in the limit as N(i) approaches zero, so the “unvisited child always wins” rule isn't a special exception bolted onto the formula — it's the formula's own limiting behavior, just computed directly (via an explicit override) instead of relying on floating-point division to blow up correctly.

Reference solution

import torch


def ucb1_select(
    q_values: torch.Tensor, visit_counts: torch.Tensor,
    parent_visits: torch.Tensor, c: float = 1.4142135623730951,
) -> torch.Tensor:
    # add a tiny eps to the denominator only to keep the divide finite for
    # visited children -- unvisited children (N == 0) get overridden below
    # regardless of what this produces, so the eps never affects the result
    exploration = c * torch.sqrt(
        torch.log(parent_visits).unsqueeze(-1) / (visit_counts + 1e-8)
    )
    scores = q_values + exploration

    # unvisited children must win unconditionally: force +inf so they beat
    # every visited child's score no matter how good the Q-value is there
    unvisited = visit_counts == 0
    scores = torch.where(unvisited, torch.full_like(scores, float("inf")), scores)

    return scores.argmax(dim=-1)

Key Functions & Tricks

• torch.log — the exploration term's ln(N_parent) factor, computed once per parent and broadcast across its children
• Tensor.unsqueeze(-1) — reshapes the per-parent parent_visits from (batch,) to (batch, 1) so it broadcasts against (batch, num_children)
• torch.where(condition, x, y) — vectorized elementwise branch that overrides unvisited children's score to +inf without a Python loop
• float("inf") — guarantees any unvisited child beats every visited child's finite score, exactly matching UCB1's asymptotic behavior
• Tensor.argmax(dim=-1) — selects the winning child index per parent row

How to Recognize This Pattern