fixes

Cargo.lock

@@ -1208,9 +1208,9 @@ dependencies = [
 
 [[package]]
 name = "crossbeam-epoch"
-version = "0.9.18"
+version = "0.9.20"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
+checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
 dependencies = [
  "crossbeam-utils",
 ]
@@ -2169,10 +2169,13 @@ dependencies = [
  "burn",
  "chess",
  "ctrlc",
+ "futures",
+ "lru 0.12.5",
  "rand 0.10.1",
  "rand_distr",
  "serde",
  "serde_json",
+ "tokio",
 ]
 
 [[package]]
@@ -2477,10 +2480,25 @@ dependencies = [
 ]
 
 [[package]]
-name = "futures-channel"
-version = "0.3.32"
+name = "futures"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d"
+checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
+dependencies = [
+ "futures-channel",
+ "futures-core",
+ "futures-executor",
+ "futures-io",
+ "futures-sink",
+ "futures-task",
+ "futures-util",
+]
+
+[[package]]
+name = "futures-channel"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
 dependencies = [
  "futures-core",
  "futures-sink",
@@ -2488,15 +2506,26 @@ dependencies = [
 
 [[package]]
 name = "futures-core"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
+checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
+
+[[package]]
+name = "futures-executor"
+version = "0.3.34"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
+dependencies = [
+ "futures-core",
+ "futures-task",
+ "futures-util",
+]
 
 [[package]]
 name = "futures-io"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
+checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
 
 [[package]]
 name = "futures-lite"
@@ -2513,26 +2542,26 @@ dependencies = [
 
 [[package]]
 name = "futures-macro"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
+checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.117",
+ "syn 3.0.4",
 ]
 
 [[package]]
 name = "futures-sink"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
+checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
 
 [[package]]
 name = "futures-task"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393"
+checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
 
 [[package]]
 name = "futures-timer"
@@ -2542,10 +2571,11 @@ checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
 
 [[package]]
 name = "futures-util"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
+checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
 dependencies = [
+ "futures-channel",
  "futures-core",
  "futures-io",
  "futures-macro",
@@ -3501,6 +3531,15 @@ dependencies = [
  "imgref",
 ]
 
+[[package]]
+name = "lru"
+version = "0.12.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38"
+dependencies = [
+ "hashbrown 0.15.5",
+]
+
 [[package]]
 name = "lru"
 version = "0.16.4"
@@ -4679,7 +4718,7 @@ dependencies = [
  "indoc",
  "itertools",
  "kasuari",
- "lru",
+ "lru 0.16.4",
  "strum 0.27.2",
  "thiserror 2.0.18",
  "unicode-segmentation",

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

@@ -11,3 +11,8 @@ 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"
+

engine/src/lib.rs

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

engine/src/main.rs

@@ -3,7 +3,7 @@
 use burn::optim::AdamConfig;
 use burn::tensor::{Device, DeviceIndex};
 use engine::mcts::MctsConfig;
-use engine::train::{train, TrainingConfig};
+use engine::train::{TrainingConfig, train};
 // fn main() {
 //     type MyBackend = Wgpu<f32, i32>;
 //
@@ -14,9 +14,20 @@ use engine::train::{train, TrainingConfig};
 // }
 
 fn main() {
+    // 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 device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(150, 1.0, 0.05, 0.25, 128);
+    // 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();
 
@@ -24,17 +35,23 @@ 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(training_config, device);

engine/src/mcts.rs

@@ -1,15 +1,16 @@
 // use crate::networks::PolicyNetwork;
-use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
-use crate::net::encoding::{encode_board_states_batch, encode_move};
-use crate::net::model::ChessModel;
-use burn::prelude::Device;
+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;
+//TODO retain tree after move selection?
 
+#[derive(Clone)]
 pub struct Node {
     pub prior: f32,
     pub children: Vec<usize>,
@@ -36,40 +37,32 @@ impl Node {
     }
 
     pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
-        let sqrt_parent = (self.visit_count as f32).sqrt();
-        let c_puct_sqrt = c_puct * sqrt_parent;
-
         self.children
             .iter()
             .copied()
             .max_by(|&a, &b| {
-                let score_a = ucb_score_fast(&arena[a], c_puct_sqrt);
-                let score_b = ucb_score_fast(&arena[b], c_puct_sqrt);
-                score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
+                ucb_score(self, &arena[a], c_puct)
+                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
+                    .unwrap()
             })
             .expect("select_child on leaf")
     }
 }
 
-#[inline]
-fn ucb_score_fast(child: &Node, c_puct_sqrt: f32) -> f32 {
-    let prior_score = c_puct_sqrt * child.prior / (1.0 + child.visit_count as f32);
+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 = if child.visit_count > 0 {
-        -child.value()
+    let value_score: f32;
+    if child.visit_count > 0 {
+        value_score = -child.value(); // value from opposing side
     } else {
-        0.0
-    };
+        value_score = 0.0;
+    }
 
     value_score + prior_score
 }
 
-impl Clone for Node {
-    fn clone(&self) -> Self {
-        Self::new(self.prior, self.board_state.clone(), self.last_move.clone())
-    }
-}
-
 #[derive(Clone, Debug)]
 pub struct MctsResults {
     pub board_state: BoardState,
@@ -88,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 {
@@ -103,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,
         }
     }
 }
@@ -127,10 +122,8 @@ pub struct Mcts {
 
 fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
     match status {
-        BoardStateStatus::Ongoing => None,
-        BoardStateStatus::Stalemate
-        | BoardStateStatus::Threefold
-        | BoardStateStatus::FiftyMove => Some(0.0),
+        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 }),
     }
@@ -171,159 +164,135 @@ fn softmax(logits: &[f32]) -> Vec<f32> {
     probs
 }
 
-const LEAF_VL_COUNT: u32 = 10;
-const LEAF_VL_WEIGHT: f32 = 10.0;
+use std::future::Future;
 
 impl Mcts {
-    pub fn search(
+    pub async fn search<F, Fut>(
         &mut self,
         board_state: &BoardState,
-        model: &ChessModel,
-        device: &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);
+        }
 
-        let mut sims_done: usize = 0;
         let num_sims = self.config.num_simulations;
+        let mut sim_count = 0;
 
-        while sims_done < num_sims {
-            let mut leaf_nodes: Vec<usize> = Vec::new();
-            let mut leaf_paths: Vec<Vec<usize>> = 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;
 
-            let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
+            let mut pending = Vec::with_capacity(current_batch_size);
 
-            while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims {
+            for _ in 0..current_batch_size {
                 let mut path = vec![root];
                 let mut current = root;
 
-                // Apply virtual loss during selection path
                 while !nodes[current].children.is_empty() {
-                    nodes[current].visit_count += 1;
-                    nodes[current].value_sum -= 1.0;
-
                     current = nodes[current].select_child(&nodes, &self.config.c_puct);
                     path.push(current);
                 }
 
-                // Strong virtual loss on unexpanded leaf so batch collector visits diverse leaves across tree
-                nodes[current].visit_count += LEAF_VL_COUNT;
-                nodes[current].value_sum -= LEAF_VL_WEIGHT;
-
-                // Check terminal status of leaf
-                let stm = nodes[current].board_state.board.side_to_move();
-                if let Some(term_val) = terminal_value(nodes[current].board_state.status, stm) {
-                    // Revert virtual loss for path
-                    for &idx in &path[..path.len() - 1] {
-                        nodes[idx].visit_count -= 1;
-                        nodes[idx].value_sum += 1.0;
-                    }
-                    nodes[current].visit_count -= LEAF_VL_COUNT;
-                    nodes[current].value_sum += LEAF_VL_WEIGHT;
-
-                    self.backpropagate(&mut nodes, &path, term_val, stm);
-                    sims_done += 1;
-                    continue;
-                }
-
-                // Prevent expanding the exact same node twice in a single batch
-                if leaf_nodes.contains(&current) {
-                    // Revert virtual loss for this duplicate path
-                    for &idx in &path[..path.len() - 1] {
-                        nodes[idx].visit_count -= 1;
-                        nodes[idx].value_sum += 1.0;
-                    }
-                    nodes[current].visit_count -= LEAF_VL_COUNT;
-                    nodes[current].value_sum += LEAF_VL_WEIGHT;
-                    break;
-                }
-
-                leaf_nodes.push(current);
-                leaf_paths.push(path);
-            }
-
-            if leaf_nodes.is_empty() {
-                if sims_done >= num_sims {
-                    break;
-                }
-                continue;
-            }
-
-            // Batch evaluate non-terminal leaf states using single CPU buffer encoding
-            let leaf_refs: Vec<&BoardState> = leaf_nodes.iter().map(|&idx| &nodes[idx].board_state).collect();
-            let batch = encode_board_states_batch(&leaf_refs, device);
-            let (policy_batch, value_batch) = model.forward(batch);
-
-            let policy_data = policy_batch.into_data().to_vec::<f32>().unwrap();
-            let value_data = value_batch.into_data().to_vec::<f32>().unwrap();
-
-            let num_moves = policy_data.len() / leaf_nodes.len();
-
-            for (i, &node_idx) in leaf_nodes.iter().enumerate() {
-                let path = &leaf_paths[i];
-
-                // Revert virtual loss along this path before applying true expansion & backpropagation
-                for &idx in &path[..path.len() - 1] {
-                    nodes[idx].visit_count -= 1;
+                // Apply virtual loss to path
+                for &idx in &path {
+                    nodes[idx].visit_count += 1;
                     nodes[idx].value_sum += 1.0;
                 }
-                nodes[node_idx].visit_count -= LEAF_VL_COUNT;
-                nodes[node_idx].value_sum += LEAF_VL_WEIGHT;
 
-                let start = i * num_moves;
-                let end = start + num_moves;
-                let logits = &policy_data[start..end];
-                let probs = softmax(logits);
+                let status = nodes[current].board_state.status;
+                let stm = nodes[current].board_state.board.side_to_move();
+                let term_val = terminal_value(status, stm);
 
-                let stm = nodes[node_idx].board_state.board.side_to_move();
-                let legal_moves: Vec<ChessMove> =
-                    MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();
+                pending.push((path, current, term_val));
+            }
 
-                let mut move_priors = Vec::with_capacity(legal_moves.len());
-                let mut prior_sum = 0.0f32;
+            // 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);
+                }
+            }
 
-                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;
+            // 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();
+
+                let results = futures::future::join_all(futures).await;
+
+                for (&idx, res) in eval_indices.iter().zip(results.into_iter()) {
+                    eval_map.insert(idx, res);
+                }
+            }
+
+            // 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;
+                }
+            }
+
+            // 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 {
+                    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)
+                };
 
-                for (mv, p) in move_priors {
-                    let normalized_prior = if prior_sum > 0.0 {
-                        p / prior_sum
-                    } else {
-                        1.0 / legal_moves.len() as f32
-                    };
-
-                    let mut new_board = nodes[node_idx].board_state.clone();
-                    new_board.apply_move(mv);
-
-                    let child_idx = nodes.len();
-                    nodes.push(Node::new(normalized_prior, new_board, Some(mv)));
-                    nodes[node_idx].children.push(child_idx);
-                }
-
-                let value = value_data[i];
-                self.backpropagate(&mut nodes, path, value, stm);
-                sims_done += 1;
+                self.backpropagate(&mut nodes, &path, value, color);
             }
         }
 
         // Build compact move distribution
         let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
         let stm = board_state.board.side_to_move();
-        let total_root_visits: u32 = nodes[root].children.iter().map(|&c| nodes[c].visit_count).sum();
-        let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 };
+        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");
             if let Some(enc) = encode_move(mv, stm) {
@@ -335,12 +304,12 @@ impl Mcts {
         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,
-        device: &Device,
+        policy_logits: &[f32],
+        value_head: f32,
     ) -> f32 {
         let status = arena[node_idx].board_state.status;
         let stm = arena[node_idx].board_state.board.side_to_move();
@@ -348,14 +317,9 @@ impl Mcts {
         if let Some(term_val) = terminal_value(status, stm) {
             term_val
         } else {
-            let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
-            let (policy_head, value_head) = model.forward(state);
-
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
-
-            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
-            let probs = softmax(&logits);
+            let probs = softmax(policy_logits);
 
             let mut move_priors = Vec::with_capacity(legal_moves.len());
             let mut prior_sum = 0.0f32;
@@ -383,10 +347,11 @@ impl Mcts {
                 arena[node_idx].children.push(child_idx);
             }
 
-            value_head.into_data().to_vec::<f32>().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];
@@ -450,8 +415,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
     samples
 }
 
-
-
 pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
     // Choose best child of root by visit count (standard MCTS behavior)
     let root_node = &nodes[root];
@@ -540,6 +503,7 @@ pub enum BoardStateStatus {
     BlackWinner,
     Threefold,
     FiftyMove,
+    InsufficientMaterial,
 }
 
 #[derive(Clone, Debug)]
@@ -555,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,6 +1,6 @@
 use crate::mcts::BoardState;
 use burn::prelude::*;
-use chess::{ChessMove, Color, File, Piece, Rank, Square};
+use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
 
 const SLIDING_DIRS: [(i8, i8); 8] = [
     (0, 1),   // N
@@ -33,14 +33,14 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
 const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
 
 pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
-    let mut planes = vec![0.0f32; 18 * 64];
+    let mut planes = vec![0.0f32; 19 * 64];
     encode_board_state_to_slice(state, &mut planes);
-    Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
+    Tensor::<3>::from_floats(TensorData::new(planes, [19, 8, 8]), device)
 }
 
-/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
+/// 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[..1152].fill(0.0);
+    out_slice[..1216].fill(0.0);
 
     let board = &state.board;
     let us = board.side_to_move();
@@ -86,23 +86,23 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
     let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
 
     if us_castle.has_kingside() {
-        fill_plane(out_slice, 12);
+        fill_plane_ones(out_slice, 12);
     }
     if us_castle.has_queenside() {
-        fill_plane(out_slice, 13);
+        fill_plane_ones(out_slice, 13);
     }
     if them_castle.has_kingside() {
-        fill_plane(out_slice, 14);
+        fill_plane_ones(out_slice, 14);
     }
     if them_castle.has_queenside() {
-        fill_plane(out_slice, 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(out_slice, 16);
+            fill_plane_ones(out_slice, 16);
         }
     }
 
@@ -116,22 +116,29 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
         }
         out_slice[idx(17, rank, file)] = 1.0;
     }
+
+    fill_plane(out_slice, 18, Some(state.halfmove_clock as f32 / 100.0));
 }
 
 pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
     let b = states.len();
-    let mut buffer = vec![0.0f32; b * 1152];
+    let mut buffer = vec![0.0f32; b * 1216];
     for (i, state) in states.iter().enumerate() {
-        let start = i * 1152;
-        encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
+        let start = i * 1216;
+        encode_board_state_to_slice(state, &mut buffer[start..start + 1216]);
     }
-    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
+    Tensor::<4>::from_floats(TensorData::new(buffer, [b, 19, 8, 8]), device)
 }
 
 #[inline]
-fn fill_plane(buffer: &mut [f32], plane: usize) {
+fn fill_plane(buffer: &mut [f32], plane: usize, value: Option<f32>) {
     let start = plane * 64;
-    buffer[start..start + 64].fill(1.0);
+    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.
@@ -205,7 +212,7 @@ 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;
     }
@@ -242,8 +249,17 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
         return None;
     }
 
-    // Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
-    if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
+    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);
     }
 
@@ -255,69 +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() {
-        use chess::{File, Rank, Square};
+        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, Color::White).unwrap();
+        let dec = decode_move(enc, &board).unwrap();
+
         assert_eq!(mv, dec);
     }
 
     #[test]
     fn underpromotion_roundtrip() {
-        use chess::{File, Rank, Square};
+        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, 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,5 +1,4 @@
 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};
@@ -21,6 +20,7 @@ Input planes:
 16: opponent queenside
 17: rep >= 2
 18: En passant square (gap pass over)
+19: No-progress plys / 100
 
 
 Output:
@@ -45,10 +45,10 @@ pub struct ResidualBlock {
     activation: Relu,
 }
 
-#[derive(Config, Debug)]
-pub struct ResidualBlockConfig {
-    channels: usize,
-}
+// #[derive(Config, Debug)]
+// pub struct ResidualBlockConfig {
+//     channels: usize,
+// }
 
 impl ResidualBlock {
     pub fn new(channels: usize, device: &Device) -> Self {
@@ -110,7 +110,7 @@ pub struct ChessModelConfig {
 impl ChessModelConfig {
     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),
@@ -205,23 +205,30 @@ pub struct ChessOutput {
     pub value_targets: Tensor<2>, // kept 2d for consistency?
     pub policy_loss: Tensor<1>,
     pub value_loss: Tensor<1>,
-    pub loss: Tensor<1>,          // [sample]
+    pub loss: Tensor<1>, // [sample]
 }
 
 impl ItemLazy for ChessOutput {
     fn sync(self) -> Self {
-        let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
-            Transaction::default()
-                .register(self.policy_logits)
-                .register(self.value)
-                .register(self.policy_targets)
-                .register(self.value_targets)
-                .register(self.policy_loss)
-                .register(self.value_loss)
-                .register(self.loss)
-                .execute()
-                .try_into()
-                .expect("Correct amount of tensor data");
+        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()
+            .expect("Correct amount of tensor data");
 
         let device = &Default::default();
 
@@ -296,44 +303,38 @@ pub struct ChessBatch {
 
 impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
     fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
-        let state_tensors = items
-            .iter()
-            .map(|item| {
-                encode_board_state_perspective(&item.board_state, device).reshape([1, 18, 8, 8])
-            })
-            .collect::<Vec<_>>();
+        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::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
-            })
-            .collect::<Vec<_>>();
+            }
 
-        let value_target_tensors = items
-            .iter()
-            .map(|item| {
-                Tensor::<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,
@@ -342,3 +343,27 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::optim::{AdamConfig, GradientsParams};
+    use burn::tensor::{Device, DeviceIndex};
+
+    #[test]
+    fn test_train_step() {
+        let device = Device::cuda(DeviceIndex::default());
+        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/train.rs

@@ -1,22 +1,27 @@
-use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
+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::module::AutodiffModule;
 use burn::optim::{AdamConfig, GradientsParams};
 use burn::prelude::{Device, Module};
 use burn::store::ModuleRecord;
-use chess::{ChessMove, Color};
+use chess::{Board, ChessMove};
 use rand::rngs::SmallRng;
 use rand::seq::SliceRandom;
 use rand::{RngExt, SeedableRng};
-use std::collections::VecDeque;
+use std::cmp::min;
+use std::collections::{HashMap, VecDeque};
 use std::fs::File;
 use std::io::BufReader;
+
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::{mpsc, Arc};
+
+use std::thread;
 use std::time::Instant;
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -29,12 +34,18 @@ pub struct TrainingConfig {
     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>) {
@@ -92,147 +103,270 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
     let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
 
-    let mut mcts: Mcts = Mcts {
-        config: training_config.mcts_config,
-    };
-
     // Create RNG once and reuse it for sampling and shuffling
     // Seed from system time (platform default entropy may be unavailable in some contexts)
-    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
-    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
+    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);
-        for episode in 0..training_config.num_episodes {
-            println!("Episode: {}", episode);
-            let mut board_state = BoardState::default();
-            let mut episode_buffer: Vec<MctsResults> = vec![];
 
-            let mut game_hist = String::new();
+        let mut episodes_dispatched: usize = 0;
+        let mut episodes_recieved: usize = 0;
 
-            let mut ply_count: 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>();
 
-            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 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,
+            );
 
-                let temp = if ply_count < 30 { 1.0 } else { 0.0 };
+            server.run_server()
+        });
 
-                let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
+        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 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(' ');
-                }
+                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...");
 
-                board_state.apply_move(mv);
-                ply_count += 1;
-                if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
-                    board_state.status = BoardStateStatus::FiftyMove;
-                }
-            }
-            if episode == 0 {
-                println!("Game history of first game of iteration: {}", game_hist);
+                    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
+        );
 
-            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,
-                };
+        while episodes_recieved < training_config.num_episodes {
+            let data = data_rx.recv().unwrap();
 
-                replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
-                    result.clone(),
-                    outcome,
-                ));
+            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() {
-            println!("Replay buffer is empty, skipping training step.");
-            continue;
+            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 (buffer size: {})...",
+            "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 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 {};
+        // 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 batch = batcher.batch(samples, &autodiff_device);
+            let mut loss_val: f32 = 0.0;
+            let mut p_loss_val: f32 = 0.0;
+            let mut v_loss_val: f32 = 0.0;
 
-        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
+            // 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 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];
+                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);
 
-        println!(
-            "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
-            iter, loss_val, p_loss_val, v_loss_val
-        );
+                if step == training_config.num_train_steps - 1 {
+                    loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+                    p_loss_val = output
+                        .policy_loss
+                        .clone()
+                        .into_data()
+                        .to_vec::<f32>()
+                        .unwrap()[0];
+                    v_loss_val = output
+                        .value_loss
+                        .clone()
+                        .into_data()
+                        .to_vec::<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 grads = output.loss.backward();
-        let grads = GradientsParams::from_grads(grads, &model);
-
-        model = optim.step(training_config.lr, model, grads);
+        // 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;
 
-        if iter % 100 == 0 {
-            println!("Completed {} iterations", iter);
-        }
+        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()
@@ -250,33 +384,20 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
         }
     }
 
-    println!("Saving model...");
-
-    let metadata = ModelMetadata {
-        name: training_config.model_name,
-        iterations: iter,
-    };
-
-    std::fs::write(
-        metadata_path,
-        serde_json::to_string_pretty(&metadata)
-            .expect("Should be able to convert metadata to JSON string"),
-    )
-    .expect("Should be able to write metadata");
-
-    model
-        .into_record()
-        .save(&model_path)
-        .expect("Should be able to save the model");
-
-    optim
-        .save(format!("{model_path}_optim"))
-        .expect("Should be able to save the optim");
-    println!("Model saved in {:?}, exiting training.", model_path);
+    // 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;
 }
 
-fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
+pub(crate) fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
     if visits.is_empty() {
         return Vec::new();
     }
@@ -312,21 +433,21 @@ fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f
     adjusted
 }
 
-fn sample_move(
+pub(crate) fn sample_move(
     dist: &[(usize, f32)],
     rng: &mut SmallRng,
-    side_to_move: Color,
+    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, side_to_move).expect("Invalid move"));
+            return Some(decode_move(*idx, board).expect("Invalid move"));
         }
     }
 
     // fallback due to floating point drift
     dist.get(0)
-        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
+        .map(|(idx, _)| decode_move(*idx, board).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().to_vec::<f32>().unwrap();
+            let v_vec = v.into_data().to_vec::<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;