first working

engine/src/mcts.rs

@@ -0,0 +1,416 @@
+// use crate::networks::PolicyNetwork;
+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::Tensor;
+use chess::BoardStatus::{Checkmate, Stalemate};
+use chess::Color::White;
+use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
+use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
+use std::collections::HashMap;
+use std::marker::PhantomData;
+use std::time::Instant;
+
+pub struct Node {
+    pub prior: f32,
+    pub children: Vec<usize>,
+    pub visit_count: u32,
+    pub value_sum: f32,
+    pub board_state: BoardState,
+    pub last_move: Option<ChessMove>, // move that produced this node from its parent
+}
+
+impl Node {
+    pub fn new(prior: f32, board_state: BoardState, last_move: Option<ChessMove>) -> Node {
+        Node {
+            prior,
+            children: vec![],
+            visit_count: 0,
+            value_sum: 0.0,
+            board_state,
+            last_move,
+        }
+    }
+
+    pub fn value(&self) -> f32 {
+        self.value_sum / self.visit_count as f32
+    }
+
+    pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
+        self.children
+            .iter()
+            .copied()
+            .max_by(|&a, &b| {
+                ucb_score(self, &arena[a], c_puct)
+                    .partial_cmp(&ucb_score(self, &arena[b], c_puct))
+                    .unwrap()
+            })
+            .expect("select_child on leaf")
+    }
+}
+
+impl Clone for Node {
+    fn clone(&self) -> Self {
+        Self::new(self.prior, self.board_state.clone(), self.last_move.clone())
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct MctsResults {
+    pub board_state: BoardState,
+    pub move_dist: HashMap<ChessMove, f32>,
+    pub value: f32,
+}
+
+impl MctsResults {
+    pub fn new(
+        board_state: BoardState,
+        move_dist: HashMap<ChessMove, f32>,
+        value: f32,
+    ) -> MctsResults {
+        MctsResults {
+            board_state,
+            move_dist,
+            value,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct MctsConfig {
+    pub num_simulations: usize,
+    pub c_puct: f32,
+    pub dirichlet_alpha: f32,
+    pub dirichlet_epsilon: f32,
+}
+
+impl MctsConfig {
+    pub fn new(
+        num_simulations: usize,
+        c_puct: f32,
+        dirichlet_alpha: f32,
+        dirichlet_epsilon: f32,
+    ) -> MctsConfig {
+        MctsConfig {
+            num_simulations,
+            c_puct,
+            dirichlet_alpha,
+            dirichlet_epsilon,
+        }
+    }
+}
+
+impl Default for MctsConfig {
+    fn default() -> MctsConfig {
+        MctsConfig::new(400, 1.0, 0.05, 0.25)
+    }
+}
+
+pub struct Mcts<B: Backend> {
+    pub config: MctsConfig,
+    pub _marker: PhantomData<B>, // if B not otherwise stored
+}
+
+impl<B: Backend> Mcts<B> {
+    pub fn search(
+        &mut self,
+        board_state: &BoardState,
+        model: &ChessModel<B>,
+        device: &B::Device,
+    ) -> MctsResults {
+        let mut nodes = Vec::<Node>::new();
+
+        let root = 0;
+        nodes.push(Node::new(0.0, board_state.clone(), None));
+
+        self.expand(root, &mut nodes, model, device);
+
+        // 👇 APPLY DIRICHLET NOISE HERE
+        self.add_dirichlet_noise(root, &mut nodes);
+
+        for i in 0..self.config.num_simulations {
+            // if i % 10 == 0 {
+            //     println!("mcts sim #{}", i);
+            // }
+            let mut path = vec![root];
+            let mut current = root;
+
+            while !nodes[current].children.is_empty() {
+                current = nodes[current].select_child(&nodes, &self.config.c_puct);
+                path.push(current);
+            }
+
+            let value: f32 = self.expand(current, &mut nodes, model, device);
+
+            let color = nodes[current].board_state.board.side_to_move();
+            self.backpropagate(&mut nodes, &path, value, color);
+        }
+
+        let mut move_dist: HashMap<ChessMove, f32> = HashMap::new(); // TODO: make vec<(Chessmove, f32)>
+        for idx in nodes[root].children.iter() {
+            move_dist.insert(
+                nodes[*idx].last_move.expect("move didnt exist"),
+                nodes[*idx].visit_count as f32 / self.config.num_simulations as f32,
+            );
+        }
+
+        MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
+    }
+
+    fn expand(
+        &mut self,
+        node_idx: usize,
+        arena: &mut Vec<Node>,
+        model: &ChessModel<B>,
+        device: &B::Device,
+    ) -> f32 {
+        let state: Tensor<B, 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();
+
+        for mv in legal_moves {
+            let stm = arena[node_idx].board_state.board.side_to_move();
+            let idx = encode_move(mv, stm);
+            let prior = policy[idx];
+
+            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[node_idx].children.push(child_idx);
+        }
+
+        value_head.into_data().to_vec().unwrap()[0]
+    }
+
+    fn backpropagate(&mut self, nodes: &mut [Node], path: &[usize], value: f32, color: Color) {
+        for &idx in path {
+            let node = &mut nodes[idx];
+
+            if node.board_state.board.side_to_move() == color {
+                node.value_sum += value;
+            } else {
+                node.value_sum -= value;
+            }
+
+            node.visit_count += 1;
+        }
+    }
+
+    fn add_dirichlet_noise(&mut self, node_id: usize, nodes: &mut Vec<Node>) {
+        let node_children = &nodes[node_id].children.clone();
+
+        let n = node_children.len();
+        if n == 0 {
+            return;
+        }
+
+        let alpha = self.config.dirichlet_alpha;
+        let epsilon = self.config.dirichlet_epsilon;
+
+        let noise = dirichlet_sample(n, alpha);
+
+        for (node_idx, noise_val) in node_children.iter().zip(noise) {
+            let prev_prior = nodes[*node_idx].prior;
+            nodes[*node_idx].prior = (1.0 - epsilon) * prev_prior + epsilon * noise_val;
+        }
+    }
+}
+
+fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
+    use rand_distr::{Distribution, Gamma};
+
+    let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
+
+    let mut samples: Vec<f32> = (0..size)
+        .map(|_| gamma.sample(&mut rand::rng()) as f32)
+        .collect();
+
+    let sum: f32 = samples.iter().sum();
+
+    for x in &mut samples {
+        *x /= sum;
+    }
+
+    samples
+}
+
+fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
+    let prior_score = c_puct * child.prior * (parent.visit_count as f32).sqrt()
+        / (1.0 + child.visit_count as f32);
+
+    let value_score: f32;
+    if child.visit_count > 0 {
+        value_score = -child.value(); // value from opposing side
+    } else {
+        value_score = 0.0;
+    }
+
+    value_score + prior_score
+}
+
+pub fn find_best_move(nodes: &[Node], root: usize) -> Option<(BoardState, String)> {
+    // Choose best child of root by visit count (standard MCTS behavior)
+    let root_node = &nodes[root];
+    if root_node.children.is_empty() {
+        return None;
+    }
+
+    let mut best_visits = 0u32;
+    let mut best_child_idx: Option<usize> = None;
+    for &child_idx in &root_node.children {
+        let child = &nodes[child_idx];
+        if child.visit_count > best_visits {
+            best_visits = child.visit_count;
+            best_child_idx = Some(child_idx);
+        }
+    }
+
+    let child_idx = match best_child_idx {
+        Some(i) => i,
+        None => return None,
+    };
+
+    let bn = &nodes[child_idx];
+
+    println!("best_child visits: {}", bn.visit_count);
+    let uci_move = bn.last_move?.to_string();
+    Some((bn.board_state.clone(), uci_move))
+}
+
+pub fn print_tree_arena(nodes: &[Node], root: usize, depth: usize) {
+    let node = &nodes[root];
+
+    println!(
+        "{}- Value: {:.2}, Visits: {}, Prior: {:.2}, Side: {:?}",
+        "  ".repeat(depth),
+        node.value_sum,
+        node.visit_count,
+        node.prior,
+        node.board_state.board.side_to_move()
+    );
+
+    for &child_idx in &node.children {
+        print_tree_arena(nodes, child_idx, depth + 1);
+    }
+}
+
+fn material_of(piece: Piece) -> f32 {
+    match piece {
+        Queen => 9.0,
+        Rook => 5.0,
+        Bishop => 3.0,
+        Knight => 3.0,
+        Pawn => 1.0,
+        _ => 0.0,
+    }
+}
+
+const MATERIAL_VALUE_DIVISOR: f32 = 40.0;
+
+pub fn heuristic_eval(board: &Board, perspective: Color) -> f32 {
+    let mut value = 0.0;
+
+    // material
+    for color in ALL_COLORS {
+        for piece in ALL_PIECES {
+            let bitboard = board.color_combined(color).0 & board.pieces(piece).0;
+            let total_val =
+                (bitboard.count_ones() as f32 * material_of(piece)) / MATERIAL_VALUE_DIVISOR;
+
+            if color == perspective {
+                value += total_val;
+            } else {
+                value -= total_val;
+            }
+        }
+    }
+
+    value
+
+    // board.checkers()
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum BoardStateStatus {
+    Ongoing,
+    Stalemate,
+    WhiteWinner,
+    BlackWinner,
+    Threefold,
+    FiftyMove,
+}
+
+#[derive(Clone, Debug)]
+pub struct BoardState {
+    pub board: Board,
+    pub halfmove_clock: u8,
+    pub repetition_table: HashMap<u64, u8>,
+    pub status: BoardStateStatus,
+}
+
+impl BoardState {
+    pub fn new(board: Board, halfmove_clock: u8, repetition_table: HashMap<u64, u8>) -> BoardState {
+        BoardState {
+            board,
+            halfmove_clock,
+            repetition_table: repetition_table,
+            status: Ongoing,
+        }
+    }
+
+    pub fn apply_move(&mut self, mv: ChessMove) {
+        self.halfmove_clock += 1;
+        if self.board.piece_on(mv.get_source()) == Some(Pawn)
+            || self.board.piece_on(mv.get_dest()).is_some()
+        {
+            self.halfmove_clock = 0;
+        }
+        self.board = self.board.make_move_new(mv);
+
+        let board_hash = self.board.get_hash();
+
+        let current_rep = self.repetition_table.get(&board_hash);
+        let mut new_rep = 1;
+        if current_rep.is_some() {
+            new_rep = current_rep.unwrap() + 1;
+        }
+        self.repetition_table.insert(board_hash, new_rep);
+
+        if self.board.status() == Checkmate {
+            if self.board.side_to_move() == White {
+                // white's move after black played the mating move
+                self.status = BoardStateStatus::BlackWinner;
+            } else {
+                self.status = BoardStateStatus::WhiteWinner;
+            }
+        } else if self.board.status() == Stalemate {
+            self.status = BoardStateStatus::Stalemate;
+        }
+
+        if self.halfmove_clock >= 100 {
+            self.status = FiftyMove;
+        }
+
+        if new_rep >= 3 {
+            self.status = Threefold;
+        }
+    }
+}
+
+impl Default for BoardState {
+    fn default() -> BoardState {
+        BoardState::new(Board::default(), 0, HashMap::new())
+    }
+}