gumbel-mcts

An implementation of Policy improvement by planning with Gumbel (Danihelka et al., ICLR 2022). Numba-accelerated, hundreds of thousands of simulations per second.

Gumbel MCTS replaces PUCT's UCB-based action selection with sequential halving over Gumbel-perturbed log-priors. This produces better actions under the same simulation budget — especially at low budgets, where PUCT wastes visits on bad branches it can't prune.

PushT — uniform prior, 64 sims	Wall-gap PushT — push through a narrow gap, 64 sims
Robot arm — 4 joints, 2 obstacles, 64 sims	Sokoban — 2-box puzzle, 32 sims

PUCT (left) vs Gumbel (right), same budget, same model. Click to play.

All demos are fully reproducible — see demo/ for scripts, sweep results, and instructions.

What's included

Three MCTS implementations sharing the same tree storage and batch interface:

See gumbel-mcts-benchmark for validation against a gold-standard MCTS.

Installation

pip install gumbel-mcts

Quick start

from gumbel_mcts import GumbelDense

tree = GumbelDense(n_games=1, max_nodes=500, logic=logic, device="cpu")
tree.initialize_roots([0], board[None], np.array([player]))
action = tree.run_simulation_batch(model, [0], num_simulations=64)[0]

The interface is the same for all three algorithms — swap GumbelDense for PUCT or GumbelSparse.

How it works

Improving MuZero using the Gumbel top-k trick, by Xavier O'Keefe

Standard PUCT explores actions proportionally to a UCB score that mixes prior probability and visit-based value estimates. When the budget is small, most actions get only a few visits, and the value estimates are noisy — PUCT can't reliably distinguish good actions from bad ones.

Gumbel MCTS takes a different approach: add Gumbel noise to the log-prior to create an initial ranking, then use sequential halving to eliminate half the candidates at each phase. Each phase allocates equal simulations to the surviving candidates. The result is that the budget concentrates on a shrinking set of promising actions, producing a much sharper signal with the same total simulation count.

Board game results

Head-to-head on 15x15 Gomoku across simulation budgets. Gumbel's advantage is largest at low budgets and persists as the model improves:

With 8 simulations on 9x9 Gomoku, Gumbel concentrates visits on the strategic moves while PUCT spreads them across irrelevant squares: