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

unbatch mcts

Browse Source
This commit is contained in:
drake
2026-08-26 14:29:58 -04:00
parent 55dc0bf737
commit c51964a443
1 changed files with 113 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/mcts.rs
+113 -176
View File
@@ -1,6 +1,6 @@
// use crate::networks::PolicyNetwork;
use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
use crate::net::encoding::{encode_board_states_batch, encode_move};
use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
use crate::net::encoding::{encode_board_state_perspective, encode_move};
use crate::net::model::ChessModel;
use burn::prelude::Device;
use chess::BoardStatus::{Checkmate, Stalemate};
@@ -10,6 +10,9 @@ use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;
//TODO retain tree after move selection?
#[derive(Clone)]
pub struct Node {
pub prior: f32,
pub children: Vec<usize>,
@@ -36,40 +39,32 @@ impl Node {
}
pub fn select_child(&self, arena: &[Node], c_puct: &f32) -> usize {
let sqrt_parent = (self.visit_count as f32).sqrt();
let c_puct_sqrt = c_puct * sqrt_parent;
self.children
.iter()
.copied()
.max_by(|&a, &b| {
let score_a = ucb_score_fast(&arena[a], c_puct_sqrt);
let score_b = ucb_score_fast(&arena[b], c_puct_sqrt);
score_a.partial_cmp(&score_b).unwrap_or(std::cmp::Ordering::Equal)
ucb_score(self, &arena[a], c_puct)
.partial_cmp(&ucb_score(self, &arena[b], c_puct))
.unwrap()
})
.expect("select_child on leaf")
}
}
#[inline]
fn ucb_score_fast(child: &Node, c_puct_sqrt: f32) -> f32 {
let prior_score = c_puct_sqrt * child.prior / (1.0 + child.visit_count as f32);
fn ucb_score(parent: &Node, child: &Node, c_puct: &f32) -> f32 {
let prior_score = c_puct * child.prior * ((parent.visit_count + 1) as f32).sqrt()
/ (1.0 + child.visit_count as f32);
let value_score = if child.visit_count > 0 {
-child.value()
let value_score: f32;
if child.visit_count > 0 {
value_score = -child.value(); // value from opposing side
} else {
0.0
};
value_score = 0.0;
}
value_score + prior_score
}
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,
@@ -127,10 +122,8 @@ pub struct Mcts {
fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
match status {
BoardStateStatus::Ongoing => None,
BoardStateStatus::Stalemate
| BoardStateStatus::Threefold
| BoardStateStatus::FiftyMove => Some(0.0),
Ongoing => None,
BoardStateStatus::Stalemate | Threefold | FiftyMove | InsufficientMaterial => 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 }),
}
@@ -171,159 +164,56 @@ fn softmax(logits: &[f32]) -> Vec<f32> {
probs
}
const LEAF_VL_COUNT: u32 = 10;
const LEAF_VL_WEIGHT: f32 = 10.0;
impl Mcts {
pub fn search(
&mut self,
board_state: &BoardState,
model: &ChessModel,
device: &Device,
training: bool,
) -> MctsResults {
let mut nodes = Vec::<Node>::new();
let root = 0;
nodes.push(Node::new(0.0, board_state.clone(), None));
// Expand root to create initial children and priors
self.expand(root, &mut nodes, model, device);
// Apply Dirichlet noise to root children
self.add_dirichlet_noise(root, &mut nodes);
if training {
// APPLY DIRICHLET NOISE HERE
self.add_dirichlet_noise(root, &mut nodes);
}
let mut sims_done: usize = 0;
let num_sims = self.config.num_simulations;
while sims_done < num_sims {
let mut leaf_nodes: Vec<usize> = Vec::new();
let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
for _ in 0..num_sims {
let mut path = vec![root];
let mut current = root;
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);
}
// Strong virtual loss on unexpanded leaf so batch collector visits diverse leaves across tree
nodes[current].visit_count += LEAF_VL_COUNT;
nodes[current].value_sum -= LEAF_VL_WEIGHT;
// 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
for &idx in &path[..path.len() - 1] {
nodes[idx].visit_count -= 1;
nodes[idx].value_sum += 1.0;
}
nodes[current].visit_count -= LEAF_VL_COUNT;
nodes[current].value_sum += LEAF_VL_WEIGHT;
self.backpropagate(&mut nodes, &path, term_val, stm);
sims_done += 1;
continue;
}
// 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;
}
nodes[current].visit_count -= LEAF_VL_COUNT;
nodes[current].value_sum += LEAF_VL_WEIGHT;
break;
}
leaf_nodes.push(current);
leaf_paths.push(path);
while !nodes[current].children.is_empty() {
current = nodes[current].select_child(&nodes, &self.config.c_puct);
path.push(current);
}
if leaf_nodes.is_empty() {
if sims_done >= num_sims {
break;
}
continue;
}
let value: f32 = self.expand(current, &mut nodes, model, device);
// Batch evaluate non-terminal leaf states using single CPU buffer encoding
let leaf_refs: Vec<&BoardState> = leaf_nodes.iter().map(|&idx| &nodes[idx].board_state).collect();
let batch = encode_board_states_batch(&leaf_refs, device);
let (policy_batch, value_batch) = model.forward(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();
for (i, &node_idx) in leaf_nodes.iter().enumerate() {
let path = &leaf_paths[i];
// 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;
}
nodes[node_idx].visit_count -= LEAF_VL_COUNT;
nodes[node_idx].value_sum += LEAF_VL_WEIGHT;
let start = i * num_moves;
let end = start + num_moves;
let logits = &policy_data[start..end];
let probs = softmax(logits);
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();
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(normalized_prior, new_board, Some(mv)));
nodes[node_idx].children.push(child_idx);
}
let value = value_data[i];
self.backpropagate(&mut nodes, path, value, stm);
sims_done += 1;
}
let color = nodes[current].board_state.board.side_to_move();
self.backpropagate(&mut nodes, &path, value, color);
}
// 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 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 };
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");
if let Some(enc) = encode_move(mv, stm) {
@@ -348,8 +238,8 @@ impl Mcts {
if let Some(term_val) = terminal_value(status, stm) {
term_val
} else {
let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
let (policy_head, value_head) = model.forward(state);
let state = encode_board_state_perspective(&arena[node_idx].board_state, device);
let (policy_head, value_head) = model.forward(state.unsqueeze::<4>());
let legal_moves: Vec<ChessMove> =
MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
@@ -450,8 +340,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
samples
}
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];
@@ -540,6 +428,7 @@ pub enum BoardStateStatus {
BlackWinner,
Threefold,
FiftyMove,
InsufficientMaterial,
}
#[derive(Clone, Debug)]
@@ -555,52 +444,100 @@ impl BoardState {
BoardState {
board,
halfmove_clock,
repetition_table: 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()
{
let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();
let is_en_passant = is_pawn_move
&& self.board.en_passant() == Some(mv.get_dest())
&& self.board.piece_on(mv.get_dest()).is_none();
self.halfmove_clock = self.halfmove_clock.saturating_add(1);
if is_pawn_move || is_normal_capture || is_en_passant {
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;
}
let new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 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;
match self.board.status() {
Checkmate => {
self.status = if self.board.side_to_move() == White {
BoardStateStatus::BlackWinner
} else {
BoardStateStatus::WhiteWinner
};
return;
}
} else if self.board.status() == Stalemate {
self.status = BoardStateStatus::Stalemate;
Stalemate => {
self.status = BoardStateStatus::Stalemate;
return;
}
_ => {}
}
if is_insufficient_material(&self.board) {
self.status = InsufficientMaterial;
return;
}
if self.halfmove_clock >= 100 {
self.status = FiftyMove;
return;
}
if new_rep >= 3 {
self.status = Threefold;
return;
}
self.status = BoardStateStatus::Ongoing;
}
}
impl Default for BoardState {
fn default() -> BoardState {
BoardState::new(Board::default(), 0, HashMap::new())
let board = Board::default();
let repetition_table = HashMap::from([(board.get_hash(), 1)]);
BoardState::new(board, 0, repetition_table)
}
}
// Could be made more comprehensive
pub fn is_insufficient_material(board: &Board) -> bool {
let pawns = board.pieces(Pawn).0;
let rooks = board.pieces(Rook).0;
let queens = board.pieces(Queen).0;
// Any pawn/rook/queen means this is not trivially dead.
if pawns != 0 || rooks != 0 || queens != 0 {
return false;
}
let bishops = board.pieces(Bishop).0;
let knights = board.pieces(Knight).0;
let bishop_count = bishops.count_ones();
let knight_count = knights.count_ones();
match (bishop_count, knight_count) {
// K vs K
(0, 0) => true,
// K+B vs K
// K+N vs K
(1, 0) | (0, 1) => true,
// Other configurations require more analysis.
_ => false,
}
}