speed up optim

engine/Cargo.toml

@@ -5,8 +5,7 @@ edition = "2024"
 
 [dependencies]
 chess = "3.2.0"
-burn = { version = "0.21.0", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
-burn-ndarray = "0.21.0"
+burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
 rand_distr = "0.6.0"
 rand = "0.10.1"
 ctrlc = "3.5.2"

engine/src/lib.rs

@@ -1,7 +1,6 @@
 pub mod mcts;
 mod net;
-mod train;
-pub mod training;
+pub mod train;
 
 use chess::Game;
 

engine/src/main.rs

@@ -1,9 +1,9 @@
 #![recursion_limit = "256"]
 
-use burn::backend::{Autodiff, Cuda};
 use burn::optim::AdamConfig;
+use burn::tensor::{Device, DeviceIndex};
 use engine::mcts::MctsConfig;
-use engine::training::train::{train, TrainingConfig};
+use engine::train::{train, TrainingConfig};
 // fn main() {
 //     type MyBackend = Wgpu<f32, i32>;
 //
@@ -14,15 +14,9 @@ use engine::training::train::{train, TrainingConfig};
 // }
 
 fn main() {
-    // 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::ndarray::NdArrayDevice::default();
-    let device = burn::backend::cuda::CudaDevice::default();
+    let device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
+    let mcts_config = MctsConfig::new(150, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -43,7 +37,7 @@ fn main() {
         seed: None,
     };
 
-    train::<MyAutodiffBackend>(training_config, device);
+    train(training_config, device);
 
     // let mut board_state: Option<BoardState> = Some(BoardState::new(
     //     Board::default(),

engine/src/mcts.rs

@@ -1,16 +1,14 @@
 // use crate::networks::PolicyNetwork;
 use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
-use crate::net::encoding::{encode_board_state_perspective, encode_move};
+use crate::net::encoding::{encode_board_states_batch, encode_move};
 use crate::net::model::ChessModel;
-use burn::prelude::Backend;
-use burn::Tensor;
+use burn::prelude::Device;
 use chess::BoardStatus::{Checkmate, Stalemate};
 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;
 
 pub struct Node {
     pub prior: f32,
@@ -38,18 +36,34 @@ 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| {
-                ucb_score(self, &arena[a], c_puct)
-                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
-                    .unwrap()
+                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)
             })
             .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);
+
+    let value_score = if child.visit_count > 0 {
+        -child.value()
+    } else {
+        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())
@@ -107,17 +121,65 @@ impl Default for MctsConfig {
     }
 }
 
-pub struct Mcts<B: Backend> {
+pub struct Mcts {
     pub config: MctsConfig,
-    pub _marker: PhantomData<B>, // if B not otherwise stored
 }
 
-impl<B: Backend> Mcts<B> {
+fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
+    match status {
+        BoardStateStatus::Ongoing => None,
+        BoardStateStatus::Stalemate
+        | 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 }),
+    }
+}
+
+fn softmax(logits: &[f32]) -> Vec<f32> {
+    if logits.is_empty() {
+        return Vec::new();
+    }
+    let mut max_logit = f32::NEG_INFINITY;
+    for &v in logits {
+        if v > max_logit {
+            max_logit = v;
+        }
+    }
+    if !max_logit.is_finite() {
+        max_logit = 0.0;
+    }
+    let mut exps_sum = 0.0f32;
+    let mut probs = vec![0.0f32; logits.len()];
+    for (i, &v) in logits.iter().enumerate() {
+        let x = (v - max_logit).exp();
+        if x.is_finite() {
+            probs[i] = x;
+            exps_sum += x;
+        }
+    }
+    if exps_sum > 0.0 && exps_sum.is_finite() {
+        for p in &mut probs {
+            *p /= exps_sum;
+        }
+    } else {
+        let uniform = 1.0 / logits.len() as f32;
+        for p in &mut probs {
+            *p = uniform;
+        }
+    }
+    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<B>,
-        device: &B::Device,
+        model: &ChessModel,
+        device: &Device,
     ) -> MctsResults {
         let mut nodes = Vec::<Node>::new();
 
@@ -130,134 +192,145 @@ impl<B: Backend> Mcts<B> {
         // Apply Dirichlet noise to root children
         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: 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 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);
                 }
 
-                // Record leaf node and its path
-                leaf_nodes.push(current);
-                leaf_paths.push(path.clone());
+                // 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;
 
-                // 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);
+                // 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;
                 }
 
-            if leaf_nodes.is_empty() {
+                // 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;
                 }
 
-            // Batch evaluate the collected leaf states
-            let batch = Tensor::cat(leaf_states, 0);
+                leaf_nodes.push(current);
+                leaf_paths.push(path);
+            }
+
+            if leaf_nodes.is_empty() {
+                if sims_done >= num_sims {
+                    break;
+                }
+                continue;
+            }
+
+            // 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);
 
-            // 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
+                // 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);
 
-                // 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 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();
-                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 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(prior, new_board, Some(mv)));
+                    nodes.push(Node::new(normalized_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);
+                self.backpropagate(&mut nodes, path, value, stm);
+                sims_done += 1;
             }
         }
 
-        // Build compact move distribution: encoded move index -> probability (visits / num_simulations)
+        // 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 denom = self.config.num_simulations as f32;
+        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");
-            let enc = encode_move(mv, stm).expect("Invalid move");
+            if let Some(enc) = encode_move(mv, stm) {
                 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())
     }
@@ -266,54 +339,51 @@ impl<B: Backend> Mcts<B> {
         &mut self,
         node_idx: usize,
         arena: &mut Vec<Node>,
-        model: &ChessModel<B>,
-        device: &B::Device,
+        model: &ChessModel,
+        device: &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
+        let status = arena[node_idx].board_state.status;
+        let stm = arena[node_idx].board_state.board.side_to_move();
+
+        if let Some(term_val) = terminal_value(status, stm) {
+            term_val
         } 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]);
-            // let start = Instant::now();
+            let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
             let (policy_head, value_head) = model.forward(state);
-            // println!("time: {:?}", start.elapsed());
 
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
 
-            let policy = policy_head.into_data().to_vec::<f32>().unwrap();
+            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
+            let probs = softmax(&logits);
 
-            for mv in legal_moves {
-                let stm = arena[node_idx].board_state.board.side_to_move();
-                let idx = encode_move(mv, stm).expect("Invalid move");
-                let prior = policy[idx];
+            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 = arena[node_idx].board_state.clone();
                 new_board.apply_move(mv);
 
                 let child_idx = arena.len();
-
-                arena.push(Node::new(prior, new_board, Some(mv)));
+                arena.push(Node::new(normalized_prior, new_board, Some(mv)));
                 arena[node_idx].children.push(child_idx);
             }
 
-            value_head.into_data().to_vec().unwrap()[0]
+            value_head.into_data().to_vec::<f32>().unwrap()[0]
         }
     }
 
@@ -356,7 +426,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
 
-    // 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();
@@ -367,26 +436,21 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let sum: f32 = samples.iter().sum();
 
+    if sum > 0.0 && sum.is_finite() {
         for x in &mut samples {
             *x /= sum;
         }
+    } else if size > 0 {
+        let uniform = 1.0 / size as f32;
+        for x in &mut samples {
+            *x = uniform;
+        }
+    }
 
     samples
 }
 
-fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
-    let prior_score = c_puct * child.prior * (parent.visit_count as f32).sqrt()
-        / (1.0 + child.visit_count as f32);
 
-    let value_score: f32;
-    if child.visit_count > 0 {
-        value_score = -child.value(); // value from opposing side
-    } else {
-        value_score = 0.0;
-    }
-
-    value_score + prior_score
-}
 
 pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
     // Choose best child of root by visit count (standard MCTS behavior)

engine/src/net/encoding.rs

@@ -1,26 +1,6 @@
 use crate::mcts::BoardState;
 use burn::prelude::*;
-use chess::{ChessMove, Color, File, Piece, Rank, Square, ALL_SQUARES};
-/*
-Input planes:
-1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
-7-12: opponent pieces
-13: your kingside
-14: your queenside
-15: opponent kingside
-16: opponent queenside
-17: rep >= 2
-18: En passant square (gap pass over)
-
-
-Output:
-4672 moves represented through planes
-1-56: sliding moves
-57-64: Knight moves
-65-73: underpromotions
-*/
-
-const PLANES_PER_SQUARE: usize = 73;
+use chess::{ChessMove, Color, File, Piece, Rank, Square};
 
 const SLIDING_DIRS: [(i8, i8); 8] = [
     (0, 1),   // N
@@ -52,104 +32,106 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
 
 const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
 
-pub fn encode_board_state_perspective<B: Backend>(
-    state: &BoardState,
-    device: &B::Device,
-) -> Tensor<B, 3> {
+pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
     let mut planes = vec![0.0f32; 18 * 64];
+    encode_board_state_to_slice(state, &mut planes);
+    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+}
+
+/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
+pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
+    out_slice[..1152].fill(0.0);
 
     let board = &state.board;
     let us = board.side_to_move();
     let them = !us;
+    let flip = us == Color::Black;
 
     let idx = |plane: usize, rank: usize, file: usize| -> usize { plane * 64 + rank * 8 + file };
 
-    let flip = us == Color::Black;
-
-    for &square in ALL_SQUARES.iter() {
-        if let Some(piece) = board.piece_on(square) {
-            let color = board.color_on(square).unwrap();
-
-            let piece_index = match piece {
-                Piece::Pawn => 0,
-                Piece::Knight => 1,
-                Piece::Bishop => 2,
-                Piece::Rook => 3,
-                Piece::Queen => 4,
-                Piece::King => 5,
-            };
-
-            // Determine plane based on perspective
-            let plane = if color == us {
-                piece_index
-            } else {
-                piece_index + 6
-            };
+    let pieces = [
+        Piece::Pawn,
+        Piece::Knight,
+        Piece::Bishop,
+        Piece::Rook,
+        Piece::Queen,
+        Piece::King,
+    ];
 
+    for (piece_index, &piece) in pieces.iter().enumerate() {
+        let us_bb = board.pieces(piece) & board.color_combined(us);
+        for square in us_bb {
             let mut rank = square.get_rank().to_index();
             let mut file = square.get_file().to_index();
-
-            // Flip board if black to move
             if flip {
                 rank = 7 - rank;
                 file = 7 - file;
             }
+            out_slice[idx(piece_index, rank, file)] = 1.0;
+        }
 
-            planes[idx(plane, rank, file)] = 1.0;
+        let them_bb = board.pieces(piece) & board.color_combined(them);
+        for square in them_bb {
+            let mut rank = square.get_rank().to_index();
+            let mut file = square.get_file().to_index();
+            if flip {
+                rank = 7 - rank;
+                file = 7 - file;
+            }
+            out_slice[idx(piece_index + 6, rank, file)] = 1.0;
         }
     }
 
-    // -------------------------
     // Castling Rights (planes 12–15)
-    // From perspective
-    // -------------------------
     let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
 
     if us_castle.has_kingside() {
-        fill_plane(&mut planes, 12);
+        fill_plane(out_slice, 12);
     }
     if us_castle.has_queenside() {
-        fill_plane(&mut planes, 13);
+        fill_plane(out_slice, 13);
     }
     if them_castle.has_kingside() {
-        fill_plane(&mut planes, 14);
+        fill_plane(out_slice, 14);
     }
     if them_castle.has_queenside() {
-        fill_plane(&mut planes, 15);
+        fill_plane(out_slice, 15);
     }
 
-    // -------------------------
     // Repetition plane (16)
-    // -------------------------
     let current_hash = board.get_hash();
     if let Some(count) = state.repetition_table.get(&current_hash) {
         if *count >= 2 {
-            fill_plane(&mut planes, 16);
+            fill_plane(out_slice, 16);
         }
     }
 
-    let en_passant = board.en_passant();
-    if let Some(ep) = en_passant {
+    // En passant plane (17)
+    if let Some(ep) = board.en_passant() {
         let mut rank = ep.get_rank().to_index();
         let mut file = ep.get_file().to_index();
-
-        // Flip board if black to move
         if flip {
             rank = 7 - rank;
             file = 7 - file;
         }
-
-        planes[idx(17, rank, file)] = 1.0;
+        out_slice[idx(17, rank, file)] = 1.0;
     }
-
-    Tensor::<B, 3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
 }
 
+pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
+    let b = states.len();
+    let mut buffer = vec![0.0f32; b * 1152];
+    for (i, state) in states.iter().enumerate() {
+        let start = i * 1152;
+        encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
+    }
+    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
+}
+
+#[inline]
 fn fill_plane(buffer: &mut [f32], plane: usize) {
     let start = plane * 64;
-    for i in 0..64 {
-        buffer[start + i] = 1.0;
-    }
+    buffer[start..start + 64].fill(1.0);
 }
 
 /// Rotate a square into the side-to-move perspective.
@@ -201,7 +183,7 @@ pub fn encode_move(mv: ChessMove, stm: Color) -> Option<usize> {
     };
 
     let from_sq = (fr as usize) * 8 + (ff as usize);
-    Some(from_sq * PLANES_PER_SQUARE + plane)
+    Some(plane * 64 + from_sq)
 }
 
 fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
@@ -228,13 +210,13 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
-    let from_idx = idx / PLANES_PER_SQUARE;
-    let plane = idx % PLANES_PER_SQUARE;
+    let plane = idx / 64;
+    let from_idx = idx % 64;
 
     let ff = (from_idx % 8) as i8;
     let fr = (from_idx / 8) as i8;
 
-    let (dx, dy, promo) = if plane < 56 {
+    let (dx, dy, mut promo) = if plane < 56 {
         let dir = plane / 7;
         let dist = (plane % 7 + 1) as i8;
 
@@ -260,6 +242,11 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
+    // Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
+    if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
+        promo = Some(Piece::Queen);
+    }
+
     let from = deorient_square(ff, fr, stm);
     let to = deorient_square(tf, tr, stm);
 
@@ -311,4 +298,26 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn queen_promotion_roundtrip() {
+        use chess::{File, Rank, Square};
+        let from = Square::make_square(Rank::Seventh, File::E);
+        let to = Square::make_square(Rank::Eighth, File::E);
+        let mv = ChessMove::new(from, to, Some(Piece::Queen));
+        let enc = encode_move(mv, Color::White).unwrap();
+        let dec = decode_move(enc, Color::White).unwrap();
+        assert_eq!(mv, dec);
+    }
+
+    #[test]
+    fn underpromotion_roundtrip() {
+        use chess::{File, Rank, Square};
+        let from = Square::make_square(Rank::Seventh, File::E);
+        let to = Square::make_square(Rank::Eighth, File::E);
+        let mv = ChessMove::new(from, to, Some(Piece::Knight));
+        let enc = encode_move(mv, Color::White).unwrap();
+        let dec = decode_move(enc, Color::White).unwrap();
+        assert_eq!(mv, dec);
+    }
 }

engine/src/net/model.rs

@@ -5,14 +5,12 @@ use burn::nn::conv::Conv2dConfig;
 use burn::nn::loss::{MseLoss, Reduction};
 use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
 use burn::tensor::activation::log_softmax;
-use burn::tensor::backend::AutodiffBackend;
 use burn::tensor::Transaction;
 use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
 use burn::{
     nn::{conv::Conv2d, Linear, Relu},
     prelude::*,
 };
-use burn_ndarray::NdArray;
 /*
 Input planes:
 1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -39,11 +37,11 @@ pub struct ModelMetadata {
 }
 
 #[derive(Module, Debug)]
-pub struct ResidualBlock<B: Backend> {
-    conv1: Conv2d<B>,
-    bn1: BatchNorm<B>,
-    conv2: Conv2d<B>,
-    bn2: BatchNorm<B>,
+pub struct ResidualBlock {
+    conv1: Conv2d,
+    bn1: BatchNorm,
+    conv2: Conv2d,
+    bn2: BatchNorm,
     activation: Relu,
 }
 
@@ -52,8 +50,8 @@ pub struct ResidualBlockConfig {
     channels: usize,
 }
 
-impl<B: Backend> ResidualBlock<B> {
-    pub fn new(channels: usize, device: &B::Device) -> Self {
+impl ResidualBlock {
+    pub fn new(channels: usize, device: &Device) -> Self {
         Self {
             conv1: Conv2dConfig::new([channels, channels], [3, 3])
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -66,7 +64,7 @@ impl<B: Backend> ResidualBlock<B> {
             activation: Relu::new(),
         }
     }
-    pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
+    pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
         let residual = x.clone();
 
         let x = self.conv1.forward(x);
@@ -81,24 +79,24 @@ impl<B: Backend> ResidualBlock<B> {
 }
 
 #[derive(Module, Debug)]
-pub struct ChessModel<B: Backend> {
+pub struct ChessModel {
     // Initial convolution
-    conv: Conv2d<B>,
-    bn: BatchNorm<B>,
+    conv: Conv2d,
+    bn: BatchNorm,
 
     // Residual tower
-    residual_blocks: Vec<ResidualBlock<B>>,
+    residual_blocks: Vec<ResidualBlock>,
 
     // Policy head
-    policy_conv: Conv2d<B>,
-    policy_bn: BatchNorm<B>,
-    policy_fc: Linear<B>,
+    policy_conv: Conv2d,
+    policy_bn: BatchNorm,
+    policy_conv2: Conv2d,
 
     // Value head
-    value_conv: Conv2d<B>,
-    value_bn: BatchNorm<B>,
-    value_fc1: Linear<B>,
-    value_fc2: Linear<B>,
+    value_conv: Conv2d,
+    value_bn: BatchNorm,
+    value_fc1: Linear,
+    value_fc2: Linear,
 
     activation: Relu,
 }
@@ -110,11 +108,7 @@ pub struct ChessModelConfig {
 }
 
 impl ChessModelConfig {
-    pub fn init<B: Backend>(
-        num_blocks: usize,
-        channels: usize,
-        device: &B::Device,
-    ) -> ChessModel<B> {
+    pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
         ChessModel {
             conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -126,14 +120,14 @@ impl ChessModelConfig {
                 .collect(),
 
             // Policy head
-            policy_conv: Conv2dConfig::new([channels, 2], [1, 1]).init(device),
-            policy_bn: BatchNormConfig::new(2).init(device),
-            policy_fc: LinearConfig::new(2 * 8 * 8, 8 * 8 * 73).init(device), // 4672 typical chess move space
+            policy_conv: Conv2dConfig::new([channels, 32], [1, 1]).init(device),
+            policy_bn: BatchNormConfig::new(32).init(device),
+            policy_conv2: Conv2dConfig::new([32, 73], [1, 1]).init(device),
 
             // Value head
             value_conv: Conv2dConfig::new([channels, 1], [1, 1]).init(device),
             value_bn: BatchNormConfig::new(1).init(device),
-            value_fc1: LinearConfig::new(1 * 8 * 8, 256).init(device),
+            value_fc1: LinearConfig::new(8 * 8, 256).init(device),
             value_fc2: LinearConfig::new(256, 1).init(device),
 
             activation: Relu::new(),
@@ -141,8 +135,8 @@ impl ChessModelConfig {
     }
 }
 
-impl<B: Backend> ChessModel<B> {
-    pub fn forward(&self, x: Tensor<B, 4>) -> (Tensor<B, 2>, Tensor<B, 2>) {
+impl ChessModel {
+    pub fn forward(&self, x: Tensor<4>) -> (Tensor<2>, Tensor<2>) {
         let mut x = self.conv.forward(x);
         x = self.bn.forward(x);
         x = self.activation.forward(x);
@@ -157,8 +151,8 @@ impl<B: Backend> ChessModel<B> {
         let mut p = self.policy_conv.forward(x.clone());
         p = self.policy_bn.forward(p);
         p = self.activation.forward(p);
-        let mut p = p.reshape([batch_size, 2 * 8 * 8]);
-        p = self.policy_fc.forward(p);
+        p = self.policy_conv2.forward(p);
+        let p = p.reshape([batch_size, 73 * 8 * 8]);
 
         // -------- Value Head --------
         let mut v = self.value_conv.forward(x);
@@ -173,10 +167,10 @@ impl<B: Backend> ChessModel<B> {
 
     pub fn forward_chess(
         &self,
-        boards: Tensor<B, 4>,         // e.g. [batch, channels, 8, 8]
-        policy_targets: Tensor<B, 2>, // move distribution (make sure is normalized)
-        value_targets: Tensor<B, 2>,  // scalar evaluation
-    ) -> ChessOutput<B> {
+        boards: Tensor<4>,         // e.g. [batch, channels, 8, 8]
+        policy_targets: Tensor<2>, // move distribution (make sure is normalized)
+        value_targets: Tensor<2>,  // scalar evaluation
+    ) -> ChessOutput {
         let (policy_logits, value) = self.forward(boards);
 
         let log_probs = log_softmax(policy_logits.clone(), 1);
@@ -190,35 +184,40 @@ impl<B: Backend> ChessModel<B> {
         let value_loss =
             MseLoss::new().forward(value.clone(), value_targets.clone(), Reduction::Mean);
 
-        let total_loss = policy_loss + 0.5 * value_loss;
+        let total_loss = policy_loss.clone() + 0.5 * value_loss.clone();
 
         ChessOutput {
             policy_logits,
             value,
             policy_targets,
             value_targets,
+            policy_loss,
+            value_loss,
             loss: total_loss,
         }
     }
 }
 
-pub struct ChessOutput<B: Backend> {
-    pub policy_logits: Tensor<B, 2>, // [sample, num_moves (4672)]
-    pub value: Tensor<B, 2>,         // [sample, 1 value]
-    pub policy_targets: Tensor<B, 2>,
-    pub value_targets: Tensor<B, 2>, // kept 2d for consistency?
-    pub loss: Tensor<B, 1>,          // [sample]
+pub struct ChessOutput {
+    pub policy_logits: Tensor<2>, // [sample, num_moves (4672)]
+    pub value: Tensor<2>,         // [sample, 1 value]
+    pub policy_targets: Tensor<2>,
+    pub value_targets: Tensor<2>, // kept 2d for consistency?
+    pub policy_loss: Tensor<1>,
+    pub value_loss: Tensor<1>,
+    pub loss: Tensor<1>,          // [sample]
 }
 
-impl<B: Backend> ItemLazy for ChessOutput<B> {
-    type ItemSync = ChessOutput<NdArray>;
-
-    fn sync(self) -> Self::ItemSync {
-        let [policy_logits, value, policy_targets, value_targets, loss] = Transaction::default()
+impl ItemLazy for ChessOutput {
+    fn sync(self) -> Self {
+        let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
+            Transaction::default()
                 .register(self.policy_logits)
                 .register(self.value)
                 .register(self.policy_targets)
                 .register(self.value_targets)
+                .register(self.policy_loss)
+                .register(self.value_loss)
                 .register(self.loss)
                 .execute()
                 .try_into()
@@ -231,27 +230,29 @@ impl<B: Backend> ItemLazy for ChessOutput<B> {
             value: Tensor::from_data(value, device),
             policy_targets: Tensor::from_data(policy_targets, device),
             value_targets: Tensor::from_data(value_targets, device),
+            policy_loss: Tensor::from_data(policy_loss, device),
+            value_loss: Tensor::from_data(value_loss, device),
             loss: Tensor::from_data(loss, device),
         }
     }
 }
 
-impl<B: AutodiffBackend> TrainStep for ChessModel<B> {
-    type Input = ChessBatch<B>;
-    type Output = ChessOutput<B>;
+impl TrainStep for ChessModel {
+    type Input = ChessBatch;
+    type Output = ChessOutput;
 
-    fn step(&self, batch: ChessBatch<B>) -> TrainOutput<ChessOutput<B>> {
+    fn step(&self, batch: ChessBatch) -> TrainOutput<ChessOutput> {
         let item = self.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
 
         TrainOutput::new(self, item.loss.backward(), item)
     }
 }
 
-impl<B: Backend> InferenceStep for ChessModel<B> {
-    type Input = ChessBatch<B>;
-    type Output = ChessOutput<B>;
+impl InferenceStep for ChessModel {
+    type Input = ChessBatch;
+    type Output = ChessOutput;
 
-    fn step(&self, batch: ChessBatch<B>) -> ChessOutput<B> {
+    fn step(&self, batch: ChessBatch) -> ChessOutput {
         self.forward_chess(batch.states, batch.policy_targets, batch.value_targets)
     }
 }
@@ -287,14 +288,14 @@ impl TrainingSample {
 pub struct ChessBatcher {}
 
 #[derive(Clone, Debug)]
-pub struct ChessBatch<B: Backend> {
-    pub states: Tensor<B, 4>,
-    pub policy_targets: Tensor<B, 2>,
-    pub value_targets: Tensor<B, 2>,
+pub struct ChessBatch {
+    pub states: Tensor<4>,
+    pub policy_targets: Tensor<2>,
+    pub value_targets: Tensor<2>,
 }
 
-impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
-    fn batch(&self, items: Vec<TrainingSample>, device: &B::Device) -> ChessBatch<B> {
+impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
+    fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
         let state_tensors = items
             .iter()
             .map(|item| {
@@ -318,17 +319,14 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
                         *p /= sum;
                     }
                 }
-                Tensor::<B, 2>::from_floats(TensorData::new(policy, [1, 4672]), device).unsqueeze()
+                Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
             })
             .collect::<Vec<_>>();
 
         let value_target_tensors = items
             .iter()
             .map(|item| {
-                Tensor::<B, 2>::from_floats(
-                    TensorData::new(vec![item.value_target], [1, 1]),
-                    device,
-                )
+                Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device)
                     .reshape([1, 1])
             })
             .collect::<Vec<_>>();

engine/src/rl/classification.rs

@@ -1,59 +0,0 @@
-/// Multi-label classification output adapted for multiple metrics.
-///
-/// Supported metrics:
-/// - HammingScore
-/// - Precision (via ConfusionStatsInput)
-/// - Recall (via ConfusionStatsInput)
-/// - FBetaScore (via ConfusionStatsInput)
-/// - Loss
-#[derive(new)]
-pub struct MultiLabelSoftClassificationOutput<B: Backend> {
-    /// The loss.
-    pub loss: Tensor<B, 1>,
-
-    /// The label logits or probabilities. Shape: \[batch_size, num_classes\].
-    pub output: Tensor<B, 2>,
-
-    /// The ground truth labels (target values). Shape: \[batch_size, num_classes\].
-    pub targets: Tensor<B, 2>,
-}
-
-impl<B: Backend> ItemLazy for MultiLabelSoftClassificationOutput<B> {
-    type ItemSync = MultiLabelSoftClassificationOutput<NdArray>;
-
-    fn sync(self) -> Self::ItemSync {
-        let [output, loss, targets] = Transaction::default()
-            .register(self.output)
-            .register(self.loss)
-            .register(self.targets)
-            .execute()
-            .try_into()
-            .expect("Correct amount of tensor data");
-
-        let device = &Default::default();
-
-        MultiLabelSoftClassificationOutput {
-            output: Tensor::from_data(output, device),
-            loss: Tensor::from_data(loss, device),
-            targets: Tensor::from_data(targets, device),
-        }
-    }
-}
-
-impl<B: Backend> Adaptor<HammingScoreInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> HammingScoreInput<B> {
-        HammingScoreInput::new(self.output.clone(), self.targets.clone())
-    }
-}
-
-impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> LossInput<B> {
-        LossInput::new(self.loss.clone())
-    }
-}
-
-impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for MultiLabelSoftClassificationOutput<B> {
-    fn adapt(&self) -> ConfusionStatsInput<B> {
-        ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool())
-    }
-}

engine/src/train.rs

@@ -1,134 +1,332 @@
-// use crate::model::{ChessBatcher, ChessModel, ChessModelConfig};
-// use burn::config::Config;
-// use burn::data::dataloader::batcher::Batcher;
-// use burn::optim::{Adam, AdamConfig, Optimizer, SimpleOptimizer};
-// use burn::optim::adaptor::OptimizerAdaptor;
-// use burn::prelude::Backend;
-// use burn::tensor::backend::AutodiffBackend;
-// use burn::train::TrainStep;
-// use crate::mcts::MctsResults;
-//
-// #[derive(Config, Debug)]
-// pub struct ChessTrainerConfig {
-//     pub model: ChessModelConfig,
-//     pub optimizer: AdamConfig,
-//     #[config(default = 10)]
-//     pub num_epochs: usize,
-//     #[config(default = 64)]
-//     pub batch_size: usize,
-//     #[config(default = 4)]
-//     pub num_workers: usize,
-//     #[config(default = 42)]
-//     pub seed: u64,
-//     #[config(default = 1.0e-4)]
-//     pub learning_rate: f64,
-// }
-//
-// impl ChessTrainerConfig {
-//     pub fn init<B: Backend>(
-//         model_config: ChessModelConfig,
-//         optimizer: AdamConfig,
-//         num_epochs: usize,
-//         batch_size: usize,
-//         num_workers: usize,
-//         seed: u64,
-//         learning_rate: f64,
-//     ) -> ChessTrainer<B> {
-//         ChessTrainer {
-//             model: model_config::init(),
-//             optimizer: optimizer.init(),
-//             num_epochs,
-//
-//
-//         }
-//     }
-// }
-//
-// pub struct ChessTrainer<B: AutodiffBackend> {
-//     pub model: ChessModel<B>,
-//     pub optimizer: Adam,
-//     learning_rate: f64,
-//     pub batcher: ChessBatcher,
-//     pub device: B::Device,
-// }
-//
-// impl<B: AutodiffBackend> ChessTrainer<B> {
-//     pub fn new(model: ChessModel<B>, device: B::Device) -> Self {
-//         let optimizer = AdamConfig::new().init();
-//
-//         Self {
-//             model,
-//             optimizer,
-//             batcher: ChessBatcher::default(),
-//             device,
-//         }
-//     }
-//
-//     pub fn train_step(&mut self, batch_data: Vec<MctsResults>) -> f32 {
-//         // 1. Convert to tensors
-//         let batch = self.batcher.batch(batch_data, &self.device);
-//
-//         // 2. Forward + backward (your TrainStep impl)
-//         let output = self.model.step(batch);
-//
-//         // 3. Update weights
-//         self.optimizer.step(, &mut self.model, output.grads);
-//
-//         // 4. Return loss (for logging)
-//         let loss_tensor = output.item.loss.clone().into_data();
-//         let loss = loss_tensor.to_vec::<f32>().unwrap()[0];
-//
-//         loss
-//     }
-// }
+use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
+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;
+use burn::optim::{AdamConfig, GradientsParams};
+use burn::prelude::{Device, Module};
+use burn::store::ModuleRecord;
+use chess::{ChessMove, Color};
+use rand::rngs::SmallRng;
+use rand::seq::SliceRandom;
+use rand::{RngExt, SeedableRng};
+use std::collections::VecDeque;
+use std::fs::File;
+use std::io::BufReader;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::time::Instant;
+use std::time::{SystemTime, UNIX_EPOCH};
 
-// fn create_artifact_dir(artifact_dir: &str) {
-//     // Remove existing artifacts before to get an accurate learner summary
-//     std::fs::remove_dir_all(artifact_dir).ok();
-//     std::fs::create_dir_all(artifact_dir).ok();
-// }
-//
-// pub fn train<B: AutodiffBackend>(
-//     artifact_dir: &str,
-//     config: ChessTrainingConfig,
-//     device: B::Device,
-// ) {
-//     create_artifact_dir(artifact_dir);
-//     config
-//         .save(format!("{artifact_dir}/config.json"))
-//         .expect("Config should be saved successfully");
-//
-//     B::seed(&device, config.seed);
-//
-//     let batcher = ChessBatcher::default();
+pub struct TrainingConfig {
+    pub max_time_s: Option<u64>,
+    pub num_iters: Option<usize>,
+    pub max_depth: u16, // unused
+    pub model_name: String,
+    pub load_model: bool,
+    pub hidden_channels: usize,
+    pub num_blocks: usize,
+    pub batch_size: usize,      // num positions sampled to a batch (2048)
+    pub num_episodes: usize,    // num games generated per iteration (5000)
+    pub buffer_max_size: usize, // max number of samples in the buffer
+    pub mcts_config: MctsConfig,
+    pub optimizer: AdamConfig,
+    pub lr: f64,
+    pub seed: Option<u64>,
+}
 
-// let dataloader_train = DataLoaderBuilder::new(batcher.clone())
-//     .batch_size(config.batch_size)
-//     .shuffle(config.seed)
-//     .num_workers(config.num_workers)
-//     .build(MnistDataset::train());
-//
-// let dataloader_test = DataLoaderBuilder::new(batcher)
-//     .batch_size(config.batch_size)
-//     .shuffle(config.seed)
-//     .num_workers(config.num_workers)
-//     .build(MnistDataset::test());
-//
-// let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test)
-//     .metrics((AccuracyMetric::new(), LossMetric::new()))
-//     .with_file_checkpointer(CompactRecorder::new())
-//     .num_epochs(config.num_epochs)
-//     .summary();
-//
-// let model = config.model.init::<B>(&device);
-// let result = training.launch(Learner::new(
-//     model,
-//     config.optimizer.init(),
-//     config.learning_rate,
-// ));
-//
-// result
-//     .model
-//     .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
-//     .expect("Trained model should be saved successfully");
-// }
+pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
+    let device = device.into();
+    if training_config.seed.is_some() {
+        device.seed(training_config.seed.unwrap());
+    }
+    let autodiff_device = device.clone().autodiff();
+
+    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 = ChessModelConfig::init(
+        training_config.num_blocks,
+        training_config.hidden_channels,
+        &autodiff_device,
+    );
+    // Initialize optimizer once so state (moments) persist across steps
+    let mut optim = training_config.optimizer.init();
+
+    if training_config.load_model {
+        println!("Loading model {}...", model_path);
+        // Load model in full precision from MessagePack file
+        let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
+        model = model.load_record(record);
+        optim = optim
+            .load(format!("{model_path}_optim"))
+            .expect("Failed to load optimizer");
+    }
+
+    let train = Arc::new(AtomicBool::new(true));
+    let train_signal = Arc::clone(&train);
+
+    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 || {
+        train_signal.store(false, Ordering::Relaxed);
+        println!("Finishing batch before exiting...");
+    })
+    .expect("Error setting Ctrl-C handler");
+
+    let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
+
+    let mut mcts: Mcts = Mcts {
+        config: training_config.mcts_config,
+    };
+
+    // 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);
+
+    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);
+        for episode in 0..training_config.num_episodes {
+            println!("Episode: {}", episode);
+            let mut board_state = BoardState::default();
+            let mut episode_buffer: Vec<MctsResults> = vec![];
+
+            let mut game_hist = String::new();
+
+            let mut ply_count: usize = 0;
+
+            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 ply_count < 30 { 1.0 } else { 0.0 };
+
+                let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
+
+                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);
+                ply_count += 1;
+                if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
+                    board_state.status = BoardStateStatus::FiftyMove;
+                }
+            }
+            if episode == 0 {
+                println!("Game history of first game of iteration: {}", game_hist);
+            }
+
+            for result in episode_buffer.iter() {
+                let stm = result.board_state.board.side_to_move();
+                let outcome = match board_state.status {
+                    BoardStateStatus::WhiteWinner => {
+                        if stm == Color::White {
+                            1.0
+                        } else {
+                            -1.0
+                        }
+                    }
+                    BoardStateStatus::BlackWinner => {
+                        if stm == Color::Black {
+                            1.0
+                        } else {
+                            -1.0
+                        }
+                    }
+                    _ => 0.0,
+                };
+
+                replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
+                    result.clone(),
+                    outcome,
+                ));
+
+                if replay_buffer.len() > training_config.buffer_max_size {
+                    replay_buffer.pop_front();
+                }
+            }
+        }
+        // train
+
+        if replay_buffer.is_empty() {
+            println!("Replay buffer is empty, skipping training step.");
+            continue;
+        }
+
+        let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
+        println!(
+            "Finished self-play! Training on {} samples (buffer size: {})...",
+            num_samples,
+            replay_buffer.len()
+        );
+        let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
+        indices.shuffle(&mut rng);
+
+        let samples = indices
+            .into_iter()
+            .take(num_samples)
+            .map(|i| replay_buffer[i].clone())
+            .collect::<Vec<TrainingSample>>();
+
+        let batcher = ChessBatcher {};
+
+        let batch = batcher.batch(samples, &autodiff_device);
+
+        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
+
+        let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+        let p_loss_val = output
+            .policy_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
+        let v_loss_val = output
+            .value_loss
+            .clone()
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
+
+        println!(
+            "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
+            iter, loss_val, p_loss_val, v_loss_val
+        );
+
+        let grads = output.loss.backward();
+        let grads = GradientsParams::from_grads(grads, &model);
+
+        model = optim.step(training_config.lr, model, grads);
+
+        iter += 1;
+
+        if iter % 100 == 0 {
+            println!("Completed {} iterations", iter);
+        }
+
+        if training_config.max_time_s.is_some()
+            && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
+        {
+            println!("Training stopping due to time limit...");
+            break;
+        }
+
+        if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
+            println!(
+                "Training stopping due to iteration limit ({} iters completed)...",
+                iter
+            );
+            break;
+        }
+    }
+
+    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");
+
+    model
+        .into_record()
+        .save(&model_path)
+        .expect("Should be able to save the model");
+
+    optim
+        .save(format!("{model_path}_optim"))
+        .expect("Should be able to save the optim");
+    println!("Model saved in {:?}, exiting training.", model_path);
+    return;
+}
+
+fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
+    if visits.is_empty() {
+        return Vec::new();
+    }
+
+    // Special case: deterministic selection
+    if temperature == 0.0 {
+        let (&best_idx, _) = visits
+            .iter()
+            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
+            .map(|(i, p)| (i, p))
+            .unwrap();
+
+        return vec![(best_idx, 1.0)];
+    }
+
+    let inv_temp = 1.0 / temperature;
+
+    // Step 1: apply exponent
+    let mut adjusted: Vec<(usize, f32)> =
+        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
+
+    // Step 2: normalize
+    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.iter_mut() {
+        *v /= sum;
+    }
+
+    adjusted
+}
+
+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 (idx, p) in dist {
+        r -= *p;
+        if r <= 0.0 {
+            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
+        }
+    }
+
+    // fallback due to floating point drift
+    dist.get(0)
+        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
+}

engine/src/training.rs

@@ -1 +0,0 @@
-pub mod train;

engine/src/training/train.rs

@@ -1,292 +0,0 @@
-use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
-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, Color};
-use rand::rngs::SmallRng;
-use rand::seq::SliceRandom;
-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<usize>,
-    pub max_depth: u16, // unused
-    pub model_name: String,
-    pub load_model: bool,
-    pub hidden_channels: usize,
-    pub num_blocks: usize,
-    pub batch_size: usize,      // num positions sampled to a batch (2048)
-    pub num_episodes: usize,    // num games generated per iteration (5000)
-    pub buffer_max_size: usize, // max number of samples in the buffer
-    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.num_blocks,
-        training_config.hidden_channels,
-        &device,
-    );
-    if training_config.load_model {
-        println!("Loading model {}...", model_path);
-        // Load model in full precision from MessagePack file
-        let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
-        model = model
-            .load_file(&model_path, &recorder, &device)
-            .expect("Should be able to load the model weights from the provided file");
-    }
-
-    let train = Arc::new(AtomicBool::new(true));
-    let train_signal = Arc::clone(&train);
-
-    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 || {
-        train_signal.store(false, Ordering::Relaxed);
-        println!("Finishing batch before exiting...");
-    })
-    .expect("Error setting Ctrl-C handler");
-
-    let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
-
-    let mut mcts: Mcts<B::InnerBackend> = Mcts {
-        config: training_config.mcts_config,
-        _marker: PhantomData,
-    };
-
-    // 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);
-        for episode in 0..training_config.num_episodes {
-            println!("Episode: {}", episode);
-            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 {
-                    1.0
-                } else {
-                    0.0
-                };
-
-                let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
-
-                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
-                    || board_state.status == BoardStateStatus::Threefold
-                    || board_state.status == BoardStateStatus::FiftyMove
-                {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        0.0,
-                    ));
-                } else if board_state.status == BoardStateStatus::WhiteWinner {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        ((result.0 % 2) as f32 * -2.0) + 1.0,
-                    ));
-                } else if board_state.status == BoardStateStatus::BlackWinner {
-                    replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                        result.1.clone(),
-                        ((result.0 % 2) as f32 * 2.0) - 1.0,
-                    ));
-                }
-                if replay_buffer.len() > training_config.buffer_max_size {
-                    replay_buffer.pop_front();
-                }
-            }
-        }
-        // train
-
-        println!(
-            "Finished! Training on {} samples...",
-            training_config.batch_size
-        );
-        let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
-        indices.shuffle(&mut rng);
-
-        let samples = indices
-            .into_iter()
-            .take(training_config.batch_size)
-            .map(|i| replay_buffer[i].clone())
-            .collect::<Vec<TrainingSample>>();
-
-        let batcher = ChessBatcher {};
-
-        let batch = batcher.batch(samples, &device);
-
-        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
-
-        let grads = output.loss.backward();
-        let grads = GradientsParams::from_grads(grads, &model);
-
-        model = optim.step(training_config.lr, model, grads);
-
-        iter += 1;
-
-        if iter % 100 == 0 {
-            println!("Completed {} iterations", iter);
-        }
-
-        if training_config.max_time_s.is_some()
-            && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
-        {
-            println!("Training stopping due to time limit...");
-            break;
-        }
-
-        if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
-            println!(
-                "Training stopping due to iteration limit ({} iters completed)...",
-                iter
-            );
-            break;
-        }
-    }
-
-    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
-        .save_file(&model_path, &recorder)
-        .expect("Should be able to save the model");
-    println!("Model saved in {:?}, exiting training.", model_path);
-    return;
-}
-
-fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
-    if visits.is_empty() {
-        return Vec::new();
-    }
-
-    // Special case: deterministic selection
-    if temperature == 0.0 {
-        let (&best_idx, _) = visits
-            .iter()
-            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
-            .map(|(i, p)| (i, p))
-            .unwrap();
-
-        return vec![(best_idx, 1.0)];
-    }
-
-    let inv_temp = 1.0 / temperature;
-
-    // Step 1: apply exponent
-    let mut adjusted: Vec<(usize, f32)> =
-        visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
-
-    // Step 2: normalize
-    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.iter_mut() {
-        *v /= sum;
-    }
-
-    adjusted
-}
-
-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 (idx, p) in dist {
-        r -= *p;
-        if r <= 0.0 {
-            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
-        }
-    }
-
-    // fallback due to floating point drift
-    dist.get(0)
-        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
-}
-