drake/chess_dragon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

improvements

Browse Source
This commit is contained in:
drake
2026-10-17 15:34:35 -04:00
parent 38f13fa299
commit 03f11e2ed8
6 changed files with 126 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/Cargo.toml
+1 -1
View File
@@ -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
+23 -12
View File
@@ -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"),
load_model: true,
hidden_channels: 48,
num_blocks: 4,
batch_size: 1024, // num positions sampled to a batch (2048)
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: 10,
num_episodes: 96, // num games generated per iteration (5000)
buffer_max_size: 200_000, // max number of samples in the buffer
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: 12,
nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB.
episodes_per_worker: 6,
nn_cache_size: 1000_000,
};
train(training_config, device);
engine/src/mcts.rs
+16 -4
View File
@@ -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
for (path, leaf_idx, term_val) in pending {
// Revert virtual loss
for &idx in &path {
// 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)
engine/src/net/model.rs
+51 -34
View File
@@ -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
.iter()
.map(|item| {
encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 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, 19, 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,
@@ -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
+24 -13
View File
@@ -4,9 +4,9 @@ use crate::train::{apply_temperature, sample_move};
use burn::tensor::Device;
use chess::Color;
use moka::sync::Cache;
use rand::prelude::SmallRng;
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};
@@ -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);
// let temp = if ply_count < 30 { 1.0 } else { 0.2 };
// 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);
@@ -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);