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

speed up optim

Browse Source
This commit is contained in:
drake
2026-08-25 13:46:00 -04:00
parent e327de00a3
commit d3bc5fbd53
4 changed files with 98 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/main.rs
+1 -1
View File
@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
fn main() {
let device = Device::cuda(DeviceIndex::default());
let mcts_config = MctsConfig::new(100, 1.0, 0.05, 0.25, 128);
let mcts_config = MctsConfig::new(150, 1.0, 0.05, 0.25, 128);
let adam_config = AdamConfig::new();
engine/src/mcts.rs
+39 -29
View File
@@ -1,9 +1,8 @@
// use crate::networks::PolicyNetwork;
use crate::mcts::BoardStateStatus::{FiftyMove, Ongoing, Threefold};
use crate::net::encoding::{encode_board_state_perspective, encode_move};
use crate::net::encoding::{encode_board_states_batch, encode_move};
use crate::net::model::ChessModel;
use burn::prelude::Device;
use burn::Tensor;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
@@ -37,18 +36,34 @@ 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| {
ucb_score(self, &arena[a], c_puct)
.partial_cmp(&ucb_score(self, &arena[b], c_puct))
.unwrap()
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)
})
.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);
let value_score = if child.visit_count > 0 {
-child.value()
} else {
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())
@@ -156,6 +171,9 @@ 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,
@@ -180,7 +198,6 @@ impl Mcts {
while sims_done < num_sims {
let mut leaf_nodes: Vec<usize> = Vec::new();
let mut leaf_paths: Vec<Vec<usize>> = Vec::new();
let mut leaf_states: Vec<Tensor<4>> = Vec::new();
let target_batch = std::cmp::min(self.config.batch_max, num_sims - sims_done);
@@ -197,14 +214,21 @@ impl Mcts {
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 except leaf itself
// 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;
@@ -217,16 +241,13 @@ impl Mcts {
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);
let state: Tensor<4> =
encode_board_state_perspective(&nodes[current].board_state, device)
.reshape([1, 18, 8, 8]);
leaf_states.push(state);
}
if leaf_nodes.is_empty() {
@@ -236,8 +257,9 @@ impl Mcts {
continue;
}
// Batch evaluate non-terminal leaf states
let batch = Tensor::cat(leaf_states, 0);
// 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();
@@ -253,6 +275,8 @@ impl Mcts {
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;
@@ -324,9 +348,7 @@ impl Mcts {
if let Some(term_val) = terminal_value(status, stm) {
term_val
} else {
let state: Tensor<4> =
encode_board_state_perspective(&arena[node_idx].board_state, device)
.reshape([1, 18, 8, 8]);
let state = encode_board_states_batch(&[&arena[node_idx].board_state], device);
let (policy_head, value_head) = model.forward(state);
let legal_moves: Vec<ChessMove> =
@@ -428,19 +450,7 @@ fn dirichlet_sample(size: usize, alpha: f32) -> Vec<f32> {
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)
engine/src/net/encoding.rs
+50 -65
View File
@@ -1,26 +1,6 @@
use crate::mcts::BoardState;
use burn::prelude::*;
use chess::{ChessMove, Color, File, Piece, Rank, Square, ALL_SQUARES};
/*
Input planes:
1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
7-12: opponent pieces
13: your kingside
14: your queenside
15: opponent kingside
16: opponent queenside
17: rep >= 2
18: En passant square (gap pass over)
Output:
4672 moves represented through planes
1-56: sliding moves
57-64: Knight moves
65-73: underpromotions
*/
const PLANES_PER_SQUARE: usize = 73;
use chess::{ChessMove, Color, File, Piece, Rank, Square};
const SLIDING_DIRS: [(i8, i8); 8] = [
(0, 1), // N
@@ -54,99 +34,104 @@ const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook
pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
let mut planes = vec![0.0f32; 18 * 64];
encode_board_state_to_slice(state, &mut planes);
Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
}
/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
out_slice[..1152].fill(0.0);
let board = &state.board;
let us = board.side_to_move();
let them = !us;
let flip = us == Color::Black;
let idx = |plane: usize, rank: usize, file: usize| -> usize { plane * 64 + rank * 8 + file };
let flip = us == Color::Black;
for &square in ALL_SQUARES.iter() {
if let Some(piece) = board.piece_on(square) {
let color = board.color_on(square).unwrap();
let piece_index = match piece {
Piece::Pawn => 0,
Piece::Knight => 1,
Piece::Bishop => 2,
Piece::Rook => 3,
Piece::Queen => 4,
Piece::King => 5,
};
// Determine plane based on perspective
let plane = if color == us {
piece_index
} else {
piece_index + 6
};
let pieces = [
Piece::Pawn,
Piece::Knight,
Piece::Bishop,
Piece::Rook,
Piece::Queen,
Piece::King,
];
for (piece_index, &piece) in pieces.iter().enumerate() {
let us_bb = board.pieces(piece) & board.color_combined(us);
for square in us_bb {
let mut rank = square.get_rank().to_index();
let mut file = square.get_file().to_index();
// Flip board if black to move
if flip {
rank = 7 - rank;
file = 7 - file;
}
out_slice[idx(piece_index, rank, file)] = 1.0;
}
planes[idx(plane, rank, file)] = 1.0;
let them_bb = board.pieces(piece) & board.color_combined(them);
for square in them_bb {
let mut rank = square.get_rank().to_index();
let mut file = square.get_file().to_index();
if flip {
rank = 7 - rank;
file = 7 - file;
}
out_slice[idx(piece_index + 6, rank, file)] = 1.0;
}
}
// -------------------------
// Castling Rights (planes 12–15)
// From perspective
// -------------------------
let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
if us_castle.has_kingside() {
fill_plane(&mut planes, 12);
fill_plane(out_slice, 12);
}
if us_castle.has_queenside() {
fill_plane(&mut planes, 13);
fill_plane(out_slice, 13);
}
if them_castle.has_kingside() {
fill_plane(&mut planes, 14);
fill_plane(out_slice, 14);
}
if them_castle.has_queenside() {
fill_plane(&mut planes, 15);
fill_plane(out_slice, 15);
}
// -------------------------
// Repetition plane (16)
// -------------------------
let current_hash = board.get_hash();
if let Some(count) = state.repetition_table.get(&current_hash) {
if *count >= 2 {
fill_plane(&mut planes, 16);
fill_plane(out_slice, 16);
}
}
let en_passant = board.en_passant();
if let Some(ep) = en_passant {
// En passant plane (17)
if let Some(ep) = board.en_passant() {
let mut rank = ep.get_rank().to_index();
let mut file = ep.get_file().to_index();
// Flip board if black to move
if flip {
rank = 7 - rank;
file = 7 - file;
}
planes[idx(17, rank, file)] = 1.0;
out_slice[idx(17, rank, file)] = 1.0;
}
}
Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
let b = states.len();
let mut buffer = vec![0.0f32; b * 1152];
for (i, state) in states.iter().enumerate() {
let start = i * 1152;
encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
}
Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
}
#[inline]
fn fill_plane(buffer: &mut [f32], plane: usize) {
let start = plane * 64;
for i in 0..64 {
buffer[start + i] = 1.0;
}
buffer[start..start + 64].fill(1.0);
}
/// Rotate a square into the side-to-move perspective.
engine/src/train.rs
+7 -7
View File
@@ -4,9 +4,9 @@ use crate::net::model::{
ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample,
};
use burn::data::dataloader::batcher::Batcher;
use burn::module::{AutodiffModule, Module};
use burn::optim::{AdamConfig, GradientsParams, Optimizer};
use burn::prelude::Device;
use burn::module::AutodiffModule;
use burn::optim::{AdamConfig, GradientsParams};
use burn::prelude::{Device, Module};
use burn::store::ModuleRecord;
use chess::{ChessMove, Color};
use rand::rngs::SmallRng;
@@ -119,9 +119,9 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
let mut ply_count: usize = 0;
while board_state.status == BoardStateStatus::Ongoing {
// let before = Instant::now();
let before = Instant::now();
let results = mcts.search(&board_state, &infer_model, &device);
// println!("{}", before.elapsed().as_millis());
println!("{}", before.elapsed().as_millis());
episode_buffer.push(results);
let temp = if ply_count < 30 { 1.0 } else { 0.0 };
@@ -130,7 +130,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
let stm = board_state.board.side_to_move();
let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
// println!("playing move: {}", mv);
println!("playing move: {}", mv);
if episode == 0 {
game_hist.push_str(mv.to_string().as_str());
game_hist.push(' ');
@@ -200,7 +200,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
let batcher = ChessBatcher {};
let batch = batcher.batch(samples, &device);
let batch = batcher.batch(samples, &autodiff_device);
let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);