diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index 8e062f3..256a9f4 100644 --- a/engine/src/mcts.rs +++ b/engine/src/mcts.rs @@ -1,6 +1,6 @@ // use crate::networks::PolicyNetwork; -use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold}; -use crate::net::encoding::{encode_board_states_batch, encode_move}; +use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold}; +use crate::net::encoding::{encode_board_state_perspective, encode_move}; use crate::net::model::ChessModel; use burn::prelude::Device; use chess::BoardStatus::{Checkmate, Stalemate}; @@ -10,6 +10,9 @@ use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES}; use rand::SeedableRng; use std::collections::HashMap; +//TODO retain tree after move selection? + +#[derive(Clone)] pub struct Node { pub prior: f32, pub children: Vec, @@ -36,40 +39,32 @@ impl Node { } pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize { - let sqrt_parent = (self.visit_count as f32).sqrt(); - let c_puct_sqrt = c_puct * sqrt_parent; - self.children .iter() .copied() .max_by(|&a, &b| { - let score_a = ucb_score_fast(&arena[a], c_puct_sqrt); - let score_b = ucb_score_fast(&arena[b], c_puct_sqrt); - score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal) + ucb_score(self, &arena[a], c_puct) + .partial_cmp(&ucb_score(self, &arena[b], c_puct)) + .unwrap() }) .expect("select_child on leaf") } } -#[inline] -fn ucb_score_fast(child: &Node, c_puct_sqrt: f32) -> f32 { - let prior_score = c_puct_sqrt * child.prior / (1.0 + child.visit_count as f32); +fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 { + let prior_score = c_puct * child.prior * ((parent.visit_count + 1) as f32).sqrt() + / (1.0 + child.visit_count as f32); - let value_score = if child.visit_count > 0 { - -child.value() + let value_score: f32; + if child.visit_count > 0 { + value_score = -child.value(); // value from opposing side } else { - 0.0 - }; + value_score = 0.0; + } value_score + prior_score } -impl Clone for Node { - fn clone(&self) -> Self { - Self::new(self.prior, self.board_state.clone(), self.last_move.clone()) - } -} - #[derive(Clone, Debug)] pub struct MctsResults { pub board_state: BoardState, @@ -127,10 +122,8 @@ pub struct Mcts { fn terminal_value(status: BoardStateStatus, stm: Color) -> Option { match status { - BoardStateStatus::Ongoing => None, - BoardStateStatus::Stalemate - | BoardStateStatus::Threefold - | BoardStateStatus::FiftyMove => Some(0.0), + Ongoing => None, + BoardStateStatus::Stalemate | Threefold | FiftyMove | InsufficientMaterial => Some(0.0), BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }), BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }), } @@ -171,159 +164,56 @@ fn softmax(logits: &[f32]) -> Vec { probs } -const LEAF_VL_COUNT: u32 = 10; -const LEAF_VL_WEIGHT: f32 = 10.0; - impl Mcts { pub fn search( &mut self, board_state: &BoardState, model: &ChessModel, device: &Device, + training: bool, ) -> MctsResults { let mut nodes = Vec::::new(); let root = 0; nodes.push(Node::new(0.0, board_state.clone(), None)); - // Expand root to create initial children and priors self.expand(root, &mut nodes, model, device); - // Apply Dirichlet noise to root children - self.add_dirichlet_noise(root, &mut nodes); + if training { + // APPLY DIRICHLET NOISE HERE + self.add_dirichlet_noise(root, &mut nodes); + } - let mut sims_done: usize = 0; let num_sims = self.config.num_simulations; - while sims_done < num_sims { - let mut leaf_nodes: Vec = Vec::new(); - let mut leaf_paths: Vec> = Vec::new(); + for _ in 0..num_sims { + let mut path = vec![root]; + let mut current = root; - let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done); - - while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims { - let mut path = vec![root]; - let mut current = root; - - // Apply virtual loss during selection path - while !nodes[current].children.is_empty() { - nodes[current].visit_count += 1; - nodes[current].value_sum -= 1.0; - - current = nodes[current].select_child(&nodes, &self.config.c_puct); - path.push(current); - } - - // Strong virtual loss on unexpanded leaf so batch collector visits diverse leaves across tree - nodes[current].visit_count += LEAF_VL_COUNT; - nodes[current].value_sum -= LEAF_VL_WEIGHT; - - // Check terminal status of leaf - let stm = nodes[current].board_state.board.side_to_move(); - if let Some(term_val) = terminal_value(nodes[current].board_state.status, stm) { - // Revert virtual loss for path - for &idx in &path[..path.len() - 1] { - nodes[idx].visit_count -= 1; - nodes[idx].value_sum += 1.0; - } - nodes[current].visit_count -= LEAF_VL_COUNT; - nodes[current].value_sum += LEAF_VL_WEIGHT; - - self.backpropagate(&mut nodes, &path, term_val, stm); - sims_done += 1; - continue; - } - - // Prevent expanding the exact same node twice in a single batch - if leaf_nodes.contains(¤t) { - // Revert virtual loss for this duplicate path - for &idx in &path[..path.len() - 1] { - nodes[idx].visit_count -= 1; - nodes[idx].value_sum += 1.0; - } - nodes[current].visit_count -= LEAF_VL_COUNT; - nodes[current].value_sum += LEAF_VL_WEIGHT; - break; - } - - leaf_nodes.push(current); - leaf_paths.push(path); + while !nodes[current].children.is_empty() { + current = nodes[current].select_child(&nodes, &self.config.c_puct); + path.push(current); } - if leaf_nodes.is_empty() { - if sims_done >= num_sims { - break; - } - continue; - } + let value: f32 = self.expand(current, &mut nodes, model, device); - // Batch evaluate non-terminal leaf states using single CPU buffer encoding - let leaf_refs: Vec<&BoardState> = leaf_nodes.iter().map(|&idx| &nodes[idx].board_state).collect(); - let batch = encode_board_states_batch(&leaf_refs, device); - let (policy_batch, value_batch) = model.forward(batch); - - let policy_data = policy_batch.into_data().to_vec::().unwrap(); - let value_data = value_batch.into_data().to_vec::().unwrap(); - - let num_moves = policy_data.len() / leaf_nodes.len(); - - for (i, &node_idx) in leaf_nodes.iter().enumerate() { - let path = &leaf_paths[i]; - - // Revert virtual loss along this path before applying true expansion & backpropagation - for &idx in &path[..path.len() - 1] { - nodes[idx].visit_count -= 1; - nodes[idx].value_sum += 1.0; - } - nodes[node_idx].visit_count -= LEAF_VL_COUNT; - nodes[node_idx].value_sum += LEAF_VL_WEIGHT; - - let start = i * num_moves; - let end = start + num_moves; - let logits = &policy_data[start..end]; - let probs = softmax(logits); - - let stm = nodes[node_idx].board_state.board.side_to_move(); - let legal_moves: Vec = - MoveGen::new_legal(&nodes[node_idx].board_state.board).collect(); - - let mut move_priors = Vec::with_capacity(legal_moves.len()); - let mut prior_sum = 0.0f32; - - for &mv in &legal_moves { - if let Some(enc_idx) = encode_move(mv, stm) { - let p = probs[enc_idx]; - move_priors.push((mv, p)); - prior_sum += p; - } - } - - for (mv, p) in move_priors { - let normalized_prior = if prior_sum > 0.0 { - p / prior_sum - } else { - 1.0 / legal_moves.len() as f32 - }; - - let mut new_board = nodes[node_idx].board_state.clone(); - new_board.apply_move(mv); - - let child_idx = nodes.len(); - nodes.push(Node::new(normalized_prior, new_board, Some(mv))); - nodes[node_idx].children.push(child_idx); - } - - let value = value_data[i]; - self.backpropagate(&mut nodes, path, value, stm); - sims_done += 1; - } + let color = nodes[current].board_state.board.side_to_move(); + self.backpropagate(&mut nodes, &path, value, color); } // Build compact move distribution let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len()); let stm = board_state.board.side_to_move(); - let total_root_visits: u32 = nodes[root].children.iter().map(|&c| nodes[c].visit_count).sum(); - let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 }; + let total_root_visits: u32 = nodes[root] + .children + .iter() + .map(|&c| nodes[c].visit_count) + .sum(); + let denom = if total_root_visits > 0 { + total_root_visits as f32 + } else { + 1.0 + }; for child_idx in nodes[root].children.iter() { let mv = nodes[*child_idx].last_move.expect("move didnt exist"); if let Some(enc) = encode_move(mv, stm) { @@ -348,8 +238,8 @@ impl Mcts { if let Some(term_val) = terminal_value(status, stm) { term_val } else { - let state = encode_board_states_batch(&[&arena[node_idx].board_state], device); - let (policy_head, value_head) = model.forward(state); + let state = encode_board_state_perspective(&arena[node_idx].board_state, device); + let (policy_head, value_head) = model.forward(state.unsqueeze::<4>()); let legal_moves: Vec = MoveGen::new_legal(&arena[node_idx].board_state.board).collect(); @@ -450,8 +340,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec { samples } - - pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> { // Choose best child of root by visit count (standard MCTS behavior) let root_node = &nodes[root]; @@ -540,6 +428,7 @@ pub enum BoardStateStatus { BlackWinner, Threefold, FiftyMove, + InsufficientMaterial, } #[derive(Clone, Debug)] @@ -555,52 +444,100 @@ impl BoardState { BoardState { board, halfmove_clock, - repetition_table: repetition_table, + repetition_table, status: Ongoing, } } pub fn apply_move(&mut self, mv: ChessMove) { - self.halfmove_clock += 1; - if self.board.piece_on(mv.get_source()) == Some(Pawn) - || self.board.piece_on(mv.get_dest()).is_some() - { + let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn); + let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some(); + + let is_en_passant = is_pawn_move + && self.board.en_passant() == Some(mv.get_dest()) + && self.board.piece_on(mv.get_dest()).is_none(); + + self.halfmove_clock = self.halfmove_clock.saturating_add(1); + + if is_pawn_move || is_normal_capture || is_en_passant { self.halfmove_clock = 0; } + self.board = self.board.make_move_new(mv); let board_hash = self.board.get_hash(); - - let current_rep = self.repetition_table.get(&board_hash); - let mut new_rep = 1; - if current_rep.is_some() { - new_rep = current_rep.unwrap() + 1; - } + let new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 1; self.repetition_table.insert(board_hash, new_rep); - if self.board.status() == Checkmate { - if self.board.side_to_move() == White { - // white's move after black played the mating move - self.status = BoardStateStatus::BlackWinner; - } else { - self.status = BoardStateStatus::WhiteWinner; + match self.board.status() { + Checkmate => { + self.status = if self.board.side_to_move() == White { + BoardStateStatus::BlackWinner + } else { + BoardStateStatus::WhiteWinner + }; + return; } - } else if self.board.status() == Stalemate { - self.status = BoardStateStatus::Stalemate; + Stalemate => { + self.status = BoardStateStatus::Stalemate; + return; + } + _ => {} + } + + if is_insufficient_material(&self.board) { + self.status = InsufficientMaterial; + return; } if self.halfmove_clock >= 100 { self.status = FiftyMove; + return; } if new_rep >= 3 { self.status = Threefold; + return; } + + self.status = BoardStateStatus::Ongoing; } } impl Default for BoardState { fn default() -> BoardState { - BoardState::new(Board::default(), 0, HashMap::new()) + let board = Board::default(); + let repetition_table = HashMap::from([(board.get_hash(), 1)]); + BoardState::new(board, 0, repetition_table) + } +} + +// Could be made more comprehensive +pub fn is_insufficient_material(board: &Board) -> bool { + let pawns = board.pieces(Pawn).0; + let rooks = board.pieces(Rook).0; + let queens = board.pieces(Queen).0; + + // Any pawn/rook/queen means this is not trivially dead. + if pawns != 0 || rooks != 0 || queens != 0 { + return false; + } + + let bishops = board.pieces(Bishop).0; + let knights = board.pieces(Knight).0; + + let bishop_count = bishops.count_ones(); + let knight_count = knights.count_ones(); + + match (bishop_count, knight_count) { + // K vs K + (0, 0) => true, + + // K+B vs K + // K+N vs K + (1, 0) | (0, 1) => true, + + // Other configurations require more analysis. + _ => false, } }