Working MCTS implementation

This is a basic working implementation of the MCTS algorithm. Though
currently the algorithm is slow compared with other implemenations, and
makes sub-optimal choices when playing tic-tac-toe. Therefore some
modifications are needed

src/config.rs

@@ -0,0 +1,67 @@
+use crate::policy::backprop::BackpropagationPolicy;
+use crate::policy::decision::DecisionPolicy;
+use crate::policy::selection::SelectionPolicy;
+use crate::policy::simulation::SimulationPolicy;
+use crate::state::GameState;
+use std::time::Duration;
+
+/// Configuration for the MCTS algorithm
+#[derive(Debug)]
+pub struct MCTSConfig<S: GameState> {
+    /// The maximum number of iterations to run when searching
+    ///
+    /// The search will stop after the given number of iterations, even if there
+    /// is search time has not exceeded `max_time`.
+    pub max_iterations: usize,
+
+    /// The maximum time to run the search
+    ///
+    /// If set, the search will stop after this duration even if the maximum
+    /// iterations hasn't been reached.
+    pub max_time: Option<Duration>,
+
+    /// The size to initially allocate for the search tree
+    ///
+    /// This pre-allocates memory for the search tree which ensures contiguous
+    /// memory and improves performance by preventing the resizing of tree
+    /// as we explore.
+    pub tree_size_allocation: usize,
+
+    /// The selection policy
+    ///
+    /// This dictates the path through which the game tree is searched. As such
+    /// the policy has a large impact on the overall aglorthm exeuction
+    pub selection_policy: SelectionPolicy<S>,
+
+    /// The simulation policy
+    ///
+    /// This dictates the game siluation when expanding and evaluating the
+    /// search tree. Random is generally a good default.
+    pub simulation_policy: SimulationPolicy<S>,
+
+    /// The backpropagation policy
+    ///
+    /// This dictates how the results of the simulation playouts are propagated
+    /// back up the tree.
+    pub backprop_policy: BackpropagationPolicy<S>,
+
+    /// The decision policy
+    ///
+    /// This dictates how the MCTS algorithm determines its final decision
+    /// after iterating through the search tree
+    pub decision_policy: DecisionPolicy,
+}
+
+impl<S: GameState> Default for MCTSConfig<S> {
+    fn default() -> Self {
+        MCTSConfig {
+            max_iterations: 10_000,
+            max_time: None,
+            tree_size_allocation: 10_000,
+            selection_policy: SelectionPolicy::UCB1Tuned(1.414),
+            simulation_policy: SimulationPolicy::Random,
+            backprop_policy: BackpropagationPolicy::Standard,
+            decision_policy: DecisionPolicy::MostVisits,
+        }
+    }
+}

src/lib.rs

@@ -0,0 +1,17 @@
+//! # rustic_mcts
+//!
+//! An extensible implementation of Monte Carlo Tree Search (MCTS) using arena allocation and
+//! configurable policies.
+
+pub mod config;
+pub mod mcts;
+pub mod policy;
+pub mod state;
+pub mod tree;
+
+pub use config::MCTSConfig;
+pub use mcts::MCTS;
+pub use state::Action;
+pub use state::GameState;
+pub use state::Player;
+pub use tree::node::RewardVal;

src/mcts.rs

@@ -0,0 +1,147 @@
+use crate::config::MCTSConfig;
+use crate::policy::backprop::backpropagate_rewards;
+use crate::policy::decision::decide_on_action;
+use crate::policy::selection::select_best_child;
+use crate::policy::simulation::simulate_reward;
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::{Node, RewardVal};
+use rand::prelude::SliceRandom;
+use std::collections::HashMap;
+use std::time::Instant;
+
+/// Monte Carlo Tree Search implementation
+///
+/// This provides the interface for performing optimal searches on a tree using
+/// the MCTS algorithm.
+pub struct MCTS<'conf, S: GameState> {
+    /// The arena used for the tree
+    arena: Arena<S>,
+
+    /// The identifier of the root node of the search tree
+    root_id: usize,
+
+    /// The configuration used for the search
+    config: &'conf MCTSConfig<S>,
+}
+
+impl<'conf, S: GameState + std::fmt::Debug> MCTS<'conf, S> {
+    /// Creates a new instance with the given initial state and configuration
+    pub fn new(initial_state: S, config: &'conf MCTSConfig<S>) -> Self {
+        let mut arena: Arena<S> = Arena::new(config.tree_size_allocation);
+        let root: Node<S> = Node::new(initial_state.clone(), None, None);
+        let root_id: usize = arena.add_node(root);
+        MCTS {
+            arena,
+            root_id,
+            config,
+        }
+    }
+
+    /// Runs the MCTS algorithm, returning the "best" action
+    ///
+    /// The search will stop once `max_iterations` or `max_time` from
+    /// the assigned configration is reached.
+    pub fn search(&mut self) -> Result<S::Action> {
+        self.search_for_iterations(self.config.max_iterations)
+    }
+
+    /// Runs the MCTS algorithm, returning the "best" action after the given iterations
+    ///
+    /// This ignores the `max_iterations` provided in the config, however will
+    /// return if `max_time` is specific and reached before the iterations are complete.
+    pub fn search_for_iterations(&mut self, iterations: usize) -> Result<S::Action> {
+        let start_time = Instant::now();
+        for _ in 0..iterations {
+            match self.config.max_time {
+                Some(max_time) => {
+                    if start_time.elapsed() >= max_time {
+                        break; // ending early due to time
+                    }
+                }
+                None => {}
+            }
+            self.execute_iteration()?;
+        }
+
+        self.best_action()
+    }
+
+    /// Runs the MCTS algorithm for a single iteration
+    fn execute_iteration(&mut self) -> Result<()> {
+        let mut selected_id: usize = self.select();
+        let selected_node: &Node<S> = self.arena.get_node(selected_id);
+        if !selected_node.state.is_terminal() {
+            self.expand(selected_id);
+            let children: &Vec<usize> = &self.arena.get_node(selected_id).children;
+            let random_child: usize = *children.choose(&mut rand::thread_rng()).unwrap();
+            selected_id = random_child;
+        }
+        let rewards = self.simulate(selected_id);
+        self.backprop(selected_id, &rewards);
+
+        Ok(())
+    }
+
+    /// MCTS Phase 1: Selection - Find the "best" node to expand
+    fn select(&mut self) -> usize {
+        let mut current_id: usize = self.root_id;
+        loop {
+            let node = &self.arena.get_node(current_id);
+            if node.is_leaf() || node.state.is_terminal() {
+                return current_id;
+            }
+            current_id = select_best_child(&self.config.selection_policy, &node, &self.arena);
+        }
+    }
+
+    /// MCTS Phase 2: Expansion - Expand the selected node on the tree
+    fn expand(&mut self, id: usize) {
+        let parent: &Node<S> = self.arena.get_node_mut(id);
+        let legal_actions: Vec<S::Action> = parent.state.get_legal_actions();
+        let parent_state: S = parent.state.clone();
+        for action in legal_actions {
+            let state = parent_state.state_after_action(&action);
+            let new_node = Node::new(state, Some(action), Some(id));
+            let new_id = self.arena.add_node(new_node);
+            self.arena.get_node_mut(id).children.push(new_id);
+        }
+    }
+
+    fn simulate(&self, id: usize) -> HashMap<S::Player, RewardVal> {
+        let node = &self.arena.get_node(id);
+        simulate_reward(&self.config.simulation_policy, &node, &self.arena)
+    }
+
+    fn backprop(&mut self, selected_id: usize, rewards: &HashMap<S::Player, RewardVal>) {
+        backpropagate_rewards(
+            &self.config.backprop_policy,
+            selected_id,
+            &mut self.arena,
+            &rewards,
+        )
+    }
+
+    fn best_action(&self) -> Result<S::Action> {
+        let root_node: &Node<S> = self.arena.get_node(self.root_id);
+        match decide_on_action(&self.config.decision_policy, &root_node, &self.arena) {
+            Some(action) => Ok(action),
+            None => Err(MCTSError::NoBestAction),
+        }
+    }
+}
+
+/// Errors returned by the MCTS algorithm
+#[derive(Debug, thiserror::Error)]
+pub enum MCTSError {
+    /// The best action doesn't exist
+    #[error("Unable to determine a best action for the game")]
+    NoBestAction,
+
+    /// The search tree was exhausted without finding a terminal node
+    #[error("Search tree exhausted without finding terminal node")]
+    NonTerminalGame,
+}
+
+/// Result returned by the MCTS algorithm
+pub type Result<T> = std::result::Result<T, MCTSError>;

src/policy/backprop/mod.rs

@@ -0,0 +1,99 @@
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::RewardVal;
+use std::collections::HashMap;
+
+/// The back propagation policy dictating the propagation of playout results
+///
+/// This policy drives how the backpropagation phase of the MCTS algorithm is
+/// executed, allowing for some minor customization.
+///
+/// Typically the Standard policy, used by most implementaions of MCTS, is
+/// sufficient
+#[derive(Debug)]
+pub enum BackpropagationPolicy<S: GameState> {
+    /// Standard back propagation
+    ///
+    /// This increments the visitation count and adds the simulated rewards
+    /// results to the aggregate values.
+    ///
+    /// This is the standard policy used in most MCTS implementations.
+    Standard,
+
+    /// Weighted back propagation
+    ///
+    /// This weights the value of the simulated rewards based on the depth,
+    /// allowing us to put more-or-less influence on deeper branches
+    /// - Positive weight factor makes deeper nodes less influential
+    /// - Negative weight factor makes deeper nodes more influential
+    Weighted(f64),
+
+    /// Custom backpropagation policy
+    Custom(Box<dyn CustomBackpropagationPolicy<S>>),
+}
+
+/// Trait for an object implementing the backpropagation logic whene exploring the MCTS
+/// search tree.
+pub trait CustomBackpropagationPolicy<S: GameState>: std::fmt::Debug {
+    /// Backpropagate the given rewards values from the node up the tree
+    fn backprop(
+        &self,
+        node_id: usize,
+        arena: &mut Arena<S>,
+        rewards: &HashMap<S::Player, RewardVal>,
+    );
+}
+
+pub fn backpropagate_rewards<S: GameState>(
+    policy: &BackpropagationPolicy<S>,
+    node_id: usize,
+    arena: &mut Arena<S>,
+    rewards: &HashMap<S::Player, RewardVal>,
+) {
+    match policy {
+        BackpropagationPolicy::Standard => standard_backprop(node_id, arena, rewards),
+        BackpropagationPolicy::Weighted(depth_factor) => {
+            weighted_backprop(*depth_factor, node_id, arena, rewards)
+        }
+        BackpropagationPolicy::Custom(custom_policy) => {
+            custom_policy.backprop(node_id, arena, rewards)
+        }
+    }
+}
+
+fn standard_backprop<S: GameState>(
+    node_id: usize,
+    arena: &mut Arena<S>,
+    rewards: &HashMap<S::Player, RewardVal>,
+) {
+    // TODO:
+    // - each node needs the perspective of the different players not just one view
+    //   - e.g. reward_sum(player), reward_avg(player), rewards(player)[], visits(player)
+    // - we could make special version for 2-player zero-sum games like below
+    let mut current_id: usize = node_id;
+    loop {
+        let node = arena.get_node_mut(current_id);
+        let player = node.state.get_current_player().clone();
+        match rewards.get(&player) {
+            Some(reward) => {
+                node.increment_visits();
+                node.record_player_reward(player, *reward);
+                if let Some(parent_id) = node.parent {
+                    current_id = parent_id;
+                } else {
+                    break;
+                }
+            }
+            None => (),
+        }
+    }
+}
+
+fn weighted_backprop<S: GameState>(
+    _depth_factor: f64,
+    _node_id: usize,
+    _arena: &mut Arena<S>,
+    _rewards: &HashMap<S::Player, RewardVal>,
+) {
+    // TODO
+}

src/policy/decision/mod.rs

@@ -0,0 +1,65 @@
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::Node;
+
+/// The decision policy when determining the action in final MCTS phase
+///
+/// This policy drives how the MCTS algorithm chooses which action is the
+/// "best" from the exploration.
+#[derive(Debug)]
+pub enum DecisionPolicy {
+    /// Decide on the action with the most visits
+    ///
+    /// This option relies on the statistical confidence drive by the MCTS
+    /// algorithm instead of the potentially more noisy value estimates.
+    ///
+    /// This is the standard policy used in most MCTS implementations, and
+    /// is a good selection when not hyper-maximizing for potential gain
+    MostVisits,
+
+    /// Decide on the action with the highest average value
+    ///
+    /// This is non-standard, but is more aggressive in attempting to gain
+    /// the highest value in a decision.
+    HighestValue,
+}
+
+pub fn decide_on_action<S: GameState>(
+    policy: &DecisionPolicy,
+    root_node: &Node<S>,
+    arena: &Arena<S>,
+) -> Option<S::Action> {
+    match policy {
+        DecisionPolicy::MostVisits => most_visits(root_node, arena),
+        DecisionPolicy::HighestValue => highest_value(root_node, arena),
+    }
+}
+
+fn most_visits<S: GameState>(root_node: &Node<S>, arena: &Arena<S>) -> Option<S::Action> {
+    let best_child_id: &usize = root_node
+        .children
+        .iter()
+        .max_by(|&a, &b| {
+            let node_a_visits = arena.get_node(*a).visits;
+            let node_b_visits = arena.get_node(*b).visits;
+            node_a_visits.partial_cmp(&node_b_visits).unwrap()
+        })
+        .unwrap();
+
+    arena.get_node(*best_child_id).action.clone()
+}
+
+fn highest_value<S: GameState>(root_node: &Node<S>, arena: &Arena<S>) -> Option<S::Action> {
+    let player = root_node.state.get_current_player();
+    let best_child_id: &usize = root_node
+        .children
+        .iter()
+        .max_by(|&a, &b| {
+            let node_a_score = arena.get_node(*a).reward_average(player);
+            let node_b_score = arena.get_node(*b).reward_average(player);
+            node_a_score.partial_cmp(&node_b_score).unwrap()
+        })
+        .unwrap();
+
+    arena.get_node(*best_child_id).action.clone()
+}

src/policy/mod.rs

@@ -0,0 +1,4 @@
+pub mod backprop;
+pub mod decision;
+pub mod selection;
+pub mod simulation;

src/policy/selection/mod.rs

@@ -0,0 +1,59 @@
+mod ucb1;
+mod ucb1_tuned;
+
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::Node;
+
+/// The selection policy used in the MCTS selection phase
+///
+/// This drives the selection of the nodes in the search tree, determining
+/// which paths are explored and evaluated.
+///
+/// In general UCB1-Tuned or UCB1 should be effective, however if necessariy
+/// a custom selection policy can be provided.
+#[derive(Debug)]
+pub enum SelectionPolicy<S: GameState> {
+    /// Upper Confidence Bound 1 (UCB1) with the given exploration constant
+    ///
+    /// The exploration constant controls the balance between exploration and
+    /// exploitation. The higher the value, the mroe likely the search will
+    /// explore less-visited nodes. A standard value is √2 ≈ 1.414.
+    UCB1(f64),
+
+    /// Upper Confidence Bound 1 Tuned (UCB1-Tuned)
+    ///
+    /// A tuned version of UCB1 instead using the empirical
+    /// standard deviation of the rewards to drive exploration.
+    ///
+    /// Auer, P., Cesa-Bianchi, N. & Fischer, P. Finite-time Analysis of the Multiarmed Bandit Problem. Machine Learning 47, 235–256 (2002). https://doi.org/10.1023/A:1013689704352
+    UCB1Tuned(f64),
+
+    /// Custom selection policy
+    Custom(Box<dyn CustomSelectionPolicy<S>>),
+}
+
+/// Trait for an object implementing the selection logic whene exploring the MCTS
+/// search tree.
+///
+/// The policy should select the child of the given node which is "best" for the current player
+pub trait CustomSelectionPolicy<S: GameState>: std::fmt::Debug {
+    /// Selects a child based on the policy, returning the node ID
+    fn select_child(&self, node: &Node<S>, arena: &Arena<S>) -> usize;
+}
+
+pub fn select_best_child<S: GameState>(
+    policy: &SelectionPolicy<S>,
+    node: &Node<S>,
+    arena: &Arena<S>,
+) -> usize {
+    match policy {
+        SelectionPolicy::UCB1(exploration_constant) => {
+            ucb1::select_best_child(*exploration_constant, node, arena)
+        }
+        SelectionPolicy::UCB1Tuned(exploration_constant) => {
+            ucb1_tuned::select_best_child(*exploration_constant, node, arena)
+        }
+        SelectionPolicy::Custom(custom_policy) => custom_policy.select_child(node, arena),
+    }
+}

src/policy/selection/ucb1.rs

@@ -0,0 +1,79 @@
+//! Upper Confidence Bound 1 (UCB1) selection policy
+//!
+//! This is the classic selection policy for MCTS, which balances
+//! exploration and exploitation using the UCB1 formula:
+//!
+//! ```text
+//! UCB1 = average_reward + exploration_constant * sqrt(ln(parent_visits) / child_visits)
+//! ```
+//!
+//! Where:
+//! - `average_reward` is the average reward from simulations through this node
+//! - `exploration_constant` controls the balance between exploration and exploitation
+//! - `parent_visits` is the number of visits to the parent node
+//! - `child_visits` is the number of visits to the child node
+//!
+//! Higher exploration constants favor exploration (trying less-visited nodes),
+//! while lower values favor exploitation (choosing nodes with higher values).
+//!
+//! The commonly used value for the exploration constant is sqrt(2) ≈ 1.414,
+//! which is the default in this implementation.
+
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::{Node, RewardVal};
+
+/// Selects the index of the "best" child using the UCB1 selection policy
+pub fn select_best_child<S: GameState>(
+    exploration_constant: f64,
+    node: &Node<S>,
+    arena: &Arena<S>,
+) -> usize {
+    if node.is_leaf() {
+        panic!("select_best_child called on leaf node");
+    }
+
+    let player = node.state.get_current_player();
+    let parent_visits = node.visits;
+    let best_child = node
+        .children
+        .iter()
+        .max_by(|&a, &b| {
+            let node_a = arena.get_node(*a);
+            let node_b = arena.get_node(*b);
+            let ucb_a = ucb1_value(
+                exploration_constant,
+                node_a.reward_average(player),
+                node_a.visits,
+                parent_visits,
+            );
+            let ucb_b = ucb1_value(
+                exploration_constant,
+                node_b.reward_average(player),
+                node_b.visits,
+                parent_visits,
+            );
+            ucb_a.partial_cmp(&ucb_b).unwrap()
+        })
+        .unwrap();
+    *best_child
+}
+
+/// Calculates the UCB1 value for a node
+pub fn ucb1_value(
+    exploration_constant: f64,
+    child_value: RewardVal,
+    child_visits: u64,
+    parent_visits: u64,
+) -> RewardVal {
+    if child_visits == 0 {
+        return f64::INFINITY; // Always explore nodes that have never been visited
+    }
+
+    // UCB1 formula: value + C * sqrt(ln(parent_visits) / child_visits)
+    let exploitation = child_value;
+    let exploration =
+        exploration_constant * ((parent_visits as f64).ln() / child_visits as f64).sqrt();
+
+    exploitation + exploration
+}

src/policy/selection/ucb1_tuned.rs

@@ -0,0 +1,97 @@
+//! Upper Confidence Bound 1 Tuned (UCB1-Tuned) selection policy
+//!
+//! This is a fine-tuned version of UCB which takes into account the
+//! empircally measured variance of the rewards to drive the exploration.
+//!
+//! This has been found to perform substantially better than UCB1 in most
+//! situations.
+//!
+//! Auer, P., Cesa-Bianchi, N. & Fischer, P.
+//!  Finite-time Analysis of the Multiarmed Bandit Problem.
+//!  Machine Learning 47, 235–256 (2002). https://doi.org/10.1023/A:1013689704352
+
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::{Node, RewardVal};
+
+/// Selects the index of the "best" child using the UCB1-Tuned selection policy
+pub fn select_best_child<S: GameState>(
+    exploration_constant: f64,
+    node: &Node<S>,
+    arena: &Arena<S>,
+) -> usize {
+    if node.is_leaf() {
+        panic!("select_best_child called on leaf node");
+    }
+
+    let player = node.state.get_current_player();
+    let parent_visits = node.visits;
+    let best_child = node
+        .children
+        .iter()
+        .max_by(|&a, &b| {
+            let node_a = arena.get_node(*a);
+            let node_b = arena.get_node(*b);
+            let ucb_a = ucb1_tuned_value(
+                exploration_constant,
+                parent_visits,
+                node_a.visits,
+                node_a.rewards(player),
+                node_a.reward_average(player),
+            );
+            let ucb_b = ucb1_tuned_value(
+                exploration_constant,
+                parent_visits,
+                node_b.visits,
+                node_b.rewards(player),
+                node_b.reward_average(player),
+            );
+            ucb_a.partial_cmp(&ucb_b).unwrap()
+        })
+        .unwrap();
+    *best_child
+}
+
+/// Calculates the UCB1-Tuned value for a node
+pub fn ucb1_tuned_value(
+    exploration_constant: f64,
+    parent_visits: u64,
+    child_visits: u64,
+    child_rewards: Option<&Vec<RewardVal>>,
+    reward_avg: RewardVal,
+) -> RewardVal {
+    match child_rewards {
+        None => {
+            RewardVal::INFINITY // Always explore nodes that have never been visited
+        }
+        Some(child_rewards) => {
+            if child_visits == 0 {
+                RewardVal::INFINITY // Always explore nodes that have never been visited
+            } else {
+                let parent_visits: RewardVal = parent_visits as RewardVal;
+                let child_visits: RewardVal = child_visits as RewardVal;
+
+                // N: number of visits to the parent node
+                // n: number of visits to the child node
+                // x_i: reward of the ith visit to the child node
+                // X: average reward of the child
+                // C: exploration constant
+                //
+                // UCB1-Tuned = X + C * sqrt(Ln(parent_visits) / child_visits * min(1/4, V_n)
+                //   V(n) = sum(x_i^2)/n - X^2 + sqrt(2*ln(N)/n)
+                let exploitation = reward_avg;
+                let mut variance = (child_rewards.iter().map(|&x| x * x).sum::<RewardVal>()
+                    / child_visits)
+                    - (reward_avg * reward_avg)
+                    + (2.0 * parent_visits.ln() / child_visits).sqrt();
+                if variance > 0.25 {
+                    variance = 0.25;
+                }
+                let exploration =
+                    exploration_constant * (parent_visits.ln() / child_visits * variance).sqrt();
+
+                exploitation + exploration
+            }
+        }
+    }
+}

src/policy/simulation/mod.rs

@@ -0,0 +1,44 @@
+mod random;
+
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::{Node, RewardVal};
+use std::collections::HashMap;
+
+/// The simulation policy used in the MCTS simulation phase
+///
+/// This policy drives the game simulations while evaluating the tree. While
+/// a random policy works well, a game-specific policy can be provided either
+/// as a custom policy.
+#[derive(Debug)]
+pub enum SimulationPolicy<S: GameState> {
+    /// Random simulation policy
+    ///
+    /// The sequential actions are selected randomly from the available actions
+    /// at each state until a terminal state is found.
+    Random,
+
+    /// Custom simulation policy
+    Custom(Box<dyn CustomSimulationPolicy<S>>),
+}
+
+/// Trait for an object implementing the simulation logic whene exploring the MCTS
+/// search tree.
+pub trait CustomSimulationPolicy<S: GameState>: std::fmt::Debug {
+    /// Simulates the gameplay from the current node onward, returning the rewards
+    ///
+    /// This should simulate the game until a terminal node is reached, returning
+    /// the final reward for each player at the terminal node
+    fn simulate(&self, node: &Node<S>, arena: &Arena<S>) -> HashMap<S::Player, RewardVal>;
+}
+
+pub fn simulate_reward<S: GameState>(
+    policy: &SimulationPolicy<S>,
+    node: &Node<S>,
+    arena: &Arena<S>,
+) -> HashMap<S::Player, RewardVal> {
+    match policy {
+        SimulationPolicy::Random => random::simulate(node),
+        SimulationPolicy::Custom(custom_policy) => custom_policy.simulate(node, arena),
+    }
+}

src/policy/simulation/random.rs

@@ -0,0 +1,18 @@
+//! Random play simulation policy
+//!
+//! Actions are chosen at random
+
+use crate::state::GameState;
+use crate::tree::node::{Node, RewardVal};
+use rand::prelude::SliceRandom;
+use std::collections::HashMap;
+
+pub fn simulate<S: GameState>(node: &Node<S>) -> HashMap<S::Player, RewardVal> {
+    let mut state: S = node.state.clone();
+    while !state.is_terminal() {
+        let legal_actions = state.get_legal_actions();
+        let action = legal_actions.choose(&mut rand::thread_rng()).unwrap();
+        state = state.state_after_action(&action);
+    }
+    state.rewards_for_players()
+}

src/state.rs

@@ -0,0 +1,90 @@
+use crate::tree::node::RewardVal;
+use std::collections::HashMap;
+use std::fmt::Debug;
+use std::hash::Hash;
+
+/// Trait for the game state used in MCTS
+///
+/// When leveraging MCTS for your game, you must implement this trait to provide
+/// the specifics for your game.
+pub trait GameState: Clone {
+    /// The type of actions that can be taken in the game
+    type Action: Action;
+
+    /// The type of players in the game
+    type Player: Player;
+
+    /// Returns if the game state is terminal, i.e. the game is over
+    ///
+    /// A game state is terminal when no other actions are possible. This can be
+    /// the result of a player winning, a draw, or because some other conditions
+    /// have been met leading to a game with no further possible states.
+    ///
+    /// The default implementation returns True if `get_legal_actions()` returns
+    /// an empty list. It is recommended to override this for a more efficient
+    /// implementation if possible.
+    fn is_terminal(&self) -> bool {
+        let actions = self.get_legal_actions();
+        actions.len() == 0
+    }
+
+    /// Returns the list of legal actions for the game state
+    ///
+    /// This method must return all possible actions that can be made from the
+    /// current game state.
+    fn get_legal_actions(&self) -> Vec<Self::Action>;
+
+    /// Returns the game state resulting from applying the action to the state
+    ///
+    /// This function should not modify the current state directly, and
+    /// instead should modify a copy of the state and return that.
+    fn state_after_action(&self, action: &Self::Action) -> Self;
+
+    /// Returns the reward from the perspective of the given player for the game state
+    ///
+    /// This evaluates the current state from the perspective of the given player, and
+    /// returns the reward indicating how good of a result the given state is for the
+    /// player.
+    ///
+    /// This is used in the MCTS backpropagation and simulation phases to evaluate
+    /// the value of a given node in the search tree.
+    ///
+    /// A general rule of thumb for values are:
+    /// - 1.0 => a win for the player
+    /// - 0.5 => a draw
+    /// - 0.0 => a loss for the player
+    ///
+    /// Other values can be used for relative wins or losses
+    fn reward_for_player(&self, player: &Self::Player) -> RewardVal;
+
+    /// Returns the rewards for all players at the current state
+    fn rewards_for_players(&self) -> HashMap<Self::Player, RewardVal>;
+
+    /// Returns the player whose turn it is for the game state
+    ///
+    /// This is used for evaluating the state, so for simultaneous games
+    /// consider the "current player" as the one from whose perspective we are
+    /// evaluating the game state from
+    fn get_current_player(&self) -> &Self::Player;
+}
+
+/// Trait used for actions that can be taken in a game
+///
+/// An action is dependent upon the specific game being defined, and includes
+/// things like moves, attacks, and other decisions.
+pub trait Action: Clone + Debug {
+    /// Returns a uniqie identifier for this action
+    fn id(&self) -> usize;
+}
+
+/// Trait used for players participating in a game
+pub trait Player: Clone + Debug + PartialEq + Eq + Hash {}
+
+/// Convenience implemnentation of a Player for usize
+impl Player for usize {}
+
+/// Convenience implemnentation of a Player for char
+impl Player for char {}
+
+/// Convenience implemnentation of a Player for String
+impl Player for String {}

src/tree/arena.rs

@@ -0,0 +1,46 @@
+use crate::state::GameState;
+use crate::tree::node::Node;
+
+/// An arena for Node allocation
+///
+/// We use an arena for node allocation to improve performance of our search.
+/// The memory is contiguous which allows for faster movement through the tree,
+/// as well as more efficient destruction as our MCTS search will destroy the
+/// entire tree at once.
+pub struct Arena<S: GameState> {
+    pub nodes: Vec<Node<S>>,
+}
+
+impl<S: GameState> Arena<S> {
+    /// Create a new Arena with the given initial capacity
+    ///
+    /// The arena creates a contiguous block. By reserving an initial capacity
+    /// that is sufficient to encapsulate a full search tree we can reduce the
+    /// number of reallocs that are required. This number is highly game
+    /// dependent.
+    pub fn new(initial_capacity: usize) -> Self {
+        Arena {
+            nodes: Vec::with_capacity(initial_capacity),
+        }
+    }
+
+    /// Adds a node to the Arena, returning its identifier
+    ///
+    /// This appends the node to the allocated Arena, and returns the nodes
+    /// index in the arena which is used as an identifier for later retrieval.
+    pub fn add_node(&mut self, node: Node<S>) -> usize {
+        let id = self.nodes.len();
+        self.nodes.push(node);
+        id
+    }
+
+    /// Retrieves a mutable reference to a Node in the Arena
+    pub fn get_node_mut(&mut self, id: usize) -> &mut Node<S> {
+        &mut self.nodes[id]
+    }
+
+    /// Retrieves a reference to a Node in the Arena
+    pub fn get_node(&self, id: usize) -> &Node<S> {
+        &self.nodes[id]
+    }
+}

src/tree/mod.rs

@@ -0,0 +1,2 @@
+pub mod arena;
+pub mod node;

src/tree/node.rs

@@ -0,0 +1,114 @@
+use std::collections::HashMap;
+use std::fmt::Debug;
+
+use crate::state::GameState;
+
+/// The type used for reward values
+pub type RewardVal = f64;
+
+/// A node in the MCTS tree
+///
+/// A node represents a given game state and, using the path from the root node,
+/// the actions that led to the given state. A node has a number of children
+/// nodes representing the game states reachable from the given state, after
+/// a given action. This creates the tree that MCTS iterates through.
+///
+/// This class is not thread safe, as the library does not provide for parallel
+/// search.
+#[derive(Debug)]
+pub struct Node<S: GameState> {
+    /// The game state at the given node, after `action`
+    pub state: S,
+
+    /// The action that led to this state from its parent
+    pub action: Option<S::Action>,
+
+    /// The identifier of the parent Node
+    pub parent: Option<usize>,
+
+    /// The number of times this node has been visited
+    pub visits: u64,
+
+    /// The player's evaluation of the node
+    pub player_view: HashMap<S::Player, PlayerNodeView>,
+
+    /// The identifiers of children nodes, states reachable from this one
+    pub children: Vec<usize>,
+}
+
+impl<S: GameState> Node<S> {
+    pub fn new(state: S, action: Option<S::Action>, parent: Option<usize>) -> Self {
+        Node {
+            state,
+            action,
+            parent,
+            visits: 0,
+            player_view: HashMap::with_capacity(2),
+            children: Vec::new(),
+        }
+    }
+
+    pub fn is_leaf(&self) -> bool {
+        self.children.is_empty()
+    }
+
+    pub fn reward_sum(&self, player: &S::Player) -> RewardVal {
+        match self.player_view.get(player) {
+            Some(pv) => pv.reward_sum,
+            None => 0.0,
+        }
+    }
+
+    pub fn reward_average(&self, player: &S::Player) -> RewardVal {
+        match self.player_view.get(player) {
+            Some(pv) => pv.reward_average,
+            None => 0.0,
+        }
+    }
+
+    pub fn rewards(&self, player: &S::Player) -> Option<&Vec<RewardVal>> {
+        match self.player_view.get(player) {
+            Some(pv) => Some(&pv.rewards),
+            None => None,
+        }
+    }
+
+    pub fn increment_visits(&mut self) {
+        self.visits += 1
+    }
+
+    pub fn record_player_reward(&mut self, player: S::Player, reward: RewardVal) {
+        let pv = self
+            .player_view
+            .entry(player)
+            .or_insert(PlayerNodeView::default());
+        pv.rewards.push(reward);
+        pv.reward_sum += reward;
+        pv.reward_average = pv.reward_sum / pv.rewards.len() as f64;
+    }
+}
+
+/// A player's specific perspective of a node's value
+///
+/// Each player has their own idea of the value of a node.
+#[derive(Debug)]
+pub struct PlayerNodeView {
+    /// The total reward from simulations through this node
+    pub reward_sum: RewardVal,
+
+    /// The average reward from simulations through this node, often called the node value
+    pub reward_average: RewardVal,
+
+    /// The rewards we have gotten so far for simulations through this node
+    pub rewards: Vec<RewardVal>,
+}
+
+impl Default for PlayerNodeView {
+    fn default() -> Self {
+        PlayerNodeView {
+            reward_sum: 0.0,
+            reward_average: 0.0,
+            rewards: Vec::new(),
+        }
+    }
+}