unbatch mcts

engine/src/mcts.rs

@@ -1,6 +1,6 @@
 // use crate::networks::PolicyNetwork;
-use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
+use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
-use crate::net::encoding::{encode_board_states_batch, encode_move};
+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<usize>,
@@ -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);
+                ucb_score(self, &arena[a], c_puct)
-                let score_b = ucb_score_fast(&arena[b], c_puct_sqrt);
+                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
-                score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
+                    .unwrap()
             })
             .expect("select_child on leaf")
     }
 }
 
-#[inline]
+fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
-fn ucb_score_fast(child: &Node, c_puct_sqrt: f32) -> f32 {
+    let prior_score = c_puct * child.prior * ((parent.visit_count + 1) as f32).sqrt()
-    let prior_score = c_puct_sqrt * child.prior / (1.0 + child.visit_count as f32);
+        / (1.0 + child.visit_count as f32);
 
-    let value_score = if child.visit_count > 0 {
+    let value_score: f32;
-        -child.value()
+    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<f32> {
     match status {
-        BoardStateStatus::Ongoing => None,
+        Ongoing => None,
-        BoardStateStatus::Stalemate
+        BoardStateStatus::Stalemate | Threefold | FiftyMove | InsufficientMaterial => Some(0.0),
-        | BoardStateStatus::Threefold
-        | BoardStateStatus::FiftyMove => 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<f32> {
     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::<Node>::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
+        if training {
-        self.add_dirichlet_noise(root, &mut nodes);
+            // 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 {
+        for _ in 0..num_sims {
-            let mut leaf_nodes: Vec<usize> = Vec::new();
+            let mut path = vec![root];
-            let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
+            let mut current = root;
 
-            let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
+            while !nodes[current].children.is_empty() {
-
+                current = nodes[current].select_child(&nodes, &self.config.c_puct);
-            while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims {
+                path.push(current);
-                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(&current) {
-                    // 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);
             }
 
-            if leaf_nodes.is_empty() {
+            let value: f32 = self.expand(current, &mut nodes, model, device);
-                if sims_done >= num_sims {
-                    break;
-                }
-                continue;
-            }
 
-            // Batch evaluate non-terminal leaf states using single CPU buffer encoding
+            let color = nodes[current].board_state.board.side_to_move();
-            let leaf_refs: Vec<&BoardState> = leaf_nodes.iter().map(|&idx| &nodes[idx].board_state).collect();
+            self.backpropagate(&mut nodes, &path, value, color);
-            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::<f32>().unwrap();
-            let value_data = value_batch.into_data().to_vec::<f32>().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<ChessMove> =
-                    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;
-            }
         }
 
         // 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 total_root_visits: u32 = nodes[root]
-        let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 };
+            .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 state = encode_board_state_perspective(&arena[node_idx].board_state, device);
-            let (policy_head, value_head) = model.forward(state);
+            let (policy_head, value_head) = model.forward(state.unsqueeze::<4>());
 
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
@@ -450,8 +340,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
     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;
+        let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
-        if self.board.piece_on(mv.get_source()) == Some(Pawn)
+        let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();
-            || 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 new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 1;
-        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;
-        }
         self.repetition_table.insert(board_hash, new_rep);
 
-        if self.board.status() == Checkmate {
+        match self.board.status() {
-            if self.board.side_to_move() == White {
+            Checkmate => {
-                // white's move after black played the mating move
+                self.status = if self.board.side_to_move() == White {
-                self.status = BoardStateStatus::BlackWinner;
+                    BoardStateStatus::BlackWinner
-            } else {
+                } else {
-                self.status = BoardStateStatus::WhiteWinner;
+                    BoardStateStatus::WhiteWinner
+                };
+                return;
             }
-        } else if self.board.status() == Stalemate {
+            Stalemate => {
-            self.status = BoardStateStatus::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,
     }
 }