4
\$\begingroup\$

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
    }
}
```
\$\endgroup\$

0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Browse other questions tagged or ask your own question.