update packages

cli/src/commands.rs

@@ -1,3 +1,4 @@
 pub mod benchmark;
 pub mod play;
+pub mod selfplay;
 pub mod train;

cli/src/commands/benchmark.rs

@@ -1 +1 @@
-
+pub fn run() {}

cli/src/commands/selfplay.rs

@@ -0,0 +1 @@
+pub fn run() {}

engine/Cargo.toml

@@ -5,10 +5,18 @@ 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", "fusion", "wgpu", "ndarray"], default-features = false }
 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"] }
+tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread", "time"] }
+
+lru = "0.12"
+futures = "0.3"
+
+[features]
+default = []
+cuda-backend = ["burn/cuda"]
+

engine/src/lib.rs

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

engine/src/main.rs

@@ -1,28 +1,60 @@
 #![recursion_limit = "256"]
 
-use burn::backend::{Autodiff, Cuda};
 use burn::optim::AdamConfig;
+use burn::tensor::Device;
 use engine::mcts::MctsConfig;
-use engine::training::train::{train, TrainingConfig};
-// fn main() {
-//     type MyBackend = Wgpu<f32, i32>;
-//
-//     let device = Default::default();
-//     let model = ChessModelConfig::new(10, 512).init::<MyBackend>(&device);
-//
-//     println!("{model}");
-// }
+use engine::train::{TrainingConfig, train};
+
+#[cfg(feature = "cuda-backend")]
+use burn::tensor::DeviceIndex;
+
+fn select_device() -> Device {
+    match std::env::var("CHESS_DRAGON_DEVICE").as_deref() {
+        Ok("cuda") | Err(_) => {
+            #[cfg(feature = "cuda-backend")]
+            {
+                eprintln!("Using CUDA device");
+                return Device::cuda(DeviceIndex::default());
+            }
+            #[cfg(not(feature = "cuda-backend"))]
+            {
+                eprintln!("CUDA requested but not compiled in. Falling back to CPU.");
+                return Device::flex();
+            }
+        }
+        Ok("cpu") => {
+            eprintln!("Using CPU device");
+            Device::flex()
+        }
+        Ok("wgpu") => {
+            eprintln!("Using WGPU device");
+            Device::wgpu(Default::default())
+        }
+        Ok(other) => {
+            eprintln!(
+                "Unknown CHESS_DRAGON_DEVICE={}. Falling back to CPU.",
+                other
+            );
+            Device::flex()
+        }
+    }
+}
 
 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();
+    // Ensure spawned background threads (e.g. CubeCL device servers like DSD-0-0) have sufficient stack space
+    unsafe {
+        if std::env::var("RUST_MIN_STACK").is_err() {
+            std::env::set_var("RUST_MIN_STACK", "33554432"); // 32MB
+        }
+    }
 
-    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
+    let device = select_device();
+
+    // Fast-validation config: 400 sims + 48 games ≈ 22 min/iteration on a 3080.
+    // Early training with a near-random model benefits more from data diversity
+    // than from deep search. Once the model has real shape (~150 iterations),
+    // bump back to MctsConfig::new(800, ..., 32) and num_episodes: 96.
+    let mcts_config = MctsConfig::new(400, 1.0, 0.3, 0.25, 16);
 
     let adam_config = AdamConfig::new();
 
@@ -30,20 +62,26 @@ fn main() {
         max_time_s: None,
         num_iters: None,
         max_depth: 100, // unused
-        model_name: String::from("Test1"),
-        load_model: false,
-        hidden_channels: 64,
-        num_blocks: 4,
-        batch_size: 1024,         // num positions sampled to a batch (2048)
-        num_episodes: 100,        // num games generated per iteration (5000)
-        buffer_max_size: 200_000, // max number of samples in the buffer
+        model_name: String::from("128x8"),
+        load_model: true, // architecture changed — must start with fresh weights
+        hidden_channels: 128,
+        num_blocks: 8,
+        batch_size: 1024,
+        minibatch_size: 256,
+        num_train_steps: 40,
+        num_episodes: 48,
+        buffer_max_size: 200_000,
         mcts_config,
         optimizer: adam_config,
         lr: 2e-4,
         seed: None,
+        num_workers: 8,
+        nn_batch_size: 64,
+        episodes_per_worker: 6,
+        nn_cache_size: 1000_000,
     };
 
-    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,17 +1,16 @@
 // use crate::networks::PolicyNetwork;
-use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
-use crate::net::encoding::{encode_board_state_perspective, encode_move};
-use crate::net::model::ChessModel;
-use burn::prelude::Backend;
-use burn::Tensor;
+use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
+use crate::net::encoding::encode_move;
+use crate::worker::NNCacheKey;
 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;
+//TODO retain tree after move selection?
 
+#[derive(Clone)]
 pub struct Node {
     pub prior: f32,
     pub children: Vec<usize>,
@@ -50,10 +49,18 @@ impl Node {
     }
 }
 
-impl Clone for Node {
-    fn clone(&self) -> Self {
-        Self::new(self.prior, self.board_state.clone(), self.last_move.clone())
+fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
+    let prior_score = c_puct * child.prior * ((parent.visit_count + 1) as f32).sqrt()
+        / (1.0 + child.visit_count as f32);
+
+    let value_score: f32;
+    if child.visit_count > 0 {
+        value_score = -child.value(); // value from opposing side
+    } else {
+        value_score = 0.0;
     }
+
+    value_score + prior_score
 }
 
 #[derive(Clone, Debug)]
@@ -74,13 +81,14 @@ impl MctsResults {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct MctsConfig {
     pub num_simulations: usize,
     pub c_puct: f32,
     pub dirichlet_alpha: f32,
     pub dirichlet_epsilon: f32,
-    pub batch_max: usize,
+    pub batch_size: usize,
+    pub num_batches: usize,
 }
 
 impl MctsConfig {
@@ -89,14 +97,15 @@ impl MctsConfig {
         c_puct: f32,
         dirichlet_alpha: f32,
         dirichlet_epsilon: f32,
-        batch_max: usize,
+        batch_size: usize,
     ) -> MctsConfig {
         MctsConfig {
             num_simulations,
             c_puct,
             dirichlet_alpha,
             dirichlet_epsilon,
-            batch_max,
+            batch_size,
+            num_batches: num_simulations / batch_size,
         }
     }
 }
@@ -107,40 +116,97 @@ 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> {
-    pub fn search(
+fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
+    match status {
+        Ongoing => None,
+        BoardStateStatus::Stalemate | Threefold | FiftyMove | InsufficientMaterial => Some(0.0),
+        BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }),
+        BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }),
+    }
+}
+
+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
+}
+
+use std::future::Future;
+
+impl Mcts {
+    pub async fn search<F, Fut>(
         &mut self,
         board_state: &BoardState,
-        model: &ChessModel<B>,
-        device: &B::Device,
-    ) -> MctsResults {
+        model_inference: F,
+        training: bool,
+    ) -> MctsResults
+    where
+        F: Fn(&BoardState) -> Fut,
+        Fut: Future<Output = (Vec<f32>, f32)>,
+    {
         let mut nodes = Vec::<Node>::new();
 
         let root = 0;
         nodes.push(Node::new(0.0, board_state.clone(), None));
 
-        // Expand root to create initial children and priors
-        self.expand(root, &mut nodes, model, device);
+        let (root_policy, root_value) = model_inference(&board_state).await;
+        self.expand_with_eval(root, &mut nodes, &root_policy, root_value);
 
-        // Apply Dirichlet noise to root children
-        self.add_dirichlet_noise(root, &mut nodes);
+        // Seed the root node with the network's own value estimate.
+        // Without this, the root starts at (visit_count=0, value_sum=0) and the
+        // first ~N simulations have no reliable value signal — the root value is
+        // otherwise silently discarded.
+        let root_color = board_state.board.side_to_move();
+        self.backpropagate(&mut nodes, &[root], root_value, root_color);
+
+        if training {
+            // APPLY DIRICHLET NOISE HERE
+            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;
+        let mut sim_count = 0;
 
-        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 sim_count < num_sims {
+            let current_batch_size = std::cmp::min(self.config.batch_size, num_sims - sim_count);
+            sim_count += current_batch_size;
 
-            while leaf_nodes.len() < std::cmp::min(self.config.batch_max, num_sims - sims_done) {
+            let mut pending = Vec::with_capacity(current_batch_size);
+
+            for _ in 0..current_batch_size {
                 let mut path = vec![root];
                 let mut current = root;
 
@@ -149,174 +215,143 @@ impl<B: Backend> Mcts<B> {
                     path.push(current);
                 }
 
-                // Record leaf node and its path
-                leaf_nodes.push(current);
-                leaf_paths.push(path.clone());
+                // Apply virtual loss to path
+                for &idx in &path {
+                    nodes[idx].visit_count += 1;
+                    nodes[idx].value_sum += 1.0;
+                }
 
-                // 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);
+                let status = nodes[current].board_state.status;
+                let stm = nodes[current].board_state.board.side_to_move();
+                let term_val = terminal_value(status, stm);
 
-                sims_done += 1;
+                pending.push((path, current, term_val));
             }
 
-            if leaf_nodes.is_empty() {
-                break;
+            // Gather unique non-terminal leaves needing evaluation
+            let mut eval_indices = Vec::new();
+            for (_, leaf_idx, term_val) in &pending {
+                if term_val.is_none() && !eval_indices.contains(leaf_idx) {
+                    eval_indices.push(*leaf_idx);
+                }
             }
 
-            // Batch evaluate the collected leaf states
-            let batch = Tensor::cat(leaf_states, 0);
-            let (policy_batch, value_batch) = model.forward(batch);
+            // Execute inference for all non-terminal leaves concurrently
+            let mut eval_map = HashMap::new();
+            if !eval_indices.is_empty() {
+                let futures: Vec<_> = eval_indices
+                    .iter()
+                    .map(|&idx| model_inference(&nodes[idx].board_state))
+                    .collect();
 
-            // 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 results = futures::future::join_all(futures).await;
 
-            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;
-                    }
+                for (&idx, res) in eval_indices.iter().zip(results.into_iter()) {
+                    eval_map.insert(idx, res);
                 }
-                if logits.is_empty() {
-                    println!("logits empty")
-                }
-                if !max_logit.is_finite() {
-                    println!("max logits not finite")
+            }
+
+            // Phase 1: Revert ALL virtual losses before touching the tree with real values.
+            // Doing this in one pass ensures that when we backpropagate simulation N,
+            // no other simulation's artificial visit counts or value sums are polluting
+            // the nodes on its path.
+            for (path, _, _) in &pending {
+                for &idx in path {
+                    nodes[idx].visit_count -= 1;
+                    nodes[idx].value_sum -= 1.0;
                 }
+            }
 
-                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;
-                    }
+            // Phase 2: Expand leaves and backpropagate with a clean tree.
+            for (path, leaf_idx, term_val) in pending {
+                let (value, color) = if let Some(tv) = term_val {
+                    let color = nodes[leaf_idx].board_state.board.side_to_move();
+                    (tv, color)
                 } else {
-                    // fallback: uniform or zero
-                    let uniform = 1.0 / num_moves as f32;
-                    for p in &mut probs {
-                        *p = uniform;
+                    let (policy, val) = &eval_map[&leaf_idx];
+                    if nodes[leaf_idx].children.is_empty() {
+                        self.expand_with_eval(leaf_idx, &mut nodes, policy, *val);
                     }
-                }
+                    let color = nodes[leaf_idx].board_state.board.side_to_move();
+                    (*val, color)
+                };
 
-                // 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);
+                self.backpropagate(&mut nodes, &path, value, color);
             }
         }
 
-        // 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");
-            let prob = nodes[*child_idx].visit_count as f32 / denom;
-            move_dist.push((enc, prob));
+            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())
     }
 
-    fn expand(
+    fn expand_with_eval(
         &mut self,
         node_idx: usize,
         arena: &mut Vec<Node>,
-        model: &ChessModel<B>,
-        device: &B::Device,
+        policy_logits: &[f32],
+        value_head: f32,
     ) -> 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]);
-            // let start = Instant::now();
-            let (policy_head, value_head) = model.forward(state);
-            // println!("time: {:?}", start.elapsed());
+        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 {
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
+            let probs = softmax(policy_logits);
 
-            let policy = policy_head.into_data().to_vec::<f32>().unwrap();
+            let mut move_priors = Vec::with_capacity(legal_moves.len());
+            let mut prior_sum = 0.0f32;
 
-            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];
+            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
         }
     }
 
+
     fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) {
         for &idx in path {
             let node = &mut nodes[idx];
@@ -356,7 +391,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,27 +401,20 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
 
     let sum: f32 = samples.iter().sum();
 
-    for x in &mut samples {
-        *x /= 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)
     let root_node = &nodes[root];
@@ -476,6 +503,7 @@ pub enum BoardStateStatus {
     BlackWinner,
     Threefold,
     FiftyMove,
+    InsufficientMaterial,
 }
 
 #[derive(Clone, Debug)]
@@ -491,52 +519,112 @@ impl BoardState {
         BoardState {
             board,
             halfmove_clock,
-            repetition_table: repetition_table,
+            repetition_table,
             status: Ongoing,
         }
     }
 
+    pub fn nn_cache_key(&self) -> NNCacheKey {
+        let board_hash = self.board.get_hash(); // whatever chess crate API you're using
+
+        let repetition_two = self.repetition_table.get(&board_hash).copied().unwrap_or(0) >= 2;
+
+        NNCacheKey {
+            board_hash,
+            halfmove_clock: self.halfmove_clock,
+            repetition_two,
+        }
+    }
+
     pub fn apply_move(&mut self, mv: ChessMove) {
-        self.halfmove_clock += 1;
-        if self.board.piece_on(mv.get_source()) == Some(Pawn)
-            || self.board.piece_on(mv.get_dest()).is_some()
-        {
+        let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
+        let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();
+
+        let is_en_passant = is_pawn_move
+            && self.board.en_passant() == Some(mv.get_dest())
+            && self.board.piece_on(mv.get_dest()).is_none();
+
+        self.halfmove_clock = self.halfmove_clock.saturating_add(1);
+
+        if is_pawn_move || is_normal_capture || is_en_passant {
             self.halfmove_clock = 0;
         }
+
         self.board = self.board.make_move_new(mv);
 
         let board_hash = self.board.get_hash();
-
-        let current_rep = self.repetition_table.get(&board_hash);
-        let mut new_rep = 1;
-        if current_rep.is_some() {
-            new_rep = current_rep.unwrap() + 1;
-        }
+        let new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 1;
         self.repetition_table.insert(board_hash, new_rep);
 
-        if self.board.status() == Checkmate {
-            if self.board.side_to_move() == White {
-                // white's move after black played the mating move
-                self.status = BoardStateStatus::BlackWinner;
-            } else {
-                self.status = BoardStateStatus::WhiteWinner;
+        match self.board.status() {
+            Checkmate => {
+                self.status = if self.board.side_to_move() == White {
+                    BoardStateStatus::BlackWinner
+                } else {
+                    BoardStateStatus::WhiteWinner
+                };
+                return;
             }
-        } else if self.board.status() == Stalemate {
-            self.status = BoardStateStatus::Stalemate;
+            Stalemate => {
+                self.status = BoardStateStatus::Stalemate;
+                return;
+            }
+            _ => {}
+        }
+
+        if is_insufficient_material(&self.board) {
+            self.status = InsufficientMaterial;
+            return;
         }
 
         if self.halfmove_clock >= 100 {
             self.status = FiftyMove;
+            return;
         }
 
         if new_rep >= 3 {
             self.status = Threefold;
+            return;
         }
+
+        self.status = BoardStateStatus::Ongoing;
     }
 }
 
 impl Default for BoardState {
     fn default() -> BoardState {
-        BoardState::new(Board::default(), 0, HashMap::new())
+        let board = Board::default();
+        let repetition_table = HashMap::from([(board.get_hash(), 1)]);
+        BoardState::new(board, 0, repetition_table)
+    }
+}
+
+// Could be made more comprehensive
+pub fn is_insufficient_material(board: &Board) -> bool {
+    let pawns = board.pieces(Pawn).0;
+    let rooks = board.pieces(Rook).0;
+    let queens = board.pieces(Queen).0;
+
+    // Any pawn/rook/queen means this is not trivially dead.
+    if pawns != 0 || rooks != 0 || queens != 0 {
+        return false;
+    }
+
+    let bishops = board.pieces(Bishop).0;
+    let knights = board.pieces(Knight).0;
+
+    let bishop_count = bishops.count_ones();
+    let knight_count = knights.count_ones();
+
+    match (bishop_count, knight_count) {
+        // K vs K
+        (0, 0) => true,
+
+        // K+B vs K
+        // K+N vs K
+        (1, 0) | (0, 1) => true,
+
+        // Other configurations require more analysis.
+        _ => false,
     }
 }

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::{Board, ChessMove, Color, File, Piece, Rank, Square};
 
 const SLIDING_DIRS: [(i8, i8); 8] = [
     (0, 1),   // N
@@ -52,104 +32,113 @@ 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> {
-    let mut planes = vec![0.0f32; 18 * 64];
+pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
+    let mut planes = vec![0.0f32; 19 * 64];
+    encode_board_state_to_slice(state, &mut planes);
+    Tensor::<3>::from_floats(TensorData::new(planes, [19, 8, 8]), device)
+}
+
+/// Write 19x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1216 elements long).
+pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
+    out_slice[..1216].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_ones(out_slice, 12);
     }
     if us_castle.has_queenside() {
-        fill_plane(&mut planes, 13);
+        fill_plane_ones(out_slice, 13);
     }
     if them_castle.has_kingside() {
-        fill_plane(&mut planes, 14);
+        fill_plane_ones(out_slice, 14);
     }
     if them_castle.has_queenside() {
-        fill_plane(&mut planes, 15);
+        fill_plane_ones(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_ones(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)
+    fill_plane(out_slice, 18, Some(state.halfmove_clock as f32 / 100.0));
 }
 
-fn fill_plane(buffer: &mut [f32], plane: usize) {
-    let start = plane * 64;
-    for i in 0..64 {
-        buffer[start + i] = 1.0;
+pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
+    let b = states.len();
+    let mut buffer = vec![0.0f32; b * 1216];
+    for (i, state) in states.iter().enumerate() {
+        let start = i * 1216;
+        encode_board_state_to_slice(state, &mut buffer[start..start + 1216]);
     }
+    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 19, 8, 8]), device)
+}
+
+#[inline]
+fn fill_plane(buffer: &mut [f32], plane: usize, value: Option<f32>) {
+    let start = plane * 64;
+    buffer[start..start + 64].fill(value.unwrap_or(1.0));
+}
+
+#[inline]
+fn fill_plane_ones(buffer: &mut [f32], plane: usize) {
+    fill_plane(buffer, plane, Some(1.0));
 }
 
 /// Rotate a square into the side-to-move perspective.
@@ -201,7 +190,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> {
@@ -223,18 +212,18 @@ fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
 }
 
 /// Decode an index in [0, 4672) back into a ChessMove.
-pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
+pub fn decode_move(idx: usize, board: &Board) -> Option<ChessMove> {
     if idx >= 4672 {
         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 +249,20 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
+    let stm = board.side_to_move();
+
+    // Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion.
+    // fr/ff are perspective coordinates, so deorient back to real squares before inspecting the board.
+    if promo.is_none()
+        && fr == 6
+        && tr == 7
+        && dy == 1
+        && dx.abs() <= 1
+        && board.piece_on(deorient_square(ff, fr, stm)) == Some(Piece::Pawn)
+    {
+        promo = Some(Piece::Queen);
+    }
+
     let from = deorient_square(ff, fr, stm);
     let to = deorient_square(tf, tr, stm);
 
@@ -268,47 +271,109 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-    use chess::ALL_COLORS;
+    use crate::net::encoding::{decode_move, encode_move};
+    use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
+    use std::str::FromStr;
 
     #[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);
+        let board = Board::default();
+        for action in 0..4672 {
+            let decoded = decode_move(action, &board);
+            if decoded.is_none() {
+                continue;
             }
+            let decoded = decoded.unwrap();
+            let encoded = encode_move(decoded, board.side_to_move()).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);
         }
     }
+
+    #[test]
+    fn queen_promotion_roundtrip() {
+        let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
+
+        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, &board).unwrap();
+
+        assert_eq!(mv, dec);
+    }
+
+    #[test]
+    fn underpromotion_roundtrip() {
+        let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
+
+        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, &board).unwrap();
+
+        assert_eq!(mv, dec);
+    }
+
+    #[test]
+    fn queen_promotion_roundtrip_black() {
+        // Black pawn on d2, white king on h1, black king on e8, black to move.
+        // d2-d1 is a legal promotion; d1 is empty.
+        let board = Board::from_str("4k3/8/8/8/8/8/3p4/7K b - - 0 1").unwrap();
+
+        let from = Square::make_square(Rank::Second, File::D);
+        let to = Square::make_square(Rank::First, File::D);
+
+        let mv = ChessMove::new(from, to, Some(Piece::Queen));
+
+        let enc = encode_move(mv, Color::Black).unwrap();
+        let dec = decode_move(enc, &board).unwrap();
+
+        assert_eq!(mv, dec);
+    }
+
+    #[test]
+    fn underpromotion_roundtrip_black() {
+        let board = Board::from_str("4k3/8/8/8/8/8/3p4/7K b - - 0 1").unwrap();
+
+        let from = Square::make_square(Rank::Second, File::D);
+        let to = Square::make_square(Rank::First, File::D);
+
+        let mv = ChessMove::new(from, to, Some(Piece::Knight));
+
+        let enc = encode_move(mv, Color::Black).unwrap();
+        let dec = decode_move(enc, &board).unwrap();
+
+        assert_eq!(mv, dec);
+    }
 }

engine/src/net/model.rs

@@ -1,18 +1,15 @@
 use crate::mcts::{BoardState, MctsResults};
-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};
 use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
-use burn::tensor::activation::log_softmax;
-use burn::tensor::backend::AutodiffBackend;
 use burn::tensor::Transaction;
+use burn::tensor::activation::log_softmax;
 use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
 use burn::{
-    nn::{conv::Conv2d, Linear, Relu},
+    nn::{Linear, Relu, conv::Conv2d},
     prelude::*,
 };
-use burn_ndarray::NdArray;
 /*
 Input planes:
 1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -23,6 +20,7 @@ Input planes:
 16: opponent queenside
 17: rep >= 2
 18: En passant square (gap pass over)
+19: No-progress plys / 100
 
 
 Output:
@@ -39,21 +37,21 @@ 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,
 }
 
-#[derive(Config, Debug)]
-pub struct ResidualBlockConfig {
-    channels: usize,
-}
+// #[derive(Config, Debug)]
+// 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,13 +108,9 @@ 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
+            conv: Conv2dConfig::new([19, channels], [3, 3]) // 19 plane input
                 .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
                 .init(device),
             bn: BatchNormConfig::new(channels).init(device),
@@ -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,47 @@ 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) -> Result<Self, burn::tensor::ExecutionError> {
+        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()
@@ -226,32 +232,34 @@ impl<B: Backend> ItemLazy for ChessOutput<B> {
 
         let device = &Default::default();
 
-        ChessOutput {
+        Ok(ChessOutput {
             policy_logits: Tensor::from_data(policy_logits, device),
             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,55 +295,46 @@ 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> {
-        let state_tensors = items
-            .iter()
-            .map(|item| {
-                encode_board_state_perspective(&item.board_state, device).reshape([1, 18, 8, 8])
-            })
-            .collect::<Vec<_>>();
+impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
+    fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
+        let b = items.len();
 
-        let policy_target_tensors = items
-            .iter()
-            .cloned()
-            .map(|item| {
-                let mut policy = vec![0.0f32; 4672];
-                for (idx, prob) in item.policy_target.iter() {
-                    policy[*idx] = *prob;
+        let mut states_buf = vec![0.0f32; b * 1216];
+        let mut policy_buf = vec![0.0f32; b * 4672];
+        let mut value_buf = vec![0.0f32; b];
+
+        for (i, item) in items.iter().enumerate() {
+            crate::net::encoding::encode_board_state_to_slice(
+                &item.board_state,
+                &mut states_buf[i * 1216..(i + 1) * 1216],
+            );
+
+            let policy_slice = &mut policy_buf[i * 4672..(i + 1) * 4672];
+            for &(idx, prob) in &item.policy_target {
+                if idx < 4672 {
+                    policy_slice[idx] = prob;
                 }
-
-                // Normalize
-                let sum: f32 = policy.iter().sum();
-                if sum > 0.0 {
-                    for p in &mut policy {
-                        *p /= sum;
-                    }
+            }
+            let sum: f32 = policy_slice.iter().sum();
+            if sum > 0.0 {
+                for p in policy_slice.iter_mut() {
+                    *p /= sum;
                 }
-                Tensor::<B, 2>::from_floats(TensorData::new(policy, [1, 4672]), device).unsqueeze()
-            })
-            .collect::<Vec<_>>();
+            }
 
-        let value_target_tensors = items
-            .iter()
-            .map(|item| {
-                Tensor::<B, 2>::from_floats(
-                    TensorData::new(vec![item.value_target], [1, 1]),
-                    device,
-                )
-                .reshape([1, 1])
-            })
-            .collect::<Vec<_>>();
+            value_buf[i] = item.value_target;
+        }
 
-        let states = Tensor::cat(state_tensors, 0); // [B, 18, 8, 8]
-        let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672]
-        let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1]
+        let states = Tensor::<4>::from_floats(TensorData::new(states_buf, [b, 19, 8, 8]), device);
+        let policy_targets =
+            Tensor::<2>::from_floats(TensorData::new(policy_buf, [b, 4672]), device);
+        let value_targets = Tensor::<2>::from_floats(TensorData::new(value_buf, [b, 1]), device);
 
         ChessBatch {
             states,
@@ -344,3 +343,33 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::optim::{AdamConfig, GradientsParams};
+    use burn::tensor::Device;
+
+    #[cfg(feature = "cuda-backend")]
+    use burn::tensor::DeviceIndex;
+
+    #[test]
+    fn test_train_step() {
+        #[cfg(feature = "cuda-backend")]
+        let device = Device::cuda(DeviceIndex::default());
+        #[cfg(not(feature = "cuda-backend"))]
+        let device = Device::flex();
+        let autodiff_device = device.clone().autodiff();
+        let model = ChessModelConfig::init(8, 128, &autodiff_device);
+        let mut optim = AdamConfig::new().init();
+        let batcher = ChessBatcher {};
+        let samples: Vec<TrainingSample> = (0..256)
+            .map(|_| TrainingSample::new(BoardState::default(), vec![(0, 1.0)], 0.0))
+            .collect();
+        let batch = batcher.batch(samples, &autodiff_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);
+        let _model = optim.step(2e-4, model, grads);
+    }
+}

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,452 @@
-// 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::MctsConfig;
+use crate::net::encoding::decode_move;
+use crate::net::model::{
+    ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
+};
+use crate::worker::{NNRequest, NNServer, NNServerCmd, Worker, WorkerCmd, WorkerData};
+use burn::data::dataloader::batcher::Batcher;
+use burn::optim::{AdamConfig, GradientsParams};
+use burn::prelude::{Device, Module};
+use burn::store::ModuleRecord;
+use chess::{Board, ChessMove};
+use rand::rngs::SmallRng;
+use rand::seq::SliceRandom;
+use rand::{RngExt, SeedableRng};
+use std::cmp::min;
+use std::collections::{HashMap, VecDeque};
+use std::fs::File;
+use std::io::BufReader;
 
-// 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();
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{mpsc, Arc};
 
-// 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");
-// }
+use std::thread;
+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 minibatch_size: usize,  // num positions sampled to a batch (2048)
+    pub num_train_steps: usize, // number of training steps of size minibatch_size
+    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 num_workers: usize,
+    pub nn_batch_size: usize,
+    pub episodes_per_worker: usize,
+    pub nn_cache_size: u64,
+}
+
+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();
+
+    // Create RNG once and reuse it for sampling and shuffling
+    // Seed from system time (platform default entropy may be unavailable in some contexts)
+    let seed = training_config.seed.unwrap_or(
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64,
+    );
+    println!("using seed: {}", seed);
+    let mut rng = SmallRng::seed_from_u64(seed);
+
+    println!("Starting training...");
+    while train.load(Ordering::Relaxed) {
+        let iter_start = Instant::now();
+        println!("Iteration: {}", iter);
+        let infer_model = model.valid();
+        // Gen samples
+        println!("Generating {} games...", training_config.num_episodes);
+
+        let mut episodes_dispatched: usize = 0;
+        let mut episodes_recieved: usize = 0;
+
+        let mut worker_map: HashMap<usize, tokio::sync::mpsc::Sender<WorkerCmd>> = HashMap::new();
+        let (server_cmd_tx, server_cmd_rx) = mpsc::channel::<NNServerCmd>();
+        let (data_tx, data_rx) = mpsc::channel::<WorkerData>();
+        let (server_request_tx, server_request_rx) = mpsc::channel::<NNRequest>();
+
+        let server_device = device.clone();
+        let server_handle = thread::spawn(move || {
+            let server = NNServer::new(
+                server_cmd_rx,
+                server_request_rx,
+                infer_model,
+                server_device,
+                training_config.nn_batch_size,
+                training_config.nn_cache_size,
+            );
+
+            server.run_server()
+        });
+
+        let mut worker_handles = Vec::new();
+        let mut workers_dispatched = 0;
+        for i in 0..training_config.num_workers {
+            // println!("initializing worker {}", i);
+            if training_config.num_episodes > episodes_dispatched {
+                let (cmd_tx, cmd_rx) = tokio::sync::mpsc::channel::<WorkerCmd>(100);
+                worker_map.insert(i, cmd_tx);
+                let episodes_to_dispatch = min(
+                    training_config.episodes_per_worker,
+                    training_config.num_episodes - episodes_dispatched,
+                );
+
+                let req_tx = server_request_tx.clone();
+                let data_tx_thread = data_tx.clone();
+                let mcts_config_thread = training_config.mcts_config.clone();
+                let device_thread = device.clone();
+                worker_handles.push(thread::spawn(move || {
+                    // println!("Starting worker thread...");
+
+                    let worker = Worker::new(
+                        i,
+                        episodes_to_dispatch,
+                        req_tx,
+                        data_tx_thread,
+                        device_thread,
+                        mcts_config_thread,
+                    );
+
+                    tokio::runtime::Builder::new_current_thread()
+                        .enable_all()
+                        .build()
+                        .unwrap()
+                        .block_on(worker.run_worker(cmd_rx));
+
+                }));
+                episodes_dispatched += episodes_to_dispatch;
+                workers_dispatched += 1;
+            }
+        }
+        println!(
+            "Started {} workers, dispatched episodes #1-{}",
+            workers_dispatched, episodes_dispatched
+        );
+
+        while episodes_recieved < training_config.num_episodes {
+            let data = data_rx.recv().unwrap();
+
+            for sample in data.data {
+                replay_buffer.push_back(sample);
+
+                if replay_buffer.len() > training_config.buffer_max_size {
+                    replay_buffer.pop_front();
+                }
+            }
+            episodes_recieved += 1;
+            println!("Recieved episode #{}", episodes_recieved);
+
+            if episodes_dispatched < training_config.num_episodes {
+                worker_map
+                    .get(&data.worker_id)
+                    .unwrap()
+                    .blocking_send(WorkerCmd::Dispatch)
+                    .expect("Expected dispatch to be sent");
+                episodes_dispatched += 1;
+                println!("Dispatched episode #{}", episodes_dispatched);
+            }
+
+            // else {
+            //     worker_map
+            //         .get(&data.worker_id)
+            //         .unwrap()
+            //         .send(WorkerCmd::Shutdown)
+            //         .expect("Expected shutdown to be sent");
+            //     // println!("Shutdown worker #{}", data.worker_id);
+            // }
+        }
+
+        server_cmd_tx
+            .send(NNServerCmd::Shutdown)
+            .expect("Expected shutdown command to be sent");
+
+        server_handle.join().unwrap();
+        for handle in worker_handles {
+            handle.join().unwrap()
+        }
+
+        // train
+
+        if replay_buffer.is_empty() {
+            panic!("Replay buffer is empty, cannot train.");
+            // 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 {} times with a minibatch size of {} (buffer size: {})...",
+            num_samples,
+            training_config.num_train_steps,
+            training_config.minibatch_size,
+            replay_buffer.len()
+        );
+
+        let batcher = ChessBatcher {};
+        // Instead of one step, do multiple epochs over the sampled batch
+        for step in 0..training_config.num_train_steps {
+            // Re-shuffle and create mini-batches within the sampled data
+            let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
+            indices.shuffle(&mut rng);
+            indices.truncate(num_samples);
+
+            let mut loss_val: f32 = 0.0;
+            let mut p_loss_val: f32 = 0.0;
+            let mut v_loss_val: f32 = 0.0;
+
+            // Process in mini-batches of e.g. 256
+            for chunk in indices.chunks(training_config.minibatch_size) {
+                let samples: Vec<TrainingSample> =
+                    chunk.iter().map(|&i| replay_buffer[i].clone()).collect();
+
+                let batch = batcher.batch(samples, &autodiff_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);
+
+                if step == training_config.num_train_steps - 1 {
+                    loss_val = output.loss.clone().into_data().try_to_vec_as::<f32>().unwrap()[0];
+                    p_loss_val = output
+                        .policy_loss
+                        .clone()
+                        .into_data()
+                        .try_to_vec_as::<f32>()
+                        .unwrap()[0];
+                    v_loss_val = output
+                        .value_loss
+                        .clone()
+                        .into_data()
+                        .try_to_vec_as::<f32>()
+                        .unwrap()[0];
+                }
+            }
+            if step == training_config.num_train_steps - 1 {
+                println!(
+                    "Iteration {}, last minibatch training loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
+                    iter, loss_val, p_loss_val, v_loss_val
+                );
+                println!("Iteration took {}s", iter_start.elapsed().as_secs());
+            }
+        }
+
+        // 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
+        // );
+        // println!("Iteration took {}s", iter_start.elapsed().as_secs());
+        //
+        // let grads = output.loss.backward();
+        // let grads = GradientsParams::from_grads(grads, &model);
+        //
+        // model = optim.step(training_config.lr, model, grads);
+
+        iter += 1;
+
+        println!("Saving model...");
+
+        let metadata = ModelMetadata {
+            name: training_config.model_name.clone(),
+            iterations: iter,
+        };
+
+        std::fs::write(
+            metadata_path.clone(),
+            serde_json::to_string_pretty(&metadata)
+                .expect("Should be able to convert metadata to JSON string"),
+        )
+        .expect("Should be able to write metadata");
+
+        model
+            .clone()
+            .into_record()
+            .save(&model_path)
+            .expect("Should be able to save the model");
+
+        optim
+            .clone()
+            .save(format!("{model_path}_optim"))
+            .expect("Should be able to save the optim");
+        println!("Model saved in {:?}.", model_path);
+
+        // 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;
+        }
+    }
+
+    // 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);
+    println!("Finished training!");
+    return;
+}
+
+pub(crate) 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
+}
+
+pub(crate) fn sample_move(
+    dist: &[(usize, f32)],
+    rng: &mut SmallRng,
+    board: &Board,
+) -> 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, board).expect("Invalid move"));
+        }
+    }
+
+    // fallback due to floating point drift
+    dist.get(0)
+        .map(|(idx, _)| decode_move(*idx, board).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"))
-}
-

engine/src/worker.rs

@@ -0,0 +1,456 @@
+use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
+use crate::net::model::{ChessModel, TrainingSample};
+use crate::train::{apply_temperature, sample_move};
+use burn::tensor::Device;
+use chess::Color;
+
+use lru::LruCache;
+use rand::SeedableRng;
+use rand::prelude::SmallRng;
+use std::collections::VecDeque;
+use std::hash::Hash;
+use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
+use std::time::Duration;
+
+pub struct Worker {
+    id: usize,
+    num_games: usize,
+    // games: Vec<BoardState>,
+    request_tx: Sender<NNRequest>,
+    data_tx: Sender<WorkerData>,
+    curr_request_id: u64,
+    device: Device,
+    mcts_config: MctsConfig,
+}
+
+pub struct WorkerData {
+    pub(crate) worker_id: usize,
+    pub(crate) data: Vec<TrainingSample>,
+}
+
+impl Worker {
+    pub fn new(
+        id: usize,
+        num_games: usize,
+        sender: Sender<NNRequest>,
+        data_tx: Sender<WorkerData>,
+        device: Device,
+        mcts_config: MctsConfig,
+    ) -> Worker {
+        Worker {
+            id,
+            // games: (0..num_games)
+            //     .map(|_| BoardState::default())
+            //     .collect::<Vec<BoardState>>(),
+            num_games,
+            request_tx: sender,
+            data_tx,
+            curr_request_id: 0,
+            device,
+            mcts_config,
+        }
+    }
+
+    async fn run_game<F, Fut>(
+        worker_id: usize,
+        mcts_config: MctsConfig,
+        model_inference: F,
+        data_tx: Sender<WorkerData>,
+    ) where
+        F: Fn(&BoardState) -> Fut,
+        Fut: Future<Output = (Vec<f32>, f32)>,
+    {
+        // println!("game started");
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap();
+        let seed = now.as_nanos() as u64;
+        let mut rng = SmallRng::seed_from_u64(seed);
+        let mut mcts: Mcts = Mcts {
+            config: mcts_config,
+        };
+
+        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, &model_inference, true).await;
+            // println!("{}", before.elapsed().as_millis());
+            episode_buffer.push(results);
+
+            // AlphaZero temperature schedule: explore freely for the first 30 plies,
+            // then play deterministically (argmax) for the rest of the game.
+            // The previous schedule never dropped below 0.3, keeping games noisy
+            // throughout and contaminating value targets with random late-game play.
+            let temp = if ply_count < 30 { 1.0 } else { 0.0 };
+            // let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3);
+
+            let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
+
+            let mv = sample_move(&adjusted, &mut rng, &board_state.board).unwrap();
+
+            game_hist.push_str(mv.to_string().as_str());
+            game_hist.push(' ');
+
+            board_state.apply_move(mv);
+            ply_count += 1;
+
+            // truncate long games
+            if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
+                board_state.status = BoardStateStatus::FiftyMove;
+            }
+        }
+
+        println!("game finished: {}", game_hist);
+
+        let mut samples: Vec<TrainingSample> = Vec::new();
+        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,
+            };
+
+            let blended_target = outcome;
+            // let blended_target = 0.75 * outcome + 0.25 * result.value;
+
+            samples.push(TrainingSample::from_mcts_with_outcome(
+                result.clone(),
+                blended_target,
+            ));
+        }
+
+        let worker_data = WorkerData {
+            worker_id,
+            data: samples,
+        };
+        data_tx.send(worker_data).unwrap();
+    }
+
+    // async fn model_inference(&self, board_state: BoardState) -> (Vec<f32>, f32) {
+    //     let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>();
+    //     let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
+    //
+    //     self.request_tx
+    //         .send(NNRequest {
+    //             request_tensor: state,
+    //             response_tx: tx,
+    //         })
+    //         .unwrap();
+    //
+    //     let response = rx.await.expect("worker dropped inference response");
+    //     let p = response
+    //         .response_tensors
+    //         .0
+    //         .into_data()
+    //         .to_vec::<f32>()
+    //         .unwrap();
+    //     let v = response
+    //         .response_tensors
+    //         .1
+    //         .into_data()
+    //         .to_vec::<f32>()
+    //         .unwrap()[0];
+    //
+    //     (p, v)
+    // }
+
+    pub async fn run_worker(self, mut cmd_rx: tokio::sync::mpsc::Receiver<WorkerCmd>) {
+        let mut tasks = Vec::new();
+        // println!("Worker started");
+        for _ in 0..self.num_games {
+            let mcts_config = self.mcts_config.clone();
+            let data_tx = self.data_tx.clone();
+            let request_tx = self.request_tx.clone();
+
+            tasks.push(tokio::spawn(async move {
+                Self::run_game(
+                    self.id,
+                    mcts_config,
+                    |board_state| {
+                        let request_tx = request_tx.clone();
+                        model_inference_static(board_state.clone(), request_tx)
+                    },
+                    data_tx,
+                )
+                .await;
+            }));
+        }
+
+        loop {
+            // Remove completed tasks.
+            tasks.retain(|task| !task.is_finished());
+
+            // If every game has completed, we're done.
+            if tasks.is_empty() {
+                break;
+            }
+
+            match tokio::time::timeout(Duration::from_millis(10), cmd_rx.recv()).await {
+                Ok(Some(WorkerCmd::Shutdown)) => break,
+
+                Ok(Some(WorkerCmd::Dispatch)) => {
+                    let mcts_config = self.mcts_config.clone();
+                    let request_tx = self.request_tx.clone();
+                    let data_tx = self.data_tx.clone();
+
+                    tasks.push(tokio::spawn(async move {
+                        Self::run_game(
+                            self.id,
+                            mcts_config,
+                            |board_state| {
+                                let request_tx = request_tx.clone();
+                                model_inference_static(board_state.clone(), request_tx)
+                            },
+                            data_tx,
+                        )
+                        .await;
+                    }));
+                }
+
+                Ok(None) => break,
+
+                Err(_) => {
+                    tokio::task::yield_now().await;
+                }
+            }
+        }
+
+        for task in tasks {
+            let _ = task.await;
+        }
+
+        println!("Worker {} shut down", self.id);
+    }
+}
+
+async fn model_inference_static(
+    board_state: BoardState,
+    request_tx: Sender<NNRequest>,
+) -> (Vec<f32>, f32) {
+    let cache_key = board_state.nn_cache_key();
+    let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
+
+    request_tx
+        .send(NNRequest {
+            cache_key,
+            board_state,
+            response_tx: tx,
+        })
+        .unwrap();
+
+    let response = rx.await.expect("worker dropped inference response");
+
+    (response.policy, response.value)
+}
+
+pub enum NNServerCmd {
+    Shutdown,
+}
+
+pub enum WorkerCmd {
+    Shutdown,
+    Dispatch,
+}
+
+pub struct NNServer {
+    command_rx: Receiver<NNServerCmd>,
+    request_rx: Receiver<NNRequest>,
+    model: ChessModel,
+    device: Device,
+    batch_size: usize,
+    cache_size: u64,
+}
+
+impl NNServer {
+    pub fn new(
+        command_rx: Receiver<NNServerCmd>,
+        request_rx: Receiver<NNRequest>,
+        model: ChessModel,
+        device: Device,
+        batch_size: usize,
+        cache_size: u64,
+    ) -> NNServer {
+        NNServer {
+            command_rx,
+            request_rx,
+            model,
+            device,
+            batch_size,
+            cache_size,
+        }
+    }
+}
+#[derive(Debug, Clone, Eq, Hash, PartialEq)]
+struct NNJobData {
+    worker_id: usize,
+    request_id: u64,
+    // game_id: u32,
+}
+pub struct NNRequest {
+    cache_key: NNCacheKey,
+    board_state: BoardState,
+    response_tx: tokio::sync::oneshot::Sender<NNResponse>,
+}
+
+// ~18.7KB
+#[derive(Debug, Clone)]
+pub struct NNResponse {
+    policy: Vec<f32>,
+    value: f32,
+}
+
+impl NNServer {
+    pub fn run_server(&self) {
+        let cap =
+            std::num::NonZeroUsize::new(self.cache_size as usize).expect("cache_size must be > 0");
+        let mut cache: LruCache<NNCacheKey, NNResponse> = LruCache::new(cap);
+        let mut queue: VecDeque<NNRequest> = VecDeque::new();
+        let mut request_counter: usize = 0;
+        let mut cache_hits: usize = 0;
+        let mut batches_ran: usize = 0;
+
+        loop {
+            // Check for shutdown before doing any work.
+            if let Ok(cmd) = self.command_rx.try_recv() {
+                match cmd {
+                    NNServerCmd::Shutdown => {
+                        println!(
+                            "NN server shutting down. Served {} requests, with {} cache hits ({:.2}%). Ran {} batches (avg {:.2} requests/batch).",
+                            request_counter,
+                            cache_hits,
+                            cache_hits as f64 / request_counter as f64 * 100.0,
+                            batches_ran,
+                            (request_counter - cache_hits) as f64 / batches_ran as f64,
+                        );
+
+                        break;
+                    } // other commands...
+                }
+            }
+
+            // If we have no queued requests, wait for either a request
+            // or periodically wake up to check for shutdown.
+            if queue.is_empty() {
+                match self.request_rx.recv_timeout(Duration::from_millis(10)) {
+                    Ok(first) => {
+                        request_counter += 1;
+                        // queue.push_back(first);
+                        if try_use_cache(first, &mut cache, &mut queue) {
+                            cache_hits += 1;
+                            continue;
+                        }
+                    }
+
+                    Err(RecvTimeoutError::Timeout) => {
+                        continue;
+                    }
+
+                    Err(RecvTimeoutError::Disconnected) => {
+                        println!("Request channel disconnected");
+                        break;
+                    }
+                }
+            }
+
+            while let Ok(req) = self.request_rx.try_recv() {
+                request_counter += 1;
+                // queue.push_back(req);
+                if try_use_cache(req, &mut cache, &mut queue) {
+                    cache_hits += 1;
+                }
+            }
+
+            let mut jobs_to_send: Vec<NNRequest> = Vec::new();
+
+            for _ in 0..self.batch_size {
+                if let Some(job) = queue.pop_front() {
+                    jobs_to_send.push(job);
+                } else {
+                    break;
+                }
+            }
+
+            if jobs_to_send.is_empty() {
+                continue;
+            }
+
+            let b_len = jobs_to_send.len();
+            let mut buffer = vec![0.0f32; b_len * 1216];
+            for (i, job) in jobs_to_send.iter().enumerate() {
+                let start = i * 1216;
+                crate::net::encoding::encode_board_state_to_slice(
+                    &job.board_state,
+                    &mut buffer[start..start + 1216],
+                );
+            }
+
+            let batched_tensor = burn::tensor::Tensor::<4>::from_floats(
+                burn::tensor::TensorData::new(buffer, [b_len, 19, 8, 8]),
+                &self.device,
+            );
+
+            let (p, v) = self.model.forward(batched_tensor);
+            batches_ran += 1;
+
+            let p_vec = p.into_data().try_to_vec_as::<f32>().unwrap();
+            let v_vec = v.into_data().try_to_vec_as::<f32>().unwrap();
+
+            const POLICY_DIM: usize = 73 * 8 * 8; // 4672
+
+            for (i, job) in jobs_to_send.into_iter().enumerate() {
+                let p_slice = p_vec[i * POLICY_DIM..(i + 1) * POLICY_DIM].to_vec();
+                let v_val = v_vec[i];
+
+                let res = NNResponse {
+                    policy: p_slice,
+                    value: v_val,
+                };
+
+                cache.put(job.cache_key, res.clone());
+
+                let _ = job.response_tx.send(res);
+            }
+        }
+    }
+}
+
+fn try_use_cache(
+    req: NNRequest,
+    cache: &mut LruCache<NNCacheKey, NNResponse>,
+    queue: &mut VecDeque<NNRequest>,
+) -> bool {
+    if let Some(cached) = cache.get(&req.cache_key) {
+        // println!("Cache hit");
+        req.response_tx.send(cached.clone()).unwrap();
+        true
+    } else {
+        queue.push_back(req);
+        false
+    }
+}
+
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+pub struct NNCacheKey {
+    pub board_hash: u64,
+    pub halfmove_clock: u8,
+    pub repetition_two: bool,
+}

uci/src/main.rs

@@ -1,5 +1,5 @@
 use chess::{ChessMove, Game};
-use engine::{legal_action_mask, Engine};
+use engine::Engine;
 use std::io;
 use std::str::FromStr;