12 Commits

Author SHA1 Message Date
krugd 3676838082 Updated rust dependencies 2026-07-07 12:29:54 -07:00
krugd 6aa9002e92 GameState::reward_for_player is no longer a trait function
We need GameState::rewards_for_players implemented, the library
currently doesn't need to specify for a singular player.
2025-06-30 21:41:49 -07:00
krugd 76051cd76b Removing the id field from the Action trait
This would in theory be useful for a transposition table, but we do not
currently support that. As such I don't want to burden the
implementation with that field until it is deemed necessary.
2025-06-30 21:41:27 -07:00
krugd a7102a0e44 Fixed missing non terminal error 2025-06-30 20:08:18 -07:00
krugd a2711f341f Cleanup debug print 2025-06-29 22:51:17 -07:00
krugd 48c93fa197 Renamed the benchmark file 2025-06-29 22:50:54 -07:00
krugd 854772a63d Improve UCB1-Tuned performance
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.
2025-06-29 18:43:52 -07:00
krugd 44ef9ebdd8 Adding really stupid benchmark 2025-06-27 22:25:56 -07:00
krugd 6a33818238 We propagate the reward for both sides
The AI now properly choses the optimal path for the active player
2025-06-27 16:12:40 -07:00
krugd 0f9d4f0c4e Use Entry::or_default 2025-06-27 15:40:58 -07:00
krugd 9f893b0005 Upgrade rand to ~0.9 2025-06-27 15:17:45 -07:00
krugd b80f039b93 Adding a basic print_tree function to visualize the MCTS search tree
This probably is not a good thing to run on a very large tree.
2025-06-27 15:00:43 -07:00
12 changed files with 517 additions and 168 deletions
Generated
+182 -37
View File
@@ -3,27 +3,113 @@
version = 4
[[package]]
name = "cfg-if"
version = "1.0.1"
name = "anstyle"
version = "1.0.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268"
checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000"
[[package]]
name = "bitflags"
version = "2.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8"
[[package]]
name = "cfg-if"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "clap"
version = "4.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51"
dependencies = [
"clap_builder",
]
[[package]]
name = "clap_builder"
version = "4.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f"
dependencies = [
"anstyle",
"clap_lex",
"terminal_size",
]
[[package]]
name = "clap_lex"
version = "1.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9"
[[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.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb"
dependencies = [
"libc",
"windows-sys",
]
[[package]]
name = "getrandom"
version = "0.2.16"
version = "0.3.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592"
checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd"
dependencies = [
"cfg-if",
"libc",
"wasi",
"r-efi",
"wasip2",
]
[[package]]
name = "libc"
version = "0.2.174"
version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53"
[[package]]
name = "ppv-lite86"
@@ -36,38 +122,43 @@ dependencies = [
[[package]]
name = "proc-macro2"
version = "1.0.95"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "02b3e5e68a3a1a02aad3ec490a98007cbc13c37cbe84a3cd7b8e406d76e7f778"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.40"
version = "1.0.46"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d"
checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368"
dependencies = [
"proc-macro2",
]
[[package]]
name = "rand"
version = "0.8.5"
name = "r-efi"
version = "5.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404"
checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f"
[[package]]
name = "rand"
version = "0.9.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea"
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,26 +166,46 @@ dependencies = [
[[package]]
name = "rand_core"
version = "0.6.4"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c"
checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c"
dependencies = [
"getrandom",
]
[[package]]
name = "regex-lite"
version = "0.1.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973"
[[package]]
name = "rustic_mcts"
version = "0.1.0"
dependencies = [
"divan",
"rand",
"thiserror",
]
[[package]]
name = "syn"
version = "2.0.103"
name = "rustix"
version = "1.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4307e30089d6fd6aff212f2da3a1f9e32f3223b1f010fb09b7c95f90f3ca1e8"
checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190"
dependencies = [
"bitflags",
"errno",
"libc",
"linux-raw-sys",
"windows-sys",
]
[[package]]
name = "syn"
version = "2.0.118"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422"
dependencies = [
"proc-macro2",
"quote",
@@ -102,19 +213,29 @@ dependencies = [
]
[[package]]
name = "thiserror"
version = "2.0.12"
name = "terminal_size"
version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708"
checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874"
dependencies = [
"rustix",
"windows-sys",
]
[[package]]
name = "thiserror"
version = "2.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4"
dependencies = [
"thiserror-impl",
]
[[package]]
name = "thiserror-impl"
version = "2.0.12"
version = "2.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d"
checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5"
dependencies = [
"proc-macro2",
"quote",
@@ -123,30 +244,54 @@ dependencies = [
[[package]]
name = "unicode-ident"
version = "1.0.18"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "wasi"
version = "0.11.1+wasi-snapshot-preview1"
name = "wasip2"
version = "1.0.4+wasi-0.2.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487"
dependencies = [
"wit-bindgen",
]
[[package]]
name = "windows-link"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
[[package]]
name = "windows-sys"
version = "0.61.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc"
dependencies = [
"windows-link",
]
[[package]]
name = "wit-bindgen"
version = "0.57.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e"
[[package]]
name = "zerocopy"
version = "0.8.26"
version = "0.8.53"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1039dd0d3c310cf05de012d8a39ff557cb0d23087fd44cad61df08fc31907a2f"
checksum = "75726053136156d419e285b9b7eddaaea9e3fea6ce32eed44a89901f0bd98de1"
dependencies = [
"zerocopy-derive",
]
[[package]]
name = "zerocopy-derive"
version = "0.8.26"
version = "0.8.53"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181"
checksum = "4714fd92cf900833d49538023a9b3915155210801d1c1169eba513b2addefd71"
dependencies = [
"proc-macro2",
"quote",
+8 -1
View File
@@ -12,5 +12,12 @@ categories = ["algorithms", "data-structures"]
[dependencies]
rand = "~0.8"
rand = "~0.9"
thiserror = "~2.0"
[dev-dependencies]
divan = "~0.1"
[[bench]]
name = "e2e"
harness = false
+187
View File
@@ -0,0 +1,187 @@
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 {}
/// 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
}
fn reward_for_player(&self, player: &Player) -> RewardVal {
if let Some(winner) = self.get_winner() {
if winner == *player {
return 1.0; // Win
} else {
return 0.0; // Loss
}
}
// Draw
0.5
}
}
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 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
}
}
+14 -18
View File
@@ -84,11 +84,7 @@ struct Move {
index: usize,
}
impl Action for Move {
fn id(&self) -> usize {
self.index
}
}
impl Action for Move {}
/// Tic-Tac-Toe game state
#[derive(Clone)]
@@ -152,6 +148,19 @@ impl TicTacToe {
None
}
fn reward_for_player(&self, player: &Player) -> RewardVal {
if let Some(winner) = self.get_winner() {
if winner == *player {
return 1.0; // Win
} else {
return 0.0; // Loss
}
}
// Draw
0.5
}
}
impl GameState for TicTacToe {
@@ -188,19 +197,6 @@ impl GameState for TicTacToe {
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)),
+14 -18
View File
@@ -119,11 +119,7 @@ struct Move {
index: usize,
}
impl Action for Move {
fn id(&self) -> usize {
self.index
}
}
impl Action for Move {}
/// Tic-Tac-Toe game state
#[derive(Clone)]
@@ -195,6 +191,19 @@ impl TicTacToe {
None
}
fn reward_for_player(&self, player: &Player) -> RewardVal {
if let Some(winner) = self.get_winner() {
if winner == *player {
return 1.0; // Win
} else {
return 0.0; // Loss
}
}
// Draw
0.5
}
}
impl GameState for TicTacToe {
@@ -231,19 +240,6 @@ impl GameState for TicTacToe {
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)),
+15 -3
View File
@@ -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,8 +75,15 @@ 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();
selected_id = random_child;
match children.choose(&mut rand::rng()) {
Some(&random_child) => {
selected_id = random_child;
}
None => {
// We ran out of nodes
return Err(MCTSError::NonTerminalGame);
}
}
}
let rewards = self.simulate(selected_id);
self.backprop(selected_id, &rewards);
@@ -130,6 +138,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
+15 -24
View File
@@ -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();
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 => (),
node.increment_visits();
for (player, reward) in rewards.iter() {
node.record_player_reward(player.clone(), *reward);
}
if let Some(parent_id) = node.parent {
current_id = parent_id;
} else {
break;
}
}
}
@@ -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) {
Some(reward) => {
node.increment_visits();
node.record_player_reward(player, (*reward) * weight);
if let Some(parent_id) = node.parent {
current_id = parent_id;
} else {
break;
}
}
None => (),
for (player, reward) in rewards.iter() {
node.record_player_reward(player.clone(), (*reward) * weight);
}
if let Some(parent_id) = node.parent {
current_id = parent_id;
} else {
break;
}
}
}
+24 -34
View File
@@ -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 {
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;
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
}
// 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 = reward_variance + (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
}
}
+2 -2
View File
@@ -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()
+6 -12
View File
@@ -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;
@@ -40,11 +40,11 @@ pub trait GameState: Clone {
/// 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
/// Returns the rewards for all players from their perspective 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 evaluates the current state from the perspective of each player, and
/// returns a HashMap mapping each player to the result of this evaluation, which
/// we call the reward.
///
/// This is used in the MCTS backpropagation and simulation phases to evaluate
/// the value of a given node in the search tree.
@@ -55,9 +55,6 @@ pub trait GameState: Clone {
/// - 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
@@ -72,10 +69,7 @@ pub trait GameState: Clone {
///
/// 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;
}
pub trait Action: Clone + Debug {}
/// Trait used for players participating in a game
pub trait Player: Clone + Debug + PartialEq + Eq + Hash {}
+25
View File
@@ -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);
}
}
}
+25 -19
View File
@@ -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),
None => None,
Some(pv) => pv.weighted_variance / self.visits as f64,
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
.player_view
.entry(player)
.or_insert(PlayerNodeView::default());
pv.rewards.push(reward);
let pv = self.player_view.entry(player).or_default();
let prev_reward_average = pv.reward_average;
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 PlayerNodeView {
pub struct PlayerRewardView {
/// 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>,
/// The weighted variance from simulations through this node
///
/// 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(())
}
}