Table of Contents
Fetching ...

Variance-Aware Prior-Based Tree Policies for Monte Carlo Tree Search

Maximilian Weichart

TL;DR

This work addresses the lack of principled prior-based UCTs beyond $UCT1$ by introducing Inverse-RPO, a pipeline that derives prior-based tree policies from any prior-free UCB. By applying it to the variance-aware base $UCB{-}V$, it derives two priors, $UCT{-}V{-}P$ and $PUCT{-}V$, that incorporate variance estimates through $\hat{\sigma}_a$ and prior weights, yielding improved exploration. Empirical results across benchmarks show variance-aware priors matching or surpassing $UCT{-}P$ and $PUCT$ with no extra computation, and the authors provide a light-weight code extension to existing MCTS implementations. Overall, the approach unifies MCTS with RPO to enable principled design of prior-based UCTs and opens avenues for further principled exploration of variance-aware tree policies.

Abstract

Monte Carlo Tree Search (MCTS) has profoundly influenced reinforcement learning (RL) by integrating planning and learning in tasks requiring long-horizon reasoning, exemplified by the AlphaZero family of algorithms. Central to MCTS is the search strategy, governed by a tree policy based on an upper confidence bound (UCB) applied to trees (UCT). A key factor in the success of AlphaZero is the introduction of a prior term in the UCB1-based tree policy PUCT, which improves exploration efficiency and thus accelerates training. While many alternative UCBs with stronger theoretical guarantees than UCB1 exist, extending them to prior-based UCTs has been challenging, since PUCT was derived empirically rather than from first principles. Recent work retrospectively justified PUCT by framing MCTS as a regularized policy optimization (RPO) problem. Building on this perspective, we introduce Inverse-RPO, a general methodology that systematically derives prior-based UCTs from any prior-free UCB. Applying this method to the variance-aware UCB-V, we obtain two new prior-based tree policies that incorporate variance estimates into the search. Experiments indicate that these variance-aware prior-based UCTs outperform PUCT across multiple benchmarks without incurring additional computational cost. We also provide an extension of the mctx library supporting variance-aware UCTs, showing that the required code changes are minimal and intended to facilitate further research on principled prior-based UCTs. Code: github.com/Max-We/inverse-rpo.

Variance-Aware Prior-Based Tree Policies for Monte Carlo Tree Search

TL;DR

This work addresses the lack of principled prior-based UCTs beyond by introducing Inverse-RPO, a pipeline that derives prior-based tree policies from any prior-free UCB. By applying it to the variance-aware base , it derives two priors, and , that incorporate variance estimates through and prior weights, yielding improved exploration. Empirical results across benchmarks show variance-aware priors matching or surpassing and with no extra computation, and the authors provide a light-weight code extension to existing MCTS implementations. Overall, the approach unifies MCTS with RPO to enable principled design of prior-based UCTs and opens avenues for further principled exploration of variance-aware tree policies.

Abstract

Monte Carlo Tree Search (MCTS) has profoundly influenced reinforcement learning (RL) by integrating planning and learning in tasks requiring long-horizon reasoning, exemplified by the AlphaZero family of algorithms. Central to MCTS is the search strategy, governed by a tree policy based on an upper confidence bound (UCB) applied to trees (UCT). A key factor in the success of AlphaZero is the introduction of a prior term in the UCB1-based tree policy PUCT, which improves exploration efficiency and thus accelerates training. While many alternative UCBs with stronger theoretical guarantees than UCB1 exist, extending them to prior-based UCTs has been challenging, since PUCT was derived empirically rather than from first principles. Recent work retrospectively justified PUCT by framing MCTS as a regularized policy optimization (RPO) problem. Building on this perspective, we introduce Inverse-RPO, a general methodology that systematically derives prior-based UCTs from any prior-free UCB. Applying this method to the variance-aware UCB-V, we obtain two new prior-based tree policies that incorporate variance estimates into the search. Experiments indicate that these variance-aware prior-based UCTs outperform PUCT across multiple benchmarks without incurring additional computational cost. We also provide an extension of the mctx library supporting variance-aware UCTs, showing that the required code changes are minimal and intended to facilitate further research on principled prior-based UCTs. Code: github.com/Max-We/inverse-rpo.
Paper Structure (30 sections, 42 equations, 2 figures, 2 tables)

This paper contains 30 sections, 42 equations, 2 figures, 2 tables.

Figures (2)

  • Figure 1: By casting the MCTS tree policy as the solution to an RPO objective, the prior becomes an explicit design term, yielding a principled prior-based UCT selection rule. This perspective resolves the otherwise opaque step from prior-free UCT to prior-based UCT and motivates our Inverse-RPO methodology.
  • Figure 2: Average returns on the MinAtar suite with $N_{\mathrm{sim}}=64$. Evaluation is performed in batches of $256$ per seed (at least $3$ seeds), using only the trained policy head without MCTS. Left:UCT-V-P vs. UCT-P. Right:PUCT vs. PUCT-V. Solid lines indicate mean returns, and shaded regions show the corresponding best--worst range across seeds.