use std::{fs, io, path::Path};
use rayon::prelude::*;
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;
/// Samples per parallel task in `gradient`. The chunks are a fixed size and summed in order, so
/// the result is the same whatever the number of threads.
const GRAD_CHUNK: usize = 32;
/// Samples per parallel task in `loss`, where each one is only a forward pass.
const LOSS_CHUNK: usize = 1024;
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
}
/// The gradient of the error summed over `batch`, and that summed error, worked out in
/// `GRAD_CHUNK`-sized pieces across all cores.
fn gradient(&self, batch: &[Sample]) -> (Vec<f32>, f32)
{
let parts: Vec<(Vec<f32>, f32)> = batch.par_chunks(GRAD_CHUNK).map(|chunk|
{
let mut grad = vec![0.0; PARAMS];
let loss = chunk.iter().map(|s| self.add_gradient(s, &mut grad)).sum();
(grad, loss)
}).collect();
let mut grad = vec![0.0; PARAMS];
let mut loss = 0.0;
for (g, l) in parts
{
grad.iter_mut().zip(&g).for_each(|(a, b)| *a += b);
loss += l;
}
(grad, loss)
}
/// 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, loss) = self.gradient(batch);
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 parts: Vec<f32> = samples.par_chunks(LOSS_CHUNK)
.map(|chunk| chunk.iter().map(|s| { let e = self.forward(&s.board, s.turn) - s.target; e * e / 2.0 }).sum())
.collect();
// A fold from 0, since an empty f32 `sum` is -0
let total = parts.into_iter().fold(0.0, |a, b| a + b);
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 parallel_gradient_matches_serial()
{
let net = Network::random(5);
// Not a multiple of GRAD_CHUNK, so the last chunk is short
let batch = samples(GRAD_CHUNK * 4 + 7, 9);
let (grad, loss) = net.gradient(&batch);
// The same chunks, one after another, give exactly the same sums
let mut serial = vec![0.0; PARAMS];
let mut serial_loss = 0.0;
for chunk in batch.chunks(GRAD_CHUNK)
{
let mut g = vec![0.0; PARAMS];
serial_loss += chunk.iter().map(|s| net.add_gradient(s, &mut g)).sum::<f32>();
serial.iter_mut().zip(&g).for_each(|(a, b)| *a += b);
}
assert_eq!(grad, serial);
assert_eq!(loss, serial_loss);
// And one running sum gives nearly the same
let mut plain = vec![0.0; PARAMS];
let plain_loss: f32 = batch.iter().map(|s| net.add_gradient(s, &mut plain)).sum();
assert!((loss - plain_loss).abs() < 1e-4 * plain_loss.abs().max(1.0));
for (a, b) in grad.iter().zip(&plain)
{
assert!((a - b).abs() < 1e-4 * b.abs().max(1.0), "{a} vs {b}");
}
}
#[test]
fn training_is_the_same_on_any_number_of_threads()
{
let train = |threads: usize|
{
let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap();
pool.install(||
{
let mut net = Network::random(2);
let data = samples(300, 4);
let mut adam = Adam::new(1e-3);
for batch in data.chunks(100)
{
net.train_batch(batch, &mut adam);
}
(net, Network::random(3).loss(&samples(3000, 6)))
})
};
assert_eq!(train(1), train(4));
}
#[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);
}
}
}
