simple-mcts

A lightweight, generic Monte Carlo Tree Search (MCTS) engine for Rust, built around a cache-friendly arena allocator.

simple-mcts gives you the low-level primitives to run MCTS/PUCT-style search (à la AlphaZero) on any turn-based, zero-sum game, without imposing a particular game representation. You bring your own state, your own evaluator (neural net, rollout, heuristic...), and the engine handles selection, expansion, and backpropagation over a flat, index-based tree.

Features

Generic over the action space — the engine is parameterized by N, the fixed number of possible actions per state (Engine<N>).
Arena-based tree — nodes live in a contiguous Vec and are addressed by a compact NonZeroU32-backed index, avoiding Rc<RefCell<...>> and reference-counting overhead.
Pluggable selection & backpropagation — provide your own selection heuristic (PUCT is included) and score transformer (e.g. negation for alternating-turn games).
Tree reuse — commit_action re-roots the tree on the move actually played, keeping the relevant subtree and discarding the rest, with amortized compaction.
Batch mode — MctsBatch manages thousands of independent Engine instances behind a slot-map allocator, useful for self-play or vectorized evaluation pipelines.
no_std-friendly core logic (the crate only depends on std::num::NonZeroU32 and standard collections — no heavyweight dependencies).

Installation

[dependencies]
simple-mcts = "0.2"

Core concepts

The library is built around a small set of pieces you assemble yourself:

Action — a thin wrapper around a usize index identifying one of the N legal moves.
StateEvaluation<N> — the result of evaluating a game state: either Active(score, policies) for an ongoing game, or Terminal(score) once the game is over. score is always expressed from the perspective of the player who just moved.
A selection function — any Fn(score, node_visits, parent_visits, policy, c) -> f32, e.g. the provided puct.
A score transformer — any FnMut(f32) -> f32 applied while backpropagating, typically negate_score for two-player alternating games. A search iteration always follows the same three-step loop:

select — walk the tree from the root using your selection heuristic until hitting a leaf or a terminal node, recording the path taken.
evaluate — replay that path on your own game state and evaluate the resulting position (rollout, heuristic, neural net...).
update (or expand + backpropagate separately) — insert the new node and propagate its score back to the root.

Quick start

use simple_mcts::{negate_score, puct, visits_to_probabilities, Action, Engine, StateEvaluation};
 
// `MyGame` is your own game state, with `N` possible actions.
const N: usize = 9;
const C: f32 = std::f32::consts::SQRT_2;
 
let mut engine = Engine::<N>::new();
let game = MyGame::new();
let mut path = Vec::new();
 
for _ in 0..10_000 {
    // 1. Select a leaf according to PUCT.
    let selection = engine.select(&mut path, &puct, C);
 
    // 2. Replay the path on a fresh copy of the game and evaluate it.
    let mut state = game.clone();
    for action in path.iter() {
        state.play(*action);
    }
    let evaluation = evaluate(&state); // your own evaluator -> StateEvaluation<N>
 
    // 3. Expand + backpropagate, negating the score at each ply.
    engine.update(evaluation, selection, negate_score).unwrap();
}
 
// Turn root visit counts into a move-selection distribution.
let move_probabilities = visits_to_probabilities(engine.scores());

See examples/tic_tac_toe.rs and examples/connect4.rs for complete, runnable implementations (Connect 4 uses random rollouts as its evaluator).

cargo run --release --example tic_tac_toe
cargo run --release --example connect4

Tree reuse across moves

Once you've picked and played a real move, call commit_action to re-root the engine's internal tree on that branch instead of throwing the whole search away:

engine.commit_action(Action::new(chosen_action));

If the branch was never explored, the tree is simply cleared. Otherwise the matching subtree becomes the new root, and unreachable nodes are compacted away once their overhead becomes significant.

Batch search with `MctsBatch`

For self-play or any workload that needs many independent trees driven by the same hyperparameters, MctsBatch wraps a slot-map of Engine<N> instances:

use simple_mcts::{MctsBatch, MctsConfig, MctsStateEvaluation, StateEvaluation};
 
let config = MctsConfig::new(1.414, puct, negate_score);
let mut batch = MctsBatch::<_, _, N>::from_config(config);
 
let ids = batch.populate(256); // 256 independent search trees
 
// Selection phase across every active engine at once.
let selections = batch.selection();
 
// ... evaluate each resulting state with your own evaluator ...
 
let mut evaluations: Vec<MctsStateEvaluation<N>> = /* build from `selections` */;
batch.update(&mut evaluations).unwrap();
 
let scores = batch.scores(); // Vec<MctsScore<N>>, one per active engine

MctsBatch also exposes remove_one / remove to free engines (their slots are recycled via a free-list), commit_action / commit_actions to re-root individual trees, and update_one_with_transformer as an escape hatch when backpropagation logic needs per-call context (e.g. asymmetric or multi-player scoring).

Utility functions

puct(score, node_visits, parent_visits, policy, c) -> f32 — the PUCT selection formula (Q + c·P·√N / (1+n)), robust to unvisited nodes.
negate_score(s) -> f32 — flips the sign of a score; the standard transformer for alternating two-player games.
visits_to_probabilities(visits) -> [f32; N] — normalizes raw visit counts into a probability distribution.
visits_to_probabilities_with_temperature(visits, tau) -> [f32; N] — AlphaZero-style temperature-scaled policy extraction; tau <= 0.0 performs greedy selection (splitting probability evenly across ties).

When not to use this crate

simple-mcts only provides the search skeleton: tree management, selection, and backpropagation. It does not implement any game logic, evaluator, or neural network inference — you're expected to bring those yourself, as shown in the examples.

License

Licensed under either of MIT license or Apache License, Version 2.0, at your option.

simple-mcts 0.2.0