544 lines
15 KiB
Rust
544 lines
15 KiB
Rust
// use crate::networks::PolicyNetwork;
|
|
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};
|
|
use chess::Color::{Black, White};
|
|
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;
|
|
|
|
//TODO retain tree after move selection?
|
|
|
|
#[derive(Clone)]
|
|
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")
|
|
}
|
|
}
|
|
|
|
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: f32;
|
|
if child.visit_count > 0 {
|
|
value_score = -child.value(); // value from opposing side
|
|
} else {
|
|
value_score = 0.0;
|
|
}
|
|
|
|
value_score + prior_score
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct MctsResults {
|
|
pub board_state: BoardState,
|
|
// Compact encoded move distribution: (encoded_move_index, probability)
|
|
pub move_dist: Vec<(usize, f32)>,
|
|
pub value: f32,
|
|
}
|
|
|
|
impl MctsResults {
|
|
pub fn new(board_state: BoardState, move_dist: Vec<(usize, 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,
|
|
pub batch_max: usize,
|
|
}
|
|
|
|
impl MctsConfig {
|
|
pub fn new(
|
|
num_simulations: usize,
|
|
c_puct: f32,
|
|
dirichlet_alpha: f32,
|
|
dirichlet_epsilon: f32,
|
|
batch_max: usize,
|
|
) -> MctsConfig {
|
|
MctsConfig {
|
|
num_simulations,
|
|
c_puct,
|
|
dirichlet_alpha,
|
|
dirichlet_epsilon,
|
|
batch_max,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for MctsConfig {
|
|
fn default() -> MctsConfig {
|
|
MctsConfig::new(400, 1.0, 0.05, 0.25, 64)
|
|
}
|
|
}
|
|
|
|
pub struct Mcts {
|
|
pub config: MctsConfig,
|
|
}
|
|
|
|
fn terminal_value(status: BoardStateStatus, stm: Color) -> Option<f32> {
|
|
match status {
|
|
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 }),
|
|
}
|
|
}
|
|
|
|
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,
|
|
device: &Device,
|
|
training: bool,
|
|
) -> 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);
|
|
|
|
if training {
|
|
// APPLY DIRICHLET NOISE HERE
|
|
self.add_dirichlet_noise(root, &mut nodes);
|
|
}
|
|
|
|
let num_sims = self.config.num_simulations;
|
|
|
|
for _ in 0..num_sims {
|
|
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);
|
|
}
|
|
|
|
// 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
|
|
};
|
|
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) {
|
|
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())
|
|
}
|
|
|
|
fn expand(
|
|
&mut self,
|
|
node_idx: usize,
|
|
arena: &mut Vec<Node>,
|
|
model: &ChessModel,
|
|
device: &Device,
|
|
) -> f32 {
|
|
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 = 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();
|
|
|
|
let logits = policy_head.into_data().to_vec::<f32>().unwrap();
|
|
let probs = softmax(&logits);
|
|
|
|
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(normalized_prior, new_board, Some(mv)));
|
|
arena[node_idx].children.push(child_idx);
|
|
}
|
|
|
|
value_head.into_data().to_vec::<f32>().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 now = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.unwrap();
|
|
let seed = now.as_nanos() as u64;
|
|
let mut rng = rand::rngs::SmallRng::seed_from_u64(seed);
|
|
|
|
let mut samples: Vec<f32> = (0..size).map(|_| gamma.sample(&mut rng) as f32).collect();
|
|
|
|
let sum: f32 = samples.iter().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
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum BoardStateStatus {
|
|
Ongoing,
|
|
Stalemate,
|
|
WhiteWinner,
|
|
BlackWinner,
|
|
Threefold,
|
|
FiftyMove,
|
|
InsufficientMaterial,
|
|
}
|
|
|
|
#[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,
|
|
status: Ongoing,
|
|
}
|
|
}
|
|
|
|
pub fn apply_move(&mut self, mv: ChessMove) {
|
|
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 new_rep = self.repetition_table.get(&board_hash).copied().unwrap_or(0) + 1;
|
|
self.repetition_table.insert(board_hash, new_rep);
|
|
|
|
match self.board.status() {
|
|
Checkmate => {
|
|
self.status = if self.board.side_to_move() == White {
|
|
BoardStateStatus::BlackWinner
|
|
} else {
|
|
BoardStateStatus::WhiteWinner
|
|
};
|
|
return;
|
|
}
|
|
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 {
|
|
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,
|
|
}
|
|
}
|