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

misc fixes

Browse Source
This commit is contained in:
drake
2026-08-24 23:19:51 -04:00
parent c40a9667bd
commit 33dc15a597
11 changed files with 1742 additions and 1262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/mcts.rs
+164 -110
View File
@@ -2,7 +2,7 @@
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::prelude::Device;
use burn::Tensor;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
@@ -10,7 +10,6 @@ use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;
use std::marker::PhantomData;
pub struct Node {
pub prior: f32,
@@ -107,17 +106,62 @@ impl Default for MctsConfig {
}
}
pub struct Mcts<B: Backend> {
pub struct Mcts {
pub config: MctsConfig,
pub _marker: PhantomData<B>, // if B not otherwise stored
}
impl<B: Backend> Mcts<B> {
fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
match status {
BoardStateStatus::Ongoing => None,
BoardStateStatus::Stalemate
| BoardStateStatus::Threefold
| BoardStateStatus::FiftyMove => 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 }),
}
}
fn softmax(logits: &[f32]) -> Vec<f32> {
if logits.is_empty() {
return Vec::new();
}
let mut max_logit = f32::NEG_INFINITY;
for &v in logits {
if v > max_logit {
max_logit = v;
}
}
if !max_logit.is_finite() {
max_logit = 0.0;
}
let mut exps_sum = 0.0f32;
let mut probs = vec![0.0f32; logits.len()];
for (i, &v) in logits.iter().enumerate() {
let x = (v - max_logit).exp();
if x.is_finite() {
probs[i] = x;
exps_sum += x;
}
}
if exps_sum > 0.0 && exps_sum.is_finite() {
for p in &mut probs {
*p /= exps_sum;
}
} else {
let uniform = 1.0 / logits.len() as f32;
for p in &mut probs {
*p = uniform;
}
}
probs
}
impl Mcts {
pub fn search(
&mut self,
board_state: &BoardState,
model: &ChessModel<B>,
device: &B::Device,
model: &ChessModel,
device: &Device,
) -> MctsResults {
let mut nodes = Vec::<Node>::new();
@@ -130,133 +174,138 @@ impl<B: Backend> Mcts<B> {
// Apply Dirichlet noise to root children
self.add_dirichlet_noise(root, &mut nodes);
// We'll batch leaf evaluations to reduce per-leaf model calls and device-host syncs.
let mut sims_done: usize = 0;
let num_sims = self.config.num_simulations;
while sims_done < num_sims {
// Collect a batch of leaf nodes (and their selection paths)
let mut leaf_nodes: Vec<usize> = Vec::new();
let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
let mut leaf_states: Vec<Tensor<B, 4>> = Vec::new();
let mut leaf_states: Vec<Tensor<4>> = Vec::new();
while leaf_nodes.len() < std::cmp::min(self.config.batch_max, num_sims - sims_done) {
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);
}
// Record leaf node and its path
leaf_nodes.push(current);
leaf_paths.push(path.clone());
// 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 except leaf itself
for &idx in &path[..path.len() - 1] {
nodes[idx].visit_count -= 1;
nodes[idx].value_sum += 1.0;
}
self.backpropagate(&mut nodes, &path, term_val, stm);
sims_done += 1;
continue;
}
// Prepare state tensor for this leaf
let state: Tensor<B, 4> =
// 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;
}
break;
}
leaf_nodes.push(current);
leaf_paths.push(path);
let state: Tensor<4> =
encode_board_state_perspective(&nodes[current].board_state, device)
.reshape([1, 18, 8, 8]);
leaf_states.push(state);
sims_done += 1;
}
if leaf_nodes.is_empty() {
break;
if sims_done >= num_sims {
break;
}
continue;
}
// Batch evaluate the collected leaf states
// Batch evaluate non-terminal leaf states
let batch = Tensor::cat(leaf_states, 0);
let (policy_batch, value_batch) = model.forward(batch);
// Move tensors to host once per 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();
// Process each evaluated leaf: expand and backpropagate
for (i, &node_idx) in leaf_nodes.iter().enumerate() {
let path = &leaf_paths[i];
// slice for this sample's logits
// 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;
}
let start = i * num_moves;
let end = start + num_moves;
let logits = &policy_data[start..end];
let probs = softmax(logits);
// Convert logits to probabilities with a numerically-stable softmax on host
let mut max_logit = f32::NEG_INFINITY;
for &v in logits.iter() {
if v > max_logit {
max_logit = v;
}
}
if logits.is_empty() {
println!("logits empty")
}
if !max_logit.is_finite() {
println!("max logits not finite")
}
let mut exps_sum = 0.0f32;
let mut probs = vec![0.0f32; num_moves];
for (j, &v) in logits.iter().enumerate() {
let x = (v - max_logit).exp();
if !x.is_finite() {
probs[j] = 0.0;
} else {
probs[j] = x;
exps_sum += x;
}
}
if exps_sum > 0.0 && exps_sum.is_finite() {
for p in &mut probs {
*p /= exps_sum;
}
} else {
// fallback: uniform or zero
let uniform = 1.0 / num_moves as f32;
for p in &mut probs {
*p = uniform;
}
}
// Expand: add legal moves as children with prior from probs
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();
for mv in legal_moves {
let stm = nodes[node_idx].board_state.board.side_to_move();
let idx = encode_move(mv, stm).expect("Invalid move");
let prior = probs[idx];
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(prior, new_board, Some(mv)));
nodes.push(Node::new(normalized_prior, new_board, Some(mv)));
nodes[node_idx].children.push(child_idx);
}
// Backpropagate the value for this leaf
let value = value_data[i];
let color = nodes[node_idx].board_state.board.side_to_move();
self.backpropagate(&mut nodes, path, value, color);
self.backpropagate(&mut nodes, path, value, stm);
sims_done += 1;
}
}
// Build compact move distribution: encoded move index -> probability (visits / num_simulations)
// 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 denom = self.config.num_simulations as f32;
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");
let enc = encode_move(mv, stm).expect("Invalid move");
let prob = nodes[*child_idx].visit_count as f32 / denom;
move_dist.push((enc, prob));
if let Some(enc) = encode_move(mv, stm) {
let prob = nodes[*child_idx].visit_count as f32 / denom;
move_dist.push((enc, prob));
}
}
MctsResults::new(board_state.clone(), move_dist, nodes[root].value())
@@ -266,54 +315,53 @@ impl<B: Backend> Mcts<B> {
&mut self,
node_idx: usize,
arena: &mut Vec<Node>,
model: &ChessModel<B>,
device: &B::Device,
model: &ChessModel,
device: &Device,
) -> f32 {
if arena[node_idx].board_state.status == BoardStateStatus::Stalemate
|| arena[node_idx].board_state.status == BoardStateStatus::Threefold
|| arena[node_idx].board_state.status == BoardStateStatus::FiftyMove
{
0.0
} else if arena[node_idx].board_state.status == BoardStateStatus::WhiteWinner {
if arena[node_idx].board_state.board.side_to_move() == Black {
1.0
} else {
-1.0
}
} else if arena[node_idx].board_state.status == BoardStateStatus::BlackWinner {
if arena[node_idx].board_state.board.side_to_move() == White {
1.0
} else {
-1.0
}
let status = arena[node_idx].board_state.status;
let stm = arena[node_idx].board_state.board.side_to_move();
if let Some(term_val) = terminal_value(status, stm) {
term_val
} else {
let state: Tensor<B, 4> =
let state: Tensor<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();
let logits = policy_head.into_data().to_vec::<f32>().unwrap();
let probs = softmax(&logits);
for mv in legal_moves {
let stm = arena[node_idx].board_state.board.side_to_move();
let idx = encode_move(mv, stm).expect("Invalid move");
let prior = policy[idx];
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 = 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.push(Node::new(normalized_prior, new_board, Some(mv)));
arena[node_idx].children.push(child_idx);
}
value_head.into_data().to_vec().unwrap()[0]
value_head.into_data().to_vec::<f32>().unwrap()[0]
}
}
@@ -356,7 +404,6 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
let gamma = Gamma::new(alpha as f64, 1.0).unwrap();
// Use a single SmallRng seeded from system time (avoid depending on thread_rng helper)
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap();
@@ -367,8 +414,15 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
let sum: f32 = samples.iter().sum();
for x in &mut samples {
*x /= sum;
if sum > 0.0 && sum.is_finite() {
for x in &mut samples {
*x /= sum;
}
} else if size > 0 {
let uniform = 1.0 / size as f32;
for x in &mut samples {
*x = uniform;
}
}
samples