Add a "threads" feature for mutli-threading

We calculate the sample variance of the rewards online storing the value
in the node. This greatly reduces the amount of summations that need to
be done to calculate the variance during the selection phase.

While this burdens other selection algorithms, the cost is not
substantial.

Cargo.lock

@@ -2,6 +2,18 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "anstyle"
+version = "1.0.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "862ed96ca487e809f1c8e5a8447f6ee2cf102f846893800b20cebdf541fc6bbd"
+
+[[package]]
+name = "bitflags"
+version = "2.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967"
+
 [[package]]
 name = "cfg-if"
 version = "1.0.1"
@@ -9,13 +21,81 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268"
 
 [[package]]
-name = "getrandom"
+name = "clap"
-version = "0.2.16"
+version = "4.5.40"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
+checksum = "40b6887a1d8685cebccf115538db5c0efe625ccac9696ad45c409d96566e910f"
+dependencies = [
+ "clap_builder",
+]
+
+[[package]]
+name = "clap_builder"
+version = "4.5.40"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e0c66c08ce9f0c698cbce5c0279d0bb6ac936d8674174fe48f736533b964f59e"
+dependencies = [
+ "anstyle",
+ "clap_lex",
+ "terminal_size",
+]
+
+[[package]]
+name = "clap_lex"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675"
+
+[[package]]
+name = "condtype"
+version = "1.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "baf0a07a401f374238ab8e2f11a104d2851bf9ce711ec69804834de8af45c7af"
+
+[[package]]
+name = "divan"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a405457ec78b8fe08b0e32b4a3570ab5dff6dd16eb9e76a5ee0a9d9cbd898933"
+dependencies = [
+ "cfg-if",
+ "clap",
+ "condtype",
+ "divan-macros",
+ "libc",
+ "regex-lite",
+]
+
+[[package]]
+name = "divan-macros"
+version = "0.1.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9556bc800956545d6420a640173e5ba7dfa82f38d3ea5a167eb555bc69ac3323"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "errno"
+version = "0.3.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad"
+dependencies = [
+ "libc",
+ "windows-sys 0.60.2",
+]
+
+[[package]]
+name = "getrandom"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4"
 dependencies = [
  "cfg-if",
  "libc",
+ "r-efi",
  "wasi",
 ]
 
@@ -25,6 +105,12 @@ version = "0.2.174"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776"
 
+[[package]]
+name = "linux-raw-sys"
+version = "0.9.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12"
+
 [[package]]
 name = "ppv-lite86"
 version = "0.2.21"
@@ -53,21 +139,26 @@ dependencies = [
 ]
 
 [[package]]
-name = "rand"
+name = "r-efi"
-version = "0.8.5"
+version = "5.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
+checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
+
+[[package]]
+name = "rand"
+version = "0.9.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9fbfd9d094a40bf3ae768db9361049ace4c0e04a4fd6b359518bd7b73a73dd97"
 dependencies = [
- "libc",
  "rand_chacha",
  "rand_core",
 ]
 
 [[package]]
 name = "rand_chacha"
-version = "0.3.1"
+version = "0.9.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb"
 dependencies = [
  "ppv-lite86",
  "rand_core",
@@ -75,21 +166,41 @@ dependencies = [
 
 [[package]]
 name = "rand_core"
-version = "0.6.4"
+version = "0.9.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
+checksum = "99d9a13982dcf210057a8a78572b2217b667c3beacbf3a0d8b454f6f82837d38"
 dependencies = [
  "getrandom",
 ]
 
+[[package]]
+name = "regex-lite"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a"
+
 [[package]]
 name = "rustic_mcts"
 version = "0.1.0"
 dependencies = [
+ "divan",
  "rand",
  "thiserror",
 ]
 
+[[package]]
+name = "rustix"
+version = "1.0.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266"
+dependencies = [
+ "bitflags",
+ "errno",
+ "libc",
+ "linux-raw-sys",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "syn"
 version = "2.0.103"
@@ -101,6 +212,16 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "terminal_size"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "45c6481c4829e4cc63825e62c49186a34538b7b2750b73b266581ffb612fb5ed"
+dependencies = [
+ "rustix",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "thiserror"
 version = "2.0.12"
@@ -129,9 +250,167 @@ checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
 
 [[package]]
 name = "wasi"
-version = "0.11.1+wasi-snapshot-preview1"
+version = "0.14.2+wasi-0.2.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
+checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3"
+dependencies = [
+ "wit-bindgen-rt",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.60.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb"
+dependencies = [
+ "windows-targets 0.53.2",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973"
+dependencies = [
+ "windows_aarch64_gnullvm 0.52.6",
+ "windows_aarch64_msvc 0.52.6",
+ "windows_i686_gnu 0.52.6",
+ "windows_i686_gnullvm 0.52.6",
+ "windows_i686_msvc 0.52.6",
+ "windows_x86_64_gnu 0.52.6",
+ "windows_x86_64_gnullvm 0.52.6",
+ "windows_x86_64_msvc 0.52.6",
+]
+
+[[package]]
+name = "windows-targets"
+version = "0.53.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c66f69fcc9ce11da9966ddb31a40968cad001c5bedeb5c2b82ede4253ab48aef"
+dependencies = [
+ "windows_aarch64_gnullvm 0.53.0",
+ "windows_aarch64_msvc 0.53.0",
+ "windows_i686_gnu 0.53.0",
+ "windows_i686_gnullvm 0.53.0",
+ "windows_i686_msvc 0.53.0",
+ "windows_x86_64_gnu 0.53.0",
+ "windows_x86_64_gnullvm 0.53.0",
+ "windows_x86_64_msvc 0.53.0",
+]
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
+
+[[package]]
+name = "windows_i686_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.52.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.53.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486"
+
+[[package]]
+name = "wit-bindgen-rt"
+version = "0.39.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1"
+dependencies = [
+ "bitflags",
+]
 
 [[package]]
 name = "zerocopy"

Cargo.toml

@@ -10,7 +10,16 @@ readme = "README.md"
 keywords = ["mcts", "rust", "monte_carlo", "tree", "ai", "ml"]
 categories = ["algorithms", "data-structures"]
 
+[features]
+threads = []
 
 [dependencies]
-rand = "~0.8"
+rand = "~0.9"
 thiserror = "~2.0"
+
+[dev-dependencies]
+divan = "0.1.21"
+
+[[bench]]
+name = "e2e"
+harness = false

benches/e2e.rs

@@ -0,0 +1,191 @@
+use std::collections::HashMap;
+
+use rustic_mcts::policy::backprop::BackpropagationPolicy;
+use rustic_mcts::policy::decision::DecisionPolicy;
+use rustic_mcts::policy::selection::SelectionPolicy;
+use rustic_mcts::policy::simulation::SimulationPolicy;
+use rustic_mcts::{Action, GameState, MCTSConfig, RewardVal, MCTS};
+
+fn main() {
+    // Run registered benchmarks.
+    divan::main();
+}
+
+#[divan::bench]
+fn tic_tac_toe() {
+    // Set up a new game
+    let mut game = TicTacToe::new();
+
+    // Create MCTS configuration
+    let config = 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,
+    };
+
+    // Main game loop
+    while !game.is_terminal() {
+        // Create a new MCTS search
+        let mut mcts = MCTS::new(game.clone(), &config);
+
+        // Find the best move
+        match mcts.search() {
+            Ok(action) => {
+                // Apply the AI's move
+                game = game.state_after_action(&action);
+            }
+            Err(_) => {
+                break;
+            }
+        }
+    }
+}
+
+/// Players in Tic-Tac-Toe
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+enum Player {
+    X,
+    O,
+}
+
+impl rustic_mcts::Player for Player {}
+
+/// Tic-Tac-Toe move
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct Move {
+    /// Board position index (0-8)
+    index: usize,
+}
+
+impl Action for Move {
+    fn id(&self) -> usize {
+        self.index
+    }
+}
+
+/// Tic-Tac-Toe game state
+#[derive(Debug, Clone)]
+struct TicTacToe {
+    /// Board representation (None = empty, Some(Player) = occupied)
+    board: [Option<Player>; 9],
+
+    /// Current player's turn
+    current_player: Player,
+
+    /// Number of moves played so far
+    moves_played: usize,
+}
+
+impl TicTacToe {
+    /// Creates a new empty Tic-Tac-Toe board
+    fn new() -> Self {
+        TicTacToe {
+            board: [None; 9],
+            current_player: Player::X,
+            moves_played: 0,
+        }
+    }
+
+    /// Returns the winner of the game, if any
+    fn get_winner(&self) -> Option<Player> {
+        // Check rows
+        for row in 0..3 {
+            let i = row * 3;
+            if self.board[i].is_some()
+                && self.board[i] == self.board[i + 1]
+                && self.board[i] == self.board[i + 2]
+            {
+                return self.board[i];
+            }
+        }
+
+        // Check columns
+        for col in 0..3 {
+            if self.board[col].is_some()
+                && self.board[col] == self.board[col + 3]
+                && self.board[col] == self.board[col + 6]
+            {
+                return self.board[col];
+            }
+        }
+
+        // Check diagonals
+        if self.board[0].is_some()
+            && self.board[0] == self.board[4]
+            && self.board[0] == self.board[8]
+        {
+            return self.board[0];
+        }
+        if self.board[2].is_some()
+            && self.board[2] == self.board[4]
+            && self.board[2] == self.board[6]
+        {
+            return self.board[2];
+        }
+
+        None
+    }
+}
+
+impl GameState for TicTacToe {
+    type Action = Move;
+    type Player = Player;
+
+    fn get_legal_actions(&self) -> Vec<Self::Action> {
+        let mut actions = Vec::new();
+        for i in 0..9 {
+            if self.board[i].is_none() {
+                actions.push(Move { index: i });
+            }
+        }
+        actions
+    }
+
+    fn state_after_action(&self, action: &Self::Action) -> Self {
+        let mut new_state = self.clone();
+
+        // Make the move
+        new_state.board[action.index] = Some(self.current_player);
+        new_state.moves_played = self.moves_played + 1;
+
+        // Switch player
+        new_state.current_player = match self.current_player {
+            Player::X => Player::O,
+            Player::O => Player::X,
+        };
+
+        new_state
+    }
+
+    fn is_terminal(&self) -> bool {
+        self.get_winner().is_some() || self.moves_played == 9
+    }
+
+    fn reward_for_player(&self, player: &Self::Player) -> RewardVal {
+        if let Some(winner) = self.get_winner() {
+            if winner == *player {
+                return 1.0; // Win
+            } else {
+                return 0.0; // Loss
+            }
+        }
+
+        // Draw
+        0.5
+    }
+
+    fn rewards_for_players(&self) -> HashMap<Self::Player, RewardVal> {
+        HashMap::from_iter(vec![
+            (Player::X, self.reward_for_player(&Player::X)),
+            (Player::O, self.reward_for_player(&Player::O)),
+        ])
+    }
+
+    fn get_current_player(&self) -> &Self::Player {
+        &self.current_player
+    }
+}

src/mcts.rs

@@ -6,7 +6,8 @@ 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 crate::tree::print_tree;
+use rand::prelude::IndexedRandom;
 use std::collections::HashMap;
 use std::time::Instant;
 
@@ -74,7 +75,7 @@ impl<'conf, S: GameState + std::fmt::Debug> MCTS<'conf, S> {
         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();
+            let random_child: usize = *children.choose(&mut rand::rng()).unwrap();
             selected_id = random_child;
         }
         let rewards = self.simulate(selected_id);
@@ -130,6 +131,10 @@ impl<'conf, S: GameState + std::fmt::Debug> MCTS<'conf, S> {
             None => Err(MCTSError::NoBestAction),
         }
     }
+
+    pub fn print_tree(&self) {
+        print_tree(self.root_id, &self.arena)
+    }
 }
 
 /// Errors returned by the MCTS algorithm

src/policy/backprop/mod.rs

@@ -69,18 +69,14 @@ fn standard_backprop<S: GameState>(
     let mut current_id: usize = node_id;
     loop {
         let node = arena.get_node_mut(current_id);
-        let player = node.state.get_current_player().clone();
+        node.increment_visits();
-        match rewards.get(&player) {
+        for (player, reward) in rewards.iter() {
-            Some(reward) => {
+            node.record_player_reward(player.clone(), *reward);
-                node.increment_visits();
+        }
-                node.record_player_reward(player, *reward);
+        if let Some(parent_id) = node.parent {
-                if let Some(parent_id) = node.parent {
+            current_id = parent_id;
-                    current_id = parent_id;
+        } else {
-                } else {
+            break;
-                    break;
-                }
-            }
-            None => (),
         }
     }
 }
@@ -94,19 +90,14 @@ fn weighted_backprop<S: GameState>(
     let mut current_id: usize = node_id;
     loop {
         let node = arena.get_node_mut(current_id);
-        let player = node.state.get_current_player().clone();
         let weight = weight_for_depth(depth_factor, node.depth);
-        match rewards.get(&player) {
+        for (player, reward) in rewards.iter() {
-            Some(reward) => {
+            node.record_player_reward(player.clone(), (*reward) * weight);
-                node.increment_visits();
+        }
-                node.record_player_reward(player, (*reward) * weight);
+        if let Some(parent_id) = node.parent {
-                if let Some(parent_id) = node.parent {
+            current_id = parent_id;
-                    current_id = parent_id;
+        } else {
-                } else {
+            break;
-                    break;
-                }
-            }
-            None => (),
         }
     }
 }

src/policy/selection/ucb1_tuned.rs

@@ -36,15 +36,15 @@ pub fn select_best_child<S: GameState>(
                 exploration_constant,
                 parent_visits,
                 node_a.visits,
-                node_a.rewards(player),
                 node_a.reward_average(player),
+                node_a.reward_variance(player),
             );
             let ucb_b = ucb1_tuned_value(
                 exploration_constant,
                 parent_visits,
                 node_b.visits,
-                node_b.rewards(player),
                 node_b.reward_average(player),
+                node_b.reward_variance(player),
             );
             ucb_a.partial_cmp(&ucb_b).unwrap()
         })
@@ -57,41 +57,31 @@ pub fn ucb1_tuned_value(
     exploration_constant: f64,
     parent_visits: u64,
     child_visits: u64,
-    child_rewards: Option<&Vec<RewardVal>>,
     reward_avg: RewardVal,
+    reward_variance: f64,
 ) -> RewardVal {
-    match child_rewards {
+    if child_visits == 0 {
-        None => {
+        RewardVal::INFINITY // Always explore nodes that have never been visited
-            RewardVal::INFINITY // Always explore nodes that have never been visited
+    } else {
-        }
+        let parent_visits: RewardVal = parent_visits as RewardVal;
-        Some(child_rewards) => {
+        let child_visits: RewardVal = child_visits as RewardVal;
-            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 parent node
-                // n: number of visits to the child node
+        // n: number of visits to the child node
-                // x_i: reward of the ith visit to the child node
+        // x_i: reward of the ith visit to the child node
-                // X: average reward of the child
+        // X: average reward of the child
-                // C: exploration constant
+        // C: exploration constant
-                //
+        //
-                // UCB1-Tuned = X + C * sqrt(Ln(parent_visits) / child_visits * min(1/4, V_n)
+        // 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)
+        //   V(n) = sum(x_i^2)/n - X^2 + sqrt(2*ln(N)/n)
-                let exploitation = reward_avg;
+        let exploitation = reward_avg;
-                let mut variance = (child_rewards.iter().map(|&x| x * x).sum::<RewardVal>()
+        let mut variance = reward_variance + (2.0 * parent_visits.ln() / child_visits).sqrt();
-                    / child_visits)
+        if variance > 0.25 {
-                    - (reward_avg * reward_avg)
+            variance = 0.25;
-                    + (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
-            }
         }
+        let exploration =
+            exploration_constant * (parent_visits.ln() / child_visits * variance).sqrt();
+
+        exploitation + exploration
     }
 }

src/policy/simulation/random.rs

@@ -4,14 +4,14 @@
 
 use crate::state::GameState;
 use crate::tree::node::{Node, RewardVal};
-use rand::prelude::SliceRandom;
+use rand::prelude::IndexedRandom;
 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();
+        let action = legal_actions.choose(&mut rand::rng()).unwrap();
         state = state.state_after_action(&action);
     }
     state.rewards_for_players()

src/state.rs

@@ -7,7 +7,7 @@ use std::hash::Hash;
 ///
 /// When leveraging MCTS for your game, you must implement this trait to provide
 /// the specifics for your game.
-pub trait GameState: Clone {
+pub trait GameState: Clone + Debug {
     /// The type of actions that can be taken in the game
     type Action: Action;
 

src/tree/mod.rs

@@ -1,2 +1,27 @@
 pub mod arena;
 pub mod node;
+
+use crate::state::GameState;
+use crate::tree::arena::Arena;
+use crate::tree::node::Node;
+
+pub fn print_tree<S: GameState>(node_id: usize, arena: &Arena<S>) {
+    let mut to_print: Vec<usize> = Vec::new();
+    to_print.push(node_id);
+    while let Some(node_id) = to_print.pop() {
+        let node: &Node<S> = arena.get_node(node_id);
+        if node.depth > 0 {
+            for _ in 0..node.depth - 1 {
+                print!("|  ")
+            }
+            print!("|- ");
+        }
+        println!(
+            "{:?} a:{:?} v:{}  {:?}",
+            node_id, node.action, node.visits, node.player_view
+        );
+        for child_id in node.children.iter().rev() {
+            to_print.push(*child_id);
+        }
+    }
+}

src/tree/node.rs

@@ -1,5 +1,4 @@
 use std::collections::HashMap;
-use std::fmt::Debug;
 
 use crate::state::GameState;
 
@@ -15,7 +14,7 @@ pub type RewardVal = f64;
 ///
 /// This class is not thread safe, as the library does not provide for parallel
 /// search.
-#[derive(Debug)]
+#[derive(std::fmt::Debug)]
 pub struct Node<S: GameState> {
     /// The game state at the given node, after `action`
     pub state: S,
@@ -30,7 +29,7 @@ pub struct Node<S: GameState> {
     pub visits: u64,
 
     /// The player's evaluation of the node
-    pub player_view: HashMap<S::Player, PlayerNodeView>,
+    pub player_view: HashMap<S::Player, PlayerRewardView>,
 
     /// The depth of the node in the tree, this is 0 for the root node
     pub depth: usize,
@@ -70,10 +69,10 @@ impl<S: GameState> Node<S> {
         }
     }
 
-    pub fn rewards(&self, player: &S::Player) -> Option<&Vec<RewardVal>> {
+    pub fn reward_variance(&self, player: &S::Player) -> f64 {
         match self.player_view.get(player) {
-            Some(pv) => Some(&pv.rewards),
+            Some(pv) => pv.weighted_variance / self.visits as f64,
-            None => None,
+            None => 0.0,
         }
     }
 
@@ -82,37 +81,44 @@ impl<S: GameState> Node<S> {
     }
 
     pub fn record_player_reward(&mut self, player: S::Player, reward: RewardVal) {
-        let pv = self
+        let pv = self.player_view.entry(player).or_default();
-            .player_view
+        let prev_reward_average = pv.reward_average;
-            .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;
+        pv.reward_average = pv.reward_sum / self.visits as f64;
+        pv.weighted_variance += (reward - prev_reward_average) * (reward - pv.reward_average);
     }
 }
 
 /// 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 PlayerRewardView {
-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
+    /// The weighted variance from simulations through this node
-    pub rewards: Vec<RewardVal>,
+    ///
+    /// This is used to calculate online sample variance.
+    /// See Donald E. Knuth. Seminumerical Algorithms, volume 2 of The Art of Computer Programming, chapter 4.2.2, page 232
+    pub weighted_variance: f64,
 }
 
-impl Default for PlayerNodeView {
+impl Default for PlayerRewardView {
     fn default() -> Self {
-        PlayerNodeView {
+        PlayerRewardView {
             reward_sum: 0.0,
             reward_average: 0.0,
-            rewards: Vec::new(),
+            weighted_variance: 0.0,
         }
     }
 }
+
+impl std::fmt::Debug for PlayerRewardView {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{{sum={} avg={}}}", self.reward_sum, self.reward_average)?;
+        Ok(())
+    }
+}