Updates

.cargo/config.toml

@@ -0,0 +1,8 @@
+# for Linux
+[target.x86_64-unknown-linux-gnu]
+linker = "clang"
+rustflags = ["-C", "link-arg=-fuse-ld=lld"]
+
+# for Windows
+[target.x86_64-pc-windows-msvc]
+linker = "rust-lld.exe"

Cargo.lock

@@ -2167,6 +2167,8 @@ dependencies = [
  "ctrlc",
  "rand 0.10.1",
  "rand_distr 0.6.0",
+ "serde",
+ "serde_json",
 ]
 
 [[package]]

engine/Cargo.toml

@@ -10,3 +10,5 @@ burn-ndarray = "0.21.0"
 rand_distr = "0.6.0"
 rand = "0.10.1"
 ctrlc = "3.5.2"
+serde_json = "1.0.150"
+serde = { version = "1.0.228", features = ["derive"] }

engine/src/main.rs

@@ -1,6 +1,6 @@
 #![recursion_limit = "256"]
 
-use burn::backend::{Autodiff, Wgpu};
+use burn::backend::{Autodiff, Cuda};
 use burn::optim::AdamConfig;
 use engine::mcts::MctsConfig;
 use engine::training::train::{train, TrainingConfig};
@@ -14,15 +14,15 @@ use engine::training::train::{train, TrainingConfig};
 // }
 
 fn main() {
-    type MyBackend = Wgpu<f32, i32>;
-    // type MyBackend = Cuda<f32, i32>;
+    // type MyBackend = Wgpu<f32, i32>;
+    type MyBackend = Cuda<f32, i32>;
     // type MyBackend = NdArray<f32, i32>;
     type MyAutodiffBackend = Autodiff<MyBackend>;
-    let device = burn::backend::wgpu::WgpuDevice::default();
+    // let device = burn::backend::wgpu::WgpuDevice::default();
     // let device = burn::backend::ndarray::NdArrayDevice::default();
-    // let device = burn::backend::cuda::CudaDevice::default();
+    let device = burn::backend::cuda::CudaDevice::default();
 
-    let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25);
+    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -40,6 +40,7 @@ fn main() {
         mcts_config,
         optimizer: adam_config,
         lr: 2e-4,
+        seed: None,
     };
 
     train::<MyAutodiffBackend>(training_config, device);

engine/src/mcts.rs

@@ -5,9 +5,10 @@ use crate::net::model::ChessModel;
 use burn::prelude::Backend;
 use burn::Tensor;
 use chess::BoardStatus::{Checkmate, Stalemate};
-use chess::Color::White;
+use chess::Color::{Black, White};
 use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
 use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
+use rand::SeedableRng;
 use std::collections::HashMap;
 use std::marker::PhantomData;
 
@@ -58,16 +59,13 @@ impl Clone for Node {
 #[derive(Clone, Debug)]
 pub struct MctsResults {
     pub board_state: BoardState,
-    pub move_dist: HashMap<ChessMove, f32>,
+    // Compact encoded move distribution: (encoded_move_index, probability)
+    pub move_dist: Vec<(usize, f32)>,
     pub value: f32,
 }
 
 impl MctsResults {
-    pub fn new(
-        board_state: BoardState,
-        move_dist: HashMap<ChessMove, f32>,
-        value: f32,
-    ) -> MctsResults {
+    pub fn new(board_state: BoardState, move_dist: Vec<(usize, f32)>, value: f32) -> MctsResults {
         MctsResults {
             board_state,
             move_dist,
@@ -82,6 +80,7 @@ pub struct MctsConfig {
     pub c_puct: f32,
     pub dirichlet_alpha: f32,
     pub dirichlet_epsilon: f32,
+    pub batch_max: usize,
 }
 
 impl MctsConfig {
@@ -90,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)
     }
 }
 
@@ -123,12 +124,23 @@ impl<B: Backend> Mcts<B> {
         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 HERE
+        // Apply Dirichlet noise to root children
         self.add_dirichlet_noise(root, &mut nodes);
 
-        for _ in 0..self.config.num_simulations {
+        // We'll batch leaf evaluations to reduce per-leaf model calls and device-host syncs.
+        let mut sims_done: usize = 0;
+        let num_sims = self.config.num_simulations;
+
+        while sims_done < num_sims {
+            // Collect a batch of leaf nodes (and their selection paths)
+            let mut leaf_nodes: Vec<usize> = Vec::new();
+            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(self.config.batch_max, num_sims - sims_done) {
                 let mut path = vec![root];
                 let mut current = root;
 
@@ -137,18 +149,114 @@ impl<B: Backend> Mcts<B> {
                     path.push(current);
                 }
 
-            let value: f32 = self.expand(current, &mut nodes, model, device);
+                // Record leaf node and its path
+                leaf_nodes.push(current);
+                leaf_paths.push(path.clone());
 
-            let color = nodes[current].board_state.board.side_to_move();
-            self.backpropagate(&mut nodes, &path, value, color);
+                // 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]);
+                leaf_states.push(state);
+
+                sims_done += 1;
             }
 
-        let mut move_dist: HashMap<ChessMove, f32> = HashMap::new(); // TODO: make vec<(Chessmove, f32)>
-        for idx in nodes[root].children.iter() {
-            move_dist.insert(
-                nodes[*idx].last_move.expect("move didnt exist"),
-                nodes[*idx].visit_count as f32 / self.config.num_simulations as f32,
-            );
+            if leaf_nodes.is_empty() {
+                break;
+            }
+
+            // Batch evaluate the collected leaf states
+            let batch = Tensor::cat(leaf_states, 0);
+            let (policy_batch, value_batch) = model.forward(batch);
+
+            // Move tensors to host once per 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();
+
+            // Process each evaluated leaf: expand and backpropagate
+            for (i, &node_idx) in leaf_nodes.iter().enumerate() {
+                let path = &leaf_paths[i];
+
+                // slice for this sample's logits
+                let start = i * num_moves;
+                let end = start + num_moves;
+                let logits = &policy_data[start..end];
+
+                // Convert logits to probabilities with a numerically-stable softmax on host
+                let mut max_logit = f32::NEG_INFINITY;
+                for &v in logits.iter() {
+                    if v > max_logit {
+                        max_logit = v;
+                    }
+                }
+                if logits.is_empty() {
+                    println!("logits empty")
+                }
+                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();
+                for mv in legal_moves {
+                    let stm = nodes[node_idx].board_state.board.side_to_move();
+                    let idx = encode_move(mv, stm).expect("Invalid move");
+                    let prior = probs[idx];
+
+                    let mut new_board = nodes[node_idx].board_state.clone();
+                    new_board.apply_move(mv);
+
+                    let child_idx = nodes.len();
+                    nodes.push(Node::new(prior, new_board, Some(mv)));
+                    nodes[node_idx].children.push(child_idx);
+                }
+
+                // Backpropagate the value for this leaf
+                let value = value_data[i];
+                let color = nodes[node_idx].board_state.board.side_to_move();
+                self.backpropagate(&mut nodes, path, value, color);
+            }
+        }
+
+        // Build compact move distribution: encoded move index -> probability (visits / num_simulations)
+        let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
+        let stm = board_state.board.side_to_move();
+        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).expect("Invalid move");
+            let prob = nodes[*child_idx].visit_count as f32 / denom;
+            move_dist.push((enc, prob));
         }
 
         MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
@@ -161,6 +269,24 @@ impl<B: Backend> Mcts<B> {
         model: &ChessModel<B>,
         device: &B::Device,
     ) -> f32 {
+        if arena[node_idx].board_state.status == BoardStateStatus::Stalemate
+            || arena[node_idx].board_state.status == BoardStateStatus::Threefold
+            || arena[node_idx].board_state.status == BoardStateStatus::FiftyMove
+        {
+            0.0
+        } else if arena[node_idx].board_state.status == BoardStateStatus::WhiteWinner {
+            if arena[node_idx].board_state.board.side_to_move() == Black {
+                1.0
+            } else {
+                -1.0
+            }
+        } else if arena[node_idx].board_state.status == BoardStateStatus::BlackWinner {
+            if arena[node_idx].board_state.board.side_to_move() == White {
+                1.0
+            } else {
+                -1.0
+            }
+        } else {
             let state: Tensor<B, 4> =
                 encode_board_state_perspective(&arena[node_idx].board_state, device)
                     .reshape([1, 18, 8, 8]);
@@ -175,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();
@@ -189,6 +315,7 @@ impl<B: Backend> Mcts<B> {
 
             value_head.into_data().to_vec().unwrap()[0]
         }
+    }
 
     fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) {
         for &idx in path {
@@ -229,9 +356,14 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
 
-    let mut samples: Vec<f32> = (0..size)
-        .map(|_| gamma.sample(&mut rand::rng()) as f32)
-        .collect();
+    // Use a single SmallRng seeded from system time (avoid depending on thread_rng helper)
+    let now = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .unwrap();
+    let seed = now.as_nanos() as u64;
+    let mut rng = rand::rngs::SmallRng::seed_from_u64(seed);
+
+    let mut samples: Vec<f32> = (0..size).map(|_| gamma.sample(&mut rng) as f32).collect();
 
     let sum: f32 = samples.iter().sum();
 
@@ -334,8 +466,6 @@ pub fn heuristic_eval(board: &Board, perspective: Color) -> f32 {
     }
 
     value
-
-    // board.checkers()
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

engine/src/net/encoding.rs

@@ -20,6 +20,38 @@ Output:
 65-73: underpromotions
 */
 
+const PLANES_PER_SQUARE: usize = 73;
+
+const SLIDING_DIRS: [(i8, i8); 8] = [
+    (0, 1),   // N
+    (1, 1),   // NE
+    (1, 0),   // E
+    (1, -1),  // SE
+    (0, -1),  // S
+    (-1, -1), // SW
+    (-1, 0),  // W
+    (-1, 1),  // NW
+];
+
+const KNIGHT_DIRS: [(i8, i8); 8] = [
+    (1, 2),
+    (2, 1),
+    (2, -1),
+    (1, -2),
+    (-1, -2),
+    (-2, -1),
+    (-2, 1),
+    (-1, 2),
+];
+
+const UNDERPROMO_DIRS: [(i8, i8); 3] = [
+    (-1, 1), // capture left
+    (0, 1),  // forward
+    (1, 1),  // capture right
+];
+
+const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
+
 pub fn encode_board_state_perspective<B: Backend>(
     state: &BoardState,
     device: &B::Device,
@@ -120,180 +152,163 @@ fn fill_plane(buffer: &mut [f32], plane: usize) {
     }
 }
 
-pub fn encode_move(mv: ChessMove, side_to_move: Color) -> usize {
-    let from = mv.get_source().to_index();
-    let to = mv.get_dest().to_index();
+/// Rotate a square into the side-to-move perspective.
+fn orient_square(sq: Square, stm: Color) -> (i8, i8) {
+    let file = sq.get_file().to_index() as i8;
+    let rank = sq.get_rank().to_index() as i8;
 
-    let mut from_rank = from / 8;
-    let mut from_file = from % 8;
-    let mut to_rank = to / 8;
-    let mut to_file = to % 8;
-
-    if side_to_move == Color::Black {
-        from_rank = 7 - from_rank;
-        from_file = 7 - from_file;
-
-        to_rank = 7 - to_rank;
-        to_file = 7 - to_file;
+    match stm {
+        Color::White => (file, rank),
+        Color::Black => (7 - file, 7 - rank),
     }
-
-    let delta_rank = to_rank as i32 - from_rank as i32;
-    let delta_file = to_file as i32 - from_file as i32;
-
-    let plane = encode_move_type(delta_rank, delta_file, mv.get_promotion());
-
-    plane * 64 + (from_rank * 8 + from_file)
 }
 
-fn encode_move_type(dr: i32, df: i32, promotion: Option<Piece>) -> usize {
-    // Knight moves
-    const KNIGHT_DELTAS: [(i32, i32); 8] = [
-        (2, 1),
-        (1, 2),
-        (-1, 2),
-        (-2, 1),
-        (-2, -1),
-        (-1, -2),
-        (1, -2),
-        (2, -1),
-    ];
+/// Convert a perspective-space coordinate back into a real square.
+fn deorient_square(file: i8, rank: i8, stm: Color) -> Square {
+    let (file, rank) = match stm {
+        Color::White => (file, rank),
+        Color::Black => (7 - file, 7 - rank),
+    };
 
-    for (i, (r, f)) in KNIGHT_DELTAS.iter().enumerate() {
-        if dr == *r && df == *f {
-            return 56 + i;
-        }
-    }
+    Square::make_square(
+        Rank::from_index(rank as usize),
+        File::from_index(file as usize),
+    )
+}
 
-    // UNDERPromotions
-    if let Some(promo) = promotion {
+/// Encode a move into [0, 4672).
+pub fn encode_move(mv: ChessMove, stm: Color) -> Option<usize> {
+    let (ff, fr) = orient_square(mv.get_source(), stm);
+    let (tf, tr) = orient_square(mv.get_dest(), stm);
+
+    let dx = tf - ff;
+    let dy = tr - fr;
+
+    let plane = if let Some(promo) = mv.get_promotion() {
         if promo != Piece::Queen {
-            let dir = if df == 0 {
-                0
-            } else if df < 0 {
-                1
+            let dir = UNDERPROMO_DIRS
+                .iter()
+                .position(|&(x, y)| x == dx && y == dy)?;
+
+            let piece = UNDERPROMO_PIECES.iter().position(|&p| p == promo)?;
+
+            64 + piece * 3 + dir
         } else {
-                2
-            };
-
-            let piece_index = match promo {
-                // Piece::Queen => 0,
-                Piece::Rook => 0,
-                Piece::Bishop => 1,
-                Piece::Knight => 2,
-                _ => unreachable!(),
-            };
-
-            return 64 + dir * 3 + piece_index;
+            encode_non_underpromo(dx, dy)?
         }
-    }
-
-    // Sliding
-    let direction_index = match (dr.signum(), df.signum()) {
-        (1, 0) => 0, // N
-        (1, 1) => 1,
-        (0, 1) => 2,
-        (-1, 1) => 3,
-        (-1, 0) => 4,
-        (-1, -1) => 5,
-        (0, -1) => 6,
-        (1, -1) => 7,
-        _ => panic!("Invalid move delta"),
+    } else {
+        encode_non_underpromo(dx, dy)?
     };
 
-    let distance = dr.abs().max(df.abs()) as usize - 1;
-
-    direction_index * 7 + distance
+    let from_sq = (fr as usize) * 8 + (ff as usize);
+    Some(from_sq * PLANES_PER_SQUARE + plane)
 }
 
-pub fn decode_move(index: usize, side_to_move: Color) -> ChessMove {
-    let from_index = index % 64;
-    let plane = index / 64;
-
-    // Perspective-space coordinates
-    let mut from_rank = from_index / 8;
-    let mut from_file = from_index % 8;
-
-    let (mut dr, mut df, promotion) = decode_move_type(plane);
-
-    // Convert from perspective coordinates back to absolute board coordinates
-    if side_to_move == Color::Black {
-        from_rank = 7 - from_rank;
-        from_file = 7 - from_file;
-
-        dr = -dr;
-        df = -df;
+fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
+    // Knight planes: 56..63
+    if let Some(idx) = KNIGHT_DIRS.iter().position(|&(x, y)| x == dx && y == dy) {
+        return Some(56 + idx);
     }
 
-    let to_rank = (from_rank as i32 + dr) as usize;
-    let to_file = (from_file as i32 + df) as usize;
+    // Sliding planes: 0..55
+    for (dir_idx, &(sx, sy)) in SLIDING_DIRS.iter().enumerate() {
+        for dist in 1..=7 {
+            if dx == sx * dist && dy == sy * dist {
+                return Some(dir_idx * 7 + (dist as usize - 1));
+            }
+        }
+    }
 
-    let from = Square::make_square(Rank::from_index(from_rank), File::from_index(from_file));
-
-    let to = Square::make_square(Rank::from_index(to_rank), File::from_index(to_file));
-
-    ChessMove::new(from, to, promotion)
+    None
 }
 
-fn decode_move_type(plane: usize) -> (i32, i32, Option<Piece>) {
-    // Knight moves
-    const KNIGHT_DELTAS: [(i32, i32); 8] = [
-        (2, 1),
-        (1, 2),
-        (-1, 2),
-        (-2, 1),
-        (-2, -1),
-        (-1, -2),
-        (1, -2),
-        (2, -1),
-    ];
-
-    // 0–55: sliding moves
-    if plane < 56 {
-        let direction = plane / 7;
-        let distance = (plane % 7) + 1;
-
-        let (dr, df) = match direction {
-            0 => (1, 0),
-            1 => (1, 1),
-            2 => (0, 1),
-            3 => (-1, 1),
-            4 => (-1, 0),
-            5 => (-1, -1),
-            6 => (0, -1),
-            7 => (1, -1),
-            _ => unreachable!(),
-        };
-
-        return (dr * distance as i32, df * distance as i32, None);
+/// Decode an index in [0, 4672) back into a ChessMove.
+pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
+    if idx >= 4672 {
+        return None;
     }
 
-    // 56–63: knight moves
-    if plane < 64 {
-        let (dr, df) = KNIGHT_DELTAS[plane - 56];
-        return (dr, df, None);
-    }
+    let from_idx = idx / PLANES_PER_SQUARE;
+    let plane = idx % PLANES_PER_SQUARE;
 
-    // 64–72: underpromotions
-    let promo_plane = plane - 64;
+    let ff = (from_idx % 8) as i8;
+    let fr = (from_idx / 8) as i8;
 
-    let dir = promo_plane / 3;
-    let piece_index = promo_plane % 3;
+    let (dx, dy, promo) = if plane < 56 {
+        let dir = plane / 7;
+        let dist = (plane % 7 + 1) as i8;
 
-    let df = match dir {
-        0 => 0,
-        1 => -1,
-        2 => 1,
-        _ => unreachable!(),
+        let (sx, sy) = SLIDING_DIRS[dir];
+        (sx * dist, sy * dist, None)
+    } else if plane < 64 {
+        let k = plane - 56;
+        let (dx, dy) = KNIGHT_DIRS[k];
+        (dx, dy, None)
+    } else {
+        let p = plane - 64;
+
+        let piece = UNDERPROMO_PIECES[p / 3];
+        let (dx, dy) = UNDERPROMO_DIRS[p % 3];
+
+        (dx, dy, Some(piece))
     };
 
-    let dr = 1; // always forward (important: assumes white perspective)
+    let tf = ff + dx;
+    let tr = fr + dy;
 
-    let promotion = Some(match piece_index {
-        0 => Piece::Rook,
-        1 => Piece::Bishop,
-        2 => Piece::Knight,
-        _ => unreachable!(),
-    });
+    if !(0..8).contains(&tf) || !(0..8).contains(&tr) {
+        return None;
+    }
 
-    (dr, df, promotion)
+    let from = deorient_square(ff, fr, stm);
+    let to = deorient_square(tf, tr, stm);
+
+    Some(ChessMove::new(from, to, promo))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use chess::ALL_COLORS;
+
+    #[test]
+    fn encoding_roundtrips() {
+        for color in ALL_COLORS {
+            for action in 0..4672 {
+                let decoded = decode_move(action, color);
+                if decoded.is_none() {
+                    continue;
+                }
+                let decoded = decoded.unwrap();
+                let encoded = encode_move(decoded, color).unwrap();
+                // eprintln!(
+                //     "orig idx = {}, plane={}, from_idx={}",
+                //     action,
+                //     action / 64,
+                //     action % 64
+                // );
+                // eprintln!(
+                //     "move = {}, from={}, to={}, promo={:?}",
+                //     decoded,
+                //     decoded.get_source(),
+                //     decoded.get_dest(),
+                //     decoded.get_promotion()
+                // );
+                // eprintln!(
+                //     "decoded = {}, from={}, to={}, promo={:?}",
+                //     decoded,
+                //     decoded.get_source(),
+                //     decoded.get_dest(),
+                //     decoded.get_promotion()
+                // );
+                // eprintln!(
+                //     "re-encoded idx = {}, plane={}, from_idx={}",
+                //     encoded,
+                //     encoded / 64,
+                //     encoded % 64
+                // );
+                assert_eq!(action, encoded);
+            }
+        }
+    }
 }

engine/src/net/model.rs

@@ -1,5 +1,5 @@
 use crate::mcts::{BoardState, MctsResults};
-use crate::net::encoding::{encode_board_state_perspective, encode_move};
+use crate::net::encoding::encode_board_state_perspective;
 use burn::data::dataloader::batcher::Batcher;
 use burn::nn::conv::Conv2dConfig;
 use burn::nn::loss::{MseLoss, Reduction};
@@ -13,8 +13,6 @@ use burn::{
     prelude::*,
 };
 use burn_ndarray::NdArray;
-use chess::ChessMove;
-use std::collections::HashMap;
 /*
 Input planes:
 1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -34,6 +32,12 @@ Output:
 65-73: underpromotions
 */
 
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct ModelMetadata {
+    pub(crate) name: String,
+    pub(crate) iterations: usize,
+}
+
 #[derive(Module, Debug)]
 pub struct ResidualBlock<B: Backend> {
     conv1: Conv2d<B>,
@@ -255,14 +259,15 @@ impl<B: Backend> InferenceStep for ChessModel<B> {
 #[derive(Clone)]
 pub struct TrainingSample {
     pub board_state: BoardState,
-    pub policy_target: HashMap<ChessMove, f32>,
+    // Compact representation: list of (encoded_move_index, probability)
+    pub policy_target: Vec<(usize, f32)>,
     pub value_target: f32,
 }
 
 impl TrainingSample {
     pub fn new(
         board_state: BoardState,
-        policy_target: HashMap<ChessMove, f32>,
+        policy_target: Vec<(usize, f32)>,
         value_target: f32,
     ) -> Self {
         TrainingSample {
@@ -273,6 +278,7 @@ impl TrainingSample {
     }
 
     pub fn from_mcts_with_outcome(mcts_results: MctsResults, outcome: f32) -> Self {
+        // move_dist is already a compact Vec<(encoded_move_index, prob)>
         TrainingSample::new(mcts_results.board_state, mcts_results.move_dist, outcome)
     }
 }
@@ -301,9 +307,8 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
             .cloned()
             .map(|item| {
                 let mut policy = vec![0.0f32; 4672];
-                let stm = item.board_state.board.side_to_move();
-                for (mv, prob) in item.policy_target {
-                    policy[encode_move(mv, stm)] = prob;
+                for (idx, prob) in item.policy_target.iter() {
+                    policy[*idx] = *prob;
                 }
 
                 // Normalize

engine/src/training/train.rs

@@ -1,23 +1,29 @@
 use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
-use crate::net::model::{ChessBatcher, ChessModel, ChessModelConfig, TrainingSample};
+use crate::net::encoding::decode_move;
+use crate::net::model::{
+    ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
+};
 use burn::data::dataloader::batcher::Batcher;
 use burn::module::{AutodiffModule, Module};
 use burn::optim::{AdamConfig, GradientsParams, Optimizer};
 use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
 use burn::tensor::backend::AutodiffBackend;
-use chess::ChessMove;
-use rand::rngs::ThreadRng;
+use chess::{ChessMove, Color};
+use rand::rngs::SmallRng;
 use rand::seq::SliceRandom;
-use rand::RngExt;
-use std::collections::{HashMap, VecDeque};
+use rand::{RngExt, SeedableRng};
+use std::collections::VecDeque;
+use std::fs::File;
+use std::io::BufReader;
 use std::marker::PhantomData;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Instant;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 pub struct TrainingConfig {
     pub max_time_s: Option<u64>,
-    pub num_iters: Option<u32>,
+    pub num_iters: Option<usize>,
     pub max_depth: u16, // unused
     pub model_name: String,
     pub load_model: bool,
@@ -29,14 +35,16 @@ pub struct TrainingConfig {
     pub mcts_config: MctsConfig,
     pub optimizer: AdamConfig,
     pub lr: f64,
+    pub seed: Option<u64>,
 }
 
 pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Device) {
     let model_path = format!("artifacts/{}", training_config.model_name.as_str());
+    let metadata_path = format!("{}.json", model_path);
     println!("Creating model...");
     let mut model: ChessModel<B> = ChessModelConfig::init(
-        training_config.hidden_channels,
         training_config.num_blocks,
+        training_config.hidden_channels,
         &device,
     );
     if training_config.load_model {
@@ -51,7 +59,20 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
     let train = Arc::new(AtomicBool::new(true));
     let train_signal = Arc::clone(&train);
 
-    let mut iter: u32 = 0;
+    let mut iter: usize = 0;
+
+    if training_config.load_model {
+        let file = File::open(&metadata_path).unwrap();
+
+        // 2. Use a buffered reader for efficiency
+        let reader = BufReader::new(file);
+
+        // 3. Deserialize JSON directly from the file reader
+        let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
+        iter = metadata.iterations;
+        println!("Loaded model had {} iters", iter);
+    }
+
     let start_time = Instant::now();
 
     ctrlc::set_handler(move || {
@@ -67,10 +88,19 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
         _marker: PhantomData,
     };
 
-    let mut rng = rand::rng();
+    // Create RNG once and reuse it for sampling and shuffling
+    // Seed from system time (platform default entropy may be unavailable in some contexts)
+    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
+    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
+    println!("using seed: {}", seed);
+    let mut rng = SmallRng::seed_from_u64(seed);
+
+    // Initialize optimizer once so state (moments) persist across steps
+    let mut optim = training_config.optimizer.init();
 
     println!("Starting training...");
     while train.load(Ordering::Relaxed) {
+        println!("Iteration: {}", iter);
         let infer_model = model.valid();
         // Gen samples
         println!("Generating {} games...", training_config.num_episodes);
@@ -79,8 +109,12 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
             let mut board_state = BoardState::default();
             let mut episode_buffer: Vec<MctsResults> = vec![];
 
+            let mut game_hist = String::new();
+
             while board_state.status == BoardStateStatus::Ongoing {
+                // let before = Instant::now();
                 let results = mcts.search(&board_state, &infer_model, &device);
+                // println!("{}", before.elapsed().as_millis());
                 episode_buffer.push(results);
 
                 let temp = if board_state.halfmove_clock < 30 {
@@ -91,10 +125,19 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
 
                 let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
 
-                let mv = sample_move(&adjusted, &mut rng).unwrap();
-                println!("playing move: {}", mv);
+                let stm = board_state.board.side_to_move();
+                let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
+                // println!("playing move: {}", mv);
+                if episode == 0 {
+                    game_hist.push_str(mv.to_string().as_str());
+                    game_hist.push(' ');
+                }
+
                 board_state.apply_move(mv)
             }
+            if episode == 0 {
+                println!("Game history of first game of iteration: {}", game_hist);
+            }
 
             for result in episode_buffer.iter().enumerate() {
                 if board_state.status == BoardStateStatus::Stalemate
@@ -140,8 +183,6 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
 
         let batch = batcher.batch(samples, &device);
 
-        let mut optim = training_config.optimizer.init();
-
         let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
 
         let grads = output.loss.backward();
@@ -172,6 +213,19 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
     }
 
     println!("Saving model...");
+
+    let metadata = ModelMetadata {
+        name: training_config.model_name,
+        iterations: iter,
+    };
+
+    std::fs::write(
+        metadata_path,
+        serde_json::to_string_pretty(&metadata)
+            .expect("Should be able to convert metadata to JSON string"),
+    )
+    .expect("Should be able to write metadata");
+
     // Save model in MessagePack format with full precision
     let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
     model
@@ -181,56 +235,58 @@ pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Dev
     return;
 }
 
-fn apply_temperature(
-    visits: &HashMap<ChessMove, f32>,
-    temperature: f32,
-) -> HashMap<ChessMove, f32> {
+fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
     if visits.is_empty() {
-        return HashMap::new();
+        return Vec::new();
     }
 
     // Special case: deterministic selection
     if temperature == 0.0 {
-        let (&best_move, _) = visits
+        let (&best_idx, _) = visits
             .iter()
-            .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
+            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
+            .map(|(i, p)| (i, p))
             .unwrap();
 
-        let mut out = HashMap::new();
-        out.insert(best_move.clone(), 1.0);
-        return out;
+        return vec![(best_idx, 1.0)];
     }
 
     let inv_temp = 1.0 / temperature;
 
     // Step 1: apply exponent
-    let mut adjusted: HashMap<ChessMove, f32> =
-        visits.iter().map(|(m, v)| (*m, v.powf(inv_temp))).collect();
+    let mut adjusted: Vec<(usize, f32)> =
+        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
 
     // Step 2: normalize
-    let sum: f32 = adjusted.values().sum();
+    let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
 
     if sum <= 0.0 {
         return adjusted; // fallback (shouldn't happen in normal MCTS)
     }
 
-    for v in adjusted.values_mut() {
+    for (_, v) in adjusted.iter_mut() {
         *v /= sum;
     }
 
     adjusted
 }
 
-fn sample_move(dist: &HashMap<ChessMove, f32>, rng: &mut ThreadRng) -> Option<ChessMove> {
+fn sample_move(
+    dist: &[(usize, f32)],
+    rng: &mut SmallRng,
+    side_to_move: Color,
+) -> Option<ChessMove> {
     let mut r: f32 = rng.random_range(0.0..1.0);
 
-    for (m, p) in dist {
-        r -= p;
+    for (idx, p) in dist {
+        r -= *p;
         if r <= 0.0 {
-            return Some(m.clone());
+            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
         }
     }
 
     // fallback due to floating point drift
-    dist.keys().next().cloned()
+    dist.get(0)
+        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
 }
+