training fixes

engine/src/main.rs

@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
 fn main() {
     let device = Device::cuda(DeviceIndex::default());
 
-    let mcts_config = MctsConfig::new(200, 1.0, 0.05, 0.25, 128);
+    let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
 
     let adam_config = AdamConfig::new();
 
@@ -29,8 +29,10 @@ fn main() {
         hidden_channels: 48,
         num_blocks: 4,
         batch_size: 1024, // num positions sampled to a batch (2048)
+        minibatch_size: 256,
+        num_train_steps: 10,
         num_episodes: 96,         // num games generated per iteration (5000)
-        buffer_max_size: 25_000, // max number of samples in the buffer
+        buffer_max_size: 200_000, // max number of samples in the buffer
         mcts_config,
         optimizer: adam_config,
         lr: 2e-4,

engine/src/train.rs

@@ -33,6 +33,8 @@ 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,
@@ -241,49 +243,101 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
         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 mut loss_val: f32 = 0.0;
+            let mut p_loss_val: f32 = 0.0;
+            let mut v_loss_val: f32 = 0.0;
+
+            // 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 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);
 
-        let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
+                if step == training_config.num_train_steps - 1 {
-
+                    loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
-        let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
+                    p_loss_val = output
-        let p_loss_val = output
                         .policy_loss
                         .clone()
                         .into_data()
                         .to_vec::<f32>()
                         .unwrap()[0];
-        let v_loss_val = output
+                    v_loss_val = output
                         .value_loss
                         .clone()
                         .into_data()
                         .to_vec::<f32>()
                         .unwrap()[0];
-
+                }
+            }
+            if step == training_config.num_train_steps - 1 {
                 println!(
-            "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
+                    "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 mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
-        let grads = GradientsParams::from_grads(grads, &model);
+        // indices.shuffle(&mut rng);
-
+        //
-        model = optim.step(training_config.lr, model, grads);
+        // 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;
 

engine/src/worker.rs

@@ -84,7 +84,8 @@ impl Worker {
             // println!("{}", before.elapsed().as_millis());
             episode_buffer.push(results);
 
-            let temp = if ply_count < 30 { 1.0 } else { 0.2 };
+            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 };
 
             let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
 
@@ -125,9 +126,11 @@ impl Worker {
                 _ => 0.0,
             };
 
+            let blended_target = 0.75 * outcome + 0.25 * result.value;
+
             samples.push(TrainingSample::from_mcts_with_outcome(
                 result.clone(),
-                outcome,
+                blended_target,
             ));
         }