improvements

Cargo.lock

@@ -2170,7 +2170,7 @@ dependencies = [
  "chess",
  "ctrlc",
  "futures",
- "moka",
+ "lru 0.12.5",
  "rand 0.10.1",
  "rand_distr",
  "serde",
@@ -3531,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"
@@ -3674,23 +3683,6 @@ version = "0.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b"
 
-[[package]]
-name = "moka"
-version = "0.12.16"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9"
-dependencies = [
- "crossbeam-channel",
- "crossbeam-epoch",
- "crossbeam-utils",
- "equivalent",
- "parking_lot",
- "portable-atomic",
- "smallvec",
- "tagptr",
- "uuid",
-]
-
 [[package]]
 name = "moxcms"
 version = "0.8.1"
@@ -4726,7 +4718,7 @@ dependencies = [
  "indoc",
  "itertools",
  "kasuari",
- "lru",
+ "lru 0.16.4",
  "strum 0.27.2",
  "thiserror 2.0.18",
  "unicode-segmentation",
@@ -5643,12 +5635,6 @@ dependencies = [
  "windows",
 ]
 
-[[package]]
-name = "tagptr"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
-
 [[package]]
 name = "tar"
 version = "0.4.46"

engine/Cargo.toml

@@ -13,6 +13,6 @@ serde_json = "1.0.150"
 serde = { version = "1.0.228", features = ["derive"] }
 tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread", "time"] }
 
-moka = { version = "0.12.16", features = ["sync"] }
+lru = "0.12"
 futures = "0.3"
 

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(400, 1.0, 0.05, 0.25, 16);
+    // 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,23 +35,23 @@ fn main() {
         max_time_s: None,
         num_iters: None,
         max_depth: 100, // unused
-        model_name: String::from("96"),
+        model_name: String::from("128x8"),
-        load_model: true,
+        load_model: true, // architecture changed — must start with fresh weights
-        hidden_channels: 48,
+        hidden_channels: 128,
-        num_blocks: 4,
+        num_blocks: 8,
-        batch_size: 1024, // num positions sampled to a batch (2048)
+        batch_size: 1024,
         minibatch_size: 256,
-        num_train_steps: 10,
+        num_train_steps: 40,
-        num_episodes: 96,         // num games generated per iteration (5000)
+        num_episodes: 48,
-        buffer_max_size: 200_000, // max number of samples in the buffer
+        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: 12,
+        episodes_per_worker: 6,
-        nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB.
+        nn_cache_size: 1000_000,
     };
 
     train(training_config, device);

engine/src/mcts.rs

@@ -185,6 +185,13 @@ impl Mcts {
         let (root_policy, root_value) = model_inference(&board_state).await;
         self.expand_with_eval(root, &mut nodes, &root_policy, root_value);
 
+        // 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);
@@ -244,14 +251,19 @@ impl Mcts {
                 }
             }
 
-            // Revert virtual loss, expand leaves, and backpropagate
+            // Phase 1: Revert ALL virtual losses before touching the tree with real values.
-            for (path, leaf_idx, term_val) in pending {
+            // Doing this in one pass ensures that when we backpropagate simulation N,
-                // Revert virtual loss
+            // no other simulation's artificial visit counts or value sums are polluting
-                for &idx in &path {
+            // 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)

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};
@@ -304,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
+        let b = items.len();
-            .iter()
-            .map(|item| {
-                encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 8, 8])
-            })
-            .collect::<Vec<_>>();
 
-        let policy_target_tensors = items
+        let mut states_buf = vec![0.0f32; b * 1216];
-            .iter()
+        let mut policy_buf = vec![0.0f32; b * 4672];
-            .cloned()
+        let mut value_buf = vec![0.0f32; b];
-            .map(|item| {
+
-                let mut policy = vec![0.0f32; 4672];
+        for (i, item) in items.iter().enumerate() {
-                for (idx, prob) in item.policy_target.iter() {
+            crate::net::encoding::encode_board_state_to_slice(
-                    policy[*idx] = *prob;
+                &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_slice.iter().sum();
-                let sum: f32 = policy.iter().sum();
+            if sum > 0.0 {
-                if sum > 0.0 {
+                for p in policy_slice.iter_mut() {
-                    for p in &mut policy {
+                    *p /= sum;
-                        *p /= sum;
-                    }
                 }
-                Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
+            }
-            })
-            .collect::<Vec<_>>();
 
-        let value_target_tensors = items
+            value_buf[i] = item.value_target;
-            .iter()
+        }
-            .map(|item| {
-                Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device)
-                    .reshape([1, 1])
-            })
-            .collect::<Vec<_>>();
 
-        let states = Tensor::cat(state_tensors, 0); // [B, 19, 8, 8]
+        let states = Tensor::<4>::from_floats(TensorData::new(states_buf, [b, 19, 8, 8]), device);
-        let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672]
+        let policy_targets =
-        let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1]
+            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,
@@ -350,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/worker.rs

@@ -4,9 +4,9 @@ use crate::train::{apply_temperature, sample_move};
 use burn::tensor::Device;
 use chess::Color;
 
-use moka::sync::Cache;
+use lru::LruCache;
-use rand::prelude::SmallRng;
 use rand::SeedableRng;
+use rand::prelude::SmallRng;
 use std::collections::VecDeque;
 use std::hash::Hash;
 use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
@@ -83,8 +83,12 @@ impl Worker {
             // println!("{}", before.elapsed().as_millis());
             episode_buffer.push(results);
 
-            let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3);
+            // AlphaZero temperature schedule: explore freely for the first 30 plies,
-            // let temp = if ply_count < 30 { 1.0 } else { 0.2 };
+            // 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);
 
@@ -125,7 +129,8 @@ impl Worker {
                 _ => 0.0,
             };
 
-            let blended_target = 0.75 * outcome + 0.25 * result.value;
+            let blended_target = outcome;
+            // let blended_target = 0.75 * outcome + 0.25 * result.value;
 
             samples.push(TrainingSample::from_mcts_with_outcome(
                 result.clone(),
@@ -235,7 +240,6 @@ impl Worker {
 
         println!("Worker {} shut down", self.id);
     }
-
 }
 
 async fn model_inference_static(
@@ -307,6 +311,7 @@ pub struct NNRequest {
     response_tx: tokio::sync::oneshot::Sender<NNResponse>,
 }
 
+// ~18.7KB
 #[derive(Debug, Clone)]
 pub struct NNResponse {
     policy: Vec<f32>,
@@ -315,7 +320,9 @@ pub struct NNResponse {
 
 impl NNServer {
     pub fn run_server(&self) {
-        let mut cache: Cache<NNCacheKey, NNResponse> = Cache::new(self.cache_size);
+        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;
@@ -327,7 +334,7 @@ impl NNServer {
                 match cmd {
                     NNServerCmd::Shutdown => {
                         println!(
-                            "NN server shutting down. Served {} requests, with {} cache hits ({:2}%). Ran {} batches (avg {}requests/batch).",
+                            "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,
@@ -346,6 +353,7 @@ impl NNServer {
                 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;
@@ -365,9 +373,9 @@ impl NNServer {
 
             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;
-                    continue;
                 }
             }
 
@@ -389,7 +397,10 @@ impl NNServer {
             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]);
+                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(
@@ -414,7 +425,7 @@ impl NNServer {
                     value: v_val,
                 };
 
-                cache.insert(job.cache_key.clone(), res.clone());
+                cache.put(job.cache_key, res.clone());
 
                 let _ = job.response_tx.send(res);
             }
@@ -424,12 +435,12 @@ impl NNServer {
 
 fn try_use_cache(
     req: NNRequest,
-    cache: &mut Cache<NNCacheKey, NNResponse>,
+    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).unwrap();
+        req.response_tx.send(cached.clone()).unwrap();
         true
     } else {
         queue.push_back(req);