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

misc fixes

Browse Source
This commit is contained in:
drake
2026-08-24 23:19:51 -04:00
parent c40a9667bd
commit 33dc15a597
11 changed files with 1742 additions and 1262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/Cargo.toml
+1 -1
View File
@@ -5,7 +5,7 @@ edition = "2024"
[dependencies]
chess = "3.2.0"
burn = { version = "0.21.0", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
burn = { version = "0.22.0-pre.2", features = ["std", "tui", "train", "cuda", "fusion", "wgpu", "ndarray"], default-features = false }
burn-ndarray = "0.21.0"
rand_distr = "0.6.0"
rand = "0.10.1"
engine/src/lib.rs
+1 -2
View File
@@ -1,7 +1,6 @@
pub mod mcts;
mod net;
mod train;
pub mod training;
pub mod train;
use chess::Game;
engine/src/main.rs
+4 -10
View File
@@ -1,9 +1,9 @@
#![recursion_limit = "256"]
use burn::backend::{Autodiff, Cuda};
use burn::optim::AdamConfig;
use burn::tensor::{Device, DeviceIndex};
use engine::mcts::MctsConfig;
use engine::training::train::{train, TrainingConfig};
use engine::train::{train, TrainingConfig};
// fn main() {
// type MyBackend = Wgpu<f32, i32>;
//
@@ -14,13 +14,7 @@ use engine::training::train::{train, TrainingConfig};
// }
fn main() {
// type MyBackend = Wgpu<f32, i32>;
type MyBackend = Cuda<f32, i32>;
// type MyBackend = NdArray<f32, i32>;
type MyAutodiffBackend = Autodiff<MyBackend>;
// let device = burn::backend::wgpu::WgpuDevice::default();
// let device = burn::backend::ndarray::NdArrayDevice::default();
let device = burn::backend::cuda::CudaDevice::default();
let device = Device::cuda(DeviceIndex::default());
let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
@@ -43,7 +37,7 @@ fn main() {
seed: None,
};
train::<MyAutodiffBackend>(training_config, device);
train(training_config, &device);
// let mut board_state: Option<BoardState> = Some(BoardState::new(
// Board::default(),
engine/src/mcts.rs
+164 -110
View File
@@ -2,7 +2,7 @@
use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
use crate::net::encoding::{encode_board_state_perspective, encode_move};
use crate::net::model::ChessModel;
use burn::prelude::Backend;
use burn::prelude::Device;
use burn::Tensor;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
@@ -10,7 +10,6 @@ use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;
use std::marker::PhantomData;
pub struct Node {
pub prior: f32,
@@ -107,17 +106,62 @@ impl Default for MctsConfig {
}
}
pub struct Mcts<B: Backend> {
pub struct Mcts {
pub config: MctsConfig,
pub _marker: PhantomData<B>, // if B not otherwise stored
}
impl<B: Backend> Mcts<B> {
fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
match status {
BoardStateStatus::Ongoing => None,
BoardStateStatus::Stalemate
| BoardStateStatus::Threefold
| BoardStateStatus::FiftyMove => Some(0.0),
BoardStateStatus::WhiteWinner => Some(if stm == White { 1.0 } else { -1.0 }),
BoardStateStatus::BlackWinner => Some(if stm == Black { 1.0 } else { -1.0 }),
}
}
fn softmax(logits: &[f32]) -> Vec<f32> {
if logits.is_empty() {
return Vec::new();
}
let mut max_logit = f32::NEG_INFINITY;
for &v in logits {
if v > max_logit {
max_logit = v;
}
}
if !max_logit.is_finite() {
max_logit = 0.0;
}
let mut exps_sum = 0.0f32;
let mut probs = vec![0.0f32; logits.len()];
for (i, &v) in logits.iter().enumerate() {
let x = (v - max_logit).exp();
if x.is_finite() {
probs[i] = x;
exps_sum += x;
}
}
if exps_sum > 0.0 && exps_sum.is_finite() {
for p in &mut probs {
*p /= exps_sum;
}
} else {
let uniform = 1.0 / logits.len() as f32;
for p in &mut probs {
*p = uniform;
}
}
probs
}
impl Mcts {
pub fn search(
&mut self,
board_state: &BoardState,
model: &ChessModel<B>,
device: &B::Device,
model: &ChessModel,
device: &Device,
) -> MctsResults {
let mut nodes = Vec::<Node>::new();
@@ -130,133 +174,138 @@ impl<B: Backend> Mcts<B> {
// Apply Dirichlet noise to root children
self.add_dirichlet_noise(root, &mut nodes);
// We'll batch leaf evaluations to reduce per-leaf model calls and device-host syncs.
let mut sims_done: usize = 0;
let num_sims = self.config.num_simulations;
while sims_done < num_sims {
// Collect a batch of leaf nodes (and their selection paths)
let mut leaf_nodes: Vec<usize> = Vec::new();
let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
let mut leaf_states: Vec<Tensor<B, 4>> = Vec::new();
let mut leaf_states: Vec<Tensor<4>> = Vec::new();
while leaf_nodes.len() < std::cmp::min(self.config.batch_max, num_sims - sims_done) {
let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
while leaf_nodes.len() < target_batch && sims_done + leaf_nodes.len() < num_sims {
let mut path = vec![root];
let mut current = root;
// Apply virtual loss during selection path
while !nodes[current].children.is_empty() {
nodes[current].visit_count += 1;
nodes[current].value_sum -= 1.0;
current = nodes[current].select_child(&nodes, &self.config.c_puct);
path.push(current);
}
// Record leaf node and its path
leaf_nodes.push(current);
leaf_paths.push(path.clone());
// 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 except leaf itself
for &idx in &path[..path.len() - 1] {
nodes[idx].visit_count -= 1;
nodes[idx].value_sum += 1.0;
}
self.backpropagate(&mut nodes, &path, term_val, stm);
sims_done += 1;
continue;
}
// Prepare state tensor for this leaf
let state: Tensor<B, 4> =
// 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;
}
break;
}
leaf_nodes.push(current);
leaf_paths.push(path);
let state: Tensor<4> =
encode_board_state_perspective(&nodes[current].board_state, device)
.reshape([1, 18, 8, 8]);
leaf_states.push(state);
sims_done += 1;
}
if leaf_nodes.is_empty() {
break;
if sims_done >= num_sims {
break;
}
continue;
}
// Batch evaluate the collected leaf states
// Batch evaluate non-terminal leaf states
let batch = Tensor::cat(leaf_states, 0);
let (policy_batch, value_batch) = model.forward(batch);
// Move tensors to host once per batch
let policy_data = policy_batch.into_data().to_vec::<f32>().unwrap();
let value_data = value_batch.into_data().to_vec::<f32>().unwrap();
let num_moves = policy_data.len() / leaf_nodes.len();
// Process each evaluated leaf: expand and backpropagate
for (i, &node_idx) in leaf_nodes.iter().enumerate() {
let path = &leaf_paths[i];
// slice for this sample's logits
// Revert virtual loss along this path before applying true expansion & backpropagation
for &idx in &path[..path.len() - 1] {
nodes[idx].visit_count -= 1;
nodes[idx].value_sum += 1.0;
}
let start = i * num_moves;
let end = start + num_moves;
let logits = &policy_data[start..end];
let probs = softmax(logits);
// Convert logits to probabilities with a numerically-stable softmax on host
let mut max_logit = f32::NEG_INFINITY;
for &v in logits.iter() {
if v > max_logit {
max_logit = v;
}
}
if logits.is_empty() {
println!("logits empty")
}
if !max_logit.is_finite() {
println!("max logits not finite")
}
let mut exps_sum = 0.0f32;
let mut probs = vec![0.0f32; num_moves];
for (j, &v) in logits.iter().enumerate() {
let x = (v - max_logit).exp();
if !x.is_finite() {
probs[j] = 0.0;
} else {
probs[j] = x;
exps_sum += x;
}
}
if exps_sum > 0.0 && exps_sum.is_finite() {
for p in &mut probs {
*p /= exps_sum;
}
} else {
// fallback: uniform or zero
let uniform = 1.0 / num_moves as f32;
for p in &mut probs {
*p = uniform;
}
}
// Expand: add legal moves as children with prior from probs
let stm = nodes[node_idx].board_state.board.side_to_move();
let legal_moves: Vec<ChessMove> =
MoveGen::new_legal(&nodes[node_idx].board_state.board).collect();
for mv in legal_moves {
let stm = nodes[node_idx].board_state.board.side_to_move();
let idx = encode_move(mv, stm).expect("Invalid move");
let prior = probs[idx];
let mut move_priors = Vec::with_capacity(legal_moves.len());
let mut prior_sum = 0.0f32;
for &mv in &legal_moves {
if let Some(enc_idx) = encode_move(mv, stm) {
let p = probs[enc_idx];
move_priors.push((mv, p));
prior_sum += p;
}
}
for (mv, p) in move_priors {
let normalized_prior = if prior_sum > 0.0 {
p / prior_sum
} else {
1.0 / legal_moves.len() as f32
};
let mut new_board = nodes[node_idx].board_state.clone();
new_board.apply_move(mv);
let child_idx = nodes.len();
nodes.push(Node::new(prior, new_board, Some(mv)));
nodes.push(Node::new(normalized_prior, new_board, Some(mv)));
nodes[node_idx].children.push(child_idx);
}
// Backpropagate the value for this leaf
let value = value_data[i];
let color = nodes[node_idx].board_state.board.side_to_move();
self.backpropagate(&mut nodes, path, value, color);
self.backpropagate(&mut nodes, path, value, stm);
sims_done += 1;
}
}
// Build compact move distribution: encoded move index -> probability (visits / num_simulations)
// Build compact move distribution
let mut move_dist: Vec<(usize, f32)> = Vec::with_capacity(nodes[root].children.len());
let stm = board_state.board.side_to_move();
let denom = self.config.num_simulations as f32;
let total_root_visits: u32 = nodes[root].children.iter().map(|&c| nodes[c].visit_count).sum();
let denom = if total_root_visits > 0 { total_root_visits as f32 } else { 1.0 };
for child_idx in nodes[root].children.iter() {
let mv = nodes[*child_idx].last_move.expect("move didnt exist");
let enc = encode_move(mv, stm).expect("Invalid move");
let prob = nodes[*child_idx].visit_count as f32 / denom;
move_dist.push((enc, prob));
if let Some(enc) = encode_move(mv, stm) {
let prob = nodes[*child_idx].visit_count as f32 / denom;
move_dist.push((enc, prob));
}
}
MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
@@ -266,54 +315,53 @@ impl<B: Backend> Mcts<B> {
&mut self,
node_idx: usize,
arena: &mut Vec<Node>,
model: &ChessModel<B>,
device: &B::Device,
model: &ChessModel,
device: &Device,
) -> f32 {
if arena[node_idx].board_state.status == BoardStateStatus::Stalemate
|| arena[node_idx].board_state.status == BoardStateStatus::Threefold
|| arena[node_idx].board_state.status == BoardStateStatus::FiftyMove
{
0.0
} else if arena[node_idx].board_state.status == BoardStateStatus::WhiteWinner {
if arena[node_idx].board_state.board.side_to_move() == Black {
1.0
} else {
-1.0
}
} else if arena[node_idx].board_state.status == BoardStateStatus::BlackWinner {
if arena[node_idx].board_state.board.side_to_move() == White {
1.0
} else {
-1.0
}
let status = arena[node_idx].board_state.status;
let stm = arena[node_idx].board_state.board.side_to_move();
if let Some(term_val) = terminal_value(status, stm) {
term_val
} else {
let state: Tensor<B, 4> =
let state: Tensor<4> =
encode_board_state_perspective(&arena[node_idx].board_state, device)
.reshape([1, 18, 8, 8]);
// let start = Instant::now();
let (policy_head, value_head) = model.forward(state);
// println!("time: {:?}", start.elapsed());
let legal_moves: Vec<ChessMove> =
MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
let policy = policy_head.into_data().to_vec::<f32>().unwrap();
let logits = policy_head.into_data().to_vec::<f32>().unwrap();
let probs = softmax(&logits);
for mv in legal_moves {
let stm = arena[node_idx].board_state.board.side_to_move();
let idx = encode_move(mv, stm).expect("Invalid move");
let prior = policy[idx];
let mut move_priors = Vec::with_capacity(legal_moves.len());
let mut prior_sum = 0.0f32;
for &mv in &legal_moves {
if let Some(enc_idx) = encode_move(mv, stm) {
let p = probs[enc_idx];
move_priors.push((mv, p));
prior_sum += p;
}
}
for (mv, p) in move_priors {
let normalized_prior = if prior_sum > 0.0 {
p / prior_sum
} else {
1.0 / legal_moves.len() as f32
};
let mut new_board = arena[node_idx].board_state.clone();
new_board.apply_move(mv);
let child_idx = arena.len();
arena.push(Node::new(prior, new_board, Some(mv)));
arena.push(Node::new(normalized_prior, new_board, Some(mv)));
arena[node_idx].children.push(child_idx);
}
value_head.into_data().to_vec().unwrap()[0]
value_head.into_data().to_vec::<f32>().unwrap()[0]
}
}
@@ -356,7 +404,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
// Use a single SmallRng seeded from system time (avoid depending on thread_rng helper)
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap();
@@ -367,8 +414,15 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
let sum: f32 = samples.iter().sum();
for x in &mut samples {
*x /= sum;
if sum > 0.0 && sum.is_finite() {
for x in &mut samples {
*x /= sum;
}
} else if size > 0 {
let uniform = 1.0 / size as f32;
for x in &mut samples {
*x = uniform;
}
}
samples
engine/src/net/encoding.rs
+33 -9
View File
@@ -52,10 +52,7 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
pub fn encode_board_state_perspective<B: Backend>(
state: &BoardState,
device: &B::Device,
) -> Tensor<B, 3> {
pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
let mut planes = vec![0.0f32; 18 * 64];
let board = &state.board;
@@ -142,7 +139,7 @@ pub fn encode_board_state_perspective<B: Backend>(
planes[idx(17, rank, file)] = 1.0;
}
Tensor::<B, 3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
}
fn fill_plane(buffer: &mut [f32], plane: usize) {
@@ -201,7 +198,7 @@ pub fn encode_move(mv: ChessMove, stm: Color) -> Option<usize> {
};
let from_sq = (fr as usize) * 8 + (ff as usize);
Some(from_sq * PLANES_PER_SQUARE + plane)
Some(plane * 64 + from_sq)
}
fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
@@ -228,13 +225,13 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
return None;
}
let from_idx = idx / PLANES_PER_SQUARE;
let plane = idx % PLANES_PER_SQUARE;
let plane = idx / 64;
let from_idx = idx % 64;
let ff = (from_idx % 8) as i8;
let fr = (from_idx / 8) as i8;
let (dx, dy, promo) = if plane < 56 {
let (dx, dy, mut promo) = if plane < 56 {
let dir = plane / 7;
let dist = (plane % 7 + 1) as i8;
@@ -260,6 +257,11 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
return None;
}
// Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
promo = Some(Piece::Queen);
}
let from = deorient_square(ff, fr, stm);
let to = deorient_square(tf, tr, stm);
@@ -311,4 +313,26 @@ mod tests {
}
}
}
#[test]
fn queen_promotion_roundtrip() {
use chess::{File, Rank, Square};
let from = Square::make_square(Rank::Seventh, File::E);
let to = Square::make_square(Rank::Eighth, File::E);
let mv = ChessMove::new(from, to, Some(Piece::Queen));
let enc = encode_move(mv, Color::White).unwrap();
let dec = decode_move(enc, Color::White).unwrap();
assert_eq!(mv, dec);
}
#[test]
fn underpromotion_roundtrip() {
use chess::{File, Rank, Square};
let from = Square::make_square(Rank::Seventh, File::E);
let to = Square::make_square(Rank::Eighth, File::E);
let mv = ChessMove::new(from, to, Some(Piece::Knight));
let enc = encode_move(mv, Color::White).unwrap();
let dec = decode_move(enc, Color::White).unwrap();
assert_eq!(mv, dec);
}
}
engine/src/net/model.rs
+76 -78
View File
@@ -5,14 +5,12 @@ use burn::nn::conv::Conv2dConfig;
use burn::nn::loss::{MseLoss, Reduction};
use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
use burn::tensor::activation::log_softmax;
use burn::tensor::backend::AutodiffBackend;
use burn::tensor::Transaction;
use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
use burn::{
nn::{conv::Conv2d, Linear, Relu},
prelude::*,
};
use burn_ndarray::NdArray;
/*
Input planes:
1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -39,11 +37,11 @@ pub struct ModelMetadata {
}
#[derive(Module, Debug)]
pub struct ResidualBlock<B: Backend> {
conv1: Conv2d<B>,
bn1: BatchNorm<B>,
conv2: Conv2d<B>,
bn2: BatchNorm<B>,
pub struct ResidualBlock {
conv1: Conv2d,
bn1: BatchNorm,
conv2: Conv2d,
bn2: BatchNorm,
activation: Relu,
}
@@ -52,8 +50,8 @@ pub struct ResidualBlockConfig {
channels: usize,
}
impl<B: Backend> ResidualBlock<B> {
pub fn new(channels: usize, device: &B::Device) -> Self {
impl ResidualBlock {
pub fn new(channels: usize, device: &Device) -> Self {
Self {
conv1: Conv2dConfig::new([channels, channels], [3, 3])
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -66,7 +64,7 @@ impl<B: Backend> ResidualBlock<B> {
activation: Relu::new(),
}
}
pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
let residual = x.clone();
let x = self.conv1.forward(x);
@@ -81,24 +79,24 @@ impl<B: Backend> ResidualBlock<B> {
}
#[derive(Module, Debug)]
pub struct ChessModel<B: Backend> {
pub struct ChessModel {
// Initial convolution
conv: Conv2d<B>,
bn: BatchNorm<B>,
conv: Conv2d,
bn: BatchNorm,
// Residual tower
residual_blocks: Vec<ResidualBlock<B>>,
residual_blocks: Vec<ResidualBlock>,
// Policy head
policy_conv: Conv2d<B>,
policy_bn: BatchNorm<B>,
policy_fc: Linear<B>,
policy_conv: Conv2d,
policy_bn: BatchNorm,
policy_conv2: Conv2d,
// Value head
value_conv: Conv2d<B>,
value_bn: BatchNorm<B>,
value_fc1: Linear<B>,
value_fc2: Linear<B>,
value_conv: Conv2d,
value_bn: BatchNorm,
value_fc1: Linear,
value_fc2: Linear,
activation: Relu,
}
@@ -110,11 +108,7 @@ pub struct ChessModelConfig {
}
impl ChessModelConfig {
pub fn init<B: Backend>(
num_blocks: usize,
channels: usize,
device: &B::Device,
) -> ChessModel<B> {
pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
ChessModel {
conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -126,14 +120,14 @@ impl ChessModelConfig {
.collect(),
// Policy head
policy_conv: Conv2dConfig::new([channels, 2], [1, 1]).init(device),
policy_bn: BatchNormConfig::new(2).init(device),
policy_fc: LinearConfig::new(2 * 8 * 8, 8 * 8 * 73).init(device), // 4672 typical chess move space
policy_conv: Conv2dConfig::new([channels, 32], [1, 1]).init(device),
policy_bn: BatchNormConfig::new(32).init(device),
policy_conv2: Conv2dConfig::new([32, 73], [1, 1]).init(device),
// Value head
value_conv: Conv2dConfig::new([channels, 1], [1, 1]).init(device),
value_bn: BatchNormConfig::new(1).init(device),
value_fc1: LinearConfig::new(1 * 8 * 8, 256).init(device),
value_fc1: LinearConfig::new(8 * 8, 256).init(device),
value_fc2: LinearConfig::new(256, 1).init(device),
activation: Relu::new(),
@@ -141,8 +135,8 @@ impl ChessModelConfig {
}
}
impl<B: Backend> ChessModel<B> {
pub fn forward(&self, x: Tensor<B, 4>) -> (Tensor<B, 2>, Tensor<B, 2>) {
impl ChessModel {
pub fn forward(&self, x: Tensor<4>) -> (Tensor<2>, Tensor<2>) {
let mut x = self.conv.forward(x);
x = self.bn.forward(x);
x = self.activation.forward(x);
@@ -157,8 +151,8 @@ impl<B: Backend> ChessModel<B> {
let mut p = self.policy_conv.forward(x.clone());
p = self.policy_bn.forward(p);
p = self.activation.forward(p);
let mut p = p.reshape([batch_size, 2 * 8 * 8]);
p = self.policy_fc.forward(p);
p = self.policy_conv2.forward(p);
let p = p.reshape([batch_size, 73 * 8 * 8]);
// -------- Value Head --------
let mut v = self.value_conv.forward(x);
@@ -173,10 +167,10 @@ impl<B: Backend> ChessModel<B> {
pub fn forward_chess(
&self,
boards: Tensor<B, 4>, // e.g. [batch, channels, 8, 8]
policy_targets: Tensor<B, 2>, // move distribution (make sure is normalized)
value_targets: Tensor<B, 2>, // scalar evaluation
) -> ChessOutput<B> {
boards: Tensor<4>, // e.g. [batch, channels, 8, 8]
policy_targets: Tensor<2>, // move distribution (make sure is normalized)
value_targets: Tensor<2>, // scalar evaluation
) -> ChessOutput {
let (policy_logits, value) = self.forward(boards);
let log_probs = log_softmax(policy_logits.clone(), 1);
@@ -190,39 +184,44 @@ impl<B: Backend> ChessModel<B> {
let value_loss =
MseLoss::new().forward(value.clone(), value_targets.clone(), Reduction::Mean);
let total_loss = policy_loss + 0.5 * value_loss;
let total_loss = policy_loss.clone() + 0.5 * value_loss.clone();
ChessOutput {
policy_logits,
value,
policy_targets,
value_targets,
policy_loss,
value_loss,
loss: total_loss,
}
}
}
pub struct ChessOutput<B: Backend> {
pub policy_logits: Tensor<B, 2>, // [sample, num_moves (4672)]
pub value: Tensor<B, 2>, // [sample, 1 value]
pub policy_targets: Tensor<B, 2>,
pub value_targets: Tensor<B, 2>, // kept 2d for consistency?
pub loss: Tensor<B, 1>, // [sample]
pub struct ChessOutput {
pub policy_logits: Tensor<2>, // [sample, num_moves (4672)]
pub value: Tensor<2>, // [sample, 1 value]
pub policy_targets: Tensor<2>,
pub value_targets: Tensor<2>, // kept 2d for consistency?
pub policy_loss: Tensor<1>,
pub value_loss: Tensor<1>,
pub loss: Tensor<1>, // [sample]
}
impl<B: Backend> ItemLazy for ChessOutput<B> {
type ItemSync = ChessOutput<NdArray>;
fn sync(self) -> Self::ItemSync {
let [policy_logits, value, policy_targets, value_targets, loss] = Transaction::default()
.register(self.policy_logits)
.register(self.value)
.register(self.policy_targets)
.register(self.value_targets)
.register(self.loss)
.execute()
.try_into()
.expect("Correct amount of tensor data");
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 device = &Default::default();
@@ -231,27 +230,29 @@ impl<B: Backend> ItemLazy for ChessOutput<B> {
value: Tensor::from_data(value, device),
policy_targets: Tensor::from_data(policy_targets, device),
value_targets: Tensor::from_data(value_targets, device),
policy_loss: Tensor::from_data(policy_loss, device),
value_loss: Tensor::from_data(value_loss, device),
loss: Tensor::from_data(loss, device),
}
}
}
impl<B: AutodiffBackend> TrainStep for ChessModel<B> {
type Input = ChessBatch<B>;
type Output = ChessOutput<B>;
impl TrainStep for ChessModel {
type Input = ChessBatch;
type Output = ChessOutput;
fn step(&self, batch: ChessBatch<B>) -> TrainOutput<ChessOutput<B>> {
fn step(&self, batch: ChessBatch) -> TrainOutput<ChessOutput> {
let item = self.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
TrainOutput::new(self, item.loss.backward(), item)
}
}
impl<B: Backend> InferenceStep for ChessModel<B> {
type Input = ChessBatch<B>;
type Output = ChessOutput<B>;
impl InferenceStep for ChessModel {
type Input = ChessBatch;
type Output = ChessOutput;
fn step(&self, batch: ChessBatch<B>) -> ChessOutput<B> {
fn step(&self, batch: ChessBatch) -> ChessOutput {
self.forward_chess(batch.states, batch.policy_targets, batch.value_targets)
}
}
@@ -287,14 +288,14 @@ impl TrainingSample {
pub struct ChessBatcher {}
#[derive(Clone, Debug)]
pub struct ChessBatch<B: Backend> {
pub states: Tensor<B, 4>,
pub policy_targets: Tensor<B, 2>,
pub value_targets: Tensor<B, 2>,
pub struct ChessBatch {
pub states: Tensor<4>,
pub policy_targets: Tensor<2>,
pub value_targets: Tensor<2>,
}
impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
fn batch(&self, items: Vec<TrainingSample>, device: &B::Device) -> ChessBatch<B> {
impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
let state_tensors = items
.iter()
.map(|item| {
@@ -318,18 +319,15 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
*p /= sum;
}
}
Tensor::<B, 2>::from_floats(TensorData::new(policy, [1, 4672]), device).unsqueeze()
Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
})
.collect::<Vec<_>>();
let value_target_tensors = items
.iter()
.map(|item| {
Tensor::<B, 2>::from_floats(
TensorData::new(vec![item.value_target], [1, 1]),
device,
)
.reshape([1, 1])
Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device)
.reshape([1, 1])
})
.collect::<Vec<_>>();
engine/src/rl/classification.rs
-59
View File
@@ -1,59 +0,0 @@
/// Multi-label classification output adapted for multiple metrics.
///
/// Supported metrics:
/// - HammingScore
/// - Precision (via ConfusionStatsInput)
/// - Recall (via ConfusionStatsInput)
/// - FBetaScore (via ConfusionStatsInput)
/// - Loss
#[derive(new)]
pub struct MultiLabelSoftClassificationOutput<B: Backend> {
/// The loss.
pub loss: Tensor<B, 1>,
/// The label logits or probabilities. Shape: \[batch_size, num_classes\].
pub output: Tensor<B, 2>,
/// The ground truth labels (target values). Shape: \[batch_size, num_classes\].
pub targets: Tensor<B, 2>,
}
impl<B: Backend> ItemLazy for MultiLabelSoftClassificationOutput<B> {
type ItemSync = MultiLabelSoftClassificationOutput<NdArray>;
fn sync(self) -> Self::ItemSync {
let [output, loss, targets] = Transaction::default()
.register(self.output)
.register(self.loss)
.register(self.targets)
.execute()
.try_into()
.expect("Correct amount of tensor data");
let device = &Default::default();
MultiLabelSoftClassificationOutput {
output: Tensor::from_data(output, device),
loss: Tensor::from_data(loss, device),
targets: Tensor::from_data(targets, device),
}
}
}
impl<B: Backend> Adaptor<HammingScoreInput<B>> for MultiLabelSoftClassificationOutput<B> {
fn adapt(&self) -> HammingScoreInput<B> {
HammingScoreInput::new(self.output.clone(), self.targets.clone())
}
}
impl<B: Backend> Adaptor<LossInput<B>> for MultiLabelSoftClassificationOutput<B> {
fn adapt(&self) -> LossInput<B> {
LossInput::new(self.loss.clone())
}
}
impl<B: Backend> Adaptor<ConfusionStatsInput<B>> for MultiLabelSoftClassificationOutput<B> {
fn adapt(&self) -> ConfusionStatsInput<B> {
ConfusionStatsInput::new(self.output.clone(), self.targets.clone().bool())
}
}
engine/src/train.rs
+312 -132
View File
@@ -1,134 +1,314 @@
// use crate::model::{ChessBatcher, ChessModel, ChessModelConfig};
// use burn::config::Config;
// use burn::data::dataloader::batcher::Batcher;
// use burn::optim::{Adam, AdamConfig, Optimizer, SimpleOptimizer};
// use burn::optim::adaptor::OptimizerAdaptor;
// use burn::prelude::Backend;
// use burn::tensor::backend::AutodiffBackend;
// use burn::train::TrainStep;
// use crate::mcts::MctsResults;
//
// #[derive(Config, Debug)]
// pub struct ChessTrainerConfig {
// pub model: ChessModelConfig,
// pub optimizer: AdamConfig,
// #[config(default = 10)]
// pub num_epochs: usize,
// #[config(default = 64)]
// pub batch_size: usize,
// #[config(default = 4)]
// pub num_workers: usize,
// #[config(default = 42)]
// pub seed: u64,
// #[config(default = 1.0e-4)]
// pub learning_rate: f64,
// }
//
// impl ChessTrainerConfig {
// pub fn init<B: Backend>(
// model_config: ChessModelConfig,
// optimizer: AdamConfig,
// num_epochs: usize,
// batch_size: usize,
// num_workers: usize,
// seed: u64,
// learning_rate: f64,
// ) -> ChessTrainer<B> {
// ChessTrainer {
// model: model_config::init(),
// optimizer: optimizer.init(),
// num_epochs,
//
//
// }
// }
// }
//
// pub struct ChessTrainer<B: AutodiffBackend> {
// pub model: ChessModel<B>,
// pub optimizer: Adam,
// learning_rate: f64,
// pub batcher: ChessBatcher,
// pub device: B::Device,
// }
//
// impl<B: AutodiffBackend> ChessTrainer<B> {
// pub fn new(model: ChessModel<B>, device: B::Device) -> Self {
// let optimizer = AdamConfig::new().init();
//
// Self {
// model,
// optimizer,
// batcher: ChessBatcher::default(),
// device,
// }
// }
//
// pub fn train_step(&mut self, batch_data: Vec<MctsResults>) -> f32 {
// // 1. Convert to tensors
// let batch = self.batcher.batch(batch_data, &self.device);
//
// // 2. Forward + backward (your TrainStep impl)
// let output = self.model.step(batch);
//
// // 3. Update weights
// self.optimizer.step(, &mut self.model, output.grads);
//
// // 4. Return loss (for logging)
// let loss_tensor = output.item.loss.clone().into_data();
// let loss = loss_tensor.to_vec::<f32>().unwrap()[0];
//
// loss
// }
// }
use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
use crate::net::encoding::decode_move;
use crate::net::model::{
ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
};
use burn::data::dataloader::batcher::Batcher;
use burn::module::{AutodiffModule, Module};
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
// use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::prelude::Device;
use burn::store::ModuleRecord;
use chess::{ChessMove, Color};
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::{RngExt, SeedableRng};
use std::collections::VecDeque;
use std::fs::File;
use std::io::BufReader;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;
use std::time::{SystemTime, UNIX_EPOCH};
// fn create_artifact_dir(artifact_dir: &str) {
// // Remove existing artifacts before to get an accurate learner summary
// std::fs::remove_dir_all(artifact_dir).ok();
// std::fs::create_dir_all(artifact_dir).ok();
// }
//
// pub fn train<B: AutodiffBackend>(
// artifact_dir: &str,
// config: ChessTrainingConfig,
// device: B::Device,
// ) {
// create_artifact_dir(artifact_dir);
// config
// .save(format!("{artifact_dir}/config.json"))
// .expect("Config should be saved successfully");
//
// B::seed(&device, config.seed);
//
// let batcher = ChessBatcher::default();
pub struct TrainingConfig {
pub max_time_s: Option<u64>,
pub num_iters: Option<usize>,
pub max_depth: u16, // unused
pub model_name: String,
pub load_model: bool,
pub hidden_channels: usize,
pub num_blocks: usize,
pub batch_size: usize, // num positions sampled to a batch (2048)
pub num_episodes: usize, // num games generated per iteration (5000)
pub buffer_max_size: usize, // max number of samples in the buffer
pub mcts_config: MctsConfig,
pub optimizer: AdamConfig,
pub lr: f64,
pub seed: Option<u64>,
}
// let dataloader_train = DataLoaderBuilder::new(batcher.clone())
// .batch_size(config.batch_size)
// .shuffle(config.seed)
// .num_workers(config.num_workers)
// .build(MnistDataset::train());
//
// let dataloader_test = DataLoaderBuilder::new(batcher)
// .batch_size(config.batch_size)
// .shuffle(config.seed)
// .num_workers(config.num_workers)
// .build(MnistDataset::test());
//
// let training = SupervisedTraining::new(artifact_dir, dataloader_train, dataloader_test)
// .metrics((AccuracyMetric::new(), LossMetric::new()))
// .with_file_checkpointer(CompactRecorder::new())
// .num_epochs(config.num_epochs)
// .summary();
//
// let model = config.model.init::<B>(&device);
// let result = training.launch(Learner::new(
// model,
// config.optimizer.init(),
// config.learning_rate,
// ));
//
// result
// .model
// .save_file(format!("{artifact_dir}/model"), &CompactRecorder::new())
// .expect("Trained model should be saved successfully");
// }
pub fn train(training_config: TrainingConfig, device: &Device) {
let model_path = format!("artifacts/{}", training_config.model_name.as_str());
let metadata_path = format!("{}.json", model_path);
println!("Creating model...");
let mut model: ChessModel = ChessModelConfig::init(
training_config.num_blocks,
training_config.hidden_channels,
device,
);
if training_config.load_model {
println!("Loading model {}...", model_path);
// Load model in full precision from MessagePack file
let record = ModuleRecord::load(&model_path).expect("Failed to load model.");
model = model.load_record(record)
}
let train = Arc::new(AtomicBool::new(true));
let train_signal = Arc::clone(&train);
let mut iter: usize = 0;
if training_config.load_model {
let file = File::open(&metadata_path).unwrap();
// 2. Use a buffered reader for efficiency
let reader = BufReader::new(file);
// 3. Deserialize JSON directly from the file reader
let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
iter = metadata.iterations;
println!("Loaded model had {} iters", iter);
}
let start_time = Instant::now();
ctrlc::set_handler(move || {
train_signal.store(false, Ordering::Relaxed);
println!("Finishing batch before exiting...");
})
.expect("Error setting Ctrl-C handler");
let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
let mut mcts: Mcts = Mcts {
config: training_config.mcts_config,
};
// Create RNG once and reuse it for sampling and shuffling
// Seed from system time (platform default entropy may be unavailable in some contexts)
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
println!("using seed: {}", seed);
let mut rng = SmallRng::seed_from_u64(seed);
// Initialize optimizer once so state (moments) persist across steps
let mut optim = training_config.optimizer.init();
println!("Starting training...");
while train.load(Ordering::Relaxed) {
println!("Iteration: {}", iter);
let infer_model = model.valid();
// Gen samples
println!("Generating {} games...", training_config.num_episodes);
for episode in 0..training_config.num_episodes {
println!("Episode: {}", episode);
let mut board_state = BoardState::default();
let mut episode_buffer: Vec<MctsResults> = vec![];
let mut game_hist = String::new();
let mut ply_count: usize = 0;
while board_state.status == BoardStateStatus::Ongoing {
// let before = Instant::now();
let results = mcts.search(&board_state, &infer_model, &device);
// println!("{}", before.elapsed().as_millis());
episode_buffer.push(results);
let temp = if ply_count < 30 {
1.0
} else {
0.0
};
let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
let stm = board_state.board.side_to_move();
let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
// println!("playing move: {}", mv);
if episode == 0 {
game_hist.push_str(mv.to_string().as_str());
game_hist.push(' ');
}
board_state.apply_move(mv);
ply_count += 1;
if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing {
board_state.status = BoardStateStatus::FiftyMove;
}
}
if episode == 0 {
println!("Game history of first game of iteration: {}", game_hist);
}
for result in episode_buffer.iter() {
let stm = result.board_state.board.side_to_move();
let outcome = match board_state.status {
BoardStateStatus::WhiteWinner => {
if stm == Color::White {
1.0
} else {
-1.0
}
}
BoardStateStatus::BlackWinner => {
if stm == Color::Black {
1.0
} else {
-1.0
}
}
_ => 0.0,
};
replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
result.clone(),
outcome,
));
if replay_buffer.len() > training_config.buffer_max_size {
replay_buffer.pop_front();
}
}
}
// train
if replay_buffer.is_empty() {
println!("Replay buffer is empty, skipping training step.");
continue;
}
let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
println!(
"Finished self-play! Training on {} samples (buffer size: {})...",
num_samples,
replay_buffer.len()
);
let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
indices.shuffle(&mut rng);
let samples = indices
.into_iter()
.take(num_samples)
.map(|i| replay_buffer[i].clone())
.collect::<Vec<TrainingSample>>();
let batcher = ChessBatcher {};
let batch = batcher.batch(samples, &device);
let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
let p_loss_val = output.policy_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
let v_loss_val = output.value_loss.clone().into_data().to_vec::<f32>().unwrap()[0];
println!(
"Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
iter, loss_val, p_loss_val, v_loss_val
);
let grads = output.loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
model = optim.step(training_config.lr, model, grads);
iter += 1;
if iter % 100 == 0 {
println!("Completed {} iterations", iter);
}
if training_config.max_time_s.is_some()
&& start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
{
println!("Training stopping due to time limit...");
break;
}
if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
println!(
"Training stopping due to iteration limit ({} iters completed)...",
iter
);
break;
}
}
println!("Saving model...");
let metadata = ModelMetadata {
name: training_config.model_name,
iterations: iter,
};
std::fs::write(
metadata_path,
serde_json::to_string_pretty(&metadata)
.expect("Should be able to convert metadata to JSON string"),
)
.expect("Should be able to write metadata");
model
.into_record()
.save(&model_path)
.expect("Should be able to save the model");
println!("Model saved in {:?}, exiting training.", model_path);
return;
}
fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
if visits.is_empty() {
return Vec::new();
}
// Special case: deterministic selection
if temperature == 0.0 {
let (&best_idx, _) = visits
.iter()
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
.map(|(i, p)| (i, p))
.unwrap();
return vec![(best_idx, 1.0)];
}
let inv_temp = 1.0 / temperature;
// Step 1: apply exponent
let mut adjusted: Vec<(usize, f32)> =
visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
// Step 2: normalize
let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
if sum <= 0.0 {
return adjusted; // fallback (shouldn't happen in normal MCTS)
}
for (_, v) in adjusted.iter_mut() {
*v /= sum;
}
adjusted
}
fn sample_move(
dist: &[(usize, f32)],
rng: &mut SmallRng,
side_to_move: Color,
) -> Option<ChessMove> {
let mut r: f32 = rng.random_range(0.0..1.0);
for (idx, p) in dist {
r -= *p;
if r <= 0.0 {
return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
}
}
// fallback due to floating point drift
dist.get(0)
.map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
}
engine/src/training.rs
-1
View File
@@ -1 +0,0 @@
pub mod train;
engine/src/training/train.rs
-292
View File
@@ -1,292 +0,0 @@
use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults};
use crate::net::encoding::decode_move;
use crate::net::model::{
ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
};
use burn::data::dataloader::batcher::Batcher;
use burn::module::{AutodiffModule, Module};
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::AutodiffBackend;
use chess::{ChessMove, Color};
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::{RngExt, SeedableRng};
use std::collections::VecDeque;
use std::fs::File;
use std::io::BufReader;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Instant;
use std::time::{SystemTime, UNIX_EPOCH};
pub struct TrainingConfig {
pub max_time_s: Option<u64>,
pub num_iters: Option<usize>,
pub max_depth: u16, // unused
pub model_name: String,
pub load_model: bool,
pub hidden_channels: usize,
pub num_blocks: usize,
pub batch_size: usize, // num positions sampled to a batch (2048)
pub num_episodes: usize, // num games generated per iteration (5000)
pub buffer_max_size: usize, // max number of samples in the buffer
pub mcts_config: MctsConfig,
pub optimizer: AdamConfig,
pub lr: f64,
pub seed: Option<u64>,
}
pub fn train<B: AutodiffBackend>(training_config: TrainingConfig, device: B::Device) {
let model_path = format!("artifacts/{}", training_config.model_name.as_str());
let metadata_path = format!("{}.json", model_path);
println!("Creating model...");
let mut model: ChessModel<B> = ChessModelConfig::init(
training_config.num_blocks,
training_config.hidden_channels,
&device,
);
if training_config.load_model {
println!("Loading model {}...", model_path);
// Load model in full precision from MessagePack file
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
model = model
.load_file(&model_path, &recorder, &device)
.expect("Should be able to load the model weights from the provided file");
}
let train = Arc::new(AtomicBool::new(true));
let train_signal = Arc::clone(&train);
let mut iter: usize = 0;
if training_config.load_model {
let file = File::open(&metadata_path).unwrap();
// 2. Use a buffered reader for efficiency
let reader = BufReader::new(file);
// 3. Deserialize JSON directly from the file reader
let metadata: ModelMetadata = serde_json::from_reader(reader).unwrap();
iter = metadata.iterations;
println!("Loaded model had {} iters", iter);
}
let start_time = Instant::now();
ctrlc::set_handler(move || {
train_signal.store(false, Ordering::Relaxed);
println!("Finishing batch before exiting...");
})
.expect("Error setting Ctrl-C handler");
let mut replay_buffer: VecDeque<TrainingSample> = VecDeque::new();
let mut mcts: Mcts<B::InnerBackend> = Mcts {
config: training_config.mcts_config,
_marker: PhantomData,
};
// Create RNG once and reuse it for sampling and shuffling
// Seed from system time (platform default entropy may be unavailable in some contexts)
let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
println!("using seed: {}", seed);
let mut rng = SmallRng::seed_from_u64(seed);
// Initialize optimizer once so state (moments) persist across steps
let mut optim = training_config.optimizer.init();
println!("Starting training...");
while train.load(Ordering::Relaxed) {
println!("Iteration: {}", iter);
let infer_model = model.valid();
// Gen samples
println!("Generating {} games...", training_config.num_episodes);
for episode in 0..training_config.num_episodes {
println!("Episode: {}", episode);
let mut board_state = BoardState::default();
let mut episode_buffer: Vec<MctsResults> = vec![];
let mut game_hist = String::new();
while board_state.status == BoardStateStatus::Ongoing {
// let before = Instant::now();
let results = mcts.search(&board_state, &infer_model, &device);
// println!("{}", before.elapsed().as_millis());
episode_buffer.push(results);
let temp = if board_state.halfmove_clock < 30 {
1.0
} else {
0.0
};
let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
let stm = board_state.board.side_to_move();
let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
// println!("playing move: {}", mv);
if episode == 0 {
game_hist.push_str(mv.to_string().as_str());
game_hist.push(' ');
}
board_state.apply_move(mv)
}
if episode == 0 {
println!("Game history of first game of iteration: {}", game_hist);
}
for result in episode_buffer.iter().enumerate() {
if board_state.status == BoardStateStatus::Stalemate
|| board_state.status == BoardStateStatus::Threefold
|| board_state.status == BoardStateStatus::FiftyMove
{
replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
result.1.clone(),
0.0,
));
} else if board_state.status == BoardStateStatus::WhiteWinner {
replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
result.1.clone(),
((result.0 % 2) as f32 * -2.0) + 1.0,
));
} else if board_state.status == BoardStateStatus::BlackWinner {
replay_buffer.push_back(TrainingSample::from_mcts_with_outcome(
result.1.clone(),
((result.0 % 2) as f32 * 2.0) - 1.0,
));
}
if replay_buffer.len() > training_config.buffer_max_size {
replay_buffer.pop_front();
}
}
}
// train
println!(
"Finished! Training on {} samples...",
training_config.batch_size
);
let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
indices.shuffle(&mut rng);
let samples = indices
.into_iter()
.take(training_config.batch_size)
.map(|i| replay_buffer[i].clone())
.collect::<Vec<TrainingSample>>();
let batcher = ChessBatcher {};
let batch = batcher.batch(samples, &device);
let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
let grads = output.loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
model = optim.step(training_config.lr, model, grads);
iter += 1;
if iter % 100 == 0 {
println!("Completed {} iterations", iter);
}
if training_config.max_time_s.is_some()
&& start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
{
println!("Training stopping due to time limit...");
break;
}
if training_config.num_iters.is_some() && iter >= training_config.num_iters.unwrap() {
println!(
"Training stopping due to iteration limit ({} iters completed)...",
iter
);
break;
}
}
println!("Saving model...");
let metadata = ModelMetadata {
name: training_config.model_name,
iterations: iter,
};
std::fs::write(
metadata_path,
serde_json::to_string_pretty(&metadata)
.expect("Should be able to convert metadata to JSON string"),
)
.expect("Should be able to write metadata");
// Save model in MessagePack format with full precision
let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
model
.save_file(&model_path, &recorder)
.expect("Should be able to save the model");
println!("Model saved in {:?}, exiting training.", model_path);
return;
}
fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> {
if visits.is_empty() {
return Vec::new();
}
// Special case: deterministic selection
if temperature == 0.0 {
let (&best_idx, _) = visits
.iter()
.max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
.map(|(i, p)| (i, p))
.unwrap();
return vec![(best_idx, 1.0)];
}
let inv_temp = 1.0 / temperature;
// Step 1: apply exponent
let mut adjusted: Vec<(usize, f32)> =
visits.iter().map(|(i, v)| (*i, v.powf(inv_temp))).collect();
// Step 2: normalize
let sum: f32 = adjusted.iter().map(|(_, v)| *v).sum();
if sum <= 0.0 {
return adjusted; // fallback (shouldn't happen in normal MCTS)
}
for (_, v) in adjusted.iter_mut() {
*v /= sum;
}
adjusted
}
fn sample_move(
dist: &[(usize, f32)],
rng: &mut SmallRng,
side_to_move: Color,
) -> Option<ChessMove> {
let mut r: f32 = rng.random_range(0.0..1.0);
for (idx, p) in dist {
r -= *p;
if r <= 0.0 {
return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
}
}
// fallback due to floating point drift
dist.get(0)
.map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
}