diff options
| -rw-r--r-- | .gitignore | 4 | ||||
| -rw-r--r-- | CLAUDE.md | 11 | ||||
| -rw-r--r-- | plans.org | 27 | ||||
| -rw-r--r-- | src/board.rs | 211 | ||||
| -rw-r--r-- | src/claude-player.rs | 2 | ||||
| -rw-r--r-- | src/game.rs | 2 | ||||
| -rw-r--r-- | src/main.rs | 59 | ||||
| -rw-r--r-- | src/network.rs | 425 | ||||
| -rw-r--r-- | src/neural-player.rs | 68 | ||||
| -rw-r--r-- | src/player.rs | 27 | ||||
| -rw-r--r-- | src/train.rs | 353 |
11 files changed, 1162 insertions, 27 deletions
@@ -9,4 +9,6 @@ /.direnv # Temporary removing github workflows -.github
\ No newline at end of file +.github +# Trained networks (cargo run --release -- train) +/weights @@ -19,6 +19,7 @@ The toolchain comes from a Nix flake (`flake.nix`, loaded via direnv `.envrc` ↠### Benchmarking and profiling the search +- Train: `cargo run --release -- train [generations] [games] [out-dir]` (defaults 50, 500, `weights/`, which is gitignored), then play it with `cargo run --release -- neural weights/latest.bin [sheep|tiger]`. - Bench: `cargo run --profile profiling -- bench [depth]` (default depth 6; 9 takes about 0.6 s) searches six fixed positions to a fixed depth without the TUI and prints the move, score and time for each. The positions come from `ClaudePlayer` playing itself to depth 4, so they are deterministic. An optimisation that only makes the search faster must leave every move and score unchanged. - Counters: add `--features search-stats` to print nodes, nodes/s, cache hit/overwrite rates and the share of cutoffs on the first move. They are thread-local, bumped with the `count!` macro in `player.rs` (also used in `board.rs`), and compile to nothing without the feature. Run `cargo test` with and without the feature. - The `profiling` Cargo profile is release plus debug info. The dev shell has `perf` and `samply`. As a non-root user, `perf record` needs `-m 16` (a small ring buffer) with `--call-graph dwarf`, otherwise mmap fails on the `perf_event_mlock_kb` limit: `perf record -e cycles:u -m 16 --call-graph dwarf target/profiling/tigers-and-sheep-rs bench 9`. @@ -27,9 +28,9 @@ The toolchain comes from a Nix flake (`flake.nix`, loaded via direnv `.envrc` ↠## Architecture - `src/board.rs` — all game logic and its unit tests (`#[cfg(test)] mod tests` at the bottom). -- `src/player.rs` — the `Player` trait. Players are polled, not blocking: the coordinator calls `update(board, turn)` before each draw, `pick_move(board, turn) -> Option<Move>` once per loop (computer players search and return `Some`), and passes clicks/keys to `handle_input(Input, board, turn) -> Option<Move>`. `highlights()` gives points for the board to highlight, `reset()` is called on both players after a revert (drop any selection and last-move comments; it is *not* called after a normal move, so players clear their own selection when they return a move), `render(area, buf)` draws the current player's panel (default: its name), and `render_last_turn(area, buf)` draws end-of-turn comments on the move it just made (default: nothing). `impl Widget for &dyn Player` makes every player a ratatui widget. Also the private-ish `MinimaxPlayer` trait: implementors supply `evaluate` (positive favours tigers) and optionally `depth()` (at most `MAX_DEPTH` = 64)/`time_limit()`/`name()`, plus `set_last_search`/`last_search` to keep a `SearchReport` (move, depth reached, score) that the blanket impl's `render_last_turn` shows; a blanket impl makes every `MinimaxPlayer` a `Player`. +- `src/player.rs` — the `Player` trait. Players are polled, not blocking: the coordinator calls `update(board, turn)` before each draw, `pick_move(board, turn) -> Option<Move>` once per loop (computer players search and return `Some`), and passes clicks/keys to `handle_input(Input, board, turn) -> Option<Move>`. `highlights()` gives points for the board to highlight, `reset()` is called on both players after a revert (drop any selection and last-move comments; it is *not* called after a normal move, so players clear their own selection when they return a move), `render(area, buf)` draws the current player's panel (default: its name), and `render_last_turn(area, buf)` draws end-of-turn comments on the move it just made (default: nothing). `impl Widget for &dyn Player` makes every player a ratatui widget. Also the private-ish `MinimaxPlayer` trait: implementors supply `evaluate(board, turn)` (positive favours tigers; `turn` is the side to move) and optionally `depth()` (at most `MAX_DEPTH` = 64)/`time_limit()`/`name()`, plus `set_last_search`/`last_search` to keep a `SearchReport` (move, depth reached, score) that the blanket impl's `render_last_turn` shows; a blanket impl makes every `MinimaxPlayer` a `Player`. - Search: alpha-beta (`minimax` checks the cache, `minimax_uncached` searches) with iterative deepening in `pick_move` until half the time limit has passed. Moves are tried in this order: the cache's best move (the `hint`), then tiger jumps (`get_moves` lists them first), then killer moves, then the rest. This is done by rotating moves to the front, because a `sort_by_key` here cost more than it saved. - - `cache()` returns an optional `ScoreCache`, keyed by board, side to move and remaining depth, storing `CacheEntry` bounds. `TranspositionTable` is the fixed-size one (`new(bits)`: `1 << bits` slots of ~24 bytes, in two-slot buckets where one slot keeps the deepest entry; `clear()` bumps a generation number instead of wiping). The `HashMap` impl is kept as a simple reference for tests. + - `cache()` returns an optional `ScoreCache`, keyed by board, side to move and remaining depth. `minimax` passes it `board.canonical()`, so the 8 symmetric boards share one entry, and turns the entry's best move between orientations with `transform_move`/`inverse_symmetry`. This assumes `evaluate` gives symmetric boards the same score. `ClaudePlayer`'s does; a `Network` only does so approximately, since it learns symmetry from augmented data. storing `CacheEntry` bounds. `TranspositionTable` is the fixed-size one (`new(bits)`: `1 << bits` slots of ~24 bytes, in two-slot buckets where one slot keeps the deepest entry; `clear()` bumps a generation number instead of wiping). The `HashMap` impl is kept as a simple reference for tests. - `killers()` returns optional `Killers`: two quiet moves per remaining depth that last caused a cutoff. `pick_move` clears both the cache and the killers. - Tests compare the search with plain minimax (`reference`) for every cache kind (`CacheKind::None`/`Map`/`TinyTable`, a table small enough to collide constantly) with and without killers. - `src/game.rs` — `GameHandler::play_game(sheep_player, tiger_player)` runs a ratatui TUI (mouse capture on) until the user quits, returning the winner if there was one. Layout: board | move history on top; below, the current player's panel (70%) | the previous player's last-turn panel (30%). Panes overlap by one cell (`Layout::spacing(-1)`) and are built with `panel()` (`MergeStrategy::Exact`) so shared borders join up. The bottom row is drawn first and keeps its titles on its bottom border, because its top border is the same line as the board's bottom border, which lists the keybinds (keep them in step with `handle_event`). History is `Vec<(Move, GameBoard)>` (board *after* each move); the side to move comes from its length's parity (sheep first). Clicking a history row calls `revert(i)`, which goes back to the board *before* move `i`. The loop only waits for input (`event::read`) when `pick_move` returns `None`. After a computer's search, `pending_action` handles any input that came in meanwhile, so quit and revert still work in computer-vs-computer games. They take effect once the current search finishes. @@ -37,7 +38,10 @@ The toolchain comes from a Nix flake (`flake.nix`, loaded via direnv `.envrc` ↠- `src/local-player.rs` (module `local_player`, pulled in via `#[path]` because of the hyphen) — `LocalPlayer`, a human using the mouse: click a point to place a sheep, or click a piece then one of its highlighted destinations. - `src/claude-player.rs` — `ClaudePlayer`, the minimax computer player: `time_limit` and `max_depth` are public fields, with a 2^20-slot `TranspositionTable` and killer moves. - `src/bench.rs` — the `bench` mode (see above). -- `src/main.rs` — `bench [depth]` runs the benchmark; otherwise it plays a game in the TUI (the sheep player first, then the tigers) and prints the winner. Which players it uses changes: `LocalPlayer` or `ClaudePlayer` on either side. +- `src/network.rs` — `Network`, a small MLP (52 inputs → 64 ReLU → 32 ReLU → tanh, positive favours tigers) with hand-written backprop and an `Adam` optimiser, all in one flat `params` Vec at the `W1`/`B1`/… offsets. Inputs are only the raw position: a bit per point for tigers and for sheep, sheep in hand / 20, and whether tigers are to move. `forward` adds up the layer-1 weight rows of the set bits rather than multiplying through the zeros. Also `Rng` (xorshift64*) and `Sample` (board, turn, target). File format (`save`/`load`): `TSNN`, then version, inputs, hidden1, hidden2 as little-endian u32s, then every param as a little-endian f32; changing the layer sizes makes old files fail to load. +- `src/neural-player.rs` — `NeuralPlayer`, a `MinimaxPlayer` like `ClaudePlayer` whose `evaluate` is a `Network`. +- `src/train.rs` — self-play training. Each generation plays `games_per_generation` games of the net against itself at a fixed depth (default 3) across all cores. Moves are random for the first `random_plies`, and with chance `epsilon` after that. A game reaching `ply_cap` is a draw, since the rules have no draw. Each position is labelled with a TD(λ) target worked backwards from the result (±1, or 0) through the searched scores. A position where a random move was played is labelled with its own search score instead, which cuts the chain there. Without that cut, the sheep's random moves (which usually hang a sheep) taught the net that tigers nearly always win. Self-play still ends mostly in tiger wins, because of those random moves, so judge progress by the match against `ClaudePlayer` (`match_games` = 100). The net then trains on the last `buffer_generations` generations, each position shown in a random one of the board's 8 symmetries (`GameBoard::transformed`). After that it plays matches against the previous generation and against `ClaudePlayer`, which is only a yardstick and never a training signal. Each generation is saved as `out_dir/gen-NNN.bin` and `latest.bin`, and a rerun carries on from the newest `gen-NNN.bin`. +- `src/main.rs` — `bench [depth]` runs the benchmark; `train [generations] [games per generation] [out dir]` trains a network (use `--release`); `neural <weights> [sheep|tiger]` plays you, on that side (default sheep), against a `NeuralPlayer`; otherwise it plays a game in the TUI (the sheep player first, then the tigers) and prints the winner. Which players it uses changes: `LocalPlayer` or `ClaudePlayer` on either side. Notation (`pos_name`, `impl Display for Move` in board.rs): files `a`–`e` are x 0–4, ranks are `5 - y` (y = 0 is drawn at the top as rank 5). Moves are `sc3` (place), `sc3c2` (sheep move), `ta5b4` (tiger step), `ta5xc3` (tiger jump). @@ -52,6 +56,7 @@ The board is two `u32` bitmasks, `tiger_mask` and `sheep_mask`. This encoding is - `Pos = (u8, u8)` is `(x, y)`; out-of-range coordinates panic in `bit_idx`. - Default layout: tigers in the four corners, 20 sheep in hand. `GameBoard::EMPTY` has no pieces and none in hand. - `key()` packs the whole board into one `u64` (used for the transposition table's hash). +- `tiger_points()`/`sheep_points()` give the board bits alone. `transformed(sym)` (sym in `0..SYMMETRIES` = 8, the point map is `transform_pos`) applies one of the board's rotations or reflections, with 5 row lookups per mask (`ROW_MAPS`). `canonical()` is the transform with the smallest `key`, plus the `sym` that reaches it. All 8 keep the lines, since `x + y` keeps its parity, so the result is an equally good position. ### Rules as implemented @@ -35,8 +35,35 @@ All at depth 9 on the TSC clock source; scores and moves the same as the baselin - Ordering: a =sort_by_key= cost ~26% of cycles and made it slower overall than no ordering; =get_moves= now lists jumps first and the search rotates killers and the hint to the front instead. - Killers on their own cut nodes 10% but made the opening search bigger (0.93 M -> 1.04 M nodes) +** DONE Symmetric cache +=minimax= looks positions up by =GameBoard::canonical= (smallest =key= of the 8 rotations and +reflections), storing best moves as played on the canonical board. Depth 9 on the HPET clock +source, so the times are only comparable with each other; moves and scores unchanged. +| Step | Time ms | Nodes | Nodes/s (stats build) | +|-------------------------+---------+--------+-----------------------| +| Before | 1170 | 1.67 M | 1.44 M | +| Canonical cache keys | 570 | 0.71 M | 1.24 M | +- The opening (fully symmetric) shrank most: 1.04 M -> 0.23 M nodes. Each probe now does 7 board + transforms of 10 row lookups each, which costs ~14% per node. *** Next candidates - Check the deadline every N nodes instead of at every node, so a slow clock source can't cost half the search - =evaluate=: tiger moves and trapped tigers from bitmasks rather than =get_moves= + 25 =get()= calls - Table lookups are now the top item (~28% of cycles), mostly cache misses into the 24 MB table * Draws +* Neural player training +** DONE Tiger bias in self-play +Before the fix, tigers won ~97% of self-play games and the net only reached ~15-20% against +=ClaudePlayer= at depth 3. 200 games each, tigers/sheep/draws: +| Players | Random opening plies | eps | Result | +|-----------+----------------------+------+------------| +| Claude d3 | 10 | 0 | 78/22/100 | +| Claude d3 | 10 | 0.05 | 171/15/14 | +| Neural d3 | 4 | 0 | 197/0/3 | +- Random sheep moves usually hang a sheep, so eps exploration alone tilts games to the tigers, and + TD targets carried those wins back to the positions before the blunders. +- Fix: a position where a random move was played is labelled with its own search score, cutting + the TD chain. Neural d3, 4 random plies, eps 0 became 126/4/70, and against =ClaudePlayer= (100 + games, depth 3) it went 12% at gen 5, 50% at gen 50, 50-62% at gens 60-100. +*** Next candidates +- Pick exploratory moves among near-best ones (softmax over root scores) rather than uniformly +- Sheep are still the weaker side: report the match score for each side separately diff --git a/src/board.rs b/src/board.rs index 1c21c3e..087ac0e 100644 --- a/src/board.rs +++ b/src/board.rs @@ -218,6 +218,94 @@ fn pos_of(idx: usize) -> Pos ((idx % 5) as u8, (idx / 5) as u8) } +/// The board's 8 symmetries (rotations and reflections) as a point map: bit 2 of `sym` swaps x and +/// y, then bit 0 mirrors x and bit 1 mirrors y. They all keep the board lines, since x + y keeps +/// its parity, so a transformed position is just as legal and just as good for either side. +pub const SYMMETRIES: usize = 8; + +pub const fn transform_pos((x, y): Pos, sym: usize) -> Pos +{ + let (x, y) = if sym & 4 != 0 { (y, x) } else { (x, y) }; + let x = if sym & 1 != 0 { 4 - x } else { x }; + let y = if sym & 2 != 0 { 4 - y } else { y }; + (x, y) +} + +/// For each symmetry, where each bit index goes. +const SYMMETRY_MAPS: [[u8; 25]; SYMMETRIES] = { + let mut maps = [[0; 25]; SYMMETRIES]; + let mut sym = 0; + while sym < SYMMETRIES + { + let mut idx = 0; + while idx < 25 + { + let (x, y) = transform_pos(((idx % 5) as u8, (idx / 5) as u8), sym); + maps[sym][idx] = x + y * 5; + idx += 1; + } + sym += 1; + } + maps +}; + +/// For each symmetry and row, where each of the 32 possible sets of points in that row go, so a +/// whole mask transforms in 5 lookups. The search does this for every cache probe. +const ROW_MAPS: [[[u32; 32]; 5]; SYMMETRIES] = { + let mut maps = [[[0; 32]; 5]; SYMMETRIES]; + let mut sym = 0; + while sym < SYMMETRIES + { + let mut row = 0; + while row < 5 + { + let mut bits = 0; + while bits < 32 + { + let mut x = 0; + while x < 5 + { + if bits & 1 << x != 0 + { + maps[sym][row][bits] |= 1 << SYMMETRY_MAPS[sym][x + row * 5]; + } + x += 1; + } + bits += 1; + } + row += 1; + } + sym += 1; + } + maps +}; + +#[inline(always)] +fn transform_mask(mask: u32, sym: usize) -> u32 +{ + let rows = &ROW_MAPS[sym]; + rows[0][(mask & 31) as usize] | rows[1][(mask >> 5 & 31) as usize] | rows[2][(mask >> 10 & 31) as usize] + | rows[3][(mask >> 15 & 31) as usize] | rows[4][(mask >> 20 & 31) as usize] +} + +/// The symmetry that undoes `sym`. The mirrors undo themselves, but with the swap of x and y +/// first, undoing them means swapping which axis is mirrored. +pub const fn inverse_symmetry(sym: usize) -> usize +{ + if sym & 4 == 0 { sym } else { 4 | (sym & 1) << 1 | (sym >> 1 & 1) } +} + +/// `mv` as played on the board transformed by `sym`. +pub fn transform_move(mv: Move, sym: usize) -> Move +{ + match mv + { + Move::PlaceSheep { pos } => Move::PlaceSheep { pos: transform_pos(pos, sym) }, + Move::MoveSheep { from, to } => Move::MoveSheep { from: transform_pos(from, sym), to: transform_pos(to, sym) }, + Move::MoveTiger { from, to } => Move::MoveTiger { from: transform_pos(from, sym), to: transform_pos(to, sym) }, + } +} + impl GameBoard { /// No pieces on the board and none in hand. @@ -266,6 +354,37 @@ impl GameBoard /// The whole board, sheep in hand included, as one number: different boards give different keys. pub fn key(&self) -> u64 { (self.sheep_mask as u64) << 32 | self.tiger_mask as u64 } pub fn sheep_count(&self) -> u32 { (self.sheep_mask & BOARD_MASK).count_ones() } + /// The points holding a tiger, as bits `x + y*5`. + pub fn tiger_points(&self) -> u32 { self.tiger_mask } + /// The points holding a sheep, as bits `x + y*5`. + pub fn sheep_points(&self) -> u32 { self.sheep_mask & BOARD_MASK } + + /// The board under symmetry `sym` (see `transform_pos`), with the same sheep in hand. + pub fn transformed(&self, sym: usize) -> GameBoard + { + GameBoard + { + tiger_mask: transform_mask(self.tiger_mask, sym), + sheep_mask: transform_mask(self.sheep_points(), sym) | (self.sheep_mask & !BOARD_MASK), + } + } + + /// One board standing for all 8 symmetric ones (the one whose `key` is smallest), and the + /// symmetry that takes `self` to it. Symmetric boards are equally good, so the search caches + /// them under one entry. + pub fn canonical(&self) -> (GameBoard, usize) + { + let mut best = (*self, 0); + for sym in 1..SYMMETRIES + { + let t = self.transformed(sym); + if t.key() < best.0.key() + { + best = (t, sym); + } + } + best + } /// Empty board points. fn empty(&self) -> u32 @@ -1246,4 +1365,96 @@ mod tests assert!(board.check_win(Turn::Sheep)); assert!(!board.check_win(Turn::Tiger)); } + + // --- symmetries --- + + #[test] + fn symmetries_are_distinct_bijections() + { + for sym in 0..SYMMETRIES + { + let mut seen = 0u32; + for idx in 0..25 + { + seen |= 1 << SYMMETRY_MAPS[sym][idx]; + } + assert_eq!(seen, BOARD_MASK, "symmetry {sym} is not a bijection"); + for other in 0..sym + { + assert_ne!(SYMMETRY_MAPS[sym], SYMMETRY_MAPS[other], "symmetries {sym} and {other}"); + } + } + assert_eq!(transform_pos((1, 0), 0), (1, 0)); + assert_eq!(GameBoard::default().transformed(5), GameBoard::default()); + } + + #[test] + fn inverse_symmetry_undoes_it() + { + for sym in 0..SYMMETRIES + { + for idx in 0..25 + { + let pos = pos_of(idx); + assert_eq!(transform_pos(transform_pos(pos, sym), inverse_symmetry(sym)), pos, "symmetry {sym}"); + } + } + } + + #[test] + fn canonical_is_shared_by_symmetric_boards() + { + let mut seed: u64 = 0x94D049BB133111EB; + let mut next = || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed }; + let mut board = GameBoard::default(); + let mut turn = Turn::Sheep; + for _ in 0..60 + { + if board.check_win(turn) { break; } + let (canon, sym) = board.canonical(); + assert_eq!(board.transformed(sym), canon); + for s in 0..SYMMETRIES + { + let t = board.transformed(s); + assert_eq!(t.canonical().0, canon, "symmetry {s} of {board:?}"); + // The fast transform agrees with moving each point + let slow = |mask: u32| indices(mask).fold(0, |m, idx| m | 1 << SYMMETRY_MAPS[s][idx]); + assert_eq!((t.tiger_points(), t.sheep_points()), (slow(board.tiger_points()), slow(board.sheep_points()))); + } + let moves = board.get_moves(turn); + board.make_move_unchecked(moves[(next() % moves.len() as u64) as usize]); + turn = turn.next_turn(); + } + } + + #[test] + fn transformed_boards_have_transformed_moves() + { + let mut seed: u64 = 0xD1B54A32D192ED03; + let mut next = || { seed ^= seed << 13; seed ^= seed >> 7; seed ^= seed << 17; seed }; + let mut board = GameBoard::default(); + let mut turn = Turn::Sheep; + for _ in 0..80 + { + if board.check_win(turn) { break; } + for sym in 0..SYMMETRIES + { + let t = board.transformed(sym); + assert_eq!(t.sheep_left(), board.sheep_left()); + assert_eq!(t.sheep_count(), board.sheep_count()); + for side in [Turn::Tiger, Turn::Sheep] + { + let mut expected: Vec<Move> = board.get_moves(side).into_iter().map(|m| transform_move(m, sym)).collect(); + expected.sort(); + let mut moves = t.get_moves(side); + moves.sort(); + assert_eq!(moves, expected, "symmetry {sym}, {side:?} on {board:?}"); + assert_eq!(t.check_win(side), board.check_win(side)); + } + } + let moves = board.get_moves(turn); + board.make_move_unchecked(moves[(next() % moves.len() as u64) as usize]); + turn = turn.next_turn(); + } + } } diff --git a/src/claude-player.rs b/src/claude-player.rs index 5cbd794..c25c5e1 100644 --- a/src/claude-player.rs +++ b/src/claude-player.rs @@ -39,7 +39,7 @@ impl MinimaxPlayer for ClaudePlayer /// - jumps the tigers could make right now (sheep left hanging); /// - tigers boxed in with no move at all, since trapping every tiger wins for the sheep; /// - tiger mobility in general, since the sheep win by squeezing it to zero. - fn evaluate(&mut self, board: &GameBoard) -> f32 + fn evaluate(&mut self, board: &GameBoard, _turn: Turn) -> f32 { let dead = TOTAL_SHEEP - board.sheep_count() - board.sheep_left(); diff --git a/src/game.rs b/src/game.rs index 11a77fd..5aef809 100644 --- a/src/game.rs +++ b/src/game.rs @@ -87,7 +87,7 @@ impl GameHandler /// Play in the terminal with `p1` as the sheep and `p2` as the tigers, until someone wins and /// the user quits. Returns the winner, or None if the user quit before the game was won. - pub fn play_game<T1: Player, T2: Player>(&mut self, p1: &mut T1, p2: &mut T2) -> io::Result<Option<Turn>> + pub fn play_game(&mut self, p1: &mut dyn Player, p2: &mut dyn Player) -> io::Result<Option<Turn>> { let mut terminal = ratatui::init(); let result = execute!(io::stdout(), EnableMouseCapture) diff --git a/src/main.rs b/src/main.rs index 16467b0..8f6968d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,6 @@ -use std::time::Duration; +use std::{path::Path, time::Duration}; -use crate::{claude_player::ClaudePlayer, game::GameHandler, local_player::LocalPlayer}; +use crate::{board::Turn, claude_player::ClaudePlayer, game::GameHandler, local_player::LocalPlayer, network::Network, neural_player::NeuralPlayer, player::Player}; mod board; mod player; @@ -12,20 +12,59 @@ mod local_player; #[path = "claude-player.rs"] mod claude_player; mod bench; +mod network; +#[path = "neural-player.rs"] +mod neural_player; +mod train; fn main() -> std::io::Result<()> { - // `bench [depth]` times the computer player's search without the TUI let args: Vec<String> = std::env::args().collect(); - if args.get(1).is_some_and(|a| a == "bench") + let arg = |i: usize| args.get(i).map(String::as_str); + let number = |i: usize, default: usize| arg(i).map_or(default, |n| n.parse().expect("expected a number")); + match arg(1) { - let depth = args.get(2).map_or(6, |d| d.parse().expect("depth should be a number")); - bench::run(depth); - return Ok(()); + // `bench [depth]` times the computer player's search without the TUI + Some("bench") => + { + bench::run(number(2, 6)); + return Ok(()); + } + // `train [generations] [games per generation] [out dir]` trains a network by self-play + Some("train") => + { + let defaults = train::Config::default(); + let cfg = train::Config + { + generations: number(2, defaults.generations), + games_per_generation: number(3, defaults.games_per_generation), + out_dir: arg(4).map_or(defaults.out_dir.clone(), Into::into), + ..defaults + }; + return train::run(&cfg); + } + // `neural <weights> [sheep|tiger]`: you play the given side against a trained network + Some("neural") => + { + let net = Network::load(Path::new(arg(2).expect("usage: neural <weights file> [sheep|tiger]")))?; + let side = match arg(3) { Some("tiger") => Turn::Tiger, _ => Turn::Sheep }; + let mut computer = NeuralPlayer::new(net, Duration::from_secs(2), 20); + let mut human = LocalPlayer::default(); + let (sheep, tiger): (&mut dyn Player, &mut dyn Player) = match side + { + Turn::Sheep => (&mut human, &mut computer), + Turn::Tiger => (&mut computer, &mut human), + }; + return play(sheep, tiger); + } + _ => (), } + play(&mut ClaudePlayer::new(Duration::from_secs(2)), &mut ClaudePlayer::new(Duration::from_secs(2))) +} + +fn play(sheep: &mut dyn Player, tiger: &mut dyn Player) -> std::io::Result<()> +{ let mut game = GameHandler::default(); - let winner = game.play_game(&mut ClaudePlayer::new(Duration::from_secs(2)), - &mut ClaudePlayer::new(Duration::from_secs(2)))?; - if let Some(winner) = winner + if let Some(winner) = game.play_game(sheep, tiger)? { println!("{winner:?} wins!"); } diff --git a/src/network.rs b/src/network.rs new file mode 100644 index 0000000..5836798 --- /dev/null +++ b/src/network.rs @@ -0,0 +1,425 @@ +use std::{fs, io, path::Path}; + +use crate::board::{GameBoard, Turn}; + +/// Inputs: a bit per point for tigers (0..25) and for sheep (25..50), the sheep in hand over 20 +/// (`HAND`) and 1 if tigers are to move (`TIGER_TO_MOVE`). Nothing else, so the network has to +/// work out for itself what makes a position good. +pub const INPUTS: usize = 52; +const HAND: usize = 50; +const TIGER_TO_MOVE: usize = 51; +pub const HIDDEN1: usize = 64; +pub const HIDDEN2: usize = 32; + +// All the weights live in one flat Vec, at these offsets, so the optimiser, the gradient and +// save/load can treat them as a single list of numbers. +/// Layer 1 weights, input-major: input i's weight to hidden unit j is at `W1 + i*HIDDEN1 + j`, +/// so an input's weights are contiguous and adding up the set bits' columns is cheap. +const W1: usize = 0; +const B1: usize = W1 + INPUTS * HIDDEN1; +/// Layer 2 weights, unit-major: hidden1 unit j's weight to hidden2 unit k is at `W2 + k*HIDDEN1 + j`. +const W2: usize = B1 + HIDDEN1; +const B2: usize = W2 + HIDDEN1 * HIDDEN2; +const W3: usize = B2 + HIDDEN2; +const B3: usize = W3 + HIDDEN2; +pub const PARAMS: usize = B3 + 1; + +const MAGIC: &[u8; 4] = b"TSNN"; +const VERSION: u32 = 1; + +/// A small, fast pseudo-random generator (xorshift64*), good enough for weight init and exploration. +#[derive(Clone, Debug)] +pub struct Rng(u64); + +impl Rng +{ + pub fn new(seed: u64) -> Self + { + // Zero is a fixed point of xorshift, and nearby seeds should still diverge quickly + Self(seed.wrapping_mul(0x9E3779B97F4A7C15) | 1) + } + + pub fn next_u64(&mut self) -> u64 + { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545F4914F6CDD1D) + } + + /// Uniform in `0.0..1.0`. + pub fn next_f32(&mut self) -> f32 + { + (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32 + } + + /// Uniform in `0..n`. + pub fn below(&mut self, n: usize) -> usize + { + (((self.next_u64() >> 32) * n as u64) >> 32) as usize + } + + /// Standard normal, by Box-Muller. + pub fn normal(&mut self) -> f32 + { + let u = 1.0 - self.next_f32(); + let v = self.next_f32(); + (-2.0 * u.ln()).sqrt() * (std::f32::consts::TAU * v).cos() + } +} + +/// A position and the score the network should learn to give it (+1 tigers win, -1 sheep win). +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct Sample +{ + pub board: GameBoard, + pub turn: Turn, + pub target: f32, +} + +/// Values of every layer for one position, kept by `forward_full` for backprop. +struct Activations +{ + h1: [f32; HIDDEN1], + h2: [f32; HIDDEN2], + out: f32, +} + +/// A position scorer: inputs → 64 ReLU → 32 ReLU → tanh, positive favouring tigers. +#[derive(Clone, Debug, PartialEq)] +pub struct Network +{ + params: Vec<f32>, +} + +/// Call `f` with the index of each input that is 1, and return the value of the `HAND` input. +#[inline(always)] +fn for_active(board: &GameBoard, turn: Turn, mut f: impl FnMut(usize)) -> f32 +{ + let mut tigers = board.tiger_points(); + while tigers != 0 + { + f(tigers.trailing_zeros() as usize); + tigers &= tigers - 1; + } + let mut sheep = board.sheep_points(); + while sheep != 0 + { + f(25 + sheep.trailing_zeros() as usize); + sheep &= sheep - 1; + } + if turn == Turn::Tiger + { + f(TIGER_TO_MOVE); + } + board.sheep_left() as f32 / 20.0 +} + +impl Network +{ + /// Random weights (He initialisation), so a new network plays roughly at random. + pub fn random(seed: u64) -> Self + { + let mut rng = Rng::new(seed); + let mut params = vec![0.0; PARAMS]; + let mut init = |range: std::ops::Range<usize>, fan_in: usize, params: &mut [f32]| + { + let scale = (2.0 / fan_in as f32).sqrt(); + for p in &mut params[range] + { + *p = rng.normal() * scale; + } + }; + // About a third of the inputs are set at once, so scale layer 1 by that rather than all 52 + init(W1..B1, 20, &mut params); + init(W2..B2, HIDDEN1, &mut params); + init(W3..B3, HIDDEN2, &mut params); + Self { params } + } + + fn forward_full(&self, board: &GameBoard, turn: Turn) -> Activations + { + let p = &self.params; + let mut h1 = [0.0; HIDDEN1]; + h1.copy_from_slice(&p[B1..B1 + HIDDEN1]); + let hand = for_active(board, turn, |i| + { + for (h, w) in h1.iter_mut().zip(&p[W1 + i * HIDDEN1..W1 + (i + 1) * HIDDEN1]) + { + *h += w; + } + }); + for (h, w) in h1.iter_mut().zip(&p[W1 + HAND * HIDDEN1..W1 + (HAND + 1) * HIDDEN1]) + { + *h = (*h + hand * w).max(0.0); + } + let mut h2 = [0.0; HIDDEN2]; + for (k, h) in h2.iter_mut().enumerate() + { + let w = &p[W2 + k * HIDDEN1..W2 + (k + 1) * HIDDEN1]; + let sum: f32 = w.iter().zip(&h1).map(|(w, x)| w * x).sum(); + *h = (p[B2 + k] + sum).max(0.0); + } + let sum: f32 = p[W3..B3].iter().zip(&h2).map(|(w, x)| w * x).sum(); + Activations { h1, h2, out: (p[B3] + sum).tanh() } + } + + /// The network's score for `board` with `turn` to move, in -1..1 (positive favours tigers). + pub fn forward(&self, board: &GameBoard, turn: Turn) -> f32 + { + self.forward_full(board, turn).out + } + + /// Add the gradient of the squared error `(out - target)² / 2` for one sample to `grad`, + /// returning that error. + fn add_gradient(&self, s: &Sample, grad: &mut [f32]) -> f32 + { + let p = &self.params; + let a = self.forward_full(&s.board, s.turn); + let err = a.out - s.target; + // Back through tanh + let d_out = err * (1.0 - a.out * a.out); + grad[B3] += d_out; + let mut d_h2 = [0.0; HIDDEN2]; + for k in 0..HIDDEN2 + { + grad[W3 + k] += d_out * a.h2[k]; + // Back through ReLU: no gradient where the unit was off + d_h2[k] = if a.h2[k] > 0.0 { d_out * p[W3 + k] } else { 0.0 }; + } + let mut d_h1 = [0.0; HIDDEN1]; + for k in 0..HIDDEN2 + { + if d_h2[k] == 0.0 + { + continue; + } + grad[B2 + k] += d_h2[k]; + let row = W2 + k * HIDDEN1; + for j in 0..HIDDEN1 + { + grad[row + j] += d_h2[k] * a.h1[j]; + d_h1[j] += d_h2[k] * p[row + j]; + } + } + for j in 0..HIDDEN1 + { + if a.h1[j] <= 0.0 + { + d_h1[j] = 0.0; + } + grad[B1 + j] += d_h1[j]; + } + let hand = for_active(&s.board, s.turn, |i| + { + for (g, d) in grad[W1 + i * HIDDEN1..W1 + (i + 1) * HIDDEN1].iter_mut().zip(&d_h1) + { + *g += d; + } + }); + for (g, d) in grad[W1 + HAND * HIDDEN1..W1 + (HAND + 1) * HIDDEN1].iter_mut().zip(&d_h1) + { + *g += hand * d; + } + err * err / 2.0 + } + + /// One optimiser step on the mean error over `batch`; returns that mean error. + pub fn train_batch(&mut self, batch: &[Sample], adam: &mut Adam) -> f32 + { + let mut grad = vec![0.0; PARAMS]; + let loss: f32 = batch.iter().map(|s| self.add_gradient(s, &mut grad)).sum(); + let n = batch.len().max(1) as f32; + for g in &mut grad + { + *g /= n; + } + adam.step(&mut self.params, &grad); + loss / n + } + + /// The mean error over `samples`, without training. + pub fn loss(&self, samples: &[Sample]) -> f32 + { + let total: f32 = samples.iter().map(|s| { let e = self.forward(&s.board, s.turn) - s.target; e * e / 2.0 }).sum(); + total / samples.len().max(1) as f32 + } + + /// Saves as "TSNN", then the version and layer sizes as little-endian u32s, then every weight as a little-endian f32. + pub fn save(&self, path: &Path) -> io::Result<()> + { + let mut bytes = Vec::with_capacity(20 + 4 * PARAMS); + bytes.extend_from_slice(MAGIC); + for n in [VERSION, INPUTS as u32, HIDDEN1 as u32, HIDDEN2 as u32] + { + bytes.extend_from_slice(&n.to_le_bytes()); + } + for p in &self.params + { + bytes.extend_from_slice(&p.to_le_bytes()); + } + fs::write(path, bytes) + } + + pub fn load(path: &Path) -> io::Result<Self> + { + let bytes = fs::read(path)?; + let bad = |msg: &str| io::Error::new(io::ErrorKind::InvalidData, format!("{}: {msg}", path.display())); + if bytes.len() != 20 + 4 * PARAMS || &bytes[..4] != MAGIC + { + return Err(bad("not a network file of the expected size")); + } + let word = |i: usize| u32::from_le_bytes(bytes[4 + 4 * i..8 + 4 * i].try_into().unwrap()); + if [word(0), word(1), word(2), word(3)] != [VERSION, INPUTS as u32, HIDDEN1 as u32, HIDDEN2 as u32] + { + return Err(bad("network has a different version or layer sizes")); + } + let params = bytes[20..].chunks_exact(4).map(|c| f32::from_le_bytes(c.try_into().unwrap())).collect(); + Ok(Self { params }) + } +} + +/// The Adam optimiser's state: running averages of each weight's gradient and squared gradient. +pub struct Adam +{ + pub learning_rate: f32, + m: Vec<f32>, + v: Vec<f32>, + steps: i32, +} + +impl Adam +{ + const BETA1: f32 = 0.9; + const BETA2: f32 = 0.999; + const EPSILON: f32 = 1e-8; + + pub fn new(learning_rate: f32) -> Self + { + Self { learning_rate, m: vec![0.0; PARAMS], v: vec![0.0; PARAMS], steps: 0 } + } + + fn step(&mut self, params: &mut [f32], grad: &[f32]) + { + self.steps += 1; + // Corrects the averages' bias towards their zero start + let c1 = 1.0 - Self::BETA1.powi(self.steps); + let c2 = 1.0 - Self::BETA2.powi(self.steps); + for i in 0..params.len() + { + self.m[i] = Self::BETA1 * self.m[i] + (1.0 - Self::BETA1) * grad[i]; + self.v[i] = Self::BETA2 * self.v[i] + (1.0 - Self::BETA2) * grad[i] * grad[i]; + params[i] -= self.learning_rate * (self.m[i] / c1) / ((self.v[i] / c2).sqrt() + Self::EPSILON); + } + } +} + +#[cfg(test)] +mod tests +{ + use super::*; + + /// Positions from a pseudo-random game, with made-up targets. + fn samples(n: usize, seed: u64) -> Vec<Sample> + { + let mut rng = Rng::new(seed); + let mut board = GameBoard::default(); + let mut turn = Turn::Sheep; + let mut out = Vec::new(); + while out.len() < n + { + if board.check_win(turn) + { + board = GameBoard::default(); + turn = Turn::Sheep; + } + out.push(Sample { board, turn, target: rng.next_f32() * 1.6 - 0.8 }); + let moves = board.get_moves(turn); + board.make_move_unchecked(moves[rng.below(moves.len())]); + turn = turn.next_turn(); + } + out + } + + #[test] + fn gradient_matches_finite_differences() + { + let mut net = Network::random(7); + let s = samples(30, 3)[29]; + let mut grad = vec![0.0; PARAMS]; + net.add_gradient(&s, &mut grad); + let loss = |net: &Network| { let e = net.forward(&s.board, s.turn) as f64 - s.target as f64; e * e / 2.0 }; + let mut rng = Rng::new(11); + // The input weights of unset inputs have no gradient, so check some of each layer's, + // plus every weight that really has one + let mut checked = 0; + for _ in 0..4000 + { + let i = rng.below(PARAMS); + let eps = 1e-2; + let old = net.params[i]; + net.params[i] = old + eps; + let up = loss(&net); + net.params[i] = old - eps; + let down = loss(&net); + net.params[i] = old; + let numeric = (up - down) / (2.0 * eps as f64); + let tolerance = 1e-3 + 2e-2 * numeric.abs().max(grad[i].abs() as f64); + // A ReLU switching on or off inside ±eps makes the numeric estimate meaningless; such + // kinks are rare, so skip a weight only where the two halves disagree wildly + let left = (loss(&net) - down) / eps as f64; + let right = (up - loss(&net)) / eps as f64; + if (left - right).abs() > 0.1 * left.abs().max(right.abs()) + 1e-4 + { + continue; + } + assert!((numeric - grad[i] as f64).abs() < tolerance, "weight {i}: numeric {numeric}, backprop {}", grad[i]); + checked += 1; + } + assert!(checked > 3000); + } + + #[test] + fn training_can_fit_a_few_samples() + { + let mut net = Network::random(1); + let data = samples(10, 5); + let mut adam = Adam::new(3e-3); + for _ in 0..2000 + { + net.train_batch(&data, &mut adam); + } + assert!(net.loss(&data) < 1e-3, "loss {}", net.loss(&data)); + } + + #[test] + fn save_and_load_round_trip() + { + let net = Network::random(9); + let path = std::env::temp_dir().join(format!("tsnn-test-{}.bin", std::process::id())); + net.save(&path).unwrap(); + let loaded = Network::load(&path); + std::fs::write(&path, b"nope").unwrap(); + let broken = Network::load(&path); + std::fs::remove_file(&path).unwrap(); + assert_eq!(loaded.unwrap(), net); + assert_eq!(broken.unwrap_err().kind(), io::ErrorKind::InvalidData); + } + + #[test] + fn forward_matches_a_plain_matrix_multiply() + { + let net = Network::random(4); + for s in samples(50, 8) + { + let mut x = [0.0f32; INPUTS]; + let hand = for_active(&s.board, s.turn, |i| x[i] = 1.0); + x[HAND] = hand; + let p = &net.params; + let h1: Vec<f32> = (0..HIDDEN1).map(|j| (p[B1 + j] + (0..INPUTS).map(|i| x[i] * p[W1 + i * HIDDEN1 + j]).sum::<f32>()).max(0.0)).collect(); + let h2: Vec<f32> = (0..HIDDEN2).map(|k| (p[B2 + k] + (0..HIDDEN1).map(|j| h1[j] * p[W2 + k * HIDDEN1 + j]).sum::<f32>()).max(0.0)).collect(); + let out = (p[B3] + (0..HIDDEN2).map(|k| h2[k] * p[W3 + k]).sum::<f32>()).tanh(); + assert!((net.forward(&s.board, s.turn) - out).abs() < 1e-5); + } + } +} diff --git a/src/neural-player.rs b/src/neural-player.rs new file mode 100644 index 0000000..3a9c4a0 --- /dev/null +++ b/src/neural-player.rs @@ -0,0 +1,68 @@ +use std::time::Duration; + +use crate::{board::{GameBoard, Turn}, network::Network, player::{Killers, MinimaxPlayer, ScoreCache, SearchReport, TranspositionTable}}; + +/// A computer player that searches like `ClaudePlayer` but scores positions with a neural network +/// trained by self-play (see `train.rs`), so it has no human idea of what a good position is. +pub struct NeuralPlayer +{ + pub net: Network, + pub time_limit: Duration, + /// The most plies to search, counting the move itself; normally `time_limit` stops the search first. + pub max_depth: usize, + cache: TranspositionTable, + killers: Killers, + last_search: Option<SearchReport>, +} + +impl NeuralPlayer +{ + /// A player with a `1 << table_bits`-slot transposition table. + pub fn new(net: Network, time_limit: Duration, table_bits: u32) -> Self + { + Self { net, time_limit, max_depth: 64, cache: TranspositionTable::new(table_bits), killers: Killers::default(), last_search: None } + } +} + +impl MinimaxPlayer for NeuralPlayer +{ + /// The network's output, in -1..1; forced wins are still ±`WIN` from the search itself. + fn evaluate(&mut self, board: &GameBoard, turn: Turn) -> f32 + { + self.net.forward(board, turn) + } + + fn name(&self) -> String { "Neural".to_string() } + + fn set_last_search(&mut self, report: Option<SearchReport>) { self.last_search = report; } + + fn last_search(&self) -> Option<SearchReport> { self.last_search } + + fn time_limit(&self) -> Duration { self.time_limit } + + fn depth(&self) -> usize { self.max_depth } + + fn cache(&mut self) -> Option<&mut dyn ScoreCache> { Some(&mut self.cache) } + + fn killers(&mut self) -> Option<&mut Killers> { Some(&mut self.killers) } +} + +#[cfg(test)] +mod tests +{ + use super::*; + use crate::player::Player; + + #[test] + fn pick_move_reports_its_search() + { + let mut player = NeuralPlayer::new(Network::random(2), Duration::from_secs(3600), 12); + player.max_depth = 3; + let board = GameBoard::default(); + let m = player.pick_move(board, Turn::Sheep).unwrap(); + assert_eq!(board.check_move(m), Ok(())); + let report = player.last_search.expect("a report"); + assert_eq!((report.mv, report.depth), (m, 3)); + assert_eq!(report.score, player.minimax(board, Turn::Sheep, 3, f32::NEG_INFINITY, f32::INFINITY, None).unwrap()); + } +} diff --git a/src/player.rs b/src/player.rs index 7a19219..0c58a58 100644 --- a/src/player.rs +++ b/src/player.rs @@ -2,7 +2,7 @@ use std::{collections::HashMap, time::{Duration, Instant}}; use ratatui::{buffer::Buffer, crossterm::event::KeyEvent, layout::Rect, widgets::{Paragraph, Widget, Wrap}}; -use crate::board::{is_jump, GameBoard, Move, Pos, Turn}; +use crate::board::{inverse_symmetry, is_jump, transform_move, GameBoard, Move, Pos, Turn}; /// Input the game coordinator passes on to the player whose turn it is. #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -89,7 +89,7 @@ pub struct CacheEntry } /// Score cache for a `MinimaxPlayer`. No defaults, so an implementor must supply all of it. -/// Entries are keyed by the side to move and the remaining search depth as well as the board, +/// `minimax` looks boards up by `GameBoard::canonical`, so symmetric boards share an entry. Entries are keyed by the side to move and the remaining search depth as well as the board, /// since the same board can be reached with either side to move, or at different depths. pub trait ScoreCache { @@ -318,7 +318,11 @@ pub trait MinimaxPlayer /// Returns None if `deadline` passed before the search finished. fn minimax(&mut self, board: GameBoard, turn: Turn, depth: usize, mut alpha: ScoreType, mut beta: ScoreType, deadline: Option<Instant>) -> Option<ScoreType> { - let known = self.cache().and_then(|c| { count!(cache_probes); c.lookup(&board, turn, depth) }); + // Symmetric boards are equally good, so they share an entry, under the canonical board. + // Entries' best moves are stored as played on that board, and turned back for this one. + let (canon, sym) = if self.cache().is_some() { board.canonical() } else { (board, 0) }; + let back = |m: Move| transform_move(m, inverse_symmetry(sym)); + let known = self.cache().and_then(|c| { count!(cache_probes); c.lookup(&canon, turn, depth) }); let mut hint = None; if let Some(k) = known { @@ -335,16 +339,17 @@ pub trait MinimaxPlayer } alpha = alpha.max(k.lower); beta = beta.min(k.upper); - hint = k.best; + hint = k.best.map(back); } // Otherwise try the best move from the previous, one-ply-shallower iteration first if hint.is_none() && depth > 0 { - hint = self.cache().and_then(|c| { count!(cache_probes); c.lookup(&board, turn, depth - 1) }).and_then(|e| e.best); + hint = self.cache().and_then(|c| { count!(cache_probes); c.lookup(&canon, turn, depth - 1) }).and_then(|e| e.best).map(back); } let (score, best) = self.minimax_uncached(board, turn, depth, alpha, beta, hint, deadline)?; if let Some(c) = self.cache() { + let best = best.map(|m| transform_move(m, sym)); // On a fail-low every move was only bounded, so none is known to be best let mut entry = if score <= alpha { @@ -364,7 +369,7 @@ pub trait MinimaxPlayer entry.upper = entry.upper.min(k.upper); entry.best = entry.best.or(k.best); } - c.insert(&board, turn, depth, entry); + c.insert(&canon, turn, depth, entry); } Some(score) } @@ -387,7 +392,7 @@ pub trait MinimaxPlayer if depth == 0 { count!(leaves); - return Some((self.evaluate(&board), None)); + return Some((self.evaluate(&board, turn), None)); } if deadline.is_some_and(|d| Instant::now() >= d) { @@ -471,8 +476,8 @@ pub trait MinimaxPlayer None } - /// More positive for tiger winning, less positive for sheep winning - fn evaluate(&mut self, board: &GameBoard) -> ScoreType; + /// More positive for tiger winning, less positive for sheep winning, with `turn` to move next + fn evaluate(&mut self, board: &GameBoard, turn: Turn) -> ScoreType; /// The most plies `pick_move` searches, counting its own move. Normally the time limit stops it first. fn depth(&self) -> usize { MAX_DEPTH } @@ -616,7 +621,7 @@ mod tests impl MinimaxPlayer for TestPlayer { - fn evaluate(&mut self, board: &GameBoard) -> ScoreType + fn evaluate(&mut self, board: &GameBoard, _turn: Turn) -> ScoreType { let dead = 20 - board.sheep_count() - board.sheep_left(); 100.0 * dead as f32 + board.get_moves(Turn::Tiger).len() as f32 @@ -651,7 +656,7 @@ mod tests } if depth == 0 { - return player.evaluate(&board); + return player.evaluate(&board, turn); } let scores = board.get_moves(turn).into_iter() .map(|m| reference(player, board.after_move_unchecked(m), turn.next_turn(), depth - 1)); diff --git a/src/train.rs b/src/train.rs new file mode 100644 index 0000000..1ee40c0 --- /dev/null +++ b/src/train.rs @@ -0,0 +1,353 @@ +use std::{collections::VecDeque, fs, path::{Path, PathBuf}, sync::atomic::{AtomicUsize, Ordering}, time::{Duration, Instant}}; + +use crate::{board::{GameBoard, Move, Turn, SYMMETRIES}, claude_player::ClaudePlayer, network::{Adam, Network, Rng, Sample}, neural_player::NeuralPlayer, player::{MinimaxPlayer, Player}}; + +/// Settings for `run`. +#[derive(Clone, Debug)] +pub struct Config +{ + pub generations: usize, + pub games_per_generation: usize, + pub out_dir: PathBuf, + /// Search depth for every move in self-play and matches. The depth cap, not a time limit, + /// ends each search, so runs are repeatable. + pub depth: usize, + /// A game this long is called a draw, since nothing in the rules ends a game that goes round in circles. + pub ply_cap: usize, + /// Self-play picks moves at random for this many plies, so games start in different places. + pub random_plies: usize, + /// After that, the chance of a random move instead of the searched one. + pub epsilon: f32, + /// TD(λ): 1 trains every position on its game's result, smaller trusts later positions' search scores more. + pub lambda: f32, + /// Train on the samples of this many of the latest generations. + pub buffer_generations: usize, + pub epochs: usize, + pub batch_size: usize, + pub learning_rate: f32, + /// Games against `ClaudePlayer`, and against the previous generation, after each generation (half as each side). + pub match_games: usize, +} + +impl Default for Config +{ + fn default() -> Self + { + Self + { + generations: 50, + games_per_generation: 500, + out_dir: PathBuf::from("weights"), + depth: 3, + ply_cap: 200, + random_plies: 10, + epsilon: 0.05, + lambda: 0.7, + buffer_generations: 5, + epochs: 2, + batch_size: 256, + learning_rate: 1e-3, + match_games: 100, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Outcome +{ + Tiger, + Sheep, + Draw, +} + +impl Outcome +{ + /// +1 for a tiger win, -1 for a sheep win, matching the sign of the scores. + fn value(self) -> f32 + { + match self { Outcome::Tiger => 1.0, Outcome::Sheep => -1.0, Outcome::Draw => 0.0 } + } +} + +/// A position from a game, with the searching player's score for it (clamped to -1..1). +#[derive(Clone, Copy, Debug)] +struct Ply +{ + board: GameBoard, + turn: Turn, + score: f32, + /// A random move was played here instead of the searched one. + random: bool, +} + +/// Plays one game from the start. `choose` gives the move and score for a position; for the +/// first `random_plies` plies, and with chance `epsilon` after, a random move is played instead +/// (the position keeps `choose`'s score, and is marked `random` unless the two moves are the same). +fn play(mut choose: impl FnMut(GameBoard, Turn) -> (Move, f32), ply_cap: usize, random_plies: usize, epsilon: f32, rng: &mut Rng) + -> (Outcome, Vec<Ply>) +{ + let mut board = GameBoard::default(); + let mut turn = Turn::Sheep; + let mut plies = Vec::new(); + while plies.len() < ply_cap + { + if board.check_win(turn) + { + // The side that just moved has won + let winner = match turn { Turn::Sheep => Outcome::Tiger, Turn::Tiger => Outcome::Sheep }; + return (winner, plies); + } + let (searched, score) = choose(board, turn); + let mut m = searched; + if plies.len() < random_plies || rng.next_f32() < epsilon + { + let moves = board.get_moves(turn); + m = moves[rng.below(moves.len())]; + } + plies.push(Ply { board, turn, score: score.clamp(-1.0, 1.0), random: m != searched }); + board.make_move_unchecked(m); + turn = turn.next_turn(); + } + (Outcome::Draw, plies) +} + +/// TD(λ) targets, worked backwards: the last position's target is the result, and each earlier +/// one's is `(1 - λ)` times the next position's score plus `λ` times the next position's target. +/// +/// A position where a random move was played is instead labelled with its own search score, which +/// also cuts the chain there, so earlier positions don't learn from the random move's result. +/// Without the cut, the sheep's random moves (which usually leave a sheep to be jumped, where a +/// tiger's random move costs little) taught the net that tigers nearly always win. +fn td_targets(plies: &[Ply], outcome: Outcome, lambda: f32) -> Vec<f32> +{ + let mut targets = vec![0.0; plies.len()]; + let mut next = outcome.value(); + for t in (0..plies.len()).rev() + { + targets[t] = if plies[t].random { plies[t].score } else { next }; + next = (1.0 - lambda) * plies[t].score + lambda * targets[t]; + } + targets +} + +/// Runs `f(0..n)` across all cores, returning the results in order. +fn parallel<T: Send>(n: usize, f: impl Fn(usize) -> T + Sync) -> Vec<T> +{ + let threads = std::thread::available_parallelism().map_or(1, |n| n.get()).min(n.max(1)); + let next = AtomicUsize::new(0); + let mut results: Vec<(usize, T)> = std::thread::scope(|s| + { + let handles: Vec<_> = (0..threads).map(|_| s.spawn(|| + { + let mut done = Vec::new(); + loop + { + let i = next.fetch_add(1, Ordering::Relaxed); + if i >= n + { + return done; + } + done.push((i, f(i))); + } + })).collect(); + handles.into_iter().flat_map(|h| h.join().expect("worker thread panicked")).collect() + }); + results.sort_by_key(|(i, _)| *i); + results.into_iter().map(|(_, r)| r).collect() +} + +/// A player that searches exactly `depth` plies. +fn neural(net: &Network, depth: usize) -> NeuralPlayer +{ + let mut player = NeuralPlayer::new(net.clone(), Duration::from_secs(1 << 20), 16); + player.max_depth = depth; + player +} + +fn claude(depth: usize) -> ClaudePlayer +{ + let mut player = ClaudePlayer::new(Duration::from_secs(1 << 20)); + player.max_depth = depth; + player +} + +/// The move a minimax player picks, and the score it gave it. +fn search(player: &mut (impl MinimaxPlayer + Player), board: GameBoard, turn: Turn) -> (Move, f32) +{ + let m = player.pick_move(board, turn).expect("the game is not over, so there is a move"); + (m, player.last_search().map_or(0.0, |r| r.score)) +} + +/// One self-play game of `net` against itself, as training samples. +fn self_play(net: &Network, cfg: &Config, seed: u64) -> (Outcome, Vec<Sample>) +{ + let mut player = neural(net, cfg.depth); + let mut rng = Rng::new(seed); + let (outcome, plies) = play(|b, t| search(&mut player, b, t), cfg.ply_cap, cfg.random_plies, cfg.epsilon, &mut rng); + let targets = td_targets(&plies, outcome, cfg.lambda); + let samples = plies.iter().zip(targets).map(|(p, target)| Sample { board: p.board, turn: p.turn, target }).collect(); + (outcome, samples) +} + +/// The share of points `net` scores over `games` against `opponent()` (a win is 1, a draw ½), +/// playing sheep in even games and tigers in odd ones. Each pair of games shares a seed, so both +/// sides get the same few random opening moves. +fn match_score<P: MinimaxPlayer + Player>(net: &Network, opponent: impl Fn() -> P + Sync, cfg: &Config, seed: u64) -> f32 +{ + let points: f32 = parallel(cfg.match_games, |i| + { + let mut me = neural(net, cfg.depth); + let mut them = opponent(); + let neural_side = if i % 2 == 0 { Turn::Sheep } else { Turn::Tiger }; + let mut rng = Rng::new(seed ^ (i / 2) as u64); + let (outcome, _) = play(|b, t| if t == neural_side { search(&mut me, b, t) } else { search(&mut them, b, t) }, + cfg.ply_cap, 4, 0.0, &mut rng); + match (outcome, neural_side) + { + (Outcome::Draw, _) => 0.5, + (Outcome::Tiger, Turn::Tiger) | (Outcome::Sheep, Turn::Sheep) => 1.0, + _ => 0.0, + } + }).into_iter().sum(); + points / cfg.match_games.max(1) as f32 +} + +/// The newest `gen-NNN.bin` in `dir`, as (generation, path). +fn latest_generation(dir: &Path) -> Option<(usize, PathBuf)> +{ + fs::read_dir(dir).ok()? + .filter_map(|e| e.ok()) + .filter_map(|e| { + let name = e.file_name().into_string().ok()?; + let n = name.strip_prefix("gen-")?.strip_suffix(".bin")?.parse().ok()?; + Some((n, e.path())) + }) + .max_by_key(|(n, _)| *n) +} + +/// Trains a network by self-play, saving each generation to `out_dir/gen-NNN.bin` and `out_dir/latest.bin`. +/// Carries on from the newest generation already in `out_dir`, if there is one. +pub fn run(cfg: &Config) -> std::io::Result<()> +{ + fs::create_dir_all(&cfg.out_dir)?; + let (first, mut net) = match latest_generation(&cfg.out_dir) + { + Some((n, path)) => + { + println!("Carrying on from {}", path.display()); + (n + 1, Network::load(&path)?) + } + None => (1, Network::random(0x5EED)), + }; + let mut adam = Adam::new(cfg.learning_rate); + let mut buffer: VecDeque<Vec<Sample>> = VecDeque::new(); + let mut rng = Rng::new(first as u64); + println!("{:>4} {:>7} {:>6} {:>6} {:>6} {:>8} {:>8} {:>7} {:>7} {:>7}", + "gen", "samples", "tigers", "sheep", "draws", "loss", "after", "vs prev", "vs Claude", "time s"); + for generation in first..first + cfg.generations + { + let start = Instant::now(); + let seed_base = (generation as u64) << 32; + let games = parallel(cfg.games_per_generation, |i| self_play(&net, cfg, seed_base | i as u64)); + let count = |o: Outcome| games.iter().filter(|(g, _)| *g == o).count(); + let (tigers, sheep, draws) = (count(Outcome::Tiger), count(Outcome::Sheep), count(Outcome::Draw)); + buffer.push_back(games.into_iter().flat_map(|(_, s)| s).collect()); + while buffer.len() > cfg.buffer_generations + { + buffer.pop_front(); + } + + let mut samples: Vec<Sample> = buffer.iter().flatten().copied().collect(); + let loss_before = net.loss(&samples); + let previous = net.clone(); + for _ in 0..cfg.epochs + { + // Shuffle, and show each position in a random one of its 8 orientations + for i in (1..samples.len()).rev() + { + samples.swap(i, rng.below(i + 1)); + } + for s in &mut samples + { + s.board = s.board.transformed(rng.below(SYMMETRIES)); + } + for batch in samples.chunks(cfg.batch_size) + { + net.train_batch(batch, &mut adam); + } + } + let loss_after = net.loss(&samples); + + let vs_prev = match_score(&net, || neural(&previous, cfg.depth), cfg, seed_base ^ 0xA); + let vs_claude = match_score(&net, || claude(cfg.depth), cfg, seed_base ^ 0xB); + net.save(&cfg.out_dir.join(format!("gen-{generation:03}.bin")))?; + net.save(&cfg.out_dir.join("latest.bin"))?; + println!("{generation:>4} {:>7} {tigers:>6} {sheep:>6} {draws:>6} {loss_before:>8.4} {loss_after:>8.4} {:>6.0}% {:>8.0}% {:>7.1}", + samples.len(), vs_prev * 100.0, vs_claude * 100.0, start.elapsed().as_secs_f64()); + } + Ok(()) +} + +#[cfg(test)] +mod tests +{ + use super::*; + + fn ply(score: f32) -> Ply + { + Ply { board: GameBoard::default(), turn: Turn::Sheep, score, random: false } + } + + #[test] + fn td_targets_blend_scores_and_the_result() + { + let plies = [ply(0.0), ply(0.2), ply(-0.4)]; + let t = td_targets(&plies, Outcome::Tiger, 0.5); + assert_eq!(t[2], 1.0); + assert_eq!(t[1], 0.5 * -0.4 + 0.5 * 1.0); + assert_eq!(t[0], 0.5 * 0.2 + 0.5 * t[1]); + // λ = 1 is just the result + assert_eq!(td_targets(&plies, Outcome::Sheep, 1.0), vec![-1.0; 3]); + } + + #[test] + fn td_targets_stop_at_random_moves() + { + let random = Ply { random: true, ..ply(0.3) }; + let plies = [ply(0.0), random, ply(-0.4)]; + let t = td_targets(&plies, Outcome::Tiger, 1.0); + // After the random move the result counts; at it and before it, only its search score does + assert_eq!(t, vec![0.3, 0.3, 1.0]); + let t = td_targets(&plies, Outcome::Tiger, 0.5); + assert_eq!(t[1], 0.3); + assert_eq!(t[0], 0.3); + } + + #[test] + fn self_play_games_are_legal_and_repeatable() + { + let cfg = Config { depth: 1, ply_cap: 60, ..Config::default() }; + let net = Network::random(3); + let (outcome, samples) = self_play(&net, &cfg, 42); + assert_eq!(self_play(&net, &cfg, 42), (outcome, samples.clone())); + assert!(!samples.is_empty() && samples.len() <= cfg.ply_cap); + assert_eq!(samples[0].board, GameBoard::default()); + // Each position follows from the one before by a legal move + for w in samples.windows(2) + { + assert_eq!(w[1].turn, w[0].turn.next_turn()); + assert!(w[0].board.get_moves(w[0].turn).iter().any(|&m| w[0].board.after_move_unchecked(m) == w[1].board)); + } + assert!(samples.iter().all(|s| (-1.0..=1.0).contains(&s.target))); + if outcome == Outcome::Draw + { + assert_eq!(samples.len(), cfg.ply_cap); + } + } + + #[test] + fn parallel_keeps_the_order() + { + assert_eq!(parallel(100, |i| i * 2), (0..100).map(|i| i * 2).collect::<Vec<_>>()); + } +} |