Generic Monte Carlo Tree Search implementation in Rust

Question

I'm working on a project where I need a strong AI for a combinatoric game, and decided to go with Monte Carlo Tree Search because the specific game will be original to the project - no good heuristics exist, so MCTS is more appropriate than something like minimax. Originally I wrote the AI in C# because that's what the rest of the project is written in, but the performance was very bad, so I thought I'd try my hand at Rust.

Since this is my first serious project in Rust, I'm mostly looking for tips on making the code more idiomatic, but any suggestions/improvements are welcome.

My primary source for implementing the algorithm was Browne et al.'s 2012 survey on the topic. While implementing it, I followed the pseudocode on page 9 fairly closely, but made the following changes:

• The AI has the capability to remember previous calculations (see the Searcher struct's previous_choice field and starting_tree method)
• When expanding nodes, the AI chooses from untried actions using the game's default policy, rather than choosing uniformly randomly

Finally, note that the plan is to expose this functionality as a cdylib, but for the purposes of sharing a minimum reproduceable example, I've converted it to a regular binary that computes the best starting move for a mock game. Also tests have been removed for brevity.

---

Project structure

.
├── Cargo.lock
├── Cargo.toml
└── src/
    ├── main.rs
    ├── game_state.rs
    └── mcts/
        ├── mod.rs
        ├── parameters.rs
        └── search.rs

search.rs contains the actual algorithm. game_state.rs contains the trait used by the search code that defines the game's behavior, while parameters.rs is just a small struct for global parameters to pass to the search algorithm.

Code

Cargo.toml

[package]
name = "mcts"
version = "0.1.0"
edition = "2021"

[dependencies]
indextree = "4.5.0"
rand = "0.8.5"

main.rs

use std::f32::consts::FRAC_1_SQRT_2;

mod game_state;
mod mcts;

use game_state::GameState;
use mcts::{SearchParameters, Searcher};

#[derive(PartialEq)]
pub struct MockGameState(i32);

impl GameState for MockGameState {
    type Move = i32;
    type Player = bool;
    type MoveIterator = std::vec::IntoIter<Self::Move>;

    fn initial_state() -> Self {
        Self(0)
    }

    fn available_moves(&self) -> Self::MoveIterator {
        vec![1, 3].into_iter()
    }

    fn next_to_play(&self) -> Self::Player {
        // `true` plays on even numbers, `false` plays on odd
        self.0 % 2 == 0
    }

    fn apply_move(&self, move_: Self::Move) -> Self {
        Self(self.0 + move_)
    }

    fn move_with_result(&self, result: &Self) -> Self::Move {
        result.0 - self.0
    }

    fn terminal_value(&self, for_player: Self::Player) -> Option<f32> {
        let score = if for_player == self.next_to_play() {
            self.0 as f32
        } else {
            -self.0 as f32
        };

        if score.abs() >= 10. {
            Some(score)
        } else {
            None
        }
    }
}

pub fn main() {
    let mut searcher: Searcher<MockGameState> = Searcher::new(SearchParameters {
        exploration_factor: FRAC_1_SQRT_2,
        search_iterations: 1000,
    });

    let move_ = searcher.search(MockGameState::initial_state());
    println!("{move_}");
}

game_state.rs

use rand::{seq::IteratorRandom, thread_rng};

pub trait GameState: PartialEq + Sized {
    type Move;
    type Player: Copy;
    type MoveIterator: Iterator<Item = Self::Move>;

    /// The state at the start of a game.
    fn initial_state() -> Self;

    /// The moves that are legal for the current player to take.
    fn available_moves(&self) -> Self::MoveIterator;

    /// In more words, the player whose turn it is.
    fn next_to_play(&self) -> Self::Player;

    /// Moves passed into this are guaranteed to be legal for this state.
    fn apply_move(&self, move_: Self::Move) -> Self;

    /// Result state is guaranteed to be legal and reachable from this state.
    fn move_with_result(&self, result: &Self) -> Self::Move;

    /// Used in simulations and expansion to choose either promising or pseudo-random
    /// moves.
    fn default_policy(&self, moves: &mut impl Iterator<Item = Self::Move>) -> Option<Self> {
        moves
            .choose(&mut thread_rng())
            .map(|move_| self.apply_move(move_))
    }

    /// Returns [None] if this state is non-terminal (ie, the game is still going on),
    /// and [Some(value)] if the game has finished and this state is worth `value` to the
    /// given player.
    fn terminal_value(&self, for_player: Self::Player) -> Option<f32>;
}

mcts/mod.rs

mod parameters;
mod search;

pub use parameters::*;
pub use search::*;

mcts/parameters.rs

pub struct SearchParameters {
    pub exploration_factor: f32,
    pub search_iterations: i32,
}

mcts/search.rs

//! Implementation of the MCTS algorithm as described by Browne et al 2012

use indextree::{Arena, NodeId};

use super::SearchParameters;
use crate::game_state::GameState;

pub struct Searcher<T>
where
    T: GameState,
{
    arena: Arena<MctsNode<T>>,
    previous_choice: Option<NodeId>,
    parameters: SearchParameters,
}

#[derive(PartialEq, Debug)]
struct MctsNode<T>
where
    T: GameState,
{
    game_state: T,
    score: f32,
    visits: i32,
    unexpanded_moves: T::MoveIterator,
}

impl<T> Searcher<T>
where
    T: GameState,
{
    pub fn new(parameters: SearchParameters) -> Self {
        Searcher {
            arena: Arena::new(),
            previous_choice: None,
            parameters,
        }
    }

    pub fn search(&mut self, starting_state: T) -> T::Move {
        let player = starting_state.next_to_play();
        let root = self.starting_tree(starting_state);

        for _ in 0..self.parameters.search_iterations {
            let leaf = self.tree_policy(root);
            let leaf_state = &self.node(leaf).game_state;
            let score = Self::rollout(leaf_state, player);
            self.backup_negamax(leaf, score);
        }

        let max_child = self.best_child(root, 0.);
        self.previous_choice = Some(max_child);

        let chosen_state = &self.node(max_child).game_state;
        self.node(root).game_state.move_with_result(chosen_state)
    }

    fn node(&self, id: NodeId) -> &MctsNode<T> {
        self.arena.get(id).unwrap().get()
    }

    fn node_mut(&mut self, id: NodeId) -> &mut MctsNode<T> {
        self.arena.get_mut(id).unwrap().get_mut()
    }

    fn starting_tree(&mut self, starting_state: T) -> NodeId {
        if self.previous_choice.is_none() {
            return self
                .arena
                .new_node(MctsNode::empty_from_state(starting_state));
        }

        let old_root = self.previous_choice.unwrap();

        let child_move = old_root
            .children(&self.arena)
            .find(|id| self.node(*id).game_state == starting_state);

        if let Some(new_root) = child_move {
            new_root.detach(&mut self.arena);
            old_root.remove_subtree(&mut self.arena);

            new_root
        } else {
            self.arena.clear();

            self.arena
                .new_node(MctsNode::empty_from_state(starting_state))
        }
    }

    fn rollout(initial_state: &T, for_player: T::Player) -> f32 {
        // wanted to do this iteratively, but was fighting the borrow checker. hopefully
        // we'll see some tail call optimization
        if let Some(val) = initial_state.terminal_value(for_player) {
            val
        } else {
            Self::rollout(
                &initial_state
                    .default_policy(&mut initial_state.available_moves())
                    .unwrap(),
                for_player,
            )
        }
    }

    fn tree_policy(&mut self, node_id: NodeId) -> NodeId {
        let mut parent = node_id;
        let mut leaf = self.expand(parent);

        while leaf.is_none() {
            parent = self.best_child(parent, self.parameters.exploration_factor);
            leaf = self.expand(parent);
        }

        leaf.unwrap()
    }

    fn expand(&mut self, node_id: NodeId) -> Option<NodeId> {
        let node = self.node_mut(node_id);

        node.game_state
            .default_policy(&mut node.unexpanded_moves)
            .map(|game_state| {
                let child = self.arena.new_node(MctsNode::empty_from_state(game_state));
                node_id.append(child, &mut self.arena);

                child
            })
    }

    /// exploration_factor is also known as c (Browne p. 9)
    fn best_child(&self, parent: NodeId, exploration_factor: f32) -> NodeId {
        let ucb1 = |id: &NodeId| {
            let parent = self.node(parent);
            let child = self.node(*id);

            let exploitation_term = child.score / child.visits_f();
            let exploration_term = (2. * parent.visits_f().ln() / child.visits_f()).sqrt();

            exploitation_term + exploration_factor * exploration_term
        };

        parent
            .children(&self.arena)
            .max_by(|a, b| {
                let a_val = ucb1(a);
                let b_val = ucb1(b);

                a_val
                    .partial_cmp(&b_val)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .unwrap()
    }

    fn backup_negamax(&mut self, node_id: NodeId, mut score: f32) {
        let mut next = self.arena.get_mut(node_id);

        while let Some(node) = next {
            let data = node.get_mut();
            data.score += score;
            data.visits += 1;

            score = -score;
            next = node.parent().and_then(|id| self.arena.get_mut(id));
        }
    }
}

impl<T> MctsNode<T>
where
    T: GameState,
{
    pub fn empty_from_state(game_state: T) -> Self {
        MctsNode {
            unexpanded_moves: game_state.available_moves(),
            game_state,
            score: 0.,
            visits: 0,
        }
    }

    #[inline]
    pub fn visits_f(&self) -> f32 {
        self.visits as f32
    }
}
```