Updates

engine/src/mcts.rs

@@ -80,6 +80,7 @@ pub struct MctsConfig {
     pub c_puct: f32,
     pub dirichlet_alpha: f32,
     pub dirichlet_epsilon: f32,
+    pub batch_max: usize,
 }
 
 impl MctsConfig {
@@ -88,19 +89,21 @@ impl MctsConfig {
         c_puct: f32,
         dirichlet_alpha: f32,
         dirichlet_epsilon: f32,
+        batch_max: usize,
     ) -> MctsConfig {
         MctsConfig {
             num_simulations,
             c_puct,
             dirichlet_alpha,
             dirichlet_epsilon,
+            batch_max,
         }
     }
 }
 
 impl Default for MctsConfig {
     fn default() -> MctsConfig {
-        MctsConfig::new(400, 1.0, 0.05, 0.25)
+        MctsConfig::new(400, 1.0, 0.05, 0.25, 64)
     }
 }
 
@@ -128,10 +131,8 @@ impl<B: Backend> Mcts<B> {
         self.add_dirichlet_noise(root, &mut nodes);
 
         // We'll batch leaf evaluations to reduce per-leaf model calls and device-host syncs.
-        let mut sims_done = 0usize;
+        let mut sims_done: usize = 0;
         let num_sims = self.config.num_simulations;
-        // Tunable batch size for NN evaluation. Small value is safe; larger values increase throughput on GPU.
-        let batch_max = 32usize;
 
         while sims_done < num_sims {
             // Collect a batch of leaf nodes (and their selection paths)
@@ -139,7 +140,7 @@ impl<B: Backend> Mcts<B> {
             let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
             let mut leaf_states: Vec<Tensor<B, 4>> = Vec::new();
 
-            while leaf_nodes.len() < std::cmp::min(batch_max, num_sims - sims_done) {
+            while leaf_nodes.len() < std::cmp::min(self.config.batch_max, num_sims - sims_done) {
                 let mut path = vec![root];
                 let mut current = root;
 
@@ -153,8 +154,9 @@ impl<B: Backend> Mcts<B> {
                 leaf_paths.push(path.clone());
 
                 // Prepare state tensor for this leaf
-                let state: Tensor<B, 4> = encode_board_state_perspective(&nodes[current].board_state, device)
-                    .reshape([1, 18, 8, 8]);
+                let state: Tensor<B, 4> =
+                    encode_board_state_perspective(&nodes[current].board_state, device)
+                        .reshape([1, 18, 8, 8]);
                 leaf_states.push(state);
 
                 sims_done += 1;
@@ -184,32 +186,51 @@ impl<B: Backend> Mcts<B> {
                 let logits = &policy_data[start..end];
 
                 // Convert logits to probabilities with a numerically-stable softmax on host
-                let mut max_logit = std::f32::NEG_INFINITY;
+                let mut max_logit = f32::NEG_INFINITY;
                 for &v in logits.iter() {
                     if v > max_logit {
                         max_logit = v;
                     }
                 }
-                let mut exps_sum = 0.0f32;
-                // We'll build a Vec<f32> of probabilities lazily when needed
-                let mut probs: Vec<f32> = Vec::new();
-                probs.resize(num_moves, 0.0);
-                for (j, &v) in logits.iter().enumerate() {
-                    let e = (v - max_logit).exp();
-                    probs[j] = e;
-                    exps_sum += e;
+                if logits.is_empty() {
+                    println!("logits empty")
                 }
-                if exps_sum > 0.0 {
-                    for p in probs.iter_mut() {
+                if !max_logit.is_finite() {
+                    println!("max logits not finite")
+                }
+
+                let mut exps_sum = 0.0f32;
+                let mut probs = vec![0.0f32; num_moves];
+
+                for (j, &v) in logits.iter().enumerate() {
+                    let x = (v - max_logit).exp();
+
+                    if !x.is_finite() {
+                        probs[j] = 0.0;
+                    } else {
+                        probs[j] = x;
+                        exps_sum += x;
+                    }
+                }
+
+                if exps_sum > 0.0 && exps_sum.is_finite() {
+                    for p in &mut probs {
                         *p /= exps_sum;
                     }
+                } else {
+                    // fallback: uniform or zero
+                    let uniform = 1.0 / num_moves as f32;
+                    for p in &mut probs {
+                        *p = uniform;
+                    }
                 }
 
                 // Expand: add legal moves as children with prior from probs
-                let legal_moves: Vec<ChessMove> = MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();
+                let legal_moves: Vec<ChessMove> =
+                    MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();
                 for mv in legal_moves {
                     let stm = nodes[node_idx].board_state.board.side_to_move();
-                    let idx = encode_move(mv, stm);
+                    let idx = encode_move(mv, stm).expect("Invalid move");
                     let prior = probs[idx];
 
                     let mut new_board = nodes[node_idx].board_state.clone();
@@ -233,7 +254,7 @@ impl<B: Backend> Mcts<B> {
         let denom = self.config.num_simulations as f32;
         for child_idx in nodes[root].children.iter() {
             let mv = nodes[*child_idx].last_move.expect("move didnt exist");
-            let enc = encode_move(mv, stm);
+            let enc = encode_move(mv, stm).expect("Invalid move");
             let prob = nodes[*child_idx].visit_count as f32 / denom;
             move_dist.push((enc, prob));
         }
@@ -280,7 +301,7 @@ impl<B: Backend> Mcts<B> {
 
             for mv in legal_moves {
                 let stm = arena[node_idx].board_state.board.side_to_move();
-                let idx = encode_move(mv, stm);
+                let idx = encode_move(mv, stm).expect("Invalid move");
                 let prior = policy[idx];
 
                 let mut new_board = arena[node_idx].board_state.clone();