batched mcts

engine/src/net/model.rs

@@ -1,5 +1,5 @@
 use crate::mcts::{BoardState, MctsResults};
-use crate::net::encoding::{encode_board_state_perspective, encode_move};
+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};
@@ -13,8 +13,6 @@ use burn::{
     prelude::*,
 };
 use burn_ndarray::NdArray;
-use chess::ChessMove;
-use std::collections::HashMap;
 /*
 Input planes:
 1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -255,14 +253,15 @@ impl<B: Backend> InferenceStep for ChessModel<B> {
 #[derive(Clone)]
 pub struct TrainingSample {
     pub board_state: BoardState,
-    pub policy_target: HashMap<ChessMove, f32>,
+    // Compact representation: list of (encoded_move_index, probability)
+    pub policy_target: Vec<(usize, f32)>,
     pub value_target: f32,
 }
 
 impl TrainingSample {
     pub fn new(
         board_state: BoardState,
-        policy_target: HashMap<ChessMove, f32>,
+        policy_target: Vec<(usize, f32)>,
         value_target: f32,
     ) -> Self {
         TrainingSample {
@@ -273,6 +272,7 @@ impl TrainingSample {
     }
 
     pub fn from_mcts_with_outcome(mcts_results: MctsResults, outcome: f32) -> Self {
+        // move_dist is already a compact Vec<(encoded_move_index, prob)>
         TrainingSample::new(mcts_results.board_state, mcts_results.move_dist, outcome)
     }
 }
@@ -301,9 +301,8 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
             .cloned()
             .map(|item| {
                 let mut policy = vec![0.0f32; 4672];
-                let stm = item.board_state.board.side_to_move();
-                for (mv, prob) in item.policy_target {
-                    policy[encode_move(mv, stm)] = prob;
+                for (idx, prob) in item.policy_target.iter() {
+                    policy[*idx] = *prob;
                 }
 
                 // Normalize