315 lines
10 KiB
Rust
315 lines
10 KiB
Rust
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};
|
|
|
|
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(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"))
|
|
}
|