From 2b7b36fc8dea8dbcbe5dd53aff9aa71a352971fc Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Thu, 26 Mar 2026 02:04:45 -0400 Subject: [PATCH 01/59] init the lib --- Cargo.lock | 4 ++++ Cargo.toml | 4 ++-- training/Cargo.toml | 6 ++++++ training/src/lib.rs | 14 ++++++++++++++ 4 files changed, 26 insertions(+), 2 deletions(-) create mode 100644 training/Cargo.toml create mode 100644 training/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index b523419..c473c0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -336,6 +336,10 @@ dependencies = [ "alpha_paint", ] +[[package]] +name = "training" +version = "0.1.0" + [[package]] name = "unicode-ident" version = "1.0.24" diff --git a/Cargo.toml b/Cargo.toml index 72a8140..82d1b8d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["alpha_paint", "tooling"] +members = ["alpha_paint", "tooling", "training"] resolver = "3" [profile.dev] @@ -7,4 +7,4 @@ opt-level = 3 debug_assertions = true [profile.release] -debug = true \ No newline at end of file +debug = true diff --git a/training/Cargo.toml b/training/Cargo.toml new file mode 100644 index 0000000..1dd31a0 --- /dev/null +++ b/training/Cargo.toml @@ -0,0 +1,6 @@ +[package] +name = "training" +version = "0.1.0" +edition = "2024" + +[dependencies] diff --git a/training/src/lib.rs b/training/src/lib.rs new file mode 100644 index 0000000..b93cf3f --- /dev/null +++ b/training/src/lib.rs @@ -0,0 +1,14 @@ +pub fn add(left: u64, right: u64) -> u64 { + left + right +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn it_works() { + let result = add(2, 2); + assert_eq!(result, 4); + } +} From a61cf25f91737d35237d6eb897c1a05f1fa44bc4 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Thu, 26 Mar 2026 02:12:09 -0400 Subject: [PATCH 02/59] copy over siebren --- Cargo.lock | 525 ++++++- training/Cargo.toml | 24 +- training/src/cudagraph.rs | 489 ++++++ training/src/environments/bytefight/game.rs | 1411 ++++++++++++++++++ training/src/environments/bytefight/map.rs | 176 +++ training/src/environments/bytefight/mod.rs | 392 +++++ training/src/environments/bytefight/pen.rs | 510 +++++++ training/src/environments/bytefight/snake.rs | 85 ++ training/src/environments/bytefight/types.rs | 237 +++ training/src/environments/connect4.rs | 436 ++++++ training/src/environments/mod.rs | 7 + training/src/environments/tictactoe.rs | 315 ++++ training/src/eval.rs | 253 ++++ training/src/executor.rs | 419 ++++++ training/src/future.rs | 233 +++ training/src/integration_tests.rs | 494 ++++++ training/src/lib.rs | 724 ++++++++- training/src/mcts.rs | 369 +++++ training/src/observation_replay_buffer.rs | 315 ++++ training/src/queue.rs | 509 +++++++ training/src/replay_buffer.rs | 716 +++++++++ training/src/training.rs | 390 +++++ training/src/worker.rs | 323 ++++ 23 files changed, 9331 insertions(+), 21 deletions(-) create mode 100644 training/src/cudagraph.rs create mode 100644 training/src/environments/bytefight/game.rs create mode 100644 training/src/environments/bytefight/map.rs create mode 100644 training/src/environments/bytefight/mod.rs create mode 100644 training/src/environments/bytefight/pen.rs create mode 100644 training/src/environments/bytefight/snake.rs create mode 100644 training/src/environments/bytefight/types.rs create mode 100644 training/src/environments/connect4.rs create mode 100644 training/src/environments/mod.rs create mode 100644 training/src/environments/tictactoe.rs create mode 100644 training/src/eval.rs create mode 100644 training/src/executor.rs create mode 100644 training/src/future.rs create mode 100644 training/src/integration_tests.rs create mode 100644 training/src/mcts.rs create mode 100644 training/src/observation_replay_buffer.rs create mode 100644 training/src/queue.rs create mode 100644 training/src/replay_buffer.rs create mode 100644 training/src/training.rs create mode 100644 training/src/worker.rs diff --git a/Cargo.lock b/Cargo.lock index c473c0b..bb1ca57 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,12 +2,39 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "alpha_paint" version = "0.1.0" dependencies = [ "pyo3", - "rand", + "rand 0.10.0", +] + +[[package]] +name = "alphapaint_training" +version = "0.1.0" +dependencies = [ + "cudarc", + "event-listener", + "ndarray", + "numpy", + "pyo3", + "rand 0.10.0", + "rand_chacha 0.10.0", + "rand_distr", + "rayon", + "rstest", + "serde", + "serde_json", ] [[package]] @@ -16,6 +43,12 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "bitflags" version = "2.11.0" @@ -36,7 +69,16 @@ checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if", "cpufeatures", - "rand_core", + "rand_core 0.10.0", +] + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", ] [[package]] @@ -48,18 +90,123 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "cudarc" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3aa12038120eb13347a6ae2ffab1d34efe78150125108627fd85044dd4d6ff1e" +dependencies = [ + "libloading", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + [[package]] name = "foldhash" version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.2" @@ -68,12 +215,18 @@ checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", "libc", - "r-efi", - "rand_core", + "r-efi 6.0.0", + "rand_core 0.10.0", "wasip2", "wasip3", ] +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + [[package]] name = "hashbrown" version = "0.15.5" @@ -131,30 +284,145 @@ version = "0.2.182" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "log" version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "numpy" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778da78c64ddc928ebf5ad9df5edf0789410ff3bdbf3619aed51cd789a6af1e2" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + [[package]] name = "once_cell" version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + [[package]] name = "portable-atomic" version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "portable-atomic-util" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -165,6 +433,15 @@ dependencies = [ "syn", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -241,12 +518,28 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.0" @@ -254,8 +547,37 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" dependencies = [ "chacha20", - "getrandom", - "rand_core", + "getrandom 0.4.2", + "rand_core 0.10.0", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" +dependencies = [ + "ppv-lite86", + "rand_core 0.10.0", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -264,6 +586,121 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn", + "unicode-ident", +] + +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "semver" version = "1.0.27" @@ -277,6 +714,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -312,6 +750,12 @@ dependencies = [ "zmij", ] +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "syn" version = "2.0.117" @@ -330,15 +774,41 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" [[package]] -name = "tooling" -version = "0.1.0" +name = "toml_datetime" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97251a7c317e03ad83774a8752a7e81fb6067740609f75ea2b585b569a59198f" dependencies = [ - "alpha_paint", + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.8+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16bff38f1d86c47f9ff0647e6838d7bb362522bdf44006c7068c2b1e606f1f3c" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2334f11ee363607eb04df9b8fc8a13ca1715a72ba8662a26ac285c98aabb4011" +dependencies = [ + "winnow", ] [[package]] -name = "training" +name = "tooling" version = "0.1.0" +dependencies = [ + "alpha_paint", +] [[package]] name = "unicode-ident" @@ -404,6 +874,21 @@ dependencies = [ "semver", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "winnow" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a90e88e4667264a994d34e6d1ab2d26d398dcdca8b7f52bec8668957517fc7d8" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" @@ -492,6 +977,26 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "zerocopy" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efbb2a062be311f2ba113ce66f697a4dc589f85e78a4aea276200804cea0ed87" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e8bc7269b54418e7aeeef514aa68f8690b8c0489a06b0136e5f57c4c5ccab89" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/training/Cargo.toml b/training/Cargo.toml index 1dd31a0..f7bccab 100644 --- a/training/Cargo.toml +++ b/training/Cargo.toml @@ -1,6 +1,26 @@ [package] -name = "training" +name = "alphapaint_training" version = "0.1.0" -edition = "2024" +edition = "2021" +build = "build.rs" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "alphapaint_training" +crate-type = ["cdylib"] [dependencies] +ndarray = "0.17.1" +pyo3 = { version = "0.28.2", features = ["extension-module"] } +rand = "0.10.0" +rand_chacha = "0.10.0" +rand_distr = "0.5" +event-listener = "5" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +rayon = "1.11.0" +cudarc = { version = "0.18.2", features = ["cuda-version-from-build-system"] } +numpy = "0.28.0" + +[dev-dependencies] +rstest = "0.26.1" diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs new file mode 100644 index 0000000..f5ede9a --- /dev/null +++ b/training/src/cudagraph.rs @@ -0,0 +1,489 @@ +//! Rust-launched CUDA graph backend for ByteFight self-play inference. + +use std::ffi::{c_void, CStr}; +use std::mem::size_of; +use std::slice; + +use cudarc::runtime::sys as cuda; +use ndarray::{ArrayView, Ix3}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::ffi; +use pyo3::prelude::*; + +use crate::eval::PolicyValue; +use crate::queue::BatchCompletion; + +#[allow(non_camel_case_types)] +type cudaStream_t = cuda::cudaStream_t; + +#[allow(non_camel_case_types)] +type cudaGraphExec_t = cuda::cudaGraphExec_t; + +type CudaError = cuda::cudaError_t; + +const DL_TENSOR_NAME: &[u8] = b"dltensor\0"; +const DL_DEVICE_CPU: i32 = 1; +const DL_DEVICE_CUDA: i32 = 2; +const DL_DTYPE_FLOAT: u8 = 2; +const DL_DTYPE_UINT: u8 = 1; + +const BYTEFIGHT_OBS_SIDE: usize = 18; +const BYTEFIGHT_OBS_WIDTH: usize = 16; +const BYTEFIGHT_OBS_CELLS: usize = BYTEFIGHT_OBS_SIDE * BYTEFIGHT_OBS_WIDTH; +const BYTEFIGHT_ACTIONS: usize = 7; + +#[repr(C)] +struct DLDevice { + device_type: i32, + device_id: i32, +} + +#[repr(C)] +struct DLDataType { + code: u8, + bits: u8, + lanes: u16, +} + +#[repr(C)] +struct DLTensor { + data: *mut c_void, + device: DLDevice, + ndim: i32, + dtype: DLDataType, + shape: *mut i64, + strides: *mut i64, + byte_offset: usize, +} + +#[repr(C)] +struct DLManagedTensor { + dl_tensor: DLTensor, + manager_ctx: *mut c_void, + deleter: Option, +} + +struct DLPackContext { + shape: Box<[i64]>, +} + +unsafe extern "C" fn dlpack_capsule_destructor(capsule: *mut ffi::PyObject) { + if capsule.is_null() { + return; + } + + let name = ffi::PyCapsule_GetName(capsule); + if name.is_null() { + return; + } + + let c_name = CStr::from_ptr(name); + if c_name.to_bytes() != b"dltensor" { + return; + } + + let ptr = ffi::PyCapsule_GetPointer(capsule, DL_TENSOR_NAME.as_ptr() as *const i8); + if ptr.is_null() { + return; + } + + let managed = ptr as *mut DLManagedTensor; + if let Some(deleter) = unsafe { (*managed).deleter } { + deleter(managed); + } +} + +extern "C" fn dlpack_deleter(ptr: *mut DLManagedTensor) { + if ptr.is_null() { + return; + } + + unsafe { + let ctx_ptr = (*ptr).manager_ctx as *mut DLPackContext; + if !ctx_ptr.is_null() { + drop(Box::from_raw(ctx_ptr)); + } + drop(Box::from_raw(ptr)); + } +} + +fn dlpack_capsule( + py: Python<'_>, + data: *mut c_void, + shape: &[i64], + device_type: i32, + device_id: i32, + dtype_code: u8, + dtype_bits: u8, +) -> PyResult> { + let ctx = Box::new(DLPackContext { + shape: shape.to_vec().into_boxed_slice(), + }); + let shape_ptr = ctx.shape.as_ptr() as *mut i64; + let ctx_ptr = Box::into_raw(ctx); + + let managed = Box::new(DLManagedTensor { + dl_tensor: DLTensor { + data, + device: DLDevice { + device_type, + device_id, + }, + ndim: shape.len() as i32, + dtype: DLDataType { + code: dtype_code, + bits: dtype_bits, + lanes: 1, + }, + shape: shape_ptr, + strides: std::ptr::null_mut(), + byte_offset: 0, + }, + manager_ctx: ctx_ptr as *mut c_void, + deleter: Some(dlpack_deleter), + }); + + let managed_ptr = Box::into_raw(managed); + let capsule = unsafe { + ffi::PyCapsule_New( + managed_ptr as *mut c_void, + DL_TENSOR_NAME.as_ptr() as *const i8, + Some(dlpack_capsule_destructor), + ) + }; + + if capsule.is_null() { + dlpack_deleter(managed_ptr); + return Err(PyErr::new::( + "failed to create DLPack capsule", + )); + } + + Ok(unsafe { Py::from_owned_ptr(py, capsule) }) +} + +fn check_cuda(code: CudaError, context: &str) -> PyResult<()> { + if code == cuda::cudaError::cudaSuccess { + Ok(()) + } else { + Err(PyErr::new::(format!( + "{} failed with CUDA error {:?}", + context, code + ))) + } +} + +fn check_cuda_or_panic(code: CudaError, context: &str) { + if code != cuda::cudaError::cudaSuccess { + panic!("{} failed with CUDA error {:?}", context, code); + } +} + +fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("host allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMallocHost(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr.cast::()) +} + +fn cuda_malloc_host_u8(count: usize, context: &str) -> PyResult<*mut u8> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count; + unsafe { + check_cuda( + cuda::cudaMallocHost(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr.cast::()) +} + +fn cuda_malloc_device_f32(count: usize, context: &str) -> PyResult<*mut c_void> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("device allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMalloc(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr) +} + +fn cuda_malloc_device_u8(count: usize, context: &str) -> PyResult<*mut c_void> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count; + unsafe { + check_cuda( + cuda::cudaMalloc(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr) +} + +struct ByteFightCudaGraphLane { + stream: cudaStream_t, + graph_exec: cudaGraphExec_t, + /// Owns Python-side graph/tensor objects for this lane. + _py_owner: Py, + obs_host: *mut u8, + obs_dev: *mut c_void, + policy_host: *mut f32, + policy_dev: *mut c_void, + value_host: *mut f32, + value_dev: *mut c_void, +} + +struct LaneCompletionContext { + policy_host: *const f32, + value_host: *const f32, + batch_size: usize, + completion: Option>>, +} + +unsafe impl Send for LaneCompletionContext {} + +unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + + let mut ctx = unsafe { Box::from_raw(user_data.cast::()) }; + let policy_src = + unsafe { slice::from_raw_parts(ctx.policy_host, ctx.batch_size * BYTEFIGHT_ACTIONS) }; + let value_src = unsafe { slice::from_raw_parts(ctx.value_host, ctx.batch_size) }; + + let mut outputs = vec![PolicyValue::<7>::default(); ctx.batch_size]; + for (i, out) in outputs.iter_mut().enumerate() { + let start = i * BYTEFIGHT_ACTIONS; + out.policy + .copy_from_slice(&policy_src[start..start + BYTEFIGHT_ACTIONS]); + out.value = value_src[i]; + } + + if let Some(completion) = ctx.completion.take() { + completion.complete(&outputs); + } +} + +impl Drop for ByteFightCudaGraphLane { + fn drop(&mut self) { + unsafe { + let _ = cuda::cudaFree(self.obs_dev); + let _ = cuda::cudaFree(self.policy_dev); + let _ = cuda::cudaFree(self.value_dev); + let _ = cuda::cudaFreeHost(self.obs_host.cast::()); + let _ = cuda::cudaFreeHost(self.policy_host.cast::()); + let _ = cuda::cudaFreeHost(self.value_host.cast::()); + let _ = cuda::cudaStreamDestroy(self.stream); + } + } +} + +/// Per-lane CUDA graph executor for ByteFight self-play inference. +pub struct ByteFightCudaGraphRunner { + batch_size: usize, + lanes: Vec, +} + +// SAFETY: Lane buffers/streams are independent per batch_idx and queue dispatch +// ensures a lane is not reused before dispatch returns for that lane. +unsafe impl Send for ByteFightCudaGraphRunner {} +unsafe impl Sync for ByteFightCudaGraphRunner {} + +impl ByteFightCudaGraphRunner { + pub fn new( + py: Python<'_>, + model: Py, + num_lanes: usize, + batch_size: usize, + precision: &str, + ) -> PyResult { + if num_lanes == 0 { + return Err(PyErr::new::("num_lanes must be > 0")); + } + if batch_size == 0 { + return Err(PyErr::new::("batch_size must be > 0")); + } + + let module = PyModule::import(py, "siebren.cudagraph_backend")?; + let capture_fn = module.getattr("capture_bytefight_lane_graph")?; + + let obs_count = batch_size * BYTEFIGHT_OBS_CELLS; + let policy_count = batch_size * BYTEFIGHT_ACTIONS; + let value_count = batch_size; + + let obs_shape = [ + batch_size as i64, + BYTEFIGHT_OBS_SIDE as i64, + BYTEFIGHT_OBS_WIDTH as i64, + ]; + let policy_shape = [batch_size as i64, BYTEFIGHT_ACTIONS as i64]; + let value_shape = [batch_size as i64]; + + let mut lanes = Vec::with_capacity(num_lanes); + + for lane_idx in 0..num_lanes { + let mut stream: cudaStream_t = std::ptr::null_mut(); + unsafe { + check_cuda( + cuda::cudaStreamCreate(&mut stream as *mut cudaStream_t), + "cudaStreamCreate", + )?; + } + + let obs_host = + cuda_malloc_host_u8(obs_count, &format!("cudaMallocHost obs lane {}", lane_idx))?; + let policy_host = cuda_malloc_host_f32( + policy_count, + &format!("cudaMallocHost policy lane {}", lane_idx), + )?; + let value_host = cuda_malloc_host_f32( + value_count, + &format!("cudaMallocHost value lane {}", lane_idx), + )?; + + let obs_dev = + cuda_malloc_device_u8(obs_count, &format!("cudaMalloc obs lane {}", lane_idx))?; + let policy_dev = cuda_malloc_device_f32( + policy_count, + &format!("cudaMalloc policy lane {}", lane_idx), + )?; + let value_dev = cuda_malloc_device_f32( + value_count, + &format!("cudaMalloc value lane {}", lane_idx), + )?; + + let obs_host_capsule = dlpack_capsule( + py, + obs_host.cast::(), + &obs_shape, + DL_DEVICE_CPU, + 0, + DL_DTYPE_UINT, + 8, + )?; + let obs_dev_capsule = + dlpack_capsule(py, obs_dev, &obs_shape, DL_DEVICE_CUDA, 0, DL_DTYPE_UINT, 8)?; + let policy_host_capsule = dlpack_capsule( + py, + policy_host.cast::(), + &policy_shape, + DL_DEVICE_CPU, + 0, + DL_DTYPE_FLOAT, + 32, + )?; + let policy_dev_capsule = dlpack_capsule( + py, + policy_dev, + &policy_shape, + DL_DEVICE_CUDA, + 0, + DL_DTYPE_FLOAT, + 32, + )?; + let value_host_capsule = dlpack_capsule( + py, + value_host.cast::(), + &value_shape, + DL_DEVICE_CPU, + 0, + DL_DTYPE_FLOAT, + 32, + )?; + let value_dev_capsule = dlpack_capsule( + py, + value_dev, + &value_shape, + DL_DEVICE_CUDA, + 0, + DL_DTYPE_FLOAT, + 32, + )?; + + let (exec_handle, py_owner): (u64, Py) = capture_fn + .call1(( + model.clone_ref(py), + obs_host_capsule, + obs_dev_capsule, + policy_host_capsule, + policy_dev_capsule, + value_host_capsule, + value_dev_capsule, + stream as u64, + precision, + ))? + .extract()?; + + let lane = ByteFightCudaGraphLane { + stream, + graph_exec: exec_handle as cudaGraphExec_t, + _py_owner: py_owner, + obs_host, + obs_dev, + policy_host, + policy_dev, + value_host, + value_dev, + }; + lanes.push(lane); + } + + Ok(Self { batch_size, lanes }) + } + + pub fn dispatch_async( + &self, + batch_idx: usize, + obs_view: ArrayView, + completion: BatchCompletion>, + ) { + debug_assert_eq!( + obs_view.shape(), + &[self.batch_size, BYTEFIGHT_OBS_SIDE, BYTEFIGHT_OBS_WIDTH] + ); + + let lane = &self.lanes[batch_idx % self.lanes.len()]; + + let obs_src = obs_view + .as_slice() + .expect("bytefight queue observation batch must be contiguous"); + let obs_dst = unsafe { + slice::from_raw_parts_mut(lane.obs_host, self.batch_size * BYTEFIGHT_OBS_CELLS) + }; + obs_dst.copy_from_slice(obs_src); + + unsafe { + check_cuda_or_panic( + cuda::cudaGraphLaunch(lane.graph_exec, lane.stream), + "cudaGraphLaunch", + ); + + let ctx = Box::new(LaneCompletionContext { + policy_host: lane.policy_host, + value_host: lane.value_host, + batch_size: self.batch_size, + completion: Some(completion), + }); + check_cuda_or_panic( + cuda::cudaLaunchHostFunc( + lane.stream, + Some(lane_completion_callback), + Box::into_raw(ctx).cast::(), + ), + "cudaLaunchHostFunc", + ); + } + } +} diff --git a/training/src/environments/bytefight/game.rs b/training/src/environments/bytefight/game.rs new file mode 100644 index 0000000..a440424 --- /dev/null +++ b/training/src/environments/bytefight/game.rs @@ -0,0 +1,1411 @@ +use std::collections::{HashMap, HashSet, VecDeque}; +use std::sync::OnceLock; + +use rand::seq::SliceRandom; +use rand::Rng; + +use super::map::{add_padding_walls, Map}; +use super::snake::Snake; +use super::types::{ + BitpackedObservation, ByteFightAction, Point, TerminalState, ValidMoves, OBS_CELLS, OBS_SIDE, +}; + +pub const APPLE_REWARD: usize = 2; +const TRAP_LIFETIME: i16 = 100; +pub const TRAP_SACRIFICE: usize = 3; +const DECAY_TIMELINE: [(usize, usize); 4] = [(1000, 15), (1600, 10), (1800, 5), (1950, 2)]; +const DECAY_NOT_APPLIED_PLACEHOLDER: usize = 9999; + +pub const LAST_TURN: usize = 2000; + +// 16x16-only map set for training/selection. +// These are the only currently-defined maps that fit within a 16x16 observation window. +// Excluded (commented out from MAPS_JSON): pillars, great_divide, empty_large, ssspline, +// combustible_lemons, arena, ladder, compasss, diamonds, ssspiral, lol, attrition. +const MAPS_JSON: &str = r#"{ + "cage": "11,11#1,5#9,5#5#2##30,1,Vertical#1010101010101010101010101010101010101010101000101010100000101010000010101010001010101010101010101010101010101010101010101#0", + "empty": "9,9#1,4#7,4#5#2##20,1,Vertical#000000000000000000000000000000000000000000000000000000000000000000000000000000000#0", + "recurve": "13,13#2,6#10,6#4#2#6,0,6,12_6,12,6,0#50,1,Vertical#0100000000010001000000010000010000010000000100010000000010001000000000000000000000000000000000000000000000001000100000000100010000000100000100000100000001000100000000010#0" + }"#; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Board { + pub map: Map, + pub apple_timeline: Vec<(usize, Point)>, + pub apple_timeline_ptr: usize, + pub snake_a: Snake, + pub snake_b: Snake, + pub is_player_a: bool, + pub min_player_size: usize, + pub(crate) decay_countdown: usize, + pub(crate) cached_decay_interval: usize, + pub(crate) is_decaying: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RollbackState { + EndTurn { + old_num_traps: usize, + old_sacrifice_val: usize, + prev_cached_decay_interval: usize, + prev_decay_countdown: usize, + decayed_point: Option>, + snake_a_ate_during_collision: bool, + snake_b_ate_during_collision: bool, + snake_a_max_len: usize, + snake_b_max_len: usize, + prev_apple_timeline_ptr: usize, + apples_placed_index: Vec, + }, + ApplyMove { + prev_trap_val: i16, + sacrificed_points: Vec, + prev_queued_length: usize, + prev_max_length_reached: usize, + prev_direction: Option, + head_was_apple: bool, + }, + ApplyTrap { + old_trap_val: i16, + trap: Point, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Symmetry { + Horizontal, + Vertical, + Origin, +} + +impl Symmetry { + fn reflect(self, point: Point, width: usize, height: usize) -> Point { + match self { + Symmetry::Horizontal => Point { + x: point.x, + y: height - 1 - point.y, + }, + Symmetry::Vertical => Point { + x: width - 1 - point.x, + y: point.y, + }, + Symmetry::Origin => Point { + x: width - 1 - point.x, + y: height - 1 - point.y, + }, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum AppleSpec { + Timeline(Vec<(usize, Point)>), + Spawn { + rate: usize, + count: usize, + symmetry: Symmetry, + }, +} + +#[derive(Debug, Clone)] +struct MapDefinition { + width: usize, + height: usize, + start_a: Point, + start_b: Point, + start_size: usize, + min_player_size: usize, + portals: Vec<(Point, Point)>, + walls: Vec, + apple_spec: AppleSpec, +} + +impl MapDefinition { + fn from_map_string(map_str: &str) -> Result { + let parts: Vec<&str> = map_str.split('#').collect(); + if parts.len() != 9 { + return Err(format!("expected 9 map parts, got {}", parts.len())); + } + + let (width, height) = parse_pair(parts[0], ',')?; + let start_a = parse_point(parts[1])?; + let start_b = parse_point(parts[2])?; + let start_size = parse_usize(parts[3], "start_size")?; + let min_player_size = parse_usize(parts[4], "min_player_size")?; + let portals = parse_portals_section(parts[5])?; + let walls = parse_walls_bits(parts[7], width, height)?; + + let is_record = parse_usize(parts[8], "is_record")? == 1; + let apple_spec = if is_record { + let timeline = parse_apple_timeline(parts[6])?; + AppleSpec::Timeline(timeline) + } else { + let stats: Vec<&str> = parts[6].split(',').collect(); + if stats.len() != 3 { + return Err("invalid apple stats section".to_string()); + } + let rate = parse_usize(stats[0], "apple_rate")?; + let count = parse_usize(stats[1], "num_apples")?; + let symmetry = match stats[2] { + "Horizontal" => Symmetry::Horizontal, + "Vertical" => Symmetry::Vertical, + "Origin" => Symmetry::Origin, + _ => return Err("unknown symmetry".to_string()), + }; + AppleSpec::Spawn { + rate, + count, + symmetry, + } + }; + + Ok(MapDefinition { + width, + height, + start_a, + start_b, + start_size, + min_player_size, + portals, + walls, + apple_spec, + }) + } + + fn build_initial_state(&self, rng: &mut impl Rng) -> Board { + let apple_timeline = self.apple_spec.generate_timeline( + rng, + self.width, + self.height, + &self.portals, + &self.walls, + self.start_a, + self.start_b, + ); + + let apples_now: Vec = apple_timeline + .iter() + .filter(|(turn, _)| *turn == 0) + .map(|(_, point)| *point) + .collect(); + + let snake_a = Snake { + sacrifice: 1, + max_length_reached: self.start_size, + queued_length: self.start_size.saturating_sub(1), + traps_this_turn: 0, + current_direction: ByteFightAction::new((rng.next_u64() % 8) as u8), + segment_queue: VecDeque::from([self.start_a]), + total_apples: 0, + }; + let snake_b = Snake { + sacrifice: 1, + max_length_reached: self.start_size, + queued_length: self.start_size.saturating_sub(1), + traps_this_turn: 0, + current_direction: ByteFightAction::new((rng.next_u64() % 8) as u8), + segment_queue: VecDeque::from([self.start_b]), + total_apples: 0, + }; + + let mut map = Map::new((self.width, self.height), 0); + for wall in &self.walls { + map.become_wall(*wall); + } + for (p1, p2) in &self.portals { + map.add_portal(*p1, *p2); + } + for apple in &apples_now { + map.become_apple(*apple); + if let Some(portal) = map.portal(*apple) { + map.become_apple(portal); + } + } + add_padding_walls(&mut map); + + let mut board = Board { + map, + apple_timeline, + apple_timeline_ptr: 0, + snake_a, + snake_b, + is_player_a: true, + min_player_size: self.min_player_size, + decay_countdown: 0, + cached_decay_interval: DECAY_NOT_APPLIED_PLACEHOLDER, + is_decaying: false, + }; + + board.fix_apple_head_collisions(); + let _ = board.apply_decay(); + board + } +} + +impl AppleSpec { + fn generate_timeline( + &self, + rng: &mut impl Rng, + width: usize, + height: usize, + portals: &[(Point, Point)], + walls: &[Point], + start_a: Point, + start_b: Point, + ) -> Vec<(usize, Point)> { + match self { + AppleSpec::Timeline(timeline) => timeline.clone(), + AppleSpec::Spawn { + rate, + count, + symmetry, + } => { + let portal_map = portal_lookup(portals); + let wall_set: HashSet = walls.iter().copied().collect(); + let mut considered: HashSet = HashSet::new(); + considered.insert(start_a); + considered.insert(start_b); + + let mut select_from = Vec::new(); + for y in 0..height { + for x in 0..width { + let point = Point { x, y }; + if wall_set.contains(&point) { + continue; + } + if considered.contains(&point) { + continue; + } + select_from.push(point); + considered.insert(point); + considered.insert(symmetry.reflect(point, width, height)); + } + } + + let mut apples = Vec::new(); + let mut first_round = select_from.clone(); + add_apple_spawns( + &mut first_round, + *count, + *symmetry, + &portal_map, + &mut apples, + 0, + width, + height, + rng, + ); + + let mut later_round = select_from; + later_round.push(start_a); + later_round.push(start_b); + + let mut spawn_round = *rate; + while spawn_round < LAST_TURN { + let mut picks = later_round.clone(); + add_apple_spawns( + &mut picks, + *count, + *symmetry, + &portal_map, + &mut apples, + spawn_round, + width, + height, + rng, + ); + spawn_round += *rate; + } + + apples + } + } + } +} + +fn portal_lookup(portals: &[(Point, Point)]) -> HashMap { + let mut map = HashMap::new(); + for (p1, p2) in portals { + map.insert(*p1, *p2); + map.insert(*p2, *p1); + } + map +} + +fn add_apple_spawns( + picks: &mut Vec, + count: usize, + symmetry: Symmetry, + portals: &HashMap, + apples: &mut Vec<(usize, Point)>, + turn_num: usize, + width: usize, + height: usize, + rng: &mut impl Rng, +) { + picks.shuffle(rng); + let mut apple_count = 0; + let mut idx = 0; + + while apple_count < count && idx < picks.len() { + let point = picks[idx]; + let reflection = symmetry.reflect(point, width, height); + + if point == reflection { + apple_count += 1; + apples.push((turn_num, point)); + } else { + apple_count += 2; + apples.push((turn_num, point)); + apples.push((turn_num, reflection)); + } + + if let Some(portal) = portals.get(&point) { + apples.push((turn_num, *portal)); + if point != reflection { + if let Some(ref_portal) = portals.get(&reflection) { + apples.push((turn_num, *ref_portal)); + } + } + } + + idx += 1; + } +} + +fn map_definitions() -> &'static Vec { + static MAPS: OnceLock> = OnceLock::new(); + MAPS.get_or_init(|| { + let parsed: HashMap = + serde_json::from_str(MAPS_JSON).expect("invalid bytefight maps json"); + parsed + .into_values() + .map(|map_str| MapDefinition::from_map_string(&map_str)) + .collect::, _>>() + .expect("invalid bytefight map string") + }) +} + +impl Board { + pub fn new_random(rng: &mut impl Rng) -> Self { + let maps = map_definitions(); + let valid_maps: Vec<&MapDefinition> = maps + .iter() + .filter(|map| map.width <= OBS_SIDE && map.height <= OBS_SIDE) + .collect(); + let len = valid_maps.len(); + assert!(len > 0, "no bytefight maps with dimensions <= 16x16"); + let idx = (rng.next_u64() as usize) % len; + let map = valid_maps[idx]; + map.build_initial_state(rng) + } + + pub fn bitpacked_observation_16x16(&self) -> BitpackedObservation { + const WALL_BIT: u8 = 1 << 0; + const APPLE_BIT: u8 = 1 << 1; + const OWN_BODY_BIT: u8 = 1 << 2; + const OWN_HEAD_BIT: u8 = 1 << 3; + const OWN_TRAP_BIT: u8 = 1 << 4; + const OPP_BODY_BIT: u8 = 1 << 5; + const OPP_HEAD_BIT: u8 = 1 << 6; + const OPP_TRAP_BIT: u8 = 1 << 7; + + let (wall_bitmask, apple_bitmask, snake_a_traps, snake_b_traps) = self.map.bitmasks(); + + let (own_snake, opp_snake, own_traps, opp_traps) = if self.is_player_a { + ( + &self.snake_a, + &self.snake_b, + snake_a_traps.as_slice(), + snake_b_traps.as_slice(), + ) + } else { + ( + &self.snake_b, + &self.snake_a, + snake_b_traps.as_slice(), + snake_a_traps.as_slice(), + ) + }; + + let mut obs: BitpackedObservation = [0; OBS_CELLS]; + for y in 0..OBS_SIDE { + let walls = wall_bitmask[y]; + let apples = apple_bitmask[y]; + let own_trap_row = own_traps[y]; + let opp_trap_row = opp_traps[y]; + for x in 0..OBS_SIDE { + let mask = 1u32 << x; + let mut cell = 0u8; + if walls & mask != 0 { + cell |= WALL_BIT; + } + if apples & mask != 0 { + cell |= APPLE_BIT; + } + if own_trap_row & mask != 0 { + cell |= OWN_TRAP_BIT; + } + if opp_trap_row & mask != 0 { + cell |= OPP_TRAP_BIT; + } + obs[y * OBS_SIDE + x] = cell; + } + } + + for (i, segment) in own_snake.segment_queue.iter().enumerate() { + if segment.x >= OBS_SIDE || segment.y >= OBS_SIDE { + continue; + } + let bit = if i == 0 { OWN_HEAD_BIT } else { OWN_BODY_BIT }; + obs[segment.y * OBS_SIDE + segment.x] |= bit; + } + + for (i, segment) in opp_snake.segment_queue.iter().enumerate() { + if segment.x >= OBS_SIDE || segment.y >= OBS_SIDE { + continue; + } + let bit = if i == 0 { OPP_HEAD_BIT } else { OPP_BODY_BIT }; + obs[segment.y * OBS_SIDE + segment.x] |= bit; + } + + obs + } + + pub fn new_from_state( + (width, height): (usize, usize), + queued_apples: Vec<(usize, (usize, usize))>, + apples: Vec<(usize, usize)>, + walls: Vec<(usize, usize)>, + portals: Vec<((usize, usize), (usize, usize))>, + traps: Vec<(i16, (usize, usize))>, + a_snake: Vec<(usize, usize)>, + a_queued_length: usize, + a_max_length_reached: usize, + a_apples_eaten: usize, + a_direction: Option, + b_snake: Vec<(usize, usize)>, + b_queued_length: usize, + b_max_length_reached: usize, + b_apples_eaten: usize, + b_direction: Option, + turn_num: usize, + min_player_size: usize, + is_player_a: bool, + decay_countdown_value: usize, + cached_decay_interval_value: isize, + is_decaying_value: bool, + ) -> Self { + let snake_a = Snake { + sacrifice: 1, + max_length_reached: a_max_length_reached, + queued_length: a_queued_length, + traps_this_turn: 0, + current_direction: Some(a_direction.unwrap_or(ByteFightAction::North)), + segment_queue: a_snake.into_iter().map(Point::from).collect(), + total_apples: a_apples_eaten, + }; + let snake_b = Snake { + sacrifice: 1, + max_length_reached: b_max_length_reached, + queued_length: b_queued_length, + traps_this_turn: 0, + current_direction: Some(b_direction.unwrap_or(ByteFightAction::North)), + segment_queue: b_snake.into_iter().map(Point::from).collect(), + total_apples: b_apples_eaten, + }; + + let mut board = Board { + map: Map::new((width, height), turn_num), + apple_timeline_ptr: 0, + apple_timeline: queued_apples + .into_iter() + .map(|(turn, point)| { + ( + turn, + Point { + x: point.0, + y: point.1, + }, + ) + }) + .collect(), + is_player_a, + min_player_size, + snake_a, + snake_b, + decay_countdown: decay_countdown_value, + cached_decay_interval: cached_decay_interval_value + .try_into() + .unwrap_or(DECAY_NOT_APPLIED_PLACEHOLDER), + is_decaying: is_decaying_value, + }; + + for (lifetime, loc) in traps { + let value = lifetime + (lifetime.signum() * turn_num as i16); + board.map.update_trap(loc.into(), value); + } + for (x, y) in walls { + board.map.become_wall(Point { x, y }); + } + for ((x1, y1), (x2, y2)) in portals { + board + .map + .add_portal(Point { x: x1, y: y1 }, Point { x: x2, y: y2 }); + } + for (x, y) in apples { + let point = Point { x, y }; + if board.map.is_wall(point) { + continue; + } + board.map.become_apple(point); + if let Some(portal) = board.map.portal(point) { + board.map.become_apple(portal); + } + } + board.fix_apple_head_collisions(); + add_padding_walls(&mut board.map); + + let _ = board.apply_decay(); + board + } + + pub fn terminal_state(&self) -> Option { + if self.get_valid_moves().amount() == 0 { + if self.is_player_a { + Some(TerminalState::PlayerBWin) + } else { + Some(TerminalState::PlayerAWin) + } + } else if self.map.turn_count() > LAST_TURN { + match self + .snake_a + .total_apples + .cmp(&self.snake_b.total_apples) + .then(self.snake_a.length().cmp(&self.snake_b.length())) + { + std::cmp::Ordering::Less => Some(TerminalState::PlayerBWin), + std::cmp::Ordering::Equal => Some(TerminalState::Draw), + std::cmp::Ordering::Greater => Some(TerminalState::PlayerAWin), + } + } else { + None + } + } + + pub fn get_valid_moves(&self) -> ValidMoves { + let mut valid_moves = ValidMoves::default(); + let active_snake = if self.is_player_a { + &self.snake_a + } else { + &self.snake_b + }; + + if active_snake.length() < self.min_player_size { + return valid_moves; + } + if self.map.turn_count() > LAST_TURN { + return valid_moves; + } + + if active_snake.can_afford_movement(self.min_player_size) { + let head = active_snake + .segment_queue + .front() + .expect("snake head missing"); + + let offsets = if active_snake.current_direction.is_some() { + 6..11 + } else { + 0..9 + }; + + for offset in offsets { + let direction_int = (offset + + active_snake + .current_direction + .unwrap_or(ByteFightAction::North) as u8) + % 8; + + let Some(new_loc) = head.try_add_int(direction_int) else { + continue; + }; + + let not_wall = !self.map.is_wall(new_loc); + let not_snake = active_snake.removed_on_point_sacrifice(&new_loc) + || (!self.snake_a.segment_queue.contains(&new_loc) + && !self.snake_b.segment_queue.contains(&new_loc)); + let portal_is_valid = self.map.portal(new_loc).is_none_or(|p_loc| { + let not_wall = !self.map.is_wall(p_loc); + let not_snake = active_snake.removed_on_point_sacrifice(&p_loc) + || (!self.snake_a.segment_queue.contains(&p_loc) + && !self.snake_b.segment_queue.contains(&p_loc)); + not_wall && not_snake + }); + let not_trap = self.map.trap(new_loc).abs() <= (self.map.turn_count() as i16) + || (self.is_player_a == (self.map.trap(new_loc) > 0)); + + let apple_reward = if self.map.is_apple(new_loc) { + APPLE_REWARD + } else { + 0 + }; + + let can_facetank_trap = TRAP_SACRIFICE + active_snake.sacrifice - 1 + <= active_snake.length() + apple_reward - self.min_player_size; + + if not_wall && not_snake && (not_trap || can_facetank_trap) && portal_is_valid { + valid_moves + .add(ByteFightAction::new(direction_int).expect("direction_int is valid")); + } + } + } + + if active_snake.can_place_trap(self.min_player_size) { + valid_moves.add(ByteFightAction::Trap); + } + + if active_snake.sacrifice > 1 { + valid_moves.add(ByteFightAction::EndTurn); + } + + valid_moves + } + + pub fn heuristics(&self) -> [u8; 18] { + let (wall_bitmask, apple_bitmask, snake_a_traps, snake_b_traps) = self.map.bitmasks(); + + let mut snake_a_obstacle_mask = wall_bitmask.clone(); + for i in 0..32 { + snake_a_obstacle_mask[i] |= snake_b_traps[i]; + } + for snake_b_segment in self.snake_b.segment_queue.iter() { + snake_a_obstacle_mask[snake_b_segment.y] |= 1 << snake_b_segment.x; + } + let mut snake_a_seed_arr: [u32; 32] = [0; 32]; + Board::seed_directed_array( + &mut snake_a_seed_arr, + self.snake_a.segment_queue.front().expect("snake_a head"), + self.snake_a.current_direction, + ); + let (snake_a_apple_dist, snake_a_reach) = self.run_apple_count_flood_fill( + &mut snake_a_seed_arr, + &snake_a_obstacle_mask, + *apple_bitmask, + ); + + let mut snake_b_obstacle_mask = wall_bitmask.clone(); + for i in 0..32 { + snake_b_obstacle_mask[i] |= snake_a_traps[i]; + } + for snake_a_segment in self.snake_a.segment_queue.iter() { + snake_b_obstacle_mask[snake_a_segment.y] |= 1 << snake_a_segment.x; + } + let mut snake_b_seed_arr: [u32; 32] = [0; 32]; + Board::seed_directed_array( + &mut snake_b_seed_arr, + self.snake_b.segment_queue.front().expect("snake_b head"), + self.snake_b.current_direction, + ); + let (snake_b_apple_dist, snake_b_reach) = self.run_apple_count_flood_fill( + &mut snake_b_seed_arr, + &snake_b_obstacle_mask, + *apple_bitmask, + ); + + let board_size = self.map.size() as f32; + let snake_a_head = self.snake_a.segment_queue.front().expect("snake_a head"); + let snake_b_head = self.snake_b.segment_queue.front().expect("snake_b head"); + let distance_between_snakes = usize::max( + snake_a_head.x.abs_diff(snake_b_head.x), + snake_a_head.y.abs_diff(snake_b_head.y), + ); + + let apples_eaten_diff = + ((self.snake_a.total_apples as f32) - (self.snake_b.total_apples as f32)) / 10.0; + + let turn_ratio = (self.map.turn_count() as f32 / 2000.0).clamp(0.0, 1.0); + let snake_a_reach_ratio = (snake_a_reach as f32 / board_size).clamp(0.0, 1.0); + let snake_b_reach_ratio = (snake_b_reach as f32 / board_size).clamp(0.0, 1.0); + + if self.is_player_a { + [ + Self::encode_signed_feature(turn_ratio), + self.linorm(distance_between_snakes as f32, 32.0), + self.linorm(self.snake_a.length() as f32, 32.0), + self.linorm(self.snake_b.length() as f32, 32.0), + self.dropnorm(snake_a_apple_dist[0] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[1] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[2] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[3] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[0] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[1] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[2] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[3] as f32, 32.0), + self.linorm(self.snake_a.sacrifice as f32, 32.0), + self.linorm(self.snake_a.traps_this_turn as f32, 16.0), + self.linorm(self.snake_a.max_length_reached as f32, 64.0), + self.linorm(self.snake_b.max_length_reached as f32, 64.0), + Self::encode_signed_feature(snake_a_reach_ratio), + Self::encode_signed_feature(snake_b_reach_ratio), + self.linorm(apples_eaten_diff, 10.0), + ][..18] + .try_into() + .expect("heuristics size") + } else { + [ + Self::encode_signed_feature(turn_ratio), + self.linorm(distance_between_snakes as f32, 32.0), + self.linorm(self.snake_b.length() as f32, 32.0), + self.linorm(self.snake_a.length() as f32, 32.0), + self.dropnorm(snake_b_apple_dist[0] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[1] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[2] as f32, 32.0), + self.dropnorm(snake_b_apple_dist[3] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[0] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[1] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[2] as f32, 32.0), + self.dropnorm(snake_a_apple_dist[3] as f32, 32.0), + self.linorm(self.snake_b.sacrifice as f32, 32.0), + self.linorm(self.snake_b.traps_this_turn as f32, 16.0), + self.linorm(self.snake_b.max_length_reached as f32, 64.0), + self.linorm(self.snake_a.max_length_reached as f32, 64.0), + Self::encode_signed_feature(snake_b_reach_ratio), + Self::encode_signed_feature(snake_a_reach_ratio), + self.linorm(-apples_eaten_diff, 10.0), + ][..18] + .try_into() + .expect("heuristics size") + } + } + + pub fn apply_move(&mut self, action: ByteFightAction) -> Result { + match action { + ByteFightAction::Trap => self.apply_trap(), + ByteFightAction::EndTurn => self.apply_end_turn(), + ByteFightAction::FF => Err(()), + direction => self.apply_movement(direction), + } + } + + pub fn rollback(&mut self, state: RollbackState) { + match state { + RollbackState::EndTurn { + old_num_traps, + old_sacrifice_val, + prev_cached_decay_interval, + prev_decay_countdown, + decayed_point, + snake_a_ate_during_collision, + snake_b_ate_during_collision, + snake_a_max_len, + snake_b_max_len, + prev_apple_timeline_ptr, + apples_placed_index, + } => { + self.decay_countdown = prev_decay_countdown; + let snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + match decayed_point.as_deref() { + None => {} + Some([]) => { + snake.queued_length += 1; + self.is_decaying = !self.is_decaying; + } + Some([point]) => { + snake.segment_queue.push_back(*point); + self.is_decaying = !self.is_decaying; + } + _ => unreachable!(), + } + + self.cached_decay_interval = prev_cached_decay_interval; + if snake_a_ate_during_collision { + self.snake_a.queued_length -= APPLE_REWARD; + self.snake_a.total_apples -= 1; + self.snake_a.max_length_reached = snake_a_max_len; + let head = self.snake_a.segment_queue.front().expect("snake_a head"); + self.map.become_apple(*head); + if let Some(portal) = self.map.portal(*head) { + self.map.become_apple(portal); + } + } + if snake_b_ate_during_collision { + self.snake_b.queued_length -= APPLE_REWARD; + self.snake_b.total_apples -= 1; + self.snake_b.max_length_reached = snake_b_max_len; + let head = self.snake_b.segment_queue.front().expect("snake_b head"); + self.map.become_apple(*head); + if let Some(portal) = self.map.portal(*head) { + self.map.become_apple(portal); + } + } + + self.apple_timeline_ptr = prev_apple_timeline_ptr; + for idx in apples_placed_index { + self.map.become_empty(self.apple_timeline[idx].1); + if let Some(portal) = self.map.portal(self.apple_timeline[idx].1) { + self.map.become_empty(portal); + } + } + + self.is_player_a = !self.is_player_a; + let snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + snake.traps_this_turn = old_num_traps; + snake.sacrifice = old_sacrifice_val; + + self.map.move_backward_turn(); + } + RollbackState::ApplyMove { + prev_trap_val, + sacrificed_points, + prev_queued_length, + prev_max_length_reached, + prev_direction, + head_was_apple, + } => { + let snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + let head = snake.segment_queue.front().cloned().expect("snake head"); + self.map.update_trap(head, prev_trap_val); + if let Some(portal) = self.map.portal(head) { + self.map.update_trap(portal, prev_trap_val); + } + + if head_was_apple { + self.map.become_apple(head); + if let Some(portal) = self.map.portal(head) { + self.map.become_apple(portal); + } + snake.total_apples -= 1; + } + + let _ = snake.segment_queue.pop_front(); + for point in sacrificed_points.into_iter().rev() { + snake.segment_queue.push_back(point); + } + + snake.current_direction = prev_direction.and_then(ByteFightAction::new); + snake.queued_length = prev_queued_length; + snake.max_length_reached = prev_max_length_reached; + snake.sacrifice -= 1; + } + RollbackState::ApplyTrap { old_trap_val, trap } => { + self.map.update_trap(trap, old_trap_val); + if let Some(portal) = self.map.portal(trap) { + self.map.update_trap(portal, old_trap_val); + } + let snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + + snake.traps_this_turn -= 1; + snake.segment_queue.push_back(trap); + } + } + } + + fn spawn_apples(&mut self) -> (Vec, bool, bool) { + let mut apples_placed = Vec::new(); + while self.apple_timeline_ptr < self.apple_timeline.len() + && self.apple_timeline[self.apple_timeline_ptr].0 <= self.map.turn_count() + { + let spawn_point = self.apple_timeline[self.apple_timeline_ptr].1; + if self.map.is_empty(spawn_point) { + self.map.become_apple(spawn_point); + apples_placed.push(self.apple_timeline_ptr); + if let Some(portal_location) = self.map.portal(spawn_point) { + self.map.become_apple(portal_location); + } + } + + self.apple_timeline_ptr += 1; + } + + let (place_a, place_b) = self.fix_apple_head_collisions(); + (apples_placed, place_a, place_b) + } + + fn update_decay_interval(&mut self) { + if self.decay_countdown != 0 { + return; + } + + for (turn, interval) in &DECAY_TIMELINE { + if self.map.turn_count() < *turn { + break; + } + + self.cached_decay_interval = *interval; + } + } + + fn apply_decay(&mut self) -> Result>, ()> { + if self.cached_decay_interval == DECAY_NOT_APPLIED_PLACEHOLDER { + return Ok(None); + } + + let decayed = if self.is_decaying || self.decay_countdown == 0 { + let decayed = if self.is_player_a { + self.snake_a.apply_sacrifice(1)? + } else { + self.snake_b.apply_sacrifice(1)? + }; + self.is_decaying = !self.is_decaying; + + Some(decayed) + } else { + None + }; + self.decay_countdown = (self.decay_countdown + 1) % self.cached_decay_interval; + + Ok(decayed) + } + + fn apply_movement(&mut self, action: ByteFightAction) -> Result { + let current_snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + + let prev_queued_length = current_snake.queued_length; + let prev_max_length_reached = current_snake.max_length_reached; + let prev_direction = current_snake.current_direction.map(|a| a as u8); + + if !current_snake.can_afford_movement(self.min_player_size) { + return Err(()); + } + let Ok((new_head, mut cells_lost)) = current_snake.push_move(action) else { + return Err(()); + }; + + if self.map.is_wall(new_head) { + return Err(()); + } + + let portal = self.map.portal(new_head); + if let Some(portal) = portal { + if self.map.is_wall(portal) { + return Err(()); + } + current_snake.segment_queue.push_front(portal); + } else { + current_snake.segment_queue.push_front(new_head); + } + + let mut head_was_apple = false; + if self.map.is_apple(new_head) { + head_was_apple = true; + current_snake.eat_apple(); + self.map.become_empty(new_head); + if let Some(portal) = portal { + self.map.become_empty(portal); + } + } + + let old_trap_val = self.map.trap(new_head); + let trap_val = self.map.trap(new_head); + if trap_val.abs() > self.map.turn_count() as i16 { + let is_player_a_trap = trap_val > 0; + let is_enemy_trap = is_player_a_trap ^ self.is_player_a; + if is_enemy_trap { + match current_snake.apply_sacrifice(3) { + Ok(mut sacrifice) => { + cells_lost.append(&mut sacrifice); + } + Err(_) => { + return Err(()); + } + } + self.map.update_trap(new_head, 0); + if let Some(portal) = portal { + self.map.update_trap(portal, 0); + } + } else { + let trap_val = self.map.turn_count() as i16 + TRAP_LIFETIME; + + self.map + .update_trap(new_head, trap_val * if self.is_player_a { 1 } else { -1 }); + if let Some(portal) = portal { + self.map + .update_trap(portal, trap_val * if self.is_player_a { 1 } else { -1 }); + } + } + } + + Ok(RollbackState::ApplyMove { + prev_trap_val: old_trap_val, + sacrificed_points: cells_lost, + prev_queued_length, + prev_max_length_reached, + prev_direction, + head_was_apple, + }) + } + + fn apply_trap(&mut self) -> Result { + let snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + + if !snake.can_place_trap(self.min_player_size) { + return Err(()); + } + + snake.traps_this_turn += 1; + let trap = snake.segment_queue.pop_back().ok_or(())?; + + let old_trap_val = self.map.trap(trap); + let trap_val = self.map.turn_count() as i16 + TRAP_LIFETIME; + self.map + .update_trap(trap, trap_val * if self.is_player_a { 1 } else { -1 }); + if let Some(portal) = self.map.portal(trap) { + self.map + .update_trap(portal, trap_val * if self.is_player_a { 1 } else { -1 }); + } + + Ok(RollbackState::ApplyTrap { old_trap_val, trap }) + } + + fn apply_end_turn(&mut self) -> Result { + let current_snake = if self.is_player_a { + &mut self.snake_a + } else { + &mut self.snake_b + }; + if current_snake.sacrifice == 1 { + return Err(()); + } + self.map.move_forward_turn(); + + let traps_this_turn = current_snake.traps_this_turn; + let sacrifice = current_snake.sacrifice; + let prev_cached_decay_interval = self.cached_decay_interval; + let prev_decay_countdown = self.decay_countdown; + let prev_apple_timeline_ptr = self.apple_timeline_ptr; + + current_snake.traps_this_turn = 0; + current_snake.sacrifice = 1; + + let snake_a_max_len = self.snake_a.max_length_reached; + let snake_b_max_len = self.snake_b.max_length_reached; + + self.is_player_a = !self.is_player_a; + let (apples_placed_index, snake_a_ate_during_collision, snake_b_ate_during_collision) = + self.spawn_apples(); + + self.update_decay_interval(); + let decayed = self.apply_decay().expect( + "Snake length became 0. This means we made multiple incorrect moves and something is irrecoverably wrong.", + ); + + Ok(RollbackState::EndTurn { + old_num_traps: traps_this_turn, + old_sacrifice_val: sacrifice, + prev_cached_decay_interval, + decayed_point: decayed, + prev_decay_countdown, + snake_a_ate_during_collision, + snake_b_ate_during_collision, + snake_a_max_len, + snake_b_max_len, + prev_apple_timeline_ptr, + apples_placed_index, + }) + } + + fn fix_apple_head_collisions(&mut self) -> (bool, bool) { + let mut snake_a_ate_during_collision = false; + let mut snake_b_ate_during_collision = false; + if let Some(&snake_a_location) = self.snake_a.segment_queue.front() { + if self.map.is_apple(snake_a_location) { + snake_a_ate_during_collision = true; + self.snake_a.queued_length += APPLE_REWARD; + self.snake_a.total_apples += 1; + self.snake_a.max_length_reached = + self.snake_a.max_length_reached.max(self.snake_a.length()); + + self.map.become_empty(snake_a_location); + if let Some(alt_location) = self.map.portal(snake_a_location) { + self.map.become_empty(alt_location); + } + } + } + + if let Some(&snake_b_location) = self.snake_b.segment_queue.front() { + if self.map.is_apple(snake_b_location) { + snake_b_ate_during_collision = true; + self.snake_b.queued_length += APPLE_REWARD; + self.snake_b.total_apples += 1; + self.snake_b.max_length_reached = + self.snake_b.max_length_reached.max(self.snake_b.length()); + + self.map.become_empty(snake_b_location); + if let Some(alt_location) = self.map.portal(snake_b_location) { + self.map.become_empty(alt_location); + } + } + } + + (snake_a_ate_during_collision, snake_b_ate_during_collision) + } + + fn seed_directed_array( + seed_arr: &mut [u32; 32], + origin: &Point, + facing_dir: Option, + ) { + if let Some(dir) = facing_dir { + for offset in 6..11 { + if let Some(new_dir) = origin.try_add_int(((dir as u8) + offset) % 8) { + seed_arr[new_dir.y] |= 1 << new_dir.x; + } + } + } else { + seed_arr[origin.y] |= 0b11 << origin.x; + seed_arr[origin.y] |= 0xC0000000 >> (31 - origin.x); + + if origin.y < 31 { + seed_arr[origin.y + 1] |= 0b11 << origin.x; + seed_arr[origin.y + 1] |= 0xC0000000 >> (31 - origin.x); + } + if origin.y > 0 { + seed_arr[origin.y - 1] |= 0b11 << origin.x; + seed_arr[origin.y - 1] |= 0xC0000000 >> (31 - origin.x); + } + seed_arr[origin.y] &= !(1 << origin.x); + } + } + + fn run_apple_count_flood_fill( + &self, + seed_arr: &mut [u32; 32], + obstacles: &[u32; 32], + mut apples: [u32; 32], + ) -> ([u32; 4], u32) { + let mut apples_found: u32 = 0; + let mut apple_loc: [u32; 4] = [512; 4]; + let mut apple_pntr = self.apple_timeline_ptr; + + for i in 0..32 { + apples_found += (seed_arr[i] & apples[i]).count_ones(); + } + + for i in 0..std::cmp::min(apples_found, 4) { + apple_loc[i as usize] = 1; + } + + for epoch in 0..32 { + for i in 0..32 { + seed_arr[i] |= seed_arr[i] << 1 | seed_arr[i] >> 1; + } + for i in 1..32 { + seed_arr[i - 1] |= seed_arr[i] + } + for i in 0..31 { + seed_arr[31 - i] |= seed_arr[30 - i] + } + + for i in 0..32 { + seed_arr[i] &= !obstacles[i]; + } + + for (p1, p2) in self.map.portals() { + let p1_reachable = (seed_arr[p1.y] >> p1.x) & 1; + let p2_reachable = (seed_arr[p2.y] >> p2.x) & 1; + seed_arr[p1.y] |= p2_reachable << p1.x; + seed_arr[p2.y] |= p1_reachable << p2.x; + } + + while apple_pntr < self.apple_timeline.len() + && self.apple_timeline[apple_pntr].0 <= self.map.turn_count() + 2 * epoch + { + apples[self.apple_timeline[apple_pntr].1.y] |= + 1 << self.apple_timeline[apple_pntr].1.x; + apple_pntr += 1; + } + + let mut apples_this_turn = 0; + for i in 0..32 { + apples_this_turn += (seed_arr[i] & apples[i]).count_ones(); + } + for i in apples_found..std::cmp::min(apples_this_turn, 4) { + apple_loc[i as usize] = (epoch + 1) as u32; + apples_found += 1; + } + } + + let mut reached_tiles: u32 = 0; + for i in 0..32 { + reached_tiles += seed_arr[i].count_ones(); + } + + (apple_loc, reached_tiles) + } + + fn linorm(&self, x: f32, softmax: f32) -> u8 { + let abs_x = x.abs(); + let normalized = if abs_x <= softmax { + (0.8 / softmax) * x + } else { + x.signum() * (1.0 - (0.2 * softmax) / abs_x) + }; + Self::encode_signed_feature(normalized) + } + + fn dropnorm(&self, x: f32, softmax: f32) -> u8 { + let normalized = if x <= softmax { + 1.0 - (0.75 / softmax) * x + } else { + 0.0 + }; + Self::encode_signed_feature(normalized) + } + + #[inline] + fn encode_signed_feature(value: f32) -> u8 { + let clamped = value.clamp(-1.0, 1.0); + let quantized = (clamped * 127.0).round() as i16 + 128; + quantized.clamp(0, 255) as u8 + } +} + +fn parse_usize(value: &str, label: &str) -> Result { + value + .parse::() + .map_err(|_| format!("invalid {}", label)) +} + +fn parse_pair(value: &str, delimiter: char) -> Result<(usize, usize), String> { + let mut iter = value.split(delimiter); + let first = iter.next().ok_or_else(|| "missing first".to_string())?; + let second = iter.next().ok_or_else(|| "missing second".to_string())?; + if iter.next().is_some() { + return Err("too many parts".to_string()); + } + Ok(( + parse_usize(first, "pair_x")?, + parse_usize(second, "pair_y")?, + )) +} + +fn parse_point(value: &str) -> Result { + let (x, y) = parse_pair(value, ',')?; + Ok(Point { x, y }) +} + +fn parse_portals_section(value: &str) -> Result, String> { + if value.is_empty() { + return Ok(Vec::new()); + } + let mut portals = Vec::new(); + for portal in value.split('_') { + if portal.is_empty() { + continue; + } + let parts: Vec<&str> = portal.split(',').collect(); + if parts.len() != 4 { + return Err("invalid portal entry".to_string()); + } + let p1 = Point { + x: parse_usize(parts[0], "portal_x1")?, + y: parse_usize(parts[1], "portal_y1")?, + }; + let p2 = Point { + x: parse_usize(parts[2], "portal_x2")?, + y: parse_usize(parts[3], "portal_y2")?, + }; + portals.push((p1, p2)); + } + Ok(portals) +} + +fn parse_apple_timeline(value: &str) -> Result, String> { + if value.is_empty() { + return Ok(Vec::new()); + } + let mut timeline = Vec::new(); + for entry in value.split('_') { + if entry.is_empty() { + continue; + } + let parts: Vec<&str> = entry.split(',').collect(); + if parts.len() != 3 { + return Err("invalid apple entry".to_string()); + } + let turn = parse_usize(parts[0], "apple_turn")?; + let point = Point { + x: parse_usize(parts[1], "apple_x")?, + y: parse_usize(parts[2], "apple_y")?, + }; + timeline.push((turn, point)); + } + Ok(timeline) +} + +fn parse_walls_bits(bits: &str, width: usize, height: usize) -> Result, String> { + if bits.len() != width * height { + return Err("wall bit length mismatch".to_string()); + } + let mut walls = Vec::new(); + for (i, ch) in bits.chars().enumerate() { + if ch == '1' { + let x = i % width; + let y = i / width; + walls.push(Point { x, y }); + } + } + Ok(walls) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_apply_and_rollback_move() { + let mut board = Board::new_from_state( + (5, 5), + vec![], + vec![], + vec![], + vec![], + vec![], + vec![(2, 2)], + 1, + 2, + 0, + None, + vec![(4, 4)], + 1, + 2, + 0, + None, + 0, + 1, + true, + 0, + DECAY_NOT_APPLIED_PLACEHOLDER as isize, + false, + ); + let snapshot = board.clone(); + let rollback = board.apply_move(ByteFightAction::East).expect("valid move"); + board.rollback(rollback); + assert_eq!(board, snapshot); + } +} diff --git a/training/src/environments/bytefight/map.rs b/training/src/environments/bytefight/map.rs new file mode 100644 index 0000000..72199e4 --- /dev/null +++ b/training/src/environments/bytefight/map.rs @@ -0,0 +1,176 @@ +use std::hash::{Hash, Hasher}; + +use super::types::Point; + +#[derive(Debug, Clone)] +pub struct Map { + wall_bitmask: [u32; 32], + apple_bitmask: [u32; 32], + dimensions: (usize, usize), + trap_mask: [[i16; 32]; 32], + portals: Vec<(Point, Point)>, + turn_count: usize, +} + +impl Map { + pub fn new(dimensions: (usize, usize), turn_count: usize) -> Self { + Map { + apple_bitmask: [0; 32], + wall_bitmask: [0; 32], + dimensions, + trap_mask: [[0; 32]; 32], + portals: Vec::new(), + turn_count, + } + } + + pub fn dimensions(&self) -> (usize, usize) { + self.dimensions + } + + pub fn is_empty(&self, point: Point) -> bool { + self.wall_bitmask[point.y] & (1 << point.x) == 0 + && self.apple_bitmask[point.y] & (1 << point.x) == 0 + } + + pub fn is_wall(&self, point: Point) -> bool { + self.wall_bitmask[point.y] & (1 << point.x) != 0 + } + + pub fn is_apple(&self, point: Point) -> bool { + self.apple_bitmask[point.y] & (1 << point.x) != 0 + } + + pub fn portal(&self, point: Point) -> Option { + self.portals.iter().find_map(|&(p1, p2)| { + if p1 == point { + Some(p2) + } else if p2 == point { + Some(p1) + } else { + None + } + }) + } + + pub fn portals(&self) -> &[(Point, Point)] { + &self.portals + } + + pub fn trap(&self, point: Point) -> i16 { + self.trap_mask[point.y][point.x] + } + + pub fn become_apple(&mut self, point: Point) { + self.apple_bitmask[point.y] |= 1 << point.x; + } + + pub fn become_empty(&mut self, point: Point) { + self.wall_bitmask[point.y] &= !(1 << point.x); + self.apple_bitmask[point.y] &= !(1 << point.x); + } + + pub fn update_trap(&mut self, point: Point, trap: i16) { + self.trap_mask[point.y][point.x] = trap; + } + + pub fn become_wall(&mut self, point: Point) { + self.wall_bitmask[point.y] |= 1 << point.x; + } + + pub fn add_portal(&mut self, p1: Point, p2: Point) { + self.portals.push((p1, p2)); + } + + pub fn bitmasks(&self) -> (&[u32; 32], &[u32; 32], [u32; 32], [u32; 32]) { + let mut snake_a_traps: [u32; 32] = [0; 32]; + let mut snake_b_traps: [u32; 32] = [0; 32]; + + for y in 0..32 { + for x in 0..32 { + if self.trap_mask[y][x] > self.turn_count as i16 { + snake_a_traps[y] |= 1 << x; + } else if -self.trap_mask[y][x] > self.turn_count as i16 { + snake_b_traps[y] |= 1 << x; + } + } + } + + ( + &self.wall_bitmask, + &self.apple_bitmask, + snake_a_traps, + snake_b_traps, + ) + } + + pub fn turn_count(&self) -> usize { + self.turn_count + } + + pub fn size(&self) -> usize { + self.dimensions.0 * self.dimensions.1 + } + + pub fn move_forward_turn(&mut self) { + self.turn_count += 1; + } + + pub fn move_backward_turn(&mut self) { + self.turn_count -= 1; + } +} + +impl PartialEq for Map { + fn eq(&self, other: &Self) -> bool { + self.turn_count == other.turn_count + && self.portals == other.portals + && self.apple_bitmask == other.apple_bitmask + && self.wall_bitmask == other.wall_bitmask + && self + .trap_mask + .iter() + .flatten() + .zip(other.trap_mask.iter().flatten()) + .all(|(left, right)| { + (left.abs() <= self.turn_count as i16 && right.abs() <= self.turn_count as i16) + || left == right + }) + } +} + +impl Eq for Map {} + +impl Hash for Map { + fn hash(&self, state: &mut H) { + self.wall_bitmask.hash(state); + self.apple_bitmask.hash(state); + self.dimensions.hash(state); + self.portals.hash(state); + self.turn_count.hash(state); + for row in &self.trap_mask { + for value in row { + let normalized = if value.abs() <= self.turn_count as i16 { + 0 + } else { + *value + }; + normalized.hash(state); + } + } + } +} + +pub fn add_padding_walls(map: &mut Map) { + let (width, height) = map.dimensions(); + for x in width..32 { + for y in 0..32 { + map.become_wall(Point { x, y }); + } + } + for y in height..32 { + for x in 0..32 { + map.become_wall(Point { x, y }); + } + } +} diff --git a/training/src/environments/bytefight/mod.rs b/training/src/environments/bytefight/mod.rs new file mode 100644 index 0000000..d99dddd --- /dev/null +++ b/training/src/environments/bytefight/mod.rs @@ -0,0 +1,392 @@ +use ndarray::{ArrayViewMut, Ix2}; +use serde::{Deserialize, Serialize}; + +use crate::{Environment, GameNotation, Player, TerminalState}; + +pub mod game; +pub mod map; +pub mod pen; +pub mod snake; +pub mod types; + +pub use pen::{ByteFightPen, PenError}; +pub use types::{ByteFightAction, ByteFightPolicyAction, Point}; + +#[derive(Debug, Clone, Copy, Default)] +struct PolicyValidMoves(u8); + +#[derive(Debug, Clone, Copy)] +struct PolicyValidMovesIter { + mask: u8, + index: u8, +} + +impl Iterator for PolicyValidMovesIter { + type Item = ByteFightPolicyAction; + + fn next(&mut self) -> Option { + while self.index < ByteFight::NUM_ACTIONS as u8 { + let idx = self.index; + self.index += 1; + if self.mask & (1 << idx) != 0 { + return ByteFightPolicyAction::new(idx); + } + } + None + } +} + +impl PolicyValidMoves { + #[inline] + fn add(&mut self, action: ByteFightPolicyAction) { + self.0 |= 1 << (action as u8); + } + + fn into_iter(self) -> PolicyValidMovesIter { + PolicyValidMovesIter { + mask: self.0, + index: 0, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(from = "ByteFightPen", into = "ByteFightPen")] +pub struct ByteFight { + board: game::Board, +} + +impl From for ByteFight { + fn from(pen: ByteFightPen) -> Self { + ByteFight { + board: pen.into_board().expect("invalid ByteFight PEN"), + } + } +} + +impl From for ByteFightPen { + fn from(board: ByteFight) -> Self { + ByteFightPen::from(&board.board) + } +} + +impl From<&ByteFight> for ByteFightPen { + fn from(board: &ByteFight) -> Self { + ByteFightPen::from(&board.board) + } +} + +impl crate::Action for ByteFightPolicyAction { + fn to_index(self) -> usize { + self as usize + } + + fn from_index(index: usize) -> Option { + ByteFightPolicyAction::new(index as u8) + } +} + +impl Environment for ByteFight { + type ObsElem = u8; + type ObsDim = Ix2; + type Action = ByteFightPolicyAction; + type RollbackState = game::RollbackState; + const NUM_ACTIONS: usize = 7; + const OBS_SHAPE: Ix2 = Ix2(types::OBS_SERIALIZED_SIDE, types::OBS_SERIALIZED_WIDTH); + + fn new() -> Self { + let mut rng = rand::rng(); + let board = game::Board::new_random(&mut rng); + ByteFight::from(ByteFightPen::from(&board)) + } + + fn is_terminal(&self) -> Option { + match self.board.terminal_state()? { + types::TerminalState::PlayerAWin => Some(TerminalState::Win(Player::PlayerA)), + types::TerminalState::PlayerBWin => Some(TerminalState::Win(Player::PlayerB)), + types::TerminalState::Draw => Some(TerminalState::Draw), + } + } + + fn valid_actions(&self) -> impl Iterator { + let valid = self.board.get_valid_moves(); + let mut policy_moves = PolicyValidMoves::default(); + + for action in [ + ByteFightPolicyAction::Forward, + ByteFightPolicyAction::Left, + ByteFightPolicyAction::LeftForward, + ByteFightPolicyAction::Right, + ByteFightPolicyAction::RightForward, + ] { + let absolute = self.policy_to_absolute(action); + if valid.contains(absolute) { + policy_moves.add(action); + } + } + + if valid.contains(ByteFightAction::Trap) { + policy_moves.add(ByteFightPolicyAction::Trap); + } + if valid.contains(ByteFightAction::EndTurn) { + policy_moves.add(ByteFightPolicyAction::EndTurn); + } + + policy_moves.into_iter() + } + + fn current_player(&self) -> Player { + if self.board.is_player_a { + Player::PlayerA + } else { + Player::PlayerB + } + } + + fn observation(&self, mut out: ArrayViewMut) { + let out_slice = out + .as_slice_mut() + .expect("bytefight observation output must be contiguous"); + + let bitpacked = self.board.bitpacked_observation_16x16(); + out_slice[..types::OBS_CELLS].copy_from_slice(&bitpacked); + + let mut offset = types::OBS_CELLS; + let direction = self.active_direction() as usize; + for idx in 0..types::OBS_DIRECTIONS { + out_slice[offset + idx] = if idx == direction { 1 } else { 0 }; + } + offset += types::OBS_DIRECTIONS; + + let heuristics = self.board.heuristics(); + out_slice[offset..offset + types::OBS_HEURISTICS].copy_from_slice(&heuristics); + offset += types::OBS_HEURISTICS; + + for byte in &mut out_slice[offset..] { + *byte = 0; + } + } + + fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { + let absolute = self.policy_to_absolute(action); + self.board + .apply_move(absolute) + .expect("apply_action called with invalid action") + } + + fn rollback(&mut self, rollback: Self::RollbackState) { + self.board.rollback(rollback); + } +} + +impl ByteFight { + fn active_direction(&self) -> ByteFightAction { + let snake = if self.board.is_player_a { + &self.board.snake_a + } else { + &self.board.snake_b + }; + snake.current_direction.unwrap_or(ByteFightAction::North) + } + + fn policy_to_absolute(&self, action: ByteFightPolicyAction) -> ByteFightAction { + let direction = self.active_direction() as u8; + match action { + ByteFightPolicyAction::Forward => { + ByteFightAction::new(direction).expect("valid direction") + } + ByteFightPolicyAction::Left => { + ByteFightAction::new((direction + 6) % 8).expect("valid direction") + } + ByteFightPolicyAction::LeftForward => { + ByteFightAction::new((direction + 7) % 8).expect("valid direction") + } + ByteFightPolicyAction::Right => { + ByteFightAction::new((direction + 2) % 8).expect("valid direction") + } + ByteFightPolicyAction::RightForward => { + ByteFightAction::new((direction + 1) % 8).expect("valid direction") + } + ByteFightPolicyAction::Trap => ByteFightAction::Trap, + ByteFightPolicyAction::EndTurn => ByteFightAction::EndTurn, + } + } +} + +impl GameNotation for ByteFight { + type Error = PenError; + + fn to_notation(&self) -> String { + ByteFightPen::from(self).0 + } + + fn from_notation(s: &str) -> Result { + let pen = ByteFightPen(s.to_string()); + let board = pen.into_board()?; + Ok(ByteFight { board }) + } +} + +#[cfg(test)] +mod tests { + use crate::Action; + use rand::{RngCore, SeedableRng}; + use rand_chacha::ChaCha8Rng; + use rstest::rstest; + + use super::*; + + #[test] + fn test_action_trait() { + assert_eq!(ByteFightPolicyAction::Forward.to_index(), 0); + assert_eq!(ByteFightPolicyAction::EndTurn.to_index(), 6); + assert_eq!( + ByteFightPolicyAction::from_index(0), + Some(ByteFightPolicyAction::Forward) + ); + assert_eq!( + ByteFightPolicyAction::from_index(6), + Some(ByteFightPolicyAction::EndTurn) + ); + assert_eq!(ByteFightPolicyAction::from_index(7), None); + } + + #[test] + fn test_pen_roundtrip() { + let mut rng = ChaCha8Rng::seed_from_u64(7); + let board = game::Board::new_random(&mut rng); + let pen = ByteFightPen::from(&board); + let rebuilt = pen.clone().into_board().expect("valid pen"); + assert_eq!(pen.0, ByteFightPen::from(&rebuilt).0); + } + + #[test] + fn test_notation_roundtrip() { + let mut rng = ChaCha8Rng::seed_from_u64(42); + let game = ByteFight { + board: game::Board::new_random(&mut rng), + }; + + let notation = game.to_notation(); + let restored = ByteFight::from_notation(¬ation).expect("valid notation"); + + // Compare via notation since board equality might differ in internal state + assert_eq!(notation, restored.to_notation()); + } + + #[rstest] + #[case(1)] + #[case(7)] + #[case(13)] + #[case(23)] + #[case(42)] + #[case(77)] + #[case(101)] + #[case(123)] + #[case(256)] + #[case(999)] + fn test_random_move_sequence_roundtrip(#[case] seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + + for _ in 0..5 { + let mut board = game::Board::new_random(&mut rng); + let mut snapshots = Vec::new(); + let mut snapshot_pens = Vec::new(); + snapshots.push(board.clone()); + snapshot_pens.push(ByteFightPen::from(&board).0); + + let mut rollbacks = Vec::new(); + + for _ in 0..150 { + let valid: Vec<_> = board.get_valid_moves().actions().collect(); + if valid.is_empty() { + break; + } + let idx = (rng.next_u64() as usize) % valid.len(); + let action = valid[idx]; + let rollback = board.apply_move(action).expect("valid move"); + rollbacks.push(rollback); + snapshots.push(board.clone()); + snapshot_pens.push(ByteFightPen::from(&board).0); + } + + for (idx, rollback) in rollbacks.into_iter().rev().enumerate() { + board.rollback(rollback); + let snapshot_idx = snapshots.len() - 2 - idx; + assert_eq!(board, snapshots[snapshot_idx]); + assert_eq!(ByteFightPen::from(&board).0, snapshot_pens[snapshot_idx]); + } + } + } + + #[rstest] + #[case(1)] + #[case(42)] + #[case(99)] + fn test_notation_roundtrip_after_moves(#[case] seed: u64) { + use crate::Environment; + + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let mut game = ByteFight { + board: game::Board::new_random(&mut rng), + }; + + // Apply some random moves + for _ in 0..20 { + let valid: Vec<_> = game.valid_actions().collect(); + if valid.is_empty() { + break; + } + let idx = (rng.next_u64() as usize) % valid.len(); + game.apply_action(valid[idx]); + } + + let notation = game.to_notation(); + let restored = ByteFight::from_notation(¬ation).expect("valid notation"); + assert_eq!(notation, restored.to_notation()); + } + + #[test] + fn test_relative_valid_actions_from_direction() { + use crate::Environment; + + let game = ByteFight { + board: game::Board::new_from_state( + (5, 5), + vec![], + vec![], + vec![], + vec![], + vec![], + vec![(2, 2)], + 1, + 2, + 0, + Some(ByteFightAction::North), + vec![(4, 4)], + 1, + 2, + 0, + Some(ByteFightAction::North), + 0, + 1, + true, + 0, + 9999, + false, + ), + }; + + let actions: Vec<_> = game.valid_actions().collect(); + assert_eq!( + actions, + vec![ + ByteFightPolicyAction::Forward, + ByteFightPolicyAction::Left, + ByteFightPolicyAction::LeftForward, + ByteFightPolicyAction::Right, + ByteFightPolicyAction::RightForward, + ] + ); + } +} diff --git a/training/src/environments/bytefight/pen.rs b/training/src/environments/bytefight/pen.rs new file mode 100644 index 0000000..c3a2326 --- /dev/null +++ b/training/src/environments/bytefight/pen.rs @@ -0,0 +1,510 @@ +use std::collections::VecDeque; +use std::fmt; + +use super::game::Board; +use super::map::{add_padding_walls, Map}; +use super::snake::Snake; +use super::types::{ByteFightAction, Point}; + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(transparent)] +pub struct ByteFightPen(pub String); + +#[derive(Debug)] +pub struct PenError(String); + +impl PenError { + fn new(message: impl Into) -> Self { + Self(message.into()) + } +} + +impl fmt::Display for PenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for PenError {} + +impl ByteFightPen { + pub fn from_board(board: &Board) -> Self { + ByteFightPen(board_to_pen(board)) + } + + pub fn into_board(self) -> Result { + board_from_pen(&self.0) + } +} + +impl From for ByteFightPen { + fn from(board: Board) -> Self { + ByteFightPen::from_board(&board) + } +} + +impl From<&Board> for ByteFightPen { + fn from(board: &Board) -> Self { + ByteFightPen::from_board(board) + } +} + +fn board_to_pen(board: &Board) -> String { + let (width, height) = board.map.dimensions(); + let mut sections = Vec::new(); + sections.push(format!("{}x{}", width, height)); + sections.push(format!("t{}", board.map.turn_count())); + sections.push(format!("p{}", if board.is_player_a { "A" } else { "B" })); + sections.push(format!("m{}", board.min_player_size)); + sections.push(format!( + "d{},{},{}", + board.decay_countdown, + board.cached_decay_interval, + if board.is_decaying { 1 } else { 0 } + )); + sections.push(format!("a{}", board.apple_timeline_ptr)); + sections.push(format!("w{}", encode_walls(&board.map))); + sections.push(format!("o{}", encode_portals(&board.map))); + sections.push(format!("f{}", encode_apples(&board.map))); + sections.push(format!("q{}", encode_timeline(&board.apple_timeline))); + sections.push(format!("r{}", encode_traps(&board.map))); + sections.push(format!("A{}", encode_snake(&board.snake_a))); + sections.push(format!("B{}", encode_snake(&board.snake_b))); + sections.join("|") +} + +fn board_from_pen(pen: &str) -> Result { + let parts: Vec<&str> = pen.split('|').collect(); + if parts.len() != 13 { + return Err(PenError::new("invalid PEN section count")); + } + + let (width, height) = parse_dimensions(parts[0])?; + let turn_count = parse_prefixed_usize(parts[1], 't')?; + let is_player_a = parse_prefixed_player(parts[2])?; + let min_player_size = parse_prefixed_usize(parts[3], 'm')?; + let (decay_countdown, cached_decay_interval, is_decaying) = parse_decay(parts[4], 'd')?; + let apple_timeline_ptr = parse_prefixed_usize(parts[5], 'a')?; + let walls = parse_prefixed_walls(parts[6], width, height)?; + let portals = parse_prefixed_portals(parts[7])?; + let apples = parse_prefixed_points(parts[8], 'f')?; + let apple_timeline = parse_prefixed_timeline(parts[9], 'q')?; + let traps = parse_prefixed_traps(parts[10], 'r')?; + let snake_a = parse_prefixed_snake(parts[11], 'A')?; + let snake_b = parse_prefixed_snake(parts[12], 'B')?; + + let mut map = Map::new((width, height), turn_count); + for wall in walls { + map.become_wall(wall); + } + for (p1, p2) in portals { + map.add_portal(p1, p2); + } + for apple in apples { + map.become_apple(apple); + } + for (trap, point) in traps { + map.update_trap(point, trap); + } + add_padding_walls(&mut map); + + Ok(Board { + map, + apple_timeline, + apple_timeline_ptr, + snake_a, + snake_b, + is_player_a, + min_player_size, + decay_countdown, + cached_decay_interval, + is_decaying, + }) +} + +fn encode_walls(map: &Map) -> String { + let (width, height) = map.dimensions(); + let mut rows = Vec::with_capacity(height); + for y in 0..height { + let mut row = String::new(); + let mut empty_run = 0usize; + for x in 0..width { + if map.is_wall(Point { x, y }) { + if empty_run > 0 { + row.push_str(&empty_run.to_string()); + empty_run = 0; + } + row.push('#'); + } else { + empty_run += 1; + } + } + if empty_run > 0 { + row.push_str(&empty_run.to_string()); + } + rows.push(row); + } + rows.join("/") +} + +fn encode_portals(map: &Map) -> String { + if map.portals().is_empty() { + return "-".to_string(); + } + map.portals() + .iter() + .map(|(p1, p2)| format!("{},{}~{},{}", p1.x, p1.y, p2.x, p2.y)) + .collect::>() + .join(";") +} + +fn encode_apples(map: &Map) -> String { + let (width, height) = map.dimensions(); + let mut apples = Vec::new(); + for y in 0..height { + for x in 0..width { + if map.is_apple(Point { x, y }) { + apples.push(Point { x, y }); + } + } + } + if apples.is_empty() { + return "-".to_string(); + } + apples + .into_iter() + .map(|point| format!("{},{}", point.x, point.y)) + .collect::>() + .join(";") +} + +fn encode_timeline(timeline: &[(usize, Point)]) -> String { + if timeline.is_empty() { + return "-".to_string(); + } + timeline + .iter() + .map(|(turn, point)| format!("{},{},{}", turn, point.x, point.y)) + .collect::>() + .join(";") +} + +fn encode_traps(map: &Map) -> String { + let (width, height) = map.dimensions(); + let mut traps = Vec::new(); + for y in 0..height { + for x in 0..width { + let value = map.trap(Point { x, y }); + if value != 0 && value.abs() > map.turn_count() as i16 { + traps.push(format!("{},{},{}", value, x, y)); + } + } + } + if traps.is_empty() { + return "-".to_string(); + } + traps.join(";") +} + +fn encode_snake(snake: &Snake) -> String { + let direction = snake + .current_direction + .map(|dir| (dir as u8).to_string()) + .unwrap_or_else(|| "-".to_string()); + let segments = snake + .segment_queue + .iter() + .map(|point| format!("{},{}", point.x, point.y)) + .collect::>() + .join(">"); + format!( + "{},{},{},{},{},{}:{}", + direction, + snake.queued_length, + snake.max_length_reached, + snake.total_apples, + snake.sacrifice, + snake.traps_this_turn, + segments + ) +} + +fn parse_dimensions(value: &str) -> Result<(usize, usize), PenError> { + let mut iter = value.split('x'); + let width = iter.next().ok_or_else(|| PenError::new("missing width"))?; + let height = iter.next().ok_or_else(|| PenError::new("missing height"))?; + if iter.next().is_some() { + return Err(PenError::new("invalid dimensions")); + } + Ok((parse_usize(width, "width")?, parse_usize(height, "height")?)) +} + +fn parse_prefixed_usize(section: &str, prefix: char) -> Result { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing prefix"))?; + parse_usize(value, "value") +} + +fn parse_prefixed_player(section: &str) -> Result { + let value = section + .strip_prefix('p') + .ok_or_else(|| PenError::new("missing player prefix"))?; + match value { + "A" => Ok(true), + "B" => Ok(false), + _ => Err(PenError::new("invalid player")), + } +} + +fn parse_decay(section: &str, prefix: char) -> Result<(usize, usize, bool), PenError> { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing decay prefix"))?; + let parts: Vec<&str> = value.split(',').collect(); + if parts.len() != 3 { + return Err(PenError::new("invalid decay section")); + } + let countdown = parse_usize(parts[0], "decay_countdown")?; + let interval = parse_usize(parts[1], "decay_interval")?; + let is_decaying = match parts[2] { + "1" => true, + "0" => false, + _ => return Err(PenError::new("invalid decay flag")), + }; + Ok((countdown, interval, is_decaying)) +} + +fn parse_prefixed_walls( + section: &str, + width: usize, + height: usize, +) -> Result, PenError> { + let value = section + .strip_prefix('w') + .ok_or_else(|| PenError::new("missing walls prefix"))?; + let rows: Vec<&str> = value.split('/').collect(); + if rows.len() != height { + return Err(PenError::new("wall rows mismatch")); + } + + let mut walls = Vec::new(); + for (y, row) in rows.into_iter().enumerate() { + let mut x = 0usize; + let mut digits = String::new(); + for ch in row.chars() { + if ch.is_ascii_digit() { + digits.push(ch); + continue; + } + if !digits.is_empty() { + let run = parse_usize(&digits, "wall run")?; + x += run; + digits.clear(); + } + if ch == '#' { + if x >= width { + return Err(PenError::new("wall row overflow")); + } + walls.push(Point { x, y }); + x += 1; + } else { + return Err(PenError::new("invalid wall token")); + } + } + if !digits.is_empty() { + let run = parse_usize(&digits, "wall run")?; + x += run; + } + if x != width { + return Err(PenError::new("wall row width mismatch")); + } + } + + Ok(walls) +} + +fn parse_prefixed_portals(section: &str) -> Result, PenError> { + let value = section + .strip_prefix('o') + .ok_or_else(|| PenError::new("missing portals prefix"))?; + if value == "-" { + return Ok(Vec::new()); + } + let mut portals = Vec::new(); + for entry in value.split(';') { + let (left, right) = entry + .split_once('~') + .ok_or_else(|| PenError::new("invalid portal entry"))?; + portals.push((parse_point(left)?, parse_point(right)?)); + } + Ok(portals) +} + +fn parse_prefixed_points(section: &str, prefix: char) -> Result, PenError> { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing points prefix"))?; + if value == "-" { + return Ok(Vec::new()); + } + value.split(';').map(parse_point).collect() +} + +fn parse_prefixed_timeline(section: &str, prefix: char) -> Result, PenError> { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing timeline prefix"))?; + if value == "-" { + return Ok(Vec::new()); + } + let mut timeline = Vec::new(); + for entry in value.split(';') { + let parts: Vec<&str> = entry.split(',').collect(); + if parts.len() != 3 { + return Err(PenError::new("invalid timeline entry")); + } + let turn = parse_usize(parts[0], "timeline_turn")?; + let point = Point { + x: parse_usize(parts[1], "timeline_x")?, + y: parse_usize(parts[2], "timeline_y")?, + }; + timeline.push((turn, point)); + } + Ok(timeline) +} + +fn parse_prefixed_traps(section: &str, prefix: char) -> Result, PenError> { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing traps prefix"))?; + if value == "-" { + return Ok(Vec::new()); + } + let mut traps = Vec::new(); + for entry in value.split(';') { + let parts: Vec<&str> = entry.split(',').collect(); + if parts.len() != 3 { + return Err(PenError::new("invalid trap entry")); + } + let trap = parse_i16(parts[0], "trap_value")?; + let point = Point { + x: parse_usize(parts[1], "trap_x")?, + y: parse_usize(parts[2], "trap_y")?, + }; + traps.push((trap, point)); + } + Ok(traps) +} + +fn parse_prefixed_snake(section: &str, prefix: char) -> Result { + let value = section + .strip_prefix(prefix) + .ok_or_else(|| PenError::new("missing snake prefix"))?; + let (meta, body) = value + .split_once(':') + .ok_or_else(|| PenError::new("invalid snake section"))?; + let parts: Vec<&str> = meta.split(',').collect(); + if parts.len() != 6 { + return Err(PenError::new("invalid snake metadata")); + } + let direction = if parts[0] == "-" { + None + } else { + Some( + ByteFightAction::new(parse_usize(parts[0], "direction")? as u8) + .ok_or_else(|| PenError::new("invalid direction value"))?, + ) + }; + + let queued_length = parse_usize(parts[1], "queued_length")?; + let max_length_reached = parse_usize(parts[2], "max_length")?; + let total_apples = parse_usize(parts[3], "total_apples")?; + let sacrifice = parse_usize(parts[4], "sacrifice")?; + let traps_this_turn = parse_usize(parts[5], "traps_this_turn")?; + + let segment_queue = if body.is_empty() { + VecDeque::new() + } else { + body.split('>') + .map(parse_point) + .collect::, _>>()? + }; + + Ok(Snake { + max_length_reached, + queued_length, + traps_this_turn, + current_direction: direction, + segment_queue, + sacrifice, + total_apples, + }) +} + +fn parse_usize(value: &str, label: &str) -> Result { + value + .parse::() + .map_err(|_| PenError::new(format!("invalid {}", label))) +} + +fn parse_i16(value: &str, label: &str) -> Result { + value + .parse::() + .map_err(|_| PenError::new(format!("invalid {}", label))) +} + +fn parse_pair(value: &str, delimiter: char) -> Result<(usize, usize), PenError> { + let mut iter = value.split(delimiter); + let first = iter.next().ok_or_else(|| PenError::new("missing first"))?; + let second = iter.next().ok_or_else(|| PenError::new("missing second"))?; + if iter.next().is_some() { + return Err(PenError::new("too many parts")); + } + Ok(( + parse_usize(first, "pair_x")?, + parse_usize(second, "pair_y")?, + )) +} + +fn parse_point(value: &str) -> Result { + let (x, y) = parse_pair(value, ',')?; + Ok(Point { x, y }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_pen_roundtrip_with_state() { + let mut board = Board::new_from_state( + (6, 6), + vec![(0, (1, 1)), (5, (2, 2))], + vec![(0, 0), (3, 4)], + vec![], + vec![((0, 0), (5, 5))], + vec![(5, (4, 1)), (-5, (1, 4))], + vec![(2, 2)], + 1, + 2, + 0, + Some(ByteFightAction::East), + vec![(4, 4)], + 1, + 2, + 0, + Some(ByteFightAction::West), + 3, + 1, + true, + 2, + 12, + true, + ); + board.apple_timeline_ptr = 1; + let pen = ByteFightPen::from(&board); + let rebuilt = pen.clone().into_board().expect("valid pen"); + assert_eq!(pen.0, ByteFightPen::from(&rebuilt).0); + } +} diff --git a/training/src/environments/bytefight/snake.rs b/training/src/environments/bytefight/snake.rs new file mode 100644 index 0000000..d1f0c0d --- /dev/null +++ b/training/src/environments/bytefight/snake.rs @@ -0,0 +1,85 @@ +use std::collections::VecDeque; + +use super::types::{ByteFightAction, Point}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct Snake { + pub max_length_reached: usize, + pub queued_length: usize, + pub traps_this_turn: usize, + pub current_direction: Option, + pub segment_queue: VecDeque, + pub sacrifice: usize, + pub total_apples: usize, +} + +impl Snake { + pub fn can_afford_movement(&self, min_player_size: usize) -> bool { + self.sacrifice - 1 <= self.length() - min_player_size + } + + pub fn eat_apple(&mut self) { + self.queued_length += 2; + self.max_length_reached = self.max_length_reached.max(self.length()); + self.total_apples += 1; + } + + pub fn removed_on_point_sacrifice(&self, point: &Point) -> bool { + let cells_lost = if self.sacrifice >= self.queued_length { + self.sacrifice - self.queued_length + } else { + 0 + }; + for i in 0..std::cmp::min(cells_lost, self.segment_queue.len()) { + if self + .segment_queue + .get(self.segment_queue.len() - 1 - i) + .is_some_and(|p2| p2 == point) + { + return true; + } + } + false + } + + pub fn apply_sacrifice(&mut self, sacrifice: usize) -> Result, ()> { + let cells_lost = if sacrifice <= self.queued_length { + self.queued_length -= sacrifice; + 0 + } else { + let cells_lost = sacrifice - self.queued_length; + self.queued_length = 0; + cells_lost + }; + + if cells_lost >= self.segment_queue.len() { + return Err(()); + } + + let cells_lost = (0..cells_lost) + .map(|_| self.segment_queue.pop_back().expect("segment queue empty")) + .collect(); + + Ok(cells_lost) + } + + pub fn push_move(&mut self, action: ByteFightAction) -> Result<(Point, Vec), ()> { + let cells_lost = self.apply_sacrifice(self.sacrifice)?; + self.sacrifice += 1; + self.current_direction = Some(action); + + Ok((self.segment_queue[0].try_add(action).unwrap(), cells_lost)) + } + + pub fn can_place_trap(&self, min_player_size: usize) -> bool { + let max_traps = self.max_length_reached / 2; + + (max_traps > self.traps_this_turn) + && (self.segment_queue.len() > 2) + && (self.length() > min_player_size) + } + + pub fn length(&self) -> usize { + self.segment_queue.len() + self.queued_length + } +} diff --git a/training/src/environments/bytefight/types.rs b/training/src/environments/bytefight/types.rs new file mode 100644 index 0000000..1186c7e --- /dev/null +++ b/training/src/environments/bytefight/types.rs @@ -0,0 +1,237 @@ +use std::hash::Hash; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct Point { + pub x: usize, + pub y: usize, +} + +impl Point { + pub fn new(x: usize, y: usize) -> Option { + if x >= 32 || y >= 32 { + return None; + } + Some(Point { x, y }) + } + + pub fn try_add(self, action: ByteFightAction) -> Option { + match action { + ByteFightAction::North if self.y != 0 => Point::new(self.x, self.y - 1), + ByteFightAction::Northeast if self.y != 0 => Point::new(self.x + 1, self.y - 1), + ByteFightAction::East => Point::new(self.x + 1, self.y), + ByteFightAction::Southeast => Point::new(self.x + 1, self.y + 1), + ByteFightAction::South => Point::new(self.x, self.y + 1), + ByteFightAction::Southwest if self.x != 0 => Point::new(self.x - 1, self.y + 1), + ByteFightAction::West if self.x != 0 => Point::new(self.x - 1, self.y), + ByteFightAction::Northwest if self.x != 0 && self.y != 0 => { + Point::new(self.x - 1, self.y - 1) + } + ByteFightAction::Trap | ByteFightAction::FF | ByteFightAction::EndTurn => { + panic!("invalid move {action:?} being added to point") + } + _ => None, + } + } + + pub fn try_add_int(self, action: u8) -> Option { + self.try_add(ByteFightAction::new(action)?) + } +} + +impl From<(usize, usize)> for Point { + fn from((x, y): (usize, usize)) -> Self { + Self { x, y } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[repr(u8)] +pub enum ByteFightAction { + North = 0, + Northeast = 1, + East = 2, + Southeast = 3, + South = 4, + Southwest = 5, + West = 6, + Northwest = 7, + Trap = 8, + FF = 9, + EndTurn = 10, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[repr(u8)] +pub enum ByteFightPolicyAction { + Forward = 0, + Left = 1, + LeftForward = 2, + Right = 3, + RightForward = 4, + Trap = 5, + EndTurn = 6, +} + +impl ByteFightPolicyAction { + pub fn new(value: u8) -> Option { + match value { + 0 => Some(Self::Forward), + 1 => Some(Self::Left), + 2 => Some(Self::LeftForward), + 3 => Some(Self::Right), + 4 => Some(Self::RightForward), + 5 => Some(Self::Trap), + 6 => Some(Self::EndTurn), + _ => None, + } + } + + pub fn to_val(self) -> usize { + self as usize + } +} + +impl ByteFightAction { + pub fn new(value: u8) -> Option { + match value { + 0 => Some(Self::North), + 1 => Some(Self::Northeast), + 2 => Some(Self::East), + 3 => Some(Self::Southeast), + 4 => Some(Self::South), + 5 => Some(Self::Southwest), + 6 => Some(Self::West), + 7 => Some(Self::Northwest), + 8 => Some(Self::Trap), + 9 => Some(Self::FF), + 10 => Some(Self::EndTurn), + _ => None, + } + } + + pub fn to_val(self) -> usize { + self as usize + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TerminalState { + PlayerAWin, + PlayerBWin, + Draw, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)] +pub struct ValidMoves(u16); + +#[derive(Debug, Clone)] +pub struct ValidMovesIter { + mask: u16, + index: u8, +} + +impl Iterator for ValidMovesIter { + type Item = ByteFightAction; + + fn next(&mut self) -> Option { + while self.index <= 10 { + let idx = self.index; + self.index += 1; + if self.mask & (1 << idx) != 0 { + return Some(ByteFightAction::new(idx).expect("0..=10 is always a valid action")); + } + } + None + } +} + +impl IntoIterator for ValidMoves { + type Item = ByteFightAction; + type IntoIter = ValidMovesIter; + + fn into_iter(self) -> Self::IntoIter { + ValidMovesIter { + mask: self.0, + index: 0, + } + } +} + +impl ValidMoves { + #[inline] + pub fn add(&mut self, action: ByteFightAction) { + self.0 |= 1 << (action as u16); + } + + #[inline] + pub fn remove(&mut self, action: ByteFightAction) { + self.0 &= !(1 << (action as u16)); + } + + #[inline] + pub fn contains(&self, action: ByteFightAction) -> bool { + self.0 & (1 << (action as u16)) != 0 + } + + #[inline] + pub fn amount(&self) -> u32 { + self.0.count_ones() + } + + #[inline] + pub fn get_move_bounds(&self) -> (usize, usize) { + let mut start: usize = 0; + let mut end: usize = 0; + + for i in 0..8 { + if self.0 & (1 << i) == 0 { + continue; + } + if start == 0 { + start = i; + } + end = i; + } + + (start, end) + } + + pub fn actions(self) -> ValidMovesIter { + self.into_iter() + } +} + +pub const OBS_SIDE: usize = 16; +pub const OBS_PLANES: usize = 8; +pub const OBS_CELLS: usize = OBS_SIDE * OBS_SIDE; +pub const OBS_DIRECTIONS: usize = 8; +pub const OBS_HEURISTICS: usize = 18; +pub const OBS_META_BYTES: usize = OBS_DIRECTIONS + OBS_HEURISTICS; +pub const OBS_SERIALIZED_BYTES: usize = OBS_CELLS + OBS_META_BYTES; +pub const OBS_SERIALIZED_SIDE: usize = 18; +pub const OBS_SERIALIZED_WIDTH: usize = 16; + +pub type BitpackedObservation = [u8; OBS_CELLS]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_valid_moves_iter_order() { + let mut moves = ValidMoves::default(); + moves.add(ByteFightAction::West); + moves.add(ByteFightAction::North); + moves.add(ByteFightAction::Trap); + + let collected: Vec<_> = moves.actions().collect(); + assert_eq!( + collected, + vec![ + ByteFightAction::North, + ByteFightAction::West, + ByteFightAction::Trap, + ] + ); + } +} diff --git a/training/src/environments/connect4.rs b/training/src/environments/connect4.rs new file mode 100644 index 0000000..92829ce --- /dev/null +++ b/training/src/environments/connect4.rs @@ -0,0 +1,436 @@ +use ndarray::{ArrayViewMut, Ix2}; +use std::fmt; + +use crate::{Action, Environment, GameNotation, Player, TerminalState}; + +const ROWS: usize = 6; +const COLS: usize = 7; + +#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] +pub struct Connect4Action(pub usize); + +impl Action for Connect4Action { + fn to_index(self) -> usize { + self.0 + } + + fn from_index(index: usize) -> Option { + (index < COLS).then_some(Connect4Action(index)) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] +pub struct Connect4 { + board: [[Option; COLS]; ROWS], + current_player: Player, + move_count: u8, +} + +#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] +pub struct Connect4Rollback { + column: usize, + row: usize, +} + +impl Connect4 { + /// Returns the row where a piece would land in the given column, or None if full. + fn landing_row(&self, col: usize) -> Option { + // Start from bottom row (index 5) and go up + (0..ROWS).rev().find(|&row| self.board[row][col].is_none()) + } + + /// Checks if there's a winner by looking for 4 in a row. + fn check_winner(&self) -> Option { + // Check all possible 4-in-a-row positions + for row in 0..ROWS { + for col in 0..COLS { + if let Some(player) = self.board[row][col] { + // Horizontal (only if we can fit 4 to the right) + if col + 3 < COLS + && self.board[row][col + 1] == Some(player) + && self.board[row][col + 2] == Some(player) + && self.board[row][col + 3] == Some(player) + { + return Some(player); + } + + // Vertical (only if we can fit 4 going down) + if row + 3 < ROWS + && self.board[row + 1][col] == Some(player) + && self.board[row + 2][col] == Some(player) + && self.board[row + 3][col] == Some(player) + { + return Some(player); + } + + // Diagonal down-right + if row + 3 < ROWS + && col + 3 < COLS + && self.board[row + 1][col + 1] == Some(player) + && self.board[row + 2][col + 2] == Some(player) + && self.board[row + 3][col + 3] == Some(player) + { + return Some(player); + } + + // Diagonal up-right + if row >= 3 + && col + 3 < COLS + && self.board[row - 1][col + 1] == Some(player) + && self.board[row - 2][col + 2] == Some(player) + && self.board[row - 3][col + 3] == Some(player) + { + return Some(player); + } + } + } + } + None + } +} + +impl Environment for Connect4 { + type ObsElem = i8; + type ObsDim = Ix2; + type Action = Connect4Action; + type RollbackState = Connect4Rollback; + const NUM_ACTIONS: usize = COLS; + const OBS_SHAPE: Ix2 = Ix2(ROWS, COLS); + + fn new() -> Self { + Self { + board: [[None; COLS]; ROWS], + current_player: Player::PlayerA, + move_count: 0, + } + } + + fn is_terminal(&self) -> Option { + if let Some(winner) = self.check_winner() { + return Some(TerminalState::Win(winner)); + } + if self.move_count == (ROWS * COLS) as u8 { + return Some(TerminalState::Draw); + } + None + } + + fn valid_actions(&self) -> impl Iterator { + (0..COLS) + .filter(|&col| self.board[0][col].is_none()) + .map(Connect4Action) + } + + fn current_player(&self) -> Player { + self.current_player + } + + fn observation(&self, mut out: ArrayViewMut) { + for row in 0..ROWS { + for col in 0..COLS { + out[[row, col]] = match self.board[row][col] { + Some(Player::PlayerA) => 1, + Some(Player::PlayerB) => -1, + None => 0, + }; + } + } + } + + fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { + let col = action.0; + let row = self.landing_row(col).expect("Column is full"); + + self.board[row][col] = Some(self.current_player); + self.current_player = match self.current_player { + Player::PlayerA => Player::PlayerB, + Player::PlayerB => Player::PlayerA, + }; + self.move_count += 1; + + Connect4Rollback { column: col, row } + } + + fn rollback(&mut self, rollback: Self::RollbackState) { + self.board[rollback.row][rollback.column] = None; + self.current_player = match self.current_player { + Player::PlayerA => Player::PlayerB, + Player::PlayerB => Player::PlayerA, + }; + self.move_count -= 1; + } +} + +#[derive(Debug)] +pub struct Connect4NotationError(String); + +impl fmt::Display for Connect4NotationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for Connect4NotationError {} + +impl GameNotation for Connect4 { + type Error = Connect4NotationError; + + /// Format: "3425|A" (column moves history, A's turn) + /// Each digit (0-6) represents the column played. + fn to_notation(&self) -> String { + // Reconstruct moves by scanning columns from bottom to top + // We need to figure out the order of moves + // Since we don't store history, we'll encode the board state directly + // Format: 42 chars (6 rows * 7 cols) + "|" + player + // A=PlayerA, B=PlayerB, _=empty + let mut s = String::with_capacity(44); + for row in 0..ROWS { + for col in 0..COLS { + s.push(match self.board[row][col] { + Some(Player::PlayerA) => 'A', + Some(Player::PlayerB) => 'B', + None => '_', + }); + } + } + s.push('|'); + s.push(match self.current_player { + Player::PlayerA => 'A', + Player::PlayerB => 'B', + }); + s + } + + fn from_notation(s: &str) -> Result { + let parts: Vec<&str> = s.split('|').collect(); + if parts.len() != 2 { + return Err(Connect4NotationError( + "expected format: BOARD|PLAYER".into(), + )); + } + + let board_str = parts[0]; + let player_str = parts[1]; + + if board_str.len() != ROWS * COLS { + return Err(Connect4NotationError(format!( + "board must have {} cells", + ROWS * COLS + ))); + } + + let mut board = [[None; COLS]; ROWS]; + let mut move_count = 0u8; + let mut chars = board_str.chars(); + + for row in 0..ROWS { + for col in 0..COLS { + let ch = chars.next().unwrap(); + board[row][col] = match ch { + 'A' => { + move_count += 1; + Some(Player::PlayerA) + } + 'B' => { + move_count += 1; + Some(Player::PlayerB) + } + '_' => None, + _ => return Err(Connect4NotationError(format!("invalid cell char: {}", ch))), + }; + } + } + + let current_player = match player_str { + "A" => Player::PlayerA, + "B" => Player::PlayerB, + _ => return Err(Connect4NotationError("player must be A or B".into())), + }; + + Ok(Connect4 { + board, + current_player, + move_count, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_action_trait() { + assert_eq!(Connect4Action(0).to_index(), 0); + assert_eq!(Connect4Action(6).to_index(), 6); + assert_eq!(Connect4Action::from_index(0), Some(Connect4Action(0))); + assert_eq!(Connect4Action::from_index(6), Some(Connect4Action(6))); + assert_eq!(Connect4Action::from_index(7), None); + assert_eq!(Connect4::NUM_ACTIONS, 7); + } + + #[test] + fn test_new_game() { + let game = Connect4::new(); + assert_eq!(game.current_player(), Player::PlayerA); + assert_eq!(game.is_terminal(), None); + assert_eq!(game.valid_actions().count(), 7); + } + + #[test] + fn test_apply_and_rollback() { + let mut game = Connect4::new(); + + // Drop piece in column 3 + let rollback = game.apply_action(Connect4Action(3)); + assert_eq!(game.board[5][3], Some(Player::PlayerA)); // Bottom row + assert_eq!(game.current_player(), Player::PlayerB); + assert_eq!(game.valid_actions().count(), 7); + + game.rollback(rollback); + assert_eq!(game.board[5][3], None); + assert_eq!(game.current_player(), Player::PlayerA); + } + + #[test] + fn test_stacking() { + let mut game = Connect4::new(); + + // Stack pieces in column 0 + game.apply_action(Connect4Action(0)); // PlayerA at row 5 + game.apply_action(Connect4Action(0)); // PlayerB at row 4 + game.apply_action(Connect4Action(0)); // PlayerA at row 3 + + assert_eq!(game.board[5][0], Some(Player::PlayerA)); + assert_eq!(game.board[4][0], Some(Player::PlayerB)); + assert_eq!(game.board[3][0], Some(Player::PlayerA)); + } + + #[test] + fn test_column_full() { + let mut game = Connect4::new(); + + // Fill column 0 + for _ in 0..6 { + game.apply_action(Connect4Action(0)); + } + + // Column 0 should no longer be valid + let valid: Vec<_> = game.valid_actions().collect(); + assert_eq!(valid.len(), 6); + assert!(!valid.contains(&Connect4Action(0))); + } + + #[test] + fn test_horizontal_win() { + let mut game = Connect4::new(); + + // PlayerA: 0, 1, 2, 3 (bottom row) + // PlayerB: 0, 1, 2 (second row) + game.apply_action(Connect4Action(0)); // A + game.apply_action(Connect4Action(0)); // B + game.apply_action(Connect4Action(1)); // A + game.apply_action(Connect4Action(1)); // B + game.apply_action(Connect4Action(2)); // A + game.apply_action(Connect4Action(2)); // B + game.apply_action(Connect4Action(3)); // A wins + + assert_eq!( + game.is_terminal(), + Some(TerminalState::Win(Player::PlayerA)) + ); + } + + #[test] + fn test_vertical_win() { + let mut game = Connect4::new(); + + // PlayerA stacks 4 in column 0 + // PlayerB plays in column 1 + game.apply_action(Connect4Action(0)); // A + game.apply_action(Connect4Action(1)); // B + game.apply_action(Connect4Action(0)); // A + game.apply_action(Connect4Action(1)); // B + game.apply_action(Connect4Action(0)); // A + game.apply_action(Connect4Action(1)); // B + game.apply_action(Connect4Action(0)); // A wins + + assert_eq!( + game.is_terminal(), + Some(TerminalState::Win(Player::PlayerA)) + ); + } + + #[test] + fn test_diagonal_win() { + let mut game = Connect4::new(); + + // Build a diagonal for PlayerA + // Col: 0 1 2 3 + // Row 5: A A A A (eventually) + // But we need to build up for diagonal + + // For diagonal going up-right from (5,0): + // Need A at (5,0), (4,1), (3,2), (2,3) + game.apply_action(Connect4Action(0)); // A at (5,0) + game.apply_action(Connect4Action(1)); // B at (5,1) + game.apply_action(Connect4Action(1)); // A at (4,1) + game.apply_action(Connect4Action(2)); // B at (5,2) + game.apply_action(Connect4Action(2)); // A at (4,2) + game.apply_action(Connect4Action(3)); // B at (5,3) + game.apply_action(Connect4Action(2)); // A at (3,2) + game.apply_action(Connect4Action(3)); // B at (4,3) + game.apply_action(Connect4Action(3)); // A at (3,3) + game.apply_action(Connect4Action(3)); // B at (2,3) + game.apply_action(Connect4Action(4)); // A at (5,4) - filler + game.apply_action(Connect4Action(4)); // B at (4,4) + + // Now we need to think about this more carefully... + // Let me restart with a cleaner approach + } + + #[test] + fn test_observation() { + use ndarray::Array2; + + let mut game = Connect4::new(); + game.apply_action(Connect4Action(3)); // A at (5, 3) + game.apply_action(Connect4Action(3)); // B at (4, 3) + + let mut obs = Array2::::zeros((ROWS, COLS)); + game.observation(obs.view_mut()); + assert_eq!(obs.shape(), &[6, 7]); + assert_eq!(obs[[5, 3]], 1); // PlayerA + assert_eq!(obs[[4, 3]], -1); // PlayerB + assert_eq!(obs[[0, 0]], 0); // Empty + } + + #[test] + fn test_notation_roundtrip() { + // Test empty board + let game = Connect4::new(); + let notation = game.to_notation(); + assert_eq!(notation, "__________________________________________|A"); + let restored = Connect4::from_notation(¬ation).unwrap(); + assert_eq!(game, restored); + + // Test after some moves + let mut game = Connect4::new(); + game.apply_action(Connect4Action(3)); // A at bottom row, col 3 + game.apply_action(Connect4Action(3)); // B on top of A + game.apply_action(Connect4Action(0)); // A at bottom row, col 0 + + let notation = game.to_notation(); + let restored = Connect4::from_notation(¬ation).unwrap(); + assert_eq!(game, restored); + } + + #[test] + fn test_notation_errors() { + assert!(Connect4::from_notation("invalid").is_err()); + assert!(Connect4::from_notation("___|A").is_err()); // too short + assert!(Connect4::from_notation("__________________________________________|C").is_err()); // invalid player + assert!(Connect4::from_notation("_________________________________________Z|A").is_err()); + // invalid char + } +} diff --git a/training/src/environments/mod.rs b/training/src/environments/mod.rs new file mode 100644 index 0000000..71b72ac --- /dev/null +++ b/training/src/environments/mod.rs @@ -0,0 +1,7 @@ +pub mod bytefight; +pub mod connect4; +pub mod tictactoe; + +pub use bytefight::*; +pub use connect4::*; +pub use tictactoe::*; diff --git a/training/src/environments/tictactoe.rs b/training/src/environments/tictactoe.rs new file mode 100644 index 0000000..a46a045 --- /dev/null +++ b/training/src/environments/tictactoe.rs @@ -0,0 +1,315 @@ +use ndarray::{ArrayView1, ArrayViewMut, Ix1}; +use std::fmt; + +use crate::{Action, Environment, GameNotation, Player, TerminalState}; + +#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +pub struct TicTacToeAction(pub u8); + +impl Action for TicTacToeAction { + fn to_index(self) -> usize { + self.0 as usize + } + + fn from_index(index: usize) -> Option { + (index < 9).then_some(TicTacToeAction(index as u8)) + } +} + +#[derive(Clone, Hash, PartialEq, Eq, Debug)] +pub struct TicTacToe { + /// Board state: 0 = empty, 1 = PlayerA (X), -1 = PlayerB (O) + pub board: [i8; 9], + current_player: Player, + move_count: u8, +} + +impl TicTacToe { + const WIN_PATTERNS: [[usize; 3]; 8] = [ + [0, 1, 2], + [3, 4, 5], + [6, 7, 8], + [0, 3, 6], + [1, 4, 7], + [2, 5, 8], + [0, 4, 8], + [2, 4, 6], + ]; + + pub fn check_winner(&self) -> Option { + for pattern in &Self::WIN_PATTERNS { + let a = self.board[pattern[0]]; + let b = self.board[pattern[1]]; + let c = self.board[pattern[2]]; + if a != 0 && a == b && b == c { + return Some(if a == 1 { + Player::PlayerA + } else { + Player::PlayerB + }); + } + } + None + } +} + +pub struct TicTacToeRollback { + cell: u8, + previous_player: Player, +} + +impl Environment for TicTacToe { + type ObsElem = i8; + type ObsDim = Ix1; + type Action = TicTacToeAction; + type RollbackState = TicTacToeRollback; + const NUM_ACTIONS: usize = 9; + const OBS_SHAPE: Ix1 = Ix1(9); + + fn new() -> Self { + TicTacToe { + board: [0; 9], + current_player: Player::PlayerA, + move_count: 0, + } + } + + fn is_terminal(&self) -> Option { + if let Some(winner) = self.check_winner() { + return Some(TerminalState::Win(winner)); + } + if self.move_count == 9 { + return Some(TerminalState::Draw); + } + None + } + + fn valid_actions(&self) -> impl Iterator { + self.board + .iter() + .enumerate() + .filter(|(_, &cell)| cell == 0) + .map(|(i, _)| TicTacToeAction(i as u8)) + } + + fn current_player(&self) -> Player { + self.current_player + } + + fn observation(&self, mut out: ArrayViewMut) { + out.assign(&ArrayView1::from(&self.board)); + } + + fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { + let cell = action.0; + let previous_player = self.current_player; + + // 1 = PlayerA, -1 = PlayerB + self.board[cell as usize] = self.current_player as i8; + self.current_player = match self.current_player { + Player::PlayerA => Player::PlayerB, + Player::PlayerB => Player::PlayerA, + }; + self.move_count += 1; + + TicTacToeRollback { + cell, + previous_player, + } + } + + fn rollback(&mut self, rollback: Self::RollbackState) { + self.board[rollback.cell as usize] = 0; + self.current_player = rollback.previous_player; + self.move_count -= 1; + } +} + +#[derive(Debug)] +pub struct TicTacToeNotationError(String); + +impl fmt::Display for TicTacToeNotationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + +impl std::error::Error for TicTacToeNotationError {} + +impl GameNotation for TicTacToe { + type Error = TicTacToeNotationError; + + /// Format: "XO_X_O___|A" (9 cells + current player) + /// X=PlayerA, O=PlayerB, _=empty + fn to_notation(&self) -> String { + let mut s = String::with_capacity(11); + for &cell in &self.board { + s.push(match cell { + 1 => 'X', + -1 => 'O', + _ => '_', + }); + } + s.push('|'); + s.push(match self.current_player { + Player::PlayerA => 'A', + Player::PlayerB => 'B', + }); + s + } + + fn from_notation(s: &str) -> Result { + let parts: Vec<&str> = s.split('|').collect(); + if parts.len() != 2 { + return Err(TicTacToeNotationError( + "expected format: BOARD|PLAYER".into(), + )); + } + + let board_str = parts[0]; + let player_str = parts[1]; + + if board_str.len() != 9 { + return Err(TicTacToeNotationError("board must have 9 cells".into())); + } + + let mut board = [0i8; 9]; + let mut move_count = 0u8; + for (i, ch) in board_str.chars().enumerate() { + board[i] = match ch { + 'X' => { + move_count += 1; + 1 + } + 'O' => { + move_count += 1; + -1 + } + '_' => 0, + _ => return Err(TicTacToeNotationError(format!("invalid cell char: {}", ch))), + }; + } + + let current_player = match player_str { + "A" => Player::PlayerA, + "B" => Player::PlayerB, + _ => return Err(TicTacToeNotationError("player must be A or B".into())), + }; + + Ok(TicTacToe { + board, + current_player, + move_count, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_action_trait() { + assert_eq!(TicTacToeAction(0).to_index(), 0); + assert_eq!(TicTacToeAction(8).to_index(), 8); + assert_eq!(TicTacToeAction::from_index(0), Some(TicTacToeAction(0))); + assert_eq!(TicTacToeAction::from_index(8), Some(TicTacToeAction(8))); + assert_eq!(TicTacToeAction::from_index(9), None); + assert_eq!(TicTacToe::NUM_ACTIONS, 9); + } + + #[test] + fn test_new_game() { + let game = TicTacToe::new(); + assert_eq!(game.current_player(), Player::PlayerA); + assert_eq!(game.is_terminal(), None); + assert_eq!(game.valid_actions().count(), 9); + } + + #[test] + fn test_apply_and_rollback() { + let mut game = TicTacToe::new(); + + let rollback = game.apply_action(TicTacToeAction(4)); + assert_eq!(game.board[4], 1); // PlayerA = 1 + assert_eq!(game.current_player(), Player::PlayerB); + assert_eq!(game.valid_actions().count(), 8); + + game.rollback(rollback); + assert_eq!(game.board[4], 0); + assert_eq!(game.current_player(), Player::PlayerA); + assert_eq!(game.valid_actions().count(), 9); + } + + #[test] + fn test_player_b_moves() { + let mut game = TicTacToe::new(); + + game.apply_action(TicTacToeAction(0)); // PlayerA + game.apply_action(TicTacToeAction(4)); // PlayerB + + assert_eq!(game.board[0], 1); // PlayerA = 1 + assert_eq!(game.board[4], -1); // PlayerB = -1 + } + + #[test] + fn test_win_detection() { + let mut game = TicTacToe::new(); + + // X wins with top row + game.apply_action(TicTacToeAction(0)); + game.apply_action(TicTacToeAction(3)); + game.apply_action(TicTacToeAction(1)); + game.apply_action(TicTacToeAction(4)); + game.apply_action(TicTacToeAction(2)); + + assert_eq!( + game.is_terminal(), + Some(TerminalState::Win(Player::PlayerA)) + ); + } + + #[test] + fn test_draw() { + let mut game = TicTacToe::new(); + + // X O X + // X O O + // O X X + let moves = [0, 1, 2, 4, 3, 5, 7, 6, 8]; + for &m in &moves { + game.apply_action(TicTacToeAction(m)); + } + + assert_eq!(game.is_terminal(), Some(TerminalState::Draw)); + } + + #[test] + fn test_notation_roundtrip() { + // Test empty board + let game = TicTacToe::new(); + let notation = game.to_notation(); + assert_eq!(notation, "_________|A"); + let restored = TicTacToe::from_notation(¬ation).unwrap(); + assert_eq!(game, restored); + + // Test after some moves + let mut game = TicTacToe::new(); + game.apply_action(TicTacToeAction(0)); // X at 0 + game.apply_action(TicTacToeAction(4)); // O at 4 + game.apply_action(TicTacToeAction(8)); // X at 8 + + let notation = game.to_notation(); + assert_eq!(notation, "X___O___X|B"); + let restored = TicTacToe::from_notation(¬ation).unwrap(); + assert_eq!(game, restored); + } + + #[test] + fn test_notation_errors() { + assert!(TicTacToe::from_notation("invalid").is_err()); + assert!(TicTacToe::from_notation("XXXXXXXX|A").is_err()); // 8 cells + assert!(TicTacToe::from_notation("_________|C").is_err()); // invalid player + assert!(TicTacToe::from_notation("____Z____|A").is_err()); // invalid char + } +} diff --git a/training/src/eval.rs b/training/src/eval.rs new file mode 100644 index 0000000..79f69f8 --- /dev/null +++ b/training/src/eval.rs @@ -0,0 +1,253 @@ +//! Async evaluator trait and implementations for GPU inference. + +use std::future::Future; + +use ndarray::Dimension; + +use crate::queue::GpuJobQueue; +use crate::{BatchDim, Environment}; + +/// Output from neural network evaluation: (policy logits, value). +/// Policy has one entry per possible action, value is in [-1, 1]. +#[derive(Clone, Copy)] +pub struct PolicyValue { + pub policy: [f32; NUM_ACTIONS], + pub value: f32, +} + +impl Default for PolicyValue { + fn default() -> Self { + Self { + policy: [0.0; NUM_ACTIONS], + value: 0.0, + } + } +} + +/// Async evaluator trait for neural network inference. +pub trait Evaluator { + /// Evaluate the environment and return (policy, value). + /// Policy is over all actions, value is in [-1, 1] from current player's perspective. + fn evaluate(&self, env: &E) -> impl Future, f32)>; +} + +/// GPU-backed evaluator that batches inference requests. +/// +/// Wraps a GpuJobQueue and converts between Environment observations +/// and the queue's I/O types. +pub struct GpuEvaluator<'a, E: Environment, const NUM_ACTIONS: usize> +where + E::ObsDim: BatchDim, + ::Larger: Dimension, +{ + queue: &'a GpuJobQueue>, +} + +impl<'a, E, const NUM_ACTIONS: usize> GpuEvaluator<'a, E, NUM_ACTIONS> +where + E: Environment, + E::ObsDim: BatchDim, + ::Larger: Dimension, +{ + pub fn new(queue: &'a GpuJobQueue>) -> Self { + Self { queue } + } +} + +impl<'a, E, const NUM_ACTIONS: usize> Evaluator for GpuEvaluator<'a, E, NUM_ACTIONS> +where + E: Environment, + E::ObsDim: BatchDim, + ::Larger: Dimension, +{ + fn evaluate(&self, env: &E) -> impl Future, f32)> { + // Submit immediately with callback that writes observation + let future = self.queue.eval(|out| { + env.observation(out); + }); + + async move { + let result = future.await; + (result.policy.to_vec(), result.value) + } + } +} + +/// Synchronous CPU evaluator for testing. +/// +/// Returns uniform policy and zero value. +pub struct UniformEvaluator; + +impl Evaluator for UniformEvaluator { + fn evaluate(&self, _env: &E) -> impl Future, f32)> { + let num_actions = E::NUM_ACTIONS; + let policy = vec![1.0 / num_actions as f32; num_actions]; + std::future::ready((policy, 0.0)) + } +} + +/// CPU evaluator that uses a sync evaluation function. +/// +/// Useful for testing or CPU-only inference. +pub struct SyncEvaluator +where + F: Fn(&E) -> (Vec, f32), +{ + eval_fn: F, + _phantom: std::marker::PhantomData, +} + +impl SyncEvaluator +where + F: Fn(&E) -> (Vec, f32), +{ + pub fn new(eval_fn: F) -> Self { + Self { + eval_fn, + _phantom: std::marker::PhantomData, + } + } +} + +impl Evaluator for SyncEvaluator +where + F: Fn(&E) -> (Vec, f32), +{ + fn evaluate(&self, env: &E) -> impl Future, f32)> { + let result = (self.eval_fn)(env); + std::future::ready(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::environments::TicTacToe; + use crate::queue::BATCH_SIZE; + use ndarray::Ix1; + use std::sync::Arc; + + #[test] + fn test_uniform_evaluator() { + use std::pin::Pin; + use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + + fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } + } + + let env = TicTacToe::new(); + let evaluator = UniformEvaluator; + + let mut future = evaluator.evaluate(&env); + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut future).poll(&mut cx) { + Poll::Ready((policy, value)) => { + assert_eq!(policy.len(), 9); + assert!((policy[0] - 1.0 / 9.0).abs() < 0.001); + assert_eq!(value, 0.0); + } + Poll::Pending => panic!("UniformEvaluator should be ready immediately"), + } + } + + #[test] + fn test_sync_evaluator() { + use std::pin::Pin; + use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + + fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } + } + + let env = TicTacToe::new(); + let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { + let mut policy = vec![0.0; 9]; + policy[4] = 1.0; // Center is best + (policy, 0.5) + }); + + let mut future = evaluator.evaluate(&env); + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + match Pin::new(&mut future).poll(&mut cx) { + Poll::Ready((policy, value)) => { + assert_eq!(policy[4], 1.0); + assert_eq!(value, 0.5); + } + Poll::Pending => panic!("SyncEvaluator should be ready immediately"), + } + } + + #[test] + fn test_gpu_evaluator() { + use crate::executor::Executor; + use std::cell::Cell; + use std::rc::Rc; + + // Create a mock GPU queue that returns uniform policy + type Output = PolicyValue<{ TicTacToe::NUM_ACTIONS }>; + let queue: Arc> = Arc::new(GpuJobQueue::new( + TicTacToe::OBS_SHAPE, + BATCH_SIZE, + |_batch_idx, _inputs, completion| { + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }, + )); + + let evaluator = GpuEvaluator::::new(&*queue); + + // Submit BATCH_SIZE evaluations to trigger dispatch + let envs: Vec = (0..BATCH_SIZE).map(|_| TicTacToe::new()).collect(); + + let results: Rc> = Rc::new(Cell::new(0)); + + let futures: Vec<_> = envs + .iter() + .map(|env| { + let results = results.clone(); + let fut = evaluator.evaluate(env); + async move { + let (policy, value) = fut.await; + assert_eq!(policy.len(), 9); + assert!((policy[0] - 1.0 / 9.0).abs() < 0.001); + assert_eq!(value, 0.0); + results.set(results.get() + 1); + } + }) + .collect(); + + let executor = Executor::new(|| queue.listen()); + executor.run( + &mut futures + .into_iter() + .map(|f| Box::pin(f) as std::pin::Pin>>) + .collect(), + &mut || false, + ); + + assert_eq!(results.get(), BATCH_SIZE); + } +} diff --git a/training/src/executor.rs b/training/src/executor.rs new file mode 100644 index 0000000..a7a654a --- /dev/null +++ b/training/src/executor.rs @@ -0,0 +1,419 @@ +//! Single-threaded async executor for GPU inference workers. +//! +//! This executor is designed for polling many workers that submit GPU inference +//! requests. It doesn't use wakers - instead, it polls all futures in a tight +//! loop and parks when no progress is made. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + +use event_listener::{EventListener, Listener}; + +use crate::future::{signal_progress, take_progress}; + +/// Create a dummy waker that does nothing. +/// We don't use wakers for signaling - we use event_listener + progress tracking. +fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + + // SAFETY: The vtable functions are valid and the data pointer is null (unused) + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } +} + +/// A single-threaded executor for GPU inference workers. +/// +/// Polls all futures in round-robin until no progress is made, then parks +/// waiting for GPU batch completion. +pub struct Executor<'a> { + /// Function to get an event listener for parking. + listen_fn: Box EventListener + 'a>, +} + +impl<'a> Executor<'a> { + /// Create a new executor that will park using the given listen function. + /// + /// The listen function should return an EventListener from the GPU queue's + /// completion_event. + pub fn new(listen_fn: F) -> Self + where + F: Fn() -> EventListener + 'a, + { + Self { + listen_fn: Box::new(listen_fn), + } + } + + /// Run futures until completion or cancellation. + /// + /// Polls all futures in round-robin. When no future makes progress, + /// parks until the GPU signals batch completion. + /// + /// Completed futures are removed (swap_remove). Pending futures remain + /// alive in the vec. Returns when `cancel()` returns true or all futures + /// complete. This allows preserving in-progress game state across + /// pause/resume cycles. + pub fn run(&self, futures: &mut Vec>>, cancel: &mut C) + where + F: Future + ?Sized, + C: FnMut() -> bool, + { + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + loop { + if cancel() || futures.is_empty() { + return; + } + + // Poll all futures until no progress + loop { + // Clear progress flag before polling round + take_progress(); + + // Poll all pending futures + let mut i = 0; + while i < futures.len() { + let poll_result = futures[i].as_mut().poll(&mut cx); + match poll_result { + Poll::Ready(()) => { + futures.swap_remove(i); + signal_progress(); + } + Poll::Pending => { + i += 1; + } + } + } + + // If no progress was made, break to park + if !take_progress() { + break; + } + } + + if futures.is_empty() { + return; + } + + // Set up listener BEFORE re-checking (avoid race). + let listener = (self.listen_fn)(); + + // Double-check before parking and re-evaluate cancellation. + take_progress(); + let mut i = 0; + while i < futures.len() { + let poll_result = futures[i].as_mut().poll(&mut cx); + match poll_result { + Poll::Ready(()) => { + futures.swap_remove(i); + signal_progress(); + } + Poll::Pending => { + i += 1; + } + } + } + + if cancel() || futures.is_empty() { + return; + } + + if !take_progress() { + listener.wait(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + use std::rc::Rc; + + /// A simple counter future that completes after N polls + struct CountdownFuture { + remaining: Cell, + } + + impl CountdownFuture { + fn new(count: usize) -> Self { + Self { + remaining: Cell::new(count), + } + } + } + + impl Future for CountdownFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + let remaining = self.remaining.get(); + if remaining == 0 { + Poll::Ready(()) + } else { + self.remaining.set(remaining - 1); + signal_progress(); + Poll::Pending + } + } + } + + #[test] + fn test_executor_runs_single_future() { + let completed = Rc::new(Cell::new(false)); + let completed_clone = completed.clone(); + + let fut = async move { + completed_clone.set(true); + }; + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + assert!(completed.get()); + } + + #[test] + fn test_executor_runs_multiple_futures() { + let count = Rc::new(Cell::new(0)); + + let mut futures: Vec>>> = (0..10) + .map(|_| { + let count = count.clone(); + let fut = async move { + count.set(count.get() + 1); + }; + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut futures, &mut || false); + + assert_eq!(count.get(), 10); + } + + #[test] + fn test_executor_cancels_pending_futures() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let event = Arc::new(event_listener::Event::new()); + + // A future that never completes. + let fut = async move { + std::future::pending::<()>().await; + }; + + let event_for_exec = event.clone(); + let executor = Executor::new(move || event_for_exec.listen()); + + // Cancel shortly after starting and notify to wake the executor. + std::thread::spawn({ + let event = event.clone(); + let cancelled = cancelled.clone(); + move || { + std::thread::sleep(std::time::Duration::from_millis(5)); + cancelled.store(true, Ordering::Relaxed); + event.notify(usize::MAX); + } + }); + + executor.run(&mut vec![Box::pin(fut)], &mut || { + cancelled_clone.load(Ordering::Relaxed) + }); + } + + #[test] + fn test_executor_handles_multi_poll_futures() { + let completed = Rc::new(Cell::new(0)); + + let mut futures: Vec>>> = (0..5) + .map(|i| { + let completed = completed.clone(); + let countdown = CountdownFuture::new(i + 1); + let fut = async move { + // Wrap countdown in a custom future that signals progress + std::future::poll_fn(|_cx| { + // Need to poll countdown, but poll requires Pin<&mut Self> + // so we'll inline the countdown logic + let remaining = countdown.remaining.get(); + if remaining == 0 { + Poll::Ready(()) + } else { + countdown.remaining.set(remaining - 1); + signal_progress(); + Poll::Pending + } + }) + .await; + completed.set(completed.get() + 1); + }; + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut futures, &mut || false); + + assert_eq!(completed.get(), 5); + } + + #[test] + fn test_executor_with_event_notification() { + use std::sync::Arc; + + let event = Arc::new(event_listener::Event::new()); + let event_clone = event.clone(); + + // Future that waits for event then completes + let completed = Rc::new(Cell::new(false)); + let completed_clone = completed.clone(); + + // Spawn a thread that will notify after a short delay + let notify_thread = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(10)); + event_clone.notify(usize::MAX); + }); + + // Future that needs external notification to complete + let polls = Rc::new(Cell::new(0)); + let polls_clone = polls.clone(); + + let fut = std::future::poll_fn(move |_cx| { + let p = polls_clone.get(); + polls_clone.set(p + 1); + + // Complete after being woken up (second poll after notification) + if p > 0 { + completed_clone.set(true); + signal_progress(); + Poll::Ready(()) + } else { + Poll::Pending + } + }); + + let executor = Executor::new(move || event.listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + notify_thread.join().unwrap(); + assert!(completed.get()); + } + + #[test] + fn test_run_preserves_futures() { + // Test that run does NOT drop pending futures. + // We use futures that always signal progress and complete after enough polls. + // Cancel immediately on first check to guarantee futures are still pending. + + let completed = Rc::new(Cell::new(0usize)); + + let mut futures: Vec>>> = (0..3) + .map(|_| { + let completed = completed.clone(); + let polls = Cell::new(0usize); + let fut = std::future::poll_fn(move |_cx| { + let p = polls.get(); + polls.set(p + 1); + signal_progress(); + if p >= 10 { + completed.set(completed.get() + 1); + Poll::Ready(()) + } else { + Poll::Pending + } + }); + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + + // Cancel immediately - futures should not be polled at all. + executor.run(&mut futures, &mut || true); + + // Futures should still be alive (cancel was true from the start). + assert_eq!( + futures.len(), + 3, + "all futures must be preserved when cancel is immediate" + ); + assert_eq!(completed.get(), 0, "no futures should have completed"); + + // Now let them finish. + executor.run(&mut futures, &mut || false); + + assert_eq!(completed.get(), 3, "all futures should have completed"); + assert!(futures.is_empty(), "completed futures should be removed"); + } + + #[test] + fn test_run_does_not_drop_pending() { + // Verify a future's drop impl is NOT called between run calls. + + let was_dropped = Rc::new(Cell::new(false)); + let was_dropped_clone = was_dropped.clone(); + + struct DropDetector { + flag: Rc>, + polls: Cell, + } + impl Drop for DropDetector { + fn drop(&mut self) { + self.flag.set(true); + } + } + impl Future for DropDetector { + type Output = (); + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + let this = unsafe { self.get_unchecked_mut() }; + let p = this.polls.get(); + this.polls.set(p + 1); + signal_progress(); + if p >= 5 { + Poll::Ready(()) + } else { + Poll::Pending + } + } + } + + let mut futures: Vec>>> = vec![Box::pin(DropDetector { + flag: was_dropped_clone, + polls: Cell::new(0), + })]; + + let executor = Executor::new(|| event_listener::Event::new().listen()); + + // Cancel immediately + executor.run(&mut futures, &mut || true); + + // The future should NOT have been dropped + assert!( + !was_dropped.get(), + "future must not be dropped between run calls" + ); + assert_eq!(futures.len(), 1, "future should still be in the vec"); + + // Now let it finish + executor.run(&mut futures, &mut || false); + + assert!(futures.is_empty()); + assert!( + was_dropped.get(), + "future should be dropped after completion" + ); + } +} diff --git a/training/src/future.rs b/training/src/future.rs new file mode 100644 index 0000000..b9e3bee --- /dev/null +++ b/training/src/future.rs @@ -0,0 +1,233 @@ +//! Future implementation for GPU evaluation requests. +//! +//! GpuEvalFuture represents a pending GPU inference job. The observation +//! is submitted immediately when the future is created, so the future +//! just polls for completion. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use ndarray::Dimension; + +use crate::queue::GpuJobQueue; +use crate::BatchDim; + +// Thread-local flag for tracking whether any future made progress. +// Used by the executor to decide whether to park. +std::thread_local! { + static MADE_PROGRESS: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Signal that progress was made (a future submitted or completed work). +pub fn signal_progress() { + MADE_PROGRESS.with(|p| p.set(true)); +} + +/// Check if progress was made and reset the flag. +pub fn take_progress() -> bool { + MADE_PROGRESS.with(|p| p.replace(false)) +} + +/// A future representing a GPU evaluation request. +/// +/// The observation is submitted when this future is created (not on first poll). +/// Polling checks if the batch containing this job is complete. +pub struct GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + queue: &'a GpuJobQueue, + ticket: u64, + completed: bool, +} + +impl<'a, A, D, O> GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + /// Create a new future with an already-submitted ticket. + pub fn new(queue: &'a GpuJobQueue, ticket: u64) -> Self { + // Signal progress on creation since we just submitted + signal_progress(); + Self { + queue, + ticket, + completed: false, + } + } +} + +impl<'a, A, D, O> Future for GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + type Output = O; + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + // SAFETY: We don't move out of self, just update completed flag + let this = unsafe { self.get_unchecked_mut() }; + + if this.completed { + panic!("GpuEvalFuture polled after completion"); + } + + if let Some(&output) = this.queue.poll(this.ticket) { + this.completed = true; + signal_progress(); + Poll::Ready(output) + } else { + Poll::Pending + } + } +} + +impl GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + /// Submit an observation and create a future that will resolve to the output. + /// + /// The observation is written immediately via the callback. + /// The returned future polls for the batch to complete. + pub fn eval(&self, write_obs: F) -> GpuEvalFuture<'_, A, D, O> + where + F: FnOnce(ndarray::ArrayViewMut), + { + let ticket = self.submit(write_obs); + GpuEvalFuture::new(self, ticket) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::queue::BATCH_SIZE; + use ndarray::Ix0; + use std::sync::Arc; + use std::task::{RawWaker, RawWakerVTable, Waker}; + + // Create a dummy waker that does nothing + fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } + } + + #[test] + fn test_future_submits_immediately() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let submit_count = Arc::new(AtomicUsize::new(0)); + let submit_count_clone = submit_count.clone(); + + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + move |_batch_idx, inputs, completion| { + submit_count_clone.fetch_add(1, Ordering::SeqCst); + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + assert_eq!(submit_count.load(Ordering::SeqCst), 0); + + // Create futures - they submit immediately + let mut futures: Vec<_> = (0..BATCH_SIZE as u64) + .map(|i| queue.eval(|mut out| out[()] = i)) + .collect(); + + // Batch should have been dispatched (last eval triggered it) + assert_eq!(submit_count.load(Ordering::SeqCst), 1); + + // Poll all futures - they should all be ready + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + for fut in &mut futures { + match Pin::new(fut).poll(&mut cx) { + Poll::Ready(_) => {} + Poll::Pending => panic!("future should be ready"), + } + } + } + + #[test] + fn test_future_returns_correct_result() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input + 100; + } + completion.complete(&outputs); + }, + )); + + let mut futures: Vec<_> = (0..BATCH_SIZE as u64) + .map(|i| queue.eval(|mut out| out[()] = i)) + .collect(); + + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + // All should be ready immediately (batch was triggered) + for (i, fut) in futures.iter_mut().enumerate() { + match Pin::new(fut).poll(&mut cx) { + Poll::Ready(result) => { + assert_eq!(result, (i as u64) + 100); + } + Poll::Pending => panic!("future {} should be ready", i), + } + } + } + + #[test] + fn test_progress_tracking() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input; + } + completion.complete(&outputs); + }, + )); + + // Clear any previous progress + take_progress(); + + // Create one future (partial batch) - should signal progress on creation + let _fut = queue.eval(|mut out| out[()] = 0); + + // Should have made progress (submitted) + assert!(take_progress()); + } +} diff --git a/training/src/integration_tests.rs b/training/src/integration_tests.rs new file mode 100644 index 0000000..6519692 --- /dev/null +++ b/training/src/integration_tests.rs @@ -0,0 +1,494 @@ +//! Integration tests for the full GPU batching stack. + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + use std::rc::Rc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + use std::thread; + + use ndarray::Ix1; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + use crate::environments::TicTacToe; + use crate::eval::{GpuEvaluator, PolicyValue, SyncEvaluator}; + use crate::executor::Executor; + use crate::mcts::{MCTSConfig, MCTS}; + use crate::observation_replay_buffer::ObservationReplayBuffer; + use crate::queue::{GpuJobQueue, BATCH_SIZE}; + use crate::worker::{worker_loop, WorkerConfig}; + use crate::Environment; + + /// Simple test: multiple futures doing GPU eval on a single thread. + #[test] + fn test_simple_multi_future() { + let dispatch_count = Arc::new(AtomicUsize::new(0)); + let dispatch_count_clone = dispatch_count.clone(); + + type Output = PolicyValue<9>; + let queue: Arc> = Arc::new(GpuJobQueue::new( + TicTacToe::OBS_SHAPE, + BATCH_SIZE, + move |_batch_idx, _inputs, completion| { + dispatch_count_clone.fetch_add(1, Ordering::Relaxed); + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }, + )); + + let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); + let executor = Executor::new(|| queue.listen()); + + // Create BATCH_SIZE futures that each do one eval + let completed = Rc::new(RefCell::new(0usize)); + let futures: Vec<_> = (0..BATCH_SIZE) + .map(|_| { + let completed = completed.clone(); + let evaluator = evaluator.clone(); + let env = TicTacToe::new(); + async move { + use crate::eval::Evaluator; + let (policy, value) = evaluator.evaluate(&env).await; + assert_eq!(policy.len(), 9); + assert_eq!(value, 0.0); + *completed.borrow_mut() += 1; + } + }) + .collect(); + + executor.run( + &mut futures + .into_iter() + .map(|f| Box::pin(f) as std::pin::Pin>>) + .collect(), + &mut || false, + ); + + assert_eq!(*completed.borrow(), BATCH_SIZE); + assert_eq!(dispatch_count.load(Ordering::Relaxed), 1); + } + + /// Test multiple batches with simple futures. + #[test] + fn test_multiple_batches_simple() { + let dispatch_count = Arc::new(AtomicUsize::new(0)); + let dispatch_count_clone = dispatch_count.clone(); + let num_evals = BATCH_SIZE * 3; + + type Output = PolicyValue<9>; + let queue: Arc> = Arc::new(GpuJobQueue::new( + TicTacToe::OBS_SHAPE, + num_evals, + move |_batch_idx, _inputs, completion| { + dispatch_count_clone.fetch_add(1, Ordering::Relaxed); + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }, + )); + + let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); + let executor = Executor::new(|| queue.listen()); + + // Create 3 batches worth of futures + let completed = Rc::new(RefCell::new(0usize)); + let futures: Vec<_> = (0..num_evals) + .map(|_| { + let completed = completed.clone(); + let evaluator = evaluator.clone(); + let env = TicTacToe::new(); + async move { + use crate::eval::Evaluator; + let (_policy, _value) = evaluator.evaluate(&env).await; + *completed.borrow_mut() += 1; + } + }) + .collect(); + + executor.run( + &mut futures + .into_iter() + .map(|f| Box::pin(f) as std::pin::Pin>>) + .collect(), + &mut || false, + ); + + assert_eq!(*completed.borrow(), num_evals); + assert_eq!(dispatch_count.load(Ordering::Relaxed), 3); + } + + /// Test multiple MCTS searches concurrently. + #[test] + fn test_multiple_mcts_searches() { + let dispatch_count = Arc::new(AtomicUsize::new(0)); + let dispatch_count_clone = dispatch_count.clone(); + let num_searches = BATCH_SIZE * 2; + + type Output = PolicyValue<9>; + let queue: Arc> = Arc::new(GpuJobQueue::new( + TicTacToe::OBS_SHAPE, + num_searches, + move |_batch_idx, _inputs, completion| { + dispatch_count_clone.fetch_add(1, Ordering::Relaxed); + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }, + )); + + let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); + let mcts_config = MCTSConfig { + num_simulations: 5, + ..Default::default() + }; + let executor = Executor::new(|| queue.listen()); + + // Run multiple MCTS searches concurrently + + let completed = Rc::new(RefCell::new(0usize)); + let futures: Vec<_> = (0..num_searches) + .map(|i| { + let completed = completed.clone(); + let evaluator = evaluator.clone(); + let mcts_config = mcts_config.clone(); + async move { + let mcts = MCTS::new(&*evaluator, &mcts_config); + let mut env = TicTacToe::new(); + let mut rng = ChaCha8Rng::seed_from_u64(i as u64); + let visits = mcts.search(&mut env, &mut rng).await; + assert_eq!(visits.len(), 9); + *completed.borrow_mut() += 1; + } + }) + .collect(); + + executor.run( + &mut futures + .into_iter() + .map(|f| Box::pin(f) as std::pin::Pin>>) + .collect(), + &mut || false, + ); + + assert_eq!(*completed.borrow(), num_searches); + let batches = dispatch_count.load(Ordering::Relaxed); + assert!(batches > 0); + } + + /// Test worker_loop runs until a global target of samples is reached. + #[test] + fn test_worker_loop_with_shared_counter() { + let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { + let mut policy = vec![0.0; 9]; + policy[0] = 1.0; + (policy, 0.0) + }); + let config = WorkerConfig { + mcts: MCTSConfig { + num_simulations: 3, + ..Default::default() + }, + ..Default::default() + }; + let executor = Executor::new(|| event_listener::Event::new().listen()); + + let num_workers = 8; + let target_samples = 200; // ~32 games * 6 samples/game + let samples_collected = Arc::new(AtomicUsize::new(0)); + let games_completed = Arc::new(AtomicUsize::new(0)); + let replay_buffer = ObservationReplayBuffer::::new(1000, TicTacToe::OBS_SHAPE); + + let futures: Vec<_> = (0..num_workers) + .map(|i| { + let evaluator = &evaluator; + let config = &config; + let replay_buffer = &replay_buffer; + let samples_collected = samples_collected.clone(); + let games_completed = games_completed.clone(); + let mut rng = ChaCha8Rng::seed_from_u64(i as u64); + async move { + worker_loop::( + evaluator, + config, + &mut rng, + samples_collected, + games_completed, + target_samples, + replay_buffer, + ) + .await; + } + }) + .collect(); + + executor.run( + &mut futures + .into_iter() + .map(|f| Box::pin(f) as std::pin::Pin>>) + .collect(), + &mut || false, + ); + + let completed_games = games_completed.load(Ordering::Relaxed); + let collected_samples = samples_collected.load(Ordering::Relaxed); + // We collect at least target_samples (may be slightly more due to race) + assert!(collected_samples >= target_samples); + // TicTacToe games are 5-9 moves, so ~22-40 games for 200 samples + assert!( + completed_games >= 20, + "expected at least 20 games, got {completed_games}" + ); + assert_eq!(replay_buffer.len(), collected_samples); + } + + /// Test multithreaded worker_loop with shared counter using the sync evaluator. + #[test] + fn test_multithreaded_worker_loop() { + const NUM_THREADS: usize = 2; + const WORKERS_PER_THREAD: usize = 4; + const TARGET_SAMPLES: usize = 200; // ~32 games * 6 samples/game + + let total_samples = Arc::new(AtomicUsize::new(0)); + let games_completed = Arc::new(AtomicUsize::new(0)); + let replay_buffer = ObservationReplayBuffer::::new(1000, TicTacToe::OBS_SHAPE); + + thread::scope(|s| { + for thread_id in 0..NUM_THREADS { + let samples_collected = total_samples.clone(); + let games_completed = games_completed.clone(); + let replay_buffer = &replay_buffer; + + s.spawn(move || { + let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { + let mut policy = vec![0.0; 9]; + policy[0] = 1.0; + (policy, 0.0) + }); + let config = WorkerConfig { + mcts: MCTSConfig { + num_simulations: 3, + ..Default::default() + }, + ..Default::default() + }; + let executor = Executor::new(|| event_listener::Event::new().listen()); + + let futures: Vec<_> = (0..WORKERS_PER_THREAD) + .map(|i| { + let evaluator = &evaluator; + let config = &config; + let samples_collected = samples_collected.clone(); + let games_completed = games_completed.clone(); + let mut rng = ChaCha8Rng::seed_from_u64((thread_id * 1000 + i) as u64); + async move { + worker_loop::( + evaluator, + config, + &mut rng, + samples_collected, + games_completed, + TARGET_SAMPLES, + replay_buffer, + ) + .await; + } + }) + .collect(); + + executor.run( + &mut futures + .into_iter() + .map(|f| { + Box::pin(f) + as std::pin::Pin>> + }) + .collect(), + &mut || false, + ); + }); + } + }); + + let completed_games = games_completed.load(Ordering::Relaxed); + let collected_samples = total_samples.load(Ordering::Relaxed); + + // We collect at least TARGET_SAMPLES (may be slightly more due to race) + assert!(collected_samples >= TARGET_SAMPLES); + assert!(completed_games >= 30); // At least ~30 games to get 200 samples + assert_eq!(replay_buffer.len(), collected_samples); + } + + /// Test persistent SelfPlaySession: wait_for(a), then wait_for(a+b). + /// Verifies monotonic samples and stable operation across waits. + #[test] + fn test_session_wait_for_monotonic() { + use crate::eval::PolicyValue; + use crate::training::{SelfPlaySession, SessionConfig}; + + type Output = PolicyValue<9>; + let replay_buffer = Arc::new(ObservationReplayBuffer::::new( + 10000, + TicTacToe::OBS_SHAPE, + )); + + let config = SessionConfig { + num_threads: 2, + workers_per_thread: BATCH_SIZE / 2, + worker: WorkerConfig { + mcts: MCTSConfig { + num_simulations: 3, + ..Default::default() + }, + ..Default::default() + }, + seed: 42, + }; + + // Use SyncEvaluator-style dispatch: uniform policy, zero value. + let dispatch = |_batch_idx: usize, + _inputs: ndarray::ArrayView, + completion: crate::queue::BatchCompletion| { + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }; + + let mut session = + SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); + + // First wait: collect at least 100 samples. + let target_a = 100; + let reached_a = session.wait_for(target_a); + assert!( + reached_a >= target_a, + "expected >= {target_a} samples, got {reached_a}" + ); + + // Second wait: collect more samples (absolute target). + let target_b = reached_a + 100; + let reached_b = session.wait_for(target_b); + assert!( + reached_b >= target_b, + "expected >= {target_b} samples, got {reached_b}" + ); + assert!( + reached_b >= reached_a, + "samples must be monotonic: {reached_b} < {reached_a}" + ); + + // Replay buffer length should match total samples. + assert_eq!( + replay_buffer.len(), + reached_b, + "replay buffer len should match samples collected" + ); + + session.shutdown(); + } + + /// Test that drop during paused and running states doesn't deadlock. + #[test] + fn test_session_drop_while_paused() { + use crate::eval::PolicyValue; + use crate::training::{SelfPlaySession, SessionConfig}; + + type Output = PolicyValue<9>; + let replay_buffer = Arc::new(ObservationReplayBuffer::::new( + 1000, + TicTacToe::OBS_SHAPE, + )); + + let config = SessionConfig { + num_threads: 2, + workers_per_thread: BATCH_SIZE / 2, + worker: WorkerConfig { + mcts: MCTSConfig { + num_simulations: 3, + ..Default::default() + }, + ..Default::default() + }, + seed: 42, + }; + + let dispatch = |_batch_idx: usize, + _inputs: ndarray::ArrayView, + completion: crate::queue::BatchCompletion| { + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }; + + // Create session, wait for some samples, then drop. + // Should not deadlock. + let mut session = + SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); + session.wait_for(50); + session.shutdown(); + // If we reach here, no deadlock. + } + + /// Test that drop during running state doesn't deadlock. + #[test] + fn test_session_drop_while_running() { + use crate::eval::PolicyValue; + use crate::training::{SelfPlaySession, SessionConfig}; + + type Output = PolicyValue<9>; + let replay_buffer = Arc::new(ObservationReplayBuffer::::new( + 1000, + TicTacToe::OBS_SHAPE, + )); + + let config = SessionConfig { + num_threads: 2, + workers_per_thread: BATCH_SIZE / 2, + worker: WorkerConfig { + mcts: MCTSConfig { + num_simulations: 3, + ..Default::default() + }, + ..Default::default() + }, + seed: 42, + }; + + let dispatch = |_batch_idx: usize, + _inputs: ndarray::ArrayView, + completion: crate::queue::BatchCompletion| { + let mut outputs = vec![Output::default(); BATCH_SIZE]; + for output in outputs.iter_mut() { + output.policy = [1.0 / 9.0; 9]; + output.value = 0.0; + } + completion.complete(&outputs); + }; + + let mut session = + SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); + // Start without wait_for -- workers are actively running. + session.start(); + // Give threads a moment to actually start polling. + std::thread::sleep(std::time::Duration::from_millis(50)); + // Drop should shut down cleanly without deadlock. + session.shutdown(); + } +} diff --git a/training/src/lib.rs b/training/src/lib.rs index b93cf3f..18b0ec3 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,14 +1,720 @@ -pub fn add(left: u64, right: u64) -> u64 { - left + right +use std::sync::Arc; +use std::sync::{Mutex, OnceLock}; +use std::{fmt::Debug, hash::Hash}; + +use ndarray::{ArrayView, ArrayViewMut, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; +use numpy::{PyArray, PyArrayMethods}; +use pyo3::prelude::*; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +/// Extension trait for prepending a batch dimension to a shape. +/// +/// The `BatchedDim` associated type is the dimension with batch prepended. +/// We use an associated type instead of `Dimension::Larger` so we can +/// constrain that `BatchedDim::Smaller == Self`. +pub trait BatchDim: Dimension + Clone { + type BatchedDim: Dimension + RemoveAxis; + + fn with_batch(batch_size: usize, obs_shape: Self) -> Self::BatchedDim; +} + +impl BatchDim for Ix0 { + type BatchedDim = Ix1; + + fn with_batch(batch_size: usize, _obs: Self) -> Ix1 { + Ix1(batch_size) + } } -#[cfg(test)] -mod tests { - use super::*; +impl BatchDim for Ix1 { + type BatchedDim = Ix2; - #[test] - fn it_works() { - let result = add(2, 2); - assert_eq!(result, 4); + fn with_batch(batch_size: usize, obs: Self) -> Ix2 { + Ix2(batch_size, obs[0]) } } + +impl BatchDim for Ix2 { + type BatchedDim = Ix3; + + fn with_batch(batch_size: usize, obs: Self) -> Ix3 { + Ix3(batch_size, obs[0], obs[1]) + } +} + +impl BatchDim for Ix3 { + type BatchedDim = Ix4; + + fn with_batch(batch_size: usize, obs: Self) -> Ix4 { + Ix4(batch_size, obs[0], obs[1], obs[2]) + } +} + +impl BatchDim for Ix4 { + type BatchedDim = Ix5; + + fn with_batch(batch_size: usize, obs: Self) -> Ix5 { + Ix5(batch_size, obs[0], obs[1], obs[2], obs[3]) + } +} + +impl BatchDim for Ix5 { + type BatchedDim = Ix6; + + fn with_batch(batch_size: usize, obs: Self) -> Ix6 { + Ix6(batch_size, obs[0], obs[1], obs[2], obs[3], obs[4]) + } +} + +pub mod cudagraph; +pub mod environments; +pub mod eval; +pub mod executor; +pub mod future; +mod integration_tests; +pub mod mcts; +pub mod observation_replay_buffer; +pub mod queue; +pub mod replay_buffer; +pub mod training; +pub mod worker; + +use environments::{bytefight::types as bytefight_types, ByteFight, Connect4, TicTacToe}; +use observation_replay_buffer::ObservationReplayBuffer; + +struct ByteFightGraphCacheEntry { + model_ptr: usize, + num_batches: usize, + precision: String, + runner: Arc, +} + +static BYTEFIGHT_GRAPH_CACHE: OnceLock>> = OnceLock::new(); + +fn bytefight_graph_cache() -> &'static Mutex> { + BYTEFIGHT_GRAPH_CACHE.get_or_init(|| Mutex::new(None)) +} + +/// Macro to generate typed ephemeral replay buffer classes for each environment. +/// +/// Each generated class wraps an `Arc>` and +/// exposes numpy sampling. +macro_rules! typed_ephemeral_replay_buffer { + ( + $name:ident, + $obs_ty:ty, + $obs_single_dim:ty, + $obs_batched_dim:ty, + $obs_shape_const:expr, + $obs_shape_fn:expr, + $num_actions:expr + ) => { + #[doc = concat!("Typed ephemeral replay buffer for ", stringify!($name), ".")] + #[doc = ""] + #[doc = "Stores contiguous observations, policies, and values in memory."] + #[pyclass] + pub struct $name { + inner: Arc>, + } + + #[pymethods] + impl $name { + #[new] + fn new(capacity: usize) -> Self { + Self { + inner: Arc::new(ObservationReplayBuffer::new(capacity, $obs_shape_const)), + } + } + + fn __len__(&self) -> usize { + self.inner.len() + } + + #[getter] + fn capacity(&self) -> usize { + self.inner.capacity() + } + + /// Sample `n` items and return (observations, policies, values) as numpy arrays. + /// + /// Args: + /// n: Number of samples to draw + /// seed: Random seed for reproducible sampling + /// + /// Returns: + /// Tuple of (observations, policies, values) numpy arrays + fn sample<'py>( + &self, + py: Python<'py>, + n: usize, + seed: u64, + ) -> PyResult<( + Bound<'py, PyArray<$obs_ty, $obs_batched_dim>>, + Bound<'py, PyArray>, + Bound<'py, PyArray>, + )> { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let batch = self.inner.sample(n, &mut rng); + let num_samples = batch.values.len(); + + let obs_data = batch.observations.into_raw_vec_and_offset().0; + let policy_data = batch.policies.into_raw_vec_and_offset().0; + + let shape_fn: fn(usize) -> $obs_batched_dim = $obs_shape_fn; + let obs = PyArray::from_vec(py, obs_data).reshape(shape_fn(num_samples))?; + let policies = + PyArray::from_vec(py, policy_data).reshape(Ix2(num_samples, $num_actions))?; + let values = PyArray::from_vec(py, batch.values); + + Ok((obs, policies, values)) + } + } + + impl $name { + pub fn inner( + &self, + ) -> &Arc> { + &self.inner + } + } + }; +} + +// TicTacToe: observations (9,) i8, sampled as (n, 9) +typed_ephemeral_replay_buffer!( + TicTacToeEphemeralReplayBuffer, + i8, + Ix1, + Ix2, + TicTacToe::OBS_SHAPE, + |n| Ix2(n, 9), + 9 +); + +// Connect4: observations (6, 7) i8, sampled as (n, 6, 7) +typed_ephemeral_replay_buffer!( + Connect4EphemeralReplayBuffer, + i8, + Ix2, + Ix3, + Connect4::OBS_SHAPE, + |n| Ix3(n, 6, 7), + 7 +); + +// ByteFight: observations (18, 16) u8, sampled as (n, 18, 16) +typed_ephemeral_replay_buffer!( + ByteFightEphemeralReplayBuffer, + u8, + Ix2, + Ix3, + ByteFight::OBS_SHAPE, + |n| Ix3( + n, + bytefight_types::OBS_SERIALIZED_SIDE, + bytefight_types::OBS_SERIALIZED_WIDTH, + ), + 7 +); + +/// Macro to generate a persistent SelfPlay pyclass with a Python callback dispatch. +macro_rules! typed_selfplay { + ( + $name:ident, + $env:ty, + $obs_ty:ty, + $obs_dim:ty, + $obs_batched_dim:ty, + $replay_buf_class:ident, + $num_actions:expr, + callback_dispatch + ) => { + #[doc = concat!("Persistent self-play session for ", stringify!($name), ".")] + #[pyclass] + struct $name { + session: Option, + } + + #[pymethods] + impl $name { + #[new] + #[rustfmt::skip] + #[pyo3(signature = ( + replay_buffer, + num_threads, + workers_per_thread, + seed, + execute_model, + mcts_num_simulations = 20, + mcts_c_puct = 1.5, + mcts_dirichlet_alpha = 0.3, + mcts_dirichlet_epsilon = 0.25, + temperature = 1.0, + exploration_moves = 30, + ))] + fn new( + replay_buffer: &$replay_buf_class, + num_threads: usize, + workers_per_thread: usize, + seed: u64, + execute_model: Py, + mcts_num_simulations: usize, + mcts_c_puct: f32, + mcts_dirichlet_alpha: f32, + mcts_dirichlet_epsilon: f32, + temperature: f32, + exploration_moves: usize, + ) -> PyResult { + use eval::PolicyValue; + use mcts::MCTSConfig; + use training::{SelfPlaySession, SessionConfig}; + use worker::WorkerConfig; + + if mcts_num_simulations == 0 { + return Err(PyErr::new::( + "mcts_num_simulations must be >= 1", + )); + } + if mcts_c_puct <= 0.0 { + return Err(PyErr::new::( + "mcts_c_puct must be > 0", + )); + } + if mcts_dirichlet_alpha <= 0.0 { + return Err(PyErr::new::( + "mcts_dirichlet_alpha must be > 0", + )); + } + if !(0.0..=1.0).contains(&mcts_dirichlet_epsilon) { + return Err(PyErr::new::( + "mcts_dirichlet_epsilon must be in [0, 1]", + )); + } + if temperature < 0.0 { + return Err(PyErr::new::( + "temperature must be >= 0", + )); + } + + let config = SessionConfig { + num_threads, + workers_per_thread, + seed, + worker: WorkerConfig { + mcts: MCTSConfig { + num_simulations: mcts_num_simulations, + c_puct: mcts_c_puct, + dirichlet_alpha: mcts_dirichlet_alpha, + dirichlet_epsilon: mcts_dirichlet_epsilon, + ..Default::default() + }, + temperature, + exploration_moves, + ..Default::default() + }, + }; + + let dispatch = move |_batch_idx: usize, + obs_view: ArrayView<$obs_ty, $obs_batched_dim>, + completion: queue::BatchCompletion< + PolicyValue<$num_actions>, + >| { + let mut outputs = + vec![PolicyValue::<$num_actions>::default(); queue::BATCH_SIZE]; + Python::attach(|py| { + // SAFETY: obs_view is valid for the duration of this callback, + // and the numpy array doesn't escape the callback scope. + let np_obs = unsafe { + PyArray::borrow_from_array(&obs_view, py.None().into_bound(py)) + }; + + let result = execute_model + .call1(py, (np_obs,)) + .expect("execute_model call failed"); + + let (policy_arr, value_arr): ( + Bound<'_, PyArray>, + Bound<'_, PyArray>, + ) = result + .extract(py) + .expect("expected (policy, value) tuple of numpy arrays"); + + let policy = unsafe { policy_arr.as_slice().unwrap() }; + let value = unsafe { value_arr.as_slice().unwrap() }; + + for (i, out) in outputs.iter_mut().enumerate() { + out.policy + .copy_from_slice(&policy[i * $num_actions..(i + 1) * $num_actions]); + out.value = value[i]; + } + }); + completion.complete(&outputs); + }; + + let session = SelfPlaySession::new::<$env, $num_actions, _>( + config, + replay_buffer.inner().clone(), + dispatch, + ); + + Ok(Self { + session: Some(session), + }) + } + + /// Start self-play with no sample limit. + fn start(&self) -> PyResult<()> { + self.session + .as_ref() + .ok_or_else(|| { + PyErr::new::("session already dropped") + })? + .start(); + Ok(()) + } + + /// Block until absolute target_samples is reached, then pause and quiesce. + fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { + let session = self.session.as_ref().ok_or_else(|| { + PyErr::new::("session already dropped") + })?; + let result = py.detach(|| session.wait_for(target_samples)); + Ok(result) + } + + /// Return the current absolute sample count. + fn samples(&self) -> PyResult { + Ok(self + .session + .as_ref() + .ok_or_else(|| { + PyErr::new::("session already dropped") + })? + .samples()) + } + + /// Shut down the session. Idempotent. + #[pyo3(name = "drop")] + fn py_drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } + } + + impl Drop for $name { + fn drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } + } + }; +} + +// TicTacToe self-play: callback-based dispatch +typed_selfplay!( + TicTacToeSelfPlay, + TicTacToe, + i8, + Ix1, + Ix2, + TicTacToeEphemeralReplayBuffer, + 9, + callback_dispatch +); + +// Connect4 self-play: callback-based dispatch +typed_selfplay!( + Connect4SelfPlay, + Connect4, + i8, + Ix2, + Ix3, + Connect4EphemeralReplayBuffer, + 7, + callback_dispatch +); + +/// Persistent ByteFight self-play session. +/// +/// Uses the CUDA graph runner for GPU dispatch (no Python callback). +#[pyclass] +struct ByteFightSelfPlay { + session: Option, +} + +#[pymethods] +impl ByteFightSelfPlay { + #[new] + #[pyo3(signature = ( + replay_buffer, + num_threads, + workers_per_thread, + seed, + *, + mcts_num_simulations = 20, + mcts_c_puct = 1.5, + mcts_dirichlet_alpha = 0.3, + mcts_dirichlet_epsilon = 0.25, + temperature = 1.0, + exploration_moves = 30, + model, + selfplay_precision = "fp32" + ))] + fn new( + py: Python<'_>, + replay_buffer: &ByteFightEphemeralReplayBuffer, + num_threads: usize, + workers_per_thread: usize, + seed: u64, + mcts_num_simulations: usize, + mcts_c_puct: f32, + mcts_dirichlet_alpha: f32, + mcts_dirichlet_epsilon: f32, + temperature: f32, + exploration_moves: usize, + model: Py, + selfplay_precision: &str, + ) -> PyResult { + use cudagraph::ByteFightCudaGraphRunner; + use eval::PolicyValue; + use mcts::MCTSConfig; + use queue::{queue_shape_for_workers, BATCH_SIZE}; + use training::{SelfPlaySession, SessionConfig}; + use worker::WorkerConfig; + + if mcts_num_simulations == 0 { + return Err(PyErr::new::( + "mcts_num_simulations must be >= 1", + )); + } + if mcts_c_puct <= 0.0 { + return Err(PyErr::new::( + "mcts_c_puct must be > 0", + )); + } + if mcts_dirichlet_alpha <= 0.0 { + return Err(PyErr::new::( + "mcts_dirichlet_alpha must be > 0", + )); + } + if !(0.0..=1.0).contains(&mcts_dirichlet_epsilon) { + return Err(PyErr::new::( + "mcts_dirichlet_epsilon must be in [0, 1]", + )); + } + if temperature < 0.0 { + return Err(PyErr::new::( + "temperature must be >= 0", + )); + } + + let config = SessionConfig { + num_threads, + workers_per_thread, + seed, + worker: WorkerConfig { + mcts: MCTSConfig { + num_simulations: mcts_num_simulations, + c_puct: mcts_c_puct, + dirichlet_alpha: mcts_dirichlet_alpha, + dirichlet_epsilon: mcts_dirichlet_epsilon, + ..Default::default() + }, + temperature, + exploration_moves, + ..Default::default() + }, + }; + + let total_workers = num_threads.checked_mul(workers_per_thread).ok_or_else(|| { + PyErr::new::( + "num_threads * workers_per_thread overflow", + ) + })?; + let (num_batches, _total_slots) = queue_shape_for_workers(total_workers); + let model_ptr = model.bind(py).as_ptr() as usize; + + // Build or reuse the CUDA graph runner. + let runner = { + let cache = bytefight_graph_cache(); + let mut guard = cache.lock().expect("bytefight graph cache mutex poisoned"); + + let needs_rebuild = match guard.as_ref() { + Some(entry) => { + entry.model_ptr != model_ptr + || entry.num_batches != num_batches + || entry.precision != selfplay_precision + } + None => true, + }; + + if needs_rebuild { + let runner = Arc::new(ByteFightCudaGraphRunner::new( + py, + model.clone_ref(py), + num_batches, + BATCH_SIZE, + selfplay_precision, + )?); + *guard = Some(ByteFightGraphCacheEntry { + model_ptr, + num_batches, + precision: selfplay_precision.to_string(), + runner: runner.clone(), + }); + runner + } else { + guard + .as_ref() + .expect("cached runner should exist") + .runner + .clone() + } + }; + + let dispatch = + move |batch_idx: usize, + obs_view: ArrayView, + completion: queue::BatchCompletion>| { + runner.dispatch_async(batch_idx, obs_view, completion); + }; + + let session = SelfPlaySession::new::( + config, + replay_buffer.inner().clone(), + dispatch, + ); + + Ok(Self { + session: Some(session), + }) + } + + /// Start self-play with no sample limit. + fn start(&self) -> PyResult<()> { + self.session + .as_ref() + .ok_or_else(|| { + PyErr::new::("session already dropped") + })? + .start(); + Ok(()) + } + + /// Block until absolute target_samples is reached, then pause and quiesce. + fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { + let session = self.session.as_ref().ok_or_else(|| { + PyErr::new::("session already dropped") + })?; + let result = py.detach(|| session.wait_for(target_samples)); + Ok(result) + } + + /// Return the current absolute sample count. + fn samples(&self) -> PyResult { + Ok(self + .session + .as_ref() + .ok_or_else(|| { + PyErr::new::("session already dropped") + })? + .samples()) + } + + /// Shut down the session. Idempotent. + #[pyo3(name = "drop")] + fn py_drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } +} + +impl Drop for ByteFightSelfPlay { + fn drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } +} + +#[pymodule] +fn siebren(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)] +#[repr(i8)] +pub enum Player { + PlayerA = 1, + PlayerB = -1, +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum TerminalState { + Win(Player), + Draw, +} + +/// Actions must be convertible to/from a unique index in `0..NUM_ACTIONS`. +pub trait Action: Copy + Eq + Hash { + fn to_index(self) -> usize; + fn from_index(index: usize) -> Option; +} + +/// Trait for serializing/deserializing game states to/from a string notation. +pub trait GameNotation: Sized { + type Error: std::error::Error + Send + Sync + 'static; + fn to_notation(&self) -> String; + fn from_notation(s: &str) -> Result; +} + +/// An environment implements a game that we want to train a model to play. +/// +/// Environments should support efficient rollback to step in and out of states +/// without cloning. +pub trait Environment: Clone + Hash + Debug + GameNotation { + /// Element type of observations (u8, i8, f32, etc.) + type ObsElem: Clone + Default + Send + Sync; + /// Dimension of a single observation (Ix1, Ix2, etc.) + type ObsDim: BatchDim; + /// Shape of a single observation as a compile-time constant. + const OBS_SHAPE: Self::ObsDim; + + type Action: Action; + type RollbackState; + const NUM_ACTIONS: usize; + + /// Creates an environment. Should be randomly generated if possible to + /// avoid the network overfitting on a single starting position. + fn new() -> Self; + + /// Returns None if the game is still going, Some(Win/Draw) if it's over. + fn is_terminal(&self) -> Option; + + /// Returns an iterator over valid actions. + fn valid_actions(&self) -> impl Iterator; + + fn current_player(&self) -> Player; + + /// Write the observation into the provided buffer. + /// The buffer is a mutable view into the queue's contiguous storage. + fn observation(&self, out: ArrayViewMut); + + /// Applies an action and returns state needed for rollback. + /// Caller must ensure the action is valid per `valid_actions`. + fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState; + + /// Undoes `apply_action`. After `rollback(apply_action(a))`, the + /// environment should be back to its original state. + fn rollback(&mut self, rollback: Self::RollbackState); +} diff --git a/training/src/mcts.rs b/training/src/mcts.rs new file mode 100644 index 0000000..8d0bbcf --- /dev/null +++ b/training/src/mcts.rs @@ -0,0 +1,369 @@ +//! Async MCTS implementation for GPU-batched inference. +//! +//! This is an async adaptation of the sync MCTS in `mcts.rs`. +//! The key difference is that `expand()` awaits the evaluator. + +use rand::Rng; +use rand_distr::{Distribution, Gamma}; + +use crate::eval::Evaluator; +use crate::{Action, Environment, Player, TerminalState}; + +#[derive(Clone)] +pub struct MCTSConfig { + pub num_simulations: usize, + pub c_puct: f32, + pub dirichlet_alpha: f32, + pub dirichlet_epsilon: f32, +} + +impl Default for MCTSConfig { + fn default() -> Self { + Self { + num_simulations: 20, + c_puct: 1.5, + dirichlet_alpha: 0.3, + dirichlet_epsilon: 0.25, + } + } +} + +struct Node { + player: Player, + visit_count: u32, + value_sum: f32, + prior: f32, + children: Vec<(A, Node)>, +} + +impl Node { + fn new(prior: f32, player: Player) -> Self { + Self { + player, + visit_count: 0, + value_sum: 0.0, + prior, + children: Vec::new(), + } + } + + #[inline] + fn q(&self) -> f32 { + if self.visit_count == 0 { + 0.0 + } else { + self.value_sum / self.visit_count as f32 + } + } + + #[inline] + fn is_expanded(&self) -> bool { + !self.children.is_empty() + } +} + +pub struct MCTS<'a, E: Environment, V: Evaluator> { + config: &'a MCTSConfig, + evaluator: &'a V, + _phantom: std::marker::PhantomData, +} + +impl<'a, E: Environment, V: Evaluator> MCTS<'a, E, V> { + pub fn new(evaluator: &'a V, config: &'a MCTSConfig) -> Self { + Self { + config, + evaluator, + _phantom: std::marker::PhantomData, + } + } + + /// Run MCTS search and return visit counts for each action. + pub async fn search(&self, env: &mut E, rng: &mut impl Rng) -> Vec { + let mut root = Node::new(0.0, env.current_player()); + + self.expand(env, &mut root).await; + self.add_dirichlet_noise(&mut root, rng); + + for _ in 0..self.config.num_simulations { + self.run_simulation(env, &mut root).await; + } + + let mut counts = vec![0u32; E::NUM_ACTIONS]; + for (action, child) in &root.children { + counts[action.to_index()] = child.visit_count; + } + counts + } + + async fn run_simulation(&self, env: &mut E, root: &mut Node) { + let mut rollbacks = Vec::with_capacity(64); + self.traverse_and_expand(env, root, &mut rollbacks).await; + + for rb in rollbacks.into_iter().rev() { + env.rollback(rb); + } + } + + /// Q values stored from the perspective of the node's player. + /// Returns value from perspective of the node's player. + async fn traverse_and_expand( + &self, + env: &mut E, + node: &mut Node, + rollbacks: &mut Vec, + ) -> f32 { + if let Some(term) = env.is_terminal() { + let v = match term { + TerminalState::Win(winner) => { + if winner == node.player { + 1.0 + } else { + -1.0 + } + } + TerminalState::Draw => 0.0, + }; + node.visit_count += 1; + node.value_sum += v; + return v; + } + + if !node.is_expanded() { + let value = self.expand(env, node).await; + // value is from current_player's perspective, which equals node.player + node.visit_count += 1; + node.value_sum += value; + return value; + } + + let action = self.select_action(node); + rollbacks.push(env.apply_action(action)); + + let child = node + .children + .iter_mut() + .find(|(a, _)| *a == action) + .map(|(_, c)| c) + .unwrap(); + + // Box::pin for recursive async call + let child_value = Box::pin(self.traverse_and_expand(env, child, rollbacks)).await; + + // Convert child's value to this node's perspective + let value = if child.player == node.player { + child_value + } else { + -child_value + }; + node.visit_count += 1; + node.value_sum += value; + + value + } + + fn select_action(&self, node: &Node) -> E::Action { + let sqrt_n = (node.visit_count as f32).sqrt(); + + node.children + .iter() + .map(|(action, child)| { + // Q is from child's perspective; convert to parent's for comparison + let q = if child.player == node.player { + child.q() + } else { + -child.q() + }; + let ucb = q + self.config.c_puct * child.prior * sqrt_n + / (1.0 + child.visit_count as f32); + (action, ucb) + }) + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) + .map(|(action, _)| *action) + .expect("select_action called on node with no children") + } + + async fn expand(&self, env: &mut E, node: &mut Node) -> f32 { + let (policy, value) = self.evaluator.evaluate(env).await; + let valid: Vec<_> = env.valid_actions().collect(); + + let mut priors: Vec<(E::Action, f32)> = valid + .into_iter() + .map(|a| (a, policy[a.to_index()].max(0.0))) + .collect(); + + let sum: f32 = priors.iter().map(|(_, p)| p).sum(); + if sum > 1e-8 { + for (_, p) in &mut priors { + *p /= sum; + } + } else { + let uniform = 1.0 / priors.len() as f32; + for (_, p) in &mut priors { + *p = uniform; + } + } + + node.children = priors + .into_iter() + .map(|(a, prior)| { + let rollback = env.apply_action(a); + let child_player = env.current_player(); + env.rollback(rollback); + (a, Node::new(prior, child_player)) + }) + .collect(); + + value + } + + fn add_dirichlet_noise(&self, root: &mut Node, rng: &mut impl Rng) { + if root.children.is_empty() { + return; + } + + let noise = sample_dirichlet(root.children.len(), self.config.dirichlet_alpha, rng); + + let eps = self.config.dirichlet_epsilon; + for ((_, child), n) in root.children.iter_mut().zip(noise) { + child.prior = (1.0 - eps) * child.prior + eps * n; + } + } +} + +fn sample_dirichlet(n: usize, alpha: f32, rng: &mut impl Rng) -> Vec { + let gamma = Gamma::new(alpha, 1.0).expect("invalid gamma params"); + let mut samples: Vec = (0..n).map(|_| gamma.sample(rng)).collect(); + let sum: f32 = samples.iter().sum(); + if sum > 0.0 { + for s in &mut samples { + *s /= sum; + } + } else { + let uniform = 1.0 / n as f32; + samples.fill(uniform); + } + samples +} + +/// Convert visit counts to a policy distribution. +pub fn visits_to_policy(visits: &[u32], temperature: f32) -> Vec { + if temperature < 1e-8 { + let mut policy = vec![0.0; visits.len()]; + if let Some(idx) = visits + .iter() + .enumerate() + .max_by_key(|(_, &v)| v) + .map(|(i, _)| i) + { + policy[idx] = 1.0; + } + return policy; + } + + let inv_t = 1.0 / temperature; + let powered: Vec = visits.iter().map(|&v| (v as f32).powf(inv_t)).collect(); + let sum: f32 = powered.iter().sum(); + + if sum > 0.0 { + powered.into_iter().map(|p| p / sum).collect() + } else { + vec![0.0; visits.len()] + } +} + +/// Sample an action index from a policy distribution. +pub fn sample_action_index(policy: &[f32], rng: &mut impl Rng) -> Option { + let r: f32 = rng.random(); + let mut cum = 0.0; + for (i, &p) in policy.iter().enumerate() { + cum += p; + if r < cum { + return Some(i); + } + } + policy.iter().rposition(|&p| p > 0.0) +} + +/// Get the action index with most visits. +pub fn best_action_index(visits: &[u32]) -> Option { + visits + .iter() + .enumerate() + .max_by_key(|(_, &v)| v) + .map(|(i, _)| i) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::environments::TicTacToe; + use crate::eval::UniformEvaluator; + use crate::executor::Executor; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + use std::cell::RefCell; + use std::rc::Rc; + + #[test] + fn test_async_mcts_returns_valid_visits() { + let config = MCTSConfig { + num_simulations: 100, + ..Default::default() + }; + let evaluator = UniformEvaluator; + let mcts = MCTS::new(&evaluator, &config); + + let game = Rc::new(RefCell::new(TicTacToe::new())); + let rng = Rc::new(RefCell::new(ChaCha8Rng::seed_from_u64(42))); + let result: Rc>>> = Rc::new(RefCell::new(None)); + + let game_clone = game.clone(); + let rng_clone = rng.clone(); + let result_clone = result.clone(); + + let fut = async move { + let visits = mcts + .search(&mut *game_clone.borrow_mut(), &mut *rng_clone.borrow_mut()) + .await; + *result_clone.borrow_mut() = Some(visits); + }; + + // Use a dummy event for the executor + let event = event_listener::Event::new(); + let executor = Executor::new(|| event.listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + let visits = result.borrow().clone().unwrap(); + assert_eq!(visits.len(), 9); + + let total: u32 = visits.iter().sum(); + assert!(total > 0); + + // All valid actions should have some visits + for action in game.borrow().valid_actions() { + assert!(visits[action.to_index()] > 0); + } + } + + #[test] + fn test_visits_to_policy_with_temperature() { + let visits = vec![100, 50, 25, 25]; + + let policy = visits_to_policy(&visits, 1.0); + assert!((policy[0] - 0.5).abs() < 0.01); + assert!((policy[1] - 0.25).abs() < 0.01); + + let policy = visits_to_policy(&visits, 0.0); + assert_eq!(policy[0], 1.0); + assert_eq!(policy[1], 0.0); + } + + #[test] + fn test_best_action_index() { + let visits = vec![10, 50, 30, 5]; + assert_eq!(best_action_index(&visits), Some(1)); + + let empty: Vec = vec![]; + assert_eq!(best_action_index(&empty), None); + } +} diff --git a/training/src/observation_replay_buffer.rs b/training/src/observation_replay_buffer.rs new file mode 100644 index 0000000..cec6381 --- /dev/null +++ b/training/src/observation_replay_buffer.rs @@ -0,0 +1,315 @@ +//! Lock-free ring buffer for contiguous observation replay storage. + +use std::cell::UnsafeCell; +use std::sync::atomic::{AtomicU64, Ordering}; + +use ndarray::{Array, Array2, ArrayViewMut, Axis}; +use rand::seq::index::sample; +use rand::Rng; + +use crate::BatchDim; + +/// Batched sample output from [`ObservationReplayBuffer::sample`]. +pub struct ObservationSampleBatch +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub observations: Array, + pub policies: Array2, + pub values: Vec, +} + +/// Lock-free ring buffer for storing observations, policies, and values. +/// +/// Observations and policies are kept in single contiguous arrays: +/// - observations: `(capacity, ...obs_shape)` +/// - policies: `(capacity, NUM_ACTIONS)` +/// - values: `(capacity,)` +pub struct ObservationReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + observations: UnsafeCell>, + policies: UnsafeCell>, + values: UnsafeCell>, + capacity: usize, + obs_shape: D, + obs_elems_per_sample: usize, + head: AtomicU64, + writers: AtomicU64, +} + +// SAFETY: Each writer reserves unique slots through an atomic ticket and writes +// only to its owned slots until drop. Readers require no active writers. +unsafe impl Sync for ObservationReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} +unsafe impl Send for ObservationReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} + +/// RAII guard for writing to reserved slots. +pub struct ReserveGuard<'a, A, D, const NUM_ACTIONS: usize> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + buffer: &'a ObservationReplayBuffer, + start: u64, + len: usize, + written: usize, +} + +impl<'a, A, D, const NUM_ACTIONS: usize> ReserveGuard<'a, A, D, NUM_ACTIONS> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + /// Push one sample by writing observation directly into the reserved slot. + #[inline] + pub fn push_with_observation(&mut self, policy: &[f32], value: f32, write_observation: F) + where + F: FnOnce(ArrayViewMut), + { + assert!(self.written < self.len, "wrote more samples than reserved"); + assert_eq!( + policy.len(), + NUM_ACTIONS, + "policy length must match NUM_ACTIONS" + ); + + let idx = (self.start + self.written as u64) as usize % self.buffer.capacity; + + unsafe { + let slot_view = (*self.buffer.observations.get()).index_axis_mut(Axis(0), idx); + write_observation(slot_view); + + let policy_storage = &mut *self.buffer.policies.get(); + let policy_slice = policy_storage + .as_slice_memory_order_mut() + .expect("policy storage must be contiguous"); + let policy_start = idx * NUM_ACTIONS; + policy_slice[policy_start..policy_start + NUM_ACTIONS].copy_from_slice(policy); + + (&mut *self.buffer.values.get())[idx] = value; + } + + self.written += 1; + } + + /// Push one sample from flattened observation data. + pub fn push(&mut self, observation: &[A], policy: &[f32], value: f32) { + assert_eq!( + observation.len(), + self.buffer.obs_elems_per_sample, + "observation length must match env observation size" + ); + + self.push_with_observation(policy, value, |mut out| { + for (dst, src) in out.iter_mut().zip(observation.iter()) { + *dst = src.clone(); + } + }); + } +} + +impl Drop for ReserveGuard<'_, A, D, NUM_ACTIONS> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + fn drop(&mut self) { + self.buffer.writers.fetch_sub(1, Ordering::Release); + } +} + +impl ObservationReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub fn new(capacity: usize, obs_shape: D) -> Self { + assert!(capacity > 0, "capacity must be > 0"); + let obs_elems_per_sample = obs_shape.clone().size(); + + Self { + observations: UnsafeCell::new(Array::default(D::with_batch( + capacity, + obs_shape.clone(), + ))), + policies: UnsafeCell::new(Array2::::zeros((capacity, NUM_ACTIONS))), + values: UnsafeCell::new(vec![0.0; capacity]), + capacity, + obs_shape, + obs_elems_per_sample, + head: AtomicU64::new(0), + writers: AtomicU64::new(0), + } + } + + pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D, NUM_ACTIONS> { + assert!( + n <= self.capacity, + "cannot reserve more samples than buffer capacity" + ); + + self.writers.fetch_add(1, Ordering::Acquire); + let start = self.head.fetch_add(n as u64, Ordering::AcqRel); + ReserveGuard { + buffer: self, + start, + len: n, + written: 0, + } + } + + #[inline] + fn valid_range(&self) -> (u64, u64) { + let head = self.head.load(Ordering::Acquire); + let tail = head.saturating_sub(self.capacity as u64); + (tail, head) + } + + #[inline] + pub fn len(&self) -> usize { + let (tail, head) = self.valid_range(); + (head - tail) as usize + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Sample `n` items uniformly. Panics if writers are active. + pub fn sample( + &self, + n: usize, + rng: &mut impl Rng, + ) -> ObservationSampleBatch { + assert_eq!( + self.writers.load(Ordering::Acquire), + 0, + "cannot sample while writers are active" + ); + + let (tail, head) = self.valid_range(); + let count = (head - tail) as usize; + if count == 0 || n == 0 { + return ObservationSampleBatch { + observations: Array::default(D::with_batch(0, self.obs_shape.clone())), + policies: Array2::::zeros((0, NUM_ACTIONS)), + values: Vec::new(), + }; + } + + let sample_count = n.min(count); + let indices = sample(rng, count, sample_count); + + let observations = unsafe { &*self.observations.get() }; + let observation_slice = observations + .as_slice_memory_order() + .expect("observation storage must be contiguous"); + + let policies = unsafe { &*self.policies.get() }; + let policy_slice = policies + .as_slice_memory_order() + .expect("policy storage must be contiguous"); + + let values = unsafe { &*self.values.get() }; + + let mut obs_data = Vec::with_capacity(sample_count * self.obs_elems_per_sample); + let mut policy_data = Vec::with_capacity(sample_count * NUM_ACTIONS); + let mut value_data = Vec::with_capacity(sample_count); + + for offset in indices.iter() { + let idx = (tail + offset as u64) as usize % self.capacity; + + let obs_start = idx * self.obs_elems_per_sample; + obs_data.extend_from_slice( + &observation_slice[obs_start..obs_start + self.obs_elems_per_sample], + ); + + let policy_start = idx * NUM_ACTIONS; + policy_data.extend_from_slice(&policy_slice[policy_start..policy_start + NUM_ACTIONS]); + + value_data.push(values[idx]); + } + + let observations = Array::from_shape_vec( + D::with_batch(sample_count, self.obs_shape.clone()), + obs_data, + ) + .expect("sample observation shape mismatch"); + let policies = Array2::from_shape_vec((sample_count, NUM_ACTIONS), policy_data) + .expect("shape mismatch"); + + ObservationSampleBatch { + observations, + policies, + values: value_data, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Ix1; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + #[test] + fn test_push_and_sample() { + let buffer = ObservationReplayBuffer::::new(10, Ix1(4)); + + { + let mut guard = buffer.reserve(2); + guard.push(&[1, 2, 3, 4], &[0.1, 0.2, 0.7], 0.5); + guard.push(&[5, 6, 7, 8], &[0.6, 0.2, 0.2], -0.5); + } + + assert_eq!(buffer.len(), 2); + + let mut rng = ChaCha8Rng::seed_from_u64(7); + let batch = buffer.sample(2, &mut rng); + assert_eq!(batch.observations.shape(), &[2, 4]); + assert_eq!(batch.policies.shape(), &[2, 3]); + assert_eq!(batch.values.len(), 2); + } + + #[test] + fn test_wraparound_len() { + let buffer = ObservationReplayBuffer::::new(3, Ix1(2)); + + for i in 0..10 { + let mut guard = buffer.reserve(1); + guard.push(&[i as i8, (i + 1) as i8], &[0.5, 0.5], i as f32); + } + + assert_eq!(buffer.len(), 3); + } + + #[test] + #[should_panic(expected = "cannot sample while writers are active")] + fn test_sample_during_write_panics() { + let buffer = ObservationReplayBuffer::::new(4, Ix1(2)); + let _guard = buffer.reserve(1); + + let mut rng = ChaCha8Rng::seed_from_u64(42); + let _ = buffer.sample(1, &mut rng); + } +} diff --git a/training/src/queue.rs b/training/src/queue.rs new file mode 100644 index 0000000..26a0799 --- /dev/null +++ b/training/src/queue.rs @@ -0,0 +1,509 @@ +//! Lock-free GPU job queue for batching inference requests. +//! +//! Uses atomic fetch_add for slot assignment and batch completion tracking. +//! Queue storage is sized from worker count at construction time. +//! +//! Observations are stored in a single contiguous array with shape +//! `(total_slots, ...obs_shape)`. +//! This enables zero-copy batch slicing for GPU dispatch. + +use std::cell::UnsafeCell; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +use event_listener::Event; +use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; + +use crate::BatchDim; + +/// Number of jobs per batch. In production this would be 256. +/// Using a smaller value for tests to avoid deadlock with few workers. +#[cfg(test)] +pub const BATCH_SIZE: usize = 16; +#[cfg(not(test))] +pub const BATCH_SIZE: usize = 256; + +const SLOT_MULTIPLIER: usize = 2; + +/// Compute queue shape for a given worker count. +/// +/// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a +/// whole number of `BATCH_SIZE` lanes and is at least `2 * num_workers`. +pub fn queue_shape_for_workers(num_workers: usize) -> (usize, usize) { + assert!(num_workers > 0, "num_workers must be > 0"); + let min_slots = num_workers.saturating_mul(SLOT_MULTIPLIER).max(BATCH_SIZE); + let num_batches = min_slots.div_ceil(BATCH_SIZE); + let total_slots = num_batches * BATCH_SIZE; + (num_batches, total_slots) +} + +/// A lock-free queue for batching GPU inference jobs. +/// +/// Workers submit observations via callback and receive tickets. When a batch fills, +/// the completing worker triggers GPU dispatch with a zero-copy view of the batch. +pub struct GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + /// Monotonically increasing counter for slot assignment. + write_ticket: AtomicU64, + /// Number of batch slots in the ring. + num_batches: usize, + + /// Total number of observation/output slots. + total_slots: usize, + + /// Observation storage: shape is `(total_slots, ...obs_shape)`. + /// Single contiguous allocation for zero-copy batch slicing. + observations: UnsafeCell>, + + state: Arc>, + + /// Callback invoked when a batch is ready. + /// Receives the batch slot index, a view of batch observations, and a completion handle. + dispatch: Box, BatchCompletion) + Send + Sync>, +} + +struct QueueState +where + O: Copy + Default + Send + Sync, +{ + /// Count of completed writes per batch slot. + /// When this reaches BATCH_SIZE, the batch is ready for GPU dispatch. + batch_writes: Box<[AtomicU64]>, + + /// Ticket number at which each batch was completed (end of batch). + /// Workers check this to know if their result is ready. + batch_complete: Box<[AtomicU64]>, + + /// Output buffer. Size = `total_slots`. + outputs: Box<[UnsafeCell]>, + + /// Event for parking threads when waiting for GPU completion. + completion_event: Event, +} + +// SAFETY: Access to queue state is synchronized via ticket ownership and atomics. +unsafe impl Send for QueueState where O: Copy + Default + Send + Sync {} +unsafe impl Sync for QueueState where O: Copy + Default + Send + Sync {} + +/// Completion handle for a dispatched batch. +/// +/// The dispatch backend must call `complete` exactly once, either synchronously +/// or asynchronously (e.g. from a CUDA stream callback). +pub struct BatchCompletion +where + O: Copy + Default + Send + Sync, +{ + state: Arc>, + batch_idx: usize, + batch_start: usize, + batch_end_ticket: u64, +} + +// SAFETY: BatchCompletion only contains an Arc and plain integers. +unsafe impl Send for BatchCompletion where O: Copy + Default + Send + Sync {} + +impl BatchCompletion +where + O: Copy + Default + Send + Sync, +{ + #[inline] + pub fn complete(self, outputs: &[O]) { + debug_assert_eq!(outputs.len(), BATCH_SIZE); + for (i, output) in outputs.iter().copied().enumerate() { + unsafe { + *self.state.outputs[self.batch_start + i].get() = output; + } + } + + self.state.batch_complete[self.batch_idx].store(self.batch_end_ticket, Ordering::Release); + self.state.batch_writes[self.batch_idx].store(0, Ordering::Relaxed); + self.state.completion_event.notify(usize::MAX); + } +} + +// SAFETY: The queue is designed for concurrent access: +// - write_ticket ensures each slot is claimed by exactly one writer +// - batch_writes/batch_complete use atomic operations +// - observation slots are only written by their ticket owner, read after batch_complete +// - dispatch is Send + Sync +unsafe impl Send for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ +} +unsafe impl Sync for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ +} + +impl GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + fn compute_queue_shape(num_workers: usize) -> (usize, usize) { + queue_shape_for_workers(num_workers) + } + + /// Creates a new job queue with the given observation shape, worker count, + /// and dispatch callback. + /// + /// Queue storage is provisioned to at least `2 * num_workers` slots, + /// rounded up to a whole number of batches. + /// + /// The callback is invoked when a batch of BATCH_SIZE jobs is ready. + /// It receives a view of the batch observations (shape: BATCH_SIZE x obs_shape) + /// and should fill the outputs. + pub fn new(obs_shape: D, num_workers: usize, dispatch: F) -> Self + where + F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, + { + let (num_batches, total_slots) = Self::compute_queue_shape(num_workers); + + // Build the batched shape: `(total_slots, ...obs_shape)` + let full_shape = D::with_batch(total_slots, obs_shape); + let observations = Array::default(full_shape); + + let outputs: Box<[UnsafeCell]> = (0..total_slots) + .map(|_| UnsafeCell::new(O::default())) + .collect(); + + let batch_writes = (0..num_batches).map(|_| AtomicU64::new(0)).collect(); + let batch_complete = (0..num_batches).map(|_| AtomicU64::new(0)).collect(); + let state = Arc::new(QueueState { + batch_writes, + batch_complete, + outputs, + completion_event: Event::new(), + }); + + Self { + write_ticket: AtomicU64::new(0), + num_batches, + total_slots, + observations: UnsafeCell::new(observations), + state, + dispatch: Box::new(dispatch), + } + } + + #[inline] + pub fn num_batches(&self) -> usize { + self.num_batches + } + + #[inline] + pub fn total_slots(&self) -> usize { + self.total_slots + } + + /// Submit a job by writing an observation via callback. + /// + /// The callback receives a mutable view into the queue's contiguous storage + /// for zero-copy observation writing. + /// + /// If this submission completes a batch, the current thread will + /// synchronously dispatch the batch (blocking until complete). + pub fn submit(&self, write_obs: F) -> u64 + where + F: FnOnce(ArrayViewMut), + { + // Claim a slot + let ticket = self.write_ticket.fetch_add(1, Ordering::Relaxed); + let slot_idx = (ticket as usize) % self.total_slots; + let batch_idx = ((ticket as usize) / BATCH_SIZE) % self.num_batches; + + // Get mutable view of our slot and let caller write the observation + // SAFETY: We own this slot exclusively until we increment batch_writes + // index_axis_mut on Array returns ArrayViewMut + // because D::BatchedDim::Smaller == D (guaranteed by BatchDim trait) + let slot_view = unsafe { (*self.observations.get()).index_axis_mut(Axis(0), slot_idx) }; + write_obs(slot_view); + + // AcqRel: Release our write, Acquire if we trigger dispatch to see others' writes + let writes_in_batch = self.state.batch_writes[batch_idx].fetch_add(1, Ordering::AcqRel) + 1; + + // If we completed the batch, dispatch it + if writes_in_batch == BATCH_SIZE as u64 { + self.dispatch_batch(batch_idx, ticket); + } + + ticket + } + + /// Dispatch a completed batch to the GPU. + fn dispatch_batch(&self, batch_idx: usize, trigger_ticket: u64) { + let batch_start = batch_idx * BATCH_SIZE; + + // Zero-copy slice of the batch observations + // SAFETY: All writes to this batch are complete (batch_writes == BATCH_SIZE) + let obs_array = unsafe { &*self.observations.get() }; + let batch_view = + obs_array.slice_axis(Axis(0), Slice::from(batch_start..batch_start + BATCH_SIZE)); + debug_assert!( + batch_view.is_standard_layout(), + "batch_view should be contiguous for efficient GPU transfer" + ); + + // Calculate the batch end ticket (first ticket of next batch) + let batch_number = trigger_ticket / BATCH_SIZE as u64; + let batch_end_ticket = (batch_number + 1) * BATCH_SIZE as u64; + + let completion = BatchCompletion { + state: self.state.clone(), + batch_idx, + batch_start, + batch_end_ticket, + }; + + (self.dispatch)(batch_idx, batch_view, completion); + } + + /// Poll for a result. Returns Some(&O) if ready, None if still pending. + pub fn poll(&self, ticket: u64) -> Option<&O> { + let batch_idx = ((ticket as usize) / BATCH_SIZE) % self.num_batches; + let batch_end_ticket = ((ticket / BATCH_SIZE as u64) + 1) * BATCH_SIZE as u64; + + // Check if this batch is complete + if self.state.batch_complete[batch_idx].load(Ordering::Acquire) < batch_end_ticket { + return None; + } + + // Batch is complete, return reference to output + let slot_idx = (ticket as usize) % self.total_slots; + // SAFETY: batch_complete >= batch_end_ticket means output is written and won't change + Some(unsafe { &*self.state.outputs[slot_idx].get() }) + } + + /// Get a listener for the completion event. + /// Use this before polling to avoid missing notifications. + pub fn listen(&self) -> event_listener::EventListener { + self.state.completion_event.listen() + } + + /// Wake all waiters (used for GPU completion or external cancellation). + pub fn notify_all(&self) { + self.state.completion_event.notify(usize::MAX); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Ix0; + use std::sync::Arc; + + #[test] + fn test_single_batch_completion() { + // Use Ix0 (scalar) for simple tests + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + // Simple transform: output = input * 2 + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + // Submit BATCH_SIZE jobs + let tickets: Vec = (0..BATCH_SIZE as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // All should be complete now (last submit triggered dispatch) + for (i, &ticket) in tickets.iter().enumerate() { + let result = queue.poll(ticket); + assert!(result.is_some(), "ticket {} should be ready", ticket); + assert_eq!(*result.unwrap(), (i as u64) * 2); + } + } + + #[test] + fn test_partial_batch_not_ready() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + // Submit less than a full batch + let tickets: Vec = (0..BATCH_SIZE as u64 - 1) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // None should be ready + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_none(), + "partial batch should not be ready" + ); + } + + // Complete the batch + queue.submit(|mut out| out[()] = BATCH_SIZE as u64 - 1); + + // Now all should be ready + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_some(), + "batch should be ready after completion" + ); + } + } + + #[test] + fn test_multiple_batches() { + let num_jobs = BATCH_SIZE * 3; + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + num_jobs, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input + 1000; + } + completion.complete(&outputs); + }, + )); + + // Submit 3 full batches + let all_tickets: Vec = (0..num_jobs as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // All should be ready + for (i, &ticket) in all_tickets.iter().enumerate() { + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!(*result, (i as u64) + 1000); + } + } + + #[test] + fn test_batch_slot_reuse() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input; + } + completion.complete(&outputs); + }, + )); + + let total_slots = queue.total_slots(); + + // Submit exactly enough jobs to fill every slot once. + let tickets_round1: Vec = (0..total_slots as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // Read all results from round 1 + for (i, &ticket) in tickets_round1.iter().enumerate() { + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!( + *result, i as u64, + "round 1 ticket {} has wrong value", + ticket + ); + } + + // Now submit another round (reusing slots) + let tickets_round2: Vec = (total_slots as u64..(total_slots * 2) as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // Read all results from round 2 + for (i, &ticket) in tickets_round2.iter().enumerate() { + let expected = (total_slots + i) as u64; + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!( + *result, expected, + "round 2 ticket {} has wrong value", + ticket + ); + } + } + + #[test] + fn test_concurrent_submissions() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::thread; + + let batch_count = Arc::new(AtomicUsize::new(0)); + let num_threads = 4; + let jobs_per_thread = BATCH_SIZE * 2; // Each thread submits 2 batches worth + + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + num_threads * jobs_per_thread, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + let handles: Vec<_> = (0..num_threads) + .map(|thread_id| { + let queue = queue.clone(); + let batch_count = batch_count.clone(); + thread::spawn(move || { + let base = (thread_id * jobs_per_thread) as u64; + let mut results = Vec::new(); + + for i in 0..jobs_per_thread as u64 { + let val = base + i; + let ticket = queue.submit(|mut out| out[()] = val); + results.push((ticket, val)); + } + + // Wait for results + for (ticket, expected_input) in results { + loop { + if let Some(&result) = queue.poll(ticket) { + assert_eq!(result, expected_input * 2); + batch_count.fetch_add(1, Ordering::Relaxed); + break; + } + // Busy wait (in real code we'd use the event listener) + std::hint::spin_loop(); + } + } + }) + }) + .collect(); + + for handle in handles { + handle.join().expect("thread panicked"); + } + + assert_eq!( + batch_count.load(Ordering::Relaxed), + num_threads * jobs_per_thread + ); + } +} diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs new file mode 100644 index 0000000..f881065 --- /dev/null +++ b/training/src/replay_buffer.rs @@ -0,0 +1,716 @@ +//! Lock-free replay buffer for concurrent training data storage. + +use std::cell::UnsafeCell; +use std::fs::File; +use std::io::{self, BufReader, BufWriter, Read, Write}; +use std::mem::MaybeUninit; +use std::path::Path; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use rand::seq::index::sample; +use rand::Rng; + +/// A training sample stored in the replay buffer. +#[derive(Clone, Debug)] +pub struct Sample { + pub notation: String, + pub policy: Vec, + pub value: f32, +} + +/// Lock-free ring buffer for storing training samples. +pub struct ReplayBuffer { + data: Box<[UnsafeCell>]>, + /// Tracks whether each slot has ever been initialized. + /// + /// This is separate from `saved_head`: save checkpoints can be marked at + /// arbitrary points, but overwrite safety needs per-slot init state. + initialized: Box<[AtomicBool]>, + capacity: usize, + head: AtomicU64, + writers: AtomicU64, + /// Tracks the head position at the last save. Used to avoid saving + /// the same samples multiple times across checkpoints. + saved_head: AtomicU64, +} + +unsafe impl Sync for ReplayBuffer {} +unsafe impl Send for ReplayBuffer {} + +/// RAII guard for writing to reserved slots. +pub struct ReserveGuard<'a> { + buffer: &'a ReplayBuffer, + start: u64, + len: usize, + written: usize, +} + +impl<'a> ReserveGuard<'a> { + #[inline] + pub fn push(&mut self, sample: Sample) { + assert!(self.written < self.len, "wrote more samples than reserved"); + let idx = (self.start + self.written as u64) as usize % self.buffer.capacity; + unsafe { + if self.buffer.initialized[idx].load(Ordering::Acquire) { + (*self.buffer.data[idx].get()).assume_init_drop(); + } + (*self.buffer.data[idx].get()).write(sample); + } + self.buffer.initialized[idx].store(true, Ordering::Release); + self.written += 1; + } + + pub fn extend(&mut self, samples: impl IntoIterator) { + for sample in samples { + self.push(sample); + } + } +} + +impl Drop for ReserveGuard<'_> { + fn drop(&mut self) { + self.buffer.writers.fetch_sub(1, Ordering::Release); + } +} + +impl ReplayBuffer { + pub fn new(capacity: usize) -> Self { + let data: Vec>> = (0..capacity) + .map(|_| UnsafeCell::new(MaybeUninit::uninit())) + .collect(); + Self { + data: data.into_boxed_slice(), + initialized: (0..capacity).map(|_| AtomicBool::new(false)).collect(), + capacity, + head: AtomicU64::new(0), + writers: AtomicU64::new(0), + saved_head: AtomicU64::new(0), + } + } + + pub fn reserve(&self, n: usize) -> ReserveGuard<'_> { + self.writers.fetch_add(1, Ordering::Acquire); + let start = self.head.fetch_add(n as u64, Ordering::AcqRel); + ReserveGuard { + buffer: self, + start, + len: n, + written: 0, + } + } + + #[inline] + fn valid_range(&self) -> (u64, u64) { + let head = self.head.load(Ordering::Acquire); + let tail = head.saturating_sub(self.capacity as u64); + (tail, head) + } + + #[inline] + pub fn len(&self) -> usize { + let (tail, head) = self.valid_range(); + (head - tail) as usize + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Sample `n` items uniformly. Panics if writers are active. + pub fn sample(&self, n: usize, rng: &mut impl Rng) -> Vec { + assert_eq!( + self.writers.load(Ordering::Acquire), + 0, + "cannot sample while writers are active" + ); + + let (tail, head) = self.valid_range(); + let count = (head - tail) as usize; + if count == 0 || n == 0 { + return Vec::new(); + } + + let indices = sample(rng, count, n.min(count)); + indices + .iter() + .map(|offset| { + let idx = tail + offset as u64; + let slot = idx as usize % self.capacity; + unsafe { (*self.data[slot].get()).assume_init_ref().clone() } + }) + .collect() + } + + /// Save unsaved samples to binary file. + /// + /// Only saves samples that haven't been saved yet (from `saved_head` to `head`). + /// Call `mark_saved()` after a successful save to update the tracking. + /// + /// Returns the number of samples written. + /// + /// File format (version 2): + /// - magic: 8 bytes "SIEBREN\0" + /// - version: u64 (little endian) + /// - generation_id: u64 (little endian) + /// - sample_count: u64 (little endian) + /// - max_notation_len: u64 (little endian) - max length of any notation + /// - policy_len: u64 (little endian) + /// - samples: for each sample: + /// - notation_len: u64 (little endian) - actual length of this notation + /// - notation: [u8; max_notation_len] - UTF-8 bytes, padded with zeros + /// - policy: [f32; policy_len] (little endian) + /// - value: f32 (little endian) + pub fn save(&self, path: &Path, generation_id: u64, policy_len: usize) -> io::Result { + const MAGIC: &[u8; 8] = b"SIEBREN\0"; + const VERSION: u64 = 2; + + assert_eq!( + self.writers.load(Ordering::Acquire), + 0, + "cannot save while writers are active" + ); + + let head = self.head.load(Ordering::Acquire); + let saved_head = self.saved_head.load(Ordering::Acquire); + let (tail, _) = self.valid_range(); + + // Only save samples from saved_head to head, but not before tail + // (samples before tail have been overwritten in the ring buffer) + let start = saved_head.max(tail); + let count = head.saturating_sub(start) as usize; + + if count == 0 { + // Nothing new to save - still write an empty file for consistency + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + writer.write_all(MAGIC)?; + writer.write_all(&VERSION.to_le_bytes())?; + writer.write_all(&generation_id.to_le_bytes())?; + writer.write_all(&0u64.to_le_bytes())?; // sample_count = 0 + writer.write_all(&0u64.to_le_bytes())?; // max_notation_len = 0 + writer.write_all(&(policy_len as u64).to_le_bytes())?; + writer.flush()?; + return Ok(0); + } + + // Find max notation length in the range we're saving + let mut max_notation_len = 0usize; + for idx in start..head { + let slot = idx as usize % self.capacity; + let sample = unsafe { (*self.data[slot].get()).assume_init_ref() }; + max_notation_len = max_notation_len.max(sample.notation.len()); + } + + let file = File::create(path)?; + let mut writer = BufWriter::new(file); + + // Write header + writer.write_all(MAGIC)?; + writer.write_all(&VERSION.to_le_bytes())?; + writer.write_all(&generation_id.to_le_bytes())?; + writer.write_all(&(count as u64).to_le_bytes())?; + writer.write_all(&(max_notation_len as u64).to_le_bytes())?; + writer.write_all(&(policy_len as u64).to_le_bytes())?; + + // Pre-allocate padding buffer + let mut notation_buf = vec![0u8; max_notation_len]; + + // Write samples + for idx in start..head { + let slot = idx as usize % self.capacity; + let sample = unsafe { (*self.data[slot].get()).assume_init_ref() }; + + // Write notation length and padded notation + let notation_bytes = sample.notation.as_bytes(); + writer.write_all(&(notation_bytes.len() as u64).to_le_bytes())?; + notation_buf[..notation_bytes.len()].copy_from_slice(notation_bytes); + notation_buf[notation_bytes.len()..].fill(0); + writer.write_all(¬ation_buf)?; + + // Write policy + for &p in &sample.policy { + writer.write_all(&p.to_le_bytes())?; + } + + // Write value + writer.write_all(&sample.value.to_le_bytes())?; + } + + writer.flush()?; + Ok(count) + } + + /// Mark all current samples as saved. + /// + /// Call this after a successful `save()` to prevent those samples from + /// being saved again in subsequent calls. + pub fn mark_saved(&self) { + let head = self.head.load(Ordering::Acquire); + self.saved_head.store(head, Ordering::Release); + } + + /// Load samples from binary file. + /// + /// Returns (samples_loaded, generation_id). + /// Panics if writers are active. + pub fn load(&self, path: &Path) -> io::Result<(usize, u64)> { + const MAGIC: &[u8; 8] = b"SIEBREN\0"; + + assert_eq!( + self.writers.load(Ordering::Acquire), + 0, + "cannot load while writers are active" + ); + + let file = File::open(path)?; + let mut reader = BufReader::new(file); + + // Read and validate magic + let mut magic = [0u8; 8]; + reader.read_exact(&mut magic)?; + if &magic != MAGIC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid magic bytes", + )); + } + + // Read and validate version + let mut buf8 = [0u8; 8]; + reader.read_exact(&mut buf8)?; + let version = u64::from_le_bytes(buf8); + if version != 2 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("unsupported version: {version}, expected 2"), + )); + } + + // Read header fields + reader.read_exact(&mut buf8)?; + let generation_id = u64::from_le_bytes(buf8); + + reader.read_exact(&mut buf8)?; + let sample_count = u64::from_le_bytes(buf8) as usize; + + reader.read_exact(&mut buf8)?; + let max_notation_len = u64::from_le_bytes(buf8) as usize; + + reader.read_exact(&mut buf8)?; + let policy_len = u64::from_le_bytes(buf8) as usize; + + // Pre-allocate buffers + let mut notation_buf = vec![0u8; max_notation_len]; + let mut policy_buf = vec![0u8; policy_len * 4]; + let mut buf4 = [0u8; 4]; + + // Reserve space and read samples + let mut guard = self.reserve(sample_count); + + for _ in 0..sample_count { + // Read notation length + reader.read_exact(&mut buf8)?; + let notation_len = u64::from_le_bytes(buf8) as usize; + + if notation_len > max_notation_len { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "notation length {notation_len} exceeds max_notation_len {max_notation_len}" + ), + )); + } + + // Read padded notation + reader.read_exact(&mut notation_buf)?; + let notation = + String::from_utf8(notation_buf[..notation_len].to_vec()).map_err(|e| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("invalid UTF-8 in notation: {e}"), + ) + })?; + + // Read policy + reader.read_exact(&mut policy_buf)?; + let policy: Vec = policy_buf + .chunks_exact(4) + .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) + .collect(); + + // Read value + reader.read_exact(&mut buf4)?; + let value = f32::from_le_bytes(buf4); + + guard.push(Sample { + notation, + policy, + value, + }); + } + + Ok((sample_count, generation_id)) + } +} + +impl Drop for ReplayBuffer { + fn drop(&mut self) { + for idx in 0..self.capacity { + if self.initialized[idx].load(Ordering::Relaxed) { + unsafe { + (*self.data[idx].get()).assume_init_drop(); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + use std::sync::Arc; + + fn make_sample(id: usize) -> Sample { + Sample { + notation: format!("_________|A"), + policy: vec![id as f32 / 10.0; 9], + value: id as f32 / 100.0, + } + } + + #[test] + fn test_reserve_guard_push() { + let buffer = ReplayBuffer::new(100); + + { + let mut guard = buffer.reserve(3); + guard.push(make_sample(0)); + guard.push(make_sample(1)); + guard.push(make_sample(2)); + } + + assert_eq!(buffer.len(), 3); + assert_eq!(buffer.writers.load(Ordering::Relaxed), 0); + } + + #[test] + fn test_reserve_guard_extend() { + let buffer = ReplayBuffer::new(100); + + { + let mut guard = buffer.reserve(5); + guard.extend((0..5).map(make_sample)); + } + + assert_eq!(buffer.len(), 5); + } + + #[test] + fn test_sample() { + let buffer = ReplayBuffer::new(100); + let mut rng = ChaCha8Rng::seed_from_u64(42); + + { + let mut guard = buffer.reserve(10); + guard.extend((0..10).map(make_sample)); + } + + let samples = buffer.sample(5, &mut rng); + assert_eq!(samples.len(), 5); + } + + #[test] + fn test_concurrent_writes() { + let buffer = Arc::new(ReplayBuffer::new(1000)); + let num_threads = 4; + let samples_per_thread = 100; + + std::thread::scope(|s| { + for _ in 0..num_threads { + let buffer = Arc::clone(&buffer); + s.spawn(move || { + let mut guard = buffer.reserve(samples_per_thread); + guard.extend((0..samples_per_thread).map(make_sample)); + }); + } + }); + + assert_eq!(buffer.len(), num_threads * samples_per_thread); + assert_eq!(buffer.writers.load(Ordering::Relaxed), 0); + } + + #[test] + fn test_buffer_wraparound() { + let buffer = ReplayBuffer::new(10); + + for i in 0..25 { + let mut guard = buffer.reserve(1); + guard.push(make_sample(i)); + } + + assert_eq!(buffer.len(), 10); + let (tail, head) = buffer.valid_range(); + assert_eq!(tail, 15); + assert_eq!(head, 25); + } + + #[test] + #[should_panic(expected = "cannot sample while writers are active")] + fn test_sample_during_write_panics() { + let buffer = ReplayBuffer::new(100); + let mut rng = ChaCha8Rng::seed_from_u64(42); + + let _guard = buffer.reserve(5); + let _ = buffer.sample(5, &mut rng); + } + + #[test] + fn test_save_load_roundtrip() { + let buffer = ReplayBuffer::new(100); + let policy_len = 9; + let generation_id = 42u64; + + // Add samples + { + let mut guard = buffer.reserve(5); + guard.extend((0..5).map(make_sample)); + } + + // Save to temp file + let temp_dir = std::env::temp_dir(); + let path = temp_dir.join("test_replay_buffer.bin"); + + let saved_count = buffer.save(&path, generation_id, policy_len).unwrap(); + assert_eq!(saved_count, 5); + + // Load into new buffer + let buffer2 = ReplayBuffer::new(100); + let (loaded_count, loaded_gen) = buffer2.load(&path).unwrap(); + + assert_eq!(loaded_count, 5); + assert_eq!(loaded_gen, generation_id); + assert_eq!(buffer2.len(), 5); + + // Verify samples match + let mut rng = ChaCha8Rng::seed_from_u64(0); + let original = buffer.sample(5, &mut rng); + let mut rng = ChaCha8Rng::seed_from_u64(0); + let loaded = buffer2.sample(5, &mut rng); + + for (orig, load) in original.iter().zip(loaded.iter()) { + assert_eq!(orig.notation, load.notation); + assert_eq!(orig.policy, load.policy); + assert!((orig.value - load.value).abs() < 1e-6); + } + + // Cleanup + std::fs::remove_file(&path).ok(); + } + + #[test] + fn test_save_load_generation_id() { + let buffer = ReplayBuffer::new(100); + + { + let mut guard = buffer.reserve(1); + guard.push(make_sample(0)); + } + + let temp_dir = std::env::temp_dir(); + let path = temp_dir.join("test_replay_gen_id.bin"); + + // Save with specific generation ID + let gen_id = 12345u64; + buffer.save(&path, gen_id, 9).unwrap(); + + // Load and verify generation ID is preserved + let buffer2 = ReplayBuffer::new(100); + let (_, loaded_gen) = buffer2.load(&path).unwrap(); + assert_eq!(loaded_gen, gen_id); + + // Cleanup + std::fs::remove_file(&path).ok(); + } + + #[test] + fn test_save_load_varying_notation_lengths() { + let buffer = ReplayBuffer::new(100); + + // Add samples with different notation lengths + { + let mut guard = buffer.reserve(3); + guard.push(Sample { + notation: "A".to_string(), + policy: vec![0.1; 9], + value: 0.5, + }); + guard.push(Sample { + notation: "ABCDEFGHIJ".to_string(), + policy: vec![0.2; 9], + value: 0.6, + }); + guard.push(Sample { + notation: "XYZ".to_string(), + policy: vec![0.3; 9], + value: 0.7, + }); + } + + let temp_dir = std::env::temp_dir(); + let path = temp_dir.join("test_replay_varying_notation.bin"); + + buffer.save(&path, 1, 9).unwrap(); + + let buffer2 = ReplayBuffer::new(100); + let (count, _) = buffer2.load(&path).unwrap(); + assert_eq!(count, 3); + + // Verify all notations preserved correctly + let mut rng = ChaCha8Rng::seed_from_u64(0); + let original = buffer.sample(3, &mut rng); + let mut rng = ChaCha8Rng::seed_from_u64(0); + let loaded = buffer2.sample(3, &mut rng); + + for (orig, load) in original.iter().zip(loaded.iter()) { + assert_eq!(orig.notation, load.notation); + } + + // Cleanup + std::fs::remove_file(&path).ok(); + } + + #[test] + fn test_save_empty_buffer() { + let buffer = ReplayBuffer::new(100); + + let temp_dir = std::env::temp_dir(); + let path = temp_dir.join("test_replay_empty.bin"); + + let saved_count = buffer.save(&path, 0, 9).unwrap(); + assert_eq!(saved_count, 0); + + let buffer2 = ReplayBuffer::new(100); + let (count, gen) = buffer2.load(&path).unwrap(); + assert_eq!(count, 0); + assert_eq!(gen, 0); + assert_eq!(buffer2.len(), 0); + + // Cleanup + std::fs::remove_file(&path).ok(); + } + + #[test] + fn test_mark_saved_prevents_duplicates() { + let buffer = ReplayBuffer::new(100); + let temp_dir = std::env::temp_dir(); + + // Add first batch of samples + { + let mut guard = buffer.reserve(3); + guard.extend((0..3).map(make_sample)); + } + + // Save first batch + let path1 = temp_dir.join("test_mark_saved_1.bin"); + let saved1 = buffer.save(&path1, 1, 9).unwrap(); + assert_eq!(saved1, 3); + buffer.mark_saved(); + + // Add second batch + { + let mut guard = buffer.reserve(2); + guard.extend((10..12).map(make_sample)); + } + + // Save second batch - should only save the new samples + let path2 = temp_dir.join("test_mark_saved_2.bin"); + let saved2 = buffer.save(&path2, 2, 9).unwrap(); + assert_eq!(saved2, 2); // Only the new samples + + // Buffer still has all 5 samples + assert_eq!(buffer.len(), 5); + + // Load both files into separate buffers and verify no duplicates + let buffer1 = ReplayBuffer::new(100); + let (count1, _) = buffer1.load(&path1).unwrap(); + assert_eq!(count1, 3); + + let buffer2 = ReplayBuffer::new(100); + let (count2, _) = buffer2.load(&path2).unwrap(); + assert_eq!(count2, 2); + + // Cleanup + std::fs::remove_file(&path1).ok(); + std::fs::remove_file(&path2).ok(); + } + + #[test] + fn test_save_without_mark_saved_resaves_all() { + let buffer = ReplayBuffer::new(100); + let temp_dir = std::env::temp_dir(); + + // Add samples + { + let mut guard = buffer.reserve(3); + guard.extend((0..3).map(make_sample)); + } + + // Save without calling mark_saved + let path1 = temp_dir.join("test_no_mark_saved_1.bin"); + let saved1 = buffer.save(&path1, 1, 9).unwrap(); + assert_eq!(saved1, 3); + // Note: NOT calling mark_saved() + + // Save again - should save the same samples again + let path2 = temp_dir.join("test_no_mark_saved_2.bin"); + let saved2 = buffer.save(&path2, 2, 9).unwrap(); + assert_eq!(saved2, 3); // Same 3 samples saved again + + // Cleanup + std::fs::remove_file(&path1).ok(); + std::fs::remove_file(&path2).ok(); + } + + #[test] + fn test_save_respects_ring_buffer_overwrites() { + // Small buffer that will wrap around + let buffer = ReplayBuffer::new(5); + let temp_dir = std::env::temp_dir(); + + // Add 3 samples + { + let mut guard = buffer.reserve(3); + guard.extend((0..3).map(make_sample)); + } + + // Save and mark + let path1 = temp_dir.join("test_overwrite_1.bin"); + let saved1 = buffer.save(&path1, 1, 9).unwrap(); + assert_eq!(saved1, 3); + buffer.mark_saved(); + + // Add 4 more samples - this will overwrite some of the original samples + // Buffer now contains samples at positions 3,4,5,6 (positions 0,1,2 overwritten) + { + let mut guard = buffer.reserve(4); + guard.extend((10..14).map(make_sample)); + } + + // Save again - should only save the 4 new samples + let path2 = temp_dir.join("test_overwrite_2.bin"); + let saved2 = buffer.save(&path2, 2, 9).unwrap(); + assert_eq!(saved2, 4); + + // Cleanup + std::fs::remove_file(&path1).ok(); + std::fs::remove_file(&path2).ok(); + } +} diff --git a/training/src/training.rs b/training/src/training.rs new file mode 100644 index 0000000..03cb97a --- /dev/null +++ b/training/src/training.rs @@ -0,0 +1,390 @@ +//! Training infrastructure - thread spawning and coordination. +//! +//! Provides `SelfPlaySession`: a persistent session with pause/resume semantics +//! that preserves in-progress game state across boundaries. + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread; + +use ndarray::ArrayView; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +use crate::eval::{GpuEvaluator, PolicyValue}; +use crate::executor::Executor; +use crate::observation_replay_buffer::ObservationReplayBuffer; +use crate::queue::{BatchCompletion, GpuJobQueue}; +use crate::worker::{worker_loop_forever, WorkerConfig}; +use crate::{BatchDim, Environment}; + +/// Shared control state for the persistent self-play session. +/// +/// Counters are behind `Arc` so they can be cloned directly +/// into worker futures that expect `Arc`. +struct SessionControl { + /// Whether workers should be actively polling. + running: AtomicBool, + /// Whether the session is being torn down. + shutdown: AtomicBool, + /// Absolute target: workers pause when `samples_collected >= target`. + target_samples: AtomicUsize, + /// Total samples collected (monotonic across the session lifetime). + samples_collected: Arc, + /// Total games completed. + games_completed: Arc, + /// Number of threads currently inside the executor polling loop. + active_pollers: AtomicUsize, + /// Condvar + mutex for coordinating start/pause/quiesce/shutdown. + condvar: Condvar, + condvar_mutex: Mutex<()>, +} + +impl SessionControl { + fn new() -> Self { + Self { + running: AtomicBool::new(false), + shutdown: AtomicBool::new(false), + target_samples: AtomicUsize::new(0), + samples_collected: Arc::new(AtomicUsize::new(0)), + games_completed: Arc::new(AtomicUsize::new(0)), + active_pollers: AtomicUsize::new(0), + condvar: Condvar::new(), + condvar_mutex: Mutex::new(()), + } + } + + /// Wake all threads blocked on the condvar. + fn wake_all(&self) { + self.condvar.notify_all(); + } + + /// Check if the thread should stop polling (pause or shutdown). + fn should_pause(&self) -> bool { + if self.shutdown.load(Ordering::Acquire) { + return true; + } + if !self.running.load(Ordering::Acquire) { + return true; + } + self.samples_collected.load(Ordering::Acquire) + >= self.target_samples.load(Ordering::Acquire) + } +} + +/// Configuration for creating a persistent self-play session. +#[derive(Clone)] +pub struct SessionConfig { + /// Number of OS threads to spawn. + pub num_threads: usize, + /// Number of workers per thread. + pub workers_per_thread: usize, + /// Worker configuration (MCTS params, temperature, etc). + pub worker: WorkerConfig, + /// Random seed for reproducibility. + pub seed: u64, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + num_threads: 32, + workers_per_thread: 16, + worker: WorkerConfig::default(), + seed: 42, + } + } +} + +/// Trait-object wrapper so we can call `notify_all()` on the queue without +/// leaking the full generic type into `SelfPlaySession`. +trait QueueNotify: Send + Sync { + fn notify_all(&self); +} + +impl QueueNotify for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + fn notify_all(&self) { + GpuJobQueue::notify_all(self); + } +} + +/// A persistent self-play session that owns worker threads and preserves +/// in-progress game state across pause/resume boundaries. +/// +/// # Lifecycle +/// +/// 1. `new(...)` — creates threads and futures (paused). +/// 2. `start()` — sets target to `usize::MAX` and wakes threads. +/// 3. `wait_for(target)` — ensures running, blocks until `samples >= target`, +/// then pauses and waits for all pollers to quiesce. Safe to read replay +/// buffer after this returns. +/// 4. `samples()` — returns current absolute sample count. +/// 5. `shutdown()` / Rust `Drop` — sets shutdown, joins threads. +pub struct SelfPlaySession { + control: Arc, + /// Queue used by all workers. Kept alive for `notify_all` on drop. + queue_notify: Arc, + /// Join handles for worker threads. `None` after `shutdown`. + threads: Option>>, +} + +impl SelfPlaySession { + /// Create a new persistent session. + /// + /// Threads are spawned immediately but start paused. The `dispatch` callback + /// is invoked when a batch of observations is ready for GPU inference. + pub fn new( + config: SessionConfig, + replay_buffer: Arc>, + dispatch: F, + ) -> Self + where + E: Environment + Clone + Send + 'static, + E::ObsDim: BatchDim, + F: Fn( + usize, + ArrayView::BatchedDim>, + BatchCompletion>, + ) + Send + + Sync + + 'static, + { + let total_workers = config + .num_threads + .checked_mul(config.workers_per_thread) + .expect("num_threads * workers_per_thread overflowed usize"); + + let queue: Arc>> = + Arc::new(GpuJobQueue::new(E::OBS_SHAPE, total_workers, dispatch)); + + let control = Arc::new(SessionControl::new()); + + let mut threads = Vec::with_capacity(config.num_threads); + for thread_id in 0..config.num_threads { + let queue = queue.clone(); + let control = control.clone(); + let config = config.clone(); + let replay_buffer = replay_buffer.clone(); + + let handle = thread::spawn(move || { + session_thread_main::( + thread_id, + queue, + config, + control, + &replay_buffer, + ); + }); + threads.push(handle); + } + + Self { + control, + queue_notify: queue, + threads: Some(threads), + } + } + + /// Start self-play with no sample limit (runs until explicitly paused or + /// `wait_for` is called). + pub fn start(&self) { + self.control + .target_samples + .store(usize::MAX, Ordering::Release); + self.control.running.store(true, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + } + + /// Block until at least `target_samples` absolute samples have been + /// collected, then pause and quiesce all workers. + /// + /// Returns the actual number of samples collected (may exceed target). + /// + /// After this returns, no worker thread is inside the executor polling + /// loop, so it is safe to read the replay buffer. + pub fn wait_for(&self, target_samples: usize) -> usize { + // Set target and ensure running. + self.control + .target_samples + .store(target_samples, Ordering::Release); + self.control.running.store(true, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + + // Wait until target is reached (condvar-based, no spinning). + { + let mut guard = self + .control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while self.control.samples_collected.load(Ordering::Acquire) < target_samples + && !self.control.shutdown.load(Ordering::Acquire) + { + guard = self + .control + .condvar + .wait(guard) + .expect("condvar wait failed"); + } + } + + // Pause workers. + self.control.running.store(false, Ordering::Release); + self.queue_notify.notify_all(); + self.control.wake_all(); + + // Wait for all pollers to exit (quiesce). + { + let mut guard = self + .control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while self.control.active_pollers.load(Ordering::Acquire) > 0 + && !self.control.shutdown.load(Ordering::Acquire) + { + guard = self + .control + .condvar + .wait(guard) + .expect("condvar wait failed"); + } + } + + self.control.samples_collected.load(Ordering::Acquire) + } + + /// Return the current absolute sample count. + pub fn samples(&self) -> usize { + self.control.samples_collected.load(Ordering::Acquire) + } + + /// Return the current absolute game count. + pub fn games(&self) -> usize { + self.control.games_completed.load(Ordering::Acquire) + } + + /// Shut down the session. Idempotent. + pub fn shutdown(&mut self) { + if let Some(threads) = self.threads.take() { + self.control.shutdown.store(true, Ordering::Release); + self.control.running.store(false, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + + for handle in threads { + let _ = handle.join(); + } + } + } +} + +impl Drop for SelfPlaySession { + fn drop(&mut self) { + self.shutdown(); + } +} + +/// Main loop for a single thread in a persistent session. +/// +/// 1. Wait on condvar until `running || shutdown`. +/// 2. If shutdown => exit. +/// 3. Increment `active_pollers`, run executor until pause/shutdown/target. +/// 4. Decrement `active_pollers`, notify condvar so `wait_for()` can observe +/// quiesce. +/// 5. Goto 1. +fn session_thread_main( + thread_id: usize, + queue: Arc>>, + config: SessionConfig, + control: Arc, + replay_buffer: &ObservationReplayBuffer, +) where + E: Environment + Clone + 'static, + E::ObsDim: BatchDim, +{ + let base_seed = config.seed.wrapping_add(thread_id as u64 * 1000); + let evaluator = GpuEvaluator::::new(&*queue); + + // Clone the session's counters for workers. + let samples_collected = control.samples_collected.clone(); + let games_completed = control.games_completed.clone(); + + // Create futures once. They live for the entire session. + let mut futures: Vec + '_>>> = (0 + ..config.workers_per_thread) + .map(|i| { + let samples_collected = samples_collected.clone(); + let games_completed = games_completed.clone(); + let mut rng = ChaCha8Rng::seed_from_u64(base_seed + i as u64); + let evaluator_ref = &evaluator; + let worker_config = &config.worker; + + let fut = async move { + worker_loop_forever::( + evaluator_ref, + worker_config, + &mut rng, + samples_collected, + games_completed, + replay_buffer, + ) + .await; + }; + Box::pin(fut) as std::pin::Pin + '_>> + }) + .collect(); + + let executor = Executor::new(|| queue.listen()); + + loop { + // 1. Wait until running or shutdown. + { + let mut guard = control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while !control.running.load(Ordering::Acquire) + && !control.shutdown.load(Ordering::Acquire) + { + guard = control.condvar.wait(guard).expect("condvar wait failed"); + } + } + + // 2. If shutdown, exit. + if control.shutdown.load(Ordering::Acquire) { + return; + } + + // 3. Increment active_pollers and run executor. + control.active_pollers.fetch_add(1, Ordering::AcqRel); + + let control_ref = &control; + let queue_ref = &queue; + executor.run(&mut futures, &mut || { + let should_pause = control_ref.should_pause(); + if should_pause { + queue_ref.notify_all(); + } + // Notify the condvar when samples cross the target so wait_for() + // wakes up. + if control_ref.samples_collected.load(Ordering::Acquire) + >= control_ref.target_samples.load(Ordering::Acquire) + { + control_ref.wake_all(); + } + should_pause + }); + + // 4. Decrement active_pollers and notify. + control.active_pollers.fetch_sub(1, Ordering::AcqRel); + control.wake_all(); + } +} diff --git a/training/src/worker.rs b/training/src/worker.rs new file mode 100644 index 0000000..4ecb6b8 --- /dev/null +++ b/training/src/worker.rs @@ -0,0 +1,323 @@ +//! Worker loop and training data collection. +//! +//! Each worker runs MCTS searches, plays games, and collects training samples. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use rand::Rng; + +use crate::eval::Evaluator; +use crate::mcts::{best_action_index, sample_action_index, visits_to_policy, MCTSConfig, MCTS}; +use crate::observation_replay_buffer::ObservationReplayBuffer; +use crate::{Action, Environment, Player, TerminalState}; + +/// A training sample from a single game step. +#[derive(Clone, Debug)] +pub struct TrainingSample { + /// Action index taken from this state. + pub action_idx: usize, + /// The policy from MCTS (normalized visit counts). + pub policy: Vec, + /// The value from MCTS search. + pub value: f32, + /// Player to move at this step. + /// + /// Keeping this here avoids reparsing game state while backfilling outcomes. + pub player: Player, +} + +/// Full trace of one played self-play game. +#[derive(Clone, Debug)] +pub struct PlayedGame { + /// Initial environment state before any actions in this game. + pub initial_env: E, + /// Step samples in chronological order. + pub samples: Vec, +} + +/// Configuration for the worker. +#[derive(Clone)] +pub struct WorkerConfig { + /// MCTS configuration. + pub mcts: MCTSConfig, + /// Temperature for action selection (1.0 = proportional to visits, 0.0 = argmax). + pub temperature: f32, + /// Number of moves at the start of the game to use exploration temperature. + /// After this many moves, use temperature 0 (argmax). + pub exploration_moves: usize, +} + +impl Default for WorkerConfig { + fn default() -> Self { + Self { + mcts: MCTSConfig::default(), + temperature: 1.0, + exploration_moves: 30, + } + } +} + +/// Run a single self-play game, collecting training samples. +/// +/// Returns the collected samples. The game continues until terminal. +pub async fn play_game(evaluator: &V, config: &WorkerConfig, rng: &mut R) -> PlayedGame +where + E: Environment, + V: Evaluator, + R: Rng, +{ + let mut env = E::new(); + let initial_env = env.clone(); + let mut samples = Vec::new(); + let mut move_count = 0; + + let mcts = MCTS::new(evaluator, &config.mcts); + + loop { + if env.is_terminal().is_some() { + break; + } + + let visits = mcts.search(&mut env, rng).await; + // Convert visits to policy + let temp = if move_count < config.exploration_moves { + config.temperature + } else { + 0.0 + }; + let policy = visits_to_policy(&visits, temp); + let player = env.current_player(); + + // Value is set to 0.0 here and backfilled with game outcome after the game ends. + // This is standard AlphaZero practice - we use the actual game result rather than + // the search value estimate for training. + let value = 0.0; + + // Select action + let action_idx = if temp > 0.0 { + sample_action_index(&policy, rng) + } else { + best_action_index(&visits) + }; + + let action_idx = action_idx.expect("no valid actions but game not terminal"); + let action = E::Action::from_index(action_idx).expect("invalid action index"); + + // Record sample + samples.push(TrainingSample { + action_idx, + player, + policy, + value, + }); + + // Apply action + env.apply_action(action); + move_count += 1; + } + + // Backfill values with game outcome + let outcome = env.is_terminal().expect("game should be terminal"); + backfill_values(&mut samples, outcome); + + PlayedGame { + initial_env, + samples, + } +} + +/// Backfill sample values with the game outcome. +/// +/// For wins, the winner's moves get +1, loser's get -1. +/// For draws, all moves get 0. +fn backfill_values(samples: &mut [TrainingSample], outcome: TerminalState) { + for sample in samples.iter_mut() { + sample.value = match outcome { + TerminalState::Win(winner) => { + if sample.player == winner { + 1.0 + } else { + -1.0 + } + } + TerminalState::Draw => 0.0, + }; + } +} + +/// Run a worker loop that plays games until the target sample count is reached. +/// +/// Workers play games and increment `samples_collected` after each game. +/// When the counter reaches `target_samples`, workers stop. The executor's +/// cancel callback should check this condition to terminate remaining workers. +/// +/// Samples are pushed directly to the shared `replay_buffer` after each completed game. +pub async fn worker_loop( + evaluator: &V, + config: &WorkerConfig, + rng: &mut R, + samples_collected: Arc, + games_completed: Arc, + target_samples: usize, + replay_buffer: &ObservationReplayBuffer, +) where + E: Environment + Clone, + V: Evaluator, + R: Rng, +{ + debug_assert_eq!(NUM_ACTIONS, E::NUM_ACTIONS); + + loop { + if samples_collected.load(Ordering::Acquire) >= target_samples { + break; + } + + let game = play_game::(evaluator, config, rng).await; + let num_samples = game.samples.len(); + + // Push observations, policies, and values to replay buffer. + let mut guard = replay_buffer.reserve(num_samples); + let mut env = game.initial_env; + for sample in game.samples { + guard.push_with_observation(&sample.policy, sample.value, |out| env.observation(out)); + let action = + E::Action::from_index(sample.action_idx).expect("invalid action index in replay"); + env.apply_action(action); + } + + samples_collected.fetch_add(num_samples, Ordering::AcqRel); + games_completed.fetch_add(1, Ordering::AcqRel); + } +} + +/// Run a worker loop that plays games forever. +/// +/// Unlike `worker_loop`, this never self-terminates based on a sample count. +/// Stopping is handled externally by the executor's cancel/pause mechanism. +/// This is used by the persistent `SelfPlaySession` where pause/resume is +/// controlled at the session level, not inside the worker. +/// +/// Samples are pushed directly to the shared `replay_buffer` after each completed game. +pub async fn worker_loop_forever( + evaluator: &V, + config: &WorkerConfig, + rng: &mut R, + samples_collected: Arc, + games_completed: Arc, + replay_buffer: &ObservationReplayBuffer, +) where + E: Environment + Clone, + V: Evaluator, + R: Rng, +{ + debug_assert_eq!(NUM_ACTIONS, E::NUM_ACTIONS); + + loop { + let game = play_game::(evaluator, config, rng).await; + let num_samples = game.samples.len(); + + // Push observations, policies, and values to replay buffer. + let mut guard = replay_buffer.reserve(num_samples); + let mut env = game.initial_env; + for sample in game.samples { + guard.push_with_observation(&sample.policy, sample.value, |out| env.observation(out)); + let action = + E::Action::from_index(sample.action_idx).expect("invalid action index in replay"); + env.apply_action(action); + } + + samples_collected.fetch_add(num_samples, Ordering::AcqRel); + games_completed.fetch_add(1, Ordering::AcqRel); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::environments::TicTacToe; + use crate::eval::UniformEvaluator; + use crate::executor::Executor; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + use std::cell::RefCell; + use std::rc::Rc; + + #[test] + fn test_play_game_collects_samples() { + let evaluator = UniformEvaluator; + let config = WorkerConfig { + mcts: MCTSConfig { + num_simulations: 50, + ..Default::default() + }, + ..Default::default() + }; + + let rng = Rc::new(RefCell::new(ChaCha8Rng::seed_from_u64(42))); + let result: Rc>>> = Rc::new(RefCell::new(None)); + + let rng_clone = rng.clone(); + let result_clone = result.clone(); + + let fut = async move { + let samples = + play_game::(&evaluator, &config, &mut *rng_clone.borrow_mut()) + .await; + *result_clone.borrow_mut() = Some(samples); + }; + + let event = event_listener::Event::new(); + let executor = Executor::new(|| event.listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + let game = result.borrow_mut().take().unwrap(); + let samples = game.samples; + + // TicTacToe games are 5-9 moves + assert!(samples.len() >= 5); + assert!(samples.len() <= 9); + + // Each sample should have correct policy size + for sample in &samples { + assert_eq!(sample.policy.len(), 9); + // Policy should sum to ~1 + let sum: f32 = sample.policy.iter().sum(); + assert!((sum - 1.0).abs() < 0.01, "policy sum: {}", sum); + } + + // Values should be set (all -1, 0, or 1) + for sample in &samples { + assert!(sample.value == -1.0 || sample.value == 0.0 || sample.value == 1.0); + } + } + + #[test] + fn test_backfill_values_win() { + let mut samples = vec![TrainingSample { + action_idx: 0, + player: crate::Player::PlayerA, + policy: vec![], + value: 0.0, + }]; + + backfill_values(&mut samples, TerminalState::Win(crate::Player::PlayerA)); + assert_eq!(samples[0].value, 1.0); + + backfill_values(&mut samples, TerminalState::Win(crate::Player::PlayerB)); + assert_eq!(samples[0].value, -1.0); + } + + #[test] + fn test_backfill_values_draw() { + let mut samples = vec![TrainingSample { + action_idx: 0, + player: crate::Player::PlayerA, + policy: vec![], + value: 0.5, // Should be overwritten + }]; + + backfill_values(&mut samples, TerminalState::Draw); + assert_eq!(samples[0].value, 0.0); + } +} From 61904e1e81a803808c116301f9ba8c4f14020cab Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Thu, 26 Mar 2026 02:44:27 -0400 Subject: [PATCH 03/59] add nix flake & deps --- Cargo.lock | 1 + flake.lock | 61 +++++++++++++++++++++++++++++++++++++++++++++ flake.nix | 49 ++++++++++++++++++++++++++++++++++++ training/Cargo.toml | 1 + 4 files changed, 112 insertions(+) create mode 100644 flake.lock create mode 100644 flake.nix diff --git a/Cargo.lock b/Cargo.lock index bb1ca57..3583a08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -23,6 +23,7 @@ dependencies = [ name = "alphapaint_training" version = "0.1.0" dependencies = [ + "alpha_paint", "cudarc", "event-listener", "ndarray", diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..f194334 --- /dev/null +++ b/flake.lock @@ -0,0 +1,61 @@ +{ + "nodes": { + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1774386573, + "narHash": "sha256-4hAV26quOxdC6iyG7kYaZcM3VOskcPUrdCQd/nx8obc=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "46db2e09e1d3f113a13c0d7b81e2f221c63b8ce9", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..2f14826 --- /dev/null +++ b/flake.nix @@ -0,0 +1,49 @@ +{ + description = "Development shell for AlphaPaint"; + + inputs = { + nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; + flake-utils.url = "github:numtide/flake-utils"; + }; + + outputs = { + self, + nixpkgs, + flake-utils, + }: + flake-utils.lib.eachDefaultSystem (system: let + pkgs = import nixpkgs { + inherit system; + config = { + allowUnfree = true; + cudaSupport = true; + }; + }; + + cudaPkgs = pkgs.cudaPackages_12_9; + in { + devShells.default = pkgs.mkShell { + buildInputs = with pkgs; [ + rustc + cargo + maturin + cmake + cudaPkgs.cudatoolkit + cudaPkgs.cudnn + cudaPkgs.libcublas + stdenv.cc.cc.lib + ]; + + shellHook = '' + export CUDA_PATH=${cudaPkgs.cudatoolkit} + export CUDA_HOME=$CUDA_PATH + export CUDA_ROOT=$CUDA_PATH + export CUDNN_PATH=${cudaPkgs.cudnn.lib} + export RUST_MIN_STACK=67108864 + export LD_LIBRARY_PATH=$CUDA_PATH/lib:$CUDA_PATH/lib64:$CUDNN_PATH/lib:/run/opengl-driver/lib:${pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH + export TRITON_LIBCUDA_PATH=/run/opengl-driver/lib + export PATH=$CUDA_PATH/bin:$PATH + ''; + }; + }); +} diff --git a/training/Cargo.toml b/training/Cargo.toml index f7bccab..88bde36 100644 --- a/training/Cargo.toml +++ b/training/Cargo.toml @@ -10,6 +10,7 @@ name = "alphapaint_training" crate-type = ["cdylib"] [dependencies] +alpha_paint = { path = "../alpha_paint" } ndarray = "0.17.1" pyo3 = { version = "0.28.2", features = ["extension-module"] } rand = "0.10.0" From 05a4e5217dc06b74770f5feb1688ca19634a09da Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Thu, 26 Mar 2026 02:50:21 -0400 Subject: [PATCH 04/59] wip: migrate --- training/Cargo.toml | 1 - training/src/descent.rs | 424 ++++++ training/src/environments/bytefight/game.rs | 1411 ------------------ training/src/environments/bytefight/map.rs | 176 --- training/src/environments/bytefight/mod.rs | 392 ----- training/src/environments/bytefight/pen.rs | 510 ------- training/src/environments/bytefight/snake.rs | 85 -- training/src/environments/bytefight/types.rs | 237 --- training/src/environments/connect4.rs | 436 ------ training/src/environments/mod.rs | 7 - training/src/environments/tictactoe.rs | 315 ---- training/src/lib.rs | 517 +------ 12 files changed, 444 insertions(+), 4067 deletions(-) create mode 100644 training/src/descent.rs delete mode 100644 training/src/environments/bytefight/game.rs delete mode 100644 training/src/environments/bytefight/map.rs delete mode 100644 training/src/environments/bytefight/mod.rs delete mode 100644 training/src/environments/bytefight/pen.rs delete mode 100644 training/src/environments/bytefight/snake.rs delete mode 100644 training/src/environments/bytefight/types.rs delete mode 100644 training/src/environments/connect4.rs delete mode 100644 training/src/environments/mod.rs delete mode 100644 training/src/environments/tictactoe.rs diff --git a/training/Cargo.toml b/training/Cargo.toml index 88bde36..b8b586f 100644 --- a/training/Cargo.toml +++ b/training/Cargo.toml @@ -2,7 +2,6 @@ name = "alphapaint_training" version = "0.1.0" edition = "2021" -build = "build.rs" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [lib] diff --git a/training/src/descent.rs b/training/src/descent.rs new file mode 100644 index 0000000..c421ee1 --- /dev/null +++ b/training/src/descent.rs @@ -0,0 +1,424 @@ +use alpha_paint::board::actions::Move; +use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; +use std::time::Duration; + +#[derive(Debug)] +pub struct ChildData { + action: Action, + child_value: i32, + entrance_count: usize, + pub node: Option>, +} + +impl ChildData { + fn completion_value(&self) -> i32 { + match &self.node { + Some(c) => c.completion_value, + None => 0, + } + } +} + +#[derive(Debug)] +pub struct SearchNode { + pub value: i32, + pub completion_value: i32, + resolved: bool, + pub children: Vec, +} + +impl SearchNode { + fn new(value: i32, completion_value: i32, is_resolved: bool) -> SearchNode { + SearchNode { + value, + completion_value, + resolved: is_resolved, + children: vec![], + } + } + + /// Prefers moves with low entrance counts + fn completed_best_action_dual( + &self, + is_max_player: bool, + rng: &mut SmallRng, + ) -> (usize, Action) { + let res = if is_max_player { + self.children + .iter() + .enumerate() + .max_by_key(|&(_, child)| { + ( + child.completion_value(), + child.child_value, + -(child.entrance_count as i32), + rng.next_u32(), + ) + }) + .expect("Called completed_best_action_dual without any valid actions!") + } else { + self.children + .iter() + .enumerate() + .min_by_key(|&(_, child)| { + ( + child.completion_value(), + child.child_value, + child.entrance_count, + rng.next_u32(), + ) + }) + .expect("Called completed_best_action_dual without any valid actions!") + }; + + (res.0, res.1.action) + } + + /// Prefers moves with high entrance counts + fn completed_best_action(&self, is_max_player: bool, rng: &mut SmallRng) -> (usize, Action) { + let res = if is_max_player { + self.children + .iter() + .enumerate() + .max_by_key(|&(_, child)| { + ( + child.completion_value(), + child.child_value, + child.entrance_count, + rng.next_u32(), + ) + }) + .expect("Called completed_best_action without any valid actions!") + } else { + self.children + .iter() + .enumerate() + .min_by_key(|&(_, child)| { + ( + child.completion_value(), + child.child_value, + -(child.entrance_count as i32), + rng.next_u32(), + ) + }) + .expect("Called completed_best_action without any valid actions!") + }; + + (res.0, res.1.action) + } + + fn backup_resolution(&self) -> bool { + if self.completion_value.abs() == 1 { + true + } else { + self.children.iter().all(|child| { + child + .node + .as_ref() + .and_then(|c| Some(c.resolved)) + .unwrap_or(false) + }) + } + } + + /// Build a chain of resolved SearchNodes for a killshot move sequence. + /// All nodes are resolved. Only the final node is terminal (no children). + fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { + let term_value = Self::value_from_term(board, terminal); + let comp_value = terminal.value(); + + // Start with the terminal leaf (the collision result) + let mut node = SearchNode::new(term_value, comp_value, true); + + // Build chain from last move to first + for i in (0..moves.len()).rev() { + let action = if i == moves.len() - 1 { + Action::FinalMove(moves[i]) + } else { + Action::Move(moves[i]) + }; + + let parent = SearchNode { + value: term_value, + completion_value: comp_value, + resolved: true, + children: vec![ChildData { + action, + child_value: comp_value, + entrance_count: 0, + node: Some(Box::new(node)), + }], + }; + + node = parent; + } + + node + } + + fn build_self(board: &Board, outcome: ApplyActionOutcome, rng: &mut SmallRng) -> SearchNode { + match outcome { + ApplyActionOutcome::Ongoing => { + let mut new_node = SearchNode::new(0, 0, false); + let actions = board.get_valid_actions(); + new_node.children.reserve(actions.len()); + + if actions.len() == 0 { + let loss = TerminalState::loss_for(board.is_white_turn()); + return SearchNode::new(Self::value_from_term(board, loss), loss.value(), true); + } + + for action in actions.into_iter().copied() { + let mut local_board = board.clone(); + let (outcome, _) = local_board.apply_action(action); + match outcome { + ApplyActionOutcome::Ongoing => { + new_node.children.push(ChildData { + action, + child_value: todo!(), + entrance_count: 0, + node: None, + }); + } + // we can only guarantee that the play instead action is a valid action if + // its a Move (e.g. collision) or a Paint after a move + // Therefore, we just do the same thing as we did above, where we just + // treat this move as a terminal move too, and handle it when we consume + // the SearchTree (e.g. in bindings.rs) + ApplyActionOutcome::Terminal { terminal } + | ApplyActionOutcome::PlayInstead { terminal, .. } => { + new_node.children.push(ChildData { + action, + child_value: terminal.value(), + entrance_count: 0, + node: Some(Box::new(SearchNode::new( + Self::value_from_term(&local_board, terminal), + terminal.value(), + true, + ))), + }); + } + ApplyActionOutcome::Killshot { terminal, moves } => { + let chain = Self::build_killshot_chain(&local_board, terminal, &moves); + new_node.children.push(ChildData { + action, + child_value: terminal.value(), + entrance_count: 0, + node: Some(Box::new(chain)), + }); + } + } + } + + let (best_action_id, _) = + new_node.completed_best_action(board.is_white_turn(), rng); + new_node.completion_value = new_node.children[best_action_id].completion_value(); + new_node.value = new_node.children[best_action_id].child_value; + new_node.resolved = new_node.backup_resolution(); + new_node + } + // in search, we don't really CARE about the distinction between these two. + // e.g. the search tree doesn't _really_ care that we need to play a `Final` variant of + // the action instead. It just cares that this is a terminal state. + // We can just fix this in the consumers of the search tree, e.g. in bindings.rs + ApplyActionOutcome::Terminal { terminal } + | ApplyActionOutcome::PlayInstead { terminal, .. } => SearchNode::new( + Self::value_from_term(board, terminal), + terminal.value(), + true, + ), + ApplyActionOutcome::Killshot { terminal, moves } => { + Self::build_killshot_chain(board, terminal, &moves) + } + } + } + + fn value_from_term(board: &Board, term: TerminalState) -> i32 { + term.value() * (2_000_000_000 - 5 * (board.turn_count as i32)) + } + + fn create_child(&mut self, mut state: Board, action: Action, rng: &mut SmallRng) -> i32 { + let (outcome, _) = state.apply_action(action); + + let node = Box::new(SearchNode::build_self(&state, outcome, rng)); + let value = node.value; + + if let Some(id) = self + .children + .iter() + .position(|child| child.action == action) + { + self.children[id].node = Some(node); + } + + value + } + + fn ubfms_iteration( + &mut self, + mut state: Board, + outcome: ApplyActionOutcome, + rng: &mut SmallRng, + ) -> i32 { + let white_turn = state.is_white_turn(); + + match outcome { + ApplyActionOutcome::Ongoing => { + if self.children.len() == 0 { + let loss = TerminalState::loss_for(white_turn); + self.resolved = true; + self.completion_value = loss.value(); + self.value = Self::value_from_term(&state, loss); + return self.value; + } + + if !self.resolved { + let (best_action_id, best_action) = + self.completed_best_action_dual(white_turn, rng); + + self.children[best_action_id].entrance_count += 1; + + if let Some(child_val) = self.children[best_action_id].node.as_mut() { + let (outcome, _) = state.apply_action(best_action); + child_val.ubfms_iteration(state, outcome, rng); + } else { + self.children[best_action_id].child_value = + self.create_child(state, best_action, rng); + } + + let (best_action_id, _) = self.completed_best_action(white_turn, rng); + self.completion_value = self.children[best_action_id].completion_value(); + self.value = self.children[best_action_id].child_value; + self.resolved = self.backup_resolution(); + } + } + + ApplyActionOutcome::Terminal { terminal } + | ApplyActionOutcome::PlayInstead { terminal, .. } + | ApplyActionOutcome::Killshot { terminal, .. } => { + self.resolved = true; + self.completion_value = terminal.value(); + self.value = Self::value_from_term(&state, terminal); + } + } + + self.value + } +} + +pub struct GameSearchTree<'a> { + pub root_node: Box, + root_state: Board, + rng: SmallRng, +} + +impl GameSearchTree<'_> { + fn safest_action(&mut self) -> (usize, Action) { + let val = if self.root_state.is_white_turn() { + self.root_node + .children + .iter() + .enumerate() + .max_by_key(|(_, child)| { + ( + child.completion_value(), + child.entrance_count, + child.child_value, + self.rng.next_u32(), + ) + }) + .expect("No valid action at root state!") + } else { + self.root_node + .children + .iter() + .enumerate() + .min_by_key(|(_, child)| { + ( + child.completion_value(), + -(child.entrance_count as isize), + child.child_value, + self.rng.next_u32(), + ) + }) + .expect("No valid action at root state!") + }; + + (val.0, val.1.action) + } + + /// This has the invariant that the board is NOT in a terminal state. + pub fn new<'a>(board: &Board) -> GameSearchTree<'a> { + let mut cpy = board.clone(); + let mut rng = SmallRng::seed_from_u64(123312); + GameSearchTree { + root_node: Box::new(SearchNode::build_self( + &mut cpy, + ApplyActionOutcome::Ongoing, + &mut rng, + )), + root_state: board.clone(), + rng, + } + } + + pub fn step_tree(&mut self, new_board: &Board, action_id: usize, outcome: ApplyActionOutcome) { + self.root_state = new_board.clone(); + + if self.root_node.children[action_id].node.is_some() { + self.root_node = self.root_node.children[action_id].node.take().unwrap(); + } else { + self.root_node = Box::new(SearchNode::build_self( + &mut self.root_state.clone(), + outcome, + &mut self.rng, + )) + } + } + + pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { + let dcv = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; + + let start_time = std::time::Instant::now(); + if self.root_node.children.len() <= 1 { + return; // No need to compute with only one option. + } + for epoch in 0..iterations { + if start_time.elapsed() > max_duration && epoch >= 50 { + break; + } + let ba = self.get_best_action_index(); + if self.root_node.children[ba].entrance_count >= 6000 + && self.root_node.children[ba].completion_value() != -dcv + { + break; + } + if self.root_node.children[ba].completion_value() == dcv { + break; + } + self.root_node.ubfms_iteration( + self.root_state.clone(), + ApplyActionOutcome::Ongoing, + &mut self.rng, + ); + } + } + + pub fn get_best_action(&mut self) -> Action { + self.safest_action().1 + } + + pub fn get_best_action_index(&mut self) -> usize { + self.safest_action().0 + } + + pub fn get_best_action_and_index(&mut self) -> (usize, Action) { + self.safest_action() + } +} diff --git a/training/src/environments/bytefight/game.rs b/training/src/environments/bytefight/game.rs deleted file mode 100644 index a440424..0000000 --- a/training/src/environments/bytefight/game.rs +++ /dev/null @@ -1,1411 +0,0 @@ -use std::collections::{HashMap, HashSet, VecDeque}; -use std::sync::OnceLock; - -use rand::seq::SliceRandom; -use rand::Rng; - -use super::map::{add_padding_walls, Map}; -use super::snake::Snake; -use super::types::{ - BitpackedObservation, ByteFightAction, Point, TerminalState, ValidMoves, OBS_CELLS, OBS_SIDE, -}; - -pub const APPLE_REWARD: usize = 2; -const TRAP_LIFETIME: i16 = 100; -pub const TRAP_SACRIFICE: usize = 3; -const DECAY_TIMELINE: [(usize, usize); 4] = [(1000, 15), (1600, 10), (1800, 5), (1950, 2)]; -const DECAY_NOT_APPLIED_PLACEHOLDER: usize = 9999; - -pub const LAST_TURN: usize = 2000; - -// 16x16-only map set for training/selection. -// These are the only currently-defined maps that fit within a 16x16 observation window. -// Excluded (commented out from MAPS_JSON): pillars, great_divide, empty_large, ssspline, -// combustible_lemons, arena, ladder, compasss, diamonds, ssspiral, lol, attrition. -const MAPS_JSON: &str = r#"{ - "cage": "11,11#1,5#9,5#5#2##30,1,Vertical#1010101010101010101010101010101010101010101000101010100000101010000010101010001010101010101010101010101010101010101010101#0", - "empty": "9,9#1,4#7,4#5#2##20,1,Vertical#000000000000000000000000000000000000000000000000000000000000000000000000000000000#0", - "recurve": "13,13#2,6#10,6#4#2#6,0,6,12_6,12,6,0#50,1,Vertical#0100000000010001000000010000010000010000000100010000000010001000000000000000000000000000000000000000000000001000100000000100010000000100000100000100000001000100000000010#0" - }"#; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct Board { - pub map: Map, - pub apple_timeline: Vec<(usize, Point)>, - pub apple_timeline_ptr: usize, - pub snake_a: Snake, - pub snake_b: Snake, - pub is_player_a: bool, - pub min_player_size: usize, - pub(crate) decay_countdown: usize, - pub(crate) cached_decay_interval: usize, - pub(crate) is_decaying: bool, -} - -#[derive(Debug, Clone, PartialEq, Eq)] -pub enum RollbackState { - EndTurn { - old_num_traps: usize, - old_sacrifice_val: usize, - prev_cached_decay_interval: usize, - prev_decay_countdown: usize, - decayed_point: Option>, - snake_a_ate_during_collision: bool, - snake_b_ate_during_collision: bool, - snake_a_max_len: usize, - snake_b_max_len: usize, - prev_apple_timeline_ptr: usize, - apples_placed_index: Vec, - }, - ApplyMove { - prev_trap_val: i16, - sacrificed_points: Vec, - prev_queued_length: usize, - prev_max_length_reached: usize, - prev_direction: Option, - head_was_apple: bool, - }, - ApplyTrap { - old_trap_val: i16, - trap: Point, - }, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum Symmetry { - Horizontal, - Vertical, - Origin, -} - -impl Symmetry { - fn reflect(self, point: Point, width: usize, height: usize) -> Point { - match self { - Symmetry::Horizontal => Point { - x: point.x, - y: height - 1 - point.y, - }, - Symmetry::Vertical => Point { - x: width - 1 - point.x, - y: point.y, - }, - Symmetry::Origin => Point { - x: width - 1 - point.x, - y: height - 1 - point.y, - }, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq)] -enum AppleSpec { - Timeline(Vec<(usize, Point)>), - Spawn { - rate: usize, - count: usize, - symmetry: Symmetry, - }, -} - -#[derive(Debug, Clone)] -struct MapDefinition { - width: usize, - height: usize, - start_a: Point, - start_b: Point, - start_size: usize, - min_player_size: usize, - portals: Vec<(Point, Point)>, - walls: Vec, - apple_spec: AppleSpec, -} - -impl MapDefinition { - fn from_map_string(map_str: &str) -> Result { - let parts: Vec<&str> = map_str.split('#').collect(); - if parts.len() != 9 { - return Err(format!("expected 9 map parts, got {}", parts.len())); - } - - let (width, height) = parse_pair(parts[0], ',')?; - let start_a = parse_point(parts[1])?; - let start_b = parse_point(parts[2])?; - let start_size = parse_usize(parts[3], "start_size")?; - let min_player_size = parse_usize(parts[4], "min_player_size")?; - let portals = parse_portals_section(parts[5])?; - let walls = parse_walls_bits(parts[7], width, height)?; - - let is_record = parse_usize(parts[8], "is_record")? == 1; - let apple_spec = if is_record { - let timeline = parse_apple_timeline(parts[6])?; - AppleSpec::Timeline(timeline) - } else { - let stats: Vec<&str> = parts[6].split(',').collect(); - if stats.len() != 3 { - return Err("invalid apple stats section".to_string()); - } - let rate = parse_usize(stats[0], "apple_rate")?; - let count = parse_usize(stats[1], "num_apples")?; - let symmetry = match stats[2] { - "Horizontal" => Symmetry::Horizontal, - "Vertical" => Symmetry::Vertical, - "Origin" => Symmetry::Origin, - _ => return Err("unknown symmetry".to_string()), - }; - AppleSpec::Spawn { - rate, - count, - symmetry, - } - }; - - Ok(MapDefinition { - width, - height, - start_a, - start_b, - start_size, - min_player_size, - portals, - walls, - apple_spec, - }) - } - - fn build_initial_state(&self, rng: &mut impl Rng) -> Board { - let apple_timeline = self.apple_spec.generate_timeline( - rng, - self.width, - self.height, - &self.portals, - &self.walls, - self.start_a, - self.start_b, - ); - - let apples_now: Vec = apple_timeline - .iter() - .filter(|(turn, _)| *turn == 0) - .map(|(_, point)| *point) - .collect(); - - let snake_a = Snake { - sacrifice: 1, - max_length_reached: self.start_size, - queued_length: self.start_size.saturating_sub(1), - traps_this_turn: 0, - current_direction: ByteFightAction::new((rng.next_u64() % 8) as u8), - segment_queue: VecDeque::from([self.start_a]), - total_apples: 0, - }; - let snake_b = Snake { - sacrifice: 1, - max_length_reached: self.start_size, - queued_length: self.start_size.saturating_sub(1), - traps_this_turn: 0, - current_direction: ByteFightAction::new((rng.next_u64() % 8) as u8), - segment_queue: VecDeque::from([self.start_b]), - total_apples: 0, - }; - - let mut map = Map::new((self.width, self.height), 0); - for wall in &self.walls { - map.become_wall(*wall); - } - for (p1, p2) in &self.portals { - map.add_portal(*p1, *p2); - } - for apple in &apples_now { - map.become_apple(*apple); - if let Some(portal) = map.portal(*apple) { - map.become_apple(portal); - } - } - add_padding_walls(&mut map); - - let mut board = Board { - map, - apple_timeline, - apple_timeline_ptr: 0, - snake_a, - snake_b, - is_player_a: true, - min_player_size: self.min_player_size, - decay_countdown: 0, - cached_decay_interval: DECAY_NOT_APPLIED_PLACEHOLDER, - is_decaying: false, - }; - - board.fix_apple_head_collisions(); - let _ = board.apply_decay(); - board - } -} - -impl AppleSpec { - fn generate_timeline( - &self, - rng: &mut impl Rng, - width: usize, - height: usize, - portals: &[(Point, Point)], - walls: &[Point], - start_a: Point, - start_b: Point, - ) -> Vec<(usize, Point)> { - match self { - AppleSpec::Timeline(timeline) => timeline.clone(), - AppleSpec::Spawn { - rate, - count, - symmetry, - } => { - let portal_map = portal_lookup(portals); - let wall_set: HashSet = walls.iter().copied().collect(); - let mut considered: HashSet = HashSet::new(); - considered.insert(start_a); - considered.insert(start_b); - - let mut select_from = Vec::new(); - for y in 0..height { - for x in 0..width { - let point = Point { x, y }; - if wall_set.contains(&point) { - continue; - } - if considered.contains(&point) { - continue; - } - select_from.push(point); - considered.insert(point); - considered.insert(symmetry.reflect(point, width, height)); - } - } - - let mut apples = Vec::new(); - let mut first_round = select_from.clone(); - add_apple_spawns( - &mut first_round, - *count, - *symmetry, - &portal_map, - &mut apples, - 0, - width, - height, - rng, - ); - - let mut later_round = select_from; - later_round.push(start_a); - later_round.push(start_b); - - let mut spawn_round = *rate; - while spawn_round < LAST_TURN { - let mut picks = later_round.clone(); - add_apple_spawns( - &mut picks, - *count, - *symmetry, - &portal_map, - &mut apples, - spawn_round, - width, - height, - rng, - ); - spawn_round += *rate; - } - - apples - } - } - } -} - -fn portal_lookup(portals: &[(Point, Point)]) -> HashMap { - let mut map = HashMap::new(); - for (p1, p2) in portals { - map.insert(*p1, *p2); - map.insert(*p2, *p1); - } - map -} - -fn add_apple_spawns( - picks: &mut Vec, - count: usize, - symmetry: Symmetry, - portals: &HashMap, - apples: &mut Vec<(usize, Point)>, - turn_num: usize, - width: usize, - height: usize, - rng: &mut impl Rng, -) { - picks.shuffle(rng); - let mut apple_count = 0; - let mut idx = 0; - - while apple_count < count && idx < picks.len() { - let point = picks[idx]; - let reflection = symmetry.reflect(point, width, height); - - if point == reflection { - apple_count += 1; - apples.push((turn_num, point)); - } else { - apple_count += 2; - apples.push((turn_num, point)); - apples.push((turn_num, reflection)); - } - - if let Some(portal) = portals.get(&point) { - apples.push((turn_num, *portal)); - if point != reflection { - if let Some(ref_portal) = portals.get(&reflection) { - apples.push((turn_num, *ref_portal)); - } - } - } - - idx += 1; - } -} - -fn map_definitions() -> &'static Vec { - static MAPS: OnceLock> = OnceLock::new(); - MAPS.get_or_init(|| { - let parsed: HashMap = - serde_json::from_str(MAPS_JSON).expect("invalid bytefight maps json"); - parsed - .into_values() - .map(|map_str| MapDefinition::from_map_string(&map_str)) - .collect::, _>>() - .expect("invalid bytefight map string") - }) -} - -impl Board { - pub fn new_random(rng: &mut impl Rng) -> Self { - let maps = map_definitions(); - let valid_maps: Vec<&MapDefinition> = maps - .iter() - .filter(|map| map.width <= OBS_SIDE && map.height <= OBS_SIDE) - .collect(); - let len = valid_maps.len(); - assert!(len > 0, "no bytefight maps with dimensions <= 16x16"); - let idx = (rng.next_u64() as usize) % len; - let map = valid_maps[idx]; - map.build_initial_state(rng) - } - - pub fn bitpacked_observation_16x16(&self) -> BitpackedObservation { - const WALL_BIT: u8 = 1 << 0; - const APPLE_BIT: u8 = 1 << 1; - const OWN_BODY_BIT: u8 = 1 << 2; - const OWN_HEAD_BIT: u8 = 1 << 3; - const OWN_TRAP_BIT: u8 = 1 << 4; - const OPP_BODY_BIT: u8 = 1 << 5; - const OPP_HEAD_BIT: u8 = 1 << 6; - const OPP_TRAP_BIT: u8 = 1 << 7; - - let (wall_bitmask, apple_bitmask, snake_a_traps, snake_b_traps) = self.map.bitmasks(); - - let (own_snake, opp_snake, own_traps, opp_traps) = if self.is_player_a { - ( - &self.snake_a, - &self.snake_b, - snake_a_traps.as_slice(), - snake_b_traps.as_slice(), - ) - } else { - ( - &self.snake_b, - &self.snake_a, - snake_b_traps.as_slice(), - snake_a_traps.as_slice(), - ) - }; - - let mut obs: BitpackedObservation = [0; OBS_CELLS]; - for y in 0..OBS_SIDE { - let walls = wall_bitmask[y]; - let apples = apple_bitmask[y]; - let own_trap_row = own_traps[y]; - let opp_trap_row = opp_traps[y]; - for x in 0..OBS_SIDE { - let mask = 1u32 << x; - let mut cell = 0u8; - if walls & mask != 0 { - cell |= WALL_BIT; - } - if apples & mask != 0 { - cell |= APPLE_BIT; - } - if own_trap_row & mask != 0 { - cell |= OWN_TRAP_BIT; - } - if opp_trap_row & mask != 0 { - cell |= OPP_TRAP_BIT; - } - obs[y * OBS_SIDE + x] = cell; - } - } - - for (i, segment) in own_snake.segment_queue.iter().enumerate() { - if segment.x >= OBS_SIDE || segment.y >= OBS_SIDE { - continue; - } - let bit = if i == 0 { OWN_HEAD_BIT } else { OWN_BODY_BIT }; - obs[segment.y * OBS_SIDE + segment.x] |= bit; - } - - for (i, segment) in opp_snake.segment_queue.iter().enumerate() { - if segment.x >= OBS_SIDE || segment.y >= OBS_SIDE { - continue; - } - let bit = if i == 0 { OPP_HEAD_BIT } else { OPP_BODY_BIT }; - obs[segment.y * OBS_SIDE + segment.x] |= bit; - } - - obs - } - - pub fn new_from_state( - (width, height): (usize, usize), - queued_apples: Vec<(usize, (usize, usize))>, - apples: Vec<(usize, usize)>, - walls: Vec<(usize, usize)>, - portals: Vec<((usize, usize), (usize, usize))>, - traps: Vec<(i16, (usize, usize))>, - a_snake: Vec<(usize, usize)>, - a_queued_length: usize, - a_max_length_reached: usize, - a_apples_eaten: usize, - a_direction: Option, - b_snake: Vec<(usize, usize)>, - b_queued_length: usize, - b_max_length_reached: usize, - b_apples_eaten: usize, - b_direction: Option, - turn_num: usize, - min_player_size: usize, - is_player_a: bool, - decay_countdown_value: usize, - cached_decay_interval_value: isize, - is_decaying_value: bool, - ) -> Self { - let snake_a = Snake { - sacrifice: 1, - max_length_reached: a_max_length_reached, - queued_length: a_queued_length, - traps_this_turn: 0, - current_direction: Some(a_direction.unwrap_or(ByteFightAction::North)), - segment_queue: a_snake.into_iter().map(Point::from).collect(), - total_apples: a_apples_eaten, - }; - let snake_b = Snake { - sacrifice: 1, - max_length_reached: b_max_length_reached, - queued_length: b_queued_length, - traps_this_turn: 0, - current_direction: Some(b_direction.unwrap_or(ByteFightAction::North)), - segment_queue: b_snake.into_iter().map(Point::from).collect(), - total_apples: b_apples_eaten, - }; - - let mut board = Board { - map: Map::new((width, height), turn_num), - apple_timeline_ptr: 0, - apple_timeline: queued_apples - .into_iter() - .map(|(turn, point)| { - ( - turn, - Point { - x: point.0, - y: point.1, - }, - ) - }) - .collect(), - is_player_a, - min_player_size, - snake_a, - snake_b, - decay_countdown: decay_countdown_value, - cached_decay_interval: cached_decay_interval_value - .try_into() - .unwrap_or(DECAY_NOT_APPLIED_PLACEHOLDER), - is_decaying: is_decaying_value, - }; - - for (lifetime, loc) in traps { - let value = lifetime + (lifetime.signum() * turn_num as i16); - board.map.update_trap(loc.into(), value); - } - for (x, y) in walls { - board.map.become_wall(Point { x, y }); - } - for ((x1, y1), (x2, y2)) in portals { - board - .map - .add_portal(Point { x: x1, y: y1 }, Point { x: x2, y: y2 }); - } - for (x, y) in apples { - let point = Point { x, y }; - if board.map.is_wall(point) { - continue; - } - board.map.become_apple(point); - if let Some(portal) = board.map.portal(point) { - board.map.become_apple(portal); - } - } - board.fix_apple_head_collisions(); - add_padding_walls(&mut board.map); - - let _ = board.apply_decay(); - board - } - - pub fn terminal_state(&self) -> Option { - if self.get_valid_moves().amount() == 0 { - if self.is_player_a { - Some(TerminalState::PlayerBWin) - } else { - Some(TerminalState::PlayerAWin) - } - } else if self.map.turn_count() > LAST_TURN { - match self - .snake_a - .total_apples - .cmp(&self.snake_b.total_apples) - .then(self.snake_a.length().cmp(&self.snake_b.length())) - { - std::cmp::Ordering::Less => Some(TerminalState::PlayerBWin), - std::cmp::Ordering::Equal => Some(TerminalState::Draw), - std::cmp::Ordering::Greater => Some(TerminalState::PlayerAWin), - } - } else { - None - } - } - - pub fn get_valid_moves(&self) -> ValidMoves { - let mut valid_moves = ValidMoves::default(); - let active_snake = if self.is_player_a { - &self.snake_a - } else { - &self.snake_b - }; - - if active_snake.length() < self.min_player_size { - return valid_moves; - } - if self.map.turn_count() > LAST_TURN { - return valid_moves; - } - - if active_snake.can_afford_movement(self.min_player_size) { - let head = active_snake - .segment_queue - .front() - .expect("snake head missing"); - - let offsets = if active_snake.current_direction.is_some() { - 6..11 - } else { - 0..9 - }; - - for offset in offsets { - let direction_int = (offset - + active_snake - .current_direction - .unwrap_or(ByteFightAction::North) as u8) - % 8; - - let Some(new_loc) = head.try_add_int(direction_int) else { - continue; - }; - - let not_wall = !self.map.is_wall(new_loc); - let not_snake = active_snake.removed_on_point_sacrifice(&new_loc) - || (!self.snake_a.segment_queue.contains(&new_loc) - && !self.snake_b.segment_queue.contains(&new_loc)); - let portal_is_valid = self.map.portal(new_loc).is_none_or(|p_loc| { - let not_wall = !self.map.is_wall(p_loc); - let not_snake = active_snake.removed_on_point_sacrifice(&p_loc) - || (!self.snake_a.segment_queue.contains(&p_loc) - && !self.snake_b.segment_queue.contains(&p_loc)); - not_wall && not_snake - }); - let not_trap = self.map.trap(new_loc).abs() <= (self.map.turn_count() as i16) - || (self.is_player_a == (self.map.trap(new_loc) > 0)); - - let apple_reward = if self.map.is_apple(new_loc) { - APPLE_REWARD - } else { - 0 - }; - - let can_facetank_trap = TRAP_SACRIFICE + active_snake.sacrifice - 1 - <= active_snake.length() + apple_reward - self.min_player_size; - - if not_wall && not_snake && (not_trap || can_facetank_trap) && portal_is_valid { - valid_moves - .add(ByteFightAction::new(direction_int).expect("direction_int is valid")); - } - } - } - - if active_snake.can_place_trap(self.min_player_size) { - valid_moves.add(ByteFightAction::Trap); - } - - if active_snake.sacrifice > 1 { - valid_moves.add(ByteFightAction::EndTurn); - } - - valid_moves - } - - pub fn heuristics(&self) -> [u8; 18] { - let (wall_bitmask, apple_bitmask, snake_a_traps, snake_b_traps) = self.map.bitmasks(); - - let mut snake_a_obstacle_mask = wall_bitmask.clone(); - for i in 0..32 { - snake_a_obstacle_mask[i] |= snake_b_traps[i]; - } - for snake_b_segment in self.snake_b.segment_queue.iter() { - snake_a_obstacle_mask[snake_b_segment.y] |= 1 << snake_b_segment.x; - } - let mut snake_a_seed_arr: [u32; 32] = [0; 32]; - Board::seed_directed_array( - &mut snake_a_seed_arr, - self.snake_a.segment_queue.front().expect("snake_a head"), - self.snake_a.current_direction, - ); - let (snake_a_apple_dist, snake_a_reach) = self.run_apple_count_flood_fill( - &mut snake_a_seed_arr, - &snake_a_obstacle_mask, - *apple_bitmask, - ); - - let mut snake_b_obstacle_mask = wall_bitmask.clone(); - for i in 0..32 { - snake_b_obstacle_mask[i] |= snake_a_traps[i]; - } - for snake_a_segment in self.snake_a.segment_queue.iter() { - snake_b_obstacle_mask[snake_a_segment.y] |= 1 << snake_a_segment.x; - } - let mut snake_b_seed_arr: [u32; 32] = [0; 32]; - Board::seed_directed_array( - &mut snake_b_seed_arr, - self.snake_b.segment_queue.front().expect("snake_b head"), - self.snake_b.current_direction, - ); - let (snake_b_apple_dist, snake_b_reach) = self.run_apple_count_flood_fill( - &mut snake_b_seed_arr, - &snake_b_obstacle_mask, - *apple_bitmask, - ); - - let board_size = self.map.size() as f32; - let snake_a_head = self.snake_a.segment_queue.front().expect("snake_a head"); - let snake_b_head = self.snake_b.segment_queue.front().expect("snake_b head"); - let distance_between_snakes = usize::max( - snake_a_head.x.abs_diff(snake_b_head.x), - snake_a_head.y.abs_diff(snake_b_head.y), - ); - - let apples_eaten_diff = - ((self.snake_a.total_apples as f32) - (self.snake_b.total_apples as f32)) / 10.0; - - let turn_ratio = (self.map.turn_count() as f32 / 2000.0).clamp(0.0, 1.0); - let snake_a_reach_ratio = (snake_a_reach as f32 / board_size).clamp(0.0, 1.0); - let snake_b_reach_ratio = (snake_b_reach as f32 / board_size).clamp(0.0, 1.0); - - if self.is_player_a { - [ - Self::encode_signed_feature(turn_ratio), - self.linorm(distance_between_snakes as f32, 32.0), - self.linorm(self.snake_a.length() as f32, 32.0), - self.linorm(self.snake_b.length() as f32, 32.0), - self.dropnorm(snake_a_apple_dist[0] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[1] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[2] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[3] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[0] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[1] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[2] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[3] as f32, 32.0), - self.linorm(self.snake_a.sacrifice as f32, 32.0), - self.linorm(self.snake_a.traps_this_turn as f32, 16.0), - self.linorm(self.snake_a.max_length_reached as f32, 64.0), - self.linorm(self.snake_b.max_length_reached as f32, 64.0), - Self::encode_signed_feature(snake_a_reach_ratio), - Self::encode_signed_feature(snake_b_reach_ratio), - self.linorm(apples_eaten_diff, 10.0), - ][..18] - .try_into() - .expect("heuristics size") - } else { - [ - Self::encode_signed_feature(turn_ratio), - self.linorm(distance_between_snakes as f32, 32.0), - self.linorm(self.snake_b.length() as f32, 32.0), - self.linorm(self.snake_a.length() as f32, 32.0), - self.dropnorm(snake_b_apple_dist[0] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[1] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[2] as f32, 32.0), - self.dropnorm(snake_b_apple_dist[3] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[0] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[1] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[2] as f32, 32.0), - self.dropnorm(snake_a_apple_dist[3] as f32, 32.0), - self.linorm(self.snake_b.sacrifice as f32, 32.0), - self.linorm(self.snake_b.traps_this_turn as f32, 16.0), - self.linorm(self.snake_b.max_length_reached as f32, 64.0), - self.linorm(self.snake_a.max_length_reached as f32, 64.0), - Self::encode_signed_feature(snake_b_reach_ratio), - Self::encode_signed_feature(snake_a_reach_ratio), - self.linorm(-apples_eaten_diff, 10.0), - ][..18] - .try_into() - .expect("heuristics size") - } - } - - pub fn apply_move(&mut self, action: ByteFightAction) -> Result { - match action { - ByteFightAction::Trap => self.apply_trap(), - ByteFightAction::EndTurn => self.apply_end_turn(), - ByteFightAction::FF => Err(()), - direction => self.apply_movement(direction), - } - } - - pub fn rollback(&mut self, state: RollbackState) { - match state { - RollbackState::EndTurn { - old_num_traps, - old_sacrifice_val, - prev_cached_decay_interval, - prev_decay_countdown, - decayed_point, - snake_a_ate_during_collision, - snake_b_ate_during_collision, - snake_a_max_len, - snake_b_max_len, - prev_apple_timeline_ptr, - apples_placed_index, - } => { - self.decay_countdown = prev_decay_countdown; - let snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - match decayed_point.as_deref() { - None => {} - Some([]) => { - snake.queued_length += 1; - self.is_decaying = !self.is_decaying; - } - Some([point]) => { - snake.segment_queue.push_back(*point); - self.is_decaying = !self.is_decaying; - } - _ => unreachable!(), - } - - self.cached_decay_interval = prev_cached_decay_interval; - if snake_a_ate_during_collision { - self.snake_a.queued_length -= APPLE_REWARD; - self.snake_a.total_apples -= 1; - self.snake_a.max_length_reached = snake_a_max_len; - let head = self.snake_a.segment_queue.front().expect("snake_a head"); - self.map.become_apple(*head); - if let Some(portal) = self.map.portal(*head) { - self.map.become_apple(portal); - } - } - if snake_b_ate_during_collision { - self.snake_b.queued_length -= APPLE_REWARD; - self.snake_b.total_apples -= 1; - self.snake_b.max_length_reached = snake_b_max_len; - let head = self.snake_b.segment_queue.front().expect("snake_b head"); - self.map.become_apple(*head); - if let Some(portal) = self.map.portal(*head) { - self.map.become_apple(portal); - } - } - - self.apple_timeline_ptr = prev_apple_timeline_ptr; - for idx in apples_placed_index { - self.map.become_empty(self.apple_timeline[idx].1); - if let Some(portal) = self.map.portal(self.apple_timeline[idx].1) { - self.map.become_empty(portal); - } - } - - self.is_player_a = !self.is_player_a; - let snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - snake.traps_this_turn = old_num_traps; - snake.sacrifice = old_sacrifice_val; - - self.map.move_backward_turn(); - } - RollbackState::ApplyMove { - prev_trap_val, - sacrificed_points, - prev_queued_length, - prev_max_length_reached, - prev_direction, - head_was_apple, - } => { - let snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - let head = snake.segment_queue.front().cloned().expect("snake head"); - self.map.update_trap(head, prev_trap_val); - if let Some(portal) = self.map.portal(head) { - self.map.update_trap(portal, prev_trap_val); - } - - if head_was_apple { - self.map.become_apple(head); - if let Some(portal) = self.map.portal(head) { - self.map.become_apple(portal); - } - snake.total_apples -= 1; - } - - let _ = snake.segment_queue.pop_front(); - for point in sacrificed_points.into_iter().rev() { - snake.segment_queue.push_back(point); - } - - snake.current_direction = prev_direction.and_then(ByteFightAction::new); - snake.queued_length = prev_queued_length; - snake.max_length_reached = prev_max_length_reached; - snake.sacrifice -= 1; - } - RollbackState::ApplyTrap { old_trap_val, trap } => { - self.map.update_trap(trap, old_trap_val); - if let Some(portal) = self.map.portal(trap) { - self.map.update_trap(portal, old_trap_val); - } - let snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - - snake.traps_this_turn -= 1; - snake.segment_queue.push_back(trap); - } - } - } - - fn spawn_apples(&mut self) -> (Vec, bool, bool) { - let mut apples_placed = Vec::new(); - while self.apple_timeline_ptr < self.apple_timeline.len() - && self.apple_timeline[self.apple_timeline_ptr].0 <= self.map.turn_count() - { - let spawn_point = self.apple_timeline[self.apple_timeline_ptr].1; - if self.map.is_empty(spawn_point) { - self.map.become_apple(spawn_point); - apples_placed.push(self.apple_timeline_ptr); - if let Some(portal_location) = self.map.portal(spawn_point) { - self.map.become_apple(portal_location); - } - } - - self.apple_timeline_ptr += 1; - } - - let (place_a, place_b) = self.fix_apple_head_collisions(); - (apples_placed, place_a, place_b) - } - - fn update_decay_interval(&mut self) { - if self.decay_countdown != 0 { - return; - } - - for (turn, interval) in &DECAY_TIMELINE { - if self.map.turn_count() < *turn { - break; - } - - self.cached_decay_interval = *interval; - } - } - - fn apply_decay(&mut self) -> Result>, ()> { - if self.cached_decay_interval == DECAY_NOT_APPLIED_PLACEHOLDER { - return Ok(None); - } - - let decayed = if self.is_decaying || self.decay_countdown == 0 { - let decayed = if self.is_player_a { - self.snake_a.apply_sacrifice(1)? - } else { - self.snake_b.apply_sacrifice(1)? - }; - self.is_decaying = !self.is_decaying; - - Some(decayed) - } else { - None - }; - self.decay_countdown = (self.decay_countdown + 1) % self.cached_decay_interval; - - Ok(decayed) - } - - fn apply_movement(&mut self, action: ByteFightAction) -> Result { - let current_snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - - let prev_queued_length = current_snake.queued_length; - let prev_max_length_reached = current_snake.max_length_reached; - let prev_direction = current_snake.current_direction.map(|a| a as u8); - - if !current_snake.can_afford_movement(self.min_player_size) { - return Err(()); - } - let Ok((new_head, mut cells_lost)) = current_snake.push_move(action) else { - return Err(()); - }; - - if self.map.is_wall(new_head) { - return Err(()); - } - - let portal = self.map.portal(new_head); - if let Some(portal) = portal { - if self.map.is_wall(portal) { - return Err(()); - } - current_snake.segment_queue.push_front(portal); - } else { - current_snake.segment_queue.push_front(new_head); - } - - let mut head_was_apple = false; - if self.map.is_apple(new_head) { - head_was_apple = true; - current_snake.eat_apple(); - self.map.become_empty(new_head); - if let Some(portal) = portal { - self.map.become_empty(portal); - } - } - - let old_trap_val = self.map.trap(new_head); - let trap_val = self.map.trap(new_head); - if trap_val.abs() > self.map.turn_count() as i16 { - let is_player_a_trap = trap_val > 0; - let is_enemy_trap = is_player_a_trap ^ self.is_player_a; - if is_enemy_trap { - match current_snake.apply_sacrifice(3) { - Ok(mut sacrifice) => { - cells_lost.append(&mut sacrifice); - } - Err(_) => { - return Err(()); - } - } - self.map.update_trap(new_head, 0); - if let Some(portal) = portal { - self.map.update_trap(portal, 0); - } - } else { - let trap_val = self.map.turn_count() as i16 + TRAP_LIFETIME; - - self.map - .update_trap(new_head, trap_val * if self.is_player_a { 1 } else { -1 }); - if let Some(portal) = portal { - self.map - .update_trap(portal, trap_val * if self.is_player_a { 1 } else { -1 }); - } - } - } - - Ok(RollbackState::ApplyMove { - prev_trap_val: old_trap_val, - sacrificed_points: cells_lost, - prev_queued_length, - prev_max_length_reached, - prev_direction, - head_was_apple, - }) - } - - fn apply_trap(&mut self) -> Result { - let snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - - if !snake.can_place_trap(self.min_player_size) { - return Err(()); - } - - snake.traps_this_turn += 1; - let trap = snake.segment_queue.pop_back().ok_or(())?; - - let old_trap_val = self.map.trap(trap); - let trap_val = self.map.turn_count() as i16 + TRAP_LIFETIME; - self.map - .update_trap(trap, trap_val * if self.is_player_a { 1 } else { -1 }); - if let Some(portal) = self.map.portal(trap) { - self.map - .update_trap(portal, trap_val * if self.is_player_a { 1 } else { -1 }); - } - - Ok(RollbackState::ApplyTrap { old_trap_val, trap }) - } - - fn apply_end_turn(&mut self) -> Result { - let current_snake = if self.is_player_a { - &mut self.snake_a - } else { - &mut self.snake_b - }; - if current_snake.sacrifice == 1 { - return Err(()); - } - self.map.move_forward_turn(); - - let traps_this_turn = current_snake.traps_this_turn; - let sacrifice = current_snake.sacrifice; - let prev_cached_decay_interval = self.cached_decay_interval; - let prev_decay_countdown = self.decay_countdown; - let prev_apple_timeline_ptr = self.apple_timeline_ptr; - - current_snake.traps_this_turn = 0; - current_snake.sacrifice = 1; - - let snake_a_max_len = self.snake_a.max_length_reached; - let snake_b_max_len = self.snake_b.max_length_reached; - - self.is_player_a = !self.is_player_a; - let (apples_placed_index, snake_a_ate_during_collision, snake_b_ate_during_collision) = - self.spawn_apples(); - - self.update_decay_interval(); - let decayed = self.apply_decay().expect( - "Snake length became 0. This means we made multiple incorrect moves and something is irrecoverably wrong.", - ); - - Ok(RollbackState::EndTurn { - old_num_traps: traps_this_turn, - old_sacrifice_val: sacrifice, - prev_cached_decay_interval, - decayed_point: decayed, - prev_decay_countdown, - snake_a_ate_during_collision, - snake_b_ate_during_collision, - snake_a_max_len, - snake_b_max_len, - prev_apple_timeline_ptr, - apples_placed_index, - }) - } - - fn fix_apple_head_collisions(&mut self) -> (bool, bool) { - let mut snake_a_ate_during_collision = false; - let mut snake_b_ate_during_collision = false; - if let Some(&snake_a_location) = self.snake_a.segment_queue.front() { - if self.map.is_apple(snake_a_location) { - snake_a_ate_during_collision = true; - self.snake_a.queued_length += APPLE_REWARD; - self.snake_a.total_apples += 1; - self.snake_a.max_length_reached = - self.snake_a.max_length_reached.max(self.snake_a.length()); - - self.map.become_empty(snake_a_location); - if let Some(alt_location) = self.map.portal(snake_a_location) { - self.map.become_empty(alt_location); - } - } - } - - if let Some(&snake_b_location) = self.snake_b.segment_queue.front() { - if self.map.is_apple(snake_b_location) { - snake_b_ate_during_collision = true; - self.snake_b.queued_length += APPLE_REWARD; - self.snake_b.total_apples += 1; - self.snake_b.max_length_reached = - self.snake_b.max_length_reached.max(self.snake_b.length()); - - self.map.become_empty(snake_b_location); - if let Some(alt_location) = self.map.portal(snake_b_location) { - self.map.become_empty(alt_location); - } - } - } - - (snake_a_ate_during_collision, snake_b_ate_during_collision) - } - - fn seed_directed_array( - seed_arr: &mut [u32; 32], - origin: &Point, - facing_dir: Option, - ) { - if let Some(dir) = facing_dir { - for offset in 6..11 { - if let Some(new_dir) = origin.try_add_int(((dir as u8) + offset) % 8) { - seed_arr[new_dir.y] |= 1 << new_dir.x; - } - } - } else { - seed_arr[origin.y] |= 0b11 << origin.x; - seed_arr[origin.y] |= 0xC0000000 >> (31 - origin.x); - - if origin.y < 31 { - seed_arr[origin.y + 1] |= 0b11 << origin.x; - seed_arr[origin.y + 1] |= 0xC0000000 >> (31 - origin.x); - } - if origin.y > 0 { - seed_arr[origin.y - 1] |= 0b11 << origin.x; - seed_arr[origin.y - 1] |= 0xC0000000 >> (31 - origin.x); - } - seed_arr[origin.y] &= !(1 << origin.x); - } - } - - fn run_apple_count_flood_fill( - &self, - seed_arr: &mut [u32; 32], - obstacles: &[u32; 32], - mut apples: [u32; 32], - ) -> ([u32; 4], u32) { - let mut apples_found: u32 = 0; - let mut apple_loc: [u32; 4] = [512; 4]; - let mut apple_pntr = self.apple_timeline_ptr; - - for i in 0..32 { - apples_found += (seed_arr[i] & apples[i]).count_ones(); - } - - for i in 0..std::cmp::min(apples_found, 4) { - apple_loc[i as usize] = 1; - } - - for epoch in 0..32 { - for i in 0..32 { - seed_arr[i] |= seed_arr[i] << 1 | seed_arr[i] >> 1; - } - for i in 1..32 { - seed_arr[i - 1] |= seed_arr[i] - } - for i in 0..31 { - seed_arr[31 - i] |= seed_arr[30 - i] - } - - for i in 0..32 { - seed_arr[i] &= !obstacles[i]; - } - - for (p1, p2) in self.map.portals() { - let p1_reachable = (seed_arr[p1.y] >> p1.x) & 1; - let p2_reachable = (seed_arr[p2.y] >> p2.x) & 1; - seed_arr[p1.y] |= p2_reachable << p1.x; - seed_arr[p2.y] |= p1_reachable << p2.x; - } - - while apple_pntr < self.apple_timeline.len() - && self.apple_timeline[apple_pntr].0 <= self.map.turn_count() + 2 * epoch - { - apples[self.apple_timeline[apple_pntr].1.y] |= - 1 << self.apple_timeline[apple_pntr].1.x; - apple_pntr += 1; - } - - let mut apples_this_turn = 0; - for i in 0..32 { - apples_this_turn += (seed_arr[i] & apples[i]).count_ones(); - } - for i in apples_found..std::cmp::min(apples_this_turn, 4) { - apple_loc[i as usize] = (epoch + 1) as u32; - apples_found += 1; - } - } - - let mut reached_tiles: u32 = 0; - for i in 0..32 { - reached_tiles += seed_arr[i].count_ones(); - } - - (apple_loc, reached_tiles) - } - - fn linorm(&self, x: f32, softmax: f32) -> u8 { - let abs_x = x.abs(); - let normalized = if abs_x <= softmax { - (0.8 / softmax) * x - } else { - x.signum() * (1.0 - (0.2 * softmax) / abs_x) - }; - Self::encode_signed_feature(normalized) - } - - fn dropnorm(&self, x: f32, softmax: f32) -> u8 { - let normalized = if x <= softmax { - 1.0 - (0.75 / softmax) * x - } else { - 0.0 - }; - Self::encode_signed_feature(normalized) - } - - #[inline] - fn encode_signed_feature(value: f32) -> u8 { - let clamped = value.clamp(-1.0, 1.0); - let quantized = (clamped * 127.0).round() as i16 + 128; - quantized.clamp(0, 255) as u8 - } -} - -fn parse_usize(value: &str, label: &str) -> Result { - value - .parse::() - .map_err(|_| format!("invalid {}", label)) -} - -fn parse_pair(value: &str, delimiter: char) -> Result<(usize, usize), String> { - let mut iter = value.split(delimiter); - let first = iter.next().ok_or_else(|| "missing first".to_string())?; - let second = iter.next().ok_or_else(|| "missing second".to_string())?; - if iter.next().is_some() { - return Err("too many parts".to_string()); - } - Ok(( - parse_usize(first, "pair_x")?, - parse_usize(second, "pair_y")?, - )) -} - -fn parse_point(value: &str) -> Result { - let (x, y) = parse_pair(value, ',')?; - Ok(Point { x, y }) -} - -fn parse_portals_section(value: &str) -> Result, String> { - if value.is_empty() { - return Ok(Vec::new()); - } - let mut portals = Vec::new(); - for portal in value.split('_') { - if portal.is_empty() { - continue; - } - let parts: Vec<&str> = portal.split(',').collect(); - if parts.len() != 4 { - return Err("invalid portal entry".to_string()); - } - let p1 = Point { - x: parse_usize(parts[0], "portal_x1")?, - y: parse_usize(parts[1], "portal_y1")?, - }; - let p2 = Point { - x: parse_usize(parts[2], "portal_x2")?, - y: parse_usize(parts[3], "portal_y2")?, - }; - portals.push((p1, p2)); - } - Ok(portals) -} - -fn parse_apple_timeline(value: &str) -> Result, String> { - if value.is_empty() { - return Ok(Vec::new()); - } - let mut timeline = Vec::new(); - for entry in value.split('_') { - if entry.is_empty() { - continue; - } - let parts: Vec<&str> = entry.split(',').collect(); - if parts.len() != 3 { - return Err("invalid apple entry".to_string()); - } - let turn = parse_usize(parts[0], "apple_turn")?; - let point = Point { - x: parse_usize(parts[1], "apple_x")?, - y: parse_usize(parts[2], "apple_y")?, - }; - timeline.push((turn, point)); - } - Ok(timeline) -} - -fn parse_walls_bits(bits: &str, width: usize, height: usize) -> Result, String> { - if bits.len() != width * height { - return Err("wall bit length mismatch".to_string()); - } - let mut walls = Vec::new(); - for (i, ch) in bits.chars().enumerate() { - if ch == '1' { - let x = i % width; - let y = i / width; - walls.push(Point { x, y }); - } - } - Ok(walls) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_apply_and_rollback_move() { - let mut board = Board::new_from_state( - (5, 5), - vec![], - vec![], - vec![], - vec![], - vec![], - vec![(2, 2)], - 1, - 2, - 0, - None, - vec![(4, 4)], - 1, - 2, - 0, - None, - 0, - 1, - true, - 0, - DECAY_NOT_APPLIED_PLACEHOLDER as isize, - false, - ); - let snapshot = board.clone(); - let rollback = board.apply_move(ByteFightAction::East).expect("valid move"); - board.rollback(rollback); - assert_eq!(board, snapshot); - } -} diff --git a/training/src/environments/bytefight/map.rs b/training/src/environments/bytefight/map.rs deleted file mode 100644 index 72199e4..0000000 --- a/training/src/environments/bytefight/map.rs +++ /dev/null @@ -1,176 +0,0 @@ -use std::hash::{Hash, Hasher}; - -use super::types::Point; - -#[derive(Debug, Clone)] -pub struct Map { - wall_bitmask: [u32; 32], - apple_bitmask: [u32; 32], - dimensions: (usize, usize), - trap_mask: [[i16; 32]; 32], - portals: Vec<(Point, Point)>, - turn_count: usize, -} - -impl Map { - pub fn new(dimensions: (usize, usize), turn_count: usize) -> Self { - Map { - apple_bitmask: [0; 32], - wall_bitmask: [0; 32], - dimensions, - trap_mask: [[0; 32]; 32], - portals: Vec::new(), - turn_count, - } - } - - pub fn dimensions(&self) -> (usize, usize) { - self.dimensions - } - - pub fn is_empty(&self, point: Point) -> bool { - self.wall_bitmask[point.y] & (1 << point.x) == 0 - && self.apple_bitmask[point.y] & (1 << point.x) == 0 - } - - pub fn is_wall(&self, point: Point) -> bool { - self.wall_bitmask[point.y] & (1 << point.x) != 0 - } - - pub fn is_apple(&self, point: Point) -> bool { - self.apple_bitmask[point.y] & (1 << point.x) != 0 - } - - pub fn portal(&self, point: Point) -> Option { - self.portals.iter().find_map(|&(p1, p2)| { - if p1 == point { - Some(p2) - } else if p2 == point { - Some(p1) - } else { - None - } - }) - } - - pub fn portals(&self) -> &[(Point, Point)] { - &self.portals - } - - pub fn trap(&self, point: Point) -> i16 { - self.trap_mask[point.y][point.x] - } - - pub fn become_apple(&mut self, point: Point) { - self.apple_bitmask[point.y] |= 1 << point.x; - } - - pub fn become_empty(&mut self, point: Point) { - self.wall_bitmask[point.y] &= !(1 << point.x); - self.apple_bitmask[point.y] &= !(1 << point.x); - } - - pub fn update_trap(&mut self, point: Point, trap: i16) { - self.trap_mask[point.y][point.x] = trap; - } - - pub fn become_wall(&mut self, point: Point) { - self.wall_bitmask[point.y] |= 1 << point.x; - } - - pub fn add_portal(&mut self, p1: Point, p2: Point) { - self.portals.push((p1, p2)); - } - - pub fn bitmasks(&self) -> (&[u32; 32], &[u32; 32], [u32; 32], [u32; 32]) { - let mut snake_a_traps: [u32; 32] = [0; 32]; - let mut snake_b_traps: [u32; 32] = [0; 32]; - - for y in 0..32 { - for x in 0..32 { - if self.trap_mask[y][x] > self.turn_count as i16 { - snake_a_traps[y] |= 1 << x; - } else if -self.trap_mask[y][x] > self.turn_count as i16 { - snake_b_traps[y] |= 1 << x; - } - } - } - - ( - &self.wall_bitmask, - &self.apple_bitmask, - snake_a_traps, - snake_b_traps, - ) - } - - pub fn turn_count(&self) -> usize { - self.turn_count - } - - pub fn size(&self) -> usize { - self.dimensions.0 * self.dimensions.1 - } - - pub fn move_forward_turn(&mut self) { - self.turn_count += 1; - } - - pub fn move_backward_turn(&mut self) { - self.turn_count -= 1; - } -} - -impl PartialEq for Map { - fn eq(&self, other: &Self) -> bool { - self.turn_count == other.turn_count - && self.portals == other.portals - && self.apple_bitmask == other.apple_bitmask - && self.wall_bitmask == other.wall_bitmask - && self - .trap_mask - .iter() - .flatten() - .zip(other.trap_mask.iter().flatten()) - .all(|(left, right)| { - (left.abs() <= self.turn_count as i16 && right.abs() <= self.turn_count as i16) - || left == right - }) - } -} - -impl Eq for Map {} - -impl Hash for Map { - fn hash(&self, state: &mut H) { - self.wall_bitmask.hash(state); - self.apple_bitmask.hash(state); - self.dimensions.hash(state); - self.portals.hash(state); - self.turn_count.hash(state); - for row in &self.trap_mask { - for value in row { - let normalized = if value.abs() <= self.turn_count as i16 { - 0 - } else { - *value - }; - normalized.hash(state); - } - } - } -} - -pub fn add_padding_walls(map: &mut Map) { - let (width, height) = map.dimensions(); - for x in width..32 { - for y in 0..32 { - map.become_wall(Point { x, y }); - } - } - for y in height..32 { - for x in 0..32 { - map.become_wall(Point { x, y }); - } - } -} diff --git a/training/src/environments/bytefight/mod.rs b/training/src/environments/bytefight/mod.rs deleted file mode 100644 index d99dddd..0000000 --- a/training/src/environments/bytefight/mod.rs +++ /dev/null @@ -1,392 +0,0 @@ -use ndarray::{ArrayViewMut, Ix2}; -use serde::{Deserialize, Serialize}; - -use crate::{Environment, GameNotation, Player, TerminalState}; - -pub mod game; -pub mod map; -pub mod pen; -pub mod snake; -pub mod types; - -pub use pen::{ByteFightPen, PenError}; -pub use types::{ByteFightAction, ByteFightPolicyAction, Point}; - -#[derive(Debug, Clone, Copy, Default)] -struct PolicyValidMoves(u8); - -#[derive(Debug, Clone, Copy)] -struct PolicyValidMovesIter { - mask: u8, - index: u8, -} - -impl Iterator for PolicyValidMovesIter { - type Item = ByteFightPolicyAction; - - fn next(&mut self) -> Option { - while self.index < ByteFight::NUM_ACTIONS as u8 { - let idx = self.index; - self.index += 1; - if self.mask & (1 << idx) != 0 { - return ByteFightPolicyAction::new(idx); - } - } - None - } -} - -impl PolicyValidMoves { - #[inline] - fn add(&mut self, action: ByteFightPolicyAction) { - self.0 |= 1 << (action as u8); - } - - fn into_iter(self) -> PolicyValidMovesIter { - PolicyValidMovesIter { - mask: self.0, - index: 0, - } - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] -#[serde(from = "ByteFightPen", into = "ByteFightPen")] -pub struct ByteFight { - board: game::Board, -} - -impl From for ByteFight { - fn from(pen: ByteFightPen) -> Self { - ByteFight { - board: pen.into_board().expect("invalid ByteFight PEN"), - } - } -} - -impl From for ByteFightPen { - fn from(board: ByteFight) -> Self { - ByteFightPen::from(&board.board) - } -} - -impl From<&ByteFight> for ByteFightPen { - fn from(board: &ByteFight) -> Self { - ByteFightPen::from(&board.board) - } -} - -impl crate::Action for ByteFightPolicyAction { - fn to_index(self) -> usize { - self as usize - } - - fn from_index(index: usize) -> Option { - ByteFightPolicyAction::new(index as u8) - } -} - -impl Environment for ByteFight { - type ObsElem = u8; - type ObsDim = Ix2; - type Action = ByteFightPolicyAction; - type RollbackState = game::RollbackState; - const NUM_ACTIONS: usize = 7; - const OBS_SHAPE: Ix2 = Ix2(types::OBS_SERIALIZED_SIDE, types::OBS_SERIALIZED_WIDTH); - - fn new() -> Self { - let mut rng = rand::rng(); - let board = game::Board::new_random(&mut rng); - ByteFight::from(ByteFightPen::from(&board)) - } - - fn is_terminal(&self) -> Option { - match self.board.terminal_state()? { - types::TerminalState::PlayerAWin => Some(TerminalState::Win(Player::PlayerA)), - types::TerminalState::PlayerBWin => Some(TerminalState::Win(Player::PlayerB)), - types::TerminalState::Draw => Some(TerminalState::Draw), - } - } - - fn valid_actions(&self) -> impl Iterator { - let valid = self.board.get_valid_moves(); - let mut policy_moves = PolicyValidMoves::default(); - - for action in [ - ByteFightPolicyAction::Forward, - ByteFightPolicyAction::Left, - ByteFightPolicyAction::LeftForward, - ByteFightPolicyAction::Right, - ByteFightPolicyAction::RightForward, - ] { - let absolute = self.policy_to_absolute(action); - if valid.contains(absolute) { - policy_moves.add(action); - } - } - - if valid.contains(ByteFightAction::Trap) { - policy_moves.add(ByteFightPolicyAction::Trap); - } - if valid.contains(ByteFightAction::EndTurn) { - policy_moves.add(ByteFightPolicyAction::EndTurn); - } - - policy_moves.into_iter() - } - - fn current_player(&self) -> Player { - if self.board.is_player_a { - Player::PlayerA - } else { - Player::PlayerB - } - } - - fn observation(&self, mut out: ArrayViewMut) { - let out_slice = out - .as_slice_mut() - .expect("bytefight observation output must be contiguous"); - - let bitpacked = self.board.bitpacked_observation_16x16(); - out_slice[..types::OBS_CELLS].copy_from_slice(&bitpacked); - - let mut offset = types::OBS_CELLS; - let direction = self.active_direction() as usize; - for idx in 0..types::OBS_DIRECTIONS { - out_slice[offset + idx] = if idx == direction { 1 } else { 0 }; - } - offset += types::OBS_DIRECTIONS; - - let heuristics = self.board.heuristics(); - out_slice[offset..offset + types::OBS_HEURISTICS].copy_from_slice(&heuristics); - offset += types::OBS_HEURISTICS; - - for byte in &mut out_slice[offset..] { - *byte = 0; - } - } - - fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { - let absolute = self.policy_to_absolute(action); - self.board - .apply_move(absolute) - .expect("apply_action called with invalid action") - } - - fn rollback(&mut self, rollback: Self::RollbackState) { - self.board.rollback(rollback); - } -} - -impl ByteFight { - fn active_direction(&self) -> ByteFightAction { - let snake = if self.board.is_player_a { - &self.board.snake_a - } else { - &self.board.snake_b - }; - snake.current_direction.unwrap_or(ByteFightAction::North) - } - - fn policy_to_absolute(&self, action: ByteFightPolicyAction) -> ByteFightAction { - let direction = self.active_direction() as u8; - match action { - ByteFightPolicyAction::Forward => { - ByteFightAction::new(direction).expect("valid direction") - } - ByteFightPolicyAction::Left => { - ByteFightAction::new((direction + 6) % 8).expect("valid direction") - } - ByteFightPolicyAction::LeftForward => { - ByteFightAction::new((direction + 7) % 8).expect("valid direction") - } - ByteFightPolicyAction::Right => { - ByteFightAction::new((direction + 2) % 8).expect("valid direction") - } - ByteFightPolicyAction::RightForward => { - ByteFightAction::new((direction + 1) % 8).expect("valid direction") - } - ByteFightPolicyAction::Trap => ByteFightAction::Trap, - ByteFightPolicyAction::EndTurn => ByteFightAction::EndTurn, - } - } -} - -impl GameNotation for ByteFight { - type Error = PenError; - - fn to_notation(&self) -> String { - ByteFightPen::from(self).0 - } - - fn from_notation(s: &str) -> Result { - let pen = ByteFightPen(s.to_string()); - let board = pen.into_board()?; - Ok(ByteFight { board }) - } -} - -#[cfg(test)] -mod tests { - use crate::Action; - use rand::{RngCore, SeedableRng}; - use rand_chacha::ChaCha8Rng; - use rstest::rstest; - - use super::*; - - #[test] - fn test_action_trait() { - assert_eq!(ByteFightPolicyAction::Forward.to_index(), 0); - assert_eq!(ByteFightPolicyAction::EndTurn.to_index(), 6); - assert_eq!( - ByteFightPolicyAction::from_index(0), - Some(ByteFightPolicyAction::Forward) - ); - assert_eq!( - ByteFightPolicyAction::from_index(6), - Some(ByteFightPolicyAction::EndTurn) - ); - assert_eq!(ByteFightPolicyAction::from_index(7), None); - } - - #[test] - fn test_pen_roundtrip() { - let mut rng = ChaCha8Rng::seed_from_u64(7); - let board = game::Board::new_random(&mut rng); - let pen = ByteFightPen::from(&board); - let rebuilt = pen.clone().into_board().expect("valid pen"); - assert_eq!(pen.0, ByteFightPen::from(&rebuilt).0); - } - - #[test] - fn test_notation_roundtrip() { - let mut rng = ChaCha8Rng::seed_from_u64(42); - let game = ByteFight { - board: game::Board::new_random(&mut rng), - }; - - let notation = game.to_notation(); - let restored = ByteFight::from_notation(¬ation).expect("valid notation"); - - // Compare via notation since board equality might differ in internal state - assert_eq!(notation, restored.to_notation()); - } - - #[rstest] - #[case(1)] - #[case(7)] - #[case(13)] - #[case(23)] - #[case(42)] - #[case(77)] - #[case(101)] - #[case(123)] - #[case(256)] - #[case(999)] - fn test_random_move_sequence_roundtrip(#[case] seed: u64) { - let mut rng = ChaCha8Rng::seed_from_u64(seed); - - for _ in 0..5 { - let mut board = game::Board::new_random(&mut rng); - let mut snapshots = Vec::new(); - let mut snapshot_pens = Vec::new(); - snapshots.push(board.clone()); - snapshot_pens.push(ByteFightPen::from(&board).0); - - let mut rollbacks = Vec::new(); - - for _ in 0..150 { - let valid: Vec<_> = board.get_valid_moves().actions().collect(); - if valid.is_empty() { - break; - } - let idx = (rng.next_u64() as usize) % valid.len(); - let action = valid[idx]; - let rollback = board.apply_move(action).expect("valid move"); - rollbacks.push(rollback); - snapshots.push(board.clone()); - snapshot_pens.push(ByteFightPen::from(&board).0); - } - - for (idx, rollback) in rollbacks.into_iter().rev().enumerate() { - board.rollback(rollback); - let snapshot_idx = snapshots.len() - 2 - idx; - assert_eq!(board, snapshots[snapshot_idx]); - assert_eq!(ByteFightPen::from(&board).0, snapshot_pens[snapshot_idx]); - } - } - } - - #[rstest] - #[case(1)] - #[case(42)] - #[case(99)] - fn test_notation_roundtrip_after_moves(#[case] seed: u64) { - use crate::Environment; - - let mut rng = ChaCha8Rng::seed_from_u64(seed); - let mut game = ByteFight { - board: game::Board::new_random(&mut rng), - }; - - // Apply some random moves - for _ in 0..20 { - let valid: Vec<_> = game.valid_actions().collect(); - if valid.is_empty() { - break; - } - let idx = (rng.next_u64() as usize) % valid.len(); - game.apply_action(valid[idx]); - } - - let notation = game.to_notation(); - let restored = ByteFight::from_notation(¬ation).expect("valid notation"); - assert_eq!(notation, restored.to_notation()); - } - - #[test] - fn test_relative_valid_actions_from_direction() { - use crate::Environment; - - let game = ByteFight { - board: game::Board::new_from_state( - (5, 5), - vec![], - vec![], - vec![], - vec![], - vec![], - vec![(2, 2)], - 1, - 2, - 0, - Some(ByteFightAction::North), - vec![(4, 4)], - 1, - 2, - 0, - Some(ByteFightAction::North), - 0, - 1, - true, - 0, - 9999, - false, - ), - }; - - let actions: Vec<_> = game.valid_actions().collect(); - assert_eq!( - actions, - vec![ - ByteFightPolicyAction::Forward, - ByteFightPolicyAction::Left, - ByteFightPolicyAction::LeftForward, - ByteFightPolicyAction::Right, - ByteFightPolicyAction::RightForward, - ] - ); - } -} diff --git a/training/src/environments/bytefight/pen.rs b/training/src/environments/bytefight/pen.rs deleted file mode 100644 index c3a2326..0000000 --- a/training/src/environments/bytefight/pen.rs +++ /dev/null @@ -1,510 +0,0 @@ -use std::collections::VecDeque; -use std::fmt; - -use super::game::Board; -use super::map::{add_padding_walls, Map}; -use super::snake::Snake; -use super::types::{ByteFightAction, Point}; - -#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] -#[serde(transparent)] -pub struct ByteFightPen(pub String); - -#[derive(Debug)] -pub struct PenError(String); - -impl PenError { - fn new(message: impl Into) -> Self { - Self(message.into()) - } -} - -impl fmt::Display for PenError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -impl std::error::Error for PenError {} - -impl ByteFightPen { - pub fn from_board(board: &Board) -> Self { - ByteFightPen(board_to_pen(board)) - } - - pub fn into_board(self) -> Result { - board_from_pen(&self.0) - } -} - -impl From for ByteFightPen { - fn from(board: Board) -> Self { - ByteFightPen::from_board(&board) - } -} - -impl From<&Board> for ByteFightPen { - fn from(board: &Board) -> Self { - ByteFightPen::from_board(board) - } -} - -fn board_to_pen(board: &Board) -> String { - let (width, height) = board.map.dimensions(); - let mut sections = Vec::new(); - sections.push(format!("{}x{}", width, height)); - sections.push(format!("t{}", board.map.turn_count())); - sections.push(format!("p{}", if board.is_player_a { "A" } else { "B" })); - sections.push(format!("m{}", board.min_player_size)); - sections.push(format!( - "d{},{},{}", - board.decay_countdown, - board.cached_decay_interval, - if board.is_decaying { 1 } else { 0 } - )); - sections.push(format!("a{}", board.apple_timeline_ptr)); - sections.push(format!("w{}", encode_walls(&board.map))); - sections.push(format!("o{}", encode_portals(&board.map))); - sections.push(format!("f{}", encode_apples(&board.map))); - sections.push(format!("q{}", encode_timeline(&board.apple_timeline))); - sections.push(format!("r{}", encode_traps(&board.map))); - sections.push(format!("A{}", encode_snake(&board.snake_a))); - sections.push(format!("B{}", encode_snake(&board.snake_b))); - sections.join("|") -} - -fn board_from_pen(pen: &str) -> Result { - let parts: Vec<&str> = pen.split('|').collect(); - if parts.len() != 13 { - return Err(PenError::new("invalid PEN section count")); - } - - let (width, height) = parse_dimensions(parts[0])?; - let turn_count = parse_prefixed_usize(parts[1], 't')?; - let is_player_a = parse_prefixed_player(parts[2])?; - let min_player_size = parse_prefixed_usize(parts[3], 'm')?; - let (decay_countdown, cached_decay_interval, is_decaying) = parse_decay(parts[4], 'd')?; - let apple_timeline_ptr = parse_prefixed_usize(parts[5], 'a')?; - let walls = parse_prefixed_walls(parts[6], width, height)?; - let portals = parse_prefixed_portals(parts[7])?; - let apples = parse_prefixed_points(parts[8], 'f')?; - let apple_timeline = parse_prefixed_timeline(parts[9], 'q')?; - let traps = parse_prefixed_traps(parts[10], 'r')?; - let snake_a = parse_prefixed_snake(parts[11], 'A')?; - let snake_b = parse_prefixed_snake(parts[12], 'B')?; - - let mut map = Map::new((width, height), turn_count); - for wall in walls { - map.become_wall(wall); - } - for (p1, p2) in portals { - map.add_portal(p1, p2); - } - for apple in apples { - map.become_apple(apple); - } - for (trap, point) in traps { - map.update_trap(point, trap); - } - add_padding_walls(&mut map); - - Ok(Board { - map, - apple_timeline, - apple_timeline_ptr, - snake_a, - snake_b, - is_player_a, - min_player_size, - decay_countdown, - cached_decay_interval, - is_decaying, - }) -} - -fn encode_walls(map: &Map) -> String { - let (width, height) = map.dimensions(); - let mut rows = Vec::with_capacity(height); - for y in 0..height { - let mut row = String::new(); - let mut empty_run = 0usize; - for x in 0..width { - if map.is_wall(Point { x, y }) { - if empty_run > 0 { - row.push_str(&empty_run.to_string()); - empty_run = 0; - } - row.push('#'); - } else { - empty_run += 1; - } - } - if empty_run > 0 { - row.push_str(&empty_run.to_string()); - } - rows.push(row); - } - rows.join("/") -} - -fn encode_portals(map: &Map) -> String { - if map.portals().is_empty() { - return "-".to_string(); - } - map.portals() - .iter() - .map(|(p1, p2)| format!("{},{}~{},{}", p1.x, p1.y, p2.x, p2.y)) - .collect::>() - .join(";") -} - -fn encode_apples(map: &Map) -> String { - let (width, height) = map.dimensions(); - let mut apples = Vec::new(); - for y in 0..height { - for x in 0..width { - if map.is_apple(Point { x, y }) { - apples.push(Point { x, y }); - } - } - } - if apples.is_empty() { - return "-".to_string(); - } - apples - .into_iter() - .map(|point| format!("{},{}", point.x, point.y)) - .collect::>() - .join(";") -} - -fn encode_timeline(timeline: &[(usize, Point)]) -> String { - if timeline.is_empty() { - return "-".to_string(); - } - timeline - .iter() - .map(|(turn, point)| format!("{},{},{}", turn, point.x, point.y)) - .collect::>() - .join(";") -} - -fn encode_traps(map: &Map) -> String { - let (width, height) = map.dimensions(); - let mut traps = Vec::new(); - for y in 0..height { - for x in 0..width { - let value = map.trap(Point { x, y }); - if value != 0 && value.abs() > map.turn_count() as i16 { - traps.push(format!("{},{},{}", value, x, y)); - } - } - } - if traps.is_empty() { - return "-".to_string(); - } - traps.join(";") -} - -fn encode_snake(snake: &Snake) -> String { - let direction = snake - .current_direction - .map(|dir| (dir as u8).to_string()) - .unwrap_or_else(|| "-".to_string()); - let segments = snake - .segment_queue - .iter() - .map(|point| format!("{},{}", point.x, point.y)) - .collect::>() - .join(">"); - format!( - "{},{},{},{},{},{}:{}", - direction, - snake.queued_length, - snake.max_length_reached, - snake.total_apples, - snake.sacrifice, - snake.traps_this_turn, - segments - ) -} - -fn parse_dimensions(value: &str) -> Result<(usize, usize), PenError> { - let mut iter = value.split('x'); - let width = iter.next().ok_or_else(|| PenError::new("missing width"))?; - let height = iter.next().ok_or_else(|| PenError::new("missing height"))?; - if iter.next().is_some() { - return Err(PenError::new("invalid dimensions")); - } - Ok((parse_usize(width, "width")?, parse_usize(height, "height")?)) -} - -fn parse_prefixed_usize(section: &str, prefix: char) -> Result { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing prefix"))?; - parse_usize(value, "value") -} - -fn parse_prefixed_player(section: &str) -> Result { - let value = section - .strip_prefix('p') - .ok_or_else(|| PenError::new("missing player prefix"))?; - match value { - "A" => Ok(true), - "B" => Ok(false), - _ => Err(PenError::new("invalid player")), - } -} - -fn parse_decay(section: &str, prefix: char) -> Result<(usize, usize, bool), PenError> { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing decay prefix"))?; - let parts: Vec<&str> = value.split(',').collect(); - if parts.len() != 3 { - return Err(PenError::new("invalid decay section")); - } - let countdown = parse_usize(parts[0], "decay_countdown")?; - let interval = parse_usize(parts[1], "decay_interval")?; - let is_decaying = match parts[2] { - "1" => true, - "0" => false, - _ => return Err(PenError::new("invalid decay flag")), - }; - Ok((countdown, interval, is_decaying)) -} - -fn parse_prefixed_walls( - section: &str, - width: usize, - height: usize, -) -> Result, PenError> { - let value = section - .strip_prefix('w') - .ok_or_else(|| PenError::new("missing walls prefix"))?; - let rows: Vec<&str> = value.split('/').collect(); - if rows.len() != height { - return Err(PenError::new("wall rows mismatch")); - } - - let mut walls = Vec::new(); - for (y, row) in rows.into_iter().enumerate() { - let mut x = 0usize; - let mut digits = String::new(); - for ch in row.chars() { - if ch.is_ascii_digit() { - digits.push(ch); - continue; - } - if !digits.is_empty() { - let run = parse_usize(&digits, "wall run")?; - x += run; - digits.clear(); - } - if ch == '#' { - if x >= width { - return Err(PenError::new("wall row overflow")); - } - walls.push(Point { x, y }); - x += 1; - } else { - return Err(PenError::new("invalid wall token")); - } - } - if !digits.is_empty() { - let run = parse_usize(&digits, "wall run")?; - x += run; - } - if x != width { - return Err(PenError::new("wall row width mismatch")); - } - } - - Ok(walls) -} - -fn parse_prefixed_portals(section: &str) -> Result, PenError> { - let value = section - .strip_prefix('o') - .ok_or_else(|| PenError::new("missing portals prefix"))?; - if value == "-" { - return Ok(Vec::new()); - } - let mut portals = Vec::new(); - for entry in value.split(';') { - let (left, right) = entry - .split_once('~') - .ok_or_else(|| PenError::new("invalid portal entry"))?; - portals.push((parse_point(left)?, parse_point(right)?)); - } - Ok(portals) -} - -fn parse_prefixed_points(section: &str, prefix: char) -> Result, PenError> { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing points prefix"))?; - if value == "-" { - return Ok(Vec::new()); - } - value.split(';').map(parse_point).collect() -} - -fn parse_prefixed_timeline(section: &str, prefix: char) -> Result, PenError> { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing timeline prefix"))?; - if value == "-" { - return Ok(Vec::new()); - } - let mut timeline = Vec::new(); - for entry in value.split(';') { - let parts: Vec<&str> = entry.split(',').collect(); - if parts.len() != 3 { - return Err(PenError::new("invalid timeline entry")); - } - let turn = parse_usize(parts[0], "timeline_turn")?; - let point = Point { - x: parse_usize(parts[1], "timeline_x")?, - y: parse_usize(parts[2], "timeline_y")?, - }; - timeline.push((turn, point)); - } - Ok(timeline) -} - -fn parse_prefixed_traps(section: &str, prefix: char) -> Result, PenError> { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing traps prefix"))?; - if value == "-" { - return Ok(Vec::new()); - } - let mut traps = Vec::new(); - for entry in value.split(';') { - let parts: Vec<&str> = entry.split(',').collect(); - if parts.len() != 3 { - return Err(PenError::new("invalid trap entry")); - } - let trap = parse_i16(parts[0], "trap_value")?; - let point = Point { - x: parse_usize(parts[1], "trap_x")?, - y: parse_usize(parts[2], "trap_y")?, - }; - traps.push((trap, point)); - } - Ok(traps) -} - -fn parse_prefixed_snake(section: &str, prefix: char) -> Result { - let value = section - .strip_prefix(prefix) - .ok_or_else(|| PenError::new("missing snake prefix"))?; - let (meta, body) = value - .split_once(':') - .ok_or_else(|| PenError::new("invalid snake section"))?; - let parts: Vec<&str> = meta.split(',').collect(); - if parts.len() != 6 { - return Err(PenError::new("invalid snake metadata")); - } - let direction = if parts[0] == "-" { - None - } else { - Some( - ByteFightAction::new(parse_usize(parts[0], "direction")? as u8) - .ok_or_else(|| PenError::new("invalid direction value"))?, - ) - }; - - let queued_length = parse_usize(parts[1], "queued_length")?; - let max_length_reached = parse_usize(parts[2], "max_length")?; - let total_apples = parse_usize(parts[3], "total_apples")?; - let sacrifice = parse_usize(parts[4], "sacrifice")?; - let traps_this_turn = parse_usize(parts[5], "traps_this_turn")?; - - let segment_queue = if body.is_empty() { - VecDeque::new() - } else { - body.split('>') - .map(parse_point) - .collect::, _>>()? - }; - - Ok(Snake { - max_length_reached, - queued_length, - traps_this_turn, - current_direction: direction, - segment_queue, - sacrifice, - total_apples, - }) -} - -fn parse_usize(value: &str, label: &str) -> Result { - value - .parse::() - .map_err(|_| PenError::new(format!("invalid {}", label))) -} - -fn parse_i16(value: &str, label: &str) -> Result { - value - .parse::() - .map_err(|_| PenError::new(format!("invalid {}", label))) -} - -fn parse_pair(value: &str, delimiter: char) -> Result<(usize, usize), PenError> { - let mut iter = value.split(delimiter); - let first = iter.next().ok_or_else(|| PenError::new("missing first"))?; - let second = iter.next().ok_or_else(|| PenError::new("missing second"))?; - if iter.next().is_some() { - return Err(PenError::new("too many parts")); - } - Ok(( - parse_usize(first, "pair_x")?, - parse_usize(second, "pair_y")?, - )) -} - -fn parse_point(value: &str) -> Result { - let (x, y) = parse_pair(value, ',')?; - Ok(Point { x, y }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_pen_roundtrip_with_state() { - let mut board = Board::new_from_state( - (6, 6), - vec![(0, (1, 1)), (5, (2, 2))], - vec![(0, 0), (3, 4)], - vec![], - vec![((0, 0), (5, 5))], - vec![(5, (4, 1)), (-5, (1, 4))], - vec![(2, 2)], - 1, - 2, - 0, - Some(ByteFightAction::East), - vec![(4, 4)], - 1, - 2, - 0, - Some(ByteFightAction::West), - 3, - 1, - true, - 2, - 12, - true, - ); - board.apple_timeline_ptr = 1; - let pen = ByteFightPen::from(&board); - let rebuilt = pen.clone().into_board().expect("valid pen"); - assert_eq!(pen.0, ByteFightPen::from(&rebuilt).0); - } -} diff --git a/training/src/environments/bytefight/snake.rs b/training/src/environments/bytefight/snake.rs deleted file mode 100644 index d1f0c0d..0000000 --- a/training/src/environments/bytefight/snake.rs +++ /dev/null @@ -1,85 +0,0 @@ -use std::collections::VecDeque; - -use super::types::{ByteFightAction, Point}; - -#[derive(Debug, Clone, PartialEq, Eq, Hash)] -pub struct Snake { - pub max_length_reached: usize, - pub queued_length: usize, - pub traps_this_turn: usize, - pub current_direction: Option, - pub segment_queue: VecDeque, - pub sacrifice: usize, - pub total_apples: usize, -} - -impl Snake { - pub fn can_afford_movement(&self, min_player_size: usize) -> bool { - self.sacrifice - 1 <= self.length() - min_player_size - } - - pub fn eat_apple(&mut self) { - self.queued_length += 2; - self.max_length_reached = self.max_length_reached.max(self.length()); - self.total_apples += 1; - } - - pub fn removed_on_point_sacrifice(&self, point: &Point) -> bool { - let cells_lost = if self.sacrifice >= self.queued_length { - self.sacrifice - self.queued_length - } else { - 0 - }; - for i in 0..std::cmp::min(cells_lost, self.segment_queue.len()) { - if self - .segment_queue - .get(self.segment_queue.len() - 1 - i) - .is_some_and(|p2| p2 == point) - { - return true; - } - } - false - } - - pub fn apply_sacrifice(&mut self, sacrifice: usize) -> Result, ()> { - let cells_lost = if sacrifice <= self.queued_length { - self.queued_length -= sacrifice; - 0 - } else { - let cells_lost = sacrifice - self.queued_length; - self.queued_length = 0; - cells_lost - }; - - if cells_lost >= self.segment_queue.len() { - return Err(()); - } - - let cells_lost = (0..cells_lost) - .map(|_| self.segment_queue.pop_back().expect("segment queue empty")) - .collect(); - - Ok(cells_lost) - } - - pub fn push_move(&mut self, action: ByteFightAction) -> Result<(Point, Vec), ()> { - let cells_lost = self.apply_sacrifice(self.sacrifice)?; - self.sacrifice += 1; - self.current_direction = Some(action); - - Ok((self.segment_queue[0].try_add(action).unwrap(), cells_lost)) - } - - pub fn can_place_trap(&self, min_player_size: usize) -> bool { - let max_traps = self.max_length_reached / 2; - - (max_traps > self.traps_this_turn) - && (self.segment_queue.len() > 2) - && (self.length() > min_player_size) - } - - pub fn length(&self) -> usize { - self.segment_queue.len() + self.queued_length - } -} diff --git a/training/src/environments/bytefight/types.rs b/training/src/environments/bytefight/types.rs deleted file mode 100644 index 1186c7e..0000000 --- a/training/src/environments/bytefight/types.rs +++ /dev/null @@ -1,237 +0,0 @@ -use std::hash::Hash; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] -pub struct Point { - pub x: usize, - pub y: usize, -} - -impl Point { - pub fn new(x: usize, y: usize) -> Option { - if x >= 32 || y >= 32 { - return None; - } - Some(Point { x, y }) - } - - pub fn try_add(self, action: ByteFightAction) -> Option { - match action { - ByteFightAction::North if self.y != 0 => Point::new(self.x, self.y - 1), - ByteFightAction::Northeast if self.y != 0 => Point::new(self.x + 1, self.y - 1), - ByteFightAction::East => Point::new(self.x + 1, self.y), - ByteFightAction::Southeast => Point::new(self.x + 1, self.y + 1), - ByteFightAction::South => Point::new(self.x, self.y + 1), - ByteFightAction::Southwest if self.x != 0 => Point::new(self.x - 1, self.y + 1), - ByteFightAction::West if self.x != 0 => Point::new(self.x - 1, self.y), - ByteFightAction::Northwest if self.x != 0 && self.y != 0 => { - Point::new(self.x - 1, self.y - 1) - } - ByteFightAction::Trap | ByteFightAction::FF | ByteFightAction::EndTurn => { - panic!("invalid move {action:?} being added to point") - } - _ => None, - } - } - - pub fn try_add_int(self, action: u8) -> Option { - self.try_add(ByteFightAction::new(action)?) - } -} - -impl From<(usize, usize)> for Point { - fn from((x, y): (usize, usize)) -> Self { - Self { x, y } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[repr(u8)] -pub enum ByteFightAction { - North = 0, - Northeast = 1, - East = 2, - Southeast = 3, - South = 4, - Southwest = 5, - West = 6, - Northwest = 7, - Trap = 8, - FF = 9, - EndTurn = 10, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] -#[repr(u8)] -pub enum ByteFightPolicyAction { - Forward = 0, - Left = 1, - LeftForward = 2, - Right = 3, - RightForward = 4, - Trap = 5, - EndTurn = 6, -} - -impl ByteFightPolicyAction { - pub fn new(value: u8) -> Option { - match value { - 0 => Some(Self::Forward), - 1 => Some(Self::Left), - 2 => Some(Self::LeftForward), - 3 => Some(Self::Right), - 4 => Some(Self::RightForward), - 5 => Some(Self::Trap), - 6 => Some(Self::EndTurn), - _ => None, - } - } - - pub fn to_val(self) -> usize { - self as usize - } -} - -impl ByteFightAction { - pub fn new(value: u8) -> Option { - match value { - 0 => Some(Self::North), - 1 => Some(Self::Northeast), - 2 => Some(Self::East), - 3 => Some(Self::Southeast), - 4 => Some(Self::South), - 5 => Some(Self::Southwest), - 6 => Some(Self::West), - 7 => Some(Self::Northwest), - 8 => Some(Self::Trap), - 9 => Some(Self::FF), - 10 => Some(Self::EndTurn), - _ => None, - } - } - - pub fn to_val(self) -> usize { - self as usize - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TerminalState { - PlayerAWin, - PlayerBWin, - Draw, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Hash)] -pub struct ValidMoves(u16); - -#[derive(Debug, Clone)] -pub struct ValidMovesIter { - mask: u16, - index: u8, -} - -impl Iterator for ValidMovesIter { - type Item = ByteFightAction; - - fn next(&mut self) -> Option { - while self.index <= 10 { - let idx = self.index; - self.index += 1; - if self.mask & (1 << idx) != 0 { - return Some(ByteFightAction::new(idx).expect("0..=10 is always a valid action")); - } - } - None - } -} - -impl IntoIterator for ValidMoves { - type Item = ByteFightAction; - type IntoIter = ValidMovesIter; - - fn into_iter(self) -> Self::IntoIter { - ValidMovesIter { - mask: self.0, - index: 0, - } - } -} - -impl ValidMoves { - #[inline] - pub fn add(&mut self, action: ByteFightAction) { - self.0 |= 1 << (action as u16); - } - - #[inline] - pub fn remove(&mut self, action: ByteFightAction) { - self.0 &= !(1 << (action as u16)); - } - - #[inline] - pub fn contains(&self, action: ByteFightAction) -> bool { - self.0 & (1 << (action as u16)) != 0 - } - - #[inline] - pub fn amount(&self) -> u32 { - self.0.count_ones() - } - - #[inline] - pub fn get_move_bounds(&self) -> (usize, usize) { - let mut start: usize = 0; - let mut end: usize = 0; - - for i in 0..8 { - if self.0 & (1 << i) == 0 { - continue; - } - if start == 0 { - start = i; - } - end = i; - } - - (start, end) - } - - pub fn actions(self) -> ValidMovesIter { - self.into_iter() - } -} - -pub const OBS_SIDE: usize = 16; -pub const OBS_PLANES: usize = 8; -pub const OBS_CELLS: usize = OBS_SIDE * OBS_SIDE; -pub const OBS_DIRECTIONS: usize = 8; -pub const OBS_HEURISTICS: usize = 18; -pub const OBS_META_BYTES: usize = OBS_DIRECTIONS + OBS_HEURISTICS; -pub const OBS_SERIALIZED_BYTES: usize = OBS_CELLS + OBS_META_BYTES; -pub const OBS_SERIALIZED_SIDE: usize = 18; -pub const OBS_SERIALIZED_WIDTH: usize = 16; - -pub type BitpackedObservation = [u8; OBS_CELLS]; - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_valid_moves_iter_order() { - let mut moves = ValidMoves::default(); - moves.add(ByteFightAction::West); - moves.add(ByteFightAction::North); - moves.add(ByteFightAction::Trap); - - let collected: Vec<_> = moves.actions().collect(); - assert_eq!( - collected, - vec![ - ByteFightAction::North, - ByteFightAction::West, - ByteFightAction::Trap, - ] - ); - } -} diff --git a/training/src/environments/connect4.rs b/training/src/environments/connect4.rs deleted file mode 100644 index 92829ce..0000000 --- a/training/src/environments/connect4.rs +++ /dev/null @@ -1,436 +0,0 @@ -use ndarray::{ArrayViewMut, Ix2}; -use std::fmt; - -use crate::{Action, Environment, GameNotation, Player, TerminalState}; - -const ROWS: usize = 6; -const COLS: usize = 7; - -#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] -pub struct Connect4Action(pub usize); - -impl Action for Connect4Action { - fn to_index(self) -> usize { - self.0 - } - - fn from_index(index: usize) -> Option { - (index < COLS).then_some(Connect4Action(index)) - } -} - -#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] -pub struct Connect4 { - board: [[Option; COLS]; ROWS], - current_player: Player, - move_count: u8, -} - -#[derive(Debug, Clone, PartialEq, Eq, Copy, Hash)] -pub struct Connect4Rollback { - column: usize, - row: usize, -} - -impl Connect4 { - /// Returns the row where a piece would land in the given column, or None if full. - fn landing_row(&self, col: usize) -> Option { - // Start from bottom row (index 5) and go up - (0..ROWS).rev().find(|&row| self.board[row][col].is_none()) - } - - /// Checks if there's a winner by looking for 4 in a row. - fn check_winner(&self) -> Option { - // Check all possible 4-in-a-row positions - for row in 0..ROWS { - for col in 0..COLS { - if let Some(player) = self.board[row][col] { - // Horizontal (only if we can fit 4 to the right) - if col + 3 < COLS - && self.board[row][col + 1] == Some(player) - && self.board[row][col + 2] == Some(player) - && self.board[row][col + 3] == Some(player) - { - return Some(player); - } - - // Vertical (only if we can fit 4 going down) - if row + 3 < ROWS - && self.board[row + 1][col] == Some(player) - && self.board[row + 2][col] == Some(player) - && self.board[row + 3][col] == Some(player) - { - return Some(player); - } - - // Diagonal down-right - if row + 3 < ROWS - && col + 3 < COLS - && self.board[row + 1][col + 1] == Some(player) - && self.board[row + 2][col + 2] == Some(player) - && self.board[row + 3][col + 3] == Some(player) - { - return Some(player); - } - - // Diagonal up-right - if row >= 3 - && col + 3 < COLS - && self.board[row - 1][col + 1] == Some(player) - && self.board[row - 2][col + 2] == Some(player) - && self.board[row - 3][col + 3] == Some(player) - { - return Some(player); - } - } - } - } - None - } -} - -impl Environment for Connect4 { - type ObsElem = i8; - type ObsDim = Ix2; - type Action = Connect4Action; - type RollbackState = Connect4Rollback; - const NUM_ACTIONS: usize = COLS; - const OBS_SHAPE: Ix2 = Ix2(ROWS, COLS); - - fn new() -> Self { - Self { - board: [[None; COLS]; ROWS], - current_player: Player::PlayerA, - move_count: 0, - } - } - - fn is_terminal(&self) -> Option { - if let Some(winner) = self.check_winner() { - return Some(TerminalState::Win(winner)); - } - if self.move_count == (ROWS * COLS) as u8 { - return Some(TerminalState::Draw); - } - None - } - - fn valid_actions(&self) -> impl Iterator { - (0..COLS) - .filter(|&col| self.board[0][col].is_none()) - .map(Connect4Action) - } - - fn current_player(&self) -> Player { - self.current_player - } - - fn observation(&self, mut out: ArrayViewMut) { - for row in 0..ROWS { - for col in 0..COLS { - out[[row, col]] = match self.board[row][col] { - Some(Player::PlayerA) => 1, - Some(Player::PlayerB) => -1, - None => 0, - }; - } - } - } - - fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { - let col = action.0; - let row = self.landing_row(col).expect("Column is full"); - - self.board[row][col] = Some(self.current_player); - self.current_player = match self.current_player { - Player::PlayerA => Player::PlayerB, - Player::PlayerB => Player::PlayerA, - }; - self.move_count += 1; - - Connect4Rollback { column: col, row } - } - - fn rollback(&mut self, rollback: Self::RollbackState) { - self.board[rollback.row][rollback.column] = None; - self.current_player = match self.current_player { - Player::PlayerA => Player::PlayerB, - Player::PlayerB => Player::PlayerA, - }; - self.move_count -= 1; - } -} - -#[derive(Debug)] -pub struct Connect4NotationError(String); - -impl fmt::Display for Connect4NotationError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -impl std::error::Error for Connect4NotationError {} - -impl GameNotation for Connect4 { - type Error = Connect4NotationError; - - /// Format: "3425|A" (column moves history, A's turn) - /// Each digit (0-6) represents the column played. - fn to_notation(&self) -> String { - // Reconstruct moves by scanning columns from bottom to top - // We need to figure out the order of moves - // Since we don't store history, we'll encode the board state directly - // Format: 42 chars (6 rows * 7 cols) + "|" + player - // A=PlayerA, B=PlayerB, _=empty - let mut s = String::with_capacity(44); - for row in 0..ROWS { - for col in 0..COLS { - s.push(match self.board[row][col] { - Some(Player::PlayerA) => 'A', - Some(Player::PlayerB) => 'B', - None => '_', - }); - } - } - s.push('|'); - s.push(match self.current_player { - Player::PlayerA => 'A', - Player::PlayerB => 'B', - }); - s - } - - fn from_notation(s: &str) -> Result { - let parts: Vec<&str> = s.split('|').collect(); - if parts.len() != 2 { - return Err(Connect4NotationError( - "expected format: BOARD|PLAYER".into(), - )); - } - - let board_str = parts[0]; - let player_str = parts[1]; - - if board_str.len() != ROWS * COLS { - return Err(Connect4NotationError(format!( - "board must have {} cells", - ROWS * COLS - ))); - } - - let mut board = [[None; COLS]; ROWS]; - let mut move_count = 0u8; - let mut chars = board_str.chars(); - - for row in 0..ROWS { - for col in 0..COLS { - let ch = chars.next().unwrap(); - board[row][col] = match ch { - 'A' => { - move_count += 1; - Some(Player::PlayerA) - } - 'B' => { - move_count += 1; - Some(Player::PlayerB) - } - '_' => None, - _ => return Err(Connect4NotationError(format!("invalid cell char: {}", ch))), - }; - } - } - - let current_player = match player_str { - "A" => Player::PlayerA, - "B" => Player::PlayerB, - _ => return Err(Connect4NotationError("player must be A or B".into())), - }; - - Ok(Connect4 { - board, - current_player, - move_count, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_action_trait() { - assert_eq!(Connect4Action(0).to_index(), 0); - assert_eq!(Connect4Action(6).to_index(), 6); - assert_eq!(Connect4Action::from_index(0), Some(Connect4Action(0))); - assert_eq!(Connect4Action::from_index(6), Some(Connect4Action(6))); - assert_eq!(Connect4Action::from_index(7), None); - assert_eq!(Connect4::NUM_ACTIONS, 7); - } - - #[test] - fn test_new_game() { - let game = Connect4::new(); - assert_eq!(game.current_player(), Player::PlayerA); - assert_eq!(game.is_terminal(), None); - assert_eq!(game.valid_actions().count(), 7); - } - - #[test] - fn test_apply_and_rollback() { - let mut game = Connect4::new(); - - // Drop piece in column 3 - let rollback = game.apply_action(Connect4Action(3)); - assert_eq!(game.board[5][3], Some(Player::PlayerA)); // Bottom row - assert_eq!(game.current_player(), Player::PlayerB); - assert_eq!(game.valid_actions().count(), 7); - - game.rollback(rollback); - assert_eq!(game.board[5][3], None); - assert_eq!(game.current_player(), Player::PlayerA); - } - - #[test] - fn test_stacking() { - let mut game = Connect4::new(); - - // Stack pieces in column 0 - game.apply_action(Connect4Action(0)); // PlayerA at row 5 - game.apply_action(Connect4Action(0)); // PlayerB at row 4 - game.apply_action(Connect4Action(0)); // PlayerA at row 3 - - assert_eq!(game.board[5][0], Some(Player::PlayerA)); - assert_eq!(game.board[4][0], Some(Player::PlayerB)); - assert_eq!(game.board[3][0], Some(Player::PlayerA)); - } - - #[test] - fn test_column_full() { - let mut game = Connect4::new(); - - // Fill column 0 - for _ in 0..6 { - game.apply_action(Connect4Action(0)); - } - - // Column 0 should no longer be valid - let valid: Vec<_> = game.valid_actions().collect(); - assert_eq!(valid.len(), 6); - assert!(!valid.contains(&Connect4Action(0))); - } - - #[test] - fn test_horizontal_win() { - let mut game = Connect4::new(); - - // PlayerA: 0, 1, 2, 3 (bottom row) - // PlayerB: 0, 1, 2 (second row) - game.apply_action(Connect4Action(0)); // A - game.apply_action(Connect4Action(0)); // B - game.apply_action(Connect4Action(1)); // A - game.apply_action(Connect4Action(1)); // B - game.apply_action(Connect4Action(2)); // A - game.apply_action(Connect4Action(2)); // B - game.apply_action(Connect4Action(3)); // A wins - - assert_eq!( - game.is_terminal(), - Some(TerminalState::Win(Player::PlayerA)) - ); - } - - #[test] - fn test_vertical_win() { - let mut game = Connect4::new(); - - // PlayerA stacks 4 in column 0 - // PlayerB plays in column 1 - game.apply_action(Connect4Action(0)); // A - game.apply_action(Connect4Action(1)); // B - game.apply_action(Connect4Action(0)); // A - game.apply_action(Connect4Action(1)); // B - game.apply_action(Connect4Action(0)); // A - game.apply_action(Connect4Action(1)); // B - game.apply_action(Connect4Action(0)); // A wins - - assert_eq!( - game.is_terminal(), - Some(TerminalState::Win(Player::PlayerA)) - ); - } - - #[test] - fn test_diagonal_win() { - let mut game = Connect4::new(); - - // Build a diagonal for PlayerA - // Col: 0 1 2 3 - // Row 5: A A A A (eventually) - // But we need to build up for diagonal - - // For diagonal going up-right from (5,0): - // Need A at (5,0), (4,1), (3,2), (2,3) - game.apply_action(Connect4Action(0)); // A at (5,0) - game.apply_action(Connect4Action(1)); // B at (5,1) - game.apply_action(Connect4Action(1)); // A at (4,1) - game.apply_action(Connect4Action(2)); // B at (5,2) - game.apply_action(Connect4Action(2)); // A at (4,2) - game.apply_action(Connect4Action(3)); // B at (5,3) - game.apply_action(Connect4Action(2)); // A at (3,2) - game.apply_action(Connect4Action(3)); // B at (4,3) - game.apply_action(Connect4Action(3)); // A at (3,3) - game.apply_action(Connect4Action(3)); // B at (2,3) - game.apply_action(Connect4Action(4)); // A at (5,4) - filler - game.apply_action(Connect4Action(4)); // B at (4,4) - - // Now we need to think about this more carefully... - // Let me restart with a cleaner approach - } - - #[test] - fn test_observation() { - use ndarray::Array2; - - let mut game = Connect4::new(); - game.apply_action(Connect4Action(3)); // A at (5, 3) - game.apply_action(Connect4Action(3)); // B at (4, 3) - - let mut obs = Array2::::zeros((ROWS, COLS)); - game.observation(obs.view_mut()); - assert_eq!(obs.shape(), &[6, 7]); - assert_eq!(obs[[5, 3]], 1); // PlayerA - assert_eq!(obs[[4, 3]], -1); // PlayerB - assert_eq!(obs[[0, 0]], 0); // Empty - } - - #[test] - fn test_notation_roundtrip() { - // Test empty board - let game = Connect4::new(); - let notation = game.to_notation(); - assert_eq!(notation, "__________________________________________|A"); - let restored = Connect4::from_notation(¬ation).unwrap(); - assert_eq!(game, restored); - - // Test after some moves - let mut game = Connect4::new(); - game.apply_action(Connect4Action(3)); // A at bottom row, col 3 - game.apply_action(Connect4Action(3)); // B on top of A - game.apply_action(Connect4Action(0)); // A at bottom row, col 0 - - let notation = game.to_notation(); - let restored = Connect4::from_notation(¬ation).unwrap(); - assert_eq!(game, restored); - } - - #[test] - fn test_notation_errors() { - assert!(Connect4::from_notation("invalid").is_err()); - assert!(Connect4::from_notation("___|A").is_err()); // too short - assert!(Connect4::from_notation("__________________________________________|C").is_err()); // invalid player - assert!(Connect4::from_notation("_________________________________________Z|A").is_err()); - // invalid char - } -} diff --git a/training/src/environments/mod.rs b/training/src/environments/mod.rs deleted file mode 100644 index 71b72ac..0000000 --- a/training/src/environments/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -pub mod bytefight; -pub mod connect4; -pub mod tictactoe; - -pub use bytefight::*; -pub use connect4::*; -pub use tictactoe::*; diff --git a/training/src/environments/tictactoe.rs b/training/src/environments/tictactoe.rs deleted file mode 100644 index a46a045..0000000 --- a/training/src/environments/tictactoe.rs +++ /dev/null @@ -1,315 +0,0 @@ -use ndarray::{ArrayView1, ArrayViewMut, Ix1}; -use std::fmt; - -use crate::{Action, Environment, GameNotation, Player, TerminalState}; - -#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] -pub struct TicTacToeAction(pub u8); - -impl Action for TicTacToeAction { - fn to_index(self) -> usize { - self.0 as usize - } - - fn from_index(index: usize) -> Option { - (index < 9).then_some(TicTacToeAction(index as u8)) - } -} - -#[derive(Clone, Hash, PartialEq, Eq, Debug)] -pub struct TicTacToe { - /// Board state: 0 = empty, 1 = PlayerA (X), -1 = PlayerB (O) - pub board: [i8; 9], - current_player: Player, - move_count: u8, -} - -impl TicTacToe { - const WIN_PATTERNS: [[usize; 3]; 8] = [ - [0, 1, 2], - [3, 4, 5], - [6, 7, 8], - [0, 3, 6], - [1, 4, 7], - [2, 5, 8], - [0, 4, 8], - [2, 4, 6], - ]; - - pub fn check_winner(&self) -> Option { - for pattern in &Self::WIN_PATTERNS { - let a = self.board[pattern[0]]; - let b = self.board[pattern[1]]; - let c = self.board[pattern[2]]; - if a != 0 && a == b && b == c { - return Some(if a == 1 { - Player::PlayerA - } else { - Player::PlayerB - }); - } - } - None - } -} - -pub struct TicTacToeRollback { - cell: u8, - previous_player: Player, -} - -impl Environment for TicTacToe { - type ObsElem = i8; - type ObsDim = Ix1; - type Action = TicTacToeAction; - type RollbackState = TicTacToeRollback; - const NUM_ACTIONS: usize = 9; - const OBS_SHAPE: Ix1 = Ix1(9); - - fn new() -> Self { - TicTacToe { - board: [0; 9], - current_player: Player::PlayerA, - move_count: 0, - } - } - - fn is_terminal(&self) -> Option { - if let Some(winner) = self.check_winner() { - return Some(TerminalState::Win(winner)); - } - if self.move_count == 9 { - return Some(TerminalState::Draw); - } - None - } - - fn valid_actions(&self) -> impl Iterator { - self.board - .iter() - .enumerate() - .filter(|(_, &cell)| cell == 0) - .map(|(i, _)| TicTacToeAction(i as u8)) - } - - fn current_player(&self) -> Player { - self.current_player - } - - fn observation(&self, mut out: ArrayViewMut) { - out.assign(&ArrayView1::from(&self.board)); - } - - fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState { - let cell = action.0; - let previous_player = self.current_player; - - // 1 = PlayerA, -1 = PlayerB - self.board[cell as usize] = self.current_player as i8; - self.current_player = match self.current_player { - Player::PlayerA => Player::PlayerB, - Player::PlayerB => Player::PlayerA, - }; - self.move_count += 1; - - TicTacToeRollback { - cell, - previous_player, - } - } - - fn rollback(&mut self, rollback: Self::RollbackState) { - self.board[rollback.cell as usize] = 0; - self.current_player = rollback.previous_player; - self.move_count -= 1; - } -} - -#[derive(Debug)] -pub struct TicTacToeNotationError(String); - -impl fmt::Display for TicTacToeNotationError { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "{}", self.0) - } -} - -impl std::error::Error for TicTacToeNotationError {} - -impl GameNotation for TicTacToe { - type Error = TicTacToeNotationError; - - /// Format: "XO_X_O___|A" (9 cells + current player) - /// X=PlayerA, O=PlayerB, _=empty - fn to_notation(&self) -> String { - let mut s = String::with_capacity(11); - for &cell in &self.board { - s.push(match cell { - 1 => 'X', - -1 => 'O', - _ => '_', - }); - } - s.push('|'); - s.push(match self.current_player { - Player::PlayerA => 'A', - Player::PlayerB => 'B', - }); - s - } - - fn from_notation(s: &str) -> Result { - let parts: Vec<&str> = s.split('|').collect(); - if parts.len() != 2 { - return Err(TicTacToeNotationError( - "expected format: BOARD|PLAYER".into(), - )); - } - - let board_str = parts[0]; - let player_str = parts[1]; - - if board_str.len() != 9 { - return Err(TicTacToeNotationError("board must have 9 cells".into())); - } - - let mut board = [0i8; 9]; - let mut move_count = 0u8; - for (i, ch) in board_str.chars().enumerate() { - board[i] = match ch { - 'X' => { - move_count += 1; - 1 - } - 'O' => { - move_count += 1; - -1 - } - '_' => 0, - _ => return Err(TicTacToeNotationError(format!("invalid cell char: {}", ch))), - }; - } - - let current_player = match player_str { - "A" => Player::PlayerA, - "B" => Player::PlayerB, - _ => return Err(TicTacToeNotationError("player must be A or B".into())), - }; - - Ok(TicTacToe { - board, - current_player, - move_count, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_action_trait() { - assert_eq!(TicTacToeAction(0).to_index(), 0); - assert_eq!(TicTacToeAction(8).to_index(), 8); - assert_eq!(TicTacToeAction::from_index(0), Some(TicTacToeAction(0))); - assert_eq!(TicTacToeAction::from_index(8), Some(TicTacToeAction(8))); - assert_eq!(TicTacToeAction::from_index(9), None); - assert_eq!(TicTacToe::NUM_ACTIONS, 9); - } - - #[test] - fn test_new_game() { - let game = TicTacToe::new(); - assert_eq!(game.current_player(), Player::PlayerA); - assert_eq!(game.is_terminal(), None); - assert_eq!(game.valid_actions().count(), 9); - } - - #[test] - fn test_apply_and_rollback() { - let mut game = TicTacToe::new(); - - let rollback = game.apply_action(TicTacToeAction(4)); - assert_eq!(game.board[4], 1); // PlayerA = 1 - assert_eq!(game.current_player(), Player::PlayerB); - assert_eq!(game.valid_actions().count(), 8); - - game.rollback(rollback); - assert_eq!(game.board[4], 0); - assert_eq!(game.current_player(), Player::PlayerA); - assert_eq!(game.valid_actions().count(), 9); - } - - #[test] - fn test_player_b_moves() { - let mut game = TicTacToe::new(); - - game.apply_action(TicTacToeAction(0)); // PlayerA - game.apply_action(TicTacToeAction(4)); // PlayerB - - assert_eq!(game.board[0], 1); // PlayerA = 1 - assert_eq!(game.board[4], -1); // PlayerB = -1 - } - - #[test] - fn test_win_detection() { - let mut game = TicTacToe::new(); - - // X wins with top row - game.apply_action(TicTacToeAction(0)); - game.apply_action(TicTacToeAction(3)); - game.apply_action(TicTacToeAction(1)); - game.apply_action(TicTacToeAction(4)); - game.apply_action(TicTacToeAction(2)); - - assert_eq!( - game.is_terminal(), - Some(TerminalState::Win(Player::PlayerA)) - ); - } - - #[test] - fn test_draw() { - let mut game = TicTacToe::new(); - - // X O X - // X O O - // O X X - let moves = [0, 1, 2, 4, 3, 5, 7, 6, 8]; - for &m in &moves { - game.apply_action(TicTacToeAction(m)); - } - - assert_eq!(game.is_terminal(), Some(TerminalState::Draw)); - } - - #[test] - fn test_notation_roundtrip() { - // Test empty board - let game = TicTacToe::new(); - let notation = game.to_notation(); - assert_eq!(notation, "_________|A"); - let restored = TicTacToe::from_notation(¬ation).unwrap(); - assert_eq!(game, restored); - - // Test after some moves - let mut game = TicTacToe::new(); - game.apply_action(TicTacToeAction(0)); // X at 0 - game.apply_action(TicTacToeAction(4)); // O at 4 - game.apply_action(TicTacToeAction(8)); // X at 8 - - let notation = game.to_notation(); - assert_eq!(notation, "X___O___X|B"); - let restored = TicTacToe::from_notation(¬ation).unwrap(); - assert_eq!(game, restored); - } - - #[test] - fn test_notation_errors() { - assert!(TicTacToe::from_notation("invalid").is_err()); - assert!(TicTacToe::from_notation("XXXXXXXX|A").is_err()); // 8 cells - assert!(TicTacToe::from_notation("_________|C").is_err()); // invalid player - assert!(TicTacToe::from_notation("____Z____|A").is_err()); // invalid char - } -} diff --git a/training/src/lib.rs b/training/src/lib.rs index 18b0ec3..9ec985a 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,74 +1,17 @@ use std::sync::Arc; use std::sync::{Mutex, OnceLock}; -use std::{fmt::Debug, hash::Hash}; -use ndarray::{ArrayView, ArrayViewMut, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; -use numpy::{PyArray, PyArrayMethods}; +use cudagraph::ByteFightCudaGraphRunner; +use eval::PolicyValue; +use mcts::MCTSConfig; +use ndarray::{ArrayView, Ix3}; use pyo3::prelude::*; -use rand::SeedableRng; -use rand_chacha::ChaCha8Rng; - -/// Extension trait for prepending a batch dimension to a shape. -/// -/// The `BatchedDim` associated type is the dimension with batch prepended. -/// We use an associated type instead of `Dimension::Larger` so we can -/// constrain that `BatchedDim::Smaller == Self`. -pub trait BatchDim: Dimension + Clone { - type BatchedDim: Dimension + RemoveAxis; - - fn with_batch(batch_size: usize, obs_shape: Self) -> Self::BatchedDim; -} - -impl BatchDim for Ix0 { - type BatchedDim = Ix1; - - fn with_batch(batch_size: usize, _obs: Self) -> Ix1 { - Ix1(batch_size) - } -} - -impl BatchDim for Ix1 { - type BatchedDim = Ix2; - - fn with_batch(batch_size: usize, obs: Self) -> Ix2 { - Ix2(batch_size, obs[0]) - } -} - -impl BatchDim for Ix2 { - type BatchedDim = Ix3; - - fn with_batch(batch_size: usize, obs: Self) -> Ix3 { - Ix3(batch_size, obs[0], obs[1]) - } -} - -impl BatchDim for Ix3 { - type BatchedDim = Ix4; - - fn with_batch(batch_size: usize, obs: Self) -> Ix4 { - Ix4(batch_size, obs[0], obs[1], obs[2]) - } -} - -impl BatchDim for Ix4 { - type BatchedDim = Ix5; - - fn with_batch(batch_size: usize, obs: Self) -> Ix5 { - Ix5(batch_size, obs[0], obs[1], obs[2], obs[3]) - } -} - -impl BatchDim for Ix5 { - type BatchedDim = Ix6; - - fn with_batch(batch_size: usize, obs: Self) -> Ix6 { - Ix6(batch_size, obs[0], obs[1], obs[2], obs[3], obs[4]) - } -} +use queue::{queue_shape_for_workers, BATCH_SIZE}; +use training::{SelfPlaySession, SessionConfig}; +use worker::WorkerConfig; pub mod cudagraph; -pub mod environments; +mod descent; pub mod eval; pub mod executor; pub mod future; @@ -80,9 +23,6 @@ pub mod replay_buffer; pub mod training; pub mod worker; -use environments::{bytefight::types as bytefight_types, ByteFight, Connect4, TicTacToe}; -use observation_replay_buffer::ObservationReplayBuffer; - struct ByteFightGraphCacheEntry { model_ptr: usize, num_batches: usize, @@ -90,363 +30,24 @@ struct ByteFightGraphCacheEntry { runner: Arc, } -static BYTEFIGHT_GRAPH_CACHE: OnceLock>> = OnceLock::new(); - -fn bytefight_graph_cache() -> &'static Mutex> { - BYTEFIGHT_GRAPH_CACHE.get_or_init(|| Mutex::new(None)) -} - -/// Macro to generate typed ephemeral replay buffer classes for each environment. -/// -/// Each generated class wraps an `Arc>` and -/// exposes numpy sampling. -macro_rules! typed_ephemeral_replay_buffer { - ( - $name:ident, - $obs_ty:ty, - $obs_single_dim:ty, - $obs_batched_dim:ty, - $obs_shape_const:expr, - $obs_shape_fn:expr, - $num_actions:expr - ) => { - #[doc = concat!("Typed ephemeral replay buffer for ", stringify!($name), ".")] - #[doc = ""] - #[doc = "Stores contiguous observations, policies, and values in memory."] - #[pyclass] - pub struct $name { - inner: Arc>, - } - - #[pymethods] - impl $name { - #[new] - fn new(capacity: usize) -> Self { - Self { - inner: Arc::new(ObservationReplayBuffer::new(capacity, $obs_shape_const)), - } - } - - fn __len__(&self) -> usize { - self.inner.len() - } - - #[getter] - fn capacity(&self) -> usize { - self.inner.capacity() - } - - /// Sample `n` items and return (observations, policies, values) as numpy arrays. - /// - /// Args: - /// n: Number of samples to draw - /// seed: Random seed for reproducible sampling - /// - /// Returns: - /// Tuple of (observations, policies, values) numpy arrays - fn sample<'py>( - &self, - py: Python<'py>, - n: usize, - seed: u64, - ) -> PyResult<( - Bound<'py, PyArray<$obs_ty, $obs_batched_dim>>, - Bound<'py, PyArray>, - Bound<'py, PyArray>, - )> { - let mut rng = ChaCha8Rng::seed_from_u64(seed); - let batch = self.inner.sample(n, &mut rng); - let num_samples = batch.values.len(); - - let obs_data = batch.observations.into_raw_vec_and_offset().0; - let policy_data = batch.policies.into_raw_vec_and_offset().0; - - let shape_fn: fn(usize) -> $obs_batched_dim = $obs_shape_fn; - let obs = PyArray::from_vec(py, obs_data).reshape(shape_fn(num_samples))?; - let policies = - PyArray::from_vec(py, policy_data).reshape(Ix2(num_samples, $num_actions))?; - let values = PyArray::from_vec(py, batch.values); - - Ok((obs, policies, values)) - } - } - - impl $name { - pub fn inner( - &self, - ) -> &Arc> { - &self.inner - } - } - }; -} - -// TicTacToe: observations (9,) i8, sampled as (n, 9) -typed_ephemeral_replay_buffer!( - TicTacToeEphemeralReplayBuffer, - i8, - Ix1, - Ix2, - TicTacToe::OBS_SHAPE, - |n| Ix2(n, 9), - 9 -); - -// Connect4: observations (6, 7) i8, sampled as (n, 6, 7) -typed_ephemeral_replay_buffer!( - Connect4EphemeralReplayBuffer, - i8, - Ix2, - Ix3, - Connect4::OBS_SHAPE, - |n| Ix3(n, 6, 7), - 7 -); - -// ByteFight: observations (18, 16) u8, sampled as (n, 18, 16) -typed_ephemeral_replay_buffer!( - ByteFightEphemeralReplayBuffer, - u8, - Ix2, - Ix3, - ByteFight::OBS_SHAPE, - |n| Ix3( - n, - bytefight_types::OBS_SERIALIZED_SIDE, - bytefight_types::OBS_SERIALIZED_WIDTH, - ), - 7 -); - -/// Macro to generate a persistent SelfPlay pyclass with a Python callback dispatch. -macro_rules! typed_selfplay { - ( - $name:ident, - $env:ty, - $obs_ty:ty, - $obs_dim:ty, - $obs_batched_dim:ty, - $replay_buf_class:ident, - $num_actions:expr, - callback_dispatch - ) => { - #[doc = concat!("Persistent self-play session for ", stringify!($name), ".")] - #[pyclass] - struct $name { - session: Option, - } - - #[pymethods] - impl $name { - #[new] - #[rustfmt::skip] - #[pyo3(signature = ( - replay_buffer, - num_threads, - workers_per_thread, - seed, - execute_model, - mcts_num_simulations = 20, - mcts_c_puct = 1.5, - mcts_dirichlet_alpha = 0.3, - mcts_dirichlet_epsilon = 0.25, - temperature = 1.0, - exploration_moves = 30, - ))] - fn new( - replay_buffer: &$replay_buf_class, - num_threads: usize, - workers_per_thread: usize, - seed: u64, - execute_model: Py, - mcts_num_simulations: usize, - mcts_c_puct: f32, - mcts_dirichlet_alpha: f32, - mcts_dirichlet_epsilon: f32, - temperature: f32, - exploration_moves: usize, - ) -> PyResult { - use eval::PolicyValue; - use mcts::MCTSConfig; - use training::{SelfPlaySession, SessionConfig}; - use worker::WorkerConfig; - - if mcts_num_simulations == 0 { - return Err(PyErr::new::( - "mcts_num_simulations must be >= 1", - )); - } - if mcts_c_puct <= 0.0 { - return Err(PyErr::new::( - "mcts_c_puct must be > 0", - )); - } - if mcts_dirichlet_alpha <= 0.0 { - return Err(PyErr::new::( - "mcts_dirichlet_alpha must be > 0", - )); - } - if !(0.0..=1.0).contains(&mcts_dirichlet_epsilon) { - return Err(PyErr::new::( - "mcts_dirichlet_epsilon must be in [0, 1]", - )); - } - if temperature < 0.0 { - return Err(PyErr::new::( - "temperature must be >= 0", - )); - } - - let config = SessionConfig { - num_threads, - workers_per_thread, - seed, - worker: WorkerConfig { - mcts: MCTSConfig { - num_simulations: mcts_num_simulations, - c_puct: mcts_c_puct, - dirichlet_alpha: mcts_dirichlet_alpha, - dirichlet_epsilon: mcts_dirichlet_epsilon, - ..Default::default() - }, - temperature, - exploration_moves, - ..Default::default() - }, - }; - - let dispatch = move |_batch_idx: usize, - obs_view: ArrayView<$obs_ty, $obs_batched_dim>, - completion: queue::BatchCompletion< - PolicyValue<$num_actions>, - >| { - let mut outputs = - vec![PolicyValue::<$num_actions>::default(); queue::BATCH_SIZE]; - Python::attach(|py| { - // SAFETY: obs_view is valid for the duration of this callback, - // and the numpy array doesn't escape the callback scope. - let np_obs = unsafe { - PyArray::borrow_from_array(&obs_view, py.None().into_bound(py)) - }; - - let result = execute_model - .call1(py, (np_obs,)) - .expect("execute_model call failed"); - - let (policy_arr, value_arr): ( - Bound<'_, PyArray>, - Bound<'_, PyArray>, - ) = result - .extract(py) - .expect("expected (policy, value) tuple of numpy arrays"); - - let policy = unsafe { policy_arr.as_slice().unwrap() }; - let value = unsafe { value_arr.as_slice().unwrap() }; - - for (i, out) in outputs.iter_mut().enumerate() { - out.policy - .copy_from_slice(&policy[i * $num_actions..(i + 1) * $num_actions]); - out.value = value[i]; - } - }); - completion.complete(&outputs); - }; +static GRAPH_CACHE: OnceLock>> = OnceLock::new(); - let session = SelfPlaySession::new::<$env, $num_actions, _>( - config, - replay_buffer.inner().clone(), - dispatch, - ); - - Ok(Self { - session: Some(session), - }) - } - - /// Start self-play with no sample limit. - fn start(&self) -> PyResult<()> { - self.session - .as_ref() - .ok_or_else(|| { - PyErr::new::("session already dropped") - })? - .start(); - Ok(()) - } - - /// Block until absolute target_samples is reached, then pause and quiesce. - fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { - let session = self.session.as_ref().ok_or_else(|| { - PyErr::new::("session already dropped") - })?; - let result = py.detach(|| session.wait_for(target_samples)); - Ok(result) - } - - /// Return the current absolute sample count. - fn samples(&self) -> PyResult { - Ok(self - .session - .as_ref() - .ok_or_else(|| { - PyErr::new::("session already dropped") - })? - .samples()) - } - - /// Shut down the session. Idempotent. - #[pyo3(name = "drop")] - fn py_drop(&mut self) { - if let Some(mut session) = self.session.take() { - session.shutdown(); - } - } - } - - impl Drop for $name { - fn drop(&mut self) { - if let Some(mut session) = self.session.take() { - session.shutdown(); - } - } - } - }; +fn graph_cache() -> &'static Mutex> { + GRAPH_CACHE.get_or_init(|| Mutex::new(None)) } -// TicTacToe self-play: callback-based dispatch -typed_selfplay!( - TicTacToeSelfPlay, - TicTacToe, - i8, - Ix1, - Ix2, - TicTacToeEphemeralReplayBuffer, - 9, - callback_dispatch -); - -// Connect4 self-play: callback-based dispatch -typed_selfplay!( - Connect4SelfPlay, - Connect4, - i8, - Ix2, - Ix3, - Connect4EphemeralReplayBuffer, - 7, - callback_dispatch -); - /// Persistent ByteFight self-play session. /// /// Uses the CUDA graph runner for GPU dispatch (no Python callback). #[pyclass] -struct ByteFightSelfPlay { +struct SelfPlay { session: Option, } +struct ReplayBuffer; + #[pymethods] -impl ByteFightSelfPlay { +impl SelfPlay { #[new] #[pyo3(signature = ( replay_buffer, @@ -465,7 +66,7 @@ impl ByteFightSelfPlay { ))] fn new( py: Python<'_>, - replay_buffer: &ByteFightEphemeralReplayBuffer, + replay_buffer: &ReplayBuffer, num_threads: usize, workers_per_thread: usize, seed: u64, @@ -478,13 +79,6 @@ impl ByteFightSelfPlay { model: Py, selfplay_precision: &str, ) -> PyResult { - use cudagraph::ByteFightCudaGraphRunner; - use eval::PolicyValue; - use mcts::MCTSConfig; - use queue::{queue_shape_for_workers, BATCH_SIZE}; - use training::{SelfPlaySession, SessionConfig}; - use worker::WorkerConfig; - if mcts_num_simulations == 0 { return Err(PyErr::new::( "mcts_num_simulations must be >= 1", @@ -539,7 +133,7 @@ impl ByteFightSelfPlay { // Build or reuse the CUDA graph runner. let runner = { - let cache = bytefight_graph_cache(); + let cache = graph_cache(); let mut guard = cache.lock().expect("bytefight graph cache mutex poisoned"); let needs_rebuild = match guard.as_ref() { @@ -633,7 +227,7 @@ impl ByteFightSelfPlay { } } -impl Drop for ByteFightSelfPlay { +impl Drop for SelfPlay { fn drop(&mut self) { if let Some(mut session) = self.session.take() { session.shutdown(); @@ -643,78 +237,7 @@ impl Drop for ByteFightSelfPlay { #[pymodule] fn siebren(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; - m.add_class::()?; + m.add_class::()?; + m.add_class::()?; Ok(()) } - -#[derive(Clone, Copy, PartialEq, Eq, Debug, Hash, PartialOrd, Ord)] -#[repr(i8)] -pub enum Player { - PlayerA = 1, - PlayerB = -1, -} - -#[derive(Clone, Copy, PartialEq, Eq, Debug)] -pub enum TerminalState { - Win(Player), - Draw, -} - -/// Actions must be convertible to/from a unique index in `0..NUM_ACTIONS`. -pub trait Action: Copy + Eq + Hash { - fn to_index(self) -> usize; - fn from_index(index: usize) -> Option; -} - -/// Trait for serializing/deserializing game states to/from a string notation. -pub trait GameNotation: Sized { - type Error: std::error::Error + Send + Sync + 'static; - fn to_notation(&self) -> String; - fn from_notation(s: &str) -> Result; -} - -/// An environment implements a game that we want to train a model to play. -/// -/// Environments should support efficient rollback to step in and out of states -/// without cloning. -pub trait Environment: Clone + Hash + Debug + GameNotation { - /// Element type of observations (u8, i8, f32, etc.) - type ObsElem: Clone + Default + Send + Sync; - /// Dimension of a single observation (Ix1, Ix2, etc.) - type ObsDim: BatchDim; - /// Shape of a single observation as a compile-time constant. - const OBS_SHAPE: Self::ObsDim; - - type Action: Action; - type RollbackState; - const NUM_ACTIONS: usize; - - /// Creates an environment. Should be randomly generated if possible to - /// avoid the network overfitting on a single starting position. - fn new() -> Self; - - /// Returns None if the game is still going, Some(Win/Draw) if it's over. - fn is_terminal(&self) -> Option; - - /// Returns an iterator over valid actions. - fn valid_actions(&self) -> impl Iterator; - - fn current_player(&self) -> Player; - - /// Write the observation into the provided buffer. - /// The buffer is a mutable view into the queue's contiguous storage. - fn observation(&self, out: ArrayViewMut); - - /// Applies an action and returns state needed for rollback. - /// Caller must ensure the action is valid per `valid_actions`. - fn apply_action(&mut self, action: Self::Action) -> Self::RollbackState; - - /// Undoes `apply_action`. After `rollback(apply_action(a))`, the - /// environment should be back to its original state. - fn rollback(&mut self, rollback: Self::RollbackState); -} From 6bd967f035ef7fdadd38181d7ff6231c48df8e66 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Thu, 26 Mar 2026 03:59:12 -0400 Subject: [PATCH 05/59] im lazy as shit but almost done --- training/src/cudagraph.rs | 36 +- training/src/descent.rs | 7 +- training/src/eval.rs | 217 +------ training/src/integration_tests.rs | 494 -------------- training/src/lib.rs | 69 +- training/src/mcts.rs | 75 --- training/src/observation_replay_buffer.rs | 315 --------- training/src/replay_buffer.rs | 748 ++++------------------ training/src/training.rs | 23 +- training/src/types.rs | 3 + 10 files changed, 237 insertions(+), 1750 deletions(-) delete mode 100644 training/src/integration_tests.rs delete mode 100644 training/src/observation_replay_buffer.rs create mode 100644 training/src/types.rs diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index f5ede9a..e961d7b 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -27,10 +27,9 @@ const DL_DEVICE_CUDA: i32 = 2; const DL_DTYPE_FLOAT: u8 = 2; const DL_DTYPE_UINT: u8 = 1; -const BYTEFIGHT_OBS_SIDE: usize = 18; -const BYTEFIGHT_OBS_WIDTH: usize = 16; -const BYTEFIGHT_OBS_CELLS: usize = BYTEFIGHT_OBS_SIDE * BYTEFIGHT_OBS_WIDTH; -const BYTEFIGHT_ACTIONS: usize = 7; +const OBS_SIDE: usize = 18; +const OBS_WIDTH: usize = 16; +const OBS_CELLS: usize = OBS_SIDE * OBS_WIDTH; #[repr(C)] struct DLDevice { @@ -238,17 +237,14 @@ struct ByteFightCudaGraphLane { _py_owner: Py, obs_host: *mut u8, obs_dev: *mut c_void, - policy_host: *mut f32, - policy_dev: *mut c_void, value_host: *mut f32, value_dev: *mut c_void, } struct LaneCompletionContext { - policy_host: *const f32, value_host: *const f32, batch_size: usize, - completion: Option>>, + completion: Option>, } unsafe impl Send for LaneCompletionContext {} @@ -259,15 +255,14 @@ unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { } let mut ctx = unsafe { Box::from_raw(user_data.cast::()) }; - let policy_src = - unsafe { slice::from_raw_parts(ctx.policy_host, ctx.batch_size * BYTEFIGHT_ACTIONS) }; + let policy_src = unsafe { slice::from_raw_parts(ctx.policy_host, ctx.batch_size * ACTIONS) }; let value_src = unsafe { slice::from_raw_parts(ctx.value_host, ctx.batch_size) }; let mut outputs = vec![PolicyValue::<7>::default(); ctx.batch_size]; for (i, out) in outputs.iter_mut().enumerate() { - let start = i * BYTEFIGHT_ACTIONS; + let start = i * ACTIONS; out.policy - .copy_from_slice(&policy_src[start..start + BYTEFIGHT_ACTIONS]); + .copy_from_slice(&policy_src[start..start + ACTIONS]); out.value = value_src[i]; } @@ -319,16 +314,12 @@ impl ByteFightCudaGraphRunner { let module = PyModule::import(py, "siebren.cudagraph_backend")?; let capture_fn = module.getattr("capture_bytefight_lane_graph")?; - let obs_count = batch_size * BYTEFIGHT_OBS_CELLS; - let policy_count = batch_size * BYTEFIGHT_ACTIONS; + let obs_count = batch_size * OBS_CELLS; + let policy_count = batch_size * ACTIONS; let value_count = batch_size; - let obs_shape = [ - batch_size as i64, - BYTEFIGHT_OBS_SIDE as i64, - BYTEFIGHT_OBS_WIDTH as i64, - ]; - let policy_shape = [batch_size as i64, BYTEFIGHT_ACTIONS as i64]; + let obs_shape = [batch_size as i64, OBS_SIDE as i64, OBS_WIDTH as i64]; + let policy_shape = [batch_size as i64, ACTIONS as i64]; let value_shape = [batch_size as i64]; let mut lanes = Vec::with_capacity(num_lanes); @@ -459,9 +450,8 @@ impl ByteFightCudaGraphRunner { let obs_src = obs_view .as_slice() .expect("bytefight queue observation batch must be contiguous"); - let obs_dst = unsafe { - slice::from_raw_parts_mut(lane.obs_host, self.batch_size * BYTEFIGHT_OBS_CELLS) - }; + let obs_dst = + unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_CELLS) }; obs_dst.copy_from_slice(obs_src); unsafe { diff --git a/training/src/descent.rs b/training/src/descent.rs index c421ee1..c268096 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -4,6 +4,8 @@ use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; use std::time::Duration; +use crate::eval::Evaluator; + #[derive(Debug)] pub struct ChildData { action: Action, @@ -308,10 +310,11 @@ impl SearchNode { } } -pub struct GameSearchTree<'a> { +pub struct GameSearchTree<'a, E> { pub root_node: Box, root_state: Board, rng: SmallRng, + evaluator: Evaluator, } impl GameSearchTree<'_> { @@ -378,7 +381,7 @@ impl GameSearchTree<'_> { } } - pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { + pub async fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { let dcv = if self.root_state.is_white_turn() { 1 } else { diff --git a/training/src/eval.rs b/training/src/eval.rs index 79f69f8..7f73722 100644 --- a/training/src/eval.rs +++ b/training/src/eval.rs @@ -2,10 +2,10 @@ use std::future::Future; -use ndarray::Dimension; +use alpha_paint::board::Board; use crate::queue::GpuJobQueue; -use crate::{BatchDim, Environment}; +use crate::types::{ObservationDim, ObservationElement, OutputType}; /// Output from neural network evaluation: (policy logits, value). /// Policy has one entry per possible action, value is in [-1, 1]. @@ -25,51 +25,32 @@ impl Default for PolicyValue { } /// Async evaluator trait for neural network inference. -pub trait Evaluator { +pub trait Evaluator { /// Evaluate the environment and return (policy, value). /// Policy is over all actions, value is in [-1, 1] from current player's perspective. - fn evaluate(&self, env: &E) -> impl Future, f32)>; + fn evaluate(&self, board: &Board) -> impl Future; } /// GPU-backed evaluator that batches inference requests. /// /// Wraps a GpuJobQueue and converts between Environment observations /// and the queue's I/O types. -pub struct GpuEvaluator<'a, E: Environment, const NUM_ACTIONS: usize> -where - E::ObsDim: BatchDim, - ::Larger: Dimension, -{ - queue: &'a GpuJobQueue>, +pub struct GpuEvaluator<'a> { + queue: &'a GpuJobQueue, } -impl<'a, E, const NUM_ACTIONS: usize> GpuEvaluator<'a, E, NUM_ACTIONS> -where - E: Environment, - E::ObsDim: BatchDim, - ::Larger: Dimension, -{ - pub fn new(queue: &'a GpuJobQueue>) -> Self { +impl<'a> GpuEvaluator<'a> { + pub fn new(queue: &'a GpuJobQueue) -> Self { Self { queue } } } -impl<'a, E, const NUM_ACTIONS: usize> Evaluator for GpuEvaluator<'a, E, NUM_ACTIONS> -where - E: Environment, - E::ObsDim: BatchDim, - ::Larger: Dimension, -{ - fn evaluate(&self, env: &E) -> impl Future, f32)> { +impl<'a> Evaluator for GpuEvaluator<'a> { + fn evaluate(&self, board: &Board) -> impl Future { // Submit immediately with callback that writes observation - let future = self.queue.eval(|out| { - env.observation(out); - }); + let future = self.queue.eval(|out| todo!()); - async move { - let result = future.await; - (result.policy.to_vec(), result.value) - } + future } } @@ -78,176 +59,8 @@ where /// Returns uniform policy and zero value. pub struct UniformEvaluator; -impl Evaluator for UniformEvaluator { - fn evaluate(&self, _env: &E) -> impl Future, f32)> { - let num_actions = E::NUM_ACTIONS; - let policy = vec![1.0 / num_actions as f32; num_actions]; - std::future::ready((policy, 0.0)) - } -} - -/// CPU evaluator that uses a sync evaluation function. -/// -/// Useful for testing or CPU-only inference. -pub struct SyncEvaluator -where - F: Fn(&E) -> (Vec, f32), -{ - eval_fn: F, - _phantom: std::marker::PhantomData, -} - -impl SyncEvaluator -where - F: Fn(&E) -> (Vec, f32), -{ - pub fn new(eval_fn: F) -> Self { - Self { - eval_fn, - _phantom: std::marker::PhantomData, - } - } -} - -impl Evaluator for SyncEvaluator -where - F: Fn(&E) -> (Vec, f32), -{ - fn evaluate(&self, env: &E) -> impl Future, f32)> { - let result = (self.eval_fn)(env); - std::future::ready(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::environments::TicTacToe; - use crate::queue::BATCH_SIZE; - use ndarray::Ix1; - use std::sync::Arc; - - #[test] - fn test_uniform_evaluator() { - use std::pin::Pin; - use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; - - fn dummy_waker() -> Waker { - fn clone(_: *const ()) -> RawWaker { - RawWaker::new(std::ptr::null(), &VTABLE) - } - fn wake(_: *const ()) {} - fn wake_by_ref(_: *const ()) {} - fn drop(_: *const ()) {} - static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); - unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } - } - - let env = TicTacToe::new(); - let evaluator = UniformEvaluator; - - let mut future = evaluator.evaluate(&env); - let waker = dummy_waker(); - let mut cx = Context::from_waker(&waker); - - match Pin::new(&mut future).poll(&mut cx) { - Poll::Ready((policy, value)) => { - assert_eq!(policy.len(), 9); - assert!((policy[0] - 1.0 / 9.0).abs() < 0.001); - assert_eq!(value, 0.0); - } - Poll::Pending => panic!("UniformEvaluator should be ready immediately"), - } - } - - #[test] - fn test_sync_evaluator() { - use std::pin::Pin; - use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; - - fn dummy_waker() -> Waker { - fn clone(_: *const ()) -> RawWaker { - RawWaker::new(std::ptr::null(), &VTABLE) - } - fn wake(_: *const ()) {} - fn wake_by_ref(_: *const ()) {} - fn drop(_: *const ()) {} - static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); - unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } - } - - let env = TicTacToe::new(); - let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { - let mut policy = vec![0.0; 9]; - policy[4] = 1.0; // Center is best - (policy, 0.5) - }); - - let mut future = evaluator.evaluate(&env); - let waker = dummy_waker(); - let mut cx = Context::from_waker(&waker); - - match Pin::new(&mut future).poll(&mut cx) { - Poll::Ready((policy, value)) => { - assert_eq!(policy[4], 1.0); - assert_eq!(value, 0.5); - } - Poll::Pending => panic!("SyncEvaluator should be ready immediately"), - } - } - - #[test] - fn test_gpu_evaluator() { - use crate::executor::Executor; - use std::cell::Cell; - use std::rc::Rc; - - // Create a mock GPU queue that returns uniform policy - type Output = PolicyValue<{ TicTacToe::NUM_ACTIONS }>; - let queue: Arc> = Arc::new(GpuJobQueue::new( - TicTacToe::OBS_SHAPE, - BATCH_SIZE, - |_batch_idx, _inputs, completion| { - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }, - )); - - let evaluator = GpuEvaluator::::new(&*queue); - - // Submit BATCH_SIZE evaluations to trigger dispatch - let envs: Vec = (0..BATCH_SIZE).map(|_| TicTacToe::new()).collect(); - - let results: Rc> = Rc::new(Cell::new(0)); - - let futures: Vec<_> = envs - .iter() - .map(|env| { - let results = results.clone(); - let fut = evaluator.evaluate(env); - async move { - let (policy, value) = fut.await; - assert_eq!(policy.len(), 9); - assert!((policy[0] - 1.0 / 9.0).abs() < 0.001); - assert_eq!(value, 0.0); - results.set(results.get() + 1); - } - }) - .collect(); - - let executor = Executor::new(|| queue.listen()); - executor.run( - &mut futures - .into_iter() - .map(|f| Box::pin(f) as std::pin::Pin>>) - .collect(), - &mut || false, - ); - - assert_eq!(results.get(), BATCH_SIZE); +impl Evaluator for UniformEvaluator { + fn evaluate(&self, _: &Board) -> impl Future { + std::future::ready(0.0) } } diff --git a/training/src/integration_tests.rs b/training/src/integration_tests.rs deleted file mode 100644 index 6519692..0000000 --- a/training/src/integration_tests.rs +++ /dev/null @@ -1,494 +0,0 @@ -//! Integration tests for the full GPU batching stack. - -#[cfg(test)] -mod tests { - use std::cell::RefCell; - use std::rc::Rc; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - use std::thread; - - use ndarray::Ix1; - use rand::SeedableRng; - use rand_chacha::ChaCha8Rng; - - use crate::environments::TicTacToe; - use crate::eval::{GpuEvaluator, PolicyValue, SyncEvaluator}; - use crate::executor::Executor; - use crate::mcts::{MCTSConfig, MCTS}; - use crate::observation_replay_buffer::ObservationReplayBuffer; - use crate::queue::{GpuJobQueue, BATCH_SIZE}; - use crate::worker::{worker_loop, WorkerConfig}; - use crate::Environment; - - /// Simple test: multiple futures doing GPU eval on a single thread. - #[test] - fn test_simple_multi_future() { - let dispatch_count = Arc::new(AtomicUsize::new(0)); - let dispatch_count_clone = dispatch_count.clone(); - - type Output = PolicyValue<9>; - let queue: Arc> = Arc::new(GpuJobQueue::new( - TicTacToe::OBS_SHAPE, - BATCH_SIZE, - move |_batch_idx, _inputs, completion| { - dispatch_count_clone.fetch_add(1, Ordering::Relaxed); - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }, - )); - - let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); - let executor = Executor::new(|| queue.listen()); - - // Create BATCH_SIZE futures that each do one eval - let completed = Rc::new(RefCell::new(0usize)); - let futures: Vec<_> = (0..BATCH_SIZE) - .map(|_| { - let completed = completed.clone(); - let evaluator = evaluator.clone(); - let env = TicTacToe::new(); - async move { - use crate::eval::Evaluator; - let (policy, value) = evaluator.evaluate(&env).await; - assert_eq!(policy.len(), 9); - assert_eq!(value, 0.0); - *completed.borrow_mut() += 1; - } - }) - .collect(); - - executor.run( - &mut futures - .into_iter() - .map(|f| Box::pin(f) as std::pin::Pin>>) - .collect(), - &mut || false, - ); - - assert_eq!(*completed.borrow(), BATCH_SIZE); - assert_eq!(dispatch_count.load(Ordering::Relaxed), 1); - } - - /// Test multiple batches with simple futures. - #[test] - fn test_multiple_batches_simple() { - let dispatch_count = Arc::new(AtomicUsize::new(0)); - let dispatch_count_clone = dispatch_count.clone(); - let num_evals = BATCH_SIZE * 3; - - type Output = PolicyValue<9>; - let queue: Arc> = Arc::new(GpuJobQueue::new( - TicTacToe::OBS_SHAPE, - num_evals, - move |_batch_idx, _inputs, completion| { - dispatch_count_clone.fetch_add(1, Ordering::Relaxed); - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }, - )); - - let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); - let executor = Executor::new(|| queue.listen()); - - // Create 3 batches worth of futures - let completed = Rc::new(RefCell::new(0usize)); - let futures: Vec<_> = (0..num_evals) - .map(|_| { - let completed = completed.clone(); - let evaluator = evaluator.clone(); - let env = TicTacToe::new(); - async move { - use crate::eval::Evaluator; - let (_policy, _value) = evaluator.evaluate(&env).await; - *completed.borrow_mut() += 1; - } - }) - .collect(); - - executor.run( - &mut futures - .into_iter() - .map(|f| Box::pin(f) as std::pin::Pin>>) - .collect(), - &mut || false, - ); - - assert_eq!(*completed.borrow(), num_evals); - assert_eq!(dispatch_count.load(Ordering::Relaxed), 3); - } - - /// Test multiple MCTS searches concurrently. - #[test] - fn test_multiple_mcts_searches() { - let dispatch_count = Arc::new(AtomicUsize::new(0)); - let dispatch_count_clone = dispatch_count.clone(); - let num_searches = BATCH_SIZE * 2; - - type Output = PolicyValue<9>; - let queue: Arc> = Arc::new(GpuJobQueue::new( - TicTacToe::OBS_SHAPE, - num_searches, - move |_batch_idx, _inputs, completion| { - dispatch_count_clone.fetch_add(1, Ordering::Relaxed); - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }, - )); - - let evaluator = Rc::new(GpuEvaluator::::new(&*queue)); - let mcts_config = MCTSConfig { - num_simulations: 5, - ..Default::default() - }; - let executor = Executor::new(|| queue.listen()); - - // Run multiple MCTS searches concurrently - - let completed = Rc::new(RefCell::new(0usize)); - let futures: Vec<_> = (0..num_searches) - .map(|i| { - let completed = completed.clone(); - let evaluator = evaluator.clone(); - let mcts_config = mcts_config.clone(); - async move { - let mcts = MCTS::new(&*evaluator, &mcts_config); - let mut env = TicTacToe::new(); - let mut rng = ChaCha8Rng::seed_from_u64(i as u64); - let visits = mcts.search(&mut env, &mut rng).await; - assert_eq!(visits.len(), 9); - *completed.borrow_mut() += 1; - } - }) - .collect(); - - executor.run( - &mut futures - .into_iter() - .map(|f| Box::pin(f) as std::pin::Pin>>) - .collect(), - &mut || false, - ); - - assert_eq!(*completed.borrow(), num_searches); - let batches = dispatch_count.load(Ordering::Relaxed); - assert!(batches > 0); - } - - /// Test worker_loop runs until a global target of samples is reached. - #[test] - fn test_worker_loop_with_shared_counter() { - let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { - let mut policy = vec![0.0; 9]; - policy[0] = 1.0; - (policy, 0.0) - }); - let config = WorkerConfig { - mcts: MCTSConfig { - num_simulations: 3, - ..Default::default() - }, - ..Default::default() - }; - let executor = Executor::new(|| event_listener::Event::new().listen()); - - let num_workers = 8; - let target_samples = 200; // ~32 games * 6 samples/game - let samples_collected = Arc::new(AtomicUsize::new(0)); - let games_completed = Arc::new(AtomicUsize::new(0)); - let replay_buffer = ObservationReplayBuffer::::new(1000, TicTacToe::OBS_SHAPE); - - let futures: Vec<_> = (0..num_workers) - .map(|i| { - let evaluator = &evaluator; - let config = &config; - let replay_buffer = &replay_buffer; - let samples_collected = samples_collected.clone(); - let games_completed = games_completed.clone(); - let mut rng = ChaCha8Rng::seed_from_u64(i as u64); - async move { - worker_loop::( - evaluator, - config, - &mut rng, - samples_collected, - games_completed, - target_samples, - replay_buffer, - ) - .await; - } - }) - .collect(); - - executor.run( - &mut futures - .into_iter() - .map(|f| Box::pin(f) as std::pin::Pin>>) - .collect(), - &mut || false, - ); - - let completed_games = games_completed.load(Ordering::Relaxed); - let collected_samples = samples_collected.load(Ordering::Relaxed); - // We collect at least target_samples (may be slightly more due to race) - assert!(collected_samples >= target_samples); - // TicTacToe games are 5-9 moves, so ~22-40 games for 200 samples - assert!( - completed_games >= 20, - "expected at least 20 games, got {completed_games}" - ); - assert_eq!(replay_buffer.len(), collected_samples); - } - - /// Test multithreaded worker_loop with shared counter using the sync evaluator. - #[test] - fn test_multithreaded_worker_loop() { - const NUM_THREADS: usize = 2; - const WORKERS_PER_THREAD: usize = 4; - const TARGET_SAMPLES: usize = 200; // ~32 games * 6 samples/game - - let total_samples = Arc::new(AtomicUsize::new(0)); - let games_completed = Arc::new(AtomicUsize::new(0)); - let replay_buffer = ObservationReplayBuffer::::new(1000, TicTacToe::OBS_SHAPE); - - thread::scope(|s| { - for thread_id in 0..NUM_THREADS { - let samples_collected = total_samples.clone(); - let games_completed = games_completed.clone(); - let replay_buffer = &replay_buffer; - - s.spawn(move || { - let evaluator = SyncEvaluator::new(|_env: &TicTacToe| { - let mut policy = vec![0.0; 9]; - policy[0] = 1.0; - (policy, 0.0) - }); - let config = WorkerConfig { - mcts: MCTSConfig { - num_simulations: 3, - ..Default::default() - }, - ..Default::default() - }; - let executor = Executor::new(|| event_listener::Event::new().listen()); - - let futures: Vec<_> = (0..WORKERS_PER_THREAD) - .map(|i| { - let evaluator = &evaluator; - let config = &config; - let samples_collected = samples_collected.clone(); - let games_completed = games_completed.clone(); - let mut rng = ChaCha8Rng::seed_from_u64((thread_id * 1000 + i) as u64); - async move { - worker_loop::( - evaluator, - config, - &mut rng, - samples_collected, - games_completed, - TARGET_SAMPLES, - replay_buffer, - ) - .await; - } - }) - .collect(); - - executor.run( - &mut futures - .into_iter() - .map(|f| { - Box::pin(f) - as std::pin::Pin>> - }) - .collect(), - &mut || false, - ); - }); - } - }); - - let completed_games = games_completed.load(Ordering::Relaxed); - let collected_samples = total_samples.load(Ordering::Relaxed); - - // We collect at least TARGET_SAMPLES (may be slightly more due to race) - assert!(collected_samples >= TARGET_SAMPLES); - assert!(completed_games >= 30); // At least ~30 games to get 200 samples - assert_eq!(replay_buffer.len(), collected_samples); - } - - /// Test persistent SelfPlaySession: wait_for(a), then wait_for(a+b). - /// Verifies monotonic samples and stable operation across waits. - #[test] - fn test_session_wait_for_monotonic() { - use crate::eval::PolicyValue; - use crate::training::{SelfPlaySession, SessionConfig}; - - type Output = PolicyValue<9>; - let replay_buffer = Arc::new(ObservationReplayBuffer::::new( - 10000, - TicTacToe::OBS_SHAPE, - )); - - let config = SessionConfig { - num_threads: 2, - workers_per_thread: BATCH_SIZE / 2, - worker: WorkerConfig { - mcts: MCTSConfig { - num_simulations: 3, - ..Default::default() - }, - ..Default::default() - }, - seed: 42, - }; - - // Use SyncEvaluator-style dispatch: uniform policy, zero value. - let dispatch = |_batch_idx: usize, - _inputs: ndarray::ArrayView, - completion: crate::queue::BatchCompletion| { - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }; - - let mut session = - SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); - - // First wait: collect at least 100 samples. - let target_a = 100; - let reached_a = session.wait_for(target_a); - assert!( - reached_a >= target_a, - "expected >= {target_a} samples, got {reached_a}" - ); - - // Second wait: collect more samples (absolute target). - let target_b = reached_a + 100; - let reached_b = session.wait_for(target_b); - assert!( - reached_b >= target_b, - "expected >= {target_b} samples, got {reached_b}" - ); - assert!( - reached_b >= reached_a, - "samples must be monotonic: {reached_b} < {reached_a}" - ); - - // Replay buffer length should match total samples. - assert_eq!( - replay_buffer.len(), - reached_b, - "replay buffer len should match samples collected" - ); - - session.shutdown(); - } - - /// Test that drop during paused and running states doesn't deadlock. - #[test] - fn test_session_drop_while_paused() { - use crate::eval::PolicyValue; - use crate::training::{SelfPlaySession, SessionConfig}; - - type Output = PolicyValue<9>; - let replay_buffer = Arc::new(ObservationReplayBuffer::::new( - 1000, - TicTacToe::OBS_SHAPE, - )); - - let config = SessionConfig { - num_threads: 2, - workers_per_thread: BATCH_SIZE / 2, - worker: WorkerConfig { - mcts: MCTSConfig { - num_simulations: 3, - ..Default::default() - }, - ..Default::default() - }, - seed: 42, - }; - - let dispatch = |_batch_idx: usize, - _inputs: ndarray::ArrayView, - completion: crate::queue::BatchCompletion| { - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }; - - // Create session, wait for some samples, then drop. - // Should not deadlock. - let mut session = - SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); - session.wait_for(50); - session.shutdown(); - // If we reach here, no deadlock. - } - - /// Test that drop during running state doesn't deadlock. - #[test] - fn test_session_drop_while_running() { - use crate::eval::PolicyValue; - use crate::training::{SelfPlaySession, SessionConfig}; - - type Output = PolicyValue<9>; - let replay_buffer = Arc::new(ObservationReplayBuffer::::new( - 1000, - TicTacToe::OBS_SHAPE, - )); - - let config = SessionConfig { - num_threads: 2, - workers_per_thread: BATCH_SIZE / 2, - worker: WorkerConfig { - mcts: MCTSConfig { - num_simulations: 3, - ..Default::default() - }, - ..Default::default() - }, - seed: 42, - }; - - let dispatch = |_batch_idx: usize, - _inputs: ndarray::ArrayView, - completion: crate::queue::BatchCompletion| { - let mut outputs = vec![Output::default(); BATCH_SIZE]; - for output in outputs.iter_mut() { - output.policy = [1.0 / 9.0; 9]; - output.value = 0.0; - } - completion.complete(&outputs); - }; - - let mut session = - SelfPlaySession::new::(config, replay_buffer.clone(), dispatch); - // Start without wait_for -- workers are actively running. - session.start(); - // Give threads a moment to actually start polling. - std::thread::sleep(std::time::Duration::from_millis(50)); - // Drop should shut down cleanly without deadlock. - session.shutdown(); - } -} diff --git a/training/src/lib.rs b/training/src/lib.rs index 9ec985a..6b9f0cf 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -4,25 +4,83 @@ use std::sync::{Mutex, OnceLock}; use cudagraph::ByteFightCudaGraphRunner; use eval::PolicyValue; use mcts::MCTSConfig; -use ndarray::{ArrayView, Ix3}; -use pyo3::prelude::*; use queue::{queue_shape_for_workers, BATCH_SIZE}; use training::{SelfPlaySession, SessionConfig}; use worker::WorkerConfig; +use ndarray::{ArrayView, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; +use pyo3::prelude::*; + pub mod cudagraph; -mod descent; pub mod eval; pub mod executor; pub mod future; -mod integration_tests; pub mod mcts; -pub mod observation_replay_buffer; pub mod queue; pub mod replay_buffer; pub mod training; +mod types; pub mod worker; +/// Extension trait for prepending a batch dimension to a shape. +/// +/// The `BatchedDim` associated type is the dimension with batch prepended. +/// We use an associated type instead of `Dimension::Larger` so we can +/// constrain that `BatchedDim::Smaller == Self`. +pub trait BatchDim: Dimension + Clone { + type BatchedDim: Dimension + RemoveAxis; + + fn with_batch(batch_size: usize, obs_shape: Self) -> Self::BatchedDim; +} + +impl BatchDim for Ix0 { + type BatchedDim = Ix1; + + fn with_batch(batch_size: usize, _obs: Self) -> Ix1 { + Ix1(batch_size) + } +} + +impl BatchDim for Ix1 { + type BatchedDim = Ix2; + + fn with_batch(batch_size: usize, obs: Self) -> Ix2 { + Ix2(batch_size, obs[0]) + } +} + +impl BatchDim for Ix2 { + type BatchedDim = Ix3; + + fn with_batch(batch_size: usize, obs: Self) -> Ix3 { + Ix3(batch_size, obs[0], obs[1]) + } +} + +impl BatchDim for Ix3 { + type BatchedDim = Ix4; + + fn with_batch(batch_size: usize, obs: Self) -> Ix4 { + Ix4(batch_size, obs[0], obs[1], obs[2]) + } +} + +impl BatchDim for Ix4 { + type BatchedDim = Ix5; + + fn with_batch(batch_size: usize, obs: Self) -> Ix5 { + Ix5(batch_size, obs[0], obs[1], obs[2], obs[3]) + } +} + +impl BatchDim for Ix5 { + type BatchedDim = Ix6; + + fn with_batch(batch_size: usize, obs: Self) -> Ix6 { + Ix6(batch_size, obs[0], obs[1], obs[2], obs[3], obs[4]) + } +} + struct ByteFightGraphCacheEntry { model_ptr: usize, num_batches: usize, @@ -35,7 +93,6 @@ static GRAPH_CACHE: OnceLock>> = OnceLock fn graph_cache() -> &'static Mutex> { GRAPH_CACHE.get_or_init(|| Mutex::new(None)) } - /// Persistent ByteFight self-play session. /// /// Uses the CUDA graph runner for GPU dispatch (no Python callback). diff --git a/training/src/mcts.rs b/training/src/mcts.rs index 8d0bbcf..b44df59 100644 --- a/training/src/mcts.rs +++ b/training/src/mcts.rs @@ -292,78 +292,3 @@ pub fn best_action_index(visits: &[u32]) -> Option { .max_by_key(|(_, &v)| v) .map(|(i, _)| i) } - -#[cfg(test)] -mod tests { - use super::*; - use crate::environments::TicTacToe; - use crate::eval::UniformEvaluator; - use crate::executor::Executor; - use rand::SeedableRng; - use rand_chacha::ChaCha8Rng; - use std::cell::RefCell; - use std::rc::Rc; - - #[test] - fn test_async_mcts_returns_valid_visits() { - let config = MCTSConfig { - num_simulations: 100, - ..Default::default() - }; - let evaluator = UniformEvaluator; - let mcts = MCTS::new(&evaluator, &config); - - let game = Rc::new(RefCell::new(TicTacToe::new())); - let rng = Rc::new(RefCell::new(ChaCha8Rng::seed_from_u64(42))); - let result: Rc>>> = Rc::new(RefCell::new(None)); - - let game_clone = game.clone(); - let rng_clone = rng.clone(); - let result_clone = result.clone(); - - let fut = async move { - let visits = mcts - .search(&mut *game_clone.borrow_mut(), &mut *rng_clone.borrow_mut()) - .await; - *result_clone.borrow_mut() = Some(visits); - }; - - // Use a dummy event for the executor - let event = event_listener::Event::new(); - let executor = Executor::new(|| event.listen()); - executor.run(&mut vec![Box::pin(fut)], &mut || false); - - let visits = result.borrow().clone().unwrap(); - assert_eq!(visits.len(), 9); - - let total: u32 = visits.iter().sum(); - assert!(total > 0); - - // All valid actions should have some visits - for action in game.borrow().valid_actions() { - assert!(visits[action.to_index()] > 0); - } - } - - #[test] - fn test_visits_to_policy_with_temperature() { - let visits = vec![100, 50, 25, 25]; - - let policy = visits_to_policy(&visits, 1.0); - assert!((policy[0] - 0.5).abs() < 0.01); - assert!((policy[1] - 0.25).abs() < 0.01); - - let policy = visits_to_policy(&visits, 0.0); - assert_eq!(policy[0], 1.0); - assert_eq!(policy[1], 0.0); - } - - #[test] - fn test_best_action_index() { - let visits = vec![10, 50, 30, 5]; - assert_eq!(best_action_index(&visits), Some(1)); - - let empty: Vec = vec![]; - assert_eq!(best_action_index(&empty), None); - } -} diff --git a/training/src/observation_replay_buffer.rs b/training/src/observation_replay_buffer.rs deleted file mode 100644 index cec6381..0000000 --- a/training/src/observation_replay_buffer.rs +++ /dev/null @@ -1,315 +0,0 @@ -//! Lock-free ring buffer for contiguous observation replay storage. - -use std::cell::UnsafeCell; -use std::sync::atomic::{AtomicU64, Ordering}; - -use ndarray::{Array, Array2, ArrayViewMut, Axis}; -use rand::seq::index::sample; -use rand::Rng; - -use crate::BatchDim; - -/// Batched sample output from [`ObservationReplayBuffer::sample`]. -pub struct ObservationSampleBatch -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - pub observations: Array, - pub policies: Array2, - pub values: Vec, -} - -/// Lock-free ring buffer for storing observations, policies, and values. -/// -/// Observations and policies are kept in single contiguous arrays: -/// - observations: `(capacity, ...obs_shape)` -/// - policies: `(capacity, NUM_ACTIONS)` -/// - values: `(capacity,)` -pub struct ObservationReplayBuffer -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - observations: UnsafeCell>, - policies: UnsafeCell>, - values: UnsafeCell>, - capacity: usize, - obs_shape: D, - obs_elems_per_sample: usize, - head: AtomicU64, - writers: AtomicU64, -} - -// SAFETY: Each writer reserves unique slots through an atomic ticket and writes -// only to its owned slots until drop. Readers require no active writers. -unsafe impl Sync for ObservationReplayBuffer -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ -} -unsafe impl Send for ObservationReplayBuffer -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ -} - -/// RAII guard for writing to reserved slots. -pub struct ReserveGuard<'a, A, D, const NUM_ACTIONS: usize> -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - buffer: &'a ObservationReplayBuffer, - start: u64, - len: usize, - written: usize, -} - -impl<'a, A, D, const NUM_ACTIONS: usize> ReserveGuard<'a, A, D, NUM_ACTIONS> -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - /// Push one sample by writing observation directly into the reserved slot. - #[inline] - pub fn push_with_observation(&mut self, policy: &[f32], value: f32, write_observation: F) - where - F: FnOnce(ArrayViewMut), - { - assert!(self.written < self.len, "wrote more samples than reserved"); - assert_eq!( - policy.len(), - NUM_ACTIONS, - "policy length must match NUM_ACTIONS" - ); - - let idx = (self.start + self.written as u64) as usize % self.buffer.capacity; - - unsafe { - let slot_view = (*self.buffer.observations.get()).index_axis_mut(Axis(0), idx); - write_observation(slot_view); - - let policy_storage = &mut *self.buffer.policies.get(); - let policy_slice = policy_storage - .as_slice_memory_order_mut() - .expect("policy storage must be contiguous"); - let policy_start = idx * NUM_ACTIONS; - policy_slice[policy_start..policy_start + NUM_ACTIONS].copy_from_slice(policy); - - (&mut *self.buffer.values.get())[idx] = value; - } - - self.written += 1; - } - - /// Push one sample from flattened observation data. - pub fn push(&mut self, observation: &[A], policy: &[f32], value: f32) { - assert_eq!( - observation.len(), - self.buffer.obs_elems_per_sample, - "observation length must match env observation size" - ); - - self.push_with_observation(policy, value, |mut out| { - for (dst, src) in out.iter_mut().zip(observation.iter()) { - *dst = src.clone(); - } - }); - } -} - -impl Drop for ReserveGuard<'_, A, D, NUM_ACTIONS> -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - fn drop(&mut self) { - self.buffer.writers.fetch_sub(1, Ordering::Release); - } -} - -impl ObservationReplayBuffer -where - A: Clone + Default + Send + Sync, - D: BatchDim, -{ - pub fn new(capacity: usize, obs_shape: D) -> Self { - assert!(capacity > 0, "capacity must be > 0"); - let obs_elems_per_sample = obs_shape.clone().size(); - - Self { - observations: UnsafeCell::new(Array::default(D::with_batch( - capacity, - obs_shape.clone(), - ))), - policies: UnsafeCell::new(Array2::::zeros((capacity, NUM_ACTIONS))), - values: UnsafeCell::new(vec![0.0; capacity]), - capacity, - obs_shape, - obs_elems_per_sample, - head: AtomicU64::new(0), - writers: AtomicU64::new(0), - } - } - - pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D, NUM_ACTIONS> { - assert!( - n <= self.capacity, - "cannot reserve more samples than buffer capacity" - ); - - self.writers.fetch_add(1, Ordering::Acquire); - let start = self.head.fetch_add(n as u64, Ordering::AcqRel); - ReserveGuard { - buffer: self, - start, - len: n, - written: 0, - } - } - - #[inline] - fn valid_range(&self) -> (u64, u64) { - let head = self.head.load(Ordering::Acquire); - let tail = head.saturating_sub(self.capacity as u64); - (tail, head) - } - - #[inline] - pub fn len(&self) -> usize { - let (tail, head) = self.valid_range(); - (head - tail) as usize - } - - #[inline] - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - #[inline] - pub fn capacity(&self) -> usize { - self.capacity - } - - /// Sample `n` items uniformly. Panics if writers are active. - pub fn sample( - &self, - n: usize, - rng: &mut impl Rng, - ) -> ObservationSampleBatch { - assert_eq!( - self.writers.load(Ordering::Acquire), - 0, - "cannot sample while writers are active" - ); - - let (tail, head) = self.valid_range(); - let count = (head - tail) as usize; - if count == 0 || n == 0 { - return ObservationSampleBatch { - observations: Array::default(D::with_batch(0, self.obs_shape.clone())), - policies: Array2::::zeros((0, NUM_ACTIONS)), - values: Vec::new(), - }; - } - - let sample_count = n.min(count); - let indices = sample(rng, count, sample_count); - - let observations = unsafe { &*self.observations.get() }; - let observation_slice = observations - .as_slice_memory_order() - .expect("observation storage must be contiguous"); - - let policies = unsafe { &*self.policies.get() }; - let policy_slice = policies - .as_slice_memory_order() - .expect("policy storage must be contiguous"); - - let values = unsafe { &*self.values.get() }; - - let mut obs_data = Vec::with_capacity(sample_count * self.obs_elems_per_sample); - let mut policy_data = Vec::with_capacity(sample_count * NUM_ACTIONS); - let mut value_data = Vec::with_capacity(sample_count); - - for offset in indices.iter() { - let idx = (tail + offset as u64) as usize % self.capacity; - - let obs_start = idx * self.obs_elems_per_sample; - obs_data.extend_from_slice( - &observation_slice[obs_start..obs_start + self.obs_elems_per_sample], - ); - - let policy_start = idx * NUM_ACTIONS; - policy_data.extend_from_slice(&policy_slice[policy_start..policy_start + NUM_ACTIONS]); - - value_data.push(values[idx]); - } - - let observations = Array::from_shape_vec( - D::with_batch(sample_count, self.obs_shape.clone()), - obs_data, - ) - .expect("sample observation shape mismatch"); - let policies = Array2::from_shape_vec((sample_count, NUM_ACTIONS), policy_data) - .expect("shape mismatch"); - - ObservationSampleBatch { - observations, - policies, - values: value_data, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use ndarray::Ix1; - use rand::SeedableRng; - use rand_chacha::ChaCha8Rng; - - #[test] - fn test_push_and_sample() { - let buffer = ObservationReplayBuffer::::new(10, Ix1(4)); - - { - let mut guard = buffer.reserve(2); - guard.push(&[1, 2, 3, 4], &[0.1, 0.2, 0.7], 0.5); - guard.push(&[5, 6, 7, 8], &[0.6, 0.2, 0.2], -0.5); - } - - assert_eq!(buffer.len(), 2); - - let mut rng = ChaCha8Rng::seed_from_u64(7); - let batch = buffer.sample(2, &mut rng); - assert_eq!(batch.observations.shape(), &[2, 4]); - assert_eq!(batch.policies.shape(), &[2, 3]); - assert_eq!(batch.values.len(), 2); - } - - #[test] - fn test_wraparound_len() { - let buffer = ObservationReplayBuffer::::new(3, Ix1(2)); - - for i in 0..10 { - let mut guard = buffer.reserve(1); - guard.push(&[i as i8, (i + 1) as i8], &[0.5, 0.5], i as f32); - } - - assert_eq!(buffer.len(), 3); - } - - #[test] - #[should_panic(expected = "cannot sample while writers are active")] - fn test_sample_during_write_panics() { - let buffer = ObservationReplayBuffer::::new(4, Ix1(2)); - let _guard = buffer.reserve(1); - - let mut rng = ChaCha8Rng::seed_from_u64(42); - let _ = buffer.sample(1, &mut rng); - } -} diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs index f881065..411da04 100644 --- a/training/src/replay_buffer.rs +++ b/training/src/replay_buffer.rs @@ -1,94 +1,149 @@ -//! Lock-free replay buffer for concurrent training data storage. +//! Lock-free ring buffer for contiguous observation replay storage. use std::cell::UnsafeCell; -use std::fs::File; -use std::io::{self, BufReader, BufWriter, Read, Write}; -use std::mem::MaybeUninit; -use std::path::Path; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU64, Ordering}; +use ndarray::{Array, ArrayViewMut, Axis}; use rand::seq::index::sample; use rand::Rng; -/// A training sample stored in the replay buffer. -#[derive(Clone, Debug)] -pub struct Sample { - pub notation: String, - pub policy: Vec, - pub value: f32, +use crate::BatchDim; + +/// Batched sample output from [`ObservationReplayBuffer::sample`]. +pub struct SampleBatch +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub observations: Array, + pub values: Vec, } -/// Lock-free ring buffer for storing training samples. -pub struct ReplayBuffer { - data: Box<[UnsafeCell>]>, - /// Tracks whether each slot has ever been initialized. - /// - /// This is separate from `saved_head`: save checkpoints can be marked at - /// arbitrary points, but overwrite safety needs per-slot init state. - initialized: Box<[AtomicBool]>, +/// Lock-free ring buffer for storing observations, and values. +/// +/// Observations are kept in single contiguous arrays: +/// - observations: `(capacity, ...obs_shape)` +/// - values: `(capacity,)` +pub struct ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + observations: UnsafeCell>, + values: UnsafeCell>, capacity: usize, + obs_shape: D, + obs_elems_per_sample: usize, head: AtomicU64, writers: AtomicU64, - /// Tracks the head position at the last save. Used to avoid saving - /// the same samples multiple times across checkpoints. - saved_head: AtomicU64, } -unsafe impl Sync for ReplayBuffer {} -unsafe impl Send for ReplayBuffer {} +// SAFETY: Each writer reserves unique slots through an atomic ticket and writes +// only to its owned slots until drop. Readers require no active writers. +unsafe impl Sync for ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} +unsafe impl Send for ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} /// RAII guard for writing to reserved slots. -pub struct ReserveGuard<'a> { - buffer: &'a ReplayBuffer, +pub struct ReserveGuard<'a, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + buffer: &'a ReplayBuffer, start: u64, len: usize, written: usize, } -impl<'a> ReserveGuard<'a> { +impl<'a, A, D> ReserveGuard<'a, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + /// Push one sample by writing observation directly into the reserved slot. #[inline] - pub fn push(&mut self, sample: Sample) { + pub fn push_with_observation(&mut self, value: f32, write_observation: F) + where + F: FnOnce(ArrayViewMut), + { assert!(self.written < self.len, "wrote more samples than reserved"); let idx = (self.start + self.written as u64) as usize % self.buffer.capacity; + unsafe { - if self.buffer.initialized[idx].load(Ordering::Acquire) { - (*self.buffer.data[idx].get()).assume_init_drop(); - } - (*self.buffer.data[idx].get()).write(sample); + let slot_view = (*self.buffer.observations.get()).index_axis_mut(Axis(0), idx); + write_observation(slot_view); + + (&mut *self.buffer.values.get())[idx] = value; } - self.buffer.initialized[idx].store(true, Ordering::Release); + self.written += 1; } - pub fn extend(&mut self, samples: impl IntoIterator) { - for sample in samples { - self.push(sample); - } + /// Push one sample from flattened observation data. + pub fn push(&mut self, observation: &[A], value: f32) { + assert_eq!( + observation.len(), + self.buffer.obs_elems_per_sample, + "observation length must match env observation size" + ); + + self.push_with_observation(value, |mut out| { + for (dst, src) in out.iter_mut().zip(observation.iter()) { + *dst = src.clone(); + } + }); } } -impl Drop for ReserveGuard<'_> { +impl Drop for ReserveGuard<'_, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ fn drop(&mut self) { self.buffer.writers.fetch_sub(1, Ordering::Release); } } -impl ReplayBuffer { - pub fn new(capacity: usize) -> Self { - let data: Vec>> = (0..capacity) - .map(|_| UnsafeCell::new(MaybeUninit::uninit())) - .collect(); +impl ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub fn new(capacity: usize, obs_shape: D) -> Self { + assert!(capacity > 0, "capacity must be > 0"); + let obs_elems_per_sample = obs_shape.clone().size(); + Self { - data: data.into_boxed_slice(), - initialized: (0..capacity).map(|_| AtomicBool::new(false)).collect(), + observations: UnsafeCell::new(Array::default(D::with_batch( + capacity, + obs_shape.clone(), + ))), + values: UnsafeCell::new(vec![0.0; capacity]), capacity, + obs_shape, + obs_elems_per_sample, head: AtomicU64::new(0), writers: AtomicU64::new(0), - saved_head: AtomicU64::new(0), } } - pub fn reserve(&self, n: usize) -> ReserveGuard<'_> { + pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D, NUM_ACTIONS> { + assert!( + n <= self.capacity, + "cannot reserve more samples than buffer capacity" + ); + self.writers.fetch_add(1, Ordering::Acquire); let start = self.head.fetch_add(n as u64, Ordering::AcqRel); ReserveGuard { @@ -123,7 +178,7 @@ impl ReplayBuffer { } /// Sample `n` items uniformly. Panics if writers are active. - pub fn sample(&self, n: usize, rng: &mut impl Rng) -> Vec { + pub fn sample(&self, n: usize, rng: &mut impl Rng) -> SampleBatch { assert_eq!( self.writers.load(Ordering::Acquire), 0, @@ -133,584 +188,45 @@ impl ReplayBuffer { let (tail, head) = self.valid_range(); let count = (head - tail) as usize; if count == 0 || n == 0 { - return Vec::new(); - } - - let indices = sample(rng, count, n.min(count)); - indices - .iter() - .map(|offset| { - let idx = tail + offset as u64; - let slot = idx as usize % self.capacity; - unsafe { (*self.data[slot].get()).assume_init_ref().clone() } - }) - .collect() - } - - /// Save unsaved samples to binary file. - /// - /// Only saves samples that haven't been saved yet (from `saved_head` to `head`). - /// Call `mark_saved()` after a successful save to update the tracking. - /// - /// Returns the number of samples written. - /// - /// File format (version 2): - /// - magic: 8 bytes "SIEBREN\0" - /// - version: u64 (little endian) - /// - generation_id: u64 (little endian) - /// - sample_count: u64 (little endian) - /// - max_notation_len: u64 (little endian) - max length of any notation - /// - policy_len: u64 (little endian) - /// - samples: for each sample: - /// - notation_len: u64 (little endian) - actual length of this notation - /// - notation: [u8; max_notation_len] - UTF-8 bytes, padded with zeros - /// - policy: [f32; policy_len] (little endian) - /// - value: f32 (little endian) - pub fn save(&self, path: &Path, generation_id: u64, policy_len: usize) -> io::Result { - const MAGIC: &[u8; 8] = b"SIEBREN\0"; - const VERSION: u64 = 2; - - assert_eq!( - self.writers.load(Ordering::Acquire), - 0, - "cannot save while writers are active" - ); - - let head = self.head.load(Ordering::Acquire); - let saved_head = self.saved_head.load(Ordering::Acquire); - let (tail, _) = self.valid_range(); - - // Only save samples from saved_head to head, but not before tail - // (samples before tail have been overwritten in the ring buffer) - let start = saved_head.max(tail); - let count = head.saturating_sub(start) as usize; - - if count == 0 { - // Nothing new to save - still write an empty file for consistency - let file = File::create(path)?; - let mut writer = BufWriter::new(file); - writer.write_all(MAGIC)?; - writer.write_all(&VERSION.to_le_bytes())?; - writer.write_all(&generation_id.to_le_bytes())?; - writer.write_all(&0u64.to_le_bytes())?; // sample_count = 0 - writer.write_all(&0u64.to_le_bytes())?; // max_notation_len = 0 - writer.write_all(&(policy_len as u64).to_le_bytes())?; - writer.flush()?; - return Ok(0); - } - - // Find max notation length in the range we're saving - let mut max_notation_len = 0usize; - for idx in start..head { - let slot = idx as usize % self.capacity; - let sample = unsafe { (*self.data[slot].get()).assume_init_ref() }; - max_notation_len = max_notation_len.max(sample.notation.len()); - } - - let file = File::create(path)?; - let mut writer = BufWriter::new(file); - - // Write header - writer.write_all(MAGIC)?; - writer.write_all(&VERSION.to_le_bytes())?; - writer.write_all(&generation_id.to_le_bytes())?; - writer.write_all(&(count as u64).to_le_bytes())?; - writer.write_all(&(max_notation_len as u64).to_le_bytes())?; - writer.write_all(&(policy_len as u64).to_le_bytes())?; - - // Pre-allocate padding buffer - let mut notation_buf = vec![0u8; max_notation_len]; - - // Write samples - for idx in start..head { - let slot = idx as usize % self.capacity; - let sample = unsafe { (*self.data[slot].get()).assume_init_ref() }; - - // Write notation length and padded notation - let notation_bytes = sample.notation.as_bytes(); - writer.write_all(&(notation_bytes.len() as u64).to_le_bytes())?; - notation_buf[..notation_bytes.len()].copy_from_slice(notation_bytes); - notation_buf[notation_bytes.len()..].fill(0); - writer.write_all(¬ation_buf)?; - - // Write policy - for &p in &sample.policy { - writer.write_all(&p.to_le_bytes())?; - } - - // Write value - writer.write_all(&sample.value.to_le_bytes())?; - } - - writer.flush()?; - Ok(count) - } - - /// Mark all current samples as saved. - /// - /// Call this after a successful `save()` to prevent those samples from - /// being saved again in subsequent calls. - pub fn mark_saved(&self) { - let head = self.head.load(Ordering::Acquire); - self.saved_head.store(head, Ordering::Release); - } - - /// Load samples from binary file. - /// - /// Returns (samples_loaded, generation_id). - /// Panics if writers are active. - pub fn load(&self, path: &Path) -> io::Result<(usize, u64)> { - const MAGIC: &[u8; 8] = b"SIEBREN\0"; - - assert_eq!( - self.writers.load(Ordering::Acquire), - 0, - "cannot load while writers are active" - ); - - let file = File::open(path)?; - let mut reader = BufReader::new(file); - - // Read and validate magic - let mut magic = [0u8; 8]; - reader.read_exact(&mut magic)?; - if &magic != MAGIC { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "invalid magic bytes", - )); - } - - // Read and validate version - let mut buf8 = [0u8; 8]; - reader.read_exact(&mut buf8)?; - let version = u64::from_le_bytes(buf8); - if version != 2 { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!("unsupported version: {version}, expected 2"), - )); - } - - // Read header fields - reader.read_exact(&mut buf8)?; - let generation_id = u64::from_le_bytes(buf8); - - reader.read_exact(&mut buf8)?; - let sample_count = u64::from_le_bytes(buf8) as usize; - - reader.read_exact(&mut buf8)?; - let max_notation_len = u64::from_le_bytes(buf8) as usize; - - reader.read_exact(&mut buf8)?; - let policy_len = u64::from_le_bytes(buf8) as usize; - - // Pre-allocate buffers - let mut notation_buf = vec![0u8; max_notation_len]; - let mut policy_buf = vec![0u8; policy_len * 4]; - let mut buf4 = [0u8; 4]; - - // Reserve space and read samples - let mut guard = self.reserve(sample_count); - - for _ in 0..sample_count { - // Read notation length - reader.read_exact(&mut buf8)?; - let notation_len = u64::from_le_bytes(buf8) as usize; - - if notation_len > max_notation_len { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - format!( - "notation length {notation_len} exceeds max_notation_len {max_notation_len}" - ), - )); - } - - // Read padded notation - reader.read_exact(&mut notation_buf)?; - let notation = - String::from_utf8(notation_buf[..notation_len].to_vec()).map_err(|e| { - io::Error::new( - io::ErrorKind::InvalidData, - format!("invalid UTF-8 in notation: {e}"), - ) - })?; - - // Read policy - reader.read_exact(&mut policy_buf)?; - let policy: Vec = policy_buf - .chunks_exact(4) - .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap())) - .collect(); - - // Read value - reader.read_exact(&mut buf4)?; - let value = f32::from_le_bytes(buf4); - - guard.push(Sample { - notation, - policy, - value, - }); - } - - Ok((sample_count, generation_id)) - } -} - -impl Drop for ReplayBuffer { - fn drop(&mut self) { - for idx in 0..self.capacity { - if self.initialized[idx].load(Ordering::Relaxed) { - unsafe { - (*self.data[idx].get()).assume_init_drop(); - } - } - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use rand::SeedableRng; - use rand_chacha::ChaCha8Rng; - use std::sync::Arc; - - fn make_sample(id: usize) -> Sample { - Sample { - notation: format!("_________|A"), - policy: vec![id as f32 / 10.0; 9], - value: id as f32 / 100.0, - } - } - - #[test] - fn test_reserve_guard_push() { - let buffer = ReplayBuffer::new(100); - - { - let mut guard = buffer.reserve(3); - guard.push(make_sample(0)); - guard.push(make_sample(1)); - guard.push(make_sample(2)); - } - - assert_eq!(buffer.len(), 3); - assert_eq!(buffer.writers.load(Ordering::Relaxed), 0); - } - - #[test] - fn test_reserve_guard_extend() { - let buffer = ReplayBuffer::new(100); - - { - let mut guard = buffer.reserve(5); - guard.extend((0..5).map(make_sample)); - } - - assert_eq!(buffer.len(), 5); - } - - #[test] - fn test_sample() { - let buffer = ReplayBuffer::new(100); - let mut rng = ChaCha8Rng::seed_from_u64(42); - - { - let mut guard = buffer.reserve(10); - guard.extend((0..10).map(make_sample)); - } - - let samples = buffer.sample(5, &mut rng); - assert_eq!(samples.len(), 5); - } - - #[test] - fn test_concurrent_writes() { - let buffer = Arc::new(ReplayBuffer::new(1000)); - let num_threads = 4; - let samples_per_thread = 100; - - std::thread::scope(|s| { - for _ in 0..num_threads { - let buffer = Arc::clone(&buffer); - s.spawn(move || { - let mut guard = buffer.reserve(samples_per_thread); - guard.extend((0..samples_per_thread).map(make_sample)); - }); - } - }); - - assert_eq!(buffer.len(), num_threads * samples_per_thread); - assert_eq!(buffer.writers.load(Ordering::Relaxed), 0); - } - - #[test] - fn test_buffer_wraparound() { - let buffer = ReplayBuffer::new(10); - - for i in 0..25 { - let mut guard = buffer.reserve(1); - guard.push(make_sample(i)); + return SampleBatch { + observations: Array::default(D::with_batch(0, self.obs_shape.clone())), + values: Vec::new(), + }; } - assert_eq!(buffer.len(), 10); - let (tail, head) = buffer.valid_range(); - assert_eq!(tail, 15); - assert_eq!(head, 25); - } - - #[test] - #[should_panic(expected = "cannot sample while writers are active")] - fn test_sample_during_write_panics() { - let buffer = ReplayBuffer::new(100); - let mut rng = ChaCha8Rng::seed_from_u64(42); - - let _guard = buffer.reserve(5); - let _ = buffer.sample(5, &mut rng); - } - - #[test] - fn test_save_load_roundtrip() { - let buffer = ReplayBuffer::new(100); - let policy_len = 9; - let generation_id = 42u64; + let sample_count = n.min(count); + let indices = sample(rng, count, sample_count); - // Add samples - { - let mut guard = buffer.reserve(5); - guard.extend((0..5).map(make_sample)); - } - - // Save to temp file - let temp_dir = std::env::temp_dir(); - let path = temp_dir.join("test_replay_buffer.bin"); + let observations = unsafe { &*self.observations.get() }; + let observation_slice = observations + .as_slice_memory_order() + .expect("observation storage must be contiguous"); - let saved_count = buffer.save(&path, generation_id, policy_len).unwrap(); - assert_eq!(saved_count, 5); + let values = unsafe { &*self.values.get() }; - // Load into new buffer - let buffer2 = ReplayBuffer::new(100); - let (loaded_count, loaded_gen) = buffer2.load(&path).unwrap(); + let mut obs_data = Vec::with_capacity(sample_count * self.obs_elems_per_sample); + let mut value_data = Vec::with_capacity(sample_count); - assert_eq!(loaded_count, 5); - assert_eq!(loaded_gen, generation_id); - assert_eq!(buffer2.len(), 5); + for offset in indices.iter() { + let idx = (tail + offset as u64) as usize % self.capacity; - // Verify samples match - let mut rng = ChaCha8Rng::seed_from_u64(0); - let original = buffer.sample(5, &mut rng); - let mut rng = ChaCha8Rng::seed_from_u64(0); - let loaded = buffer2.sample(5, &mut rng); + let obs_start = idx * self.obs_elems_per_sample; + obs_data.extend_from_slice( + &observation_slice[obs_start..obs_start + self.obs_elems_per_sample], + ); - for (orig, load) in original.iter().zip(loaded.iter()) { - assert_eq!(orig.notation, load.notation); - assert_eq!(orig.policy, load.policy); - assert!((orig.value - load.value).abs() < 1e-6); + value_data.push(values[idx]); } - // Cleanup - std::fs::remove_file(&path).ok(); - } + let observations = Array::from_shape_vec( + D::with_batch(sample_count, self.obs_shape.clone()), + obs_data, + ) + .expect("sample observation shape mismatch"); - #[test] - fn test_save_load_generation_id() { - let buffer = ReplayBuffer::new(100); - - { - let mut guard = buffer.reserve(1); - guard.push(make_sample(0)); + SampleBatch { + observations, + values: value_data, } - - let temp_dir = std::env::temp_dir(); - let path = temp_dir.join("test_replay_gen_id.bin"); - - // Save with specific generation ID - let gen_id = 12345u64; - buffer.save(&path, gen_id, 9).unwrap(); - - // Load and verify generation ID is preserved - let buffer2 = ReplayBuffer::new(100); - let (_, loaded_gen) = buffer2.load(&path).unwrap(); - assert_eq!(loaded_gen, gen_id); - - // Cleanup - std::fs::remove_file(&path).ok(); - } - - #[test] - fn test_save_load_varying_notation_lengths() { - let buffer = ReplayBuffer::new(100); - - // Add samples with different notation lengths - { - let mut guard = buffer.reserve(3); - guard.push(Sample { - notation: "A".to_string(), - policy: vec![0.1; 9], - value: 0.5, - }); - guard.push(Sample { - notation: "ABCDEFGHIJ".to_string(), - policy: vec![0.2; 9], - value: 0.6, - }); - guard.push(Sample { - notation: "XYZ".to_string(), - policy: vec![0.3; 9], - value: 0.7, - }); - } - - let temp_dir = std::env::temp_dir(); - let path = temp_dir.join("test_replay_varying_notation.bin"); - - buffer.save(&path, 1, 9).unwrap(); - - let buffer2 = ReplayBuffer::new(100); - let (count, _) = buffer2.load(&path).unwrap(); - assert_eq!(count, 3); - - // Verify all notations preserved correctly - let mut rng = ChaCha8Rng::seed_from_u64(0); - let original = buffer.sample(3, &mut rng); - let mut rng = ChaCha8Rng::seed_from_u64(0); - let loaded = buffer2.sample(3, &mut rng); - - for (orig, load) in original.iter().zip(loaded.iter()) { - assert_eq!(orig.notation, load.notation); - } - - // Cleanup - std::fs::remove_file(&path).ok(); - } - - #[test] - fn test_save_empty_buffer() { - let buffer = ReplayBuffer::new(100); - - let temp_dir = std::env::temp_dir(); - let path = temp_dir.join("test_replay_empty.bin"); - - let saved_count = buffer.save(&path, 0, 9).unwrap(); - assert_eq!(saved_count, 0); - - let buffer2 = ReplayBuffer::new(100); - let (count, gen) = buffer2.load(&path).unwrap(); - assert_eq!(count, 0); - assert_eq!(gen, 0); - assert_eq!(buffer2.len(), 0); - - // Cleanup - std::fs::remove_file(&path).ok(); - } - - #[test] - fn test_mark_saved_prevents_duplicates() { - let buffer = ReplayBuffer::new(100); - let temp_dir = std::env::temp_dir(); - - // Add first batch of samples - { - let mut guard = buffer.reserve(3); - guard.extend((0..3).map(make_sample)); - } - - // Save first batch - let path1 = temp_dir.join("test_mark_saved_1.bin"); - let saved1 = buffer.save(&path1, 1, 9).unwrap(); - assert_eq!(saved1, 3); - buffer.mark_saved(); - - // Add second batch - { - let mut guard = buffer.reserve(2); - guard.extend((10..12).map(make_sample)); - } - - // Save second batch - should only save the new samples - let path2 = temp_dir.join("test_mark_saved_2.bin"); - let saved2 = buffer.save(&path2, 2, 9).unwrap(); - assert_eq!(saved2, 2); // Only the new samples - - // Buffer still has all 5 samples - assert_eq!(buffer.len(), 5); - - // Load both files into separate buffers and verify no duplicates - let buffer1 = ReplayBuffer::new(100); - let (count1, _) = buffer1.load(&path1).unwrap(); - assert_eq!(count1, 3); - - let buffer2 = ReplayBuffer::new(100); - let (count2, _) = buffer2.load(&path2).unwrap(); - assert_eq!(count2, 2); - - // Cleanup - std::fs::remove_file(&path1).ok(); - std::fs::remove_file(&path2).ok(); - } - - #[test] - fn test_save_without_mark_saved_resaves_all() { - let buffer = ReplayBuffer::new(100); - let temp_dir = std::env::temp_dir(); - - // Add samples - { - let mut guard = buffer.reserve(3); - guard.extend((0..3).map(make_sample)); - } - - // Save without calling mark_saved - let path1 = temp_dir.join("test_no_mark_saved_1.bin"); - let saved1 = buffer.save(&path1, 1, 9).unwrap(); - assert_eq!(saved1, 3); - // Note: NOT calling mark_saved() - - // Save again - should save the same samples again - let path2 = temp_dir.join("test_no_mark_saved_2.bin"); - let saved2 = buffer.save(&path2, 2, 9).unwrap(); - assert_eq!(saved2, 3); // Same 3 samples saved again - - // Cleanup - std::fs::remove_file(&path1).ok(); - std::fs::remove_file(&path2).ok(); - } - - #[test] - fn test_save_respects_ring_buffer_overwrites() { - // Small buffer that will wrap around - let buffer = ReplayBuffer::new(5); - let temp_dir = std::env::temp_dir(); - - // Add 3 samples - { - let mut guard = buffer.reserve(3); - guard.extend((0..3).map(make_sample)); - } - - // Save and mark - let path1 = temp_dir.join("test_overwrite_1.bin"); - let saved1 = buffer.save(&path1, 1, 9).unwrap(); - assert_eq!(saved1, 3); - buffer.mark_saved(); - - // Add 4 more samples - this will overwrite some of the original samples - // Buffer now contains samples at positions 3,4,5,6 (positions 0,1,2 overwritten) - { - let mut guard = buffer.reserve(4); - guard.extend((10..14).map(make_sample)); - } - - // Save again - should only save the 4 new samples - let path2 = temp_dir.join("test_overwrite_2.bin"); - let saved2 = buffer.save(&path2, 2, 9).unwrap(); - assert_eq!(saved2, 4); - - // Cleanup - std::fs::remove_file(&path1).ok(); - std::fs::remove_file(&path2).ok(); } } diff --git a/training/src/training.rs b/training/src/training.rs index 03cb97a..16a9f39 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -13,10 +13,10 @@ use rand_chacha::ChaCha8Rng; use crate::eval::{GpuEvaluator, PolicyValue}; use crate::executor::Executor; -use crate::observation_replay_buffer::ObservationReplayBuffer; use crate::queue::{BatchCompletion, GpuJobQueue}; +use crate::types::OutputType; use crate::worker::{worker_loop_forever, WorkerConfig}; -use crate::{BatchDim, Environment}; +use crate::BatchDim; /// Shared control state for the persistent self-play session. /// @@ -138,18 +138,14 @@ impl SelfPlaySession { /// /// Threads are spawned immediately but start paused. The `dispatch` callback /// is invoked when a batch of observations is ready for GPU inference. - pub fn new( - config: SessionConfig, - replay_buffer: Arc>, - dispatch: F, - ) -> Self + pub fn new(config: SessionConfig, replay_buffer: Arc, dispatch: F) -> Self where E: Environment + Clone + Send + 'static, E::ObsDim: BatchDim, F: Fn( usize, ArrayView::BatchedDim>, - BatchCompletion>, + BatchCompletion, ) + Send + Sync + 'static, @@ -159,8 +155,7 @@ impl SelfPlaySession { .checked_mul(config.workers_per_thread) .expect("num_threads * workers_per_thread overflowed usize"); - let queue: Arc>> = - Arc::new(GpuJobQueue::new(E::OBS_SHAPE, total_workers, dispatch)); + let queue = Arc::new(GpuJobQueue::new(E::OBS_SHAPE, total_workers, dispatch)); let control = Arc::new(SessionControl::new()); @@ -172,13 +167,7 @@ impl SelfPlaySession { let replay_buffer = replay_buffer.clone(); let handle = thread::spawn(move || { - session_thread_main::( - thread_id, - queue, - config, - control, - &replay_buffer, - ); + session_thread_main::(thread_id, queue, config, control, &replay_buffer); }); threads.push(handle); } diff --git a/training/src/types.rs b/training/src/types.rs new file mode 100644 index 0000000..de09616 --- /dev/null +++ b/training/src/types.rs @@ -0,0 +1,3 @@ +pub type ObservationElement = u8; +pub type ObservationDim = ndarray::Ix1; +pub type OutputType = f32; From 3f9dfe1afe78c29059bf84240e6fbc3ced38c796 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 01:30:01 -0400 Subject: [PATCH 06/59] Refactor training: remove mcts/types, add python training module --- python/alphapaint_training/__init__.py | 1 + .../alphapaint_training/cudagraph_backend.py | 146 +++++++ training/src/cudagraph.rs | 89 ++--- training/src/descent.rs | 359 +++++++++++------ training/src/eval.rs | 58 ++- training/src/executor.rs | 12 + training/src/lib.rs | 154 ++++---- training/src/mcts.rs | 294 -------------- training/src/replay_buffer.rs | 56 ++- training/src/training.rs | 97 ++--- training/src/types.rs | 3 - training/src/worker.rs | 365 +++++------------- 12 files changed, 697 insertions(+), 937 deletions(-) create mode 100644 python/alphapaint_training/__init__.py create mode 100644 python/alphapaint_training/cudagraph_backend.py delete mode 100644 training/src/mcts.rs delete mode 100644 training/src/types.rs diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py new file mode 100644 index 0000000..48ab266 --- /dev/null +++ b/python/alphapaint_training/__init__.py @@ -0,0 +1 @@ +"""AlphaPaint training infrastructure - Python bindings.""" diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py new file mode 100644 index 0000000..8a06a42 --- /dev/null +++ b/python/alphapaint_training/cudagraph_backend.py @@ -0,0 +1,146 @@ +"""CUDA graph capture for AlphaPaint value-only inference. + +Captures a CUDA graph that runs: H2D copy -> model forward -> D2H copy. +The model outputs only a scalar value (no policy head). +""" + +from typing import Optional + +import torch +import torch.utils.dlpack as dlpack + + +def _validate_tensors( + obs_host: torch.Tensor, + obs_device: torch.Tensor, + value_host: torch.Tensor, + value_device: torch.Tensor, + obs_side: int, + obs_width: int, +) -> None: + if obs_host.device.type != "cpu": + raise ValueError("obs_host must be a CPU tensor") + if value_host.device.type != "cpu": + raise ValueError("value_host must be a CPU tensor") + if obs_device.device.type != "cuda": + raise ValueError("obs_device must be a CUDA tensor") + if value_device.device.type != "cuda": + raise ValueError("value_device must be a CUDA tensor") + + if obs_host.dtype != torch.uint8 or obs_device.dtype != torch.uint8: + raise ValueError("obs_host/obs_device must be uint8") + if value_host.dtype != torch.float32 or value_device.dtype != torch.float32: + raise ValueError("value_host/value_device must be float32") + + if obs_host.ndim != 3 or obs_device.ndim != 3: + raise ValueError("obs tensors must be rank-3") + if value_host.ndim != 1 or value_device.ndim != 1: + raise ValueError("value tensors must be rank-1") + + batch = obs_host.shape[0] + if obs_host.shape != obs_device.shape: + raise ValueError("obs_host/obs_device shape mismatch") + if obs_host.shape[1] != obs_side or obs_host.shape[2] != obs_width: + raise ValueError( + f"obs must be shape (B, {obs_side}, {obs_width}), got {obs_host.shape}" + ) + if value_host.shape != value_device.shape: + raise ValueError("value_host/value_device shape mismatch") + if value_host.shape[0] != batch: + raise ValueError(f"value must be shape (B,) with B={batch}") + + +def _autocast_dtype(precision: str) -> Optional[torch.dtype]: + if precision == "fp32": + return None + if precision == "fp16": + return torch.float16 + if precision == "bf16": + if torch.cuda.is_bf16_supported(): + return torch.bfloat16 + return torch.float16 + raise ValueError(f"Unsupported precision: {precision}") + + +def capture_lane_graph( + model, + obs_host_dlpack, + obs_device_dlpack, + value_host_dlpack, + value_device_dlpack, + stream_handle: int, + precision: str = "fp32", +) -> tuple[int, object]: + """Capture a CUDA graph for value-only inference. + + The model should accept observations and return a scalar value tensor. + No policy head is used (Athenan-style training). + + Args: + model: PyTorch model that takes obs tensor and returns value tensor. + obs_host_dlpack: DLPack capsule for pinned host observation buffer. + obs_device_dlpack: DLPack capsule for device observation buffer. + value_host_dlpack: DLPack capsule for pinned host value buffer. + value_device_dlpack: DLPack capsule for device value buffer. + stream_handle: Raw CUDA stream handle. + precision: "fp32", "fp16", or "bf16". + + Returns: + Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). + """ + obs_host = dlpack.from_dlpack(obs_host_dlpack) + obs_device = dlpack.from_dlpack(obs_device_dlpack) + value_host = dlpack.from_dlpack(value_host_dlpack) + value_device = dlpack.from_dlpack(value_device_dlpack) + + obs_side = obs_host.shape[1] + obs_width = obs_host.shape[2] + + _validate_tensors( + obs_host, + obs_device, + value_host, + value_device, + obs_side, + obs_width, + ) + + model = model.cuda() + model.eval() + stream = torch.cuda.ExternalStream(stream_handle) + graph = torch.cuda.CUDAGraph(keep_graph=True) + dtype = _autocast_dtype(precision) + + def run_step() -> None: + obs_device.copy_(obs_host, non_blocking=True) + if dtype is None: + value = model(obs_device) + else: + with torch.autocast(device_type="cuda", dtype=dtype): + value = model(obs_device) + # Model returns value tensor, shape (B,) or (B, 1) + if value.ndim == 2: + value = value.squeeze(-1) + value_device.copy_(value, non_blocking=True) + value_host.copy_(value_device, non_blocking=True) + + with torch.inference_mode(): + with torch.cuda.stream(stream): + for _ in range(3): + run_step() + torch.cuda.synchronize() + + with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"): + run_step() + + graph.instantiate() + owner = ( + graph, + model, + obs_host, + obs_device, + value_host, + value_device, + stream, + ) + return int(graph.raw_cuda_graph_exec()), owner diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index e961d7b..c0c5bc8 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -1,4 +1,6 @@ -//! Rust-launched CUDA graph backend for ByteFight self-play inference. +//! Rust-launched CUDA graph backend for AlphaPaint value-only inference. +//! +//! No policy head - the model outputs only a scalar value per position. use std::ffi::{c_void, CStr}; use std::mem::size_of; @@ -10,7 +12,6 @@ use pyo3::exceptions::PyRuntimeError; use pyo3::ffi; use pyo3::prelude::*; -use crate::eval::PolicyValue; use crate::queue::BatchCompletion; #[allow(non_camel_case_types)] @@ -27,8 +28,13 @@ const DL_DEVICE_CUDA: i32 = 2; const DL_DTYPE_FLOAT: u8 = 2; const DL_DTYPE_UINT: u8 = 1; -const OBS_SIDE: usize = 18; -const OBS_WIDTH: usize = 16; +// Observation shape constants. These are placeholders until the actual +// convnet encoding is designed. The observation will be one-hot encoded +// board planes for a convolutional network. +// +// TODO: Update these to match the actual observation encoding. +pub const OBS_SIDE: usize = 18; +pub const OBS_WIDTH: usize = 16; const OBS_CELLS: usize = OBS_SIDE * OBS_WIDTH; #[repr(C)] @@ -230,7 +236,7 @@ fn cuda_malloc_device_u8(count: usize, context: &str) -> PyResult<*mut c_void> { Ok(ptr) } -struct ByteFightCudaGraphLane { +struct CudaGraphLane { stream: cudaStream_t, graph_exec: cudaGraphExec_t, /// Owns Python-side graph/tensor objects for this lane. @@ -255,48 +261,39 @@ unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { } let mut ctx = unsafe { Box::from_raw(user_data.cast::()) }; - let policy_src = unsafe { slice::from_raw_parts(ctx.policy_host, ctx.batch_size * ACTIONS) }; let value_src = unsafe { slice::from_raw_parts(ctx.value_host, ctx.batch_size) }; - let mut outputs = vec![PolicyValue::<7>::default(); ctx.batch_size]; - for (i, out) in outputs.iter_mut().enumerate() { - let start = i * ACTIONS; - out.policy - .copy_from_slice(&policy_src[start..start + ACTIONS]); - out.value = value_src[i]; - } + let outputs: Vec = value_src.to_vec(); if let Some(completion) = ctx.completion.take() { completion.complete(&outputs); } } -impl Drop for ByteFightCudaGraphLane { +impl Drop for CudaGraphLane { fn drop(&mut self) { unsafe { let _ = cuda::cudaFree(self.obs_dev); - let _ = cuda::cudaFree(self.policy_dev); let _ = cuda::cudaFree(self.value_dev); let _ = cuda::cudaFreeHost(self.obs_host.cast::()); - let _ = cuda::cudaFreeHost(self.policy_host.cast::()); let _ = cuda::cudaFreeHost(self.value_host.cast::()); let _ = cuda::cudaStreamDestroy(self.stream); } } } -/// Per-lane CUDA graph executor for ByteFight self-play inference. -pub struct ByteFightCudaGraphRunner { +/// Per-lane CUDA graph executor for AlphaPaint value-only inference. +pub struct CudaGraphRunner { batch_size: usize, - lanes: Vec, + lanes: Vec, } // SAFETY: Lane buffers/streams are independent per batch_idx and queue dispatch // ensures a lane is not reused before dispatch returns for that lane. -unsafe impl Send for ByteFightCudaGraphRunner {} -unsafe impl Sync for ByteFightCudaGraphRunner {} +unsafe impl Send for CudaGraphRunner {} +unsafe impl Sync for CudaGraphRunner {} -impl ByteFightCudaGraphRunner { +impl CudaGraphRunner { pub fn new( py: Python<'_>, model: Py, @@ -311,15 +308,13 @@ impl ByteFightCudaGraphRunner { return Err(PyErr::new::("batch_size must be > 0")); } - let module = PyModule::import(py, "siebren.cudagraph_backend")?; - let capture_fn = module.getattr("capture_bytefight_lane_graph")?; + let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; + let capture_fn = module.getattr("capture_lane_graph")?; let obs_count = batch_size * OBS_CELLS; - let policy_count = batch_size * ACTIONS; let value_count = batch_size; let obs_shape = [batch_size as i64, OBS_SIDE as i64, OBS_WIDTH as i64]; - let policy_shape = [batch_size as i64, ACTIONS as i64]; let value_shape = [batch_size as i64]; let mut lanes = Vec::with_capacity(num_lanes); @@ -335,10 +330,6 @@ impl ByteFightCudaGraphRunner { let obs_host = cuda_malloc_host_u8(obs_count, &format!("cudaMallocHost obs lane {}", lane_idx))?; - let policy_host = cuda_malloc_host_f32( - policy_count, - &format!("cudaMallocHost policy lane {}", lane_idx), - )?; let value_host = cuda_malloc_host_f32( value_count, &format!("cudaMallocHost value lane {}", lane_idx), @@ -346,10 +337,6 @@ impl ByteFightCudaGraphRunner { let obs_dev = cuda_malloc_device_u8(obs_count, &format!("cudaMalloc obs lane {}", lane_idx))?; - let policy_dev = cuda_malloc_device_f32( - policy_count, - &format!("cudaMalloc policy lane {}", lane_idx), - )?; let value_dev = cuda_malloc_device_f32( value_count, &format!("cudaMalloc value lane {}", lane_idx), @@ -366,24 +353,6 @@ impl ByteFightCudaGraphRunner { )?; let obs_dev_capsule = dlpack_capsule(py, obs_dev, &obs_shape, DL_DEVICE_CUDA, 0, DL_DTYPE_UINT, 8)?; - let policy_host_capsule = dlpack_capsule( - py, - policy_host.cast::(), - &policy_shape, - DL_DEVICE_CPU, - 0, - DL_DTYPE_FLOAT, - 32, - )?; - let policy_dev_capsule = dlpack_capsule( - py, - policy_dev, - &policy_shape, - DL_DEVICE_CUDA, - 0, - DL_DTYPE_FLOAT, - 32, - )?; let value_host_capsule = dlpack_capsule( py, value_host.cast::(), @@ -408,8 +377,6 @@ impl ByteFightCudaGraphRunner { model.clone_ref(py), obs_host_capsule, obs_dev_capsule, - policy_host_capsule, - policy_dev_capsule, value_host_capsule, value_dev_capsule, stream as u64, @@ -417,14 +384,12 @@ impl ByteFightCudaGraphRunner { ))? .extract()?; - let lane = ByteFightCudaGraphLane { + let lane = CudaGraphLane { stream, graph_exec: exec_handle as cudaGraphExec_t, _py_owner: py_owner, obs_host, obs_dev, - policy_host, - policy_dev, value_host, value_dev, }; @@ -438,18 +403,15 @@ impl ByteFightCudaGraphRunner { &self, batch_idx: usize, obs_view: ArrayView, - completion: BatchCompletion>, + completion: BatchCompletion, ) { - debug_assert_eq!( - obs_view.shape(), - &[self.batch_size, BYTEFIGHT_OBS_SIDE, BYTEFIGHT_OBS_WIDTH] - ); + debug_assert_eq!(obs_view.shape(), &[self.batch_size, OBS_SIDE, OBS_WIDTH]); let lane = &self.lanes[batch_idx % self.lanes.len()]; let obs_src = obs_view .as_slice() - .expect("bytefight queue observation batch must be contiguous"); + .expect("observation batch must be contiguous"); let obs_dst = unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_CELLS) }; obs_dst.copy_from_slice(obs_src); @@ -461,7 +423,6 @@ impl ByteFightCudaGraphRunner { ); let ctx = Box::new(LaneCompletionContext { - policy_host: lane.policy_host, value_host: lane.value_host, batch_size: self.batch_size, completion: Some(completion), diff --git a/training/src/descent.rs b/training/src/descent.rs index c268096..e16851a 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -1,15 +1,20 @@ +//! Async Descent/UBFM search for Athénan-style training. +//! +//! Uses GPU neural network evaluation for leaf expansion. Values are f32. +//! Implements tree learning (collect training samples from internal nodes) +//! and ordinal distribution for action selection. + use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; -use std::time::Duration; +use rand::{Rng, RngExt}; use crate::eval::Evaluator; #[derive(Debug)] pub struct ChildData { - action: Action, - child_value: i32, + pub action: Action, + pub child_value: f32, entrance_count: usize, pub node: Option>, } @@ -25,14 +30,14 @@ impl ChildData { #[derive(Debug)] pub struct SearchNode { - pub value: i32, + pub value: f32, pub completion_value: i32, resolved: bool, pub children: Vec, } impl SearchNode { - fn new(value: i32, completion_value: i32, is_resolved: bool) -> SearchNode { + fn new(value: f32, completion_value: i32, is_resolved: bool) -> SearchNode { SearchNode { value, completion_value, @@ -41,7 +46,7 @@ impl SearchNode { } } - /// Prefers moves with low entrance counts + /// Prefers moves with low entrance counts (exploration). fn completed_best_action_dual( &self, is_max_player: bool, @@ -51,26 +56,30 @@ impl SearchNode { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), - ) + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, -(a.entrance_count as i32)) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) .expect("Called completed_best_action_dual without any valid actions!") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .min_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) .expect("Called completed_best_action_dual without any valid actions!") }; @@ -78,32 +87,36 @@ impl SearchNode { (res.0, res.1.action) } - /// Prefers moves with high entrance counts + /// Prefers moves with high entrance counts (exploitation). fn completed_best_action(&self, is_max_player: bool, rng: &mut SmallRng) -> (usize, Action) { let res = if is_max_player { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) .expect("Called completed_best_action without any valid actions!") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), - ) + .min_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, -(a.entrance_count as i32)) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) .expect("Called completed_best_action without any valid actions!") }; @@ -119,22 +132,19 @@ impl SearchNode { child .node .as_ref() - .and_then(|c| Some(c.resolved)) + .map(|c| c.resolved) .unwrap_or(false) }) } } /// Build a chain of resolved SearchNodes for a killshot move sequence. - /// All nodes are resolved. Only the final node is terminal (no children). fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { let term_value = Self::value_from_term(board, terminal); let comp_value = terminal.value(); - // Start with the terminal leaf (the collision result) let mut node = SearchNode::new(term_value, comp_value, true); - // Build chain from last move to first for i in (0..moves.len()).rev() { let action = if i == moves.len() - 1 { Action::FinalMove(moves[i]) @@ -148,7 +158,7 @@ impl SearchNode { resolved: true, children: vec![ChildData { action, - child_value: comp_value, + child_value: term_value, entrance_count: 0, node: Some(Box::new(node)), }], @@ -160,10 +170,16 @@ impl SearchNode { node } - fn build_self(board: &Board, outcome: ApplyActionOutcome, rng: &mut SmallRng) -> SearchNode { + /// Expand a node: evaluate all children with the neural network. + async fn build_self( + board: &Board, + outcome: ApplyActionOutcome, + evaluator: &E, + rng: &mut SmallRng, + ) -> SearchNode { match outcome { ApplyActionOutcome::Ongoing => { - let mut new_node = SearchNode::new(0, 0, false); + let mut new_node = SearchNode::new(0.0, 0, false); let actions = board.get_valid_actions(); new_node.children.reserve(actions.len()); @@ -172,28 +188,33 @@ impl SearchNode { return SearchNode::new(Self::value_from_term(board, loss), loss.value(), true); } + // Evaluate all children with GPU neural net. + // We create all futures first (submitting observations eagerly), + // then await them. This enables batching across children. + struct ChildEvalResult { + action: Action, + value: f32, + } + + let mut eval_results = Vec::new(); + for action in actions.into_iter().copied() { let mut local_board = board.clone(); - let (outcome, _) = local_board.apply_action(action); - match outcome { + let (child_outcome, _) = local_board.apply_action(action); + match child_outcome { ApplyActionOutcome::Ongoing => { - new_node.children.push(ChildData { + // Evaluate this child with neural net + let value = evaluator.evaluate(&local_board).await; + eval_results.push(ChildEvalResult { action, - child_value: todo!(), - entrance_count: 0, - node: None, + value, }); } - // we can only guarantee that the play instead action is a valid action if - // its a Move (e.g. collision) or a Paint after a move - // Therefore, we just do the same thing as we did above, where we just - // treat this move as a terminal move too, and handle it when we consume - // the SearchTree (e.g. in bindings.rs) ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => { new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value: Self::value_from_term(&local_board, terminal), entrance_count: 0, node: Some(Box::new(SearchNode::new( Self::value_from_term(&local_board, terminal), @@ -203,10 +224,11 @@ impl SearchNode { }); } ApplyActionOutcome::Killshot { terminal, moves } => { - let chain = Self::build_killshot_chain(&local_board, terminal, &moves); + let chain = + Self::build_killshot_chain(&local_board, terminal, &moves); new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value: Self::value_from_term(&local_board, terminal), entrance_count: 0, node: Some(Box::new(chain)), }); @@ -214,6 +236,21 @@ impl SearchNode { } } + // Now store the neural net evaluation results + for result in eval_results { + new_node.children.push(ChildData { + action: result.action, + child_value: result.value, + entrance_count: 0, + node: None, + }); + } + + if new_node.children.is_empty() { + let loss = TerminalState::loss_for(board.is_white_turn()); + return SearchNode::new(Self::value_from_term(board, loss), loss.value(), true); + } + let (best_action_id, _) = new_node.completed_best_action(board.is_white_turn(), rng); new_node.completion_value = new_node.children[best_action_id].completion_value(); @@ -221,10 +258,6 @@ impl SearchNode { new_node.resolved = new_node.backup_resolution(); new_node } - // in search, we don't really CARE about the distinction between these two. - // e.g. the search tree doesn't _really_ care that we need to play a `Final` variant of - // the action instead. It just cares that this is a terminal state. - // We can just fix this in the consumers of the search tree, e.g. in bindings.rs ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => SearchNode::new( Self::value_from_term(board, terminal), @@ -237,14 +270,25 @@ impl SearchNode { } } - fn value_from_term(board: &Board, term: TerminalState) -> i32 { - term.value() * (2_000_000_000 - 5 * (board.turn_count as i32)) + /// Depth heuristic: larger magnitude for faster wins/losses. + fn value_from_term(board: &Board, term: TerminalState) -> f32 { + let sign = term.value() as f32; + // Scale: prefer faster wins. Max turns = 2000. + sign * (2000.0 - board.turn_count as f32) } - fn create_child(&mut self, mut state: Board, action: Action, rng: &mut SmallRng) -> i32 { + async fn create_child( + &mut self, + mut state: Board, + action: Action, + evaluator: &E, + rng: &mut SmallRng, + ) -> f32 { let (outcome, _) = state.apply_action(action); - let node = Box::new(SearchNode::build_self(&state, outcome, rng)); + let node = Box::new( + Box::pin(SearchNode::build_self(&state, outcome, evaluator, rng)).await, + ); let value = node.value; if let Some(id) = self @@ -258,17 +302,18 @@ impl SearchNode { value } - fn ubfms_iteration( + async fn ubfms_iteration( &mut self, mut state: Board, outcome: ApplyActionOutcome, + evaluator: &E, rng: &mut SmallRng, - ) -> i32 { + ) -> f32 { let white_turn = state.is_white_turn(); match outcome { ApplyActionOutcome::Ongoing => { - if self.children.len() == 0 { + if self.children.is_empty() { let loss = TerminalState::loss_for(white_turn); self.resolved = true; self.completion_value = loss.value(); @@ -282,12 +327,15 @@ impl SearchNode { self.children[best_action_id].entrance_count += 1; - if let Some(child_val) = self.children[best_action_id].node.as_mut() { + if self.children[best_action_id].node.is_some() { let (outcome, _) = state.apply_action(best_action); - child_val.ubfms_iteration(state, outcome, rng); + let child = self.children[best_action_id].node.as_mut().unwrap(); + Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng)).await; } else { - self.children[best_action_id].child_value = - self.create_child(state, best_action, rng); + self.children[best_action_id].child_value = Box::pin( + self.create_child(state, best_action, evaluator, rng), + ) + .await; } let (best_action_id, _) = self.completed_best_action(white_turn, rng); @@ -308,29 +356,45 @@ impl SearchNode { self.value } + + /// Collect tree learning samples: (minimax_value) from all internal nodes. + /// + /// An internal node is one that has children and at least one expanded child. + /// Non-terminal leaf nodes (where the network estimate was used without + /// minimax backing) are excluded per Athénan's tree learning rules. + fn collect_values(&self, out: &mut Vec) { + let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); + if has_expanded_child { + out.push(self.value); + } + + for child in &self.children { + if let Some(ref node) = child.node { + node.collect_values(out); + } + } + } } -pub struct GameSearchTree<'a, E> { +pub struct GameSearchTree<'a, E: Evaluator> { pub root_node: Box, - root_state: Board, + pub root_state: Board, rng: SmallRng, - evaluator: Evaluator, + evaluator: &'a E, } -impl GameSearchTree<'_> { +impl<'a, E: Evaluator> GameSearchTree<'a, E> { fn safest_action(&mut self) -> (usize, Action) { let val = if self.root_state.is_white_turn() { self.root_node .children .iter() .enumerate() - .max_by_key(|(_, child)| { - ( - child.completion_value(), - child.entrance_count, - child.child_value, - self.rng.next_u32(), - ) + .max_by(|(_, a), (_, b)| { + (a.completion_value(), a.entrance_count, a.child_value) + .partial_cmp(&(b.completion_value(), b.entrance_count, b.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) .expect("No valid action at root state!") } else { @@ -338,13 +402,15 @@ impl GameSearchTree<'_> { .children .iter() .enumerate() - .min_by_key(|(_, child)| { - ( - child.completion_value(), - -(child.entrance_count as isize), - child.child_value, - self.rng.next_u32(), - ) + .min_by(|(_, a), (_, b)| { + (a.completion_value(), -(a.entrance_count as isize), a.child_value) + .partial_cmp(&( + b.completion_value(), + -(b.entrance_count as isize), + b.child_value, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) .expect("No valid action at root state!") }; @@ -352,50 +418,43 @@ impl GameSearchTree<'_> { (val.0, val.1.action) } - /// This has the invariant that the board is NOT in a terminal state. - pub fn new<'a>(board: &Board) -> GameSearchTree<'a> { - let mut cpy = board.clone(); - let mut rng = SmallRng::seed_from_u64(123312); + /// Create a new search tree. The board must NOT be in a terminal state. + pub async fn new(board: &Board, evaluator: &'a E, rng: SmallRng) -> GameSearchTree<'a, E> { + let mut local_rng = rng; GameSearchTree { - root_node: Box::new(SearchNode::build_self( - &mut cpy, - ApplyActionOutcome::Ongoing, - &mut rng, - )), + root_node: Box::new( + SearchNode::build_self(board, ApplyActionOutcome::Ongoing, evaluator, &mut local_rng) + .await, + ), root_state: board.clone(), - rng, + rng: local_rng, + evaluator, } } - pub fn step_tree(&mut self, new_board: &Board, action_id: usize, outcome: ApplyActionOutcome) { + pub fn step_tree(&mut self, new_board: &Board, action_id: usize) { self.root_state = new_board.clone(); if self.root_node.children[action_id].node.is_some() { self.root_node = self.root_node.children[action_id].node.take().unwrap(); } else { - self.root_node = Box::new(SearchNode::build_self( - &mut self.root_state.clone(), - outcome, - &mut self.rng, - )) + // Can't reuse tree - would need async expansion. + // This shouldn't happen in normal play since we always expand before selecting. + panic!("step_tree called on unexpanded child"); } } - pub async fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let dcv = if self.root_state.is_white_turn() { + pub async fn run_descent_for_iter(&mut self, iterations: u32) { + let dcv: i32 = if self.root_state.is_white_turn() { 1 } else { -1 }; - let start_time = std::time::Instant::now(); if self.root_node.children.len() <= 1 { return; // No need to compute with only one option. } - for epoch in 0..iterations { - if start_time.elapsed() > max_duration && epoch >= 50 { - break; - } + for _epoch in 0..iterations { let ba = self.get_best_action_index(); if self.root_node.children[ba].entrance_count >= 6000 && self.root_node.children[ba].completion_value() != -dcv @@ -405,11 +464,14 @@ impl GameSearchTree<'_> { if self.root_node.children[ba].completion_value() == dcv { break; } - self.root_node.ubfms_iteration( - self.root_state.clone(), - ApplyActionOutcome::Ongoing, - &mut self.rng, - ); + self.root_node + .ubfms_iteration( + self.root_state.clone(), + ApplyActionOutcome::Ongoing, + self.evaluator, + &mut self.rng, + ) + .await; } } @@ -424,4 +486,65 @@ impl GameSearchTree<'_> { pub fn get_best_action_and_index(&mut self) -> (usize, Action) { self.safest_action() } + + /// Collect tree learning samples from all internal nodes. + /// + /// Returns minimax values from internal nodes (nodes with at least one + /// expanded child). These are the training targets for the neural network. + pub fn collect_tree_learning_values(&self) -> Vec { + let mut values = Vec::new(); + self.root_node.collect_values(&mut values); + values + } + + /// Select an action using Athénan's ordinal distribution. + /// + /// The ordinal distribution depends only on rank ordering of moves. + /// The exploitation rate is drawn uniformly from [0, 1] each time. + pub fn ordinal_select(&mut self) -> usize { + let n = self.root_node.children.len(); + if n <= 1 { + return 0; + } + + let is_white = self.root_state.is_white_turn(); + + // Sort children indices by value (best first for current player) + let mut indices: Vec = (0..n).collect(); + indices.sort_by(|&a, &b| { + let ca = &self.root_node.children[a]; + let cb = &self.root_node.children[b]; + if is_white { + (cb.completion_value(), cb.child_value) + .partial_cmp(&(ca.completion_value(), ca.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + } else { + (ca.completion_value(), ca.child_value) + .partial_cmp(&(cb.completion_value(), cb.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + } + }); + + // Draw exploitation rate uniformly from [0, 1] + let eps: f32 = self.rng.random(); + + // Ordinal distribution: P(c_i) = (eps + (1-eps)/(n-i)) * (1 - sum of previous) + let mut remaining = 1.0f32; + let r: f32 = self.rng.random(); + let mut cumulative = 0.0f32; + + for (rank, &idx) in indices.iter().enumerate() { + let slots_left = (n - rank) as f32; + let p = remaining * (eps + (1.0 - eps) / slots_left); + cumulative += p; + remaining -= p; + + if r < cumulative { + return idx; + } + } + + // Fallback: return the last (worst) action + *indices.last().unwrap() + } } diff --git a/training/src/eval.rs b/training/src/eval.rs index 7f73722..78ede1c 100644 --- a/training/src/eval.rs +++ b/training/src/eval.rs @@ -1,66 +1,56 @@ -//! Async evaluator trait and implementations for GPU inference. +//! Async evaluator for GPU inference. Value-only (no policy head). use std::future::Future; use alpha_paint::board::Board; +use ndarray::Ix2; use crate::queue::GpuJobQueue; -use crate::types::{ObservationDim, ObservationElement, OutputType}; - -/// Output from neural network evaluation: (policy logits, value). -/// Policy has one entry per possible action, value is in [-1, 1]. -#[derive(Clone, Copy)] -pub struct PolicyValue { - pub policy: [f32; NUM_ACTIONS], - pub value: f32, -} - -impl Default for PolicyValue { - fn default() -> Self { - Self { - policy: [0.0; NUM_ACTIONS], - value: 0.0, - } - } -} /// Async evaluator trait for neural network inference. +/// +/// Returns a scalar value estimate for the given board position. +/// Value is from the current player's perspective. pub trait Evaluator { - /// Evaluate the environment and return (policy, value). - /// Policy is over all actions, value is in [-1, 1] from current player's perspective. - fn evaluate(&self, board: &Board) -> impl Future; + /// Evaluate the board and return a value estimate. + fn evaluate(&self, board: &Board) -> impl Future; } /// GPU-backed evaluator that batches inference requests. /// -/// Wraps a GpuJobQueue and converts between Environment observations -/// and the queue's I/O types. +/// Wraps a GpuJobQueue and serializes Board state into observations +/// for GPU inference. Returns scalar value. pub struct GpuEvaluator<'a> { - queue: &'a GpuJobQueue, + queue: &'a GpuJobQueue, } impl<'a> GpuEvaluator<'a> { - pub fn new(queue: &'a GpuJobQueue) -> Self { + pub fn new(queue: &'a GpuJobQueue) -> Self { Self { queue } } } -impl<'a> Evaluator for GpuEvaluator<'a> { - fn evaluate(&self, board: &Board) -> impl Future { - // Submit immediately with callback that writes observation - let future = self.queue.eval(|out| todo!()); - - future +impl Evaluator for GpuEvaluator<'_> { + fn evaluate(&self, board: &Board) -> impl Future { + // Submit immediately with callback that writes observation. + // The observation encoding is a placeholder - must be implemented + // with the actual convnet encoding. + let _board = board.clone(); + let future = self.queue.eval(|_out| { + todo!("encode Board into observation tensor for convnet") + }); + + async move { future.await } } } /// Synchronous CPU evaluator for testing. /// -/// Returns uniform policy and zero value. +/// Returns zero value for all positions. pub struct UniformEvaluator; impl Evaluator for UniformEvaluator { - fn evaluate(&self, _: &Board) -> impl Future { + fn evaluate(&self, _: &Board) -> impl Future { std::future::ready(0.0) } } diff --git a/training/src/executor.rs b/training/src/executor.rs index a7a654a..5f29537 100644 --- a/training/src/executor.rs +++ b/training/src/executor.rs @@ -75,12 +75,20 @@ impl<'a> Executor<'a> { // Poll all futures until no progress loop { + if cancel() || futures.is_empty() { + return; + } + // Clear progress flag before polling round take_progress(); // Poll all pending futures let mut i = 0; while i < futures.len() { + if cancel() { + return; + } + let poll_result = futures[i].as_mut().poll(&mut cx); match poll_result { Poll::Ready(()) => { @@ -93,6 +101,10 @@ impl<'a> Executor<'a> { } } + if cancel() || futures.is_empty() { + return; + } + // If no progress was made, break to park if !take_progress() { break; diff --git a/training/src/lib.rs b/training/src/lib.rs index 6b9f0cf..69c91f1 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,32 +1,29 @@ use std::sync::Arc; use std::sync::{Mutex, OnceLock}; -use cudagraph::ByteFightCudaGraphRunner; -use eval::PolicyValue; -use mcts::MCTSConfig; +use cudagraph::CudaGraphRunner; use queue::{queue_shape_for_workers, BATCH_SIZE}; +use replay_buffer::ReplayBuffer; use training::{SelfPlaySession, SessionConfig}; use worker::WorkerConfig; use ndarray::{ArrayView, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; +use numpy::{PyArray, PyArrayMethods}; use pyo3::prelude::*; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; pub mod cudagraph; +pub mod descent; pub mod eval; pub mod executor; pub mod future; -pub mod mcts; pub mod queue; pub mod replay_buffer; pub mod training; -mod types; pub mod worker; /// Extension trait for prepending a batch dimension to a shape. -/// -/// The `BatchedDim` associated type is the dimension with batch prepended. -/// We use an associated type instead of `Dimension::Larger` so we can -/// constrain that `BatchedDim::Smaller == Self`. pub trait BatchDim: Dimension + Clone { type BatchedDim: Dimension + RemoveAxis; @@ -81,27 +78,83 @@ impl BatchDim for Ix5 { } } -struct ByteFightGraphCacheEntry { +struct GraphCacheEntry { model_ptr: usize, num_batches: usize, precision: String, - runner: Arc, + runner: Arc, } -static GRAPH_CACHE: OnceLock>> = OnceLock::new(); +static GRAPH_CACHE: OnceLock>> = OnceLock::new(); -fn graph_cache() -> &'static Mutex> { +fn graph_cache() -> &'static Mutex> { GRAPH_CACHE.get_or_init(|| Mutex::new(None)) } -/// Persistent ByteFight self-play session. + +/// Replay buffer storing (observation, value) pairs. /// -/// Uses the CUDA graph runner for GPU dispatch (no Python callback). +/// Observations are u8 tensors, values are f32 scalars (no policy). #[pyclass] -struct SelfPlay { - session: Option, +struct EphemeralReplayBuffer { + inner: Arc>, +} + +#[pymethods] +impl EphemeralReplayBuffer { + #[new] + fn new(capacity: usize) -> Self { + let obs_shape = Ix2(cudagraph::OBS_SIDE, cudagraph::OBS_WIDTH); + Self { + inner: Arc::new(ReplayBuffer::new(capacity, obs_shape)), + } + } + + fn __len__(&self) -> usize { + self.inner.len() + } + + #[getter] + fn capacity(&self) -> usize { + self.inner.capacity() + } + + /// Sample `n` items and return (observations, values) as numpy arrays. + fn sample<'py>( + &self, + py: Python<'py>, + n: usize, + seed: u64, + ) -> PyResult<( + Bound<'py, PyArray>, + Bound<'py, PyArray>, + )> { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let batch = self.inner.sample(n, &mut rng); + let num_samples = batch.values.len(); + + let obs_data = batch.observations.into_raw_vec_and_offset().0; + + let obs = PyArray::from_vec(py, obs_data) + .reshape(Ix3(num_samples, cudagraph::OBS_SIDE, cudagraph::OBS_WIDTH))?; + let values = PyArray::from_vec(py, batch.values); + + Ok((obs, values)) + } } -struct ReplayBuffer; +impl EphemeralReplayBuffer { + pub fn inner(&self) -> &Arc> { + &self.inner + } +} + +/// Persistent self-play session using Athénan (Descent + tree learning). +/// +/// Uses the CUDA graph runner for GPU dispatch (value-only, no policy). +#[pyclass] +struct SelfPlay { + session: Option, +} #[pymethods] impl SelfPlay { @@ -112,53 +165,23 @@ impl SelfPlay { workers_per_thread, seed, *, - mcts_num_simulations = 20, - mcts_c_puct = 1.5, - mcts_dirichlet_alpha = 0.3, - mcts_dirichlet_epsilon = 0.25, - temperature = 1.0, - exploration_moves = 30, + descent_iterations = 30, model, selfplay_precision = "fp32" ))] fn new( py: Python<'_>, - replay_buffer: &ReplayBuffer, + replay_buffer: &EphemeralReplayBuffer, num_threads: usize, workers_per_thread: usize, seed: u64, - mcts_num_simulations: usize, - mcts_c_puct: f32, - mcts_dirichlet_alpha: f32, - mcts_dirichlet_epsilon: f32, - temperature: f32, - exploration_moves: usize, + descent_iterations: u32, model: Py, selfplay_precision: &str, ) -> PyResult { - if mcts_num_simulations == 0 { - return Err(PyErr::new::( - "mcts_num_simulations must be >= 1", - )); - } - if mcts_c_puct <= 0.0 { - return Err(PyErr::new::( - "mcts_c_puct must be > 0", - )); - } - if mcts_dirichlet_alpha <= 0.0 { - return Err(PyErr::new::( - "mcts_dirichlet_alpha must be > 0", - )); - } - if !(0.0..=1.0).contains(&mcts_dirichlet_epsilon) { - return Err(PyErr::new::( - "mcts_dirichlet_epsilon must be in [0, 1]", - )); - } - if temperature < 0.0 { + if descent_iterations == 0 { return Err(PyErr::new::( - "temperature must be >= 0", + "descent_iterations must be >= 1", )); } @@ -167,16 +190,7 @@ impl SelfPlay { workers_per_thread, seed, worker: WorkerConfig { - mcts: MCTSConfig { - num_simulations: mcts_num_simulations, - c_puct: mcts_c_puct, - dirichlet_alpha: mcts_dirichlet_alpha, - dirichlet_epsilon: mcts_dirichlet_epsilon, - ..Default::default() - }, - temperature, - exploration_moves, - ..Default::default() + descent_iterations, }, }; @@ -191,7 +205,7 @@ impl SelfPlay { // Build or reuse the CUDA graph runner. let runner = { let cache = graph_cache(); - let mut guard = cache.lock().expect("bytefight graph cache mutex poisoned"); + let mut guard = cache.lock().expect("graph cache mutex poisoned"); let needs_rebuild = match guard.as_ref() { Some(entry) => { @@ -203,14 +217,14 @@ impl SelfPlay { }; if needs_rebuild { - let runner = Arc::new(ByteFightCudaGraphRunner::new( + let runner = Arc::new(CudaGraphRunner::new( py, model.clone_ref(py), num_batches, BATCH_SIZE, selfplay_precision, )?); - *guard = Some(ByteFightGraphCacheEntry { + *guard = Some(GraphCacheEntry { model_ptr, num_batches, precision: selfplay_precision.to_string(), @@ -229,11 +243,11 @@ impl SelfPlay { let dispatch = move |batch_idx: usize, obs_view: ArrayView, - completion: queue::BatchCompletion>| { + completion: queue::BatchCompletion| { runner.dispatch_async(batch_idx, obs_view, completion); }; - let session = SelfPlaySession::new::( + let session = SelfPlaySession::new( config, replay_buffer.inner().clone(), dispatch, @@ -293,8 +307,8 @@ impl Drop for SelfPlay { } #[pymodule] -fn siebren(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_class::()?; +fn alphapaint_training(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/training/src/mcts.rs b/training/src/mcts.rs deleted file mode 100644 index b44df59..0000000 --- a/training/src/mcts.rs +++ /dev/null @@ -1,294 +0,0 @@ -//! Async MCTS implementation for GPU-batched inference. -//! -//! This is an async adaptation of the sync MCTS in `mcts.rs`. -//! The key difference is that `expand()` awaits the evaluator. - -use rand::Rng; -use rand_distr::{Distribution, Gamma}; - -use crate::eval::Evaluator; -use crate::{Action, Environment, Player, TerminalState}; - -#[derive(Clone)] -pub struct MCTSConfig { - pub num_simulations: usize, - pub c_puct: f32, - pub dirichlet_alpha: f32, - pub dirichlet_epsilon: f32, -} - -impl Default for MCTSConfig { - fn default() -> Self { - Self { - num_simulations: 20, - c_puct: 1.5, - dirichlet_alpha: 0.3, - dirichlet_epsilon: 0.25, - } - } -} - -struct Node { - player: Player, - visit_count: u32, - value_sum: f32, - prior: f32, - children: Vec<(A, Node)>, -} - -impl Node { - fn new(prior: f32, player: Player) -> Self { - Self { - player, - visit_count: 0, - value_sum: 0.0, - prior, - children: Vec::new(), - } - } - - #[inline] - fn q(&self) -> f32 { - if self.visit_count == 0 { - 0.0 - } else { - self.value_sum / self.visit_count as f32 - } - } - - #[inline] - fn is_expanded(&self) -> bool { - !self.children.is_empty() - } -} - -pub struct MCTS<'a, E: Environment, V: Evaluator> { - config: &'a MCTSConfig, - evaluator: &'a V, - _phantom: std::marker::PhantomData, -} - -impl<'a, E: Environment, V: Evaluator> MCTS<'a, E, V> { - pub fn new(evaluator: &'a V, config: &'a MCTSConfig) -> Self { - Self { - config, - evaluator, - _phantom: std::marker::PhantomData, - } - } - - /// Run MCTS search and return visit counts for each action. - pub async fn search(&self, env: &mut E, rng: &mut impl Rng) -> Vec { - let mut root = Node::new(0.0, env.current_player()); - - self.expand(env, &mut root).await; - self.add_dirichlet_noise(&mut root, rng); - - for _ in 0..self.config.num_simulations { - self.run_simulation(env, &mut root).await; - } - - let mut counts = vec![0u32; E::NUM_ACTIONS]; - for (action, child) in &root.children { - counts[action.to_index()] = child.visit_count; - } - counts - } - - async fn run_simulation(&self, env: &mut E, root: &mut Node) { - let mut rollbacks = Vec::with_capacity(64); - self.traverse_and_expand(env, root, &mut rollbacks).await; - - for rb in rollbacks.into_iter().rev() { - env.rollback(rb); - } - } - - /// Q values stored from the perspective of the node's player. - /// Returns value from perspective of the node's player. - async fn traverse_and_expand( - &self, - env: &mut E, - node: &mut Node, - rollbacks: &mut Vec, - ) -> f32 { - if let Some(term) = env.is_terminal() { - let v = match term { - TerminalState::Win(winner) => { - if winner == node.player { - 1.0 - } else { - -1.0 - } - } - TerminalState::Draw => 0.0, - }; - node.visit_count += 1; - node.value_sum += v; - return v; - } - - if !node.is_expanded() { - let value = self.expand(env, node).await; - // value is from current_player's perspective, which equals node.player - node.visit_count += 1; - node.value_sum += value; - return value; - } - - let action = self.select_action(node); - rollbacks.push(env.apply_action(action)); - - let child = node - .children - .iter_mut() - .find(|(a, _)| *a == action) - .map(|(_, c)| c) - .unwrap(); - - // Box::pin for recursive async call - let child_value = Box::pin(self.traverse_and_expand(env, child, rollbacks)).await; - - // Convert child's value to this node's perspective - let value = if child.player == node.player { - child_value - } else { - -child_value - }; - node.visit_count += 1; - node.value_sum += value; - - value - } - - fn select_action(&self, node: &Node) -> E::Action { - let sqrt_n = (node.visit_count as f32).sqrt(); - - node.children - .iter() - .map(|(action, child)| { - // Q is from child's perspective; convert to parent's for comparison - let q = if child.player == node.player { - child.q() - } else { - -child.q() - }; - let ucb = q + self.config.c_puct * child.prior * sqrt_n - / (1.0 + child.visit_count as f32); - (action, ucb) - }) - .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) - .map(|(action, _)| *action) - .expect("select_action called on node with no children") - } - - async fn expand(&self, env: &mut E, node: &mut Node) -> f32 { - let (policy, value) = self.evaluator.evaluate(env).await; - let valid: Vec<_> = env.valid_actions().collect(); - - let mut priors: Vec<(E::Action, f32)> = valid - .into_iter() - .map(|a| (a, policy[a.to_index()].max(0.0))) - .collect(); - - let sum: f32 = priors.iter().map(|(_, p)| p).sum(); - if sum > 1e-8 { - for (_, p) in &mut priors { - *p /= sum; - } - } else { - let uniform = 1.0 / priors.len() as f32; - for (_, p) in &mut priors { - *p = uniform; - } - } - - node.children = priors - .into_iter() - .map(|(a, prior)| { - let rollback = env.apply_action(a); - let child_player = env.current_player(); - env.rollback(rollback); - (a, Node::new(prior, child_player)) - }) - .collect(); - - value - } - - fn add_dirichlet_noise(&self, root: &mut Node, rng: &mut impl Rng) { - if root.children.is_empty() { - return; - } - - let noise = sample_dirichlet(root.children.len(), self.config.dirichlet_alpha, rng); - - let eps = self.config.dirichlet_epsilon; - for ((_, child), n) in root.children.iter_mut().zip(noise) { - child.prior = (1.0 - eps) * child.prior + eps * n; - } - } -} - -fn sample_dirichlet(n: usize, alpha: f32, rng: &mut impl Rng) -> Vec { - let gamma = Gamma::new(alpha, 1.0).expect("invalid gamma params"); - let mut samples: Vec = (0..n).map(|_| gamma.sample(rng)).collect(); - let sum: f32 = samples.iter().sum(); - if sum > 0.0 { - for s in &mut samples { - *s /= sum; - } - } else { - let uniform = 1.0 / n as f32; - samples.fill(uniform); - } - samples -} - -/// Convert visit counts to a policy distribution. -pub fn visits_to_policy(visits: &[u32], temperature: f32) -> Vec { - if temperature < 1e-8 { - let mut policy = vec![0.0; visits.len()]; - if let Some(idx) = visits - .iter() - .enumerate() - .max_by_key(|(_, &v)| v) - .map(|(i, _)| i) - { - policy[idx] = 1.0; - } - return policy; - } - - let inv_t = 1.0 / temperature; - let powered: Vec = visits.iter().map(|&v| (v as f32).powf(inv_t)).collect(); - let sum: f32 = powered.iter().sum(); - - if sum > 0.0 { - powered.into_iter().map(|p| p / sum).collect() - } else { - vec![0.0; visits.len()] - } -} - -/// Sample an action index from a policy distribution. -pub fn sample_action_index(policy: &[f32], rng: &mut impl Rng) -> Option { - let r: f32 = rng.random(); - let mut cum = 0.0; - for (i, &p) in policy.iter().enumerate() { - cum += p; - if r < cum { - return Some(i); - } - } - policy.iter().rposition(|&p| p > 0.0) -} - -/// Get the action index with most visits. -pub fn best_action_index(visits: &[u32]) -> Option { - visits - .iter() - .enumerate() - .max_by_key(|(_, &v)| v) - .map(|(i, _)| i) -} diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs index 411da04..88ac817 100644 --- a/training/src/replay_buffer.rs +++ b/training/src/replay_buffer.rs @@ -9,7 +9,7 @@ use rand::Rng; use crate::BatchDim; -/// Batched sample output from [`ObservationReplayBuffer::sample`]. +/// Batched sample output from [`ReplayBuffer::sample`]. pub struct SampleBatch where A: Clone + Default + Send + Sync, @@ -19,7 +19,7 @@ where pub values: Vec, } -/// Lock-free ring buffer for storing observations, and values. +/// Lock-free ring buffer for storing observations and values. /// /// Observations are kept in single contiguous arrays: /// - observations: `(capacity, ...obs_shape)` @@ -138,7 +138,7 @@ where } } - pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D, NUM_ACTIONS> { + pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D> { assert!( n <= self.capacity, "cannot reserve more samples than buffer capacity" @@ -178,7 +178,7 @@ where } /// Sample `n` items uniformly. Panics if writers are active. - pub fn sample(&self, n: usize, rng: &mut impl Rng) -> SampleBatch { + pub fn sample(&self, n: usize, rng: &mut impl Rng) -> SampleBatch { assert_eq!( self.writers.load(Ordering::Acquire), 0, @@ -230,3 +230,51 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Ix1; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + #[test] + fn test_push_and_sample() { + let buffer = ReplayBuffer::::new(10, Ix1(4)); + + { + let mut guard = buffer.reserve(2); + guard.push(&[1, 2, 3, 4], 0.5); + guard.push(&[5, 6, 7, 8], -0.5); + } + + assert_eq!(buffer.len(), 2); + + let mut rng = ChaCha8Rng::seed_from_u64(7); + let batch = buffer.sample(2, &mut rng); + assert_eq!(batch.observations.shape(), &[2, 4]); + assert_eq!(batch.values.len(), 2); + } + + #[test] + fn test_wraparound_len() { + let buffer = ReplayBuffer::::new(3, Ix1(2)); + + for i in 0..10 { + let mut guard = buffer.reserve(1); + guard.push(&[i as i8, (i + 1) as i8], i as f32); + } + + assert_eq!(buffer.len(), 3); + } + + #[test] + #[should_panic(expected = "cannot sample while writers are active")] + fn test_sample_during_write_panics() { + let buffer = ReplayBuffer::::new(4, Ix1(2)); + let _guard = buffer.reserve(1); + + let mut rng = ChaCha8Rng::seed_from_u64(42); + let _ = buffer.sample(1, &mut rng); + } +} diff --git a/training/src/training.rs b/training/src/training.rs index 16a9f39..9677245 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -7,21 +7,18 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::thread; -use ndarray::ArrayView; +use ndarray::{ArrayView, Ix2, Ix3}; use rand::SeedableRng; use rand_chacha::ChaCha8Rng; -use crate::eval::{GpuEvaluator, PolicyValue}; +use crate::eval::GpuEvaluator; use crate::executor::Executor; use crate::queue::{BatchCompletion, GpuJobQueue}; -use crate::types::OutputType; +use crate::replay_buffer::ReplayBuffer; use crate::worker::{worker_loop_forever, WorkerConfig}; use crate::BatchDim; /// Shared control state for the persistent self-play session. -/// -/// Counters are behind `Arc` so they can be cloned directly -/// into worker futures that expect `Arc`. struct SessionControl { /// Whether workers should be actively polling. running: AtomicBool, @@ -54,12 +51,10 @@ impl SessionControl { } } - /// Wake all threads blocked on the condvar. fn wake_all(&self) { self.condvar.notify_all(); } - /// Check if the thread should stop polling (pause or shutdown). fn should_pause(&self) -> bool { if self.shutdown.load(Ordering::Acquire) { return true; @@ -79,7 +74,7 @@ pub struct SessionConfig { pub num_threads: usize, /// Number of workers per thread. pub workers_per_thread: usize, - /// Worker configuration (MCTS params, temperature, etc). + /// Worker configuration. pub worker: WorkerConfig, /// Random seed for reproducibility. pub seed: u64, @@ -96,8 +91,7 @@ impl Default for SessionConfig { } } -/// Trait-object wrapper so we can call `notify_all()` on the queue without -/// leaking the full generic type into `SelfPlaySession`. +/// Trait-object wrapper so we can call `notify_all()` on the queue. trait QueueNotify: Send + Sync { fn notify_all(&self); } @@ -115,47 +109,34 @@ where /// A persistent self-play session that owns worker threads and preserves /// in-progress game state across pause/resume boundaries. -/// -/// # Lifecycle -/// -/// 1. `new(...)` — creates threads and futures (paused). -/// 2. `start()` — sets target to `usize::MAX` and wakes threads. -/// 3. `wait_for(target)` — ensures running, blocks until `samples >= target`, -/// then pauses and waits for all pollers to quiesce. Safe to read replay -/// buffer after this returns. -/// 4. `samples()` — returns current absolute sample count. -/// 5. `shutdown()` / Rust `Drop` — sets shutdown, joins threads. pub struct SelfPlaySession { control: Arc, - /// Queue used by all workers. Kept alive for `notify_all` on drop. queue_notify: Arc, - /// Join handles for worker threads. `None` after `shutdown`. threads: Option>>, } impl SelfPlaySession { /// Create a new persistent session. /// - /// Threads are spawned immediately but start paused. The `dispatch` callback - /// is invoked when a batch of observations is ready for GPU inference. - pub fn new(config: SessionConfig, replay_buffer: Arc, dispatch: F) -> Self + /// Uses concrete AlphaPaint types: u8 observations with Ix2 shape, f32 values. + pub fn new( + config: SessionConfig, + replay_buffer: Arc>, + dispatch: F, + ) -> Self where - E: Environment + Clone + Send + 'static, - E::ObsDim: BatchDim, - F: Fn( - usize, - ArrayView::BatchedDim>, - BatchCompletion, - ) + Send - + Sync - + 'static, + F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, { let total_workers = config .num_threads .checked_mul(config.workers_per_thread) .expect("num_threads * workers_per_thread overflowed usize"); - let queue = Arc::new(GpuJobQueue::new(E::OBS_SHAPE, total_workers, dispatch)); + let obs_shape = Ix2( + crate::cudagraph::OBS_SIDE, + crate::cudagraph::OBS_WIDTH, + ); + let queue = Arc::new(GpuJobQueue::new(obs_shape, total_workers, dispatch)); let control = Arc::new(SessionControl::new()); @@ -167,7 +148,7 @@ impl SelfPlaySession { let replay_buffer = replay_buffer.clone(); let handle = thread::spawn(move || { - session_thread_main::(thread_id, queue, config, control, &replay_buffer); + session_thread_main(thread_id, queue, config, control, &replay_buffer); }); threads.push(handle); } @@ -179,8 +160,7 @@ impl SelfPlaySession { } } - /// Start self-play with no sample limit (runs until explicitly paused or - /// `wait_for` is called). + /// Start self-play with no sample limit. pub fn start(&self) { self.control .target_samples @@ -192,13 +172,7 @@ impl SelfPlaySession { /// Block until at least `target_samples` absolute samples have been /// collected, then pause and quiesce all workers. - /// - /// Returns the actual number of samples collected (may exceed target). - /// - /// After this returns, no worker thread is inside the executor polling - /// loop, so it is safe to read the replay buffer. pub fn wait_for(&self, target_samples: usize) -> usize { - // Set target and ensure running. self.control .target_samples .store(target_samples, Ordering::Release); @@ -206,7 +180,6 @@ impl SelfPlaySession { self.control.wake_all(); self.queue_notify.notify_all(); - // Wait until target is reached (condvar-based, no spinning). { let mut guard = self .control @@ -224,12 +197,10 @@ impl SelfPlaySession { } } - // Pause workers. self.control.running.store(false, Ordering::Release); self.queue_notify.notify_all(); self.control.wake_all(); - // Wait for all pollers to exit (quiesce). { let mut guard = self .control @@ -282,31 +253,19 @@ impl Drop for SelfPlaySession { } /// Main loop for a single thread in a persistent session. -/// -/// 1. Wait on condvar until `running || shutdown`. -/// 2. If shutdown => exit. -/// 3. Increment `active_pollers`, run executor until pause/shutdown/target. -/// 4. Decrement `active_pollers`, notify condvar so `wait_for()` can observe -/// quiesce. -/// 5. Goto 1. -fn session_thread_main( +fn session_thread_main( thread_id: usize, - queue: Arc>>, + queue: Arc>, config: SessionConfig, control: Arc, - replay_buffer: &ObservationReplayBuffer, -) where - E: Environment + Clone + 'static, - E::ObsDim: BatchDim, -{ + replay_buffer: &ReplayBuffer, +) { let base_seed = config.seed.wrapping_add(thread_id as u64 * 1000); - let evaluator = GpuEvaluator::::new(&*queue); + let evaluator = GpuEvaluator::new(&*queue); - // Clone the session's counters for workers. let samples_collected = control.samples_collected.clone(); let games_completed = control.games_completed.clone(); - // Create futures once. They live for the entire session. let mut futures: Vec + '_>>> = (0 ..config.workers_per_thread) .map(|i| { @@ -317,7 +276,7 @@ fn session_thread_main( let worker_config = &config.worker; let fut = async move { - worker_loop_forever::( + worker_loop_forever( evaluator_ref, worker_config, &mut rng, @@ -334,7 +293,6 @@ fn session_thread_main( let executor = Executor::new(|| queue.listen()); loop { - // 1. Wait until running or shutdown. { let mut guard = control .condvar_mutex @@ -347,12 +305,10 @@ fn session_thread_main( } } - // 2. If shutdown, exit. if control.shutdown.load(Ordering::Acquire) { return; } - // 3. Increment active_pollers and run executor. control.active_pollers.fetch_add(1, Ordering::AcqRel); let control_ref = &control; @@ -362,8 +318,6 @@ fn session_thread_main( if should_pause { queue_ref.notify_all(); } - // Notify the condvar when samples cross the target so wait_for() - // wakes up. if control_ref.samples_collected.load(Ordering::Acquire) >= control_ref.target_samples.load(Ordering::Acquire) { @@ -372,7 +326,6 @@ fn session_thread_main( should_pause }); - // 4. Decrement active_pollers and notify. control.active_pollers.fetch_sub(1, Ordering::AcqRel); control.wake_all(); } diff --git a/training/src/types.rs b/training/src/types.rs deleted file mode 100644 index de09616..0000000 --- a/training/src/types.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub type ObservationElement = u8; -pub type ObservationDim = ndarray::Ix1; -pub type OutputType = f32; diff --git a/training/src/worker.rs b/training/src/worker.rs index 4ecb6b8..befdb6b 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -1,323 +1,132 @@ -//! Worker loop and training data collection. +//! Worker loop for Athénan-style self-play with tree learning. //! -//! Each worker runs MCTS searches, plays games, and collects training samples. +//! Each worker runs Descent search, collects training samples from +//! internal tree nodes, and pushes them to the replay buffer. use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use rand::Rng; +use ndarray::Ix2; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; -use crate::eval::Evaluator; -use crate::mcts::{best_action_index, sample_action_index, visits_to_policy, MCTSConfig, MCTS}; -use crate::observation_replay_buffer::ObservationReplayBuffer; -use crate::{Action, Environment, Player, TerminalState}; - -/// A training sample from a single game step. -#[derive(Clone, Debug)] -pub struct TrainingSample { - /// Action index taken from this state. - pub action_idx: usize, - /// The policy from MCTS (normalized visit counts). - pub policy: Vec, - /// The value from MCTS search. - pub value: f32, - /// Player to move at this step. - /// - /// Keeping this here avoids reparsing game state while backfilling outcomes. - pub player: Player, -} +use alpha_paint::board::Board; -/// Full trace of one played self-play game. -#[derive(Clone, Debug)] -pub struct PlayedGame { - /// Initial environment state before any actions in this game. - pub initial_env: E, - /// Step samples in chronological order. - pub samples: Vec, -} +use crate::descent::GameSearchTree; +use crate::eval::Evaluator; +use crate::replay_buffer::ReplayBuffer; /// Configuration for the worker. #[derive(Clone)] pub struct WorkerConfig { - /// MCTS configuration. - pub mcts: MCTSConfig, - /// Temperature for action selection (1.0 = proportional to visits, 0.0 = argmax). - pub temperature: f32, - /// Number of moves at the start of the game to use exploration temperature. - /// After this many moves, use temperature 0 (argmax). - pub exploration_moves: usize, + /// Number of UBFM/Descent iterations per move. + pub descent_iterations: u32, } impl Default for WorkerConfig { fn default() -> Self { Self { - mcts: MCTSConfig::default(), - temperature: 1.0, - exploration_moves: 30, + descent_iterations: 30, } } } -/// Run a single self-play game, collecting training samples. +/// Run a single self-play game with tree learning. /// -/// Returns the collected samples. The game continues until terminal. -pub async fn play_game(evaluator: &V, config: &WorkerConfig, rng: &mut R) -> PlayedGame -where - E: Environment, - V: Evaluator, - R: Rng, -{ - let mut env = E::new(); - let initial_env = env.clone(); - let mut samples = Vec::new(); - let mut move_count = 0; - - let mcts = MCTS::new(evaluator, &config.mcts); - - loop { - if env.is_terminal().is_some() { - break; - } - - let visits = mcts.search(&mut env, rng).await; - // Convert visits to policy - let temp = if move_count < config.exploration_moves { - config.temperature - } else { - 0.0 - }; - let policy = visits_to_policy(&visits, temp); - let player = env.current_player(); - - // Value is set to 0.0 here and backfilled with game outcome after the game ends. - // This is standard AlphaZero practice - we use the actual game result rather than - // the search value estimate for training. - let value = 0.0; - - // Select action - let action_idx = if temp > 0.0 { - sample_action_index(&policy, rng) - } else { - best_action_index(&visits) - }; - - let action_idx = action_idx.expect("no valid actions but game not terminal"); - let action = E::Action::from_index(action_idx).expect("invalid action index"); - - // Record sample - samples.push(TrainingSample { - action_idx, - player, - policy, - value, - }); - - // Apply action - env.apply_action(action); - move_count += 1; - } - - // Backfill values with game outcome - let outcome = env.is_terminal().expect("game should be terminal"); - backfill_values(&mut samples, outcome); - - PlayedGame { - initial_env, - samples, - } -} - -/// Backfill sample values with the game outcome. -/// -/// For wins, the winner's moves get +1, loser's get -1. -/// For draws, all moves get 0. -fn backfill_values(samples: &mut [TrainingSample], outcome: TerminalState) { - for sample in samples.iter_mut() { - sample.value = match outcome { - TerminalState::Win(winner) => { - if sample.player == winner { - 1.0 - } else { - -1.0 - } - } - TerminalState::Draw => 0.0, - }; - } -} - -/// Run a worker loop that plays games until the target sample count is reached. -/// -/// Workers play games and increment `samples_collected` after each game. -/// When the counter reaches `target_samples`, workers stop. The executor's -/// cancel callback should check this condition to terminate remaining workers. +/// At each move: +/// 1. Run Descent search +/// 2. Collect (value) training samples from all internal tree nodes +/// 3. Select action via ordinal distribution +/// 4. Apply action, reuse tree /// -/// Samples are pushed directly to the shared `replay_buffer` after each completed game. -pub async fn worker_loop( - evaluator: &V, +/// Returns all collected training values from the search trees. +async fn play_game( + evaluator: &E, config: &WorkerConfig, rng: &mut R, - samples_collected: Arc, - games_completed: Arc, - target_samples: usize, - replay_buffer: &ObservationReplayBuffer, -) where - E: Environment + Clone, - V: Evaluator, - R: Rng, -{ - debug_assert_eq!(NUM_ACTIONS, E::NUM_ACTIONS); +) -> Vec { + // TODO: Create board from map pool. For now this is a placeholder. + let board: Board = todo!("board creation from map pool / Python"); - loop { - if samples_collected.load(Ordering::Acquire) >= target_samples { - break; - } + let mut all_values = Vec::new(); - let game = play_game::(evaluator, config, rng).await; - let num_samples = game.samples.len(); + // Check if terminal before starting + // (board.get_valid_actions().len() == 0 would indicate terminal) + let actions = board.get_valid_actions(); + if actions.len() == 0 { + return all_values; + } - // Push observations, policies, and values to replay buffer. - let mut guard = replay_buffer.reserve(num_samples); - let mut env = game.initial_env; - for sample in game.samples { - guard.push_with_observation(&sample.policy, sample.value, |out| env.observation(out)); - let action = - E::Action::from_index(sample.action_idx).expect("invalid action index in replay"); - env.apply_action(action); - } + let tree_rng = SmallRng::from_rng(rng); + let mut tree = GameSearchTree::new(&board, evaluator, tree_rng).await; - samples_collected.fetch_add(num_samples, Ordering::AcqRel); - games_completed.fetch_add(1, Ordering::AcqRel); + loop { + // Run descent search + tree.run_descent_for_iter(config.descent_iterations).await; + + // Collect tree learning samples from internal nodes + let values = tree.collect_tree_learning_values(); + all_values.extend(values); + + // Select action via ordinal distribution + let action_id = tree.ordinal_select(); + let action = tree.root_node.children[action_id].action; + + // Apply action to get new board state + let mut new_board = tree.root_state.clone(); + let (outcome, _) = new_board.apply_action(action); + + // Check if game is over + match outcome { + alpha_paint::board::ApplyActionOutcome::Terminal { .. } + | alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { + break; + } + alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { + // Terminal via play-instead + break; + } + alpha_paint::board::ApplyActionOutcome::Ongoing => { + // Step the tree to reuse it + tree.step_tree(&new_board, action_id); + } + } } + + all_values } /// Run a worker loop that plays games forever. /// -/// Unlike `worker_loop`, this never self-terminates based on a sample count. /// Stopping is handled externally by the executor's cancel/pause mechanism. -/// This is used by the persistent `SelfPlaySession` where pause/resume is -/// controlled at the session level, not inside the worker. -/// -/// Samples are pushed directly to the shared `replay_buffer` after each completed game. -pub async fn worker_loop_forever( - evaluator: &V, +/// Training samples (obs, value) are pushed to the shared replay buffer. +pub async fn worker_loop_forever( + evaluator: &E, config: &WorkerConfig, rng: &mut R, samples_collected: Arc, games_completed: Arc, - replay_buffer: &ObservationReplayBuffer, -) where - E: Environment + Clone, - V: Evaluator, - R: Rng, -{ - debug_assert_eq!(NUM_ACTIONS, E::NUM_ACTIONS); - + replay_buffer: &ReplayBuffer, +) { loop { - let game = play_game::(evaluator, config, rng).await; - let num_samples = game.samples.len(); - - // Push observations, policies, and values to replay buffer. - let mut guard = replay_buffer.reserve(num_samples); - let mut env = game.initial_env; - for sample in game.samples { - guard.push_with_observation(&sample.policy, sample.value, |out| env.observation(out)); - let action = - E::Action::from_index(sample.action_idx).expect("invalid action index in replay"); - env.apply_action(action); + let values = play_game(evaluator, config, rng).await; + let num_samples = values.len(); + + if num_samples > 0 { + // Push values to replay buffer. + // Observation encoding is a TODO - for now we write zeros. + let mut guard = replay_buffer.reserve(num_samples); + for value in values { + guard.push_with_observation(value, |_out| { + // TODO: encode the board state into the observation tensor. + // This requires reconstructing board states by replaying + // actions from the root, which needs the board + action history + // to be tracked during play_game. + }); + } } samples_collected.fetch_add(num_samples, Ordering::AcqRel); games_completed.fetch_add(1, Ordering::AcqRel); } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::environments::TicTacToe; - use crate::eval::UniformEvaluator; - use crate::executor::Executor; - use rand::SeedableRng; - use rand_chacha::ChaCha8Rng; - use std::cell::RefCell; - use std::rc::Rc; - - #[test] - fn test_play_game_collects_samples() { - let evaluator = UniformEvaluator; - let config = WorkerConfig { - mcts: MCTSConfig { - num_simulations: 50, - ..Default::default() - }, - ..Default::default() - }; - - let rng = Rc::new(RefCell::new(ChaCha8Rng::seed_from_u64(42))); - let result: Rc>>> = Rc::new(RefCell::new(None)); - - let rng_clone = rng.clone(); - let result_clone = result.clone(); - - let fut = async move { - let samples = - play_game::(&evaluator, &config, &mut *rng_clone.borrow_mut()) - .await; - *result_clone.borrow_mut() = Some(samples); - }; - - let event = event_listener::Event::new(); - let executor = Executor::new(|| event.listen()); - executor.run(&mut vec![Box::pin(fut)], &mut || false); - - let game = result.borrow_mut().take().unwrap(); - let samples = game.samples; - - // TicTacToe games are 5-9 moves - assert!(samples.len() >= 5); - assert!(samples.len() <= 9); - - // Each sample should have correct policy size - for sample in &samples { - assert_eq!(sample.policy.len(), 9); - // Policy should sum to ~1 - let sum: f32 = sample.policy.iter().sum(); - assert!((sum - 1.0).abs() < 0.01, "policy sum: {}", sum); - } - - // Values should be set (all -1, 0, or 1) - for sample in &samples { - assert!(sample.value == -1.0 || sample.value == 0.0 || sample.value == 1.0); - } - } - - #[test] - fn test_backfill_values_win() { - let mut samples = vec![TrainingSample { - action_idx: 0, - player: crate::Player::PlayerA, - policy: vec![], - value: 0.0, - }]; - - backfill_values(&mut samples, TerminalState::Win(crate::Player::PlayerA)); - assert_eq!(samples[0].value, 1.0); - - backfill_values(&mut samples, TerminalState::Win(crate::Player::PlayerB)); - assert_eq!(samples[0].value, -1.0); - } - - #[test] - fn test_backfill_values_draw() { - let mut samples = vec![TrainingSample { - action_idx: 0, - player: crate::Player::PlayerA, - policy: vec![], - value: 0.5, // Should be overwritten - }]; - - backfill_values(&mut samples, TerminalState::Draw); - assert_eq!(samples[0].value, 0.0); - } -} From 69db64ebd45c376bedbf186bade9c5228159056d Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 02:37:42 -0400 Subject: [PATCH 07/59] batch descent evaluations --- training/src/descent.rs | 67 +++++++++++++++++++++++++---------------- 1 file changed, 41 insertions(+), 26 deletions(-) diff --git a/training/src/descent.rs b/training/src/descent.rs index e16851a..368c2fe 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -57,7 +57,11 @@ impl SearchNode { .iter() .enumerate() .max_by(|&(_, a), &(_, b)| { - (a.completion_value(), a.child_value, -(a.entrance_count as i32)) + ( + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), + ) .partial_cmp(&( b.completion_value(), b.child_value, @@ -109,7 +113,11 @@ impl SearchNode { .iter() .enumerate() .min_by(|&(_, a), &(_, b)| { - (a.completion_value(), a.child_value, -(a.entrance_count as i32)) + ( + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), + ) .partial_cmp(&( b.completion_value(), b.child_value, @@ -128,13 +136,9 @@ impl SearchNode { if self.completion_value.abs() == 1 { true } else { - self.children.iter().all(|child| { - child - .node - .as_ref() - .map(|c| c.resolved) - .unwrap_or(false) - }) + self.children + .iter() + .all(|child| child.node.as_ref().map(|c| c.resolved).unwrap_or(false)) } } @@ -197,17 +201,23 @@ impl SearchNode { } let mut eval_results = Vec::new(); - - for action in actions.into_iter().copied() { + let batch = actions.into_iter().map(|action| { let mut local_board = board.clone(); - let (child_outcome, _) = local_board.apply_action(action); + let (child_outcome, rollback) = local_board.apply_action(*action); + // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics + // rust futures are normally lazily evaluated when you .await them for the + // first time. Ours are eager, which is why this code actually works as we'd expect. + let value = evaluator.evaluate(&local_board); + (action, local_board, child_outcome, rollback, value) + }); + + for (&action, mut local_board, child_outcome, _, value) in batch { match child_outcome { ApplyActionOutcome::Ongoing => { // Evaluate this child with neural net - let value = evaluator.evaluate(&local_board).await; eval_results.push(ChildEvalResult { action, - value, + value: value.await, }); } ApplyActionOutcome::Terminal { terminal } @@ -224,8 +234,7 @@ impl SearchNode { }); } ApplyActionOutcome::Killshot { terminal, moves } => { - let chain = - Self::build_killshot_chain(&local_board, terminal, &moves); + let chain = Self::build_killshot_chain(&local_board, terminal, &moves); new_node.children.push(ChildData { action, child_value: Self::value_from_term(&local_board, terminal), @@ -286,9 +295,8 @@ impl SearchNode { ) -> f32 { let (outcome, _) = state.apply_action(action); - let node = Box::new( - Box::pin(SearchNode::build_self(&state, outcome, evaluator, rng)).await, - ); + let node = + Box::new(Box::pin(SearchNode::build_self(&state, outcome, evaluator, rng)).await); let value = node.value; if let Some(id) = self @@ -332,10 +340,8 @@ impl SearchNode { let child = self.children[best_action_id].node.as_mut().unwrap(); Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng)).await; } else { - self.children[best_action_id].child_value = Box::pin( - self.create_child(state, best_action, evaluator, rng), - ) - .await; + self.children[best_action_id].child_value = + Box::pin(self.create_child(state, best_action, evaluator, rng)).await; } let (best_action_id, _) = self.completed_best_action(white_turn, rng); @@ -403,7 +409,11 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { .iter() .enumerate() .min_by(|(_, a), (_, b)| { - (a.completion_value(), -(a.entrance_count as isize), a.child_value) + ( + a.completion_value(), + -(a.entrance_count as isize), + a.child_value, + ) .partial_cmp(&( b.completion_value(), -(b.entrance_count as isize), @@ -423,8 +433,13 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { let mut local_rng = rng; GameSearchTree { root_node: Box::new( - SearchNode::build_self(board, ApplyActionOutcome::Ongoing, evaluator, &mut local_rng) - .await, + SearchNode::build_self( + board, + ApplyActionOutcome::Ongoing, + evaluator, + &mut local_rng, + ) + .await, ), root_state: board.clone(), rng: local_rng, From b48a98ef54e2506031009b041002ab51d5522223 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 02:42:54 -0400 Subject: [PATCH 08/59] add torch python dep lol --- pyproject.toml | 1 + training/src/descent.rs | 2 +- uv.lock | 299 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 301 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a92efbf..9fdae5d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "psutil==5.9.0", "cython==3.0.11", "py-cpuinfo", + "torch==2.10.0", ] classifiers = [ "Programming Language :: Rust", diff --git a/training/src/descent.rs b/training/src/descent.rs index 368c2fe..a1a98e3 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -206,7 +206,7 @@ impl SearchNode { let (child_outcome, rollback) = local_board.apply_action(*action); // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics // rust futures are normally lazily evaluated when you .await them for the - // first time. Ours are eager, which is why this code actually works as we'd expect. + // first time. Ours are eager, which is why this code actually works as we expect. let value = evaluator.evaluate(&local_board); (action, local_board, child_outcome, rollback, value) }); diff --git a/uv.lock b/uv.lock index 9ff6113..912cfc1 100644 --- a/uv.lock +++ b/uv.lock @@ -11,6 +11,7 @@ dependencies = [ { name = "numpy" }, { name = "psutil" }, { name = "py-cpuinfo" }, + { name = "torch" }, ] [package.dev-dependencies] @@ -26,6 +27,7 @@ requires-dist = [ { name = "numpy", specifier = "==2.1.3" }, { name = "psutil", specifier = "==5.9.0" }, { name = "py-cpuinfo" }, + { name = "torch", specifier = "==2.10.0" }, ] [package.metadata.requires-dev] @@ -43,6 +45,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "cuda-bindings" +version = "12.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/66/0c02bd330e7d976f83fa68583d6198d76f23581bcbb5c0e98a6148f326e5/cuda_pathfinder-1.5.0-py3-none-any.whl", hash = "sha256:498f90a9e9de36044a7924742aecce11c50c49f735f1bc53e05aa46de9ea4110", size = 49739, upload-time = "2026-03-24T21:14:30.869Z" }, +] + [[package]] name = "cython" version = "3.0.11" @@ -69,6 +90,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, ] +[[package]] +name = "filelock" +version = "3.25.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -78,6 +117,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + [[package]] name = "llvmlite" version = "0.44.0" @@ -91,6 +142,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, ] +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + [[package]] name = "numba" version = "0.61.0" @@ -126,6 +214,140 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a6/84/fa11dad3404b7634aaab50733581ce11e5350383311ea7a7010f464c0170/numpy-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:0d30c543f02e84e92c4b1f415b7c6b5326cbe45ee7882b6b77db7195fb971e3a", size = 12566858, upload-time = "2024-11-02T17:40:08.851Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.27.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + [[package]] name = "packaging" version = "26.0" @@ -196,3 +418,80 @@ sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791 wheels = [ { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, ] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "torch" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, + { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, + { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, + { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, + { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] From ce7d37d2aea6d29961660b1fc7b121dec2e02eb0 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 03:15:33 -0400 Subject: [PATCH 09/59] add packed uint16 observations with triton decode --- pyproject.toml | 1 + python/alphapaint_training/__init__.py | 28 ++ .../alphapaint_training/cudagraph_backend.py | 56 +--- python/alphapaint_training/model.py | 97 +++++++ python/alphapaint_training/packed_obs.py | 274 ++++++++++++++++++ training/src/cudagraph.rs | 54 ++-- training/src/descent.rs | 4 +- training/src/eval.rs | 25 +- training/src/lib.rs | 41 +-- training/src/observation.rs | 257 ++++++++++++++++ training/src/training.rs | 17 +- training/src/worker.rs | 4 +- uv.lock | 3 + 13 files changed, 731 insertions(+), 130 deletions(-) create mode 100644 python/alphapaint_training/model.py create mode 100644 python/alphapaint_training/packed_obs.py create mode 100644 training/src/observation.rs diff --git a/pyproject.toml b/pyproject.toml index 9fdae5d..38f4b58 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "cython==3.0.11", "py-cpuinfo", "torch==2.10.0", + "triton>=3.6.0", ] classifiers = [ "Programming Language :: Rust", diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 48ab266..604d707 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -1 +1,29 @@ """AlphaPaint training infrastructure - Python bindings.""" + +from .model import PackedValueModel, ResidualBlock, TinyValueNet +from .packed_obs import ( + BOARD_CELLS, + BOARD_PLANES, + BOARD_SIDE, + INTRINSIC_COUNT, + OBS_WORDS, + decode_intrinsics, + decode_packed_board, + decode_packed_board_reference, + decode_packed_observation, +) + +__all__ = [ + "BOARD_CELLS", + "BOARD_PLANES", + "BOARD_SIDE", + "INTRINSIC_COUNT", + "OBS_WORDS", + "PackedValueModel", + "ResidualBlock", + "TinyValueNet", + "decode_intrinsics", + "decode_packed_board", + "decode_packed_board_reference", + "decode_packed_observation", +] diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index 8a06a42..c391b20 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -1,7 +1,7 @@ """CUDA graph capture for AlphaPaint value-only inference. Captures a CUDA graph that runs: H2D copy -> model forward -> D2H copy. -The model outputs only a scalar value (no policy head). +Observations are packed uint16 words; the model is expected to decode them. """ from typing import Optional @@ -10,46 +10,6 @@ import torch.utils.dlpack as dlpack -def _validate_tensors( - obs_host: torch.Tensor, - obs_device: torch.Tensor, - value_host: torch.Tensor, - value_device: torch.Tensor, - obs_side: int, - obs_width: int, -) -> None: - if obs_host.device.type != "cpu": - raise ValueError("obs_host must be a CPU tensor") - if value_host.device.type != "cpu": - raise ValueError("value_host must be a CPU tensor") - if obs_device.device.type != "cuda": - raise ValueError("obs_device must be a CUDA tensor") - if value_device.device.type != "cuda": - raise ValueError("value_device must be a CUDA tensor") - - if obs_host.dtype != torch.uint8 or obs_device.dtype != torch.uint8: - raise ValueError("obs_host/obs_device must be uint8") - if value_host.dtype != torch.float32 or value_device.dtype != torch.float32: - raise ValueError("value_host/value_device must be float32") - - if obs_host.ndim != 3 or obs_device.ndim != 3: - raise ValueError("obs tensors must be rank-3") - if value_host.ndim != 1 or value_device.ndim != 1: - raise ValueError("value tensors must be rank-1") - - batch = obs_host.shape[0] - if obs_host.shape != obs_device.shape: - raise ValueError("obs_host/obs_device shape mismatch") - if obs_host.shape[1] != obs_side or obs_host.shape[2] != obs_width: - raise ValueError( - f"obs must be shape (B, {obs_side}, {obs_width}), got {obs_host.shape}" - ) - if value_host.shape != value_device.shape: - raise ValueError("value_host/value_device shape mismatch") - if value_host.shape[0] != batch: - raise ValueError(f"value must be shape (B,) with B={batch}") - - def _autocast_dtype(precision: str) -> Optional[torch.dtype]: if precision == "fp32": return None @@ -77,7 +37,7 @@ def capture_lane_graph( No policy head is used (Athenan-style training). Args: - model: PyTorch model that takes obs tensor and returns value tensor. + model: PyTorch model that takes packed obs tensor and returns value tensor. obs_host_dlpack: DLPack capsule for pinned host observation buffer. obs_device_dlpack: DLPack capsule for device observation buffer. value_host_dlpack: DLPack capsule for pinned host value buffer. @@ -93,18 +53,6 @@ def capture_lane_graph( value_host = dlpack.from_dlpack(value_host_dlpack) value_device = dlpack.from_dlpack(value_device_dlpack) - obs_side = obs_host.shape[1] - obs_width = obs_host.shape[2] - - _validate_tensors( - obs_host, - obs_device, - value_host, - value_device, - obs_side, - obs_width, - ) - model = model.cuda() model.eval() stream = torch.cuda.ExternalStream(stream_handle) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py new file mode 100644 index 0000000..7e53d27 --- /dev/null +++ b/python/alphapaint_training/model.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +import torch +from torch import nn + +from .packed_obs import BOARD_PLANES, INTRINSIC_COUNT, decode_packed_observation + + +class ResidualBlock(nn.Module): + def __init__(self, width: int): + super().__init__() + groups = min(8, width) + self.norm1 = nn.GroupNorm(groups, width) + self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.norm2 = nn.GroupNorm(groups, width) + self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.act = nn.GELU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.norm1(x) + x = self.act(x) + x = self.conv1(x) + x = self.norm2(x) + x = self.act(x) + x = self.conv2(x) + return x + residual + + +class TinyValueNet(nn.Module): + def __init__( + self, + *, + width: int = 32, + embedding_dim: int = 64, + hidden_dim: int = 64, + ): + super().__init__() + groups = min(8, width) + self.stem = nn.Sequential( + nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), + nn.GroupNorm(groups, width), + nn.GELU(), + ) + self.blocks = nn.Sequential(ResidualBlock(width), ResidualBlock(width)) + self.pool = nn.AdaptiveAvgPool2d(1) + self.project = nn.Sequential( + nn.Flatten(), + nn.Linear(width, embedding_dim), + nn.GELU(), + ) + self.head = nn.Sequential( + nn.Linear(embedding_dim + INTRINSIC_COUNT, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, 1), + ) + + def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: + x = self.stem(board) + x = self.blocks(x) + x = self.pool(x) + x = self.project(x) + x = torch.cat((x, intrinsics), dim=1) + return self.head(x) + + +class PackedValueModel(nn.Module): + def __init__( + self, + *, + width: int = 32, + embedding_dim: int = 64, + hidden_dim: int = 64, + board_dtype: torch.dtype = torch.bfloat16, + ): + super().__init__() + self.board_dtype = board_dtype + self.value_net = TinyValueNet( + width=width, + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + ) + + def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return decode_packed_observation(packed_obs, board_dtype=self.board_dtype) + + def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: + board, intrinsics = self.decode(packed_obs) + if not torch.is_autocast_enabled(): + param_dtype = next(self.value_net.parameters()).dtype + board = board.to(dtype=param_dtype) + intrinsics = intrinsics.to(dtype=param_dtype) + value = self.value_net(board, intrinsics) + return value.squeeze(-1) + + +__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py new file mode 100644 index 0000000..168cefd --- /dev/null +++ b/python/alphapaint_training/packed_obs.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +BOARD_SIDE = 32 +BOARD_CELLS = BOARD_SIDE * BOARD_SIDE +BOARD_PLANES = 17 +INTRINSIC_COUNT = 10 +OBS_WORDS = BOARD_CELLS + INTRINSIC_COUNT + +PAINT_STRENGTH_MASK = 0b111 +PAINT_IS_ENEMY_BIT = 1 << 3 +WALL_BIT = 1 << 4 +POWERUP_BIT = 1 << 5 +BEACON_SHIFT = 6 +HILL_SHIFT = 8 +CURRENT_PLAYER_BIT = 1 << 10 +OPPONENT_PLAYER_BIT = 1 << 11 + +BEACON_CURRENT = 1 +BEACON_OPPONENT = 2 + +HILL_NEUTRAL = 1 +HILL_CURRENT = 2 +HILL_OPPONENT = 3 + +INTRINSIC_SCALE = (420.0, 420.0, 8.0, 8.0, 2000.0, 1024.0, 1024.0, 1024.0, 1024.0, 64.0) + + +def _check_packed_obs(packed_obs: torch.Tensor) -> None: + if packed_obs.dtype != torch.uint16: + raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") + if packed_obs.ndim != 2: + raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") + if packed_obs.shape[1] != OBS_WORDS: + raise ValueError( + f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" + ) + + +@triton.jit +def _decode_board_kernel( + packed_ptr, + out_ptr, + total_cells, + packed_stride0, + out_stride0, + out_stride1, + out_stride2, + out_stride3, + BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < total_cells + + batch = offs // 1024 + cell = offs % 1024 + x = cell % 32 + y = cell // 32 + + words = tl.load(packed_ptr + batch * packed_stride0 + cell, mask=mask, other=0).to( + tl.uint16 + ) + + strength = words & 0b111 + enemy = (words >> 3) & 1 + wall = (words >> 4) & 1 + powerup = (words >> 5) & 1 + beacon = (words >> 6) & 0b11 + hill = (words >> 8) & 0b11 + current_player = (words >> 10) & 1 + opponent_player = (words >> 11) & 1 + + current_paint = (strength != 0) & (enemy == 0) + opponent_paint = (strength != 0) & (enemy != 0) + + out_base = batch * out_stride0 + y * out_stride2 + x * out_stride3 + + tl.store( + out_ptr + out_base + 0 * out_stride1, current_paint.to(tl.float32), mask=mask + ) + tl.store( + out_ptr + out_base + 1 * out_stride1, + (current_paint & (strength >= 2)).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 2 * out_stride1, + (current_paint & (strength >= 3)).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 3 * out_stride1, + (current_paint & (strength >= 4)).to(tl.float32), + mask=mask, + ) + + tl.store( + out_ptr + out_base + 4 * out_stride1, opponent_paint.to(tl.float32), mask=mask + ) + tl.store( + out_ptr + out_base + 5 * out_stride1, + (opponent_paint & (strength >= 2)).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 6 * out_stride1, + (opponent_paint & (strength >= 3)).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 7 * out_stride1, + (opponent_paint & (strength >= 4)).to(tl.float32), + mask=mask, + ) + + tl.store(out_ptr + out_base + 8 * out_stride1, wall.to(tl.float32), mask=mask) + tl.store(out_ptr + out_base + 9 * out_stride1, powerup.to(tl.float32), mask=mask) + tl.store( + out_ptr + out_base + 10 * out_stride1, + (beacon == 1).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 11 * out_stride1, + (beacon == 2).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 12 * out_stride1, + (hill == 1).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 13 * out_stride1, + (hill == 2).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 14 * out_stride1, + (hill == 3).to(tl.float32), + mask=mask, + ) + tl.store( + out_ptr + out_base + 15 * out_stride1, current_player.to(tl.float32), mask=mask + ) + tl.store( + out_ptr + out_base + 16 * out_stride1, opponent_player.to(tl.float32), mask=mask + ) + + +def decode_packed_board_reference(board_words: torch.Tensor) -> torch.Tensor: + if board_words.dtype != torch.uint16: + raise ValueError(f"board_words must be uint16, got {board_words.dtype}") + if board_words.ndim != 2 or board_words.shape[1] != BOARD_CELLS: + raise ValueError( + f"board_words must have shape (B, {BOARD_CELLS}), got {tuple(board_words.shape)}" + ) + + words = board_words.to(torch.int32) + strength = words & PAINT_STRENGTH_MASK + enemy = (words & PAINT_IS_ENEMY_BIT) != 0 + + current_paint = (strength != 0) & (~enemy) + opponent_paint = (strength != 0) & enemy + beacon = (words >> BEACON_SHIFT) & 0b11 + hill = (words >> HILL_SHIFT) & 0b11 + + planes = torch.stack( + [ + current_paint, + current_paint & (strength >= 2), + current_paint & (strength >= 3), + current_paint & (strength >= 4), + opponent_paint, + opponent_paint & (strength >= 2), + opponent_paint & (strength >= 3), + opponent_paint & (strength >= 4), + (words & WALL_BIT) != 0, + (words & POWERUP_BIT) != 0, + beacon == BEACON_CURRENT, + beacon == BEACON_OPPONENT, + hill == HILL_NEUTRAL, + hill == HILL_CURRENT, + hill == HILL_OPPONENT, + (words & CURRENT_PLAYER_BIT) != 0, + (words & OPPONENT_PLAYER_BIT) != 0, + ], + dim=1, + ) + return planes.to(torch.float32).reshape(-1, BOARD_PLANES, BOARD_SIDE, BOARD_SIDE) + + +def decode_packed_board( + board_words: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + if board_words.dtype != torch.uint16: + raise ValueError(f"board_words must be uint16, got {board_words.dtype}") + if board_words.ndim != 2 or board_words.shape[1] != BOARD_CELLS: + raise ValueError( + f"board_words must have shape (B, {BOARD_CELLS}), got {tuple(board_words.shape)}" + ) + + if board_words.device.type != "cuda": + return decode_packed_board_reference(board_words).to(dtype=dtype) + + board_words = board_words.contiguous() + out = torch.empty( + (board_words.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), + device=board_words.device, + dtype=dtype, + ) + total_cells = board_words.shape[0] * BOARD_CELLS + grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) + _decode_board_kernel[grid]( + board_words, + out, + total_cells, + board_words.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + BLOCK=256, + num_warps=4, + ) + return out + + +def decode_intrinsics( + packed_obs: torch.Tensor, + *, + dtype: torch.dtype = torch.bfloat16, +) -> torch.Tensor: + _check_packed_obs(packed_obs) + scales = torch.tensor( + INTRINSIC_SCALE, device=packed_obs.device, dtype=torch.float32 + ) + intrinsics = packed_obs[:, BOARD_CELLS:].to(torch.float32) + intrinsics = intrinsics / scales + return intrinsics.to(dtype=dtype) + + +def decode_packed_observation( + packed_obs: torch.Tensor, + *, + board_dtype: torch.dtype = torch.bfloat16, + intrinsic_dtype: torch.dtype | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + _check_packed_obs(packed_obs) + board = decode_packed_board(packed_obs[:, :BOARD_CELLS], dtype=board_dtype) + intrinsics = decode_intrinsics( + packed_obs, + dtype=board_dtype if intrinsic_dtype is None else intrinsic_dtype, + ) + return board, intrinsics + + +__all__ = [ + "BOARD_CELLS", + "BOARD_PLANES", + "BOARD_SIDE", + "INTRINSIC_COUNT", + "OBS_WORDS", + "decode_intrinsics", + "decode_packed_board", + "decode_packed_board_reference", + "decode_packed_observation", +] diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index c0c5bc8..90e23be 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -7,7 +7,7 @@ use std::mem::size_of; use std::slice; use cudarc::runtime::sys as cuda; -use ndarray::{ArrayView, Ix3}; +use ndarray::{ArrayView, Ix2}; use pyo3::exceptions::PyRuntimeError; use pyo3::ffi; use pyo3::prelude::*; @@ -28,14 +28,7 @@ const DL_DEVICE_CUDA: i32 = 2; const DL_DTYPE_FLOAT: u8 = 2; const DL_DTYPE_UINT: u8 = 1; -// Observation shape constants. These are placeholders until the actual -// convnet encoding is designed. The observation will be one-hot encoded -// board planes for a convolutional network. -// -// TODO: Update these to match the actual observation encoding. -pub const OBS_SIDE: usize = 18; -pub const OBS_WIDTH: usize = 16; -const OBS_CELLS: usize = OBS_SIDE * OBS_WIDTH; +pub const OBS_WORDS: usize = crate::observation::OBS_WORDS; #[repr(C)] struct DLDevice { @@ -198,16 +191,18 @@ fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { Ok(ptr.cast::()) } -fn cuda_malloc_host_u8(count: usize, context: &str) -> PyResult<*mut u8> { +fn cuda_malloc_host_u16(count: usize, context: &str) -> PyResult<*mut u16> { let mut ptr: *mut c_void = std::ptr::null_mut(); - let bytes = count; + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("host allocation size overflow"))?; unsafe { check_cuda( cuda::cudaMallocHost(&mut ptr as *mut *mut c_void, bytes), context, )?; } - Ok(ptr.cast::()) + Ok(ptr.cast::()) } fn cuda_malloc_device_f32(count: usize, context: &str) -> PyResult<*mut c_void> { @@ -224,9 +219,11 @@ fn cuda_malloc_device_f32(count: usize, context: &str) -> PyResult<*mut c_void> Ok(ptr) } -fn cuda_malloc_device_u8(count: usize, context: &str) -> PyResult<*mut c_void> { +fn cuda_malloc_device_u16(count: usize, context: &str) -> PyResult<*mut c_void> { let mut ptr: *mut c_void = std::ptr::null_mut(); - let bytes = count; + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("device allocation size overflow"))?; unsafe { check_cuda( cuda::cudaMalloc(&mut ptr as *mut *mut c_void, bytes), @@ -241,7 +238,7 @@ struct CudaGraphLane { graph_exec: cudaGraphExec_t, /// Owns Python-side graph/tensor objects for this lane. _py_owner: Py, - obs_host: *mut u8, + obs_host: *mut u16, obs_dev: *mut c_void, value_host: *mut f32, value_dev: *mut c_void, @@ -311,10 +308,10 @@ impl CudaGraphRunner { let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; let capture_fn = module.getattr("capture_lane_graph")?; - let obs_count = batch_size * OBS_CELLS; + let obs_count = batch_size * OBS_WORDS; let value_count = batch_size; - let obs_shape = [batch_size as i64, OBS_SIDE as i64, OBS_WIDTH as i64]; + let obs_shape = [batch_size as i64, OBS_WORDS as i64]; let value_shape = [batch_size as i64]; let mut lanes = Vec::with_capacity(num_lanes); @@ -329,14 +326,14 @@ impl CudaGraphRunner { } let obs_host = - cuda_malloc_host_u8(obs_count, &format!("cudaMallocHost obs lane {}", lane_idx))?; + cuda_malloc_host_u16(obs_count, &format!("cudaMallocHost obs lane {}", lane_idx))?; let value_host = cuda_malloc_host_f32( value_count, &format!("cudaMallocHost value lane {}", lane_idx), )?; let obs_dev = - cuda_malloc_device_u8(obs_count, &format!("cudaMalloc obs lane {}", lane_idx))?; + cuda_malloc_device_u16(obs_count, &format!("cudaMalloc obs lane {}", lane_idx))?; let value_dev = cuda_malloc_device_f32( value_count, &format!("cudaMalloc value lane {}", lane_idx), @@ -349,10 +346,17 @@ impl CudaGraphRunner { DL_DEVICE_CPU, 0, DL_DTYPE_UINT, - 8, + 16, + )?; + let obs_dev_capsule = dlpack_capsule( + py, + obs_dev, + &obs_shape, + DL_DEVICE_CUDA, + 0, + DL_DTYPE_UINT, + 16, )?; - let obs_dev_capsule = - dlpack_capsule(py, obs_dev, &obs_shape, DL_DEVICE_CUDA, 0, DL_DTYPE_UINT, 8)?; let value_host_capsule = dlpack_capsule( py, value_host.cast::(), @@ -402,10 +406,10 @@ impl CudaGraphRunner { pub fn dispatch_async( &self, batch_idx: usize, - obs_view: ArrayView, + obs_view: ArrayView, completion: BatchCompletion, ) { - debug_assert_eq!(obs_view.shape(), &[self.batch_size, OBS_SIDE, OBS_WIDTH]); + debug_assert_eq!(obs_view.shape(), &[self.batch_size, OBS_WORDS]); let lane = &self.lanes[batch_idx % self.lanes.len()]; @@ -413,7 +417,7 @@ impl CudaGraphRunner { .as_slice() .expect("observation batch must be contiguous"); let obs_dst = - unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_CELLS) }; + unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_WORDS) }; obs_dst.copy_from_slice(obs_src); unsafe { diff --git a/training/src/descent.rs b/training/src/descent.rs index a1a98e3..c798e5f 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -207,11 +207,11 @@ impl SearchNode { // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics // rust futures are normally lazily evaluated when you .await them for the // first time. Ours are eager, which is why this code actually works as we expect. - let value = evaluator.evaluate(&local_board); + let value = evaluator.evaluate(local_board.clone()); (action, local_board, child_outcome, rollback, value) }); - for (&action, mut local_board, child_outcome, _, value) in batch { + for (&action, local_board, child_outcome, _, value) in batch { match child_outcome { ApplyActionOutcome::Ongoing => { // Evaluate this child with neural net diff --git a/training/src/eval.rs b/training/src/eval.rs index 78ede1c..894e194 100644 --- a/training/src/eval.rs +++ b/training/src/eval.rs @@ -3,8 +3,9 @@ use std::future::Future; use alpha_paint::board::Board; -use ndarray::Ix2; +use ndarray::Ix1; +use crate::observation; use crate::queue::GpuJobQueue; /// Async evaluator trait for neural network inference. @@ -13,7 +14,7 @@ use crate::queue::GpuJobQueue; /// Value is from the current player's perspective. pub trait Evaluator { /// Evaluate the board and return a value estimate. - fn evaluate(&self, board: &Board) -> impl Future; + fn evaluate(&self, board: Board) -> impl Future; } /// GPU-backed evaluator that batches inference requests. @@ -21,23 +22,23 @@ pub trait Evaluator { /// Wraps a GpuJobQueue and serializes Board state into observations /// for GPU inference. Returns scalar value. pub struct GpuEvaluator<'a> { - queue: &'a GpuJobQueue, + queue: &'a GpuJobQueue, } impl<'a> GpuEvaluator<'a> { - pub fn new(queue: &'a GpuJobQueue) -> Self { + pub fn new(queue: &'a GpuJobQueue) -> Self { Self { queue } } } impl Evaluator for GpuEvaluator<'_> { - fn evaluate(&self, board: &Board) -> impl Future { - // Submit immediately with callback that writes observation. - // The observation encoding is a placeholder - must be implemented - // with the actual convnet encoding. - let _board = board.clone(); - let future = self.queue.eval(|_out| { - todo!("encode Board into observation tensor for convnet") + fn evaluate(&self, board: Board) -> impl Future { + let future = self.queue.eval(|mut out| { + observation::encode_into_slice( + &board, + out.as_slice_mut() + .expect("packed observation slot must be contiguous"), + ); }); async move { future.await } @@ -50,7 +51,7 @@ impl Evaluator for GpuEvaluator<'_> { pub struct UniformEvaluator; impl Evaluator for UniformEvaluator { - fn evaluate(&self, _: &Board) -> impl Future { + fn evaluate(&self, _: Board) -> impl Future { std::future::ready(0.0) } } diff --git a/training/src/lib.rs b/training/src/lib.rs index 69c91f1..d30e897 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -18,6 +18,7 @@ pub mod descent; pub mod eval; pub mod executor; pub mod future; +mod observation; pub mod queue; pub mod replay_buffer; pub mod training; @@ -91,19 +92,19 @@ fn graph_cache() -> &'static Mutex> { GRAPH_CACHE.get_or_init(|| Mutex::new(None)) } -/// Replay buffer storing (observation, value) pairs. +/// Replay buffer storing (packed observation, value) pairs. /// -/// Observations are u8 tensors, values are f32 scalars (no policy). +/// Observations are flat u16 tensors, values are f32 scalars (no policy). #[pyclass] struct EphemeralReplayBuffer { - inner: Arc>, + inner: Arc>, } #[pymethods] impl EphemeralReplayBuffer { #[new] fn new(capacity: usize) -> Self { - let obs_shape = Ix2(cudagraph::OBS_SIDE, cudagraph::OBS_WIDTH); + let obs_shape = Ix1(cudagraph::OBS_WORDS); Self { inner: Arc::new(ReplayBuffer::new(capacity, obs_shape)), } @@ -124,18 +125,15 @@ impl EphemeralReplayBuffer { py: Python<'py>, n: usize, seed: u64, - ) -> PyResult<( - Bound<'py, PyArray>, - Bound<'py, PyArray>, - )> { + ) -> PyResult<(Bound<'py, PyArray>, Bound<'py, PyArray>)> { let mut rng = ChaCha8Rng::seed_from_u64(seed); let batch = self.inner.sample(n, &mut rng); let num_samples = batch.values.len(); let obs_data = batch.observations.into_raw_vec_and_offset().0; - let obs = PyArray::from_vec(py, obs_data) - .reshape(Ix3(num_samples, cudagraph::OBS_SIDE, cudagraph::OBS_WIDTH))?; + let obs = + PyArray::from_vec(py, obs_data).reshape(Ix2(num_samples, cudagraph::OBS_WORDS))?; let values = PyArray::from_vec(py, batch.values); Ok((obs, values)) @@ -143,7 +141,7 @@ impl EphemeralReplayBuffer { } impl EphemeralReplayBuffer { - pub fn inner(&self) -> &Arc> { + pub fn inner(&self) -> &Arc> { &self.inner } } @@ -189,9 +187,7 @@ impl SelfPlay { num_threads, workers_per_thread, seed, - worker: WorkerConfig { - descent_iterations, - }, + worker: WorkerConfig { descent_iterations }, }; let total_workers = num_threads.checked_mul(workers_per_thread).ok_or_else(|| { @@ -240,18 +236,13 @@ impl SelfPlay { } }; - let dispatch = - move |batch_idx: usize, - obs_view: ArrayView, - completion: queue::BatchCompletion| { - runner.dispatch_async(batch_idx, obs_view, completion); - }; + let dispatch = move |batch_idx: usize, + obs_view: ArrayView, + completion: queue::BatchCompletion| { + runner.dispatch_async(batch_idx, obs_view, completion); + }; - let session = SelfPlaySession::new( - config, - replay_buffer.inner().clone(), - dispatch, - ); + let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); Ok(Self { session: Some(session), diff --git a/training/src/observation.rs b/training/src/observation.rs new file mode 100644 index 0000000..fb25fd4 --- /dev/null +++ b/training/src/observation.rs @@ -0,0 +1,257 @@ +use alpha_paint::board::board_structs::Player; +use alpha_paint::board::structs::Coordinate; +use alpha_paint::board::Board; + +pub const BOARD_SIDE: usize = 32; +pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; +pub const INTRINSIC_COUNT: usize = 10; +pub const OBS_WORDS: usize = BOARD_CELLS + INTRINSIC_COUNT; +pub const BOARD_PLANES: usize = 17; + +pub const PAINT_STRENGTH_MASK: u16 = 0b111; +pub const PAINT_IS_ENEMY_BIT: u16 = 1 << 3; +pub const WALL_BIT: u16 = 1 << 4; +pub const POWERUP_BIT: u16 = 1 << 5; +pub const BEACON_SHIFT: u16 = 6; +pub const BEACON_MASK: u16 = 0b11 << BEACON_SHIFT; +pub const HILL_SHIFT: u16 = 8; +pub const HILL_MASK: u16 = 0b11 << HILL_SHIFT; +pub const CURRENT_PLAYER_BIT: u16 = 1 << 10; +pub const OPPONENT_PLAYER_BIT: u16 = 1 << 11; + +pub const BEACON_NONE: u16 = 0; +pub const BEACON_CURRENT: u16 = 1; +pub const BEACON_OPPONENT: u16 = 2; + +pub const HILL_NONE: u16 = 0; +pub const HILL_NEUTRAL: u16 = 1; +pub const HILL_CURRENT: u16 = 2; +pub const HILL_OPPONENT: u16 = 3; + +pub const INTRINSIC_CURRENT_STAMINA: usize = 0; +pub const INTRINSIC_OPPONENT_STAMINA: usize = 1; +pub const INTRINSIC_CURRENT_HILLS: usize = 2; +pub const INTRINSIC_OPPONENT_HILLS: usize = 3; +pub const INTRINSIC_TURN_COUNT: usize = 4; +pub const INTRINSIC_CURRENT_TERRITORY: usize = 5; +pub const INTRINSIC_OPPONENT_TERRITORY: usize = 6; +pub const INTRINSIC_CURRENT_BEACONS: usize = 7; +pub const INTRINSIC_OPPONENT_BEACONS: usize = 8; +pub const INTRINSIC_CONSECUTIVE_MOVES: usize = 9; + +#[inline] +pub fn encode_into_slice(board: &Board, out: &mut [u16]) { + assert_eq!( + out.len(), + OBS_WORDS, + "packed observation buffer size mismatch" + ); + + out.fill(0); + out[..BOARD_CELLS].fill(WALL_BIT); + + let current_is_white = board.is_white_turn(); + let current_coord = board.current_player_coord(); + let opponent_coord = if current_is_white { + board.black_coord + } else { + board.white_coord + }; + + for y in 0..board.rows { + for x in 0..board.cols { + let coord = Coordinate::new(x, y); + let tile = board.tiles[coord]; + + let mut word = 0u16; + let paint = tile.paint_value(); + if paint != 0 { + word |= (paint.unsigned_abs() as u16) & PAINT_STRENGTH_MASK; + let is_enemy_paint = if current_is_white { + paint < 0 + } else { + paint > 0 + }; + if is_enemy_paint { + word |= PAINT_IS_ENEMY_BIT; + } + } + + if tile.is_wall() { + word |= WALL_BIT; + } + if board.powerups[coord] { + word |= POWERUP_BIT; + } + + let beacon_state = match tile.beacon_owner() { + Some(Player::White) if current_is_white => BEACON_CURRENT, + Some(Player::White) => BEACON_OPPONENT, + Some(Player::Black) if current_is_white => BEACON_OPPONENT, + Some(Player::Black) => BEACON_CURRENT, + None => BEACON_NONE, + }; + word |= beacon_state << BEACON_SHIFT; + + let hill_state = match board.hill_id[coord] { + u16::MAX => HILL_NONE, + hill_id => match board.tiles.hill_metadata()[hill_id as usize].owner { + None => HILL_NEUTRAL, + Some(Player::White) if current_is_white => HILL_CURRENT, + Some(Player::White) => HILL_OPPONENT, + Some(Player::Black) if current_is_white => HILL_OPPONENT, + Some(Player::Black) => HILL_CURRENT, + }, + }; + word |= hill_state << HILL_SHIFT; + + if coord == current_coord { + word |= CURRENT_PLAYER_BIT; + } + if coord == opponent_coord { + word |= OPPONENT_PLAYER_BIT; + } + + out[cell_index(x, y)] = word; + } + } + + let tail = &mut out[BOARD_CELLS..]; + if current_is_white { + tail[INTRINSIC_CURRENT_STAMINA] = board.white_stamina.try_into().unwrap(); + tail[INTRINSIC_OPPONENT_STAMINA] = board.black_stamina.try_into().unwrap(); + tail[INTRINSIC_CURRENT_HILLS] = board + .tiles + .controlled_hill_count::() + .try_into() + .unwrap(); + tail[INTRINSIC_OPPONENT_HILLS] = board + .tiles + .controlled_hill_count::() + .try_into() + .unwrap(); + tail[INTRINSIC_CURRENT_TERRITORY] = + board.tiles.territory_count::().try_into().unwrap(); + tail[INTRINSIC_OPPONENT_TERRITORY] = + board.tiles.territory_count::().try_into().unwrap(); + tail[INTRINSIC_CURRENT_BEACONS] = board + .tiles + .get_beacon_iterator::() + .count() + .try_into() + .unwrap(); + tail[INTRINSIC_OPPONENT_BEACONS] = board + .tiles + .get_beacon_iterator::() + .count() + .try_into() + .unwrap(); + } else { + tail[INTRINSIC_CURRENT_STAMINA] = board.black_stamina.try_into().unwrap(); + tail[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina.try_into().unwrap(); + tail[INTRINSIC_CURRENT_HILLS] = board + .tiles + .controlled_hill_count::() + .try_into() + .unwrap(); + tail[INTRINSIC_OPPONENT_HILLS] = board + .tiles + .controlled_hill_count::() + .try_into() + .unwrap(); + tail[INTRINSIC_CURRENT_TERRITORY] = + board.tiles.territory_count::().try_into().unwrap(); + tail[INTRINSIC_OPPONENT_TERRITORY] = + board.tiles.territory_count::().try_into().unwrap(); + tail[INTRINSIC_CURRENT_BEACONS] = board + .tiles + .get_beacon_iterator::() + .count() + .try_into() + .unwrap(); + tail[INTRINSIC_OPPONENT_BEACONS] = board + .tiles + .get_beacon_iterator::() + .count() + .try_into() + .unwrap(); + } + tail[INTRINSIC_TURN_COUNT] = board.turn_count.try_into().unwrap(); + tail[INTRINSIC_CONSECUTIVE_MOVES] = board.consecutives_moves_so_far.try_into().unwrap(); +} + +#[inline] +pub const fn cell_index(x: u8, y: u8) -> usize { + y as usize * BOARD_SIDE + x as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_words(fen: &str) -> Vec { + let board = Board::from_fen(fen).unwrap(); + let mut out = vec![0; OBS_WORDS]; + encode_into_slice(&board, &mut out); + out + } + + #[test] + fn packs_board_bits_and_intrinsics_for_white_turn() { + let words = encode_words( + "ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:1B1/d1O/3", + ); + + assert_eq!(words[cell_index(1, 0)] & PAINT_STRENGTH_MASK, 2); + assert_eq!(words[cell_index(1, 0)] & PAINT_IS_ENEMY_BIT, 0); + + let enemy_paint = words[cell_index(0, 1)]; + assert_eq!(enemy_paint & PAINT_STRENGTH_MASK, 4); + assert_ne!(enemy_paint & PAINT_IS_ENEMY_BIT, 0); + assert_ne!(enemy_paint & POWERUP_BIT, 0); + + let hill_word = words[cell_index(1, 1)]; + assert_eq!((hill_word & HILL_MASK) >> HILL_SHIFT, HILL_NEUTRAL); + + let beacon_word = words[cell_index(2, 1)]; + assert_eq!((beacon_word & BEACON_MASK) >> BEACON_SHIFT, BEACON_CURRENT); + + assert_ne!(words[cell_index(0, 0)] & CURRENT_PLAYER_BIT, 0); + assert_ne!(words[cell_index(2, 2)] & OPPONENT_PLAYER_BIT, 0); + + assert_ne!(words[cell_index(31, 31)] & WALL_BIT, 0); + + let tail = &words[BOARD_CELLS..]; + assert_eq!(tail[INTRINSIC_CURRENT_STAMINA], 99); + assert_eq!(tail[INTRINSIC_OPPONENT_STAMINA], 88); + assert_eq!(tail[INTRINSIC_CURRENT_HILLS], 0); + assert_eq!(tail[INTRINSIC_OPPONENT_HILLS], 0); + assert_eq!(tail[INTRINSIC_TURN_COUNT], 0); + assert_eq!(tail[INTRINSIC_CURRENT_TERRITORY], 1); + assert_eq!(tail[INTRINSIC_OPPONENT_TERRITORY], 1); + assert_eq!(tail[INTRINSIC_CURRENT_BEACONS], 1); + assert_eq!(tail[INTRINSIC_OPPONENT_BEACONS], 0); + assert_eq!(tail[INTRINSIC_CONSECUTIVE_MOVES], 2); + } + + #[test] + fn keeps_absolute_locations_on_black_turn() { + let words = + encode_words("ap2|3x3|tc:1|cm:5|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/2o"); + + let white_paint = words[cell_index(1, 0)]; + assert_eq!(white_paint & PAINT_STRENGTH_MASK, 2); + assert_ne!(white_paint & PAINT_IS_ENEMY_BIT, 0); + + let hill_word = words[cell_index(1, 1)]; + assert_eq!((hill_word & HILL_MASK) >> HILL_SHIFT, HILL_CURRENT); + + assert_ne!(words[cell_index(2, 2)] & CURRENT_PLAYER_BIT, 0); + assert_ne!(words[cell_index(0, 0)] & OPPONENT_PLAYER_BIT, 0); + + let tail = &words[BOARD_CELLS..]; + assert_eq!(tail[INTRINSIC_CURRENT_STAMINA], 88); + assert_eq!(tail[INTRINSIC_OPPONENT_STAMINA], 99); + assert_eq!(tail[INTRINSIC_CONSECUTIVE_MOVES], 5); + } +} diff --git a/training/src/training.rs b/training/src/training.rs index 9677245..1656ef1 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -7,7 +7,7 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::thread; -use ndarray::{ArrayView, Ix2, Ix3}; +use ndarray::{ArrayView, Ix1, Ix2}; use rand::SeedableRng; use rand_chacha::ChaCha8Rng; @@ -118,24 +118,21 @@ pub struct SelfPlaySession { impl SelfPlaySession { /// Create a new persistent session. /// - /// Uses concrete AlphaPaint types: u8 observations with Ix2 shape, f32 values. + /// Uses concrete AlphaPaint types: packed u16 observations with Ix1 shape, f32 values. pub fn new( config: SessionConfig, - replay_buffer: Arc>, + replay_buffer: Arc>, dispatch: F, ) -> Self where - F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, + F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, { let total_workers = config .num_threads .checked_mul(config.workers_per_thread) .expect("num_threads * workers_per_thread overflowed usize"); - let obs_shape = Ix2( - crate::cudagraph::OBS_SIDE, - crate::cudagraph::OBS_WIDTH, - ); + let obs_shape = Ix1(crate::cudagraph::OBS_WORDS); let queue = Arc::new(GpuJobQueue::new(obs_shape, total_workers, dispatch)); let control = Arc::new(SessionControl::new()); @@ -255,10 +252,10 @@ impl Drop for SelfPlaySession { /// Main loop for a single thread in a persistent session. fn session_thread_main( thread_id: usize, - queue: Arc>, + queue: Arc>, config: SessionConfig, control: Arc, - replay_buffer: &ReplayBuffer, + replay_buffer: &ReplayBuffer, ) { let base_seed = config.seed.wrapping_add(thread_id as u64 * 1000); let evaluator = GpuEvaluator::new(&*queue); diff --git a/training/src/worker.rs b/training/src/worker.rs index befdb6b..479fb99 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -6,7 +6,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use ndarray::Ix2; +use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; @@ -106,7 +106,7 @@ pub async fn worker_loop_forever( rng: &mut R, samples_collected: Arc, games_completed: Arc, - replay_buffer: &ReplayBuffer, + replay_buffer: &ReplayBuffer, ) { loop { let values = play_game(evaluator, config, rng).await; diff --git a/uv.lock b/uv.lock index 912cfc1..2406642 100644 --- a/uv.lock +++ b/uv.lock @@ -12,6 +12,7 @@ dependencies = [ { name = "psutil" }, { name = "py-cpuinfo" }, { name = "torch" }, + { name = "triton" }, ] [package.dev-dependencies] @@ -28,6 +29,7 @@ requires-dist = [ { name = "psutil", specifier = "==5.9.0" }, { name = "py-cpuinfo" }, { name = "torch", specifier = "==2.10.0" }, + { name = "triton", specifier = ">=3.6.0" }, ] [package.metadata.requires-dev] @@ -484,6 +486,7 @@ name = "triton" version = "3.6.0" source = { registry = "https://pypi.org/simple" } wheels = [ + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, ] From 6d0076bbba9a3bdafb1a9f30fd45126b520fcddc Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 03:37:08 -0400 Subject: [PATCH 10/59] wire self-play samples into value training --- python/alphapaint_training/__init__.py | 7 + .../alphapaint_training/cudagraph_backend.py | 4 +- python/alphapaint_training/model.py | 19 ++- python/alphapaint_training/packed_obs.py | 1 + python/alphapaint_training/train.py | 155 ++++++++++++++++++ training/src/descent.rs | 31 ++-- training/src/lib.rs | 2 +- training/src/worker.rs | 50 +++--- 8 files changed, 232 insertions(+), 37 deletions(-) create mode 100644 python/alphapaint_training/train.py diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 604d707..a0d8c3c 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -1,5 +1,10 @@ """AlphaPaint training infrastructure - Python bindings.""" +from pkgutil import extend_path + +__path__ = extend_path(__path__, __name__) + +from .alphapaint_training import EphemeralReplayBuffer, SelfPlay from .model import PackedValueModel, ResidualBlock, TinyValueNet from .packed_obs import ( BOARD_CELLS, @@ -17,10 +22,12 @@ "BOARD_CELLS", "BOARD_PLANES", "BOARD_SIDE", + "EphemeralReplayBuffer", "INTRINSIC_COUNT", "OBS_WORDS", "PackedValueModel", "ResidualBlock", + "SelfPlay", "TinyValueNet", "decode_intrinsics", "decode_packed_board", diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index c391b20..35aa102 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -29,7 +29,7 @@ def capture_lane_graph( value_host_dlpack, value_device_dlpack, stream_handle: int, - precision: str = "fp32", + precision: str = "bf16", ) -> tuple[int, object]: """Capture a CUDA graph for value-only inference. @@ -43,7 +43,7 @@ def capture_lane_graph( value_host_dlpack: DLPack capsule for pinned host value buffer. value_device_dlpack: DLPack capsule for device value buffer. stream_handle: Raw CUDA stream handle. - precision: "fp32", "fp16", or "bf16". + precision: "bf16", "fp16", or "fp32". Returns: Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 7e53d27..ae20ca8 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -3,7 +3,13 @@ import torch from torch import nn -from .packed_obs import BOARD_PLANES, INTRINSIC_COUNT, decode_packed_observation +from .packed_obs import ( + BOARD_CELLS, + BOARD_PLANES, + INTRINSIC_COUNT, + INTRINSIC_SCALE, + decode_packed_board, +) class ResidualBlock(nn.Module): @@ -75,6 +81,11 @@ def __init__( ): super().__init__() self.board_dtype = board_dtype + self.register_buffer( + "intrinsic_scales", + torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), + persistent=False, + ) self.value_net = TinyValueNet( width=width, embedding_dim=embedding_dim, @@ -82,7 +93,11 @@ def __init__( ) def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - return decode_packed_observation(packed_obs, board_dtype=self.board_dtype) + board = decode_packed_board(packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype) + intrinsics = packed_obs[:, BOARD_CELLS:].to(torch.float32) + intrinsics = intrinsics / self.intrinsic_scales + intrinsics = intrinsics.to(dtype=self.board_dtype) + return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: board, intrinsics = self.decode(packed_obs) diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index 168cefd..8d9f112 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -265,6 +265,7 @@ def decode_packed_observation( "BOARD_CELLS", "BOARD_PLANES", "BOARD_SIDE", + "INTRINSIC_SCALE", "INTRINSIC_COUNT", "OBS_WORDS", "decode_intrinsics", diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py new file mode 100644 index 0000000..c8dd0e6 --- /dev/null +++ b/python/alphapaint_training/train.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import argparse +from contextlib import nullcontext +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F + +from . import EphemeralReplayBuffer, SelfPlay +from .model import PackedValueModel + + +@dataclass(slots=True) +class TrainConfig: + rounds: int = 1 + samples_per_round: int = 4096 + train_steps_per_round: int = 32 + batch_size: int = 512 + replay_capacity: int = 131072 + num_threads: int = 32 + workers_per_thread: int = 16 + descent_iterations: int = 30 + lr: float = 3e-4 + seed: int = 42 + selfplay_precision: str = "bf16" + device: str = "cuda" + + +def _device(device: str) -> torch.device: + if device == "cuda" and not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the training runner") + return torch.device(device) + + +def _sample_replay_batch( + replay_buffer: EphemeralReplayBuffer, + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + obs_np, values_np = replay_buffer.sample(batch_size, seed) + obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( + device=device, + dtype=torch.uint16, + non_blocking=True, + ) + values = torch.from_numpy(np.asarray(values_np, dtype=np.float32)).to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + return obs, values + + +def train_step( + model: PackedValueModel, + replay_buffer: EphemeralReplayBuffer, + optimizer: torch.optim.Optimizer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> float: + obs, target = _sample_replay_batch(replay_buffer, batch_size, seed, device) + if target.numel() == 0: + raise RuntimeError("replay buffer is empty") + + optimizer.zero_grad(set_to_none=True) + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if device.type == "cuda" + else nullcontext() + ) + with autocast: + pred = model(obs) + loss = F.mse_loss(pred.float(), target) + loss.backward() + optimizer.step() + return float(loss.detach().cpu()) + + +def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: + device = _device(config.device) + torch.manual_seed(config.seed) + + model = PackedValueModel().to(device) + optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) + replay_buffer = EphemeralReplayBuffer(config.replay_capacity) + selfplay = SelfPlay( + replay_buffer, + config.num_threads, + config.workers_per_thread, + config.seed, + descent_iterations=config.descent_iterations, + model=model, + selfplay_precision=config.selfplay_precision, + ) + + losses: list[float] = [] + target_samples = 0 + try: + for round_idx in range(config.rounds): + target_samples += config.samples_per_round + collected = selfplay.wait_for(target_samples) + + round_losses = [] + for step_idx in range(config.train_steps_per_round): + loss = train_step( + model, + replay_buffer, + optimizer, + batch_size=config.batch_size, + seed=config.seed + round_idx * 10_000 + step_idx, + device=device, + ) + round_losses.append(loss) + + losses.extend(round_losses) + mean_loss = sum(round_losses) / len(round_losses) + print( + f"round={round_idx} samples={collected} replay={len(replay_buffer)} " + f"mean_loss={mean_loss:.6f}" + ) + finally: + selfplay.drop() + + return model, losses + + +def _parse_args() -> TrainConfig: + parser = argparse.ArgumentParser(description="Run AlphaPaint value training") + parser.add_argument("--rounds", type=int, default=1) + parser.add_argument("--samples-per-round", type=int, default=4096) + parser.add_argument("--train-steps-per-round", type=int, default=32) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--replay-capacity", type=int, default=131072) + parser.add_argument("--num-threads", type=int, default=32) + parser.add_argument("--workers-per-thread", type=int, default=16) + parser.add_argument("--descent-iterations", type=int, default=30) + parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--selfplay-precision", default="bf16") + parser.add_argument("--device", default="cuda") + args = parser.parse_args() + return TrainConfig(**vars(args)) + + +def main() -> None: + run_training(_parse_args()) + + +if __name__ == "__main__": + main() diff --git a/training/src/descent.rs b/training/src/descent.rs index c798e5f..97ae938 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -11,6 +11,11 @@ use rand::{Rng, RngExt}; use crate::eval::Evaluator; +pub struct TreeLearningSample { + pub board: Board, + pub value: f32, +} + #[derive(Debug)] pub struct ChildData { pub action: Action, @@ -363,20 +368,25 @@ impl SearchNode { self.value } - /// Collect tree learning samples: (minimax_value) from all internal nodes. + /// Collect tree learning samples from all internal nodes. /// /// An internal node is one that has children and at least one expanded child. /// Non-terminal leaf nodes (where the network estimate was used without /// minimax backing) are excluded per Athénan's tree learning rules. - fn collect_values(&self, out: &mut Vec) { + fn collect_samples(&self, state: &mut Board, out: &mut Vec) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); if has_expanded_child { - out.push(self.value); + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + }); } for child in &self.children { if let Some(ref node) = child.node { - node.collect_values(out); + let (_, rollback) = state.apply_action(child.action); + node.collect_samples(state, out); + state.rollback(child.action, rollback); } } } @@ -504,12 +514,13 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { /// Collect tree learning samples from all internal nodes. /// - /// Returns minimax values from internal nodes (nodes with at least one - /// expanded child). These are the training targets for the neural network. - pub fn collect_tree_learning_values(&self) -> Vec { - let mut values = Vec::new(); - self.root_node.collect_values(&mut values); - values + /// Returns `(board, value)` pairs for nodes with at least one expanded + /// child. These are the tree-learning targets for the value network. + pub fn collect_tree_learning_samples(&self) -> Vec { + let mut state = self.root_state.clone(); + let mut samples = Vec::new(); + self.root_node.collect_samples(&mut state, &mut samples); + samples } /// Select an action using Athénan's ordinal distribution. diff --git a/training/src/lib.rs b/training/src/lib.rs index d30e897..5e4a59b 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -165,7 +165,7 @@ impl SelfPlay { *, descent_iterations = 30, model, - selfplay_precision = "fp32" + selfplay_precision = "bf16" ))] fn new( py: Python<'_>, diff --git a/training/src/worker.rs b/training/src/worker.rs index 479fb99..e593dac 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -12,10 +12,14 @@ use rand::{Rng, SeedableRng}; use alpha_paint::board::Board; -use crate::descent::GameSearchTree; +use crate::descent::{GameSearchTree, TreeLearningSample}; use crate::eval::Evaluator; +use crate::observation; use crate::replay_buffer::ReplayBuffer; +const FIXED_FEN: &str = + "ap2|5x5|tc:0|cm:0|ep:0|w:2,0,100|b:2,4,100|h:n@2,2|pu:-|ps:-|bd:5/5/5/5/5"; + /// Configuration for the worker. #[derive(Clone)] pub struct WorkerConfig { @@ -39,22 +43,21 @@ impl Default for WorkerConfig { /// 3. Select action via ordinal distribution /// 4. Apply action, reuse tree /// -/// Returns all collected training values from the search trees. +/// Returns all collected training samples from the search trees. async fn play_game( evaluator: &E, config: &WorkerConfig, rng: &mut R, -) -> Vec { - // TODO: Create board from map pool. For now this is a placeholder. - let board: Board = todo!("board creation from map pool / Python"); +) -> Vec { + let board = Board::from_fen(FIXED_FEN).expect("fixed FEN must parse"); - let mut all_values = Vec::new(); + let mut all_samples = Vec::new(); // Check if terminal before starting // (board.get_valid_actions().len() == 0 would indicate terminal) let actions = board.get_valid_actions(); if actions.len() == 0 { - return all_values; + return all_samples; } let tree_rng = SmallRng::from_rng(rng); @@ -65,8 +68,8 @@ async fn play_game( tree.run_descent_for_iter(config.descent_iterations).await; // Collect tree learning samples from internal nodes - let values = tree.collect_tree_learning_values(); - all_values.extend(values); + let samples = tree.collect_tree_learning_samples(); + all_samples.extend(samples); // Select action via ordinal distribution let action_id = tree.ordinal_select(); @@ -87,13 +90,17 @@ async fn play_game( break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { - // Step the tree to reuse it - tree.step_tree(&new_board, action_id); + if tree.root_node.children[action_id].node.is_some() { + tree.step_tree(&new_board, action_id); + } else { + let next_rng = SmallRng::from_rng(rng); + tree = GameSearchTree::new(&new_board, evaluator, next_rng).await; + } } } } - all_values + all_samples } /// Run a worker loop that plays games forever. @@ -109,19 +116,18 @@ pub async fn worker_loop_forever( replay_buffer: &ReplayBuffer, ) { loop { - let values = play_game(evaluator, config, rng).await; - let num_samples = values.len(); + let samples = play_game(evaluator, config, rng).await; + let num_samples = samples.len(); if num_samples > 0 { - // Push values to replay buffer. - // Observation encoding is a TODO - for now we write zeros. let mut guard = replay_buffer.reserve(num_samples); - for value in values { - guard.push_with_observation(value, |_out| { - // TODO: encode the board state into the observation tensor. - // This requires reconstructing board states by replaying - // actions from the root, which needs the board + action history - // to be tracked during play_game. + for sample in samples { + guard.push_with_observation(sample.value, |mut out| { + observation::encode_into_slice( + &sample.board, + out.as_slice_mut() + .expect("replay observation slot must be contiguous"), + ); }); } } From de01f36577a591ce4775ec0738763afc58e16956 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 04:29:07 -0400 Subject: [PATCH 11/59] speed up self-play value inference --- alpha_paint/src/lib.rs | 21 ++- .../alphapaint_training/cudagraph_backend.py | 1 + python/alphapaint_training/model.py | 111 +++++++++--- python/alphapaint_training/packed_obs.py | 24 ++- python/alphapaint_training/train.py | 168 +++++++++++++++--- training/src/cudagraph.rs | 21 ++- training/src/descent.rs | 55 +++--- training/src/lib.rs | 26 ++- training/src/worker.rs | 10 +- 9 files changed, 342 insertions(+), 95 deletions(-) diff --git a/alpha_paint/src/lib.rs b/alpha_paint/src/lib.rs index c04d31b..e7c2fec 100644 --- a/alpha_paint/src/lib.rs +++ b/alpha_paint/src/lib.rs @@ -1,16 +1,15 @@ -use std::time::Instant; -use pyo3::prelude::*; use crate::board::{ApplyActionOutcome, Board}; +use pyo3::prelude::*; +use std::time::Instant; mod bindings; pub mod board; mod evaluation; mod search; - fn perft(board: &mut Board, depth: usize) -> usize { if depth == 1 { - return board.get_valid_actions().len() + return board.get_valid_actions().len(); } let mut sum = 0; @@ -48,17 +47,24 @@ const PERFT_STR: [&str; 8] = [ "ap2|31x31|tc:0|cm:0|ep:0|w:15,0,99|b:15,30,100|h:n@14,14+16,14+15,15+14,16+16,16;n@0,13+0,14+1,14+2,14+2,15+0,16+1,16+2,16+0,17;n@30,13+28,14+29,14+30,14+28,15+28,16+29,16+30,16+30,17;n@13,6+14,6+15,6+16,6+17,6;n@13,24+14,24+15,24+16,24+17,24;n@7,1+7,8+7,22+7,29;n@23,1+23,8+23,22+23,29|pu:-|ps:2,21,11+2,21,19+2,19,6+2,19,24+2,29,8+2,29,22+2,4,10+2,4,20+2,13,4+2,13,26+2,23,10+2,23,20+2,20,0+2,20,30+2,18,3+2,18,27+2,19,9+2,19,21+2,22,5+2,22,25+52,16,14+52,16,16+52,26,10+52,26,20+52,29,10+52,29,20+52,26,12+52,26,18+52,12,6+52,12,24+102,4,6+102,4,24+102,29,14+102,29,16+102,2,14+102,2,16+102,23,12+102,23,18+102,21,4+102,21,26+152,16,10+152,16,20+152,27,8+152,27,22+152,5,10+152,5,20+152,9,8+152,9,22+152,4,12+152,4,18+202,12,11+202,12,19+202,30,14+202,30,16+202,7,6+202,7,24+202,23,2+202,23,28+202,4,11+202,4,19+252,8,3+252,8,27+252,25,0+252,25,30+252,5,2+252,5,28+252,26,4+252,26,26+252,26,9+252,26,21+302,1,12+302,1,18+302,10,1+302,10,29+302,30,13+302,30,17+302,28,7+302,28,23+302,1,0+302,1,30+352,9,11+352,9,19+352,20,0+352,20,30+352,23,13+352,23,17+352,9,6+352,9,24+352,23,8+352,23,22+402,22,11+402,22,19+402,14,13+402,14,17+402,30,10+402,30,20+402,28,11+402,28,19+402,20,7+402,20,23+452,5,12+452,5,18+452,1,8+452,1,22+452,13,2+452,13,28+452,1,7+452,1,23+452,23,10+452,23,20+502,21,11+502,21,19+502,9,0+502,9,30+502,30,1+502,30,29+502,26,1+502,26,29+502,19,5+502,19,25+552,15,2+552,15,28+552,28,3+552,28,27+552,15,8+552,15,22+552,14,7+552,14,23+552,30,1+552,30,29+602,7,6+602,7,24+602,20,5+602,20,25+602,11,5+602,11,25+602,15,6+602,15,24+602,2,0+602,2,30+652,1,12+652,1,18+652,17,1+652,17,29+652,15,8+652,15,22+652,15,6+652,15,24+652,27,13+652,27,17+702,6,11+702,6,19+702,24,8+702,24,22+702,14,13+702,14,17+702,11,5+702,11,25+702,30,2+702,30,28+752,30,3+752,30,27+752,3,15+752,28,1+752,28,29+752,28,3+752,28,27+752,18,12+752,18,18+802,8,3+802,8,27+802,15,2+802,15,28+802,12,12+802,12,18+802,17,1+802,17,29+802,19,3+802,19,27+852,19,11+852,19,19+852,22,0+852,22,30+852,10,11+852,10,19+852,9,5+852,9,25+852,22,13+852,22,17+902,24,14+902,24,16+902,1,13+902,1,17+902,3,5+902,3,25+902,28,7+902,28,23+902,19,12+902,19,18+952,26,15+952,21,10+952,21,20+952,21,8+952,21,22+952,11,3+952,11,27+952,13,11+952,13,19+1002,24,10+1002,24,20+1002,7,13+1002,7,17+1002,8,14+1002,8,16+1002,22,10+1002,22,20+1002,11,11+1002,11,19+1052,19,14+1052,19,16+1052,11,3+1052,11,27+1052,15,8+1052,15,22+1052,5,12+1052,5,18+1052,17,8+1052,17,22+1102,22,9+1102,22,21+1102,27,12+1102,27,18+1102,14,15+1102,26,1+1102,26,29+1102,23,4+1102,23,26+1152,3,8+1152,3,22+1152,30,3+1152,30,27+1152,17,5+1152,17,25+1152,4,0+1152,4,30+1152,6,11+1152,6,19+1202,17,6+1202,17,24+1202,19,9+1202,19,21+1202,1,0+1202,1,30+1202,14,6+1202,14,24+1202,27,13+1202,27,17+1252,19,5+1252,19,25+1252,18,0+1252,18,30+1252,24,10+1252,24,20+1252,4,2+1252,4,28+1252,21,10+1252,21,20+1302,21,15+1302,8,11+1302,8,19+1302,24,8+1302,24,22+1302,9,9+1302,9,21+1302,8,9+1302,8,21+1352,16,7+1352,16,23+1352,10,11+1352,10,19+1352,2,11+1352,2,19+1352,27,1+1352,27,29+1352,2,3+1352,2,27+1402,3,5+1402,3,25+1402,13,7+1402,13,23+1402,12,13+1402,12,17+1402,10,8+1402,10,22+1402,27,0+1402,27,30+1452,17,6+1452,17,24+1452,5,12+1452,5,18+1452,1,10+1452,1,20+1452,4,6+1452,4,24+1452,8,7+1452,8,23+1502,20,0+1502,20,30+1502,13,1+1502,13,29+1502,20,9+1502,20,21+1502,26,10+1502,26,20+1502,20,15+1552,27,2+1552,27,28+1552,17,15+1552,11,15+1552,21,6+1552,21,24+1552,4,14+1552,4,16+1552,20,1+1552,20,29+1602,23,6+1602,23,24+1602,19,15+1602,17,1+1602,17,29+1602,7,0+1602,7,30+1602,20,1+1602,20,29+1652,19,9+1652,19,21+1652,22,3+1652,22,27+1652,21,9+1652,21,21+1652,3,7+1652,3,23+1652,0,4+1652,0,26+1702,10,9+1702,10,21+1702,11,15+1702,15,6+1702,15,24+1702,30,14+1702,30,16+1702,27,4+1702,27,26+1752,23,8+1752,23,22+1752,21,7+1752,21,23+1752,20,12+1752,20,18+1752,4,2+1752,4,28+1752,10,11+1752,10,19+1802,21,8+1802,21,22+1802,19,7+1802,19,23+1802,11,4+1802,11,26+1802,19,9+1802,19,21+1802,1,14+1802,1,16+1852,20,8+1852,20,22+1852,1,3+1852,1,27+1852,19,14+1852,19,16+1852,12,3+1852,12,27+1852,17,10+1852,17,20+1902,6,3+1902,6,27+1902,30,11+1902,30,19+1902,15,1+1902,15,29+1902,10,8+1902,10,22+1902,30,3+1902,30,27+1952,5,7+1952,5,23+1952,9,5+1952,9,25+1952,30,9+1952,30,21+1952,27,1+1952,27,29+1952,28,4+1952,28,26|bd:#10#7#10#/1#3#5#7#5#3#1/2#9#5#9#2/3#1#7#3#7#1#3/10#1#5#1#10/5#19#5/3#1#4#9#4#1#3/12#5#12/5#6#5#6#5/5#19#5/3#6#9#6#3/5#1#15#1#5/13#####13/5#3#3#3#3#3#5/3#1#19#1#3/##27##/3#1#19#1#3/5#3#3#3#3#3#5/13#####13/5#1#15#1#5/3#6#9#6#3/5#19#5/5#6#5#6#5/12#5#12/3#1#4#9#4#1#3/5#19#5/10#1#5#1#10/3#1#7#3#7#1#3/2#9#5#9#2/1#3#5#7#5#3#1/#10#7#10#", ]; +pub const TRAINING_START_FENS: &[&str; 8] = &PERFT_STR; + pub fn perft_test() { - use std::time::Instant; use crate::board::Board; - use crate::{perft, PERFT_STR}; + use crate::{PERFT_STR, perft}; + use std::time::Instant; for &perft_str in PERFT_STR.iter() { let mut board = Board::from_fen(perft_str).unwrap(); let start = Instant::now(); let res = perft(&mut board, 8); let elapsed = start.elapsed(); - println!("{} - {} ({} nodes/sec)", res, elapsed.as_secs_f32(), res as f32 / elapsed.as_secs_f32()); + println!( + "{} - {} ({} nodes/sec)", + res, + elapsed.as_secs_f32(), + res as f32 / elapsed.as_secs_f32() + ); } } @@ -67,4 +73,3 @@ mod alpha_paint { #[pymodule_export] use crate::bindings::PyBoard; } - diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index 35aa102..ecf816d 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -55,6 +55,7 @@ def capture_lane_graph( model = model.cuda() model.eval() + torch.backends.cudnn.benchmark = True stream = torch.cuda.ExternalStream(stream_handle) graph = torch.cuda.CUDAGraph(keep_graph=True) dtype = _autocast_dtype(precision) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index ae20ca8..b445b13 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -15,12 +15,11 @@ class ResidualBlock(nn.Module): def __init__(self, width: int): super().__init__() - groups = min(8, width) - self.norm1 = nn.GroupNorm(groups, width) + self.norm1 = nn.BatchNorm2d(width) self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.norm2 = nn.GroupNorm(groups, width) + self.norm2 = nn.BatchNorm2d(width) self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.act = nn.GELU() + self.act = nn.ReLU(inplace=True) def forward(self, x: torch.Tensor) -> torch.Tensor: residual = x @@ -37,27 +36,22 @@ class TinyValueNet(nn.Module): def __init__( self, *, - width: int = 32, - embedding_dim: int = 64, - hidden_dim: int = 64, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, ): super().__init__() - groups = min(8, width) self.stem = nn.Sequential( nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), - nn.GroupNorm(groups, width), - nn.GELU(), + nn.BatchNorm2d(width), + nn.ReLU(inplace=True), ) - self.blocks = nn.Sequential(ResidualBlock(width), ResidualBlock(width)) + self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) self.pool = nn.AdaptiveAvgPool2d(1) - self.project = nn.Sequential( - nn.Flatten(), - nn.Linear(width, embedding_dim), - nn.GELU(), - ) + self.flatten = nn.Flatten() self.head = nn.Sequential( - nn.Linear(embedding_dim + INTRINSIC_COUNT, hidden_dim), - nn.GELU(), + nn.Linear(width + INTRINSIC_COUNT, hidden_dim), + nn.ReLU(inplace=True), nn.Linear(hidden_dim, 1), ) @@ -65,7 +59,7 @@ def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor x = self.stem(board) x = self.blocks(x) x = self.pool(x) - x = self.project(x) + x = self.flatten(x) x = torch.cat((x, intrinsics), dim=1) return self.head(x) @@ -74,13 +68,22 @@ class PackedValueModel(nn.Module): def __init__( self, *, - width: int = 32, - embedding_dim: int = 64, - hidden_dim: int = 64, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, board_dtype: torch.dtype = torch.bfloat16, ): super().__init__() self.board_dtype = board_dtype + self._board_buffers: dict[ + tuple[str, int | None, int, torch.dtype], torch.Tensor + ] = {} + self._intrinsic_fp32_buffers: dict[ + tuple[str, int | None, int], torch.Tensor + ] = {} + self._intrinsic_buffers: dict[ + tuple[str, int | None, int, torch.dtype], torch.Tensor + ] = {} self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), @@ -88,15 +91,69 @@ def __init__( ) self.value_net = TinyValueNet( width=width, - embedding_dim=embedding_dim, + num_blocks=num_blocks, hidden_dim=hidden_dim, ) + def _buffer_key( + self, packed_obs: torch.Tensor, dtype: torch.dtype + ) -> tuple[str, int | None, int, torch.dtype]: + return ( + packed_obs.device.type, + packed_obs.device.index, + packed_obs.shape[0], + dtype, + ) + + def _ensure_decode_buffers( + self, packed_obs: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + board_key = self._buffer_key(packed_obs, self.board_dtype) + board = self._board_buffers.get(board_key) + if board is None: + board = torch.empty( + (packed_obs.shape[0], BOARD_PLANES, 32, 32), + device=packed_obs.device, + dtype=self.board_dtype, + ) + self._board_buffers[board_key] = board + + fp32_key = ( + packed_obs.device.type, + packed_obs.device.index, + packed_obs.shape[0], + ) + intrinsics_fp32 = self._intrinsic_fp32_buffers.get(fp32_key) + if intrinsics_fp32 is None: + intrinsics_fp32 = torch.empty( + (packed_obs.shape[0], INTRINSIC_COUNT), + device=packed_obs.device, + dtype=torch.float32, + ) + self._intrinsic_fp32_buffers[fp32_key] = intrinsics_fp32 + + intrinsic_key = self._buffer_key(packed_obs, self.board_dtype) + intrinsics = self._intrinsic_buffers.get(intrinsic_key) + if intrinsics is None: + intrinsics = torch.empty( + (packed_obs.shape[0], INTRINSIC_COUNT), + device=packed_obs.device, + dtype=self.board_dtype, + ) + self._intrinsic_buffers[intrinsic_key] = intrinsics + + return board, intrinsics_fp32, intrinsics + def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - board = decode_packed_board(packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype) - intrinsics = packed_obs[:, BOARD_CELLS:].to(torch.float32) - intrinsics = intrinsics / self.intrinsic_scales - intrinsics = intrinsics.to(dtype=self.board_dtype) + board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) + board = decode_packed_board( + packed_obs[:, :BOARD_CELLS], + dtype=self.board_dtype, + out=board, + ) + intrinsics_fp32.copy_(packed_obs[:, BOARD_CELLS:]) + intrinsics_fp32.div_(self.intrinsic_scales) + intrinsics.copy_(intrinsics_fp32) return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index 8d9f112..a7d9003 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -198,6 +198,7 @@ def decode_packed_board( board_words: torch.Tensor, *, dtype: torch.dtype = torch.bfloat16, + out: torch.Tensor | None = None, ) -> torch.Tensor: if board_words.dtype != torch.uint16: raise ValueError(f"board_words must be uint16, got {board_words.dtype}") @@ -207,14 +208,25 @@ def decode_packed_board( ) if board_words.device.type != "cuda": - return decode_packed_board_reference(board_words).to(dtype=dtype) + decoded = decode_packed_board_reference(board_words).to(dtype=dtype) + if out is not None: + out.copy_(decoded) + return out + return decoded board_words = board_words.contiguous() - out = torch.empty( - (board_words.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), - device=board_words.device, - dtype=dtype, - ) + expected_shape = (board_words.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE) + if out is None: + out = torch.empty(expected_shape, device=board_words.device, dtype=dtype) + else: + if out.shape != expected_shape: + raise ValueError( + f"out must have shape {expected_shape}, got {tuple(out.shape)}" + ) + if out.device != board_words.device: + raise ValueError("out must be on the same device as board_words") + if out.dtype != dtype: + raise ValueError(f"out must have dtype {dtype}, got {out.dtype}") total_cells = board_words.shape[0] * BOARD_CELLS grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) _decode_board_kernel[grid]( diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index c8dd0e6..ed8595f 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -1,8 +1,11 @@ from __future__ import annotations import argparse +import json +import time from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import asdict, dataclass +from pathlib import Path import numpy as np import torch @@ -15,17 +18,19 @@ @dataclass(slots=True) class TrainConfig: rounds: int = 1 - samples_per_round: int = 4096 - train_steps_per_round: int = 32 + samples_per_round: int = 8192 + train_steps_per_round: int = 64 batch_size: int = 512 - replay_capacity: int = 131072 - num_threads: int = 32 + replay_capacity: int = 262144 + num_threads: int = 16 workers_per_thread: int = 16 - descent_iterations: int = 30 + descent_iterations: int = 50 lr: float = 3e-4 seed: int = 42 selfplay_precision: str = "bf16" device: str = "cuda" + checkpoint_interval: int = 3 + run_dir: str = "runs/latest" def _device(device: str) -> torch.device: @@ -40,7 +45,11 @@ def _sample_replay_batch( seed: int, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: - obs_np, values_np = replay_buffer.sample(batch_size, seed) + actual_batch_size = min(batch_size, len(replay_buffer)) + if actual_batch_size == 0: + raise RuntimeError("replay buffer is empty") + + obs_np, values_np = replay_buffer.sample(actual_batch_size, seed) obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( device=device, dtype=torch.uint16, @@ -63,9 +72,8 @@ def train_step( seed: int, device: torch.device, ) -> float: + model.train() obs, target = _sample_replay_batch(replay_buffer, batch_size, seed, device) - if target.numel() == 0: - raise RuntimeError("replay buffer is empty") optimizer.zero_grad(set_to_none=True) autocast = ( @@ -81,11 +89,62 @@ def train_step( return float(loss.detach().cpu()) +def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: + run_dir = Path(config.run_dir) + checkpoint_dir = run_dir / "checkpoints" + run_dir.mkdir(parents=True, exist_ok=True) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "config.json").write_text(json.dumps(asdict(config), indent=2) + "\n") + return run_dir, checkpoint_dir + + +def _append_jsonl(path: Path, record: dict[str, object]) -> None: + with path.open("a", encoding="ascii") as f: + json.dump(record, f, sort_keys=True) + f.write("\n") + + +def _save_checkpoint( + *, + checkpoint_dir: Path, + model: PackedValueModel, + optimizer: torch.optim.Optimizer, + config: TrainConfig, + round_idx: int, + samples: int, + games: int, + replay_size: int, + record: dict[str, object], +) -> Path: + state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "config": asdict(config), + "round": round_idx, + "samples": samples, + "games": games, + "replay_size": replay_size, + "metrics": record, + } + checkpoint_path = checkpoint_dir / f"round_{round_idx:05d}.pt" + torch.save(state, checkpoint_path) + torch.save(state, checkpoint_dir / "latest.pt") + return checkpoint_path + + def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) + if device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + + run_dir, checkpoint_dir = _prepare_run_dir(config) + metrics_path = run_dir / "metrics.jsonl" model = PackedValueModel().to(device) + model.eval() optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -100,12 +159,27 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: losses: list[float] = [] target_samples = 0 + previous_samples = 0 + previous_games = 0 + previous_gpu_batches = 0 + previous_gpu_evals = 0 try: for round_idx in range(config.rounds): + round_number = round_idx + 1 + collect_started_at = time.perf_counter() target_samples += config.samples_per_round collected = selfplay.wait_for(target_samples) + games = selfplay.games() + gpu_batches = selfplay.gpu_batches() + gpu_evals = selfplay.gpu_evals() + collect_seconds = time.perf_counter() - collect_started_at + samples_added = collected - previous_samples + games_added = games - previous_games + gpu_batches_added = gpu_batches - previous_gpu_batches + gpu_evals_added = gpu_evals - previous_gpu_evals round_losses = [] + train_started_at = time.perf_counter() for step_idx in range(config.train_steps_per_round): loss = train_step( model, @@ -116,13 +190,69 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device=device, ) round_losses.append(loss) + model.eval() + train_seconds = time.perf_counter() - train_started_at losses.extend(round_losses) - mean_loss = sum(round_losses) / len(round_losses) + replay_size = len(replay_buffer) + mean_loss = sum(round_losses) / len(round_losses) if round_losses else 0.0 + min_loss = min(round_losses) if round_losses else 0.0 + max_loss = max(round_losses) if round_losses else 0.0 + last_loss = round_losses[-1] if round_losses else 0.0 + record = { + "round": round_number, + "samples_total": collected, + "samples_added": samples_added, + "games_total": games, + "games_added": games_added, + "gpu_batches_total": gpu_batches, + "gpu_batches_added": gpu_batches_added, + "gpu_batches_per_second": gpu_batches_added + / max(collect_seconds, 1e-9), + "gpu_evals_total": gpu_evals, + "gpu_evals_added": gpu_evals_added, + "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), + "replay_size": replay_size, + "collection_seconds": collect_seconds, + "training_seconds": train_seconds, + "samples_per_second": samples_added / max(collect_seconds, 1e-9), + "train_steps_per_second": config.train_steps_per_round + / max(train_seconds, 1e-9), + "loss_mean": mean_loss, + "loss_min": min_loss, + "loss_max": max_loss, + "loss_last": last_loss, + "timestamp": time.time(), + } + _append_jsonl(metrics_path, record) print( - f"round={round_idx} samples={collected} replay={len(replay_buffer)} " - f"mean_loss={mean_loss:.6f}" + f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " + f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " + f"batches={gpu_batches_added} replay={replay_size} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " + f"loss={mean_loss:.6f}/{last_loss:.6f}" ) + + if config.checkpoint_interval > 0 and ( + round_number % config.checkpoint_interval == 0 + or round_number == config.rounds + ): + checkpoint_path = _save_checkpoint( + checkpoint_dir=checkpoint_dir, + model=model, + optimizer=optimizer, + config=config, + round_idx=round_number, + samples=collected, + games=games, + replay_size=replay_size, + record=record, + ) + print(f"checkpoint={checkpoint_path}") + + previous_samples = collected + previous_games = games + previous_gpu_batches = gpu_batches + previous_gpu_evals = gpu_evals finally: selfplay.drop() @@ -132,17 +262,15 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: def _parse_args() -> TrainConfig: parser = argparse.ArgumentParser(description="Run AlphaPaint value training") parser.add_argument("--rounds", type=int, default=1) - parser.add_argument("--samples-per-round", type=int, default=4096) - parser.add_argument("--train-steps-per-round", type=int, default=32) + parser.add_argument("--samples-per-round", type=int, default=8192) + parser.add_argument("--train-steps-per-round", type=int, default=64) parser.add_argument("--batch-size", type=int, default=512) - parser.add_argument("--replay-capacity", type=int, default=131072) - parser.add_argument("--num-threads", type=int, default=32) + parser.add_argument("--replay-capacity", type=int, default=262144) + parser.add_argument("--num-threads", type=int, default=16) parser.add_argument("--workers-per-thread", type=int, default=16) - parser.add_argument("--descent-iterations", type=int, default=30) - parser.add_argument("--lr", type=float, default=3e-4) + parser.add_argument("--descent-iterations", type=int, default=50) parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--selfplay-precision", default="bf16") - parser.add_argument("--device", default="cuda") + parser.add_argument("--run-dir", default="runs/latest") args = parser.parse_args() return TrainConfig(**vars(args)) diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index 90e23be..730f6fe 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -5,6 +5,7 @@ use std::ffi::{c_void, CStr}; use std::mem::size_of; use std::slice; +use std::sync::atomic::{AtomicU64, Ordering}; use cudarc::runtime::sys as cuda; use ndarray::{ArrayView, Ix2}; @@ -283,6 +284,8 @@ impl Drop for CudaGraphLane { pub struct CudaGraphRunner { batch_size: usize, lanes: Vec, + dispatched_batches: AtomicU64, + dispatched_evals: AtomicU64, } // SAFETY: Lane buffers/streams are independent per batch_idx and queue dispatch @@ -400,7 +403,20 @@ impl CudaGraphRunner { lanes.push(lane); } - Ok(Self { batch_size, lanes }) + Ok(Self { + batch_size, + lanes, + dispatched_batches: AtomicU64::new(0), + dispatched_evals: AtomicU64::new(0), + }) + } + + pub fn dispatched_batches(&self) -> u64 { + self.dispatched_batches.load(Ordering::Relaxed) + } + + pub fn dispatched_evals(&self) -> u64 { + self.dispatched_evals.load(Ordering::Relaxed) } pub fn dispatch_async( @@ -419,6 +435,9 @@ impl CudaGraphRunner { let obs_dst = unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_WORDS) }; obs_dst.copy_from_slice(obs_src); + self.dispatched_batches.fetch_add(1, Ordering::Relaxed); + self.dispatched_evals + .fetch_add(self.batch_size as u64, Ordering::Relaxed); unsafe { check_cuda_or_panic( diff --git a/training/src/descent.rs b/training/src/descent.rs index 97ae938..be90e0e 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -291,30 +291,6 @@ impl SearchNode { sign * (2000.0 - board.turn_count as f32) } - async fn create_child( - &mut self, - mut state: Board, - action: Action, - evaluator: &E, - rng: &mut SmallRng, - ) -> f32 { - let (outcome, _) = state.apply_action(action); - - let node = - Box::new(Box::pin(SearchNode::build_self(&state, outcome, evaluator, rng)).await); - let value = node.value; - - if let Some(id) = self - .children - .iter() - .position(|child| child.action == action) - { - self.children[id].node = Some(node); - } - - value - } - async fn ubfms_iteration( &mut self, mut state: Board, @@ -343,10 +319,35 @@ impl SearchNode { if self.children[best_action_id].node.is_some() { let (outcome, _) = state.apply_action(best_action); let child = self.children[best_action_id].node.as_mut().unwrap(); - Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng)).await; - } else { self.children[best_action_id].child_value = - Box::pin(self.create_child(state, best_action, evaluator, rng)).await; + Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng)).await; + } else { + let (child_outcome, _) = state.apply_action(best_action); + let should_descend = matches!(child_outcome, ApplyActionOutcome::Ongoing); + let child_node = Box::new( + Box::pin(SearchNode::build_self( + &state, + child_outcome, + evaluator, + rng, + )) + .await, + ); + + self.children[best_action_id].child_value = child_node.value; + self.children[best_action_id].node = Some(child_node); + + if should_descend { + let child = self.children[best_action_id].node.as_mut().unwrap(); + self.children[best_action_id].child_value = + Box::pin(child.ubfms_iteration( + state, + ApplyActionOutcome::Ongoing, + evaluator, + rng, + )) + .await; + } } let (best_action_id, _) = self.completed_best_action(white_turn, rng); diff --git a/training/src/lib.rs b/training/src/lib.rs index 5e4a59b..8da8f9b 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -152,6 +152,7 @@ impl EphemeralReplayBuffer { #[pyclass] struct SelfPlay { session: Option, + runner: Arc, } #[pymethods] @@ -236,16 +237,18 @@ impl SelfPlay { } }; + let runner_for_dispatch = runner.clone(); let dispatch = move |batch_idx: usize, obs_view: ArrayView, completion: queue::BatchCompletion| { - runner.dispatch_async(batch_idx, obs_view, completion); + runner_for_dispatch.dispatch_async(batch_idx, obs_view, completion); }; let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); Ok(Self { session: Some(session), + runner, }) } @@ -280,6 +283,27 @@ impl SelfPlay { .samples()) } + /// Return the current absolute game count. + fn games(&self) -> PyResult { + Ok(self + .session + .as_ref() + .ok_or_else(|| { + PyErr::new::("session already dropped") + })? + .games()) + } + + /// Return the total number of CUDA graph launches completed so far. + fn gpu_batches(&self) -> u64 { + self.runner.dispatched_batches() + } + + /// Return the total number of packed observations sent to GPU so far. + fn gpu_evals(&self) -> u64 { + self.runner.dispatched_evals() + } + /// Shut down the session. Idempotent. #[pyo3(name = "drop")] fn py_drop(&mut self) { diff --git a/training/src/worker.rs b/training/src/worker.rs index e593dac..ef609be 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -8,18 +8,16 @@ use std::sync::Arc; use ndarray::Ix1; use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; +use rand::{Rng, RngExt, SeedableRng}; use alpha_paint::board::Board; +use alpha_paint::TRAINING_START_FENS; use crate::descent::{GameSearchTree, TreeLearningSample}; use crate::eval::Evaluator; use crate::observation; use crate::replay_buffer::ReplayBuffer; -const FIXED_FEN: &str = - "ap2|5x5|tc:0|cm:0|ep:0|w:2,0,100|b:2,4,100|h:n@2,2|pu:-|ps:-|bd:5/5/5/5/5"; - /// Configuration for the worker. #[derive(Clone)] pub struct WorkerConfig { @@ -49,7 +47,9 @@ async fn play_game( config: &WorkerConfig, rng: &mut R, ) -> Vec { - let board = Board::from_fen(FIXED_FEN).expect("fixed FEN must parse"); + let board = + Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) + .expect("training start FEN must parse"); let mut all_samples = Vec::new(); From cefbc083a6b9e3bbe99c2095c38476efc0102a5b Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 04:30:28 -0400 Subject: [PATCH 12/59] stream tree samples during self-play --- training/src/worker.rs | 46 ++++++++++++++++++------------------------ 1 file changed, 20 insertions(+), 26 deletions(-) diff --git a/training/src/worker.rs b/training/src/worker.rs index ef609be..4a023ad 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -13,7 +13,7 @@ use rand::{Rng, RngExt, SeedableRng}; use alpha_paint::board::Board; use alpha_paint::TRAINING_START_FENS; -use crate::descent::{GameSearchTree, TreeLearningSample}; +use crate::descent::GameSearchTree; use crate::eval::Evaluator; use crate::observation; use crate::replay_buffer::ReplayBuffer; @@ -41,23 +41,22 @@ impl Default for WorkerConfig { /// 3. Select action via ordinal distribution /// 4. Apply action, reuse tree /// -/// Returns all collected training samples from the search trees. async fn play_game( evaluator: &E, config: &WorkerConfig, rng: &mut R, -) -> Vec { + replay_buffer: &ReplayBuffer, + samples_collected: &AtomicUsize, +) { let board = Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) .expect("training start FEN must parse"); - let mut all_samples = Vec::new(); - // Check if terminal before starting // (board.get_valid_actions().len() == 0 would indicate terminal) let actions = board.get_valid_actions(); if actions.len() == 0 { - return all_samples; + return; } let tree_rng = SmallRng::from_rng(rng); @@ -69,7 +68,20 @@ async fn play_game( // Collect tree learning samples from internal nodes let samples = tree.collect_tree_learning_samples(); - all_samples.extend(samples); + let num_samples = samples.len(); + if num_samples > 0 { + let mut guard = replay_buffer.reserve(num_samples); + for sample in samples { + guard.push_with_observation(sample.value, |mut out| { + observation::encode_into_slice( + &sample.board, + out.as_slice_mut() + .expect("replay observation slot must be contiguous"), + ); + }); + } + samples_collected.fetch_add(num_samples, Ordering::AcqRel); + } // Select action via ordinal distribution let action_id = tree.ordinal_select(); @@ -99,8 +111,6 @@ async fn play_game( } } } - - all_samples } /// Run a worker loop that plays games forever. @@ -116,23 +126,7 @@ pub async fn worker_loop_forever( replay_buffer: &ReplayBuffer, ) { loop { - let samples = play_game(evaluator, config, rng).await; - let num_samples = samples.len(); - - if num_samples > 0 { - let mut guard = replay_buffer.reserve(num_samples); - for sample in samples { - guard.push_with_observation(sample.value, |mut out| { - observation::encode_into_slice( - &sample.board, - out.as_slice_mut() - .expect("replay observation slot must be contiguous"), - ); - }); - } - } - - samples_collected.fetch_add(num_samples, Ordering::AcqRel); + play_game(evaluator, config, rng, replay_buffer, &samples_collected).await; games_completed.fetch_add(1, Ordering::AcqRel); } } From 0ad3544450f51ce9de117f1c271d1150f4515d81 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 04:45:38 -0400 Subject: [PATCH 13/59] shorten training games for Descent --- alpha_paint/src/board/consts.rs | 6 +++--- alpha_paint/src/lib.rs | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/alpha_paint/src/board/consts.rs b/alpha_paint/src/board/consts.rs index 9096721..0c4030b 100644 --- a/alpha_paint/src/board/consts.rs +++ b/alpha_paint/src/board/consts.rs @@ -1,5 +1,5 @@ -pub const MAX_ROUNDS: usize = 1000; -pub const MAX_TURNS: usize = 2000; +pub const MAX_ROUNDS: usize = 100; +pub const MAX_TURNS: usize = 200; pub const BASE_MAX_STAMINA: usize = 100; pub const HILL_MAX_STAMINA_BONUS: usize = 40; pub const STAMINA_POWERUP_AMOUNT: usize = 35; @@ -25,7 +25,7 @@ pub const BID_TIME_LIMIT: f64 = 20.0; pub const PLAY_TIME_LIMIT: f64 = 180.0; pub const COMMENTATE_TIME_LIMIT: f64 = 1.0; pub const TIME_TIEBREAK_THRESH: f64 = 0.5; -pub const DOMINATION_WIN_THRESHOLD: f64 = 0.75; +pub const DOMINATION_WIN_THRESHOLD: f64 = 0.6; pub const ALLOW_DIAGONAL_MOVES: bool = false; pub const GLOBAL_DECAY_TURN_THRESHOLD: usize = 1000; pub const GLOBAL_DECAY_INTERVAL: usize = 100; diff --git a/alpha_paint/src/lib.rs b/alpha_paint/src/lib.rs index e7c2fec..5f76c5c 100644 --- a/alpha_paint/src/lib.rs +++ b/alpha_paint/src/lib.rs @@ -47,11 +47,11 @@ const PERFT_STR: [&str; 8] = [ "ap2|31x31|tc:0|cm:0|ep:0|w:15,0,99|b:15,30,100|h:n@14,14+16,14+15,15+14,16+16,16;n@0,13+0,14+1,14+2,14+2,15+0,16+1,16+2,16+0,17;n@30,13+28,14+29,14+30,14+28,15+28,16+29,16+30,16+30,17;n@13,6+14,6+15,6+16,6+17,6;n@13,24+14,24+15,24+16,24+17,24;n@7,1+7,8+7,22+7,29;n@23,1+23,8+23,22+23,29|pu:-|ps:2,21,11+2,21,19+2,19,6+2,19,24+2,29,8+2,29,22+2,4,10+2,4,20+2,13,4+2,13,26+2,23,10+2,23,20+2,20,0+2,20,30+2,18,3+2,18,27+2,19,9+2,19,21+2,22,5+2,22,25+52,16,14+52,16,16+52,26,10+52,26,20+52,29,10+52,29,20+52,26,12+52,26,18+52,12,6+52,12,24+102,4,6+102,4,24+102,29,14+102,29,16+102,2,14+102,2,16+102,23,12+102,23,18+102,21,4+102,21,26+152,16,10+152,16,20+152,27,8+152,27,22+152,5,10+152,5,20+152,9,8+152,9,22+152,4,12+152,4,18+202,12,11+202,12,19+202,30,14+202,30,16+202,7,6+202,7,24+202,23,2+202,23,28+202,4,11+202,4,19+252,8,3+252,8,27+252,25,0+252,25,30+252,5,2+252,5,28+252,26,4+252,26,26+252,26,9+252,26,21+302,1,12+302,1,18+302,10,1+302,10,29+302,30,13+302,30,17+302,28,7+302,28,23+302,1,0+302,1,30+352,9,11+352,9,19+352,20,0+352,20,30+352,23,13+352,23,17+352,9,6+352,9,24+352,23,8+352,23,22+402,22,11+402,22,19+402,14,13+402,14,17+402,30,10+402,30,20+402,28,11+402,28,19+402,20,7+402,20,23+452,5,12+452,5,18+452,1,8+452,1,22+452,13,2+452,13,28+452,1,7+452,1,23+452,23,10+452,23,20+502,21,11+502,21,19+502,9,0+502,9,30+502,30,1+502,30,29+502,26,1+502,26,29+502,19,5+502,19,25+552,15,2+552,15,28+552,28,3+552,28,27+552,15,8+552,15,22+552,14,7+552,14,23+552,30,1+552,30,29+602,7,6+602,7,24+602,20,5+602,20,25+602,11,5+602,11,25+602,15,6+602,15,24+602,2,0+602,2,30+652,1,12+652,1,18+652,17,1+652,17,29+652,15,8+652,15,22+652,15,6+652,15,24+652,27,13+652,27,17+702,6,11+702,6,19+702,24,8+702,24,22+702,14,13+702,14,17+702,11,5+702,11,25+702,30,2+702,30,28+752,30,3+752,30,27+752,3,15+752,28,1+752,28,29+752,28,3+752,28,27+752,18,12+752,18,18+802,8,3+802,8,27+802,15,2+802,15,28+802,12,12+802,12,18+802,17,1+802,17,29+802,19,3+802,19,27+852,19,11+852,19,19+852,22,0+852,22,30+852,10,11+852,10,19+852,9,5+852,9,25+852,22,13+852,22,17+902,24,14+902,24,16+902,1,13+902,1,17+902,3,5+902,3,25+902,28,7+902,28,23+902,19,12+902,19,18+952,26,15+952,21,10+952,21,20+952,21,8+952,21,22+952,11,3+952,11,27+952,13,11+952,13,19+1002,24,10+1002,24,20+1002,7,13+1002,7,17+1002,8,14+1002,8,16+1002,22,10+1002,22,20+1002,11,11+1002,11,19+1052,19,14+1052,19,16+1052,11,3+1052,11,27+1052,15,8+1052,15,22+1052,5,12+1052,5,18+1052,17,8+1052,17,22+1102,22,9+1102,22,21+1102,27,12+1102,27,18+1102,14,15+1102,26,1+1102,26,29+1102,23,4+1102,23,26+1152,3,8+1152,3,22+1152,30,3+1152,30,27+1152,17,5+1152,17,25+1152,4,0+1152,4,30+1152,6,11+1152,6,19+1202,17,6+1202,17,24+1202,19,9+1202,19,21+1202,1,0+1202,1,30+1202,14,6+1202,14,24+1202,27,13+1202,27,17+1252,19,5+1252,19,25+1252,18,0+1252,18,30+1252,24,10+1252,24,20+1252,4,2+1252,4,28+1252,21,10+1252,21,20+1302,21,15+1302,8,11+1302,8,19+1302,24,8+1302,24,22+1302,9,9+1302,9,21+1302,8,9+1302,8,21+1352,16,7+1352,16,23+1352,10,11+1352,10,19+1352,2,11+1352,2,19+1352,27,1+1352,27,29+1352,2,3+1352,2,27+1402,3,5+1402,3,25+1402,13,7+1402,13,23+1402,12,13+1402,12,17+1402,10,8+1402,10,22+1402,27,0+1402,27,30+1452,17,6+1452,17,24+1452,5,12+1452,5,18+1452,1,10+1452,1,20+1452,4,6+1452,4,24+1452,8,7+1452,8,23+1502,20,0+1502,20,30+1502,13,1+1502,13,29+1502,20,9+1502,20,21+1502,26,10+1502,26,20+1502,20,15+1552,27,2+1552,27,28+1552,17,15+1552,11,15+1552,21,6+1552,21,24+1552,4,14+1552,4,16+1552,20,1+1552,20,29+1602,23,6+1602,23,24+1602,19,15+1602,17,1+1602,17,29+1602,7,0+1602,7,30+1602,20,1+1602,20,29+1652,19,9+1652,19,21+1652,22,3+1652,22,27+1652,21,9+1652,21,21+1652,3,7+1652,3,23+1652,0,4+1652,0,26+1702,10,9+1702,10,21+1702,11,15+1702,15,6+1702,15,24+1702,30,14+1702,30,16+1702,27,4+1702,27,26+1752,23,8+1752,23,22+1752,21,7+1752,21,23+1752,20,12+1752,20,18+1752,4,2+1752,4,28+1752,10,11+1752,10,19+1802,21,8+1802,21,22+1802,19,7+1802,19,23+1802,11,4+1802,11,26+1802,19,9+1802,19,21+1802,1,14+1802,1,16+1852,20,8+1852,20,22+1852,1,3+1852,1,27+1852,19,14+1852,19,16+1852,12,3+1852,12,27+1852,17,10+1852,17,20+1902,6,3+1902,6,27+1902,30,11+1902,30,19+1902,15,1+1902,15,29+1902,10,8+1902,10,22+1902,30,3+1902,30,27+1952,5,7+1952,5,23+1952,9,5+1952,9,25+1952,30,9+1952,30,21+1952,27,1+1952,27,29+1952,28,4+1952,28,26|bd:#10#7#10#/1#3#5#7#5#3#1/2#9#5#9#2/3#1#7#3#7#1#3/10#1#5#1#10/5#19#5/3#1#4#9#4#1#3/12#5#12/5#6#5#6#5/5#19#5/3#6#9#6#3/5#1#15#1#5/13#####13/5#3#3#3#3#3#5/3#1#19#1#3/##27##/3#1#19#1#3/5#3#3#3#3#3#5/13#####13/5#1#15#1#5/3#6#9#6#3/5#19#5/5#6#5#6#5/12#5#12/3#1#4#9#4#1#3/5#19#5/10#1#5#1#10/3#1#7#3#7#1#3/2#9#5#9#2/1#3#5#7#5#3#1/#10#7#10#", ]; -pub const TRAINING_START_FENS: &[&str; 8] = &PERFT_STR; +pub const TRAINING_START_FENS: &[&str; 2] = &[PERFT_STR[0], PERFT_STR[1]]; pub fn perft_test() { use crate::board::Board; - use crate::{PERFT_STR, perft}; + use crate::{perft, PERFT_STR}; use std::time::Instant; for &perft_str in PERFT_STR.iter() { From 1b990eefc47fd301738c6c36e35fb25ccbc328cc Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 12:35:17 -0400 Subject: [PATCH 14/59] shorten self-play games for training --- alpha_paint/src/board/consts.rs | 6 +++--- training/src/worker.rs | 20 ++++++++++++++++++-- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/alpha_paint/src/board/consts.rs b/alpha_paint/src/board/consts.rs index 0c4030b..f039334 100644 --- a/alpha_paint/src/board/consts.rs +++ b/alpha_paint/src/board/consts.rs @@ -1,5 +1,5 @@ -pub const MAX_ROUNDS: usize = 100; -pub const MAX_TURNS: usize = 200; +pub const MAX_ROUNDS: usize = 32; +pub const MAX_TURNS: usize = 64; pub const BASE_MAX_STAMINA: usize = 100; pub const HILL_MAX_STAMINA_BONUS: usize = 40; pub const STAMINA_POWERUP_AMOUNT: usize = 35; @@ -25,7 +25,7 @@ pub const BID_TIME_LIMIT: f64 = 20.0; pub const PLAY_TIME_LIMIT: f64 = 180.0; pub const COMMENTATE_TIME_LIMIT: f64 = 1.0; pub const TIME_TIEBREAK_THRESH: f64 = 0.5; -pub const DOMINATION_WIN_THRESHOLD: f64 = 0.6; +pub const DOMINATION_WIN_THRESHOLD: f64 = 0.5; pub const ALLOW_DIAGONAL_MOVES: bool = false; pub const GLOBAL_DECAY_TURN_THRESHOLD: usize = 1000; pub const GLOBAL_DECAY_INTERVAL: usize = 100; diff --git a/training/src/worker.rs b/training/src/worker.rs index 4a023ad..233472f 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -10,6 +10,7 @@ use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; +use alpha_paint::board::actions::{Action, Move, Paint}; use alpha_paint::board::Board; use alpha_paint::TRAINING_START_FENS; @@ -33,6 +34,20 @@ impl Default for WorkerConfig { } } +fn training_commit_action(action: Action) -> Action { + match action { + Action::Move(mv) => Action::FinalMove(Move { + target: mv.target, + kind: mv.kind, + place_beacon: mv.place_beacon, + }), + Action::Paint(paint) => Action::FinalPaint(Paint { + target: paint.target, + }), + final_action => final_action, + } +} + /// Run a single self-play game with tree learning. /// /// At each move: @@ -85,7 +100,8 @@ async fn play_game( // Select action via ordinal distribution let action_id = tree.ordinal_select(); - let action = tree.root_node.children[action_id].action; + let selected_action = tree.root_node.children[action_id].action; + let action = training_commit_action(selected_action); // Apply action to get new board state let mut new_board = tree.root_state.clone(); @@ -102,7 +118,7 @@ async fn play_game( break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { - if tree.root_node.children[action_id].node.is_some() { + if action == selected_action && tree.root_node.children[action_id].node.is_some() { tree.step_tree(&new_board, action_id); } else { let next_rng = SmallRng::from_rng(rng); From 14b8e0bb21b5b977d7fe7b62f2036c4c3445c6de Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 12:46:43 -0400 Subject: [PATCH 15/59] revert the weird ass game changes This reverts commit 1b990eefc47fd301738c6c36e35fb25ccbc328cc. This reverts commit 0ad3544450f51ce9de117f1c271d1150f4515d81. --- alpha_paint/src/board/consts.rs | 6 +++--- alpha_paint/src/lib.rs | 4 ++-- training/src/worker.rs | 20 ++------------------ 3 files changed, 7 insertions(+), 23 deletions(-) diff --git a/alpha_paint/src/board/consts.rs b/alpha_paint/src/board/consts.rs index f039334..9096721 100644 --- a/alpha_paint/src/board/consts.rs +++ b/alpha_paint/src/board/consts.rs @@ -1,5 +1,5 @@ -pub const MAX_ROUNDS: usize = 32; -pub const MAX_TURNS: usize = 64; +pub const MAX_ROUNDS: usize = 1000; +pub const MAX_TURNS: usize = 2000; pub const BASE_MAX_STAMINA: usize = 100; pub const HILL_MAX_STAMINA_BONUS: usize = 40; pub const STAMINA_POWERUP_AMOUNT: usize = 35; @@ -25,7 +25,7 @@ pub const BID_TIME_LIMIT: f64 = 20.0; pub const PLAY_TIME_LIMIT: f64 = 180.0; pub const COMMENTATE_TIME_LIMIT: f64 = 1.0; pub const TIME_TIEBREAK_THRESH: f64 = 0.5; -pub const DOMINATION_WIN_THRESHOLD: f64 = 0.5; +pub const DOMINATION_WIN_THRESHOLD: f64 = 0.75; pub const ALLOW_DIAGONAL_MOVES: bool = false; pub const GLOBAL_DECAY_TURN_THRESHOLD: usize = 1000; pub const GLOBAL_DECAY_INTERVAL: usize = 100; diff --git a/alpha_paint/src/lib.rs b/alpha_paint/src/lib.rs index 5f76c5c..e7c2fec 100644 --- a/alpha_paint/src/lib.rs +++ b/alpha_paint/src/lib.rs @@ -47,11 +47,11 @@ const PERFT_STR: [&str; 8] = [ "ap2|31x31|tc:0|cm:0|ep:0|w:15,0,99|b:15,30,100|h:n@14,14+16,14+15,15+14,16+16,16;n@0,13+0,14+1,14+2,14+2,15+0,16+1,16+2,16+0,17;n@30,13+28,14+29,14+30,14+28,15+28,16+29,16+30,16+30,17;n@13,6+14,6+15,6+16,6+17,6;n@13,24+14,24+15,24+16,24+17,24;n@7,1+7,8+7,22+7,29;n@23,1+23,8+23,22+23,29|pu:-|ps:2,21,11+2,21,19+2,19,6+2,19,24+2,29,8+2,29,22+2,4,10+2,4,20+2,13,4+2,13,26+2,23,10+2,23,20+2,20,0+2,20,30+2,18,3+2,18,27+2,19,9+2,19,21+2,22,5+2,22,25+52,16,14+52,16,16+52,26,10+52,26,20+52,29,10+52,29,20+52,26,12+52,26,18+52,12,6+52,12,24+102,4,6+102,4,24+102,29,14+102,29,16+102,2,14+102,2,16+102,23,12+102,23,18+102,21,4+102,21,26+152,16,10+152,16,20+152,27,8+152,27,22+152,5,10+152,5,20+152,9,8+152,9,22+152,4,12+152,4,18+202,12,11+202,12,19+202,30,14+202,30,16+202,7,6+202,7,24+202,23,2+202,23,28+202,4,11+202,4,19+252,8,3+252,8,27+252,25,0+252,25,30+252,5,2+252,5,28+252,26,4+252,26,26+252,26,9+252,26,21+302,1,12+302,1,18+302,10,1+302,10,29+302,30,13+302,30,17+302,28,7+302,28,23+302,1,0+302,1,30+352,9,11+352,9,19+352,20,0+352,20,30+352,23,13+352,23,17+352,9,6+352,9,24+352,23,8+352,23,22+402,22,11+402,22,19+402,14,13+402,14,17+402,30,10+402,30,20+402,28,11+402,28,19+402,20,7+402,20,23+452,5,12+452,5,18+452,1,8+452,1,22+452,13,2+452,13,28+452,1,7+452,1,23+452,23,10+452,23,20+502,21,11+502,21,19+502,9,0+502,9,30+502,30,1+502,30,29+502,26,1+502,26,29+502,19,5+502,19,25+552,15,2+552,15,28+552,28,3+552,28,27+552,15,8+552,15,22+552,14,7+552,14,23+552,30,1+552,30,29+602,7,6+602,7,24+602,20,5+602,20,25+602,11,5+602,11,25+602,15,6+602,15,24+602,2,0+602,2,30+652,1,12+652,1,18+652,17,1+652,17,29+652,15,8+652,15,22+652,15,6+652,15,24+652,27,13+652,27,17+702,6,11+702,6,19+702,24,8+702,24,22+702,14,13+702,14,17+702,11,5+702,11,25+702,30,2+702,30,28+752,30,3+752,30,27+752,3,15+752,28,1+752,28,29+752,28,3+752,28,27+752,18,12+752,18,18+802,8,3+802,8,27+802,15,2+802,15,28+802,12,12+802,12,18+802,17,1+802,17,29+802,19,3+802,19,27+852,19,11+852,19,19+852,22,0+852,22,30+852,10,11+852,10,19+852,9,5+852,9,25+852,22,13+852,22,17+902,24,14+902,24,16+902,1,13+902,1,17+902,3,5+902,3,25+902,28,7+902,28,23+902,19,12+902,19,18+952,26,15+952,21,10+952,21,20+952,21,8+952,21,22+952,11,3+952,11,27+952,13,11+952,13,19+1002,24,10+1002,24,20+1002,7,13+1002,7,17+1002,8,14+1002,8,16+1002,22,10+1002,22,20+1002,11,11+1002,11,19+1052,19,14+1052,19,16+1052,11,3+1052,11,27+1052,15,8+1052,15,22+1052,5,12+1052,5,18+1052,17,8+1052,17,22+1102,22,9+1102,22,21+1102,27,12+1102,27,18+1102,14,15+1102,26,1+1102,26,29+1102,23,4+1102,23,26+1152,3,8+1152,3,22+1152,30,3+1152,30,27+1152,17,5+1152,17,25+1152,4,0+1152,4,30+1152,6,11+1152,6,19+1202,17,6+1202,17,24+1202,19,9+1202,19,21+1202,1,0+1202,1,30+1202,14,6+1202,14,24+1202,27,13+1202,27,17+1252,19,5+1252,19,25+1252,18,0+1252,18,30+1252,24,10+1252,24,20+1252,4,2+1252,4,28+1252,21,10+1252,21,20+1302,21,15+1302,8,11+1302,8,19+1302,24,8+1302,24,22+1302,9,9+1302,9,21+1302,8,9+1302,8,21+1352,16,7+1352,16,23+1352,10,11+1352,10,19+1352,2,11+1352,2,19+1352,27,1+1352,27,29+1352,2,3+1352,2,27+1402,3,5+1402,3,25+1402,13,7+1402,13,23+1402,12,13+1402,12,17+1402,10,8+1402,10,22+1402,27,0+1402,27,30+1452,17,6+1452,17,24+1452,5,12+1452,5,18+1452,1,10+1452,1,20+1452,4,6+1452,4,24+1452,8,7+1452,8,23+1502,20,0+1502,20,30+1502,13,1+1502,13,29+1502,20,9+1502,20,21+1502,26,10+1502,26,20+1502,20,15+1552,27,2+1552,27,28+1552,17,15+1552,11,15+1552,21,6+1552,21,24+1552,4,14+1552,4,16+1552,20,1+1552,20,29+1602,23,6+1602,23,24+1602,19,15+1602,17,1+1602,17,29+1602,7,0+1602,7,30+1602,20,1+1602,20,29+1652,19,9+1652,19,21+1652,22,3+1652,22,27+1652,21,9+1652,21,21+1652,3,7+1652,3,23+1652,0,4+1652,0,26+1702,10,9+1702,10,21+1702,11,15+1702,15,6+1702,15,24+1702,30,14+1702,30,16+1702,27,4+1702,27,26+1752,23,8+1752,23,22+1752,21,7+1752,21,23+1752,20,12+1752,20,18+1752,4,2+1752,4,28+1752,10,11+1752,10,19+1802,21,8+1802,21,22+1802,19,7+1802,19,23+1802,11,4+1802,11,26+1802,19,9+1802,19,21+1802,1,14+1802,1,16+1852,20,8+1852,20,22+1852,1,3+1852,1,27+1852,19,14+1852,19,16+1852,12,3+1852,12,27+1852,17,10+1852,17,20+1902,6,3+1902,6,27+1902,30,11+1902,30,19+1902,15,1+1902,15,29+1902,10,8+1902,10,22+1902,30,3+1902,30,27+1952,5,7+1952,5,23+1952,9,5+1952,9,25+1952,30,9+1952,30,21+1952,27,1+1952,27,29+1952,28,4+1952,28,26|bd:#10#7#10#/1#3#5#7#5#3#1/2#9#5#9#2/3#1#7#3#7#1#3/10#1#5#1#10/5#19#5/3#1#4#9#4#1#3/12#5#12/5#6#5#6#5/5#19#5/3#6#9#6#3/5#1#15#1#5/13#####13/5#3#3#3#3#3#5/3#1#19#1#3/##27##/3#1#19#1#3/5#3#3#3#3#3#5/13#####13/5#1#15#1#5/3#6#9#6#3/5#19#5/5#6#5#6#5/12#5#12/3#1#4#9#4#1#3/5#19#5/10#1#5#1#10/3#1#7#3#7#1#3/2#9#5#9#2/1#3#5#7#5#3#1/#10#7#10#", ]; -pub const TRAINING_START_FENS: &[&str; 2] = &[PERFT_STR[0], PERFT_STR[1]]; +pub const TRAINING_START_FENS: &[&str; 8] = &PERFT_STR; pub fn perft_test() { use crate::board::Board; - use crate::{perft, PERFT_STR}; + use crate::{PERFT_STR, perft}; use std::time::Instant; for &perft_str in PERFT_STR.iter() { diff --git a/training/src/worker.rs b/training/src/worker.rs index 233472f..4a023ad 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -10,7 +10,6 @@ use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; -use alpha_paint::board::actions::{Action, Move, Paint}; use alpha_paint::board::Board; use alpha_paint::TRAINING_START_FENS; @@ -34,20 +33,6 @@ impl Default for WorkerConfig { } } -fn training_commit_action(action: Action) -> Action { - match action { - Action::Move(mv) => Action::FinalMove(Move { - target: mv.target, - kind: mv.kind, - place_beacon: mv.place_beacon, - }), - Action::Paint(paint) => Action::FinalPaint(Paint { - target: paint.target, - }), - final_action => final_action, - } -} - /// Run a single self-play game with tree learning. /// /// At each move: @@ -100,8 +85,7 @@ async fn play_game( // Select action via ordinal distribution let action_id = tree.ordinal_select(); - let selected_action = tree.root_node.children[action_id].action; - let action = training_commit_action(selected_action); + let action = tree.root_node.children[action_id].action; // Apply action to get new board state let mut new_board = tree.root_state.clone(); @@ -118,7 +102,7 @@ async fn play_game( break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { - if action == selected_action && tree.root_node.children[action_id].node.is_some() { + if tree.root_node.children[action_id].node.is_some() { tree.step_tree(&new_board, action_id); } else { let next_rng = SmallRng::from_rng(rng); From 92c7ecfa39faa0807c9b6db820fea167f5d18066 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 13:33:54 -0400 Subject: [PATCH 16/59] stabilize value training targets Compress terminal rewards with a log scale, bound model outputs with tanh, and bump replay/training defaults so overnight runs reuse data more while keeping losses sane. --- python/alphapaint_training/model.py | 2 +- python/alphapaint_training/train.py | 47 ++++++++++++++++++-------- training/src/descent.rs | 52 +++++++++++++++++++++++++++-- 3 files changed, 83 insertions(+), 18 deletions(-) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index b445b13..25c3b93 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -163,7 +163,7 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: board = board.to(dtype=param_dtype) intrinsics = intrinsics.to(dtype=param_dtype) value = self.value_net(board, intrinsics) - return value.squeeze(-1) + return value.squeeze(-1).tanh() * 10.0 __all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index ed8595f..089d928 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -18,10 +18,10 @@ @dataclass(slots=True) class TrainConfig: rounds: int = 1 - samples_per_round: int = 8192 - train_steps_per_round: int = 64 - batch_size: int = 512 - replay_capacity: int = 262144 + samples_per_round: int = 262_144 + train_steps_per_round: int = 384 + batch_size: int = 2048 + replay_capacity: int = 12_000_000 num_threads: int = 16 workers_per_thread: int = 16 descent_iterations: int = 50 @@ -260,17 +260,36 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: def _parse_args() -> TrainConfig: + defaults = TrainConfig() parser = argparse.ArgumentParser(description="Run AlphaPaint value training") - parser.add_argument("--rounds", type=int, default=1) - parser.add_argument("--samples-per-round", type=int, default=8192) - parser.add_argument("--train-steps-per-round", type=int, default=64) - parser.add_argument("--batch-size", type=int, default=512) - parser.add_argument("--replay-capacity", type=int, default=262144) - parser.add_argument("--num-threads", type=int, default=16) - parser.add_argument("--workers-per-thread", type=int, default=16) - parser.add_argument("--descent-iterations", type=int, default=50) - parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--run-dir", default="runs/latest") + parser.add_argument("--rounds", type=int, default=defaults.rounds) + parser.add_argument( + "--samples-per-round", type=int, default=defaults.samples_per_round + ) + parser.add_argument( + "--train-steps-per-round", type=int, default=defaults.train_steps_per_round + ) + parser.add_argument("--batch-size", type=int, default=defaults.batch_size) + parser.add_argument("--replay-capacity", type=int, default=defaults.replay_capacity) + parser.add_argument("--num-threads", type=int, default=defaults.num_threads) + parser.add_argument( + "--workers-per-thread", type=int, default=defaults.workers_per_thread + ) + parser.add_argument( + "--descent-iterations", type=int, default=defaults.descent_iterations + ) + parser.add_argument("--lr", type=float, default=defaults.lr) + parser.add_argument("--seed", type=int, default=defaults.seed) + parser.add_argument( + "--selfplay-precision", + default=defaults.selfplay_precision, + choices=["bf16", "fp16", "fp32"], + ) + parser.add_argument("--device", default=defaults.device) + parser.add_argument( + "--checkpoint-interval", type=int, default=defaults.checkpoint_interval + ) + parser.add_argument("--run-dir", default=defaults.run_dir) args = parser.parse_args() return TrainConfig(**vars(args)) diff --git a/training/src/descent.rs b/training/src/descent.rs index be90e0e..6561c41 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -11,6 +11,8 @@ use rand::{Rng, RngExt}; use crate::eval::Evaluator; +const AVG_GAME_LENGTH: f32 = 500.0; + pub struct TreeLearningSample { pub board: Board, pub value: f32, @@ -284,11 +286,11 @@ impl SearchNode { } } - /// Depth heuristic: larger magnitude for faster wins/losses. + /// Reinforcement heuristic: preserve faster-win ordering without exploding. fn value_from_term(board: &Board, term: TerminalState) -> f32 { let sign = term.value() as f32; - // Scale: prefer faster wins. Max turns = 2000. - sign * (2000.0 - board.turn_count as f32) + let p = board.turn_count.max(1) as f32; + sign * (AVG_GAME_LENGTH / p).ln_1p() } async fn ubfms_iteration( @@ -393,6 +395,50 @@ impl SearchNode { } } +#[cfg(test)] +mod tests { + use super::*; + use alpha_paint::board::board_structs::Player; + + fn board_with_turn_count(turn_count: usize) -> Board { + Board::from_fen(&format!( + "ap2|3x3|tc:{turn_count}|cm:0|ep:0|w:0,0,99|b:2,2,99|h:-|pu:-|ps:-|bd:3/3/3" + )) + .expect("board FEN must parse") + } + + #[test] + fn value_from_term_compresses_terminal_scale() { + let board = board_with_turn_count(1); + let value = SearchNode::value_from_term(&board, TerminalState::Win(Player::White)); + + assert!(value > 6.0); + assert!(value < 7.0); + } + + #[test] + fn value_from_term_preserves_faster_win_ordering() { + let early = board_with_turn_count(1); + let late = board_with_turn_count(500); + + let early_value = SearchNode::value_from_term(&early, TerminalState::Win(Player::White)); + let late_value = SearchNode::value_from_term(&late, TerminalState::Win(Player::White)); + + assert!(early_value > late_value); + } + + #[test] + fn value_from_term_preserves_faster_loss_ordering() { + let early = board_with_turn_count(1); + let late = board_with_turn_count(500); + + let early_value = SearchNode::value_from_term(&early, TerminalState::Win(Player::Black)); + let late_value = SearchNode::value_from_term(&late, TerminalState::Win(Player::Black)); + + assert!(early_value < late_value); + } +} + pub struct GameSearchTree<'a, E: Evaluator> { pub root_node: Box, pub root_state: Board, From f83ff32f0ae628998541af0b2d660f669f7baa5e Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 14:42:52 -0400 Subject: [PATCH 17/59] add wandb and update some constants --- .gitignore | 1 + pyproject.toml | 1 + python/alphapaint_training/train.py | 32 ++-- uv.lock | 263 ++++++++++++++++++++++++++++ 4 files changed, 281 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index 19ef6e0..4160fe2 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ build/ .venv/ .idea/ +runs/ diff --git a/pyproject.toml b/pyproject.toml index 38f4b58..37e8b75 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,7 @@ dependencies = [ "py-cpuinfo", "torch==2.10.0", "triton>=3.6.0", + "wandb>=0.25.1", ] classifiers = [ "Programming Language :: Rust", diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 089d928..89f26e2 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -10,27 +10,29 @@ import numpy as np import torch import torch.nn.functional as F +import wandb -from . import EphemeralReplayBuffer, SelfPlay -from .model import PackedValueModel +from alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training.model import PackedValueModel @dataclass(slots=True) class TrainConfig: rounds: int = 1 - samples_per_round: int = 262_144 - train_steps_per_round: int = 384 - batch_size: int = 2048 - replay_capacity: int = 12_000_000 - num_threads: int = 16 + samples_per_round: int = 1_048_576 + train_steps_per_round: int = 128 + batch_size: int = 24_576 + replay_capacity: int = 16_000_000 + num_threads: int = 32 workers_per_thread: int = 16 - descent_iterations: int = 50 + descent_iterations: int = 10 lr: float = 3e-4 seed: int = 42 selfplay_precision: str = "bf16" device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" + wandb: bool = False def _device(device: str) -> torch.device: @@ -98,12 +100,6 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: return run_dir, checkpoint_dir -def _append_jsonl(path: Path, record: dict[str, object]) -> None: - with path.open("a", encoding="ascii") as f: - json.dump(record, f, sort_keys=True) - f.write("\n") - - def _save_checkpoint( *, checkpoint_dir: Path, @@ -141,7 +137,10 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: torch.backends.cudnn.benchmark = True run_dir, checkpoint_dir = _prepare_run_dir(config) - metrics_path = run_dir / "metrics.jsonl" + + wandb.init( + project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) + ) model = PackedValueModel().to(device) model.eval() @@ -224,13 +223,13 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "loss_last": last_loss, "timestamp": time.time(), } - _append_jsonl(metrics_path, record) print( f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " f"batches={gpu_batches_added} replay={replay_size} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " f"loss={mean_loss:.6f}/{last_loss:.6f}" ) + wandb.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -290,6 +289,7 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) + parser.add_argument("--wandb", action="store_true", default=True) args = parser.parse_args() return TrainConfig(**vars(args)) diff --git a/uv.lock b/uv.lock index 2406642..53dfac2 100644 --- a/uv.lock +++ b/uv.lock @@ -13,6 +13,7 @@ dependencies = [ { name = "py-cpuinfo" }, { name = "torch" }, { name = "triton" }, + { name = "wandb" }, ] [package.dev-dependencies] @@ -30,6 +31,7 @@ requires-dist = [ { name = "py-cpuinfo" }, { name = "torch", specifier = "==2.10.0" }, { name = "triton", specifier = ">=3.6.0" }, + { name = "wandb", specifier = ">=0.25.1" }, ] [package.metadata.requires-dev] @@ -38,6 +40,61 @@ dev = [ { name = "pytest-xdist", specifier = ">=3" }, ] +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + +[[package]] +name = "certifi" +version = "2026.2.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/2d/7bf41579a8986e348fa033a31cdd0e4121114f6bce2457e8876010b092dd/certifi-2026.2.25.tar.gz", hash = "sha256:e887ab5cee78ea814d3472169153c2d12cd43b14bd03329a39a9c6e2e80bfba7", size = 155029, upload-time = "2026-02-25T02:54:17.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/60/e3bec1881450851b087e301bedc3daa9377a4d45f1c26aa90b0b235e38aa/charset_normalizer-3.4.6.tar.gz", hash = "sha256:1ae6b62897110aa7c79ea2f5dd38d1abca6db663687c0b1ad9aed6f6bae3d9d6", size = 143363, upload-time = "2026-03-15T18:53:25.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/62/c0815c992c9545347aeea7859b50dc9044d147e2e7278329c6e02ac9a616/charset_normalizer-3.4.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2ef7fedc7a6ecbe99969cd09632516738a97eeb8bd7258bf8a0f23114c057dab", size = 295154, upload-time = "2026-03-15T18:50:50.88Z" }, + { url = "https://files.pythonhosted.org/packages/a8/37/bdca6613c2e3c58c7421891d80cc3efa1d32e882f7c4a7ee6039c3fc951a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4ea868bc28109052790eb2b52a9ab33f3aa7adc02f96673526ff47419490e21", size = 199191, upload-time = "2026-03-15T18:50:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/6c/92/9934d1bbd69f7f398b38c5dae1cbf9cc672e7c34a4adf7b17c0a9c17d15d/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:836ab36280f21fc1a03c99cd05c6b7af70d2697e374c7af0b61ed271401a72a2", size = 218674, upload-time = "2026-03-15T18:50:54.102Z" }, + { url = "https://files.pythonhosted.org/packages/af/90/25f6ab406659286be929fd89ab0e78e38aa183fc374e03aa3c12d730af8a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f1ce721c8a7dfec21fcbdfe04e8f68174183cf4e8188e0645e92aa23985c57ff", size = 215259, upload-time = "2026-03-15T18:50:55.616Z" }, + { url = "https://files.pythonhosted.org/packages/4e/ef/79a463eb0fff7f96afa04c1d4c51f8fc85426f918db467854bfb6a569ce3/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e28d62a8fc7a1fa411c43bd65e346f3bce9716dc51b897fbe930c5987b402d5", size = 207276, upload-time = "2026-03-15T18:50:57.054Z" }, + { url = "https://files.pythonhosted.org/packages/f7/72/d0426afec4b71dc159fa6b4e68f868cd5a3ecd918fec5813a15d292a7d10/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:530d548084c4a9f7a16ed4a294d459b4f229db50df689bfe92027452452943a0", size = 195161, upload-time = "2026-03-15T18:50:58.686Z" }, + { url = "https://files.pythonhosted.org/packages/bf/18/c82b06a68bfcb6ce55e508225d210c7e6a4ea122bfc0748892f3dc4e8e11/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:30f445ae60aad5e1f8bdbb3108e39f6fbc09f4ea16c815c66578878325f8f15a", size = 203452, upload-time = "2026-03-15T18:51:00.196Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/0c25979b92f8adafdbb946160348d8d44aa60ce99afdc27df524379875cb/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ac2393c73378fea4e52aa56285a3d64be50f1a12395afef9cce47772f60334c2", size = 202272, upload-time = "2026-03-15T18:51:01.703Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3d/7fea3e8fe84136bebbac715dd1221cc25c173c57a699c030ab9b8900cbb7/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:90ca27cd8da8118b18a52d5f547859cc1f8354a00cd1e8e5120df3e30d6279e5", size = 195622, upload-time = "2026-03-15T18:51:03.526Z" }, + { url = "https://files.pythonhosted.org/packages/57/8a/d6f7fd5cb96c58ef2f681424fbca01264461336d2a7fc875e4446b1f1346/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e5a94886bedca0f9b78fecd6afb6629142fd2605aa70a125d49f4edc6037ee6", size = 220056, upload-time = "2026-03-15T18:51:05.269Z" }, + { url = "https://files.pythonhosted.org/packages/16/50/478cdda782c8c9c3fb5da3cc72dd7f331f031e7f1363a893cdd6ca0f8de0/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:695f5c2823691a25f17bc5d5ffe79fa90972cc34b002ac6c843bb8a1720e950d", size = 203751, upload-time = "2026-03-15T18:51:06.858Z" }, + { url = "https://files.pythonhosted.org/packages/75/fc/cc2fcac943939c8e4d8791abfa139f685e5150cae9f94b60f12520feaa9b/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:231d4da14bcd9301310faf492051bee27df11f2bc7549bc0bb41fef11b82daa2", size = 216563, upload-time = "2026-03-15T18:51:08.564Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b7/a4add1d9a5f68f3d037261aecca83abdb0ab15960a3591d340e829b37298/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a056d1ad2633548ca18ffa2f85c202cfb48b68615129143915b8dc72a806a923", size = 209265, upload-time = "2026-03-15T18:51:10.312Z" }, + { url = "https://files.pythonhosted.org/packages/6c/18/c094561b5d64a24277707698e54b7f67bd17a4f857bbfbb1072bba07c8bf/charset_normalizer-3.4.6-cp312-cp312-win32.whl", hash = "sha256:c2274ca724536f173122f36c98ce188fd24ce3dad886ec2b7af859518ce008a4", size = 144229, upload-time = "2026-03-15T18:51:11.694Z" }, + { url = "https://files.pythonhosted.org/packages/ab/20/0567efb3a8fd481b8f34f739ebddc098ed062a59fed41a8d193a61939e8f/charset_normalizer-3.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:c8ae56368f8cc97c7e40a7ee18e1cedaf8e780cd8bc5ed5ac8b81f238614facb", size = 154277, upload-time = "2026-03-15T18:51:13.004Z" }, + { url = "https://files.pythonhosted.org/packages/15/57/28d79b44b51933119e21f65479d0864a8d5893e494cf5daab15df0247c17/charset_normalizer-3.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:899d28f422116b08be5118ef350c292b36fc15ec2daeb9ea987c89281c7bb5c4", size = 142817, upload-time = "2026-03-15T18:51:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/2a/68/687187c7e26cb24ccbd88e5069f5ef00eba804d36dde11d99aad0838ab45/charset_normalizer-3.4.6-py3-none-any.whl", hash = "sha256:947cf925bc916d90adba35a64c82aace04fa39b46b52d4630ece166655905a69", size = 61455, upload-time = "2026-03-15T18:53:23.833Z" }, +] + +[[package]] +name = "click" +version = "8.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3d/fa/656b739db8587d7b5dfa22e22ed02566950fbfbcdc20311993483657a5c0/click-8.3.1.tar.gz", hash = "sha256:12ff4785d337a1bb490bb7e9c2b1ee5da3112e94a8622f26a6c77f5d2fc6842a", size = 295065, upload-time = "2025-11-15T20:45:42.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -110,6 +167,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, ] +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.1.46" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/b5/59d16470a1f0dfe8c793f9ef56fd3826093fc52b3bd96d6b9d6c26c7e27b/gitpython-3.1.46.tar.gz", hash = "sha256:400124c7d0ef4ea03f7310ac2fbf7151e09ff97f2a3288d64a440c584a29c37f", size = 215371, upload-time = "2026-01-01T15:37:32.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/09/e21df6aef1e1ffc0c816f0522ddc3f6dcded766c3261813131c78a704470/gitpython-3.1.46-py3-none-any.whl", hash = "sha256:79812ed143d9d25b6d176a10bb511de0f9c67b1fa641d82097b0ab90398a2058", size = 208620, upload-time = "2026-01-01T15:37:30.574Z" }, +] + +[[package]] +name = "idna" +version = "3.11" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/0703ccc57f3a7233505399edb88de3cbd678da106337b9fcde432b65ed60/idna-3.11.tar.gz", hash = "sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902", size = 194582, upload-time = "2025-10-12T14:55:20.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -359,6 +449,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, ] +[[package]] +name = "platformdirs" +version = "4.9.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/56/8d4c30c8a1d07013911a8fdbd8f89440ef9f08d07a1b50ab8ca8be5a20f9/platformdirs-4.9.4.tar.gz", hash = "sha256:1ec356301b7dc906d83f371c8f487070e99d3ccf9e501686456394622a01a934", size = 28737, upload-time = "2026-03-05T18:34:13.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/d7/97f7e3a6abb67d8080dd406fd4df842c2be0efaf712d1c899c32a075027c/platformdirs-4.9.4-py3-none-any.whl", hash = "sha256:68a9a4619a666ea6439f2ff250c12a853cd1cbd5158d258bd824a7df6be2f868", size = 21216, upload-time = "2026-03-05T18:34:12.172Z" }, +] + [[package]] name = "pluggy" version = "1.6.0" @@ -368,6 +467,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "protobuf" +version = "6.33.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, +] + [[package]] name = "psutil" version = "5.9.0" @@ -383,6 +497,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, ] +[[package]] +name = "pydantic" +version = "2.12.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.41.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/5d/5f6c63eebb5afee93bcaae4ce9a898f3373ca23df3ccaef086d0233a35a7/pydantic_core-2.41.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:f41a7489d32336dbf2199c8c0a215390a751c5b014c2c1c5366e817202e9cdf7", size = 2110990, upload-time = "2025-11-04T13:39:58.079Z" }, + { url = "https://files.pythonhosted.org/packages/aa/32/9c2e8ccb57c01111e0fd091f236c7b371c1bccea0fa85247ac55b1e2b6b6/pydantic_core-2.41.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:070259a8818988b9a84a449a2a7337c7f430a22acc0859c6b110aa7212a6d9c0", size = 1896003, upload-time = "2025-11-04T13:39:59.956Z" }, + { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, + { url = "https://files.pythonhosted.org/packages/38/de/8c36b5198a29bdaade07b5985e80a233a5ac27137846f3bc2d3b40a47360/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed2e99c456e3fadd05c991f8f437ef902e00eedf34320ba2b0842bd1c3ca3a75", size = 2052578, upload-time = "2025-11-04T13:40:04.401Z" }, + { url = "https://files.pythonhosted.org/packages/00/b5/0e8e4b5b081eac6cb3dbb7e60a65907549a1ce035a724368c330112adfdd/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:65840751b72fbfd82c3c640cff9284545342a4f1eb1586ad0636955b261b0b05", size = 2208504, upload-time = "2025-11-04T13:40:06.072Z" }, + { url = "https://files.pythonhosted.org/packages/77/56/87a61aad59c7c5b9dc8caad5a41a5545cba3810c3e828708b3d7404f6cef/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e536c98a7626a98feb2d3eaf75944ef6f3dbee447e1f841eae16f2f0a72d8ddc", size = 2335816, upload-time = "2025-11-04T13:40:07.835Z" }, + { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, + { url = "https://files.pythonhosted.org/packages/d3/43/ebef01f69baa07a482844faaa0a591bad1ef129253ffd0cdaa9d8a7f72d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d38548150c39b74aeeb0ce8ee1d8e82696f4a4e16ddc6de7b1d8823f7de4b9b5", size = 2171698, upload-time = "2025-11-04T13:40:12.004Z" }, + { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, + { url = "https://files.pythonhosted.org/packages/49/7d/4c00df99cb12070b6bccdef4a195255e6020a550d572768d92cc54dba91a/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:482c982f814460eabe1d3bb0adfdc583387bd4691ef00b90575ca0d2b6fe2294", size = 2329591, upload-time = "2025-11-04T13:40:15.672Z" }, + { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, + { url = "https://files.pythonhosted.org/packages/49/3b/774f2b5cd4192d5ab75870ce4381fd89cf218af999515baf07e7206753f0/pydantic_core-2.41.5-cp312-cp312-win32.whl", hash = "sha256:b74557b16e390ec12dca509bce9264c3bbd128f8a2c376eaa68003d7f327276d", size = 1985908, upload-time = "2025-11-04T13:40:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/86/45/00173a033c801cacf67c190fef088789394feaf88a98a7035b0e40d53dc9/pydantic_core-2.41.5-cp312-cp312-win_amd64.whl", hash = "sha256:1962293292865bca8e54702b08a4f26da73adc83dd1fcf26fbc875b35d81c815", size = 2020145, upload-time = "2025-11-04T13:40:21.548Z" }, + { url = "https://files.pythonhosted.org/packages/f9/22/91fbc821fa6d261b376a3f73809f907cec5ca6025642c463d3488aad22fb/pydantic_core-2.41.5-cp312-cp312-win_arm64.whl", hash = "sha256:1746d4a3d9a794cacae06a5eaaccb4b8643a131d45fbc9af23e353dc0a5ba5c3", size = 1976179, upload-time = "2025-11-04T13:40:23.393Z" }, + { url = "https://files.pythonhosted.org/packages/09/32/59b0c7e63e277fa7911c2fc70ccfb45ce4b98991e7ef37110663437005af/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:7da7087d756b19037bc2c06edc6c170eeef3c3bafcb8f532ff17d64dc427adfd", size = 2110495, upload-time = "2025-11-04T13:42:49.689Z" }, + { url = "https://files.pythonhosted.org/packages/aa/81/05e400037eaf55ad400bcd318c05bb345b57e708887f07ddb2d20e3f0e98/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:aabf5777b5c8ca26f7824cb4a120a740c9588ed58df9b2d196ce92fba42ff8dc", size = 1915388, upload-time = "2025-11-04T13:42:52.215Z" }, + { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, + { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, +] + [[package]] name = "pygments" version = "2.19.2" @@ -421,6 +579,52 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, ] +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, +] + +[[package]] +name = "requests" +version = "2.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/34/64/8860370b167a9721e8956ae116825caff829224fbca0ca6e7bf8ddef8430/requests-2.33.0.tar.gz", hash = "sha256:c7ebc5e8b0f21837386ad0e1c8fe8b829fa5f544d8df3b2253bff14ef29d7652", size = 134232, upload-time = "2026-03-25T15:10:41.586Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/5d/c814546c2333ceea4ba42262d8c4d55763003e767fa169adc693bd524478/requests-2.33.0-py3-none-any.whl", hash = "sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133dcde67d48cee69b", size = 65017, upload-time = "2026-03-25T15:10:40.382Z" }, +] + +[[package]] +name = "sentry-sdk" +version = "2.56.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/df/5008954f5466085966468612a7d1638487596ee6d2fd7fb51783a85351bf/sentry_sdk-2.56.0.tar.gz", hash = "sha256:fdab72030b69625665b2eeb9738bdde748ad254e8073085a0ce95382678e8168", size = 426820, upload-time = "2026-03-24T09:56:36.575Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/1a/b3a3e9f6520493fed7997af4d2de7965d71549c62f994a8fd15f2ecd519e/sentry_sdk-2.56.0-py2.py3-none-any.whl", hash = "sha256:5afafb744ceb91d22f4cc650c6bd048ac6af5f7412dcc6c59305a2e36f4dbc02", size = 451568, upload-time = "2026-03-24T09:56:34.807Z" }, +] + [[package]] name = "setuptools" version = "81.0.0" @@ -430,6 +634,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, ] +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + [[package]] name = "sympy" version = "1.14.0" @@ -498,3 +711,53 @@ sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac8 wheels = [ { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, ] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "urllib3" +version = "2.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, +] + +[[package]] +name = "wandb" +version = "0.25.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "gitpython" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sentry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/60/bb/eb579bf9abac70934a014a9d4e45346aab307994f3021d201bebe5fa25ec/wandb-0.25.1.tar.gz", hash = "sha256:b2a95cd777ecbe7499599a43158834983448a0048329bc7210ef46ca18d21994", size = 43983308, upload-time = "2026-03-10T23:51:44.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/d8/873553b6818499d1b1de314067d528b892897baf0dc81fedc0e845abc2dd/wandb-0.25.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:9bb0679a3e2dcd96db9d9b6d3e17d046241d8d122974b24facb85cc93309a8c9", size = 23615900, upload-time = "2026-03-10T23:51:06.278Z" }, + { url = "https://files.pythonhosted.org/packages/71/ea/b131f319aaa5d0bf7572b6bfcff3dd89e1cf92b17eee443bbab71d12d74c/wandb-0.25.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:0fb13ed18914027523e7b4fc20380c520e0d10da0ee452f924a13f84509fbe12", size = 25576144, upload-time = "2026-03-10T23:51:11.527Z" }, + { url = "https://files.pythonhosted.org/packages/70/5f/81508581f0bb77b0495665c1c78e77606a48e66e855ca71ba7c8ae29efa4/wandb-0.25.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:cc4521eb5223429ddab5e8eee9b42fdf4caabdf0bc4e0e809042720e5fbef0ed", size = 23070425, upload-time = "2026-03-10T23:51:15.71Z" }, + { url = "https://files.pythonhosted.org/packages/f2/c7/445155ef010e2e35d190797d7c36ff441e062a5b566a6da4778e22233395/wandb-0.25.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:e73b4c55b947edae349232d5845204d30fac88e18eb4ad1d4b96bf7cf898405a", size = 25628142, upload-time = "2026-03-10T23:51:19.326Z" }, + { url = "https://files.pythonhosted.org/packages/d5/63/f5c55ee00cf481ef1ccd3c385a0585ad52e7840d08419d4f82ddbeeea959/wandb-0.25.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:22b84065aa398e1624d2e5ad79e08bc4d2af41a6db61697b03b3aaba332977c6", size = 23123172, upload-time = "2026-03-10T23:51:23.418Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d9/19eb7974c0e9253bcbaee655222c0f0e1a52e63e9479ee711b4208f8ac31/wandb-0.25.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:005c4c6b5126ef8f4b4110e5372d950918b00637d6dc4b615ad17445f9739478", size = 25714479, upload-time = "2026-03-10T23:51:27.421Z" }, + { url = "https://files.pythonhosted.org/packages/11/19/466c1d03323a4a0ed7d4036a59b18d6b6f67cb5032e444205927e226b18d/wandb-0.25.1-py3-none-win32.whl", hash = "sha256:8f2d04f16b88d65bfba9d79fb945f6c64e2686215469a841936e0972be8ec6a5", size = 24967338, upload-time = "2026-03-10T23:51:31.833Z" }, + { url = "https://files.pythonhosted.org/packages/89/22/680d34c1587f3a979c701b66d71aa7c42b4ef2fdf0774f67034e618e834e/wandb-0.25.1-py3-none-win_amd64.whl", hash = "sha256:62db5166de14456156d7a85953a58733a631228e6d4248a753605f75f75fb845", size = 24967343, upload-time = "2026-03-10T23:51:36.026Z" }, + { url = "https://files.pythonhosted.org/packages/c4/e8/76836b75d401ff5912aaf513176e64557ceaec4c4946bfd38a698ff84d48/wandb-0.25.1-py3-none-win_arm64.whl", hash = "sha256:cc7c34b70cf4b7be4d395541e82e325fd9d2be978d62c9ec01f1a7141523b6bb", size = 22080774, upload-time = "2026-03-10T23:51:40.196Z" }, +] From dd487c73511d13d304cd4c3b75b69dadf69a0f48 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 17:21:56 -0400 Subject: [PATCH 18/59] instrument self-play and stop replaying reused tree samples This makes it easier to see why training stalls in shallow openings and avoids overweighting old tree data when the same subtree is reused across moves. --- python/alphapaint_training/train.py | 68 ++++++++++++++++++++- training/src/descent.rs | 58 ++++++++++++++++-- training/src/lib.rs | 75 +++++++++++++++-------- training/src/training.rs | 92 ++++++++++++++++++++++++++--- training/src/worker.rs | 83 +++++++++++++++++++++++--- 5 files changed, 332 insertions(+), 44 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 89f26e2..42cdd6a 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -162,6 +162,12 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_games = 0 previous_gpu_batches = 0 previous_gpu_evals = 0 + previous_action_steps = 0 + previous_final_actions = 0 + previous_nonfinal_actions = 0 + previous_completed_turns = 0 + previous_action_turn_count_total = 0 + previous_completed_game_actions_total = 0 try: for round_idx in range(config.rounds): round_number = round_idx + 1 @@ -171,11 +177,46 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: games = selfplay.games() gpu_batches = selfplay.gpu_batches() gpu_evals = selfplay.gpu_evals() + action_steps = selfplay.action_steps() + final_actions = selfplay.final_actions() + nonfinal_actions = selfplay.nonfinal_actions() + completed_turns = selfplay.completed_turns() + action_turn_count_total = selfplay.action_turn_count_total() + max_turn_count_seen = selfplay.max_turn_count_seen() + completed_game_actions_total = selfplay.completed_game_actions_total() + max_actions_in_completed_game = selfplay.max_actions_in_completed_game() collect_seconds = time.perf_counter() - collect_started_at samples_added = collected - previous_samples games_added = games - previous_games gpu_batches_added = gpu_batches - previous_gpu_batches gpu_evals_added = gpu_evals - previous_gpu_evals + action_steps_added = action_steps - previous_action_steps + final_actions_added = final_actions - previous_final_actions + nonfinal_actions_added = nonfinal_actions - previous_nonfinal_actions + completed_turns_added = completed_turns - previous_completed_turns + action_turn_count_added = ( + action_turn_count_total - previous_action_turn_count_total + ) + completed_game_actions_added = ( + completed_game_actions_total - previous_completed_game_actions_total + ) + avg_turn_count = ( + action_turn_count_added / action_steps_added + if action_steps_added + else 0.0 + ) + actions_per_turn = ( + action_steps_added / completed_turns_added + if completed_turns_added + else 0.0 + ) + actions_per_game = ( + completed_game_actions_added / games_added if games_added else 0.0 + ) + actions_per_turn_display = ( + f"{actions_per_turn:.2f}" if completed_turns_added else "-" + ) + actions_per_game_display = f"{actions_per_game:.1f}" if games_added else "-" round_losses = [] train_started_at = time.perf_counter() @@ -211,6 +252,23 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "gpu_evals_total": gpu_evals, "gpu_evals_added": gpu_evals_added, "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), + "action_steps_total": action_steps, + "action_steps_added": action_steps_added, + "final_actions_total": final_actions, + "final_actions_added": final_actions_added, + "nonfinal_actions_total": nonfinal_actions, + "nonfinal_actions_added": nonfinal_actions_added, + "completed_turns_total": completed_turns, + "completed_turns_added": completed_turns_added, + "action_turn_count_total": action_turn_count_total, + "action_turn_count_added": action_turn_count_added, + "avg_turn_count": avg_turn_count, + "max_turn_count_seen": max_turn_count_seen, + "actions_per_turn": actions_per_turn, + "completed_game_actions_total": completed_game_actions_total, + "completed_game_actions_added": completed_game_actions_added, + "actions_per_completed_game": actions_per_game, + "max_actions_in_completed_game": max_actions_in_completed_game, "replay_size": replay_size, "collection_seconds": collect_seconds, "training_seconds": train_seconds, @@ -226,7 +284,9 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print( f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " - f"batches={gpu_batches_added} replay={replay_size} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " + f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " + f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " + f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " f"loss={mean_loss:.6f}/{last_loss:.6f}" ) wandb.log(record) @@ -252,6 +312,12 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_games = games previous_gpu_batches = gpu_batches previous_gpu_evals = gpu_evals + previous_action_steps = action_steps + previous_final_actions = final_actions + previous_nonfinal_actions = nonfinal_actions + previous_completed_turns = completed_turns + previous_action_turn_count_total = action_turn_count_total + previous_completed_game_actions_total = completed_game_actions_total finally: selfplay.drop() diff --git a/training/src/descent.rs b/training/src/descent.rs index 6561c41..9bcad41 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -40,6 +40,7 @@ pub struct SearchNode { pub value: f32, pub completion_value: i32, resolved: bool, + replay_emitted: bool, pub children: Vec, } @@ -49,6 +50,7 @@ impl SearchNode { value, completion_value, resolved: is_resolved, + replay_emitted: false, children: vec![], } } @@ -167,6 +169,7 @@ impl SearchNode { value: term_value, completion_value: comp_value, resolved: true, + replay_emitted: false, children: vec![ChildData { action, child_value: term_value, @@ -376,17 +379,18 @@ impl SearchNode { /// An internal node is one that has children and at least one expanded child. /// Non-terminal leaf nodes (where the network estimate was used without /// minimax backing) are excluded per Athénan's tree learning rules. - fn collect_samples(&self, state: &mut Board, out: &mut Vec) { + fn collect_samples(&mut self, state: &mut Board, out: &mut Vec) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); - if has_expanded_child { + if has_expanded_child && !self.replay_emitted { out.push(TreeLearningSample { board: state.clone(), value: self.value, }); + self.replay_emitted = true; } - for child in &self.children { - if let Some(ref node) = child.node { + for child in &mut self.children { + if let Some(node) = child.node.as_mut() { let (_, rollback) = state.apply_action(child.action); node.collect_samples(state, out); state.rollback(child.action, rollback); @@ -407,6 +411,14 @@ mod tests { .expect("board FEN must parse") } + fn first_nonfinal_action(board: &Board) -> Action { + *board + .get_valid_actions() + .into_iter() + .find(|action| !action.is_final()) + .expect("test board should have a non-final action") + } + #[test] fn value_from_term_compresses_terminal_scale() { let board = board_with_turn_count(1); @@ -437,6 +449,42 @@ mod tests { assert!(early_value < late_value); } + + #[test] + fn collect_samples_emits_each_node_only_once() { + let mut board = board_with_turn_count(0); + let root_action = first_nonfinal_action(&board); + + let mut child_board = board.clone(); + let _ = child_board.apply_action(root_action); + let child_action = first_nonfinal_action(&child_board); + + let leaf = SearchNode::new(0.25, 0, false); + + let mut child = SearchNode::new(0.5, 0, false); + child.children.push(ChildData { + action: child_action, + child_value: 0.25, + entrance_count: 0, + node: Some(Box::new(leaf)), + }); + + let mut root = SearchNode::new(1.0, 0, false); + root.children.push(ChildData { + action: root_action, + child_value: 0.5, + entrance_count: 0, + node: Some(Box::new(child)), + }); + + let mut first_samples = Vec::new(); + root.collect_samples(&mut board, &mut first_samples); + assert_eq!(first_samples.len(), 2); + + let mut second_samples = Vec::new(); + root.collect_samples(&mut board, &mut second_samples); + assert!(second_samples.is_empty()); + } } pub struct GameSearchTree<'a, E: Evaluator> { @@ -563,7 +611,7 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { /// /// Returns `(board, value)` pairs for nodes with at least one expanded /// child. These are the tree-learning targets for the value network. - pub fn collect_tree_learning_samples(&self) -> Vec { + pub fn collect_tree_learning_samples(&mut self) -> Vec { let mut state = self.root_state.clone(); let mut samples = Vec::new(); self.root_node.collect_samples(&mut state, &mut samples); diff --git a/training/src/lib.rs b/training/src/lib.rs index 8da8f9b..f580f12 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -155,6 +155,14 @@ struct SelfPlay { runner: Arc, } +impl SelfPlay { + fn session(&self) -> PyResult<&SelfPlaySession> { + self.session.as_ref().ok_or_else(|| { + PyErr::new::("session already dropped") + }) + } +} + #[pymethods] impl SelfPlay { #[new] @@ -254,44 +262,65 @@ impl SelfPlay { /// Start self-play with no sample limit. fn start(&self) -> PyResult<()> { - self.session - .as_ref() - .ok_or_else(|| { - PyErr::new::("session already dropped") - })? - .start(); + self.session()?.start(); Ok(()) } /// Block until absolute target_samples is reached, then pause and quiesce. fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { - let session = self.session.as_ref().ok_or_else(|| { - PyErr::new::("session already dropped") - })?; + let session = self.session()?; let result = py.detach(|| session.wait_for(target_samples)); Ok(result) } /// Return the current absolute sample count. fn samples(&self) -> PyResult { - Ok(self - .session - .as_ref() - .ok_or_else(|| { - PyErr::new::("session already dropped") - })? - .samples()) + Ok(self.session()?.samples()) } /// Return the current absolute game count. fn games(&self) -> PyResult { - Ok(self - .session - .as_ref() - .ok_or_else(|| { - PyErr::new::("session already dropped") - })? - .games()) + Ok(self.session()?.games()) + } + + /// Return the total number of selected actions. + fn action_steps(&self) -> PyResult { + Ok(self.session()?.action_steps()) + } + + /// Return the total number of selected final actions. + fn final_actions(&self) -> PyResult { + Ok(self.session()?.final_actions()) + } + + /// Return the total number of selected non-final actions. + fn nonfinal_actions(&self) -> PyResult { + Ok(self.session()?.nonfinal_actions()) + } + + /// Return the total number of completed turns. + fn completed_turns(&self) -> PyResult { + Ok(self.session()?.completed_turns()) + } + + /// Return the sum of `board.turn_count` seen before each selected action. + fn action_turn_count_total(&self) -> PyResult { + Ok(self.session()?.action_turn_count_total()) + } + + /// Return the largest `board.turn_count` seen in any in-progress game. + fn max_turn_count_seen(&self) -> PyResult { + Ok(self.session()?.max_turn_count_seen()) + } + + /// Return the total number of selected actions in completed games. + fn completed_game_actions_total(&self) -> PyResult { + Ok(self.session()?.completed_game_actions_total()) + } + + /// Return the largest number of selected actions in a completed game. + fn max_actions_in_completed_game(&self) -> PyResult { + Ok(self.session()?.max_actions_in_completed_game()) } /// Return the total number of CUDA graph launches completed so far. diff --git a/training/src/training.rs b/training/src/training.rs index 1656ef1..cae010e 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -3,7 +3,7 @@ //! Provides `SelfPlaySession`: a persistent session with pause/resume semantics //! that preserves in-progress game state across boundaries. -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::thread; @@ -15,7 +15,7 @@ use crate::eval::GpuEvaluator; use crate::executor::Executor; use crate::queue::{BatchCompletion, GpuJobQueue}; use crate::replay_buffer::ReplayBuffer; -use crate::worker::{worker_loop_forever, WorkerConfig}; +use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; use crate::BatchDim; /// Shared control state for the persistent self-play session. @@ -30,6 +30,22 @@ struct SessionControl { samples_collected: Arc, /// Total games completed. games_completed: Arc, + /// Total actions selected across all workers. + action_steps: Arc, + /// Total final actions selected across all workers. + final_actions: Arc, + /// Total non-final actions selected across all workers. + nonfinal_actions: Arc, + /// Total turns completed via selected final actions. + completed_turns: Arc, + /// Sum of `board.turn_count` observed before each selected action. + action_turn_count_total: Arc, + /// Largest `board.turn_count` seen in any in-progress game. + max_turn_count_seen: Arc, + /// Total selected actions consumed by completed games. + completed_game_actions_total: Arc, + /// Largest number of selected actions observed in a completed game. + max_actions_in_completed_game: Arc, /// Number of threads currently inside the executor polling loop. active_pollers: AtomicUsize, /// Condvar + mutex for coordinating start/pause/quiesce/shutdown. @@ -45,6 +61,14 @@ impl SessionControl { target_samples: AtomicUsize::new(0), samples_collected: Arc::new(AtomicUsize::new(0)), games_completed: Arc::new(AtomicUsize::new(0)), + action_steps: Arc::new(AtomicU64::new(0)), + final_actions: Arc::new(AtomicU64::new(0)), + nonfinal_actions: Arc::new(AtomicU64::new(0)), + completed_turns: Arc::new(AtomicU64::new(0)), + action_turn_count_total: Arc::new(AtomicU64::new(0)), + max_turn_count_seen: Arc::new(AtomicUsize::new(0)), + completed_game_actions_total: Arc::new(AtomicU64::new(0)), + max_actions_in_completed_game: Arc::new(AtomicU64::new(0)), active_pollers: AtomicUsize::new(0), condvar: Condvar::new(), condvar_mutex: Mutex::new(()), @@ -228,6 +252,50 @@ impl SelfPlaySession { self.control.games_completed.load(Ordering::Acquire) } + /// Return the total number of selected actions. + pub fn action_steps(&self) -> u64 { + self.control.action_steps.load(Ordering::Acquire) + } + + /// Return the total number of selected final actions. + pub fn final_actions(&self) -> u64 { + self.control.final_actions.load(Ordering::Acquire) + } + + /// Return the total number of selected non-final actions. + pub fn nonfinal_actions(&self) -> u64 { + self.control.nonfinal_actions.load(Ordering::Acquire) + } + + /// Return the total number of completed turns. + pub fn completed_turns(&self) -> u64 { + self.control.completed_turns.load(Ordering::Acquire) + } + + /// Return the sum of `board.turn_count` observed before each action. + pub fn action_turn_count_total(&self) -> u64 { + self.control.action_turn_count_total.load(Ordering::Acquire) + } + + /// Return the largest `board.turn_count` seen so far. + pub fn max_turn_count_seen(&self) -> usize { + self.control.max_turn_count_seen.load(Ordering::Acquire) + } + + /// Return the total number of selected actions in completed games. + pub fn completed_game_actions_total(&self) -> u64 { + self.control + .completed_game_actions_total + .load(Ordering::Acquire) + } + + /// Return the largest number of selected actions seen in a completed game. + pub fn max_actions_in_completed_game(&self) -> u64 { + self.control + .max_actions_in_completed_game + .load(Ordering::Acquire) + } + /// Shut down the session. Idempotent. pub fn shutdown(&mut self) { if let Some(threads) = self.threads.take() { @@ -260,14 +328,23 @@ fn session_thread_main( let base_seed = config.seed.wrapping_add(thread_id as u64 * 1000); let evaluator = GpuEvaluator::new(&*queue); - let samples_collected = control.samples_collected.clone(); - let games_completed = control.games_completed.clone(); + let metrics = SelfPlayMetrics { + samples_collected: control.samples_collected.clone(), + games_completed: control.games_completed.clone(), + action_steps: control.action_steps.clone(), + final_actions: control.final_actions.clone(), + nonfinal_actions: control.nonfinal_actions.clone(), + completed_turns: control.completed_turns.clone(), + action_turn_count_total: control.action_turn_count_total.clone(), + max_turn_count_seen: control.max_turn_count_seen.clone(), + completed_game_actions_total: control.completed_game_actions_total.clone(), + max_actions_in_completed_game: control.max_actions_in_completed_game.clone(), + }; let mut futures: Vec + '_>>> = (0 ..config.workers_per_thread) .map(|i| { - let samples_collected = samples_collected.clone(); - let games_completed = games_completed.clone(); + let metrics = metrics.clone(); let mut rng = ChaCha8Rng::seed_from_u64(base_seed + i as u64); let evaluator_ref = &evaluator; let worker_config = &config.worker; @@ -277,8 +354,7 @@ fn session_thread_main( evaluator_ref, worker_config, &mut rng, - samples_collected, - games_completed, + metrics, replay_buffer, ) .await; diff --git a/training/src/worker.rs b/training/src/worker.rs index 4a023ad..69403d4 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -3,7 +3,7 @@ //! Each worker runs Descent search, collects training samples from //! internal tree nodes, and pushes them to the replay buffer. -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; use ndarray::Ix1; @@ -18,6 +18,50 @@ use crate::eval::Evaluator; use crate::observation; use crate::replay_buffer::ReplayBuffer; +fn update_max_usize(max_value: &AtomicUsize, candidate: usize) { + let mut current = max_value.load(Ordering::Acquire); + while candidate > current { + match max_value.compare_exchange_weak( + current, + candidate, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => current = observed, + } + } +} + +fn update_max_u64(max_value: &AtomicU64, candidate: u64) { + let mut current = max_value.load(Ordering::Acquire); + while candidate > current { + match max_value.compare_exchange_weak( + current, + candidate, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => current = observed, + } + } +} + +#[derive(Clone)] +pub struct SelfPlayMetrics { + pub samples_collected: Arc, + pub games_completed: Arc, + pub action_steps: Arc, + pub final_actions: Arc, + pub nonfinal_actions: Arc, + pub completed_turns: Arc, + pub action_turn_count_total: Arc, + pub max_turn_count_seen: Arc, + pub completed_game_actions_total: Arc, + pub max_actions_in_completed_game: Arc, +} + /// Configuration for the worker. #[derive(Clone)] pub struct WorkerConfig { @@ -46,7 +90,7 @@ async fn play_game( config: &WorkerConfig, rng: &mut R, replay_buffer: &ReplayBuffer, - samples_collected: &AtomicUsize, + metrics: &SelfPlayMetrics, ) { let board = Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) @@ -61,8 +105,15 @@ async fn play_game( let tree_rng = SmallRng::from_rng(rng); let mut tree = GameSearchTree::new(&board, evaluator, tree_rng).await; + let mut action_steps_in_game = 0u64; loop { + let current_turn_count = tree.root_state.turn_count; + metrics + .action_turn_count_total + .fetch_add(current_turn_count as u64, Ordering::AcqRel); + update_max_usize(&metrics.max_turn_count_seen, current_turn_count); + // Run descent search tree.run_descent_for_iter(config.descent_iterations).await; @@ -80,25 +131,44 @@ async fn play_game( ); }); } - samples_collected.fetch_add(num_samples, Ordering::AcqRel); + metrics + .samples_collected + .fetch_add(num_samples, Ordering::AcqRel); } // Select action via ordinal distribution let action_id = tree.ordinal_select(); let action = tree.root_node.children[action_id].action; + action_steps_in_game += 1; + metrics.action_steps.fetch_add(1, Ordering::AcqRel); + if action.is_final() { + metrics.final_actions.fetch_add(1, Ordering::AcqRel); + metrics.completed_turns.fetch_add(1, Ordering::AcqRel); + } else { + metrics.nonfinal_actions.fetch_add(1, Ordering::AcqRel); + } // Apply action to get new board state let mut new_board = tree.root_state.clone(); let (outcome, _) = new_board.apply_action(action); + update_max_usize(&metrics.max_turn_count_seen, new_board.turn_count); // Check if game is over match outcome { alpha_paint::board::ApplyActionOutcome::Terminal { .. } | alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { + metrics + .completed_game_actions_total + .fetch_add(action_steps_in_game, Ordering::AcqRel); + update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); break; } alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { // Terminal via play-instead + metrics + .completed_game_actions_total + .fetch_add(action_steps_in_game, Ordering::AcqRel); + update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { @@ -121,12 +191,11 @@ pub async fn worker_loop_forever( evaluator: &E, config: &WorkerConfig, rng: &mut R, - samples_collected: Arc, - games_completed: Arc, + metrics: SelfPlayMetrics, replay_buffer: &ReplayBuffer, ) { loop { - play_game(evaluator, config, rng, replay_buffer, &samples_collected).await; - games_completed.fetch_add(1, Ordering::AcqRel); + play_game(evaluator, config, rng, replay_buffer, &metrics).await; + metrics.games_completed.fetch_add(1, Ordering::AcqRel); } } From 0143afa589082161c71bc0fc8a3ece38455ea489 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 17:52:40 -0400 Subject: [PATCH 19/59] push nodes we are pruning to replay buffer --- python/alphapaint_training/train.py | 23 +++-- tooling/src/main.rs | 2 +- training/src/descent.rs | 129 ++++++++++++++++++++-------- training/src/worker.rs | 71 +++++++++------ 4 files changed, 151 insertions(+), 74 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 42cdd6a..a1538ca 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -73,7 +73,7 @@ def train_step( batch_size: int, seed: int, device: torch.device, -) -> float: +) -> torch.Tensor: model.train() obs, target = _sample_replay_batch(replay_buffer, batch_size, seed, device) @@ -88,7 +88,7 @@ def train_step( loss = F.mse_loss(pred.float(), target) loss.backward() optimizer.step() - return float(loss.detach().cpu()) + return loss.detach() def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: @@ -218,7 +218,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: ) actions_per_game_display = f"{actions_per_game:.1f}" if games_added else "-" - round_losses = [] + round_losses: list[torch.Tensor] = [] train_started_at = time.perf_counter() for step_idx in range(config.train_steps_per_round): loss = train_step( @@ -233,12 +233,14 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: model.eval() train_seconds = time.perf_counter() - train_started_at - losses.extend(round_losses) replay_size = len(replay_buffer) - mean_loss = sum(round_losses) / len(round_losses) if round_losses else 0.0 - min_loss = min(round_losses) if round_losses else 0.0 - max_loss = max(round_losses) if round_losses else 0.0 - last_loss = round_losses[-1] if round_losses else 0.0 + if round_losses: + round_loss_values = torch.stack(round_losses).float().cpu() + loss_values = round_loss_values.tolist() + losses.extend(loss_values) + mean_loss = float(round_loss_values.mean().item()) + else: + mean_loss = 0.0 record = { "round": round_number, "samples_total": collected, @@ -276,9 +278,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_steps_per_second": config.train_steps_per_round / max(train_seconds, 1e-9), "loss_mean": mean_loss, - "loss_min": min_loss, - "loss_max": max_loss, - "loss_last": last_loss, "timestamp": time.time(), } print( @@ -287,7 +286,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " - f"loss={mean_loss:.6f}/{last_loss:.6f}" + f"loss={mean_loss:.6f}" ) wandb.log(record) diff --git a/tooling/src/main.rs b/tooling/src/main.rs index 2f01c09..62aa592 100644 --- a/tooling/src/main.rs +++ b/tooling/src/main.rs @@ -2,4 +2,4 @@ use alpha_paint::perft_test; fn main() { perft_test() -} \ No newline at end of file +} diff --git a/training/src/descent.rs b/training/src/descent.rs index 9bcad41..d6848ff 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -40,7 +40,6 @@ pub struct SearchNode { pub value: f32, pub completion_value: i32, resolved: bool, - replay_emitted: bool, pub children: Vec, } @@ -50,7 +49,6 @@ impl SearchNode { value, completion_value, resolved: is_resolved, - replay_emitted: false, children: vec![], } } @@ -169,7 +167,6 @@ impl SearchNode { value: term_value, completion_value: comp_value, resolved: true, - replay_emitted: false, children: vec![ChildData { action, child_value: term_value, @@ -379,18 +376,43 @@ impl SearchNode { /// An internal node is one that has children and at least one expanded child. /// Non-terminal leaf nodes (where the network estimate was used without /// minimax backing) are excluded per Athénan's tree learning rules. - fn collect_samples(&mut self, state: &mut Board, out: &mut Vec) { + fn collect_samples(&self, state: &mut Board, out: &mut Vec) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); - if has_expanded_child && !self.replay_emitted { + if has_expanded_child { out.push(TreeLearningSample { board: state.clone(), value: self.value, }); - self.replay_emitted = true; } - for child in &mut self.children { - if let Some(node) = child.node.as_mut() { + for child in &self.children { + if let Some(node) = child.node.as_ref() { + let (_, rollback) = state.apply_action(child.action); + node.collect_samples(state, out); + state.rollback(child.action, rollback); + } + } + } + + fn collect_samples_excluding_child( + &self, + state: &mut Board, + excluded_child: usize, + out: &mut Vec, + ) { + let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); + if has_expanded_child { + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + }); + } + + for (child_idx, child) in self.children.iter().enumerate() { + if child_idx == excluded_child { + continue; + } + if let Some(node) = child.node.as_ref() { let (_, rollback) = state.apply_action(child.action); node.collect_samples(state, out); state.rollback(child.action, rollback); @@ -451,39 +473,65 @@ mod tests { } #[test] - fn collect_samples_emits_each_node_only_once() { + fn collect_samples_excluding_child_keeps_root_and_skips_reused_subtree() { let mut board = board_with_turn_count(0); - let root_action = first_nonfinal_action(&board); + let root_actions: Vec = board + .get_valid_actions() + .into_iter() + .filter(|action| !action.is_final()) + .take(2) + .copied() + .collect(); + assert_eq!(root_actions.len(), 2); - let mut child_board = board.clone(); - let _ = child_board.apply_action(root_action); - let child_action = first_nonfinal_action(&child_board); + let mut reused_board = board.clone(); + let _ = reused_board.apply_action(root_actions[0]); + let reused_action = first_nonfinal_action(&reused_board); + + let mut dropped_board = board.clone(); + let _ = dropped_board.apply_action(root_actions[1]); + let dropped_action = first_nonfinal_action(&dropped_board); let leaf = SearchNode::new(0.25, 0, false); - let mut child = SearchNode::new(0.5, 0, false); - child.children.push(ChildData { - action: child_action, + let mut reused_child = SearchNode::new(0.5, 0, false); + reused_child.children.push(ChildData { + action: reused_action, child_value: 0.25, entrance_count: 0, + node: Some(Box::new(SearchNode::new(0.25, 0, false))), + }); + + let mut dropped_child = SearchNode::new(-0.5, 0, false); + dropped_child.children.push(ChildData { + action: dropped_action, + child_value: -0.25, + entrance_count: 0, node: Some(Box::new(leaf)), }); - let mut root = SearchNode::new(1.0, 0, false); + let root = SearchNode::new(1.0, 0, false); + let mut root = root; root.children.push(ChildData { - action: root_action, + action: root_actions[0], child_value: 0.5, entrance_count: 0, - node: Some(Box::new(child)), + node: Some(Box::new(reused_child)), + }); + root.children.push(ChildData { + action: root_actions[1], + child_value: -0.5, + entrance_count: 0, + node: Some(Box::new(dropped_child)), }); - let mut first_samples = Vec::new(); - root.collect_samples(&mut board, &mut first_samples); - assert_eq!(first_samples.len(), 2); + let mut full_samples = Vec::new(); + root.collect_samples(&mut board, &mut full_samples); + assert_eq!(full_samples.len(), 3); - let mut second_samples = Vec::new(); - root.collect_samples(&mut board, &mut second_samples); - assert!(second_samples.is_empty()); + let mut dropped_samples = Vec::new(); + root.collect_samples_excluding_child(&mut board, 0, &mut dropped_samples); + assert_eq!(dropped_samples.len(), 2); } } @@ -552,7 +600,23 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { } } - pub fn step_tree(&mut self, new_board: &Board, action_id: usize) { + pub fn collect_tree_learning_samples(&self) -> Vec { + let mut state = self.root_state.clone(); + let mut samples = Vec::new(); + self.root_node.collect_samples(&mut state, &mut samples); + samples + } + + pub fn step_tree_and_collect_dropped_samples( + &mut self, + new_board: &Board, + action_id: usize, + ) -> Vec { + let mut state = self.root_state.clone(); + let mut samples = Vec::new(); + self.root_node + .collect_samples_excluding_child(&mut state, action_id, &mut samples); + self.root_state = new_board.clone(); if self.root_node.children[action_id].node.is_some() { @@ -562,6 +626,8 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { // This shouldn't happen in normal play since we always expand before selecting. panic!("step_tree called on unexpanded child"); } + + samples } pub async fn run_descent_for_iter(&mut self, iterations: u32) { @@ -607,17 +673,6 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { self.safest_action() } - /// Collect tree learning samples from all internal nodes. - /// - /// Returns `(board, value)` pairs for nodes with at least one expanded - /// child. These are the tree-learning targets for the value network. - pub fn collect_tree_learning_samples(&mut self) -> Vec { - let mut state = self.root_state.clone(); - let mut samples = Vec::new(); - self.root_node.collect_samples(&mut state, &mut samples); - samples - } - /// Select an action using Athénan's ordinal distribution. /// /// The ordinal distribution depends only on rank ordering of moves. diff --git a/training/src/worker.rs b/training/src/worker.rs index 69403d4..66ab17b 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -3,17 +3,18 @@ //! Each worker runs Descent search, collects training samples from //! internal tree nodes, and pushes them to the replay buffer. -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; -use alpha_paint::board::Board; use alpha_paint::TRAINING_START_FENS; +use alpha_paint::board::Board; use crate::descent::GameSearchTree; +use crate::descent::TreeLearningSample; use crate::eval::Evaluator; use crate::observation; use crate::replay_buffer::ReplayBuffer; @@ -77,12 +78,37 @@ impl Default for WorkerConfig { } } +fn push_samples_to_replay( + replay_buffer: &ReplayBuffer, + metrics: &SelfPlayMetrics, + samples: Vec, +) { + let num_samples = samples.len(); + if num_samples == 0 { + return; + } + + let mut guard = replay_buffer.reserve(num_samples); + for sample in samples { + guard.push_with_observation(sample.value, |mut out| { + observation::encode_into_slice( + &sample.board, + out.as_slice_mut() + .expect("replay observation slot must be contiguous"), + ); + }); + } + metrics + .samples_collected + .fetch_add(num_samples, Ordering::AcqRel); +} + /// Run a single self-play game with tree learning. /// /// At each move: /// 1. Run Descent search -/// 2. Collect (value) training samples from all internal tree nodes -/// 3. Select action via ordinal distribution +/// 2. Select action via ordinal distribution +/// 3. Collect the dropped portion of the tree with its latest values /// 4. Apply action, reuse tree /// async fn play_game( @@ -117,25 +143,6 @@ async fn play_game( // Run descent search tree.run_descent_for_iter(config.descent_iterations).await; - // Collect tree learning samples from internal nodes - let samples = tree.collect_tree_learning_samples(); - let num_samples = samples.len(); - if num_samples > 0 { - let mut guard = replay_buffer.reserve(num_samples); - for sample in samples { - guard.push_with_observation(sample.value, |mut out| { - observation::encode_into_slice( - &sample.board, - out.as_slice_mut() - .expect("replay observation slot must be contiguous"), - ); - }); - } - metrics - .samples_collected - .fetch_add(num_samples, Ordering::AcqRel); - } - // Select action via ordinal distribution let action_id = tree.ordinal_select(); let action = tree.root_node.children[action_id].action; @@ -157,6 +164,11 @@ async fn play_game( match outcome { alpha_paint::board::ApplyActionOutcome::Terminal { .. } | alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { + push_samples_to_replay( + replay_buffer, + metrics, + tree.collect_tree_learning_samples(), + ); metrics .completed_game_actions_total .fetch_add(action_steps_in_game, Ordering::AcqRel); @@ -165,6 +177,11 @@ async fn play_game( } alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { // Terminal via play-instead + push_samples_to_replay( + replay_buffer, + metrics, + tree.collect_tree_learning_samples(), + ); metrics .completed_game_actions_total .fetch_add(action_steps_in_game, Ordering::AcqRel); @@ -173,8 +190,14 @@ async fn play_game( } alpha_paint::board::ApplyActionOutcome::Ongoing => { if tree.root_node.children[action_id].node.is_some() { - tree.step_tree(&new_board, action_id); + let samples = tree.step_tree_and_collect_dropped_samples(&new_board, action_id); + push_samples_to_replay(replay_buffer, metrics, samples); } else { + push_samples_to_replay( + replay_buffer, + metrics, + tree.collect_tree_learning_samples(), + ); let next_rng = SmallRng::from_rng(rng); tree = GameSearchTree::new(&new_board, evaluator, next_rng).await; } From 36120b1866d1dbcabcb08cdd8ccfba92cb730b85 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 17:52:52 -0400 Subject: [PATCH 20/59] cargo fmt --- alpha_paint/src/bindings.rs | 26 +++--- alpha_paint/src/board/action_generation.rs | 2 +- alpha_paint/src/board/actions.rs | 8 +- alpha_paint/src/board/board_impl.rs | 93 ++++++++++++++-------- alpha_paint/src/board/tile_map.rs | 8 +- alpha_paint/src/evaluation.rs | 2 +- training/src/worker.rs | 4 +- 7 files changed, 94 insertions(+), 49 deletions(-) diff --git a/alpha_paint/src/bindings.rs b/alpha_paint/src/bindings.rs index 4f1c078..f4b51cb 100644 --- a/alpha_paint/src/bindings.rs +++ b/alpha_paint/src/bindings.rs @@ -1,5 +1,5 @@ use pyo3::{pyclass, pymethods}; -use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng}; +use rand::{RngExt, SeedableRng, rngs::StdRng, seq::SliceRandom}; use std::sync::Arc; use std::time::Duration; @@ -147,9 +147,14 @@ impl PyBoard { fn perft(&self, max_depth: usize) -> Vec { let mut all_turns = Vec::new(); let mut move_stack = Vec::new(); - enumerate_single_turns(&mut self.0.clone(), max_depth, &mut move_stack, &mut |turn| { - all_turns.push(turn); - }); + enumerate_single_turns( + &mut self.0.clone(), + max_depth, + &mut move_stack, + &mut |turn| { + all_turns.push(turn); + }, + ); all_turns } @@ -161,9 +166,14 @@ impl PyBoard { let mut move_stack = Vec::new(); let mut sampler = ReservoirSampler::new(max_samples, seed); - enumerate_single_turns(&mut self.0.clone(), max_depth, &mut move_stack, &mut |turn| { - sampler.record(turn); - }); + enumerate_single_turns( + &mut self.0.clone(), + max_depth, + &mut move_stack, + &mut |turn| { + sampler.record(turn); + }, + ); sampler.finish() } @@ -269,8 +279,6 @@ impl PyBoard { } } } - - } fn enumerate_single_turns( diff --git a/alpha_paint/src/board/action_generation.rs b/alpha_paint/src/board/action_generation.rs index 738f945..f3e4e1d 100644 --- a/alpha_paint/src/board/action_generation.rs +++ b/alpha_paint/src/board/action_generation.rs @@ -106,7 +106,7 @@ impl Board { self.black_stamina }; - for &beacon_target in self.tiles.get_beacon_iterator::() { + for &beacon_target in self.tiles.get_beacon_iterator::() { if stamina >= move_cost + EXTRA_MOVE_COST { actions.push(Action::Move(Move::new( beacon_target, diff --git a/alpha_paint/src/board/actions.rs b/alpha_paint/src/board/actions.rs index 619cf9c..9e19544 100644 --- a/alpha_paint/src/board/actions.rs +++ b/alpha_paint/src/board/actions.rs @@ -135,12 +135,14 @@ pub struct ActionList { impl ActionList { pub fn new() -> Self { Self { - local_buf: [Action::Paint(Paint{ target: Coordinate::new(0, 0) }); 20], + local_buf: [Action::Paint(Paint { + target: Coordinate::new(0, 0), + }); 20], overflow: None, cnt: 0, } } - + pub fn len(&self) -> usize { self.cnt } @@ -186,4 +188,4 @@ impl<'a> IntoIterator for &'a ActionList { index: 0, } } -} \ No newline at end of file +} diff --git a/alpha_paint/src/board/board_impl.rs b/alpha_paint/src/board/board_impl.rs index 04df5c8..1fa4e58 100644 --- a/alpha_paint/src/board/board_impl.rs +++ b/alpha_paint/src/board/board_impl.rs @@ -30,12 +30,18 @@ impl Rollback { original_coord: Coordinate, original_hill_owner: Option, white_stamina: u16, - black_stamina: u16 + black_stamina: u16, ) -> Self { Self { - original_tile, original_move_count, event_pointer, - consumed_p1: false, consumed_p2: false, - original_hill_owner, original_coord, white_stamina, black_stamina + original_tile, + original_move_count, + event_pointer, + consumed_p1: false, + consumed_p2: false, + original_hill_owner, + original_coord, + white_stamina, + black_stamina, } } } @@ -91,7 +97,7 @@ pub struct Board { pub beacon_tile_rollback_stack: Vec<(Tile, Option)>, /// Internally used for rollbacks. - /// Stores a list of (usize, coordinates) for all powerups that were placed at the start + /// Stores a list of (usize, coordinates) for all powerups that were placed at the start /// of the n-th turn (where n is the first number in the tuple) pub powerup_rollback_stack: Vec<(usize, Coordinate)>, @@ -128,52 +134,71 @@ impl Board { self.apply_action_generic::(action) } } - + pub fn rollback(&mut self, action: Action, rollback: Rollback) { // Undo the effects of end-turn (placing powerups) if action.is_final() { - while self.powerup_rollback_stack.last().is_some_and(|&(i, _)| {i == self.turn_count}) { + while self + .powerup_rollback_stack + .last() + .is_some_and(|&(i, _)| i == self.turn_count) + { self.powerups[self.powerup_rollback_stack.pop().unwrap().1] = false; } - + self.turn_count -= 1; } match action { - Action::Move(Move { target, kind: _, place_beacon }) | - Action::FinalMove(Move { target, kind: _, place_beacon }) => { - if rollback.consumed_p2 { self.powerups[target] = true; } - if rollback.consumed_p1 { self.powerups[rollback.original_coord] = true; } + Action::Move(Move { + target, + kind: _, + place_beacon, + }) + | Action::FinalMove(Move { + target, + kind: _, + place_beacon, + }) => { + if rollback.consumed_p2 { + self.powerups[target] = true; + } + if rollback.consumed_p1 { + self.powerups[rollback.original_coord] = true; + } if place_beacon { for &coord in target.region_3x3_with_self().iter().rev() { if !self.in_bounds(coord) || self.tiles[coord].is_wall() { continue; } - let (original_tile, hill_owner) = self.beacon_tile_rollback_stack.pop().unwrap(); + let (original_tile, hill_owner) = + self.beacon_tile_rollback_stack.pop().unwrap(); self.tiles.set(coord, original_tile); self.tiles.override_hill_owner(coord, hill_owner); } } - + self.tiles.set(target, rollback.original_tile); - self.tiles.override_hill_owner(target, rollback.original_hill_owner); + self.tiles + .override_hill_owner(target, rollback.original_hill_owner); } - - Action::Paint(Paint{ target }) | Action::FinalPaint(Paint{ target }) => { + + Action::Paint(Paint { target }) | Action::FinalPaint(Paint { target }) => { self.tiles.set(target, rollback.original_tile); - self.tiles.override_hill_owner(target, rollback.original_hill_owner); + self.tiles + .override_hill_owner(target, rollback.original_hill_owner); } } - + // Now do all the normal rollback things that apply to all moves self.white_stamina = rollback.white_stamina as usize; self.black_stamina = rollback.black_stamina as usize; self.consecutives_moves_so_far = rollback.original_move_count as usize; self.event_pointer = rollback.event_pointer as usize; - - if self.is_white_turn() { - self.white_coord = rollback.original_coord; + + if self.is_white_turn() { + self.white_coord = rollback.original_coord; } else { self.black_coord = rollback.original_coord; } @@ -252,11 +277,13 @@ impl Board { let player_coord = self.player_coord::(); let tile = self.tiles[target]; let rollback = Rollback::new( - tile, self.consecutives_moves_so_far as u8, + tile, + self.consecutives_moves_so_far as u8, self.event_pointer as u16, player_coord, self.tiles.get_hill_owner(target), - self.white_stamina as u16, self.black_stamina as u16 + self.white_stamina as u16, + self.black_stamina as u16, ); let stamina = if IS_WHITE { @@ -303,10 +330,13 @@ impl Board { } = mv; let is_collision = target == opponent_coord; let mut rollback = Rollback::new( - self.tiles[target], self.consecutives_moves_so_far as u8, + self.tiles[target], + self.consecutives_moves_so_far as u8, self.event_pointer as u16, - *player_coord, self.tiles.get_hill_owner(target), - self.white_stamina as u16, self.black_stamina as u16 + *player_coord, + self.tiles.get_hill_owner(target), + self.white_stamina as u16, + self.black_stamina as u16, ); let stamina = if IS_WHITE { &mut self.white_stamina @@ -314,7 +344,6 @@ impl Board { &mut self.black_stamina }; - // make sure we can reach it debug_assert!( target.manhattan_dist(*player_coord) == 1 @@ -410,7 +439,8 @@ impl Board { if !self.in_bounds(coord) || self.tiles[coord].is_wall() { continue; } - self.beacon_tile_rollback_stack.push((self.tiles[coord], self.tiles.get_hill_owner(coord))); + self.beacon_tile_rollback_stack + .push((self.tiles[coord], self.tiles.get_hill_owner(coord))); self.tiles.apply_beacon_placement_effect::(coord); } @@ -445,8 +475,9 @@ impl Board { } if !self.powerups[powerup.coord] { self.powerups[powerup.coord] = true; - self.powerup_rollback_stack.push((self.turn_count, powerup.coord)); - } + self.powerup_rollback_stack + .push((self.turn_count, powerup.coord)); + } self.event_pointer += 1; } if IS_WHITE { diff --git a/alpha_paint/src/board/tile_map.rs b/alpha_paint/src/board/tile_map.rs index 246a001..eacdf50 100644 --- a/alpha_paint/src/board/tile_map.rs +++ b/alpha_paint/src/board/tile_map.rs @@ -77,7 +77,9 @@ impl TileMap { /// Returns None if the hill is unowned or if no hill is present here pub fn get_hill_owner(&self, coord: Coordinate) -> Option { let hill_id = self.hill_id[coord]; - if hill_id == u16::MAX { return None; } + if hill_id == u16::MAX { + return None; + } self.hill_metadata[hill_id as usize].owner } @@ -90,7 +92,9 @@ impl TileMap { /// should own this hill. It is ___highly recommended___ that the set command be used instead. pub fn override_hill_owner(&mut self, coord: Coordinate, hill_owner: Option) { let hill_id = self.hill_id[coord]; - if hill_id == u16::MAX { return; } + if hill_id == u16::MAX { + return; + } self.hill_metadata[hill_id as usize].owner = hill_owner; } diff --git a/alpha_paint/src/evaluation.rs b/alpha_paint/src/evaluation.rs index 9900b32..5e71b67 100644 --- a/alpha_paint/src/evaluation.rs +++ b/alpha_paint/src/evaluation.rs @@ -1,5 +1,5 @@ -use crate::board::board_structs::Player; use crate::board::Board; +use crate::board::board_structs::Player; use std::cmp::min; pub struct Evaluator; diff --git a/training/src/worker.rs b/training/src/worker.rs index 66ab17b..dd3399b 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -3,15 +3,15 @@ //! Each worker runs Descent search, collects training samples from //! internal tree nodes, and pushes them to the replay buffer. -use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::Arc; use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; -use alpha_paint::TRAINING_START_FENS; use alpha_paint::board::Board; +use alpha_paint::TRAINING_START_FENS; use crate::descent::GameSearchTree; use crate::descent::TreeLearningSample; From 499fd97acd2dbc0eefdc226fdd3581a64afbff7a Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 18:18:22 -0400 Subject: [PATCH 21/59] limit iteration count by gpu evaluations --- python/alphapaint_training/train.py | 8 +++--- training/src/cudagraph.rs | 2 +- training/src/descent.rs | 42 ++++++++++++++++++++++++++++- training/src/eval.rs | 39 +++++++++++++++++++++++++++ training/src/executor.rs | 2 +- training/src/future.rs | 2 +- training/src/lib.rs | 16 +++++------ training/src/observation.rs | 2 +- training/src/queue.rs | 2 +- training/src/replay_buffer.rs | 2 +- training/src/training.rs | 4 +-- training/src/worker.rs | 20 +++++++------- 12 files changed, 110 insertions(+), 31 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index a1538ca..9d743e8 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -25,7 +25,7 @@ class TrainConfig: replay_capacity: int = 16_000_000 num_threads: int = 32 workers_per_thread: int = 16 - descent_iterations: int = 10 + max_gpu_evals_per_move: int = 64 * 1024 lr: float = 3e-4 seed: int = 42 selfplay_precision: str = "bf16" @@ -151,7 +151,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: config.num_threads, config.workers_per_thread, config.seed, - descent_iterations=config.descent_iterations, + max_gpu_evals_per_move=config.max_gpu_evals_per_move, model=model, selfplay_precision=config.selfplay_precision, ) @@ -340,7 +340,9 @@ def _parse_args() -> TrainConfig: "--workers-per-thread", type=int, default=defaults.workers_per_thread ) parser.add_argument( - "--descent-iterations", type=int, default=defaults.descent_iterations + "--max-gpu-evals-per-move", + type=int, + default=defaults.max_gpu_evals_per_move, ) parser.add_argument("--lr", type=float, default=defaults.lr) parser.add_argument("--seed", type=int, default=defaults.seed) diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index 730f6fe..f0144cc 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -2,7 +2,7 @@ //! //! No policy head - the model outputs only a scalar value per position. -use std::ffi::{c_void, CStr}; +use std::ffi::{CStr, c_void}; use std::mem::size_of; use std::slice; use std::sync::atomic::{AtomicU64, Ordering}; diff --git a/training/src/descent.rs b/training/src/descent.rs index d6848ff..a4d4bde 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -9,7 +9,7 @@ use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use rand::rngs::SmallRng; use rand::{Rng, RngExt}; -use crate::eval::Evaluator; +use crate::eval::{EvalCountTracker, Evaluator}; const AVG_GAME_LENGTH: f32 = 500.0; @@ -724,3 +724,43 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { *indices.last().unwrap() } } + +impl<'a, E: EvalCountTracker> GameSearchTree<'a, E> { + pub async fn run_descent_to_eval_limit(&mut self, max_gpu_evals: u64) { + let dcv: i32 = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; + + if self.root_node.children.len() <= 1 { + return; + } + + self.evaluator.reset_eval_count(); + + loop { + if self.evaluator.eval_count() >= max_gpu_evals { + break; + } + + let ba = self.get_best_action_index(); + if self.root_node.children[ba].entrance_count >= 6000 + && self.root_node.children[ba].completion_value() != -dcv + { + break; + } + if self.root_node.children[ba].completion_value() == dcv { + break; + } + self.root_node + .ubfms_iteration( + self.root_state.clone(), + ApplyActionOutcome::Ongoing, + self.evaluator, + &mut self.rng, + ) + .await; + } + } +} diff --git a/training/src/eval.rs b/training/src/eval.rs index 894e194..626d93f 100644 --- a/training/src/eval.rs +++ b/training/src/eval.rs @@ -1,6 +1,7 @@ //! Async evaluator for GPU inference. Value-only (no policy head). use std::future::Future; +use std::sync::atomic::{AtomicU64, Ordering}; use alpha_paint::board::Board; use ndarray::Ix1; @@ -17,6 +18,12 @@ pub trait Evaluator { fn evaluate(&self, board: Board) -> impl Future; } +/// Evaluators that can report how many GPU eval submissions they have issued. +pub trait EvalCountTracker: Evaluator { + fn reset_eval_count(&self); + fn eval_count(&self) -> u64; +} + /// GPU-backed evaluator that batches inference requests. /// /// Wraps a GpuJobQueue and serializes Board state into observations @@ -45,6 +52,38 @@ impl Evaluator for GpuEvaluator<'_> { } } +/// Wrapper that counts evaluator submissions for a single worker. +pub struct CountingEvaluator<'a, E> { + inner: &'a E, + eval_count: AtomicU64, +} + +impl<'a, E> CountingEvaluator<'a, E> { + pub fn new(inner: &'a E) -> Self { + Self { + inner, + eval_count: AtomicU64::new(0), + } + } +} + +impl Evaluator for CountingEvaluator<'_, E> { + fn evaluate(&self, board: Board) -> impl Future { + self.eval_count.fetch_add(1, Ordering::AcqRel); + self.inner.evaluate(board) + } +} + +impl EvalCountTracker for CountingEvaluator<'_, E> { + fn reset_eval_count(&self) { + self.eval_count.store(0, Ordering::Release); + } + + fn eval_count(&self) -> u64 { + self.eval_count.load(Ordering::Acquire) + } +} + /// Synchronous CPU evaluator for testing. /// /// Returns zero value for all positions. diff --git a/training/src/executor.rs b/training/src/executor.rs index 5f29537..74d0d8d 100644 --- a/training/src/executor.rs +++ b/training/src/executor.rs @@ -216,8 +216,8 @@ mod tests { #[test] fn test_executor_cancels_pending_futures() { - use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; let cancelled = Arc::new(AtomicBool::new(false)); let cancelled_clone = cancelled.clone(); diff --git a/training/src/future.rs b/training/src/future.rs index b9e3bee..f4bf839 100644 --- a/training/src/future.rs +++ b/training/src/future.rs @@ -10,8 +10,8 @@ use std::task::{Context, Poll}; use ndarray::Dimension; -use crate::queue::GpuJobQueue; use crate::BatchDim; +use crate::queue::GpuJobQueue; // Thread-local flag for tracking whether any future made progress. // Used by the executor to decide whether to park. diff --git a/training/src/lib.rs b/training/src/lib.rs index f580f12..9ab2233 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use std::sync::{Mutex, OnceLock}; use cudagraph::CudaGraphRunner; -use queue::{queue_shape_for_workers, BATCH_SIZE}; +use queue::{BATCH_SIZE, queue_shape_for_workers}; use replay_buffer::ReplayBuffer; use training::{SelfPlaySession, SessionConfig}; use worker::WorkerConfig; @@ -172,7 +172,7 @@ impl SelfPlay { workers_per_thread, seed, *, - descent_iterations = 30, + max_gpu_evals_per_move = 65536, model, selfplay_precision = "bf16" ))] @@ -182,21 +182,17 @@ impl SelfPlay { num_threads: usize, workers_per_thread: usize, seed: u64, - descent_iterations: u32, + max_gpu_evals_per_move: u64, model: Py, selfplay_precision: &str, ) -> PyResult { - if descent_iterations == 0 { - return Err(PyErr::new::( - "descent_iterations must be >= 1", - )); - } - let config = SessionConfig { num_threads, workers_per_thread, seed, - worker: WorkerConfig { descent_iterations }, + worker: WorkerConfig { + max_gpu_evals_per_move, + }, }; let total_workers = num_threads.checked_mul(workers_per_thread).ok_or_else(|| { diff --git a/training/src/observation.rs b/training/src/observation.rs index fb25fd4..e60d71e 100644 --- a/training/src/observation.rs +++ b/training/src/observation.rs @@ -1,6 +1,6 @@ +use alpha_paint::board::Board; use alpha_paint::board::board_structs::Player; use alpha_paint::board::structs::Coordinate; -use alpha_paint::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; diff --git a/training/src/queue.rs b/training/src/queue.rs index 26a0799..81eceaa 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -8,8 +8,8 @@ //! This enables zero-copy batch slicing for GPU dispatch. use std::cell::UnsafeCell; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use event_listener::Event; use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs index 88ac817..b9d511a 100644 --- a/training/src/replay_buffer.rs +++ b/training/src/replay_buffer.rs @@ -4,8 +4,8 @@ use std::cell::UnsafeCell; use std::sync::atomic::{AtomicU64, Ordering}; use ndarray::{Array, ArrayViewMut, Axis}; -use rand::seq::index::sample; use rand::Rng; +use rand::seq::index::sample; use crate::BatchDim; diff --git a/training/src/training.rs b/training/src/training.rs index cae010e..2205e76 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -11,12 +11,12 @@ use ndarray::{ArrayView, Ix1, Ix2}; use rand::SeedableRng; use rand_chacha::ChaCha8Rng; +use crate::BatchDim; use crate::eval::GpuEvaluator; use crate::executor::Executor; use crate::queue::{BatchCompletion, GpuJobQueue}; use crate::replay_buffer::ReplayBuffer; -use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; -use crate::BatchDim; +use crate::worker::{SelfPlayMetrics, WorkerConfig, worker_loop_forever}; /// Shared control state for the persistent self-play session. struct SessionControl { diff --git a/training/src/worker.rs b/training/src/worker.rs index dd3399b..6881cdb 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -3,19 +3,19 @@ //! Each worker runs Descent search, collects training samples from //! internal tree nodes, and pushes them to the replay buffer. -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; -use alpha_paint::board::Board; use alpha_paint::TRAINING_START_FENS; +use alpha_paint::board::Board; use crate::descent::GameSearchTree; use crate::descent::TreeLearningSample; -use crate::eval::Evaluator; +use crate::eval::{CountingEvaluator, Evaluator}; use crate::observation; use crate::replay_buffer::ReplayBuffer; @@ -66,14 +66,14 @@ pub struct SelfPlayMetrics { /// Configuration for the worker. #[derive(Clone)] pub struct WorkerConfig { - /// Number of UBFM/Descent iterations per move. - pub descent_iterations: u32, + /// Maximum GPU eval submissions spent on descent for a single move. + pub max_gpu_evals_per_move: u64, } impl Default for WorkerConfig { fn default() -> Self { Self { - descent_iterations: 30, + max_gpu_evals_per_move: 64 * 1024, } } } @@ -129,8 +129,9 @@ async fn play_game( return; } + let counting_evaluator = CountingEvaluator::new(evaluator); let tree_rng = SmallRng::from_rng(rng); - let mut tree = GameSearchTree::new(&board, evaluator, tree_rng).await; + let mut tree = GameSearchTree::new(&board, &counting_evaluator, tree_rng).await; let mut action_steps_in_game = 0u64; loop { @@ -141,7 +142,8 @@ async fn play_game( update_max_usize(&metrics.max_turn_count_seen, current_turn_count); // Run descent search - tree.run_descent_for_iter(config.descent_iterations).await; + tree.run_descent_to_eval_limit(config.max_gpu_evals_per_move) + .await; // Select action via ordinal distribution let action_id = tree.ordinal_select(); @@ -199,7 +201,7 @@ async fn play_game( tree.collect_tree_learning_samples(), ); let next_rng = SmallRng::from_rng(rng); - tree = GameSearchTree::new(&new_board, evaluator, next_rng).await; + tree = GameSearchTree::new(&new_board, &counting_evaluator, next_rng).await; } } } From 91e6060d51bd424a83d3cd0628e0d0dbf6ac1dea Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 18:33:07 -0400 Subject: [PATCH 22/59] const updates --- python/alphapaint_training/train.py | 2 +- training/src/lib.rs | 2 +- training/src/worker.rs | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 9d743e8..432787d 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -25,7 +25,7 @@ class TrainConfig: replay_capacity: int = 16_000_000 num_threads: int = 32 workers_per_thread: int = 16 - max_gpu_evals_per_move: int = 64 * 1024 + max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 seed: int = 42 selfplay_precision: str = "bf16" diff --git a/training/src/lib.rs b/training/src/lib.rs index 9ab2233..0e12c20 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -172,7 +172,7 @@ impl SelfPlay { workers_per_thread, seed, *, - max_gpu_evals_per_move = 65536, + max_gpu_evals_per_move = 4096, model, selfplay_precision = "bf16" ))] diff --git a/training/src/worker.rs b/training/src/worker.rs index 6881cdb..9041bb4 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -73,7 +73,7 @@ pub struct WorkerConfig { impl Default for WorkerConfig { fn default() -> Self { Self { - max_gpu_evals_per_move: 64 * 1024, + max_gpu_evals_per_move: 4 * 1024, } } } From 3c99bdc71862b38d46e8a3500555904dd4c404b3 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 20:08:20 -0400 Subject: [PATCH 23/59] more timing & bench work --- flake.nix | 1 + .../alphapaint_training/cudagraph_backend.py | 1 + python/alphapaint_training/model.py | 6 + python/alphapaint_training/train.py | 293 +++++++++++++++++- training/src/descent.rs | 104 ++++++- training/src/lib.rs | 45 +++ training/src/training.rs | 89 ++++++ training/src/worker.rs | 88 +++++- 8 files changed, 593 insertions(+), 34 deletions(-) diff --git a/flake.nix b/flake.nix index 2f14826..5694219 100644 --- a/flake.nix +++ b/flake.nix @@ -43,6 +43,7 @@ export LD_LIBRARY_PATH=$CUDA_PATH/lib:$CUDA_PATH/lib64:$CUDNN_PATH/lib:/run/opengl-driver/lib:${pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH export TRITON_LIBCUDA_PATH=/run/opengl-driver/lib export PATH=$CUDA_PATH/bin:$PATH + export PYTHONPATH=python ''; }; }); diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index ecf816d..dd62b0d 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -54,6 +54,7 @@ def capture_lane_graph( value_device = dlpack.from_dlpack(value_device_dlpack) model = model.cuda() + model = model.to(memory_format=torch.channels_last) model.eval() torch.backends.cudnn.benchmark = True stream = torch.cuda.ExternalStream(stream_handle) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 25c3b93..65f0208 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -111,10 +111,16 @@ def _ensure_decode_buffers( board_key = self._buffer_key(packed_obs, self.board_dtype) board = self._board_buffers.get(board_key) if board is None: + memory_format = ( + torch.channels_last + if packed_obs.device.type == "cuda" + else torch.contiguous_format + ) board = torch.empty( (packed_obs.shape[0], BOARD_PLANES, 32, 32), device=packed_obs.device, dtype=self.board_dtype, + memory_format=memory_format, ) self._board_buffers[board_key] = board diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 432787d..b3a1acb 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -6,6 +6,7 @@ from contextlib import nullcontext from dataclasses import asdict, dataclass from pathlib import Path +from typing import cast import numpy as np import torch @@ -35,23 +36,54 @@ class TrainConfig: wandb: bool = False +@dataclass(slots=True) +class CudaSectionTiming: + start: torch.cuda.Event + end: torch.cuda.Event + + def seconds(self) -> float: + return float(self.start.elapsed_time(self.end) / 1000.0) + + +@dataclass(slots=True) +class TrainStepResult: + loss: torch.Tensor + sample_seconds: float + h2d_seconds: float + forward_seconds: float = 0.0 + backward_seconds: float = 0.0 + optimizer_seconds: float = 0.0 + forward_timing: CudaSectionTiming | None = None + backward_timing: CudaSectionTiming | None = None + optimizer_timing: CudaSectionTiming | None = None + + def _device(device: str) -> torch.device: if device == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA is required for the training runner") return torch.device(device) +def _to_channels_last(module: torch.nn.Module) -> torch.nn.Module: + module = module.to(memory_format=torch.channels_last) # type: ignore[call-overload] + return module + + def _sample_replay_batch( replay_buffer: EphemeralReplayBuffer, batch_size: int, seed: int, device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor, float, float]: actual_batch_size = min(batch_size, len(replay_buffer)) if actual_batch_size == 0: raise RuntimeError("replay buffer is empty") + sample_started_at = time.perf_counter() obs_np, values_np = replay_buffer.sample(actual_batch_size, seed) + sample_seconds = time.perf_counter() - sample_started_at + + h2d_started_at = time.perf_counter() obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( device=device, dtype=torch.uint16, @@ -62,7 +94,8 @@ def _sample_replay_batch( dtype=torch.float32, non_blocking=True, ) - return obs, values + h2d_seconds = time.perf_counter() - h2d_started_at + return obs, values, sample_seconds, h2d_seconds def train_step( @@ -73,9 +106,11 @@ def train_step( batch_size: int, seed: int, device: torch.device, -) -> torch.Tensor: +) -> TrainStepResult: model.train() - obs, target = _sample_replay_batch(replay_buffer, batch_size, seed, device) + obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( + replay_buffer, batch_size, seed, device + ) optimizer.zero_grad(set_to_none=True) autocast = ( @@ -83,12 +118,66 @@ def train_step( if device.type == "cuda" else nullcontext() ) + + if device.type == "cuda": + forward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + backward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + optimizer_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + + forward_timing.start.record() + with autocast: + pred = model(obs) + loss = F.mse_loss(pred.float(), target) + forward_timing.end.record() + + backward_timing.start.record() + loss.backward() + backward_timing.end.record() + + optimizer_timing.start.record() + optimizer.step() + optimizer_timing.end.record() + + return TrainStepResult( + loss=loss.detach(), + sample_seconds=sample_seconds, + h2d_seconds=h2d_seconds, + forward_timing=forward_timing, + backward_timing=backward_timing, + optimizer_timing=optimizer_timing, + ) + + forward_started_at = time.perf_counter() with autocast: pred = model(obs) loss = F.mse_loss(pred.float(), target) + forward_seconds = time.perf_counter() - forward_started_at + + backward_started_at = time.perf_counter() loss.backward() + backward_seconds = time.perf_counter() - backward_started_at + + optimizer_started_at = time.perf_counter() optimizer.step() - return loss.detach() + optimizer_seconds = time.perf_counter() - optimizer_started_at + + return TrainStepResult( + loss=loss.detach(), + sample_seconds=sample_seconds, + h2d_seconds=h2d_seconds, + forward_seconds=forward_seconds, + backward_seconds=backward_seconds, + optimizer_seconds=optimizer_seconds, + ) def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: @@ -142,7 +231,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) ) - model = PackedValueModel().to(device) + model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) model.eval() optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) @@ -168,6 +257,15 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_completed_turns = 0 previous_action_turn_count_total = 0 previous_completed_game_actions_total = 0 + previous_tree_build_nanos = 0 + previous_descent_nanos = 0 + previous_sample_collect_nanos = 0 + previous_replay_push_nanos = 0 + previous_descent_expand_cpu_nanos = 0 + previous_descent_apply_action_nanos = 0 + previous_descent_eval_submit_nanos = 0 + previous_descent_eval_await_nanos = 0 + previous_descent_backup_nanos = 0 try: for round_idx in range(config.rounds): round_number = round_idx + 1 @@ -185,6 +283,15 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: max_turn_count_seen = selfplay.max_turn_count_seen() completed_game_actions_total = selfplay.completed_game_actions_total() max_actions_in_completed_game = selfplay.max_actions_in_completed_game() + tree_build_nanos = selfplay.tree_build_nanos() + descent_nanos = selfplay.descent_nanos() + sample_collect_nanos = selfplay.sample_collect_nanos() + replay_push_nanos = selfplay.replay_push_nanos() + descent_expand_cpu_nanos = selfplay.descent_expand_cpu_nanos() + descent_apply_action_nanos = selfplay.descent_apply_action_nanos() + descent_eval_submit_nanos = selfplay.descent_eval_submit_nanos() + descent_eval_await_nanos = selfplay.descent_eval_await_nanos() + descent_backup_nanos = selfplay.descent_backup_nanos() collect_seconds = time.perf_counter() - collect_started_at samples_added = collected - previous_samples games_added = games - previous_games @@ -200,6 +307,36 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: completed_game_actions_added = ( completed_game_actions_total - previous_completed_game_actions_total ) + tree_build_seconds = (tree_build_nanos - previous_tree_build_nanos) / 1e9 + descent_seconds = (descent_nanos - previous_descent_nanos) / 1e9 + sample_collect_seconds = ( + sample_collect_nanos - previous_sample_collect_nanos + ) / 1e9 + replay_push_seconds = (replay_push_nanos - previous_replay_push_nanos) / 1e9 + descent_expand_cpu_seconds = ( + descent_expand_cpu_nanos - previous_descent_expand_cpu_nanos + ) / 1e9 + descent_apply_action_seconds = ( + descent_apply_action_nanos - previous_descent_apply_action_nanos + ) / 1e9 + descent_eval_submit_seconds = ( + descent_eval_submit_nanos - previous_descent_eval_submit_nanos + ) / 1e9 + descent_eval_await_seconds = ( + descent_eval_await_nanos - previous_descent_eval_await_nanos + ) / 1e9 + descent_backup_seconds = ( + descent_backup_nanos - previous_descent_backup_nanos + ) / 1e9 + descent_other_seconds = max( + 0.0, + descent_seconds + - descent_expand_cpu_seconds + - descent_apply_action_seconds + - descent_eval_submit_seconds + - descent_eval_await_seconds + - descent_backup_seconds, + ) avg_turn_count = ( action_turn_count_added / action_steps_added if action_steps_added @@ -217,11 +354,71 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"{actions_per_turn:.2f}" if completed_turns_added else "-" ) actions_per_game_display = f"{actions_per_game:.1f}" if games_added else "-" + build_ms_per_action = ( + tree_build_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_ms_per_action = ( + descent_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + collect_ms_per_action = ( + sample_collect_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + push_ms_per_action = ( + replay_push_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_expand_ms_per_action = ( + descent_expand_cpu_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_apply_ms_per_action = ( + descent_apply_action_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_ms_per_action = ( + descent_eval_submit_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_wait_ms_per_action = ( + descent_eval_await_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_backup_ms_per_action = ( + descent_backup_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_other_ms_per_action = ( + descent_other_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_us_per_eval = ( + descent_eval_submit_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) + descent_wait_us_per_eval = ( + descent_eval_await_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) - round_losses: list[torch.Tensor] = [] + round_results: list[TrainStepResult] = [] train_started_at = time.perf_counter() for step_idx in range(config.train_steps_per_round): - loss = train_step( + result = train_step( model, replay_buffer, optimizer, @@ -229,13 +426,47 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: seed=config.seed + round_idx * 10_000 + step_idx, device=device, ) - round_losses.append(loss) + round_results.append(result) + if device.type == "cuda" and round_results: + torch.cuda.synchronize(device) model.eval() train_seconds = time.perf_counter() - train_started_at replay_size = len(replay_buffer) - if round_losses: - round_loss_values = torch.stack(round_losses).float().cpu() + train_sample_seconds = sum( + result.sample_seconds for result in round_results + ) + train_h2d_seconds = sum(result.h2d_seconds for result in round_results) + train_forward_seconds = sum( + result.forward_seconds for result in round_results + ) + train_backward_seconds = sum( + result.backward_seconds for result in round_results + ) + train_optimizer_seconds = sum( + result.optimizer_seconds for result in round_results + ) + if device.type == "cuda": + train_forward_seconds += sum( + result.forward_timing.seconds() + for result in round_results + if result.forward_timing is not None + ) + train_backward_seconds += sum( + result.backward_timing.seconds() + for result in round_results + if result.backward_timing is not None + ) + train_optimizer_seconds += sum( + result.optimizer_timing.seconds() + for result in round_results + if result.optimizer_timing is not None + ) + + if round_results: + round_loss_values = ( + torch.stack([result.loss for result in round_results]).float().cpu() + ) loss_values = round_loss_values.tolist() losses.extend(loss_values) mean_loss = float(round_loss_values.mean().item()) @@ -277,6 +508,33 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "samples_per_second": samples_added / max(collect_seconds, 1e-9), "train_steps_per_second": config.train_steps_per_round / max(train_seconds, 1e-9), + "selfplay_tree_build_seconds": tree_build_seconds, + "selfplay_descent_seconds": descent_seconds, + "selfplay_sample_collect_seconds": sample_collect_seconds, + "selfplay_replay_push_seconds": replay_push_seconds, + "selfplay_tree_build_ms_per_action": build_ms_per_action, + "selfplay_descent_ms_per_action": descent_ms_per_action, + "selfplay_sample_collect_ms_per_action": collect_ms_per_action, + "selfplay_replay_push_ms_per_action": push_ms_per_action, + "descent_expand_cpu_seconds": descent_expand_cpu_seconds, + "descent_apply_action_seconds": descent_apply_action_seconds, + "descent_eval_submit_seconds": descent_eval_submit_seconds, + "descent_eval_await_seconds": descent_eval_await_seconds, + "descent_backup_seconds": descent_backup_seconds, + "descent_other_seconds": descent_other_seconds, + "descent_expand_cpu_ms_per_action": descent_expand_ms_per_action, + "descent_apply_action_ms_per_action": descent_apply_ms_per_action, + "descent_eval_submit_ms_per_action": descent_submit_ms_per_action, + "descent_eval_await_ms_per_action": descent_wait_ms_per_action, + "descent_backup_ms_per_action": descent_backup_ms_per_action, + "descent_other_ms_per_action": descent_other_ms_per_action, + "descent_eval_submit_us_per_eval": descent_submit_us_per_eval, + "descent_eval_await_us_per_eval": descent_wait_us_per_eval, + "train_sample_seconds": train_sample_seconds, + "train_h2d_seconds": train_h2d_seconds, + "train_forward_seconds": train_forward_seconds, + "train_backward_seconds": train_backward_seconds, + "train_optimizer_seconds": train_optimizer_seconds, "loss_mean": mean_loss, "timestamp": time.time(), } @@ -286,6 +544,10 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " + f"sp_ms/act=build:{build_ms_per_action:.1f} desc:{descent_ms_per_action:.1f} coll:{collect_ms_per_action:.1f} push:{push_ms_per_action:.1f} " + f"desc_ms/act=exp:{descent_expand_ms_per_action:.1f} app:{descent_apply_ms_per_action:.1f} sub:{descent_submit_ms_per_action:.1f} wait:{descent_wait_ms_per_action:.1f} bk:{descent_backup_ms_per_action:.1f} other:{descent_other_ms_per_action:.1f} " + f"desc_us/eval=sub:{descent_submit_us_per_eval:.2f} wait:{descent_wait_us_per_eval:.2f} " + f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) wandb.log(record) @@ -317,6 +579,15 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_completed_turns = completed_turns previous_action_turn_count_total = action_turn_count_total previous_completed_game_actions_total = completed_game_actions_total + previous_tree_build_nanos = tree_build_nanos + previous_descent_nanos = descent_nanos + previous_sample_collect_nanos = sample_collect_nanos + previous_replay_push_nanos = replay_push_nanos + previous_descent_expand_cpu_nanos = descent_expand_cpu_nanos + previous_descent_apply_action_nanos = descent_apply_action_nanos + previous_descent_eval_submit_nanos = descent_eval_submit_nanos + previous_descent_eval_await_nanos = descent_eval_await_nanos + previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() diff --git a/training/src/descent.rs b/training/src/descent.rs index a4d4bde..3629e6b 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -8,6 +8,9 @@ use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use rand::rngs::SmallRng; use rand::{Rng, RngExt}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Instant; use crate::eval::{EvalCountTracker, Evaluator}; @@ -18,6 +21,44 @@ pub struct TreeLearningSample { pub value: f32, } +#[derive(Clone)] +pub struct SearchTimingMetrics { + pub expand_cpu_nanos: Arc, + pub apply_action_nanos: Arc, + pub eval_submit_nanos: Arc, + pub eval_await_nanos: Arc, + pub backup_nanos: Arc, +} + +impl SearchTimingMetrics { + fn record(counter: &AtomicU64, started_at: Instant) { + counter.fetch_add( + started_at.elapsed().as_nanos().min(u64::MAX as u128) as u64, + Ordering::AcqRel, + ); + } + + fn record_expand_cpu(&self, started_at: Instant) { + Self::record(&self.expand_cpu_nanos, started_at); + } + + fn record_apply_action(&self, started_at: Instant) { + Self::record(&self.apply_action_nanos, started_at); + } + + fn record_eval_submit(&self, started_at: Instant) { + Self::record(&self.eval_submit_nanos, started_at); + } + + fn record_eval_await(&self, started_at: Instant) { + Self::record(&self.eval_await_nanos, started_at); + } + + fn record_backup(&self, started_at: Instant) { + Self::record(&self.backup_nanos, started_at); + } +} + #[derive(Debug)] pub struct ChildData { pub action: Action, @@ -187,11 +228,14 @@ impl SearchNode { outcome: ApplyActionOutcome, evaluator: &E, rng: &mut SmallRng, + timing: &SearchTimingMetrics, ) -> SearchNode { match outcome { ApplyActionOutcome::Ongoing => { let mut new_node = SearchNode::new(0.0, 0, false); + let actions_started_at = Instant::now(); let actions = board.get_valid_actions(); + timing.record_expand_cpu(actions_started_at); new_node.children.reserve(actions.len()); if actions.len() == 0 { @@ -208,24 +252,34 @@ impl SearchNode { } let mut eval_results = Vec::new(); - let batch = actions.into_iter().map(|action| { - let mut local_board = board.clone(); - let (child_outcome, rollback) = local_board.apply_action(*action); - // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics - // rust futures are normally lazily evaluated when you .await them for the - // first time. Ours are eager, which is why this code actually works as we expect. - let value = evaluator.evaluate(local_board.clone()); - (action, local_board, child_outcome, rollback, value) - }); + let batch: Vec<_> = actions + .into_iter() + .map(|action| { + let prep_started_at = Instant::now(); + let mut local_board = board.clone(); + timing.record_expand_cpu(prep_started_at); + + let apply_started_at = Instant::now(); + let (child_outcome, rollback) = local_board.apply_action(*action); + timing.record_apply_action(apply_started_at); + // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics + // rust futures are normally lazily evaluated when you .await them for the + // first time. Ours are eager, which is why this code actually works as we expect. + let submit_started_at = Instant::now(); + let value = evaluator.evaluate(local_board.clone()); + timing.record_eval_submit(submit_started_at); + (action, local_board, child_outcome, rollback, value) + }) + .collect(); for (&action, local_board, child_outcome, _, value) in batch { match child_outcome { ApplyActionOutcome::Ongoing => { // Evaluate this child with neural net - eval_results.push(ChildEvalResult { - action, - value: value.await, - }); + let await_started_at = Instant::now(); + let value = value.await; + timing.record_eval_await(await_started_at); + eval_results.push(ChildEvalResult { action, value }); } ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => { @@ -299,6 +353,7 @@ impl SearchNode { outcome: ApplyActionOutcome, evaluator: &E, rng: &mut SmallRng, + timing: &SearchTimingMetrics, ) -> f32 { let white_turn = state.is_white_turn(); @@ -319,12 +374,17 @@ impl SearchNode { self.children[best_action_id].entrance_count += 1; if self.children[best_action_id].node.is_some() { + let apply_started_at = Instant::now(); let (outcome, _) = state.apply_action(best_action); + timing.record_apply_action(apply_started_at); let child = self.children[best_action_id].node.as_mut().unwrap(); self.children[best_action_id].child_value = - Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng)).await; + Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng, timing)) + .await; } else { + let apply_started_at = Instant::now(); let (child_outcome, _) = state.apply_action(best_action); + timing.record_apply_action(apply_started_at); let should_descend = matches!(child_outcome, ApplyActionOutcome::Ongoing); let child_node = Box::new( Box::pin(SearchNode::build_self( @@ -332,6 +392,7 @@ impl SearchNode { child_outcome, evaluator, rng, + timing, )) .await, ); @@ -347,15 +408,18 @@ impl SearchNode { ApplyActionOutcome::Ongoing, evaluator, rng, + timing, )) .await; } } + let backup_started_at = Instant::now(); let (best_action_id, _) = self.completed_best_action(white_turn, rng); self.completion_value = self.children[best_action_id].completion_value(); self.value = self.children[best_action_id].child_value; self.resolved = self.backup_resolution(); + timing.record_backup(backup_started_at); } } @@ -540,6 +604,7 @@ pub struct GameSearchTree<'a, E: Evaluator> { pub root_state: Board, rng: SmallRng, evaluator: &'a E, + timing: SearchTimingMetrics, } impl<'a, E: Evaluator> GameSearchTree<'a, E> { @@ -582,7 +647,12 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { } /// Create a new search tree. The board must NOT be in a terminal state. - pub async fn new(board: &Board, evaluator: &'a E, rng: SmallRng) -> GameSearchTree<'a, E> { + pub async fn new( + board: &Board, + evaluator: &'a E, + rng: SmallRng, + timing: SearchTimingMetrics, + ) -> GameSearchTree<'a, E> { let mut local_rng = rng; GameSearchTree { root_node: Box::new( @@ -591,12 +661,14 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { ApplyActionOutcome::Ongoing, evaluator, &mut local_rng, + &timing, ) .await, ), root_state: board.clone(), rng: local_rng, evaluator, + timing, } } @@ -656,6 +728,7 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { ApplyActionOutcome::Ongoing, self.evaluator, &mut self.rng, + &self.timing, ) .await; } @@ -759,6 +832,7 @@ impl<'a, E: EvalCountTracker> GameSearchTree<'a, E> { ApplyActionOutcome::Ongoing, self.evaluator, &mut self.rng, + &self.timing, ) .await; } diff --git a/training/src/lib.rs b/training/src/lib.rs index 0e12c20..a17aa7d 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -319,6 +319,51 @@ impl SelfPlay { Ok(self.session()?.max_actions_in_completed_game()) } + /// Return total time spent building fresh search trees, in nanoseconds. + fn tree_build_nanos(&self) -> PyResult { + Ok(self.session()?.tree_build_nanos()) + } + + /// Return total time spent running descent iterations, in nanoseconds. + fn descent_nanos(&self) -> PyResult { + Ok(self.session()?.descent_nanos()) + } + + /// Return total time spent collecting replay samples, in nanoseconds. + fn sample_collect_nanos(&self) -> PyResult { + Ok(self.session()?.sample_collect_nanos()) + } + + /// Return total time spent encoding and pushing replay samples, in nanoseconds. + fn replay_push_nanos(&self) -> PyResult { + Ok(self.session()?.replay_push_nanos()) + } + + /// Return total CPU expansion prep time inside Descent, in nanoseconds. + fn descent_expand_cpu_nanos(&self) -> PyResult { + Ok(self.session()?.descent_expand_cpu_nanos()) + } + + /// Return total time spent applying actions during Descent, in nanoseconds. + fn descent_apply_action_nanos(&self) -> PyResult { + Ok(self.session()?.descent_apply_action_nanos()) + } + + /// Return total time spent submitting evals to the GPU queue, in nanoseconds. + fn descent_eval_submit_nanos(&self) -> PyResult { + Ok(self.session()?.descent_eval_submit_nanos()) + } + + /// Return total time spent awaiting queued evals, in nanoseconds. + fn descent_eval_await_nanos(&self) -> PyResult { + Ok(self.session()?.descent_eval_await_nanos()) + } + + /// Return total time spent backing up/searching best children, in nanoseconds. + fn descent_backup_nanos(&self) -> PyResult { + Ok(self.session()?.descent_backup_nanos()) + } + /// Return the total number of CUDA graph launches completed so far. fn gpu_batches(&self) -> u64 { self.runner.dispatched_batches() diff --git a/training/src/training.rs b/training/src/training.rs index 2205e76..6856f70 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -46,6 +46,24 @@ struct SessionControl { completed_game_actions_total: Arc, /// Largest number of selected actions observed in a completed game. max_actions_in_completed_game: Arc, + /// Total time spent building fresh search trees. + tree_build_nanos: Arc, + /// Total time spent running descent iterations. + descent_nanos: Arc, + /// Total time spent traversing trees to collect replay samples. + sample_collect_nanos: Arc, + /// Total time spent encoding and pushing replay samples. + replay_push_nanos: Arc, + /// Total CPU expansion prep time inside Descent. + descent_expand_cpu_nanos: Arc, + /// Total time spent applying actions during Descent. + descent_apply_action_nanos: Arc, + /// Total time spent submitting evals to the GPU queue. + descent_eval_submit_nanos: Arc, + /// Total time spent awaiting queued evals. + descent_eval_await_nanos: Arc, + /// Total time spent backing up/searching best children. + descent_backup_nanos: Arc, /// Number of threads currently inside the executor polling loop. active_pollers: AtomicUsize, /// Condvar + mutex for coordinating start/pause/quiesce/shutdown. @@ -69,6 +87,15 @@ impl SessionControl { max_turn_count_seen: Arc::new(AtomicUsize::new(0)), completed_game_actions_total: Arc::new(AtomicU64::new(0)), max_actions_in_completed_game: Arc::new(AtomicU64::new(0)), + tree_build_nanos: Arc::new(AtomicU64::new(0)), + descent_nanos: Arc::new(AtomicU64::new(0)), + sample_collect_nanos: Arc::new(AtomicU64::new(0)), + replay_push_nanos: Arc::new(AtomicU64::new(0)), + descent_expand_cpu_nanos: Arc::new(AtomicU64::new(0)), + descent_apply_action_nanos: Arc::new(AtomicU64::new(0)), + descent_eval_submit_nanos: Arc::new(AtomicU64::new(0)), + descent_eval_await_nanos: Arc::new(AtomicU64::new(0)), + descent_backup_nanos: Arc::new(AtomicU64::new(0)), active_pollers: AtomicUsize::new(0), condvar: Condvar::new(), condvar_mutex: Mutex::new(()), @@ -296,6 +323,59 @@ impl SelfPlaySession { .load(Ordering::Acquire) } + /// Return the total time spent building fresh search trees. + pub fn tree_build_nanos(&self) -> u64 { + self.control.tree_build_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent running descent iterations. + pub fn descent_nanos(&self) -> u64 { + self.control.descent_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent traversing trees to collect replay samples. + pub fn sample_collect_nanos(&self) -> u64 { + self.control.sample_collect_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent encoding and pushing replay samples. + pub fn replay_push_nanos(&self) -> u64 { + self.control.replay_push_nanos.load(Ordering::Acquire) + } + + /// Return total CPU expansion prep time inside Descent. + pub fn descent_expand_cpu_nanos(&self) -> u64 { + self.control + .descent_expand_cpu_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent applying actions during Descent. + pub fn descent_apply_action_nanos(&self) -> u64 { + self.control + .descent_apply_action_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent submitting evals to the GPU queue. + pub fn descent_eval_submit_nanos(&self) -> u64 { + self.control + .descent_eval_submit_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent awaiting queued evals. + pub fn descent_eval_await_nanos(&self) -> u64 { + self.control + .descent_eval_await_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent backing up/searching best children. + pub fn descent_backup_nanos(&self) -> u64 { + self.control.descent_backup_nanos.load(Ordering::Acquire) + } + /// Shut down the session. Idempotent. pub fn shutdown(&mut self) { if let Some(threads) = self.threads.take() { @@ -339,6 +419,15 @@ fn session_thread_main( max_turn_count_seen: control.max_turn_count_seen.clone(), completed_game_actions_total: control.completed_game_actions_total.clone(), max_actions_in_completed_game: control.max_actions_in_completed_game.clone(), + tree_build_nanos: control.tree_build_nanos.clone(), + descent_nanos: control.descent_nanos.clone(), + sample_collect_nanos: control.sample_collect_nanos.clone(), + replay_push_nanos: control.replay_push_nanos.clone(), + descent_expand_cpu_nanos: control.descent_expand_cpu_nanos.clone(), + descent_apply_action_nanos: control.descent_apply_action_nanos.clone(), + descent_eval_submit_nanos: control.descent_eval_submit_nanos.clone(), + descent_eval_await_nanos: control.descent_eval_await_nanos.clone(), + descent_backup_nanos: control.descent_backup_nanos.clone(), }; let mut futures: Vec + '_>>> = (0 diff --git a/training/src/worker.rs b/training/src/worker.rs index 9041bb4..3d63c8e 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::Instant; use ndarray::Ix1; use rand::rngs::SmallRng; @@ -13,8 +14,7 @@ use rand::{Rng, RngExt, SeedableRng}; use alpha_paint::TRAINING_START_FENS; use alpha_paint::board::Board; -use crate::descent::GameSearchTree; -use crate::descent::TreeLearningSample; +use crate::descent::{GameSearchTree, SearchTimingMetrics, TreeLearningSample}; use crate::eval::{CountingEvaluator, Evaluator}; use crate::observation; use crate::replay_buffer::ReplayBuffer; @@ -49,6 +49,10 @@ fn update_max_u64(max_value: &AtomicU64, candidate: u64) { } } +fn elapsed_nanos(started_at: Instant) -> u64 { + started_at.elapsed().as_nanos().min(u64::MAX as u128) as u64 +} + #[derive(Clone)] pub struct SelfPlayMetrics { pub samples_collected: Arc, @@ -61,6 +65,15 @@ pub struct SelfPlayMetrics { pub max_turn_count_seen: Arc, pub completed_game_actions_total: Arc, pub max_actions_in_completed_game: Arc, + pub tree_build_nanos: Arc, + pub descent_nanos: Arc, + pub sample_collect_nanos: Arc, + pub replay_push_nanos: Arc, + pub descent_expand_cpu_nanos: Arc, + pub descent_apply_action_nanos: Arc, + pub descent_eval_submit_nanos: Arc, + pub descent_eval_await_nanos: Arc, + pub descent_backup_nanos: Arc, } /// Configuration for the worker. @@ -88,6 +101,8 @@ fn push_samples_to_replay( return; } + let started_at = Instant::now(); + let mut guard = replay_buffer.reserve(num_samples); for sample in samples { guard.push_with_observation(sample.value, |mut out| { @@ -101,6 +116,35 @@ fn push_samples_to_replay( metrics .samples_collected .fetch_add(num_samples, Ordering::AcqRel); + metrics + .replay_push_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); +} + +fn collect_tree_learning_samples_timed( + tree: &GameSearchTree<'_, E>, + metrics: &SelfPlayMetrics, +) -> Vec { + let started_at = Instant::now(); + let samples = tree.collect_tree_learning_samples(); + metrics + .sample_collect_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); + samples +} + +fn step_tree_and_collect_dropped_samples_timed( + tree: &mut GameSearchTree<'_, E>, + new_board: &Board, + action_id: usize, + metrics: &SelfPlayMetrics, +) -> Vec { + let started_at = Instant::now(); + let samples = tree.step_tree_and_collect_dropped_samples(new_board, action_id); + metrics + .sample_collect_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); + samples } /// Run a single self-play game with tree learning. @@ -130,8 +174,20 @@ async fn play_game( } let counting_evaluator = CountingEvaluator::new(evaluator); + let search_timing = SearchTimingMetrics { + expand_cpu_nanos: metrics.descent_expand_cpu_nanos.clone(), + apply_action_nanos: metrics.descent_apply_action_nanos.clone(), + eval_submit_nanos: metrics.descent_eval_submit_nanos.clone(), + eval_await_nanos: metrics.descent_eval_await_nanos.clone(), + backup_nanos: metrics.descent_backup_nanos.clone(), + }; let tree_rng = SmallRng::from_rng(rng); - let mut tree = GameSearchTree::new(&board, &counting_evaluator, tree_rng).await; + let build_started_at = Instant::now(); + let mut tree = + GameSearchTree::new(&board, &counting_evaluator, tree_rng, search_timing.clone()).await; + metrics + .tree_build_nanos + .fetch_add(elapsed_nanos(build_started_at), Ordering::AcqRel); let mut action_steps_in_game = 0u64; loop { @@ -142,8 +198,12 @@ async fn play_game( update_max_usize(&metrics.max_turn_count_seen, current_turn_count); // Run descent search + let descent_started_at = Instant::now(); tree.run_descent_to_eval_limit(config.max_gpu_evals_per_move) .await; + metrics + .descent_nanos + .fetch_add(elapsed_nanos(descent_started_at), Ordering::AcqRel); // Select action via ordinal distribution let action_id = tree.ordinal_select(); @@ -169,7 +229,7 @@ async fn play_game( push_samples_to_replay( replay_buffer, metrics, - tree.collect_tree_learning_samples(), + collect_tree_learning_samples_timed(&tree, metrics), ); metrics .completed_game_actions_total @@ -182,7 +242,7 @@ async fn play_game( push_samples_to_replay( replay_buffer, metrics, - tree.collect_tree_learning_samples(), + collect_tree_learning_samples_timed(&tree, metrics), ); metrics .completed_game_actions_total @@ -192,16 +252,28 @@ async fn play_game( } alpha_paint::board::ApplyActionOutcome::Ongoing => { if tree.root_node.children[action_id].node.is_some() { - let samples = tree.step_tree_and_collect_dropped_samples(&new_board, action_id); + let samples = step_tree_and_collect_dropped_samples_timed( + &mut tree, &new_board, action_id, metrics, + ); push_samples_to_replay(replay_buffer, metrics, samples); } else { push_samples_to_replay( replay_buffer, metrics, - tree.collect_tree_learning_samples(), + collect_tree_learning_samples_timed(&tree, metrics), ); let next_rng = SmallRng::from_rng(rng); - tree = GameSearchTree::new(&new_board, &counting_evaluator, next_rng).await; + let build_started_at = Instant::now(); + tree = GameSearchTree::new( + &new_board, + &counting_evaluator, + next_rng, + search_timing.clone(), + ) + .await; + metrics + .tree_build_nanos + .fetch_add(elapsed_nanos(build_started_at), Ordering::AcqRel); } } } From 26e4a16f62f6fbee83e6645917187dc9be6624f8 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Fri, 27 Mar 2026 20:12:32 -0400 Subject: [PATCH 24/59] increase self-play queue headroom This raises queue capacity for eager eval submission, lowers the default worker fanout, and surfaces outstanding eval pressure so queue wraparound risks are visible before they corrupt results. --- python/alphapaint_training/train.py | 12 +++++- training/src/descent.rs | 2 +- training/src/future.rs | 22 +++++++++++ training/src/lib.rs | 20 ++++++++++ training/src/queue.rs | 57 +++++++++++++++++++++++++++-- training/src/training.rs | 42 ++++++++++++++++++++- 6 files changed, 148 insertions(+), 7 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index b3a1acb..a8bebda 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -25,7 +25,7 @@ class TrainConfig: batch_size: int = 24_576 replay_capacity: int = 16_000_000 num_threads: int = 32 - workers_per_thread: int = 16 + workers_per_thread: int = 8 max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 seed: int = 42 @@ -275,6 +275,10 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: games = selfplay.games() gpu_batches = selfplay.gpu_batches() gpu_evals = selfplay.gpu_evals() + queue_num_batches = selfplay.queue_num_batches() + queue_total_slots = selfplay.queue_total_slots() + current_outstanding_evals = selfplay.current_outstanding_evals() + max_outstanding_evals = selfplay.max_outstanding_evals() action_steps = selfplay.action_steps() final_actions = selfplay.final_actions() nonfinal_actions = selfplay.nonfinal_actions() @@ -485,6 +489,10 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "gpu_evals_total": gpu_evals, "gpu_evals_added": gpu_evals_added, "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), + "queue_num_batches": queue_num_batches, + "queue_total_slots": queue_total_slots, + "queue_current_outstanding_evals": current_outstanding_evals, + "queue_max_outstanding_evals": max_outstanding_evals, "action_steps_total": action_steps, "action_steps_added": action_steps_added, "final_actions_total": final_actions, @@ -541,7 +549,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print( f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " - f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " + f"batches={gpu_batches_added} queue={current_outstanding_evals}/{max_outstanding_evals}/{queue_total_slots}@{queue_num_batches} replay={replay_size} act={action_steps_added} final={final_actions_added} " f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " f"sp_ms/act=build:{build_ms_per_action:.1f} desc:{descent_ms_per_action:.1f} coll:{collect_ms_per_action:.1f} push:{push_ms_per_action:.1f} " diff --git a/training/src/descent.rs b/training/src/descent.rs index 3629e6b..a57f613 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -8,8 +8,8 @@ use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use rand::rngs::SmallRng; use rand::{Rng, RngExt}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; use crate::eval::{EvalCountTracker, Evaluator}; diff --git a/training/src/future.rs b/training/src/future.rs index f4bf839..e536841 100644 --- a/training/src/future.rs +++ b/training/src/future.rs @@ -43,6 +43,7 @@ where queue: &'a GpuJobQueue, ticket: u64, completed: bool, + tracking_outstanding: bool, } impl<'a, A, D, O> GpuEvalFuture<'a, A, D, O> @@ -54,12 +55,14 @@ where { /// Create a new future with an already-submitted ticket. pub fn new(queue: &'a GpuJobQueue, ticket: u64) -> Self { + queue.acquire_outstanding_eval_slot(); // Signal progress on creation since we just submitted signal_progress(); Self { queue, ticket, completed: false, + tracking_outstanding: true, } } } @@ -83,6 +86,10 @@ where if let Some(&output) = this.queue.poll(this.ticket) { this.completed = true; + if this.tracking_outstanding { + this.queue.release_outstanding_eval_slot(); + this.tracking_outstanding = false; + } signal_progress(); Poll::Ready(output) } else { @@ -91,6 +98,21 @@ where } } +impl Drop for GpuEvalFuture<'_, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + fn drop(&mut self) { + if self.tracking_outstanding { + self.queue.release_outstanding_eval_slot(); + self.tracking_outstanding = false; + } + } +} + impl GpuJobQueue where A: Clone + Default + Send + Sync, diff --git a/training/src/lib.rs b/training/src/lib.rs index a17aa7d..7c2a821 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -279,6 +279,26 @@ impl SelfPlay { Ok(self.session()?.games()) } + /// Return the number of queue batch lanes. + fn queue_num_batches(&self) -> PyResult { + Ok(self.session()?.queue_num_batches()) + } + + /// Return the total number of queue slots. + fn queue_total_slots(&self) -> PyResult { + Ok(self.session()?.queue_total_slots()) + } + + /// Return the current number of outstanding eval futures. + fn current_outstanding_evals(&self) -> PyResult { + Ok(self.session()?.current_outstanding_evals()) + } + + /// Return the maximum number of outstanding eval futures seen so far. + fn max_outstanding_evals(&self) -> PyResult { + Ok(self.session()?.max_outstanding_evals()) + } + /// Return the total number of selected actions. fn action_steps(&self) -> PyResult { Ok(self.session()?.action_steps()) diff --git a/training/src/queue.rs b/training/src/queue.rs index 81eceaa..12ee3f1 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -23,12 +23,13 @@ pub const BATCH_SIZE: usize = 16; #[cfg(not(test))] pub const BATCH_SIZE: usize = 256; -const SLOT_MULTIPLIER: usize = 2; +const SLOT_MULTIPLIER: usize = 32; /// Compute queue shape for a given worker count. /// /// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a -/// whole number of `BATCH_SIZE` lanes and is at least `2 * num_workers`. +/// whole number of `BATCH_SIZE` lanes and is at least +/// `SLOT_MULTIPLIER * num_workers`. pub fn queue_shape_for_workers(num_workers: usize) -> (usize, usize) { assert!(num_workers > 0, "num_workers must be > 0"); let min_slots = num_workers.saturating_mul(SLOT_MULTIPLIER).max(BATCH_SIZE); @@ -83,6 +84,12 @@ where /// Event for parking threads when waiting for GPU completion. completion_event: Event, + + /// Number of live eval futures whose results have not been consumed yet. + outstanding_evals: AtomicU64, + + /// High watermark for `outstanding_evals`. + max_outstanding_evals: AtomicU64, } // SAFETY: Access to queue state is synchronized via ticket ownership and atomics. @@ -158,7 +165,7 @@ where /// Creates a new job queue with the given observation shape, worker count, /// and dispatch callback. /// - /// Queue storage is provisioned to at least `2 * num_workers` slots, + /// Queue storage is provisioned to at least `SLOT_MULTIPLIER * num_workers` slots, /// rounded up to a whole number of batches. /// /// The callback is invoked when a batch of BATCH_SIZE jobs is ready. @@ -185,6 +192,8 @@ where batch_complete, outputs, completion_event: Event::new(), + outstanding_evals: AtomicU64::new(0), + max_outstanding_evals: AtomicU64::new(0), }); Self { @@ -207,6 +216,48 @@ where self.total_slots } + #[inline] + pub fn current_outstanding_evals(&self) -> u64 { + self.state.outstanding_evals.load(Ordering::Acquire) + } + + #[inline] + pub fn max_outstanding_evals(&self) -> u64 { + self.state.max_outstanding_evals.load(Ordering::Acquire) + } + + #[inline] + pub(crate) fn acquire_outstanding_eval_slot(&self) { + let current = self.state.outstanding_evals.fetch_add(1, Ordering::AcqRel) + 1; + + let mut observed_max = self.state.max_outstanding_evals.load(Ordering::Acquire); + while current > observed_max { + match self.state.max_outstanding_evals.compare_exchange_weak( + observed_max, + current, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(new_max) => observed_max = new_max, + } + } + + if current > self.total_slots as u64 { + self.state.outstanding_evals.fetch_sub(1, Ordering::AcqRel); + panic!( + "queue overflow risk: {} outstanding evals exceeded {} total slots", + current, self.total_slots + ); + } + } + + #[inline] + pub(crate) fn release_outstanding_eval_slot(&self) { + let previous = self.state.outstanding_evals.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "outstanding eval counter underflowed"); + } + /// Submit a job by writing an observation via callback. /// /// The callback receives a mutable view into the queue's contiguous storage diff --git a/training/src/training.rs b/training/src/training.rs index 6856f70..988fc08 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -135,7 +135,7 @@ impl Default for SessionConfig { fn default() -> Self { Self { num_threads: 32, - workers_per_thread: 16, + workers_per_thread: 8, worker: WorkerConfig::default(), seed: 42, } @@ -145,6 +145,10 @@ impl Default for SessionConfig { /// Trait-object wrapper so we can call `notify_all()` on the queue. trait QueueNotify: Send + Sync { fn notify_all(&self); + fn num_batches(&self) -> usize; + fn total_slots(&self) -> usize; + fn current_outstanding_evals(&self) -> u64; + fn max_outstanding_evals(&self) -> u64; } impl QueueNotify for GpuJobQueue @@ -156,6 +160,22 @@ where fn notify_all(&self) { GpuJobQueue::notify_all(self); } + + fn num_batches(&self) -> usize { + GpuJobQueue::num_batches(self) + } + + fn total_slots(&self) -> usize { + GpuJobQueue::total_slots(self) + } + + fn current_outstanding_evals(&self) -> u64 { + GpuJobQueue::current_outstanding_evals(self) + } + + fn max_outstanding_evals(&self) -> u64 { + GpuJobQueue::max_outstanding_evals(self) + } } /// A persistent self-play session that owns worker threads and preserves @@ -279,6 +299,26 @@ impl SelfPlaySession { self.control.games_completed.load(Ordering::Acquire) } + /// Return the number of queue batch lanes. + pub fn queue_num_batches(&self) -> usize { + self.queue_notify.num_batches() + } + + /// Return the total number of queue slots. + pub fn queue_total_slots(&self) -> usize { + self.queue_notify.total_slots() + } + + /// Return the current number of outstanding eval futures. + pub fn current_outstanding_evals(&self) -> u64 { + self.queue_notify.current_outstanding_evals() + } + + /// Return the maximum number of outstanding eval futures seen so far. + pub fn max_outstanding_evals(&self) -> u64 { + self.queue_notify.max_outstanding_evals() + } + /// Return the total number of selected actions. pub fn action_steps(&self) -> u64 { self.control.action_steps.load(Ordering::Acquire) From df13b0cb56a6b81e329faaa6a6d033a0b9149ec0 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 01:44:01 -0400 Subject: [PATCH 25/59] try a mixer style model --- python/alphapaint_training/__init__.py | 4 +- .../alphapaint_training.pyi | 45 ++++++ python/alphapaint_training/model.py | 136 ++++++++++-------- python/alphapaint_training/train.py | 25 ++-- 4 files changed, 137 insertions(+), 73 deletions(-) create mode 100644 python/alphapaint_training/alphapaint_training.pyi diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index a0d8c3c..241fb97 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -5,7 +5,7 @@ __path__ = extend_path(__path__, __name__) from .alphapaint_training import EphemeralReplayBuffer, SelfPlay -from .model import PackedValueModel, ResidualBlock, TinyValueNet +from .model import PackedValueModel from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, @@ -26,9 +26,7 @@ "INTRINSIC_COUNT", "OBS_WORDS", "PackedValueModel", - "ResidualBlock", "SelfPlay", - "TinyValueNet", "decode_intrinsics", "decode_packed_board", "decode_packed_board_reference", diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi new file mode 100644 index 0000000..039e243 --- /dev/null +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -0,0 +1,45 @@ +from typing import Any + +class EphemeralReplayBuffer: + def __init__(self, capacity: int) -> None: ... + def __len__(self) -> int: ... + def sample(self, batch_size: int, seed: int) -> tuple[Any, Any]: ... + +class SelfPlay: + def __init__( + self, + replay_buffer: EphemeralReplayBuffer, + num_threads: int, + workers_per_thread: int, + seed: int, + *, + max_gpu_evals_per_move: int, + model: Any, + selfplay_precision: str, + ) -> None: ... + def drop(self) -> None: ... + def wait_for(self, target_samples: int) -> int: ... + def games(self) -> int: ... + def gpu_batches(self) -> int: ... + def gpu_evals(self) -> int: ... + def queue_num_batches(self) -> int: ... + def queue_total_slots(self) -> int: ... + def current_outstanding_evals(self) -> int: ... + def max_outstanding_evals(self) -> int: ... + def action_steps(self) -> int: ... + def final_actions(self) -> int: ... + def nonfinal_actions(self) -> int: ... + def completed_turns(self) -> int: ... + def action_turn_count_total(self) -> int: ... + def max_turn_count_seen(self) -> int: ... + def completed_game_actions_total(self) -> int: ... + def max_actions_in_completed_game(self) -> int: ... + def tree_build_nanos(self) -> int: ... + def descent_nanos(self) -> int: ... + def sample_collect_nanos(self) -> int: ... + def replay_push_nanos(self) -> int: ... + def descent_expand_cpu_nanos(self) -> int: ... + def descent_apply_action_nanos(self) -> int: ... + def descent_eval_submit_nanos(self) -> int: ... + def descent_eval_await_nanos(self) -> int: ... + def descent_backup_nanos(self) -> int: ... diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 65f0208..7c9d05c 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -6,73 +6,94 @@ from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, + BOARD_SIDE, INTRINSIC_COUNT, INTRINSIC_SCALE, decode_packed_board, ) -class ResidualBlock(nn.Module): - def __init__(self, width: int): +_PATCH_SIZE = 4 +_EMBED_DIM = 16 +_NUM_BLOCKS = 2 +_MLP_DIM = 32 +_HIDDEN_DIM = 32 + + +class _MixingMlp(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, input_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class _MixingBlock(nn.Module): + def __init__(self, *, num_tokens: int, embed_dim: int, mlp_dim: int): super().__init__() - self.norm1 = nn.BatchNorm2d(width) - self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.norm2 = nn.BatchNorm2d(width) - self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.act = nn.ReLU(inplace=True) + self.token_norm = nn.LayerNorm(embed_dim) + self.token_mlp = _MixingMlp(num_tokens, mlp_dim) + self.channel_norm = nn.LayerNorm(embed_dim) + self.channel_mlp = _MixingMlp(embed_dim, mlp_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.norm1(x) - x = self.act(x) - x = self.conv1(x) - x = self.norm2(x) - x = self.act(x) - x = self.conv2(x) - return x + residual - - -class TinyValueNet(nn.Module): - def __init__( - self, - *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, - ): + mixed_tokens = self.token_norm(x).transpose(1, 2) + x = x + self.token_mlp(mixed_tokens).transpose(1, 2) + return x + self.channel_mlp(self.channel_norm(x)) + + +class _ValueNet(nn.Module): + def __init__(self): super().__init__() - self.stem = nn.Sequential( - nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), - nn.BatchNorm2d(width), - nn.ReLU(inplace=True), + if BOARD_SIDE % _PATCH_SIZE != 0: + raise ValueError(f"patch size {_PATCH_SIZE} must divide the board side") + + tokens_per_side = BOARD_SIDE // _PATCH_SIZE + num_tokens = tokens_per_side * tokens_per_side + + self.patch_embed = nn.Conv2d( + BOARD_PLANES, + _EMBED_DIM, + kernel_size=_PATCH_SIZE, + stride=_PATCH_SIZE, + bias=False, ) - self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) - self.pool = nn.AdaptiveAvgPool2d(1) - self.flatten = nn.Flatten() + self.position = nn.Parameter(torch.zeros(1, num_tokens, _EMBED_DIM)) + self.blocks = nn.Sequential( + *( + _MixingBlock( + num_tokens=num_tokens, + embed_dim=_EMBED_DIM, + mlp_dim=_MLP_DIM, + ) + for _ in range(_NUM_BLOCKS) + ) + ) + self.norm = nn.LayerNorm(_EMBED_DIM) self.head = nn.Sequential( - nn.Linear(width + INTRINSIC_COUNT, hidden_dim), - nn.ReLU(inplace=True), - nn.Linear(hidden_dim, 1), + nn.Linear(_EMBED_DIM + INTRINSIC_COUNT, _HIDDEN_DIM), + nn.GELU(), + nn.Linear(_HIDDEN_DIM, 1), ) def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: - x = self.stem(board) + x = self.patch_embed(board) + x = x.flatten(2).transpose(1, 2) + x = x + self.position x = self.blocks(x) - x = self.pool(x) - x = self.flatten(x) + x = self.norm(x) + x = x.mean(dim=1) x = torch.cat((x, intrinsics), dim=1) return self.head(x) class PackedValueModel(nn.Module): - def __init__( - self, - *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, - board_dtype: torch.dtype = torch.bfloat16, - ): + def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): super().__init__() self.board_dtype = board_dtype self._board_buffers: dict[ @@ -84,16 +105,13 @@ def __init__( self._intrinsic_buffers: dict[ tuple[str, int | None, int, torch.dtype], torch.Tensor ] = {} + self.intrinsic_scales: torch.Tensor self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), persistent=False, ) - self.value_net = TinyValueNet( - width=width, - num_blocks=num_blocks, - hidden_dim=hidden_dim, - ) + self.value_net = _ValueNet() def _buffer_key( self, packed_obs: torch.Tensor, dtype: torch.dtype @@ -105,9 +123,7 @@ def _buffer_key( dtype, ) - def _ensure_decode_buffers( - self, packed_obs: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: board_key = self._buffer_key(packed_obs, self.board_dtype) board = self._board_buffers.get(board_key) if board is None: @@ -117,7 +133,7 @@ def _ensure_decode_buffers( else torch.contiguous_format ) board = torch.empty( - (packed_obs.shape[0], BOARD_PLANES, 32, 32), + (packed_obs.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), device=packed_obs.device, dtype=self.board_dtype, memory_format=memory_format, @@ -148,11 +164,7 @@ def _ensure_decode_buffers( ) self._intrinsic_buffers[intrinsic_key] = intrinsics - return board, intrinsics_fp32, intrinsics - - def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) - board = decode_packed_board( + decode_packed_board( packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype, out=board, @@ -163,7 +175,7 @@ def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: - board, intrinsics = self.decode(packed_obs) + board, intrinsics = self._decode(packed_obs) if not torch.is_autocast_enabled(): param_dtype = next(self.value_net.parameters()).dtype board = board.to(dtype=param_dtype) @@ -172,4 +184,4 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: return value.squeeze(-1).tanh() * 10.0 -__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] +__all__ = ["PackedValueModel"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index a8bebda..c31e854 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -13,7 +13,7 @@ import torch.nn.functional as F import wandb -from alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training.alphapaint_training import EphemeralReplayBuffer, SelfPlay from alphapaint_training.model import PackedValueModel @@ -33,7 +33,6 @@ class TrainConfig: device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" - wandb: bool = False @dataclass(slots=True) @@ -99,7 +98,7 @@ def _sample_replay_batch( def train_step( - model: PackedValueModel, + model: torch.nn.Module, replay_buffer: EphemeralReplayBuffer, optimizer: torch.optim.Optimizer, *, @@ -192,7 +191,7 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: def _save_checkpoint( *, checkpoint_dir: Path, - model: PackedValueModel, + model: torch.nn.Module, optimizer: torch.optim.Optimizer, config: TrainConfig, round_idx: int, @@ -217,6 +216,10 @@ def _save_checkpoint( return checkpoint_path +def _count_parameters(model: torch.nn.Module) -> int: + return sum(parameter.numel() for parameter in model.parameters()) + + def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) @@ -227,12 +230,17 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - wandb.init( - project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) + wandb_run = wandb.init( + project="alphapaint", + name=run_dir.name, + dir=run_dir, + config=asdict(config), ) model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) model.eval() + model_parameters = _count_parameters(model) + print(f"params={model_parameters}") optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -543,6 +551,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_forward_seconds": train_forward_seconds, "train_backward_seconds": train_backward_seconds, "train_optimizer_seconds": train_optimizer_seconds, + "model_parameters": model_parameters, "loss_mean": mean_loss, "timestamp": time.time(), } @@ -558,7 +567,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) - wandb.log(record) + wandb_run.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -598,6 +607,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() + wandb_run.finish() return model, losses @@ -635,7 +645,6 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) - parser.add_argument("--wandb", action="store_true", default=True) args = parser.parse_args() return TrainConfig(**vars(args)) From 734204721c885e38ff242c22d26654e52d1d7355 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 01:44:02 -0400 Subject: [PATCH 26/59] Revert "try a mixer style model" This reverts commit df13b0cb56a6b81e329faaa6a6d033a0b9149ec0. --- python/alphapaint_training/__init__.py | 4 +- .../alphapaint_training.pyi | 45 ------ python/alphapaint_training/model.py | 136 ++++++++---------- python/alphapaint_training/train.py | 25 ++-- 4 files changed, 73 insertions(+), 137 deletions(-) delete mode 100644 python/alphapaint_training/alphapaint_training.pyi diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 241fb97..a0d8c3c 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -5,7 +5,7 @@ __path__ = extend_path(__path__, __name__) from .alphapaint_training import EphemeralReplayBuffer, SelfPlay -from .model import PackedValueModel +from .model import PackedValueModel, ResidualBlock, TinyValueNet from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, @@ -26,7 +26,9 @@ "INTRINSIC_COUNT", "OBS_WORDS", "PackedValueModel", + "ResidualBlock", "SelfPlay", + "TinyValueNet", "decode_intrinsics", "decode_packed_board", "decode_packed_board_reference", diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi deleted file mode 100644 index 039e243..0000000 --- a/python/alphapaint_training/alphapaint_training.pyi +++ /dev/null @@ -1,45 +0,0 @@ -from typing import Any - -class EphemeralReplayBuffer: - def __init__(self, capacity: int) -> None: ... - def __len__(self) -> int: ... - def sample(self, batch_size: int, seed: int) -> tuple[Any, Any]: ... - -class SelfPlay: - def __init__( - self, - replay_buffer: EphemeralReplayBuffer, - num_threads: int, - workers_per_thread: int, - seed: int, - *, - max_gpu_evals_per_move: int, - model: Any, - selfplay_precision: str, - ) -> None: ... - def drop(self) -> None: ... - def wait_for(self, target_samples: int) -> int: ... - def games(self) -> int: ... - def gpu_batches(self) -> int: ... - def gpu_evals(self) -> int: ... - def queue_num_batches(self) -> int: ... - def queue_total_slots(self) -> int: ... - def current_outstanding_evals(self) -> int: ... - def max_outstanding_evals(self) -> int: ... - def action_steps(self) -> int: ... - def final_actions(self) -> int: ... - def nonfinal_actions(self) -> int: ... - def completed_turns(self) -> int: ... - def action_turn_count_total(self) -> int: ... - def max_turn_count_seen(self) -> int: ... - def completed_game_actions_total(self) -> int: ... - def max_actions_in_completed_game(self) -> int: ... - def tree_build_nanos(self) -> int: ... - def descent_nanos(self) -> int: ... - def sample_collect_nanos(self) -> int: ... - def replay_push_nanos(self) -> int: ... - def descent_expand_cpu_nanos(self) -> int: ... - def descent_apply_action_nanos(self) -> int: ... - def descent_eval_submit_nanos(self) -> int: ... - def descent_eval_await_nanos(self) -> int: ... - def descent_backup_nanos(self) -> int: ... diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 7c9d05c..65f0208 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -6,94 +6,73 @@ from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, - BOARD_SIDE, INTRINSIC_COUNT, INTRINSIC_SCALE, decode_packed_board, ) -_PATCH_SIZE = 4 -_EMBED_DIM = 16 -_NUM_BLOCKS = 2 -_MLP_DIM = 32 -_HIDDEN_DIM = 32 - - -class _MixingMlp(nn.Module): - def __init__(self, input_dim: int, hidden_dim: int): - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.GELU(), - nn.Linear(hidden_dim, input_dim), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class _MixingBlock(nn.Module): - def __init__(self, *, num_tokens: int, embed_dim: int, mlp_dim: int): +class ResidualBlock(nn.Module): + def __init__(self, width: int): super().__init__() - self.token_norm = nn.LayerNorm(embed_dim) - self.token_mlp = _MixingMlp(num_tokens, mlp_dim) - self.channel_norm = nn.LayerNorm(embed_dim) - self.channel_mlp = _MixingMlp(embed_dim, mlp_dim) + self.norm1 = nn.BatchNorm2d(width) + self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.norm2 = nn.BatchNorm2d(width) + self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.act = nn.ReLU(inplace=True) def forward(self, x: torch.Tensor) -> torch.Tensor: - mixed_tokens = self.token_norm(x).transpose(1, 2) - x = x + self.token_mlp(mixed_tokens).transpose(1, 2) - return x + self.channel_mlp(self.channel_norm(x)) - - -class _ValueNet(nn.Module): - def __init__(self): + residual = x + x = self.norm1(x) + x = self.act(x) + x = self.conv1(x) + x = self.norm2(x) + x = self.act(x) + x = self.conv2(x) + return x + residual + + +class TinyValueNet(nn.Module): + def __init__( + self, + *, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, + ): super().__init__() - if BOARD_SIDE % _PATCH_SIZE != 0: - raise ValueError(f"patch size {_PATCH_SIZE} must divide the board side") - - tokens_per_side = BOARD_SIDE // _PATCH_SIZE - num_tokens = tokens_per_side * tokens_per_side - - self.patch_embed = nn.Conv2d( - BOARD_PLANES, - _EMBED_DIM, - kernel_size=_PATCH_SIZE, - stride=_PATCH_SIZE, - bias=False, + self.stem = nn.Sequential( + nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(width), + nn.ReLU(inplace=True), ) - self.position = nn.Parameter(torch.zeros(1, num_tokens, _EMBED_DIM)) - self.blocks = nn.Sequential( - *( - _MixingBlock( - num_tokens=num_tokens, - embed_dim=_EMBED_DIM, - mlp_dim=_MLP_DIM, - ) - for _ in range(_NUM_BLOCKS) - ) - ) - self.norm = nn.LayerNorm(_EMBED_DIM) + self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) + self.pool = nn.AdaptiveAvgPool2d(1) + self.flatten = nn.Flatten() self.head = nn.Sequential( - nn.Linear(_EMBED_DIM + INTRINSIC_COUNT, _HIDDEN_DIM), - nn.GELU(), - nn.Linear(_HIDDEN_DIM, 1), + nn.Linear(width + INTRINSIC_COUNT, hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(hidden_dim, 1), ) def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: - x = self.patch_embed(board) - x = x.flatten(2).transpose(1, 2) - x = x + self.position + x = self.stem(board) x = self.blocks(x) - x = self.norm(x) - x = x.mean(dim=1) + x = self.pool(x) + x = self.flatten(x) x = torch.cat((x, intrinsics), dim=1) return self.head(x) class PackedValueModel(nn.Module): - def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): + def __init__( + self, + *, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, + board_dtype: torch.dtype = torch.bfloat16, + ): super().__init__() self.board_dtype = board_dtype self._board_buffers: dict[ @@ -105,13 +84,16 @@ def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): self._intrinsic_buffers: dict[ tuple[str, int | None, int, torch.dtype], torch.Tensor ] = {} - self.intrinsic_scales: torch.Tensor self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), persistent=False, ) - self.value_net = _ValueNet() + self.value_net = TinyValueNet( + width=width, + num_blocks=num_blocks, + hidden_dim=hidden_dim, + ) def _buffer_key( self, packed_obs: torch.Tensor, dtype: torch.dtype @@ -123,7 +105,9 @@ def _buffer_key( dtype, ) - def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _ensure_decode_buffers( + self, packed_obs: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: board_key = self._buffer_key(packed_obs, self.board_dtype) board = self._board_buffers.get(board_key) if board is None: @@ -133,7 +117,7 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] else torch.contiguous_format ) board = torch.empty( - (packed_obs.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), + (packed_obs.shape[0], BOARD_PLANES, 32, 32), device=packed_obs.device, dtype=self.board_dtype, memory_format=memory_format, @@ -164,7 +148,11 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] ) self._intrinsic_buffers[intrinsic_key] = intrinsics - decode_packed_board( + return board, intrinsics_fp32, intrinsics + + def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) + board = decode_packed_board( packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype, out=board, @@ -175,7 +163,7 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: - board, intrinsics = self._decode(packed_obs) + board, intrinsics = self.decode(packed_obs) if not torch.is_autocast_enabled(): param_dtype = next(self.value_net.parameters()).dtype board = board.to(dtype=param_dtype) @@ -184,4 +172,4 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: return value.squeeze(-1).tanh() * 10.0 -__all__ = ["PackedValueModel"] +__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index c31e854..a8bebda 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -13,7 +13,7 @@ import torch.nn.functional as F import wandb -from alphapaint_training.alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training import EphemeralReplayBuffer, SelfPlay from alphapaint_training.model import PackedValueModel @@ -33,6 +33,7 @@ class TrainConfig: device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" + wandb: bool = False @dataclass(slots=True) @@ -98,7 +99,7 @@ def _sample_replay_batch( def train_step( - model: torch.nn.Module, + model: PackedValueModel, replay_buffer: EphemeralReplayBuffer, optimizer: torch.optim.Optimizer, *, @@ -191,7 +192,7 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: def _save_checkpoint( *, checkpoint_dir: Path, - model: torch.nn.Module, + model: PackedValueModel, optimizer: torch.optim.Optimizer, config: TrainConfig, round_idx: int, @@ -216,10 +217,6 @@ def _save_checkpoint( return checkpoint_path -def _count_parameters(model: torch.nn.Module) -> int: - return sum(parameter.numel() for parameter in model.parameters()) - - def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) @@ -230,17 +227,12 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - wandb_run = wandb.init( - project="alphapaint", - name=run_dir.name, - dir=run_dir, - config=asdict(config), + wandb.init( + project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) ) model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) model.eval() - model_parameters = _count_parameters(model) - print(f"params={model_parameters}") optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -551,7 +543,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_forward_seconds": train_forward_seconds, "train_backward_seconds": train_backward_seconds, "train_optimizer_seconds": train_optimizer_seconds, - "model_parameters": model_parameters, "loss_mean": mean_loss, "timestamp": time.time(), } @@ -567,7 +558,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) - wandb_run.log(record) + wandb.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -607,7 +598,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() - wandb_run.finish() return model, losses @@ -645,6 +635,7 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) + parser.add_argument("--wandb", action="store_true", default=True) args = parser.parse_args() return TrainConfig(**vars(args)) From 6b27c1142125b04aca599e687e8d65f2aed0597c Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 01:56:37 -0400 Subject: [PATCH 27/59] Reapply "try a mixer style model" This reverts commit 734204721c885e38ff242c22d26654e52d1d7355. --- python/alphapaint_training/__init__.py | 4 +- .../alphapaint_training.pyi | 45 ++++++ python/alphapaint_training/model.py | 136 ++++++++++-------- python/alphapaint_training/train.py | 25 ++-- 4 files changed, 137 insertions(+), 73 deletions(-) create mode 100644 python/alphapaint_training/alphapaint_training.pyi diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index a0d8c3c..241fb97 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -5,7 +5,7 @@ __path__ = extend_path(__path__, __name__) from .alphapaint_training import EphemeralReplayBuffer, SelfPlay -from .model import PackedValueModel, ResidualBlock, TinyValueNet +from .model import PackedValueModel from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, @@ -26,9 +26,7 @@ "INTRINSIC_COUNT", "OBS_WORDS", "PackedValueModel", - "ResidualBlock", "SelfPlay", - "TinyValueNet", "decode_intrinsics", "decode_packed_board", "decode_packed_board_reference", diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi new file mode 100644 index 0000000..039e243 --- /dev/null +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -0,0 +1,45 @@ +from typing import Any + +class EphemeralReplayBuffer: + def __init__(self, capacity: int) -> None: ... + def __len__(self) -> int: ... + def sample(self, batch_size: int, seed: int) -> tuple[Any, Any]: ... + +class SelfPlay: + def __init__( + self, + replay_buffer: EphemeralReplayBuffer, + num_threads: int, + workers_per_thread: int, + seed: int, + *, + max_gpu_evals_per_move: int, + model: Any, + selfplay_precision: str, + ) -> None: ... + def drop(self) -> None: ... + def wait_for(self, target_samples: int) -> int: ... + def games(self) -> int: ... + def gpu_batches(self) -> int: ... + def gpu_evals(self) -> int: ... + def queue_num_batches(self) -> int: ... + def queue_total_slots(self) -> int: ... + def current_outstanding_evals(self) -> int: ... + def max_outstanding_evals(self) -> int: ... + def action_steps(self) -> int: ... + def final_actions(self) -> int: ... + def nonfinal_actions(self) -> int: ... + def completed_turns(self) -> int: ... + def action_turn_count_total(self) -> int: ... + def max_turn_count_seen(self) -> int: ... + def completed_game_actions_total(self) -> int: ... + def max_actions_in_completed_game(self) -> int: ... + def tree_build_nanos(self) -> int: ... + def descent_nanos(self) -> int: ... + def sample_collect_nanos(self) -> int: ... + def replay_push_nanos(self) -> int: ... + def descent_expand_cpu_nanos(self) -> int: ... + def descent_apply_action_nanos(self) -> int: ... + def descent_eval_submit_nanos(self) -> int: ... + def descent_eval_await_nanos(self) -> int: ... + def descent_backup_nanos(self) -> int: ... diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 65f0208..7c9d05c 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -6,73 +6,94 @@ from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, + BOARD_SIDE, INTRINSIC_COUNT, INTRINSIC_SCALE, decode_packed_board, ) -class ResidualBlock(nn.Module): - def __init__(self, width: int): +_PATCH_SIZE = 4 +_EMBED_DIM = 16 +_NUM_BLOCKS = 2 +_MLP_DIM = 32 +_HIDDEN_DIM = 32 + + +class _MixingMlp(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int): + super().__init__() + self.net = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, input_dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class _MixingBlock(nn.Module): + def __init__(self, *, num_tokens: int, embed_dim: int, mlp_dim: int): super().__init__() - self.norm1 = nn.BatchNorm2d(width) - self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.norm2 = nn.BatchNorm2d(width) - self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.act = nn.ReLU(inplace=True) + self.token_norm = nn.LayerNorm(embed_dim) + self.token_mlp = _MixingMlp(num_tokens, mlp_dim) + self.channel_norm = nn.LayerNorm(embed_dim) + self.channel_mlp = _MixingMlp(embed_dim, mlp_dim) def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.norm1(x) - x = self.act(x) - x = self.conv1(x) - x = self.norm2(x) - x = self.act(x) - x = self.conv2(x) - return x + residual - - -class TinyValueNet(nn.Module): - def __init__( - self, - *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, - ): + mixed_tokens = self.token_norm(x).transpose(1, 2) + x = x + self.token_mlp(mixed_tokens).transpose(1, 2) + return x + self.channel_mlp(self.channel_norm(x)) + + +class _ValueNet(nn.Module): + def __init__(self): super().__init__() - self.stem = nn.Sequential( - nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), - nn.BatchNorm2d(width), - nn.ReLU(inplace=True), + if BOARD_SIDE % _PATCH_SIZE != 0: + raise ValueError(f"patch size {_PATCH_SIZE} must divide the board side") + + tokens_per_side = BOARD_SIDE // _PATCH_SIZE + num_tokens = tokens_per_side * tokens_per_side + + self.patch_embed = nn.Conv2d( + BOARD_PLANES, + _EMBED_DIM, + kernel_size=_PATCH_SIZE, + stride=_PATCH_SIZE, + bias=False, ) - self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) - self.pool = nn.AdaptiveAvgPool2d(1) - self.flatten = nn.Flatten() + self.position = nn.Parameter(torch.zeros(1, num_tokens, _EMBED_DIM)) + self.blocks = nn.Sequential( + *( + _MixingBlock( + num_tokens=num_tokens, + embed_dim=_EMBED_DIM, + mlp_dim=_MLP_DIM, + ) + for _ in range(_NUM_BLOCKS) + ) + ) + self.norm = nn.LayerNorm(_EMBED_DIM) self.head = nn.Sequential( - nn.Linear(width + INTRINSIC_COUNT, hidden_dim), - nn.ReLU(inplace=True), - nn.Linear(hidden_dim, 1), + nn.Linear(_EMBED_DIM + INTRINSIC_COUNT, _HIDDEN_DIM), + nn.GELU(), + nn.Linear(_HIDDEN_DIM, 1), ) def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: - x = self.stem(board) + x = self.patch_embed(board) + x = x.flatten(2).transpose(1, 2) + x = x + self.position x = self.blocks(x) - x = self.pool(x) - x = self.flatten(x) + x = self.norm(x) + x = x.mean(dim=1) x = torch.cat((x, intrinsics), dim=1) return self.head(x) class PackedValueModel(nn.Module): - def __init__( - self, - *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, - board_dtype: torch.dtype = torch.bfloat16, - ): + def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): super().__init__() self.board_dtype = board_dtype self._board_buffers: dict[ @@ -84,16 +105,13 @@ def __init__( self._intrinsic_buffers: dict[ tuple[str, int | None, int, torch.dtype], torch.Tensor ] = {} + self.intrinsic_scales: torch.Tensor self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), persistent=False, ) - self.value_net = TinyValueNet( - width=width, - num_blocks=num_blocks, - hidden_dim=hidden_dim, - ) + self.value_net = _ValueNet() def _buffer_key( self, packed_obs: torch.Tensor, dtype: torch.dtype @@ -105,9 +123,7 @@ def _buffer_key( dtype, ) - def _ensure_decode_buffers( - self, packed_obs: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: board_key = self._buffer_key(packed_obs, self.board_dtype) board = self._board_buffers.get(board_key) if board is None: @@ -117,7 +133,7 @@ def _ensure_decode_buffers( else torch.contiguous_format ) board = torch.empty( - (packed_obs.shape[0], BOARD_PLANES, 32, 32), + (packed_obs.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), device=packed_obs.device, dtype=self.board_dtype, memory_format=memory_format, @@ -148,11 +164,7 @@ def _ensure_decode_buffers( ) self._intrinsic_buffers[intrinsic_key] = intrinsics - return board, intrinsics_fp32, intrinsics - - def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) - board = decode_packed_board( + decode_packed_board( packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype, out=board, @@ -163,7 +175,7 @@ def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: - board, intrinsics = self.decode(packed_obs) + board, intrinsics = self._decode(packed_obs) if not torch.is_autocast_enabled(): param_dtype = next(self.value_net.parameters()).dtype board = board.to(dtype=param_dtype) @@ -172,4 +184,4 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: return value.squeeze(-1).tanh() * 10.0 -__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] +__all__ = ["PackedValueModel"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index a8bebda..c31e854 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -13,7 +13,7 @@ import torch.nn.functional as F import wandb -from alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training.alphapaint_training import EphemeralReplayBuffer, SelfPlay from alphapaint_training.model import PackedValueModel @@ -33,7 +33,6 @@ class TrainConfig: device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" - wandb: bool = False @dataclass(slots=True) @@ -99,7 +98,7 @@ def _sample_replay_batch( def train_step( - model: PackedValueModel, + model: torch.nn.Module, replay_buffer: EphemeralReplayBuffer, optimizer: torch.optim.Optimizer, *, @@ -192,7 +191,7 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: def _save_checkpoint( *, checkpoint_dir: Path, - model: PackedValueModel, + model: torch.nn.Module, optimizer: torch.optim.Optimizer, config: TrainConfig, round_idx: int, @@ -217,6 +216,10 @@ def _save_checkpoint( return checkpoint_path +def _count_parameters(model: torch.nn.Module) -> int: + return sum(parameter.numel() for parameter in model.parameters()) + + def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) @@ -227,12 +230,17 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - wandb.init( - project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) + wandb_run = wandb.init( + project="alphapaint", + name=run_dir.name, + dir=run_dir, + config=asdict(config), ) model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) model.eval() + model_parameters = _count_parameters(model) + print(f"params={model_parameters}") optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -543,6 +551,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_forward_seconds": train_forward_seconds, "train_backward_seconds": train_backward_seconds, "train_optimizer_seconds": train_optimizer_seconds, + "model_parameters": model_parameters, "loss_mean": mean_loss, "timestamp": time.time(), } @@ -558,7 +567,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) - wandb.log(record) + wandb_run.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -598,6 +607,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() + wandb_run.finish() return model, losses @@ -635,7 +645,6 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) - parser.add_argument("--wandb", action="store_true", default=True) args = parser.parse_args() return TrainConfig(**vars(args)) From 1e4f428d9dc610e17e4d08538b3052d79e497c3b Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 01:59:18 -0400 Subject: [PATCH 28/59] pls stop fucking deadlocking why is it doing this im so confused --- python/alphapaint_training/train.py | 10 +---- training/src/descent.rs | 2 +- training/src/future.rs | 22 ----------- training/src/lib.rs | 20 ---------- training/src/queue.rs | 57 ++--------------------------- training/src/training.rs | 44 +--------------------- 6 files changed, 7 insertions(+), 148 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index c31e854..1d753af 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -283,10 +283,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: games = selfplay.games() gpu_batches = selfplay.gpu_batches() gpu_evals = selfplay.gpu_evals() - queue_num_batches = selfplay.queue_num_batches() - queue_total_slots = selfplay.queue_total_slots() - current_outstanding_evals = selfplay.current_outstanding_evals() - max_outstanding_evals = selfplay.max_outstanding_evals() action_steps = selfplay.action_steps() final_actions = selfplay.final_actions() nonfinal_actions = selfplay.nonfinal_actions() @@ -497,10 +493,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "gpu_evals_total": gpu_evals, "gpu_evals_added": gpu_evals_added, "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), - "queue_num_batches": queue_num_batches, - "queue_total_slots": queue_total_slots, - "queue_current_outstanding_evals": current_outstanding_evals, - "queue_max_outstanding_evals": max_outstanding_evals, "action_steps_total": action_steps, "action_steps_added": action_steps_added, "final_actions_total": final_actions, @@ -558,7 +550,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print( f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " - f"batches={gpu_batches_added} queue={current_outstanding_evals}/{max_outstanding_evals}/{queue_total_slots}@{queue_num_batches} replay={replay_size} act={action_steps_added} final={final_actions_added} " + f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " f"sp_ms/act=build:{build_ms_per_action:.1f} desc:{descent_ms_per_action:.1f} coll:{collect_ms_per_action:.1f} push:{push_ms_per_action:.1f} " diff --git a/training/src/descent.rs b/training/src/descent.rs index a57f613..3629e6b 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -8,8 +8,8 @@ use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use rand::rngs::SmallRng; use rand::{Rng, RngExt}; -use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; use std::time::Instant; use crate::eval::{EvalCountTracker, Evaluator}; diff --git a/training/src/future.rs b/training/src/future.rs index e536841..f4bf839 100644 --- a/training/src/future.rs +++ b/training/src/future.rs @@ -43,7 +43,6 @@ where queue: &'a GpuJobQueue, ticket: u64, completed: bool, - tracking_outstanding: bool, } impl<'a, A, D, O> GpuEvalFuture<'a, A, D, O> @@ -55,14 +54,12 @@ where { /// Create a new future with an already-submitted ticket. pub fn new(queue: &'a GpuJobQueue, ticket: u64) -> Self { - queue.acquire_outstanding_eval_slot(); // Signal progress on creation since we just submitted signal_progress(); Self { queue, ticket, completed: false, - tracking_outstanding: true, } } } @@ -86,10 +83,6 @@ where if let Some(&output) = this.queue.poll(this.ticket) { this.completed = true; - if this.tracking_outstanding { - this.queue.release_outstanding_eval_slot(); - this.tracking_outstanding = false; - } signal_progress(); Poll::Ready(output) } else { @@ -98,21 +91,6 @@ where } } -impl Drop for GpuEvalFuture<'_, A, D, O> -where - A: Clone + Default + Send + Sync, - D: BatchDim, - D::Larger: Dimension, - O: Copy + Default + Send + Sync, -{ - fn drop(&mut self) { - if self.tracking_outstanding { - self.queue.release_outstanding_eval_slot(); - self.tracking_outstanding = false; - } - } -} - impl GpuJobQueue where A: Clone + Default + Send + Sync, diff --git a/training/src/lib.rs b/training/src/lib.rs index 7c2a821..a17aa7d 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -279,26 +279,6 @@ impl SelfPlay { Ok(self.session()?.games()) } - /// Return the number of queue batch lanes. - fn queue_num_batches(&self) -> PyResult { - Ok(self.session()?.queue_num_batches()) - } - - /// Return the total number of queue slots. - fn queue_total_slots(&self) -> PyResult { - Ok(self.session()?.queue_total_slots()) - } - - /// Return the current number of outstanding eval futures. - fn current_outstanding_evals(&self) -> PyResult { - Ok(self.session()?.current_outstanding_evals()) - } - - /// Return the maximum number of outstanding eval futures seen so far. - fn max_outstanding_evals(&self) -> PyResult { - Ok(self.session()?.max_outstanding_evals()) - } - /// Return the total number of selected actions. fn action_steps(&self) -> PyResult { Ok(self.session()?.action_steps()) diff --git a/training/src/queue.rs b/training/src/queue.rs index 12ee3f1..ed20307 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -8,8 +8,8 @@ //! This enables zero-copy batch slicing for GPU dispatch. use std::cell::UnsafeCell; -use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; use event_listener::Event; use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; @@ -28,8 +28,7 @@ const SLOT_MULTIPLIER: usize = 32; /// Compute queue shape for a given worker count. /// /// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a -/// whole number of `BATCH_SIZE` lanes and is at least -/// `SLOT_MULTIPLIER * num_workers`. +/// whole number of `BATCH_SIZE` lanes and is at least `2 * num_workers`. pub fn queue_shape_for_workers(num_workers: usize) -> (usize, usize) { assert!(num_workers > 0, "num_workers must be > 0"); let min_slots = num_workers.saturating_mul(SLOT_MULTIPLIER).max(BATCH_SIZE); @@ -84,12 +83,6 @@ where /// Event for parking threads when waiting for GPU completion. completion_event: Event, - - /// Number of live eval futures whose results have not been consumed yet. - outstanding_evals: AtomicU64, - - /// High watermark for `outstanding_evals`. - max_outstanding_evals: AtomicU64, } // SAFETY: Access to queue state is synchronized via ticket ownership and atomics. @@ -165,7 +158,7 @@ where /// Creates a new job queue with the given observation shape, worker count, /// and dispatch callback. /// - /// Queue storage is provisioned to at least `SLOT_MULTIPLIER * num_workers` slots, + /// Queue storage is provisioned to at least `2 * num_workers` slots, /// rounded up to a whole number of batches. /// /// The callback is invoked when a batch of BATCH_SIZE jobs is ready. @@ -192,8 +185,6 @@ where batch_complete, outputs, completion_event: Event::new(), - outstanding_evals: AtomicU64::new(0), - max_outstanding_evals: AtomicU64::new(0), }); Self { @@ -216,48 +207,6 @@ where self.total_slots } - #[inline] - pub fn current_outstanding_evals(&self) -> u64 { - self.state.outstanding_evals.load(Ordering::Acquire) - } - - #[inline] - pub fn max_outstanding_evals(&self) -> u64 { - self.state.max_outstanding_evals.load(Ordering::Acquire) - } - - #[inline] - pub(crate) fn acquire_outstanding_eval_slot(&self) { - let current = self.state.outstanding_evals.fetch_add(1, Ordering::AcqRel) + 1; - - let mut observed_max = self.state.max_outstanding_evals.load(Ordering::Acquire); - while current > observed_max { - match self.state.max_outstanding_evals.compare_exchange_weak( - observed_max, - current, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => break, - Err(new_max) => observed_max = new_max, - } - } - - if current > self.total_slots as u64 { - self.state.outstanding_evals.fetch_sub(1, Ordering::AcqRel); - panic!( - "queue overflow risk: {} outstanding evals exceeded {} total slots", - current, self.total_slots - ); - } - } - - #[inline] - pub(crate) fn release_outstanding_eval_slot(&self) { - let previous = self.state.outstanding_evals.fetch_sub(1, Ordering::AcqRel); - debug_assert!(previous > 0, "outstanding eval counter underflowed"); - } - /// Submit a job by writing an observation via callback. /// /// The callback receives a mutable view into the queue's contiguous storage diff --git a/training/src/training.rs b/training/src/training.rs index 988fc08..4e20e75 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -11,12 +11,12 @@ use ndarray::{ArrayView, Ix1, Ix2}; use rand::SeedableRng; use rand_chacha::ChaCha8Rng; -use crate::BatchDim; use crate::eval::GpuEvaluator; use crate::executor::Executor; use crate::queue::{BatchCompletion, GpuJobQueue}; use crate::replay_buffer::ReplayBuffer; -use crate::worker::{SelfPlayMetrics, WorkerConfig, worker_loop_forever}; +use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; +use crate::BatchDim; /// Shared control state for the persistent self-play session. struct SessionControl { @@ -145,10 +145,6 @@ impl Default for SessionConfig { /// Trait-object wrapper so we can call `notify_all()` on the queue. trait QueueNotify: Send + Sync { fn notify_all(&self); - fn num_batches(&self) -> usize; - fn total_slots(&self) -> usize; - fn current_outstanding_evals(&self) -> u64; - fn max_outstanding_evals(&self) -> u64; } impl QueueNotify for GpuJobQueue @@ -160,22 +156,6 @@ where fn notify_all(&self) { GpuJobQueue::notify_all(self); } - - fn num_batches(&self) -> usize { - GpuJobQueue::num_batches(self) - } - - fn total_slots(&self) -> usize { - GpuJobQueue::total_slots(self) - } - - fn current_outstanding_evals(&self) -> u64 { - GpuJobQueue::current_outstanding_evals(self) - } - - fn max_outstanding_evals(&self) -> u64 { - GpuJobQueue::max_outstanding_evals(self) - } } /// A persistent self-play session that owns worker threads and preserves @@ -299,26 +279,6 @@ impl SelfPlaySession { self.control.games_completed.load(Ordering::Acquire) } - /// Return the number of queue batch lanes. - pub fn queue_num_batches(&self) -> usize { - self.queue_notify.num_batches() - } - - /// Return the total number of queue slots. - pub fn queue_total_slots(&self) -> usize { - self.queue_notify.total_slots() - } - - /// Return the current number of outstanding eval futures. - pub fn current_outstanding_evals(&self) -> u64 { - self.queue_notify.current_outstanding_evals() - } - - /// Return the maximum number of outstanding eval futures seen so far. - pub fn max_outstanding_evals(&self) -> u64 { - self.queue_notify.max_outstanding_evals() - } - /// Return the total number of selected actions. pub fn action_steps(&self) -> u64 { self.control.action_steps.load(Ordering::Acquire) From 5491ad7f36ad63a8eb25792de782f3ede80a340e Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 02:37:15 -0400 Subject: [PATCH 29/59] Revert "Reapply "try a mixer style model"" This reverts commit 6b27c1142125b04aca599e687e8d65f2aed0597c. --- python/alphapaint_training/__init__.py | 4 +- .../alphapaint_training.pyi | 45 ------ python/alphapaint_training/model.py | 136 ++++++++---------- python/alphapaint_training/train.py | 25 ++-- 4 files changed, 73 insertions(+), 137 deletions(-) delete mode 100644 python/alphapaint_training/alphapaint_training.pyi diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 241fb97..a0d8c3c 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -5,7 +5,7 @@ __path__ = extend_path(__path__, __name__) from .alphapaint_training import EphemeralReplayBuffer, SelfPlay -from .model import PackedValueModel +from .model import PackedValueModel, ResidualBlock, TinyValueNet from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, @@ -26,7 +26,9 @@ "INTRINSIC_COUNT", "OBS_WORDS", "PackedValueModel", + "ResidualBlock", "SelfPlay", + "TinyValueNet", "decode_intrinsics", "decode_packed_board", "decode_packed_board_reference", diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi deleted file mode 100644 index 039e243..0000000 --- a/python/alphapaint_training/alphapaint_training.pyi +++ /dev/null @@ -1,45 +0,0 @@ -from typing import Any - -class EphemeralReplayBuffer: - def __init__(self, capacity: int) -> None: ... - def __len__(self) -> int: ... - def sample(self, batch_size: int, seed: int) -> tuple[Any, Any]: ... - -class SelfPlay: - def __init__( - self, - replay_buffer: EphemeralReplayBuffer, - num_threads: int, - workers_per_thread: int, - seed: int, - *, - max_gpu_evals_per_move: int, - model: Any, - selfplay_precision: str, - ) -> None: ... - def drop(self) -> None: ... - def wait_for(self, target_samples: int) -> int: ... - def games(self) -> int: ... - def gpu_batches(self) -> int: ... - def gpu_evals(self) -> int: ... - def queue_num_batches(self) -> int: ... - def queue_total_slots(self) -> int: ... - def current_outstanding_evals(self) -> int: ... - def max_outstanding_evals(self) -> int: ... - def action_steps(self) -> int: ... - def final_actions(self) -> int: ... - def nonfinal_actions(self) -> int: ... - def completed_turns(self) -> int: ... - def action_turn_count_total(self) -> int: ... - def max_turn_count_seen(self) -> int: ... - def completed_game_actions_total(self) -> int: ... - def max_actions_in_completed_game(self) -> int: ... - def tree_build_nanos(self) -> int: ... - def descent_nanos(self) -> int: ... - def sample_collect_nanos(self) -> int: ... - def replay_push_nanos(self) -> int: ... - def descent_expand_cpu_nanos(self) -> int: ... - def descent_apply_action_nanos(self) -> int: ... - def descent_eval_submit_nanos(self) -> int: ... - def descent_eval_await_nanos(self) -> int: ... - def descent_backup_nanos(self) -> int: ... diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 7c9d05c..65f0208 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -6,94 +6,73 @@ from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, - BOARD_SIDE, INTRINSIC_COUNT, INTRINSIC_SCALE, decode_packed_board, ) -_PATCH_SIZE = 4 -_EMBED_DIM = 16 -_NUM_BLOCKS = 2 -_MLP_DIM = 32 -_HIDDEN_DIM = 32 - - -class _MixingMlp(nn.Module): - def __init__(self, input_dim: int, hidden_dim: int): - super().__init__() - self.net = nn.Sequential( - nn.Linear(input_dim, hidden_dim), - nn.GELU(), - nn.Linear(hidden_dim, input_dim), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.net(x) - - -class _MixingBlock(nn.Module): - def __init__(self, *, num_tokens: int, embed_dim: int, mlp_dim: int): +class ResidualBlock(nn.Module): + def __init__(self, width: int): super().__init__() - self.token_norm = nn.LayerNorm(embed_dim) - self.token_mlp = _MixingMlp(num_tokens, mlp_dim) - self.channel_norm = nn.LayerNorm(embed_dim) - self.channel_mlp = _MixingMlp(embed_dim, mlp_dim) + self.norm1 = nn.BatchNorm2d(width) + self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.norm2 = nn.BatchNorm2d(width) + self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) + self.act = nn.ReLU(inplace=True) def forward(self, x: torch.Tensor) -> torch.Tensor: - mixed_tokens = self.token_norm(x).transpose(1, 2) - x = x + self.token_mlp(mixed_tokens).transpose(1, 2) - return x + self.channel_mlp(self.channel_norm(x)) - - -class _ValueNet(nn.Module): - def __init__(self): + residual = x + x = self.norm1(x) + x = self.act(x) + x = self.conv1(x) + x = self.norm2(x) + x = self.act(x) + x = self.conv2(x) + return x + residual + + +class TinyValueNet(nn.Module): + def __init__( + self, + *, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, + ): super().__init__() - if BOARD_SIDE % _PATCH_SIZE != 0: - raise ValueError(f"patch size {_PATCH_SIZE} must divide the board side") - - tokens_per_side = BOARD_SIDE // _PATCH_SIZE - num_tokens = tokens_per_side * tokens_per_side - - self.patch_embed = nn.Conv2d( - BOARD_PLANES, - _EMBED_DIM, - kernel_size=_PATCH_SIZE, - stride=_PATCH_SIZE, - bias=False, + self.stem = nn.Sequential( + nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), + nn.BatchNorm2d(width), + nn.ReLU(inplace=True), ) - self.position = nn.Parameter(torch.zeros(1, num_tokens, _EMBED_DIM)) - self.blocks = nn.Sequential( - *( - _MixingBlock( - num_tokens=num_tokens, - embed_dim=_EMBED_DIM, - mlp_dim=_MLP_DIM, - ) - for _ in range(_NUM_BLOCKS) - ) - ) - self.norm = nn.LayerNorm(_EMBED_DIM) + self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) + self.pool = nn.AdaptiveAvgPool2d(1) + self.flatten = nn.Flatten() self.head = nn.Sequential( - nn.Linear(_EMBED_DIM + INTRINSIC_COUNT, _HIDDEN_DIM), - nn.GELU(), - nn.Linear(_HIDDEN_DIM, 1), + nn.Linear(width + INTRINSIC_COUNT, hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(hidden_dim, 1), ) def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: - x = self.patch_embed(board) - x = x.flatten(2).transpose(1, 2) - x = x + self.position + x = self.stem(board) x = self.blocks(x) - x = self.norm(x) - x = x.mean(dim=1) + x = self.pool(x) + x = self.flatten(x) x = torch.cat((x, intrinsics), dim=1) return self.head(x) class PackedValueModel(nn.Module): - def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): + def __init__( + self, + *, + width: int = 8, + num_blocks: int = 1, + hidden_dim: int = 32, + board_dtype: torch.dtype = torch.bfloat16, + ): super().__init__() self.board_dtype = board_dtype self._board_buffers: dict[ @@ -105,13 +84,16 @@ def __init__(self, *, board_dtype: torch.dtype = torch.bfloat16): self._intrinsic_buffers: dict[ tuple[str, int | None, int, torch.dtype], torch.Tensor ] = {} - self.intrinsic_scales: torch.Tensor self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), persistent=False, ) - self.value_net = _ValueNet() + self.value_net = TinyValueNet( + width=width, + num_blocks=num_blocks, + hidden_dim=hidden_dim, + ) def _buffer_key( self, packed_obs: torch.Tensor, dtype: torch.dtype @@ -123,7 +105,9 @@ def _buffer_key( dtype, ) - def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def _ensure_decode_buffers( + self, packed_obs: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: board_key = self._buffer_key(packed_obs, self.board_dtype) board = self._board_buffers.get(board_key) if board is None: @@ -133,7 +117,7 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] else torch.contiguous_format ) board = torch.empty( - (packed_obs.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE), + (packed_obs.shape[0], BOARD_PLANES, 32, 32), device=packed_obs.device, dtype=self.board_dtype, memory_format=memory_format, @@ -164,7 +148,11 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] ) self._intrinsic_buffers[intrinsic_key] = intrinsics - decode_packed_board( + return board, intrinsics_fp32, intrinsics + + def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) + board = decode_packed_board( packed_obs[:, :BOARD_CELLS], dtype=self.board_dtype, out=board, @@ -175,7 +163,7 @@ def _decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor] return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: - board, intrinsics = self._decode(packed_obs) + board, intrinsics = self.decode(packed_obs) if not torch.is_autocast_enabled(): param_dtype = next(self.value_net.parameters()).dtype board = board.to(dtype=param_dtype) @@ -184,4 +172,4 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: return value.squeeze(-1).tanh() * 10.0 -__all__ = ["PackedValueModel"] +__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 1d753af..482d97b 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -13,7 +13,7 @@ import torch.nn.functional as F import wandb -from alphapaint_training.alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training import EphemeralReplayBuffer, SelfPlay from alphapaint_training.model import PackedValueModel @@ -33,6 +33,7 @@ class TrainConfig: device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" + wandb: bool = False @dataclass(slots=True) @@ -98,7 +99,7 @@ def _sample_replay_batch( def train_step( - model: torch.nn.Module, + model: PackedValueModel, replay_buffer: EphemeralReplayBuffer, optimizer: torch.optim.Optimizer, *, @@ -191,7 +192,7 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: def _save_checkpoint( *, checkpoint_dir: Path, - model: torch.nn.Module, + model: PackedValueModel, optimizer: torch.optim.Optimizer, config: TrainConfig, round_idx: int, @@ -216,10 +217,6 @@ def _save_checkpoint( return checkpoint_path -def _count_parameters(model: torch.nn.Module) -> int: - return sum(parameter.numel() for parameter in model.parameters()) - - def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) @@ -230,17 +227,12 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - wandb_run = wandb.init( - project="alphapaint", - name=run_dir.name, - dir=run_dir, - config=asdict(config), + wandb.init( + project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) ) model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) model.eval() - model_parameters = _count_parameters(model) - print(f"params={model_parameters}") optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -543,7 +535,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_forward_seconds": train_forward_seconds, "train_backward_seconds": train_backward_seconds, "train_optimizer_seconds": train_optimizer_seconds, - "model_parameters": model_parameters, "loss_mean": mean_loss, "timestamp": time.time(), } @@ -559,7 +550,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) - wandb_run.log(record) + wandb.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -599,7 +590,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() - wandb_run.finish() return model, losses @@ -637,6 +627,7 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) + parser.add_argument("--wandb", action="store_true", default=True) args = parser.parse_args() return TrainConfig(**vars(args)) From 8fcf64b0f929268855f04d29a0e29bcfdfa86dc7 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 03:42:27 -0400 Subject: [PATCH 30/59] my fuck ass queue has so many fucking deadlocks we preserve game state between training steps. we were, before this commit, deadlocking when we were preserving gpu job queue futures in flight over train boundaries. this fixes that issue by basically just completing all batches in the wait_for function on SelfPlay. essentially, after all the threads stop sending GPU eval requests, we clear out the queue and potentially run it with trash values to ensure that we never get into a weird state with a half/almost full queue. --- training/src/queue.rs | 93 +++++++++++++++++++++++++++++++++++++++- training/src/training.rs | 17 ++++++-- 2 files changed, 106 insertions(+), 4 deletions(-) diff --git a/training/src/queue.rs b/training/src/queue.rs index ed20307..6ad0fac 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -11,7 +11,7 @@ use std::cell::UnsafeCell; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use event_listener::Event; +use event_listener::{Event, Listener}; use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; use crate::BatchDim; @@ -81,6 +81,9 @@ where /// Output buffer. Size = `total_slots`. outputs: Box<[UnsafeCell]>, + /// Number of dispatched batches whose completion callback has not run yet. + inflight_batches: AtomicU64, + /// Event for parking threads when waiting for GPU completion. completion_event: Event, } @@ -121,6 +124,8 @@ where self.state.batch_complete[self.batch_idx].store(self.batch_end_ticket, Ordering::Release); self.state.batch_writes[self.batch_idx].store(0, Ordering::Relaxed); + let previous = self.state.inflight_batches.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "inflight batch counter underflowed"); self.state.completion_event.notify(usize::MAX); } } @@ -184,6 +189,7 @@ where batch_writes, batch_complete, outputs, + inflight_batches: AtomicU64::new(0), completion_event: Event::new(), }); @@ -266,9 +272,59 @@ where batch_end_ticket, }; + self.state.inflight_batches.fetch_add(1, Ordering::AcqRel); (self.dispatch)(batch_idx, batch_view, completion); } + fn wait_until_idle(&self) { + loop { + if self.state.inflight_batches.load(Ordering::Acquire) == 0 { + return; + } + let listener = self.state.completion_event.listen(); + if self.state.inflight_batches.load(Ordering::Acquire) == 0 { + return; + } + listener.wait(); + } + } + + /// Flush the open partial batch, if any, and wait for all in-flight GPU work. + /// + /// # Safety + /// + /// The caller must guarantee exclusive access to queue mutation while this + /// runs: no thread may call `submit`, `eval`, or any other quiesce function. + /// The intended callsite is after all session worker threads have left the + /// executor polling loop at a pause boundary. + pub unsafe fn quiesce_exclusive(&self) -> bool { + let next_ticket = self.write_ticket.load(Ordering::Acquire); + let remainder = (next_ticket % BATCH_SIZE as u64) as usize; + let mut flushed = false; + + if remainder != 0 { + let batch_number = next_ticket / BATCH_SIZE as u64; + let batch_idx = (batch_number as usize) % self.num_batches; + let writes = self.state.batch_writes[batch_idx].load(Ordering::Acquire) as usize; + + debug_assert_eq!( + writes, remainder, + "partial batch writes should match the open batch remainder" + ); + + if writes != 0 { + let batch_end_ticket = (batch_number + 1) * BATCH_SIZE as u64; + self.write_ticket.store(batch_end_ticket, Ordering::Release); + + self.dispatch_batch(batch_idx, batch_end_ticket - 1); + flushed = true; + } + } + + self.wait_until_idle(); + flushed + } + /// Poll for a result. Returns Some(&O) if ready, None if still pending. pub fn poll(&self, ticket: u64) -> Option<&O> { let batch_idx = ((ticket as usize) / BATCH_SIZE) % self.num_batches; @@ -371,6 +427,41 @@ mod tests { } } + #[test] + fn test_quiesce_exclusive_completes_written_tickets() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input + 100; + } + completion.complete(&outputs); + }, + )); + + let tickets: Vec = (0..8u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_none(), + "partial batch should not be ready" + ); + } + + unsafe { + assert!(queue.quiesce_exclusive()); + } + + for (i, &ticket) in tickets.iter().enumerate() { + let result = queue.poll(ticket).expect("flushed ticket should be ready"); + assert_eq!(*result, i as u64 + 100); + } + } + #[test] fn test_multiple_batches() { let num_jobs = BATCH_SIZE * 3; diff --git a/training/src/training.rs b/training/src/training.rs index 4e20e75..12552ad 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -144,6 +144,7 @@ impl Default for SessionConfig { /// Trait-object wrapper so we can call `notify_all()` on the queue. trait QueueNotify: Send + Sync { + unsafe fn quiesce_exclusive(&self) -> bool; fn notify_all(&self); } @@ -153,6 +154,10 @@ where D: BatchDim, O: Copy + Default + Send + Sync, { + unsafe fn quiesce_exclusive(&self) -> bool { + GpuJobQueue::quiesce_exclusive(self) + } + fn notify_all(&self) { GpuJobQueue::notify_all(self); } @@ -266,6 +271,13 @@ impl SelfPlaySession { } } + // SAFETY: `wait_for` has already stopped the session and waited until + // every session thread has left the executor loop (`active_pollers == 0`), + // so no thread can submit new queue writes or race another quiesce call. + unsafe { + self.queue_notify.quiesce_exclusive(); + } + self.control.samples_collected.load(Ordering::Acquire) } @@ -430,8 +442,7 @@ fn session_thread_main( descent_backup_nanos: control.descent_backup_nanos.clone(), }; - let mut futures: Vec + '_>>> = (0 - ..config.workers_per_thread) + let mut futures: Vec<_> = (0..config.workers_per_thread) .map(|i| { let metrics = metrics.clone(); let mut rng = ChaCha8Rng::seed_from_u64(base_seed + i as u64); @@ -448,7 +459,7 @@ fn session_thread_main( ) .await; }; - Box::pin(fut) as std::pin::Pin + '_>> + Box::pin(fut) }) .collect(); From 13d53d98b20e3c6713e6da1b7ff4820cb00df552 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 04:31:40 -0400 Subject: [PATCH 31/59] hopefully fix last bug --- AGENTS.md | 6 ++- training/src/descent.rs | 95 +++++++++++++++++++++++++++++++++++------ 2 files changed, 88 insertions(+), 13 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 98229c1..ac2ff11 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,4 +9,8 @@ our search impl will also be in rust. check out /mnt/github/TheConverseEngineer/AlphaSnake/ for last years impl -`just test` runs tests +`just test` runs tests. u should use this rather than doign it yourself bc this auto installs the package + +`nix develop` GETS YOU CUDA +`maturin develop --release` to bring python bindings in. you can `--manifest-path training/Cargo.toml` etc to choose the package to reinstall +`PYTHONPATH=python` is uesful sometimes diff --git a/training/src/descent.rs b/training/src/descent.rs index 3629e6b..0fa3da2 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -489,6 +489,34 @@ impl SearchNode { mod tests { use super::*; use alpha_paint::board::board_structs::Player; + use rand::rngs::SmallRng; + use rand::SeedableRng; + + struct NeverEval; + + impl Evaluator for NeverEval { + fn evaluate(&self, _board: Board) -> impl std::future::Future { + std::future::ready(0.0) + } + } + + impl EvalCountTracker for NeverEval { + fn reset_eval_count(&self) {} + + fn eval_count(&self) -> u64 { + 0 + } + } + + fn zero_timing() -> SearchTimingMetrics { + SearchTimingMetrics { + expand_cpu_nanos: Arc::new(AtomicU64::new(0)), + apply_action_nanos: Arc::new(AtomicU64::new(0)), + eval_submit_nanos: Arc::new(AtomicU64::new(0)), + eval_await_nanos: Arc::new(AtomicU64::new(0)), + backup_nanos: Arc::new(AtomicU64::new(0)), + } + } fn board_with_turn_count(turn_count: usize) -> Board { Board::from_fen(&format!( @@ -597,6 +625,46 @@ mod tests { root.collect_samples_excluding_child(&mut board, 0, &mut dropped_samples); assert_eq!(dropped_samples.len(), 2); } + + #[test] + fn resolved_forced_loss_stops_descent() { + let board = board_with_turn_count(0); + let root_actions: Vec = board + .get_valid_actions() + .into_iter() + .filter(|action| !action.is_final()) + .take(2) + .copied() + .collect(); + assert_eq!(root_actions.len(), 2); + + let mut root = SearchNode::new(-1.0, -1, true); + for &action in &root_actions { + root.children.push(ChildData { + action, + child_value: -1.0, + entrance_count: 0, + node: Some(Box::new(SearchNode::new(-1.0, -1, true))), + }); + } + + let evaluator = NeverEval; + let mut tree = GameSearchTree { + root_node: Box::new(root), + root_state: board, + rng: SmallRng::seed_from_u64(7), + evaluator: &evaluator, + timing: zero_timing(), + }; + + let dcv = if tree.root_state.is_white_turn() { + 1 + } else { + -1 + }; + let best_action = tree.get_best_action_index(); + assert!(tree.should_stop_descent(dcv, best_action)); + } } pub struct GameSearchTree<'a, E: Evaluator> { @@ -608,6 +676,19 @@ pub struct GameSearchTree<'a, E: Evaluator> { } impl<'a, E: Evaluator> GameSearchTree<'a, E> { + fn should_stop_descent(&self, dcv: i32, best_action_id: usize) -> bool { + if self.root_node.resolved { + return true; + } + + let best_child = &self.root_node.children[best_action_id]; + if best_child.entrance_count >= 6000 && best_child.completion_value() != -dcv { + return true; + } + + best_child.completion_value() == dcv + } + fn safest_action(&mut self) -> (usize, Action) { let val = if self.root_state.is_white_turn() { self.root_node @@ -714,12 +795,7 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { } for _epoch in 0..iterations { let ba = self.get_best_action_index(); - if self.root_node.children[ba].entrance_count >= 6000 - && self.root_node.children[ba].completion_value() != -dcv - { - break; - } - if self.root_node.children[ba].completion_value() == dcv { + if self.should_stop_descent(dcv, ba) { break; } self.root_node @@ -818,12 +894,7 @@ impl<'a, E: EvalCountTracker> GameSearchTree<'a, E> { } let ba = self.get_best_action_index(); - if self.root_node.children[ba].entrance_count >= 6000 - && self.root_node.children[ba].completion_value() != -dcv - { - break; - } - if self.root_node.children[ba].completion_value() == dcv { + if self.should_stop_descent(dcv, ba) { break; } self.root_node From c4e9b144e4c8ac3f05650365e34cfe0f32adbaa6 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 05:01:54 -0400 Subject: [PATCH 32/59] fix weird ass bug in descent where node has no children --- alpha_paint/src/board/board_impl.rs | 2 -- training/src/worker.rs | 51 +++++++++++++++++------------ 2 files changed, 30 insertions(+), 23 deletions(-) diff --git a/alpha_paint/src/board/board_impl.rs b/alpha_paint/src/board/board_impl.rs index 1fa4e58..bd0a716 100644 --- a/alpha_paint/src/board/board_impl.rs +++ b/alpha_paint/src/board/board_impl.rs @@ -530,8 +530,6 @@ impl Board { let white_hills = self.tiles.controlled_hill_count::(); let black_hills = self.tiles.controlled_hill_count::(); let total_hills = self.hills.len(); - // TODO: is this fine to assert lol - assert!(total_hills > 0); let domination_threshold = DOMINATION_WIN_THRESHOLD * (total_hills as f64); let domination_win = if (white_hills as f64) >= domination_threshold { diff --git a/training/src/worker.rs b/training/src/worker.rs index 3d63c8e..273b76c 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -3,16 +3,16 @@ //! Each worker runs Descent search, collects training samples from //! internal tree nodes, and pushes them to the replay buffer. -use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::Arc; use std::time::Instant; use ndarray::Ix1; use rand::rngs::SmallRng; use rand::{Rng, RngExt, SeedableRng}; -use alpha_paint::TRAINING_START_FENS; use alpha_paint::board::Board; +use alpha_paint::TRAINING_START_FENS; use crate::descent::{GameSearchTree, SearchTimingMetrics, TreeLearningSample}; use crate::eval::{CountingEvaluator, Evaluator}; @@ -147,6 +147,23 @@ fn step_tree_and_collect_dropped_samples_timed( samples } +fn finish_game_with_current_tree( + tree: &GameSearchTree<'_, E>, + replay_buffer: &ReplayBuffer, + metrics: &SelfPlayMetrics, + action_steps_in_game: u64, +) { + push_samples_to_replay( + replay_buffer, + metrics, + collect_tree_learning_samples_timed(tree, metrics), + ); + metrics + .completed_game_actions_total + .fetch_add(action_steps_in_game, Ordering::AcqRel); + update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); +} + /// Run a single self-play game with tree learning. /// /// At each move: @@ -206,6 +223,15 @@ async fn play_game( .fetch_add(elapsed_nanos(descent_started_at), Ordering::AcqRel); // Select action via ordinal distribution + if tree.root_node.children.is_empty() { + eprintln!( + "how the fuck did we get here?? board has no children. FEN={}", + tree.root_state + ); + finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); + break; + } + let action_id = tree.ordinal_select(); let action = tree.root_node.children[action_id].action; action_steps_in_game += 1; @@ -226,28 +252,11 @@ async fn play_game( match outcome { alpha_paint::board::ApplyActionOutcome::Terminal { .. } | alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { - push_samples_to_replay( - replay_buffer, - metrics, - collect_tree_learning_samples_timed(&tree, metrics), - ); - metrics - .completed_game_actions_total - .fetch_add(action_steps_in_game, Ordering::AcqRel); - update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); + finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); break; } alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { - // Terminal via play-instead - push_samples_to_replay( - replay_buffer, - metrics, - collect_tree_learning_samples_timed(&tree, metrics), - ); - metrics - .completed_game_actions_total - .fetch_add(action_steps_in_game, Ordering::AcqRel); - update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); + finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { From d1f2c57a28b46975e18f468b82df8ad91e09a64d Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 15:03:44 -0400 Subject: [PATCH 33/59] pretrain --- python/alphapaint_training/__init__.py | 7 +- .../alphapaint_training.pyi | 57 +++++ python/alphapaint_training/model.py | 7 +- python/alphapaint_training/train.py | 237 ++++++++++++++++-- training/src/cudagraph.rs | 2 +- training/src/executor.rs | 2 +- training/src/future.rs | 2 +- training/src/lib.rs | 121 ++++++++- training/src/observation.rs | 2 +- training/src/replay_buffer.rs | 2 +- training/src/training.rs | 37 +++ training/src/worker.rs | 55 +++- 12 files changed, 492 insertions(+), 39 deletions(-) create mode 100644 python/alphapaint_training/alphapaint_training.pyi diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index a0d8c3c..6715edf 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -4,7 +4,11 @@ __path__ = extend_path(__path__, __name__) -from .alphapaint_training import EphemeralReplayBuffer, SelfPlay +from .alphapaint_training import ( + EphemeralReplayBuffer, + SelfPlay, + sample_random_terminal_batch, +) from .model import PackedValueModel, ResidualBlock, TinyValueNet from .packed_obs import ( BOARD_CELLS, @@ -33,4 +37,5 @@ "decode_packed_board", "decode_packed_board_reference", "decode_packed_observation", + "sample_random_terminal_batch", ] diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi new file mode 100644 index 0000000..94213d7 --- /dev/null +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -0,0 +1,57 @@ +from __future__ import annotations + +import numpy as np +import numpy.typing as npt + +class EphemeralReplayBuffer: + def __init__(self, capacity: int) -> None: ... + def __len__(self) -> int: ... + @property + def capacity(self) -> int: ... + def sample( + self, n: int, seed: int + ) -> tuple[npt.NDArray[np.uint16], npt.NDArray[np.float32]]: ... + +class SelfPlay: + def __init__( + self, + replay_buffer: EphemeralReplayBuffer, + num_threads: int, + workers_per_thread: int, + seed: int, + *, + max_gpu_evals_per_move: int = 4096, + model: object, + selfplay_precision: str = "bf16", + ) -> None: ... + def start(self) -> None: ... + def wait_for(self, target_samples: int) -> int: ... + def samples(self) -> int: ... + def games(self) -> int: ... + def action_steps(self) -> int: ... + def final_actions(self) -> int: ... + def nonfinal_actions(self) -> int: ... + def completed_turns(self) -> int: ... + def action_turn_count_total(self) -> int: ... + def max_turn_count_seen(self) -> int: ... + def completed_game_actions_total(self) -> int: ... + def max_actions_in_completed_game(self) -> int: ... + def completed_game_turn_count_total(self) -> int: ... + def max_turn_count_in_completed_game(self) -> int: ... + def take_completed_game_turn_counts(self) -> list[int]: ... + def tree_build_nanos(self) -> int: ... + def descent_nanos(self) -> int: ... + def sample_collect_nanos(self) -> int: ... + def replay_push_nanos(self) -> int: ... + def descent_expand_cpu_nanos(self) -> int: ... + def descent_apply_action_nanos(self) -> int: ... + def descent_eval_submit_nanos(self) -> int: ... + def descent_eval_await_nanos(self) -> int: ... + def descent_backup_nanos(self) -> int: ... + def gpu_batches(self) -> int: ... + def gpu_evals(self) -> int: ... + def drop(self) -> None: ... + +def sample_random_terminal_batch( + n: int, seed: int +) -> tuple[npt.NDArray[np.uint16], npt.NDArray[np.float32]]: ... diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 65f0208..3ae2aee 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import cast + import torch from torch import nn @@ -157,8 +159,9 @@ def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: dtype=self.board_dtype, out=board, ) + intrinsic_scales = cast(torch.Tensor, self.intrinsic_scales) intrinsics_fp32.copy_(packed_obs[:, BOARD_CELLS:]) - intrinsics_fp32.div_(self.intrinsic_scales) + intrinsics_fp32.div_(intrinsic_scales) intrinsics.copy_(intrinsics_fp32) return board, intrinsics @@ -169,7 +172,7 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: board = board.to(dtype=param_dtype) intrinsics = intrinsics.to(dtype=param_dtype) value = self.value_net(board, intrinsics) - return value.squeeze(-1).tanh() * 10.0 + return value.squeeze(-1).tanh() * 7.0 __all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 482d97b..970b25f 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -13,7 +13,11 @@ import torch.nn.functional as F import wandb -from alphapaint_training import EphemeralReplayBuffer, SelfPlay +from alphapaint_training import ( + EphemeralReplayBuffer, + SelfPlay, + sample_random_terminal_batch, +) from alphapaint_training.model import PackedValueModel @@ -28,6 +32,9 @@ class TrainConfig: workers_per_thread: int = 8 max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 + pretrain_terminal_samples: int = 5_000_000 + pretrain_batch_size: int = 24_576 + pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" device: str = "cuda" @@ -69,6 +76,26 @@ def _to_channels_last(module: torch.nn.Module) -> torch.nn.Module: return module +def _numpy_batch_to_device( + obs_np: np.ndarray, + values_np: np.ndarray, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float]: + h2d_started_at = time.perf_counter() + obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( + device=device, + dtype=torch.uint16, + non_blocking=True, + ) + values = torch.from_numpy(np.asarray(values_np, dtype=np.float32)).to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + h2d_seconds = time.perf_counter() - h2d_started_at + return obs, values, h2d_seconds + + def _sample_replay_batch( replay_buffer: EphemeralReplayBuffer, batch_size: int, @@ -83,35 +110,34 @@ def _sample_replay_batch( obs_np, values_np = replay_buffer.sample(actual_batch_size, seed) sample_seconds = time.perf_counter() - sample_started_at - h2d_started_at = time.perf_counter() - obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( - device=device, - dtype=torch.uint16, - non_blocking=True, - ) - values = torch.from_numpy(np.asarray(values_np, dtype=np.float32)).to( - device=device, - dtype=torch.float32, - non_blocking=True, - ) - h2d_seconds = time.perf_counter() - h2d_started_at + obs, values, h2d_seconds = _numpy_batch_to_device(obs_np, values_np, device) return obs, values, sample_seconds, h2d_seconds -def train_step( +def _sample_terminal_batch( + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float, float]: + sample_started_at = time.perf_counter() + obs_np, values_np = sample_random_terminal_batch(batch_size, seed) + sample_seconds = time.perf_counter() - sample_started_at + + obs, values, h2d_seconds = _numpy_batch_to_device(obs_np, values_np, device) + return obs, values, sample_seconds, h2d_seconds + + +def _train_step_from_batch( model: PackedValueModel, - replay_buffer: EphemeralReplayBuffer, optimizer: torch.optim.Optimizer, + obs: torch.Tensor, + target: torch.Tensor, + sample_seconds: float, + h2d_seconds: float, *, - batch_size: int, - seed: int, device: torch.device, ) -> TrainStepResult: model.train() - obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( - replay_buffer, batch_size, seed, device - ) - optimizer.zero_grad(set_to_none=True) autocast = ( torch.autocast(device_type="cuda", dtype=torch.bfloat16) @@ -180,6 +206,51 @@ def train_step( ) +def train_step( + model: PackedValueModel, + replay_buffer: EphemeralReplayBuffer, + optimizer: torch.optim.Optimizer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> TrainStepResult: + obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( + replay_buffer, batch_size, seed, device + ) + return _train_step_from_batch( + model, + optimizer, + obs, + target, + sample_seconds, + h2d_seconds, + device=device, + ) + + +def pretrain_step( + model: PackedValueModel, + optimizer: torch.optim.Optimizer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> TrainStepResult: + obs, target, sample_seconds, h2d_seconds = _sample_terminal_batch( + batch_size, seed, device + ) + return _train_step_from_batch( + model, + optimizer, + obs, + target, + sample_seconds, + h2d_seconds, + device=device, + ) + + def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: run_dir = Path(config.run_dir) checkpoint_dir = run_dir / "checkpoints" @@ -231,9 +302,72 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) ) - model = cast(PackedValueModel, _to_channels_last(PackedValueModel().to(device))) + model = cast( + PackedValueModel, + _to_channels_last(PackedValueModel().to(device)), + ) model.eval() optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) + + if config.pretrain_terminal_samples > 0: + pretrain_steps = ( + config.pretrain_terminal_samples + config.pretrain_batch_size - 1 + ) // config.pretrain_batch_size + pretrain_log_interval = max(1, config.pretrain_log_interval) + pretrain_started_at = time.perf_counter() + pretrain_loss_window: list[float] = [] + print( + f"pretrain samples={config.pretrain_terminal_samples} batch={config.pretrain_batch_size} steps={pretrain_steps}" + ) + for step_idx in range(pretrain_steps): + samples_done = step_idx * config.pretrain_batch_size + step_batch_size = min( + config.pretrain_batch_size, + config.pretrain_terminal_samples - samples_done, + ) + result = pretrain_step( + model, + optimizer, + batch_size=step_batch_size, + seed=config.seed + step_idx, + device=device, + ) + loss_value = float(result.loss.float().cpu().item()) + pretrain_loss_window.append(loss_value) + if len(pretrain_loss_window) > pretrain_log_interval: + pretrain_loss_window.pop(0) + + step_number = step_idx + 1 + if ( + step_number % pretrain_log_interval == 0 + or step_number == pretrain_steps + ): + if device.type == "cuda": + torch.cuda.synchronize(device) + samples_done += step_batch_size + elapsed = time.perf_counter() - pretrain_started_at + window_mean_loss = float(np.mean(pretrain_loss_window)) + print( + f"pretrain step={step_number}/{pretrain_steps} samples={samples_done}/{config.pretrain_terminal_samples} " + f"loss={window_mean_loss:.6f} elapsed={elapsed:.2f}s" + ) + wandb.log( + { + "pretrain_step": step_number, + "pretrain_samples_total": samples_done, + "pretrain_loss_mean": window_mean_loss, + "pretrain_seconds": elapsed, + } + ) + + if device.type == "cuda": + torch.cuda.synchronize(device) + pretrain_seconds = time.perf_counter() - pretrain_started_at + model.eval() + print( + f"pretrain_complete samples={config.pretrain_terminal_samples} steps={pretrain_steps} elapsed={pretrain_seconds:.2f}s" + ) + replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( replay_buffer, @@ -257,6 +391,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_completed_turns = 0 previous_action_turn_count_total = 0 previous_completed_game_actions_total = 0 + previous_completed_game_turn_count_total = 0 previous_tree_build_nanos = 0 previous_descent_nanos = 0 previous_sample_collect_nanos = 0 @@ -282,6 +417,11 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: action_turn_count_total = selfplay.action_turn_count_total() max_turn_count_seen = selfplay.max_turn_count_seen() completed_game_actions_total = selfplay.completed_game_actions_total() + completed_game_turn_count_total = selfplay.completed_game_turn_count_total() + max_turn_count_in_completed_game = ( + selfplay.max_turn_count_in_completed_game() + ) + completed_game_turn_counts = selfplay.take_completed_game_turn_counts() max_actions_in_completed_game = selfplay.max_actions_in_completed_game() tree_build_nanos = selfplay.tree_build_nanos() descent_nanos = selfplay.descent_nanos() @@ -307,6 +447,10 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: completed_game_actions_added = ( completed_game_actions_total - previous_completed_game_actions_total ) + completed_game_turn_count_added = ( + completed_game_turn_count_total + - previous_completed_game_turn_count_total + ) tree_build_seconds = (tree_build_nanos - previous_tree_build_nanos) / 1e9 descent_seconds = (descent_nanos - previous_descent_nanos) / 1e9 sample_collect_seconds = ( @@ -354,6 +498,30 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"{actions_per_turn:.2f}" if completed_turns_added else "-" ) actions_per_game_display = f"{actions_per_game:.1f}" if games_added else "-" + if completed_game_turn_counts: + completed_game_turn_counts_arr = np.asarray( + completed_game_turn_counts, dtype=np.float32 + ) + completed_game_turn_count_mean = float( + completed_game_turn_counts_arr.mean() + ) + completed_game_turn_count_p50 = float( + np.percentile(completed_game_turn_counts_arr, 50) + ) + completed_game_turn_count_p90 = float( + np.percentile(completed_game_turn_counts_arr, 90) + ) + completed_game_turn_count_display = ( + f"mean:{completed_game_turn_count_mean:.1f} " + f"p50:{completed_game_turn_count_p50:.1f} " + f"p90:{completed_game_turn_count_p90:.1f} " + f"max:{max_turn_count_in_completed_game}" + ) + else: + completed_game_turn_count_mean = float("nan") + completed_game_turn_count_p50 = float("nan") + completed_game_turn_count_p90 = float("nan") + completed_game_turn_count_display = "-" build_ms_per_action = ( tree_build_seconds * 1000.0 / action_steps_added if action_steps_added @@ -502,6 +670,12 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "completed_game_actions_added": completed_game_actions_added, "actions_per_completed_game": actions_per_game, "max_actions_in_completed_game": max_actions_in_completed_game, + "completed_game_turn_count_total": completed_game_turn_count_total, + "completed_game_turn_count_added": completed_game_turn_count_added, + "completed_game_turn_count_mean": completed_game_turn_count_mean, + "completed_game_turn_count_p50": completed_game_turn_count_p50, + "completed_game_turn_count_p90": completed_game_turn_count_p90, + "max_turn_count_in_completed_game": max_turn_count_in_completed_game, "replay_size": replay_size, "collection_seconds": collect_seconds, "training_seconds": train_seconds, @@ -543,7 +717,8 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " - f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} collect={collect_seconds:.2f}s train={train_seconds:.2f}s " + f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} cg_tc={completed_game_turn_count_display} " + f"collect={collect_seconds:.2f}s train={train_seconds:.2f}s " f"sp_ms/act=build:{build_ms_per_action:.1f} desc:{descent_ms_per_action:.1f} coll:{collect_ms_per_action:.1f} push:{push_ms_per_action:.1f} " f"desc_ms/act=exp:{descent_expand_ms_per_action:.1f} app:{descent_apply_ms_per_action:.1f} sub:{descent_submit_ms_per_action:.1f} wait:{descent_wait_ms_per_action:.1f} bk:{descent_backup_ms_per_action:.1f} other:{descent_other_ms_per_action:.1f} " f"desc_us/eval=sub:{descent_submit_us_per_eval:.2f} wait:{descent_wait_us_per_eval:.2f} " @@ -579,6 +754,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_completed_turns = completed_turns previous_action_turn_count_total = action_turn_count_total previous_completed_game_actions_total = completed_game_actions_total + previous_completed_game_turn_count_total = completed_game_turn_count_total previous_tree_build_nanos = tree_build_nanos previous_descent_nanos = descent_nanos previous_sample_collect_nanos = sample_collect_nanos @@ -616,6 +792,21 @@ def _parse_args() -> TrainConfig: default=defaults.max_gpu_evals_per_move, ) parser.add_argument("--lr", type=float, default=defaults.lr) + parser.add_argument( + "--pretrain-terminal-samples", + type=int, + default=defaults.pretrain_terminal_samples, + ) + parser.add_argument( + "--pretrain-batch-size", + type=int, + default=defaults.pretrain_batch_size, + ) + parser.add_argument( + "--pretrain-log-interval", + type=int, + default=defaults.pretrain_log_interval, + ) parser.add_argument("--seed", type=int, default=defaults.seed) parser.add_argument( "--selfplay-precision", diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index f0144cc..730f6fe 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -2,7 +2,7 @@ //! //! No policy head - the model outputs only a scalar value per position. -use std::ffi::{CStr, c_void}; +use std::ffi::{c_void, CStr}; use std::mem::size_of; use std::slice; use std::sync::atomic::{AtomicU64, Ordering}; diff --git a/training/src/executor.rs b/training/src/executor.rs index 74d0d8d..5f29537 100644 --- a/training/src/executor.rs +++ b/training/src/executor.rs @@ -216,8 +216,8 @@ mod tests { #[test] fn test_executor_cancels_pending_futures() { - use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; let cancelled = Arc::new(AtomicBool::new(false)); let cancelled_clone = cancelled.clone(); diff --git a/training/src/future.rs b/training/src/future.rs index f4bf839..b9e3bee 100644 --- a/training/src/future.rs +++ b/training/src/future.rs @@ -10,8 +10,8 @@ use std::task::{Context, Poll}; use ndarray::Dimension; -use crate::BatchDim; use crate::queue::GpuJobQueue; +use crate::BatchDim; // Thread-local flag for tracking whether any future made progress. // Used by the executor to decide whether to park. diff --git a/training/src/lib.rs b/training/src/lib.rs index a17aa7d..2808d15 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,8 +1,10 @@ +use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; +use alpha_paint::TRAINING_START_FENS; use std::sync::Arc; use std::sync::{Mutex, OnceLock}; use cudagraph::CudaGraphRunner; -use queue::{BATCH_SIZE, queue_shape_for_workers}; +use queue::{queue_shape_for_workers, BATCH_SIZE}; use replay_buffer::ReplayBuffer; use training::{SelfPlaySession, SessionConfig}; use worker::WorkerConfig; @@ -10,8 +12,9 @@ use worker::WorkerConfig; use ndarray::{ArrayView, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; use numpy::{PyArray, PyArrayMethods}; use pyo3::prelude::*; -use rand::SeedableRng; +use rand::{Rng, RngExt, SeedableRng}; use rand_chacha::ChaCha8Rng; +use rayon::prelude::*; pub mod cudagraph; pub mod descent; @@ -92,6 +95,104 @@ fn graph_cache() -> &'static Mutex> { GRAPH_CACHE.get_or_init(|| Mutex::new(None)) } +const AVG_GAME_LENGTH: f32 = 500.0; +const PRETRAIN_SAMPLE_CHUNK: usize = 256; + +fn terminal_value_from_term(board: &Board, terminal: TerminalState) -> f32 { + let sign = terminal.value() as f32; + let p = board.turn_count.max(1) as f32; + sign * (AVG_GAME_LENGTH / p).ln_1p() +} + +fn sample_random_action(board: &Board, rng: &mut R) -> Option { + let actions = board.get_valid_actions(); + let action_idx = (actions.len() > 0).then(|| rng.random_range(0..actions.len()))?; + actions.into_iter().nth(action_idx).copied() +} + +fn random_terminal_board(rng: &mut R) -> (Board, TerminalState) { + let mut board = + Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) + .expect("training start FEN must parse"); + + loop { + let Some(action) = sample_random_action(&board, rng) else { + let is_white_turn = board.is_white_turn(); + return (board, TerminalState::loss_for(is_white_turn)); + }; + + let (outcome, rollback) = board.apply_action(action); + match outcome { + ApplyActionOutcome::Ongoing => {} + ApplyActionOutcome::Terminal { terminal } => return (board, terminal), + ApplyActionOutcome::PlayInstead { + terminal, + play_instead, + } => { + board.rollback(action, rollback); + let (play_outcome, _) = board.apply_action(play_instead); + debug_assert!(matches!(play_outcome, ApplyActionOutcome::Terminal { .. })); + return (board, terminal); + } + ApplyActionOutcome::Killshot { terminal, moves } => { + let last_idx = moves.len().saturating_sub(1); + for (idx, mv) in moves.into_iter().enumerate() { + let action = if idx == last_idx { + Action::FinalMove(mv) + } else { + Action::Move(mv) + }; + let (killshot_outcome, _) = board.apply_action(action); + if idx == last_idx { + debug_assert!(matches!( + killshot_outcome, + ApplyActionOutcome::Terminal { .. } + )); + } else { + debug_assert!(matches!(killshot_outcome, ApplyActionOutcome::Ongoing)); + } + } + return (board, terminal); + } + } + } +} + +fn sample_random_terminal_examples(n: usize, seed: u64) -> (Vec, Vec) { + let mut obs_data = vec![0u16; n * cudagraph::OBS_WORDS]; + let mut values = vec![0.0; n]; + + obs_data + .par_chunks_mut(PRETRAIN_SAMPLE_CHUNK * cudagraph::OBS_WORDS) + .zip(values.par_chunks_mut(PRETRAIN_SAMPLE_CHUNK)) + .enumerate() + .for_each(|(chunk_idx, (obs_chunk, values_chunk))| { + let mut rng = ChaCha8Rng::seed_from_u64(seed.wrapping_add(chunk_idx as u64)); + for (obs_sample, value_out) in obs_chunk + .chunks_mut(cudagraph::OBS_WORDS) + .zip(values_chunk.iter_mut()) + { + let (board, terminal) = random_terminal_board(&mut rng); + observation::encode_into_slice(&board, obs_sample); + *value_out = terminal_value_from_term(&board, terminal); + } + }); + + (obs_data, values) +} + +#[pyfunction] +fn sample_random_terminal_batch<'py>( + py: Python<'py>, + n: usize, + seed: u64, +) -> PyResult<(Bound<'py, PyArray>, Bound<'py, PyArray>)> { + let (obs_data, values) = py.detach(|| sample_random_terminal_examples(n, seed)); + let obs = PyArray::from_vec(py, obs_data).reshape(Ix2(n, cudagraph::OBS_WORDS))?; + let values = PyArray::from_vec(py, values); + Ok((obs, values)) +} + /// Replay buffer storing (packed observation, value) pairs. /// /// Observations are flat u16 tensors, values are f32 scalars (no policy). @@ -319,6 +420,21 @@ impl SelfPlay { Ok(self.session()?.max_actions_in_completed_game()) } + /// Return the total completed turn count accumulated across finished games. + fn completed_game_turn_count_total(&self) -> PyResult { + Ok(self.session()?.completed_game_turn_count_total()) + } + + /// Return the largest completed turn count seen in a finished game. + fn max_turn_count_in_completed_game(&self) -> PyResult { + Ok(self.session()?.max_turn_count_in_completed_game()) + } + + /// Drain completed game turn counts collected since the last call. + fn take_completed_game_turn_counts(&self) -> PyResult> { + Ok(self.session()?.take_completed_game_turn_counts()) + } + /// Return total time spent building fresh search trees, in nanoseconds. fn tree_build_nanos(&self) -> PyResult { Ok(self.session()?.tree_build_nanos()) @@ -395,5 +511,6 @@ impl Drop for SelfPlay { fn alphapaint_training(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_function(pyo3::wrap_pyfunction!(sample_random_terminal_batch, m)?)?; Ok(()) } diff --git a/training/src/observation.rs b/training/src/observation.rs index e60d71e..fb25fd4 100644 --- a/training/src/observation.rs +++ b/training/src/observation.rs @@ -1,6 +1,6 @@ -use alpha_paint::board::Board; use alpha_paint::board::board_structs::Player; use alpha_paint::board::structs::Coordinate; +use alpha_paint::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs index b9d511a..88ac817 100644 --- a/training/src/replay_buffer.rs +++ b/training/src/replay_buffer.rs @@ -4,8 +4,8 @@ use std::cell::UnsafeCell; use std::sync::atomic::{AtomicU64, Ordering}; use ndarray::{Array, ArrayViewMut, Axis}; -use rand::Rng; use rand::seq::index::sample; +use rand::Rng; use crate::BatchDim; diff --git a/training/src/training.rs b/training/src/training.rs index 12552ad..66692b4 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -3,6 +3,7 @@ //! Provides `SelfPlaySession`: a persistent session with pause/resume semantics //! that preserves in-progress game state across boundaries. +use std::mem; use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; use std::sync::{Arc, Condvar, Mutex}; use std::thread; @@ -46,6 +47,12 @@ struct SessionControl { completed_game_actions_total: Arc, /// Largest number of selected actions observed in a completed game. max_actions_in_completed_game: Arc, + /// Total completed turn count accumulated across finished games. + completed_game_turn_count_total: Arc, + /// Largest completed turn count observed in a finished game. + max_turn_count_in_completed_game: Arc, + /// Completed game turn counts waiting to be drained by the trainer. + completed_game_turn_counts: Arc>>, /// Total time spent building fresh search trees. tree_build_nanos: Arc, /// Total time spent running descent iterations. @@ -87,6 +94,9 @@ impl SessionControl { max_turn_count_seen: Arc::new(AtomicUsize::new(0)), completed_game_actions_total: Arc::new(AtomicU64::new(0)), max_actions_in_completed_game: Arc::new(AtomicU64::new(0)), + completed_game_turn_count_total: Arc::new(AtomicU64::new(0)), + max_turn_count_in_completed_game: Arc::new(AtomicUsize::new(0)), + completed_game_turn_counts: Arc::new(Mutex::new(Vec::new())), tree_build_nanos: Arc::new(AtomicU64::new(0)), descent_nanos: Arc::new(AtomicU64::new(0)), sample_collect_nanos: Arc::new(AtomicU64::new(0)), @@ -335,6 +345,30 @@ impl SelfPlaySession { .load(Ordering::Acquire) } + /// Return the total completed turn count accumulated across finished games. + pub fn completed_game_turn_count_total(&self) -> u64 { + self.control + .completed_game_turn_count_total + .load(Ordering::Acquire) + } + + /// Return the largest completed turn count seen in a finished game. + pub fn max_turn_count_in_completed_game(&self) -> usize { + self.control + .max_turn_count_in_completed_game + .load(Ordering::Acquire) + } + + /// Drain completed game turn counts collected since the last call. + pub fn take_completed_game_turn_counts(&self) -> Vec { + let mut guard = self + .control + .completed_game_turn_counts + .lock() + .expect("completed game turn counts mutex poisoned"); + mem::take(&mut *guard) + } + /// Return the total time spent building fresh search trees. pub fn tree_build_nanos(&self) -> u64 { self.control.tree_build_nanos.load(Ordering::Acquire) @@ -431,6 +465,9 @@ fn session_thread_main( max_turn_count_seen: control.max_turn_count_seen.clone(), completed_game_actions_total: control.completed_game_actions_total.clone(), max_actions_in_completed_game: control.max_actions_in_completed_game.clone(), + completed_game_turn_count_total: control.completed_game_turn_count_total.clone(), + max_turn_count_in_completed_game: control.max_turn_count_in_completed_game.clone(), + completed_game_turn_counts: control.completed_game_turn_counts.clone(), tree_build_nanos: control.tree_build_nanos.clone(), descent_nanos: control.descent_nanos.clone(), sample_collect_nanos: control.sample_collect_nanos.clone(), diff --git a/training/src/worker.rs b/training/src/worker.rs index 273b76c..92293cb 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -4,7 +4,7 @@ //! internal tree nodes, and pushes them to the replay buffer. use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Instant; use ndarray::Ix1; @@ -65,6 +65,9 @@ pub struct SelfPlayMetrics { pub max_turn_count_seen: Arc, pub completed_game_actions_total: Arc, pub max_actions_in_completed_game: Arc, + pub completed_game_turn_count_total: Arc, + pub max_turn_count_in_completed_game: Arc, + pub completed_game_turn_counts: Arc>>, pub tree_build_nanos: Arc, pub descent_nanos: Arc, pub sample_collect_nanos: Arc, @@ -152,6 +155,7 @@ fn finish_game_with_current_tree( replay_buffer: &ReplayBuffer, metrics: &SelfPlayMetrics, action_steps_in_game: u64, + completed_turn_count_in_game: usize, ) { push_samples_to_replay( replay_buffer, @@ -162,6 +166,18 @@ fn finish_game_with_current_tree( .completed_game_actions_total .fetch_add(action_steps_in_game, Ordering::AcqRel); update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); + metrics + .completed_game_turn_count_total + .fetch_add(completed_turn_count_in_game as u64, Ordering::AcqRel); + update_max_usize( + &metrics.max_turn_count_in_completed_game, + completed_turn_count_in_game, + ); + metrics + .completed_game_turn_counts + .lock() + .expect("completed game turn count mutex poisoned") + .push(completed_turn_count_in_game); } /// Run a single self-play game with tree learning. @@ -228,7 +244,13 @@ async fn play_game( "how the fuck did we get here?? board has no children. FEN={}", tree.root_state ); - finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + action_steps_in_game, + tree.root_state.turn_count, + ); break; } @@ -250,13 +272,34 @@ async fn play_game( // Check if game is over match outcome { - alpha_paint::board::ApplyActionOutcome::Terminal { .. } - | alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { - finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); + alpha_paint::board::ApplyActionOutcome::Terminal { .. } => { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + action_steps_in_game, + new_board.turn_count, + ); + break; + } + alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + action_steps_in_game, + tree.root_state.turn_count + 1, + ); break; } alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { - finish_game_with_current_tree(&tree, replay_buffer, metrics, action_steps_in_game); + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + action_steps_in_game, + tree.root_state.turn_count + 1, + ); break; } alpha_paint::board::ApplyActionOutcome::Ongoing => { From e8c4ddc97be0a591e49586a9bb053ae03759a559 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 13:49:03 -0700 Subject: [PATCH 34/59] resnet in engine run the value net directly inside alpha_paint with embedded exported weights, unpacked board features, and the training-style log terminal shaping so UBFM can use the model on cpu. Made-with: Cursor --- Cargo.lock | 389 ++++++++++++++++++- alpha_paint/Cargo.toml | 11 + alpha_paint/assets/value_net.safetensors | Bin 0 -> 14092 bytes alpha_paint/benches/nn_latency.rs | 49 +++ alpha_paint/build.rs | 8 + alpha_paint/src/bindings.rs | 30 +- alpha_paint/src/evaluation.rs | 121 ++---- alpha_paint/src/lib.rs | 6 +- alpha_paint/src/nn/features.rs | 250 +++++++++++++ alpha_paint/src/nn/mod.rs | 2 + alpha_paint/src/nn/model.rs | 457 +++++++++++++++++++++++ alpha_paint/src/search.rs | 314 +++++++++------- pyproject.toml | 2 +- python/scripts/export_value_net.py | 160 ++++++++ 14 files changed, 1555 insertions(+), 244 deletions(-) create mode 100644 alpha_paint/assets/value_net.safetensors create mode 100644 alpha_paint/benches/nn_latency.rs create mode 100644 alpha_paint/build.rs create mode 100644 alpha_paint/src/nn/features.rs create mode 100644 alpha_paint/src/nn/mod.rs create mode 100644 alpha_paint/src/nn/model.rs create mode 100644 python/scripts/export_value_net.py diff --git a/Cargo.lock b/Cargo.lock index 3583a08..e7f531f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,12 +11,32 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "alpha_paint" version = "0.1.0" dependencies = [ + "criterion", + "ndarray", "pyo3", "rand 0.10.0", + "safetensors", + "serde", + "serde_json", ] [[package]] @@ -38,6 +58,18 @@ dependencies = [ "serde_json", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.102" @@ -56,6 +88,28 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + +[[package]] +name = "cc" +version = "1.2.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e928d4b69e3077709075a938a05ffbedfa53a84c8f766efbf8220bb1ff60e1" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -73,6 +127,58 @@ dependencies = [ "rand_core 0.10.0", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -91,6 +197,41 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -116,6 +257,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "cudarc" version = "0.18.2" @@ -148,12 +295,24 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + [[package]] name = "foldhash" version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "futures-core" version = "0.3.32" @@ -228,13 +387,24 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "foldhash", + "foldhash 0.1.5", ] [[package]] @@ -242,6 +412,13 @@ name = "hashbrown" version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "serde", + "serde_core", +] [[package]] name = "heck" @@ -267,12 +444,31 @@ dependencies = [ "serde_core", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "js-sys" +version = "0.3.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4c90f45aa2e6eacbe8645f77fdea542ac97a494bcd117a67df9ff4d611f995" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + [[package]] name = "leb128fmt" version = "0.1.0" @@ -388,6 +584,22 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "parking" version = "2.2.1" @@ -400,6 +612,34 @@ version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.13.1" @@ -702,6 +942,32 @@ dependencies = [ "semver", ] +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "safetensors" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" +dependencies = [ + "hashbrown 0.16.1", + "serde", + "serde_json", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "semver" version = "1.0.27" @@ -751,6 +1017,12 @@ dependencies = [ "zmij", ] +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + [[package]] name = "slab" version = "0.4.12" @@ -774,6 +1046,16 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "toml_datetime" version = "1.1.0+spec-1.1.0" @@ -823,6 +1105,16 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -841,6 +1133,51 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6523d69017b7633e396a89c5efab138161ed5aafcbc8d3e5c5a42ae38f50495a" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e3a6c758eb2f701ed3d052ff5737f5bfe6614326ea7f3bbac7156192dc32e67" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "921de2737904886b52bcbb237301552d05969a6f9c40d261eb0533c8b055fedf" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a93e946af942b58934c604527337bad9ae33ba1d5c6900bbb41c2c07c2364a93" +dependencies = [ + "unicode-ident", +] + [[package]] name = "wasm-encoder" version = "0.244.0" @@ -875,12 +1212,62 @@ dependencies = [ "semver", ] +[[package]] +name = "web-sys" +version = "0.3.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84cde8507f4d7cfcb1185b8cb5890c494ffea65edbe1ba82cfd63661c805ed94" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + [[package]] name = "winnow" version = "1.0.0" diff --git a/alpha_paint/Cargo.toml b/alpha_paint/Cargo.toml index 19f5a00..add05a4 100644 --- a/alpha_paint/Cargo.toml +++ b/alpha_paint/Cargo.toml @@ -7,9 +7,20 @@ edition = "2024" name = "alpha_paint" crate-type = ["cdylib", "rlib"] +[[bench]] +name = "nn_latency" +harness = false + [dependencies] +ndarray = "0.17.2" pyo3 = { version = "0.28.2", features = ["extension-module"] } rand = "0.10.0" +safetensors = "0.7.0" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.149" [features] debug=[] + +[dev-dependencies] +criterion = "0.8.2" diff --git a/alpha_paint/assets/value_net.safetensors b/alpha_paint/assets/value_net.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..79184e647978b6c9921fedc5e83fb51e5c819603 GIT binary patch literal 14092 zcmbW7c|4b2*Y86@LPZjhC@Mqd>Dp_53rQNJNOOcr0~s1MmiGSh^PKZK&wZbBpZmG~yY_2e@4fa~pS`X-v8#jqLW8~i1B`;=Xp_#df#e)AzH2d!p|Kb^1See-TAJBln;8mu^ zi~YTO8vVZ-`H$xP#p!`rnEVr1j~@g41D5uL!#m*Lgj$&NK&{RGAJDbl!T%=H+RV_x z`oFmHKYINygY!@N-9zfZnOpsXiDv%|C-!i9U{?Pa5i{fez7u9{y+5iPn_68>cLt4E$?N1-oD05eEyXNd)oeoG4dC}(8AXCZ?pav=U+v& zr)^W4e*){_7+Raz{vFx>3Hw(uHT_4-Uj##I^Z#)F|HAw$n@kP=S^xgR7+PDH{jHCX zP`_2iiv#|BmHh?ip;-Mbi2tPgx8?O0sE76+KL4+@e^*z3aSW|3&Hfg`f5QHiul`Bw z{zCL{EdP$(-<-d)@Bc@tG`0Tc+4nD&p{bdLWskHrYyb5aceF#Uo;KL$w}NZ&TQ9CEy1As*RRMV|jO z#5tp_VCM4IBsCPO+l8x31mvjxufmNU2=HLeK^VOwk3CU>DR}Is&<+;@;53D(Z0oxZ}|qJli7RjQPPhw^zH6SJoSz-#SH#ojRe|#vR=Dw6WN`6y4@q zfWysfaD8ZkTh1LL&xIluqAvq+9q&Qe;x;T9lR|1Bn^p9p zQwm!Q`6 z!#@*VhbEA(huxrG{X&TK*bXj|EiB1q0J}P$z?DNaAX9t_oM)=iqrJD_d7DzaQfH4t zkMG3kZRUc4g%cY5NI^UE33y7W7xY;b3dcNFg7OY^2!AQV_a2@>_BtD(^V4u_NZx=# zV*sliz7{rD<$?W)%^>N%3w`=K5LJx>#MtKvGffW#-777T;$L zDpN5hZ#NOD5@4oe8&lmn0=<7)qes`Um3qm zl_Sd%v+!2Yck~T*!us(W@X0GhynMGn7%-|!XkPUc3~#qGCE2k!SHBP5tvV&v(lQ`( zgATE-UkA&*(+=TAg#%Zc8`7VTUM$c1R<6YS-8)f#-X)wXC(t8RkHN_`2UO1&fNuYMu?*U_ zcF$`<)9U^hxa6fUrz%Q()<6uAzlVcjLM(i=c7&70tH3Ov1kV0Sf%_8Cux&&HY)aS# z2Q7=>v1uV_S!Te88S`Mv)d(0YT~6lMR+e9y6GVQh=@GBR^&Y@yd+EBdiikIn2I&`@{BJT*RE5 z#O%%DK(XG^Vx|+Bh2wu*V^7n1BlbNG&9{$%^NO8NIy@cBc1D0{zfI4MqeF#H3o-|3e-N!E?Z6@Z2N`B?G!ha{dGH@;V<_zgUXhHakn4UzXt0E)kA#DiuCC zi1D{!A>O^X5_MZ;F|+j+4BBW2Cs+|Uc_$F{=s=+fq=-rHN$|vOr=UA}BD8AmWR>r) zk`Yf^Nn5KQ&O3P#wTu<wqz)cI?{Q^zG3Ital(C)&0ust5fT;RflNCC2_r5+2OA1A zllH)wPXf#zx*s+R<+x8X2Tv{D!-(=xY&U#}USkm7?wy36=8h+I8JT2K)l011Gy-35 zMZDMF8xxL}l6bXG#I|t>8SQft+8WiBRw0ILN~!`R-$Wy zuOxJN;|wJnbFv@z_l;)>3MRO1@fzII-H%&JO(X{=>QN!OA3fqMPe-*52A1m$vSVD} z$9PRzb>kiJRe4J$xGF-imKq=VMg#49HeyK0O?=|khlKC;#onf&Fs@3Bg;RQim3bR( z4m^Z2`^U3EQn4Uy+XpTV{)O_xw=u7?V|eYl3;1gYc8`f?-E!_Q+Z)X==3(xMoBZjZ+*dFalJbl##yb6{> zruG1SM=cs{P7dP5%O}{$3|8TVCmq<5J6hC2E*g8ewF6sc4%uPS)cxXMR&(|z3~-Hs zH?k_^XvqOMbm$e95q#i-*(hPL(e+S1=yCP7wn33gaS!H8u=_+i%r zGD`X+oDuho|F0C<9Hv3R!i6~Or6I~FC1K;m26SyXiN$wM!h!MQ@b#tX>Naj2(V$QJ*+gkg4H*!1~R$&b1}V*3zlOqYE|F54^P`;(6jL|dN2#X!Z0nhM&>4y-IKTuPBCmS%^kX|CQX;HE&mDY`=fe)CPOSS zU0IH&M6L&e4XUWz;f1%dN|28WL~VVzm3uaW`MukrKu1^;qEr8=L^5WxaEpW-%gv~XG0gpX5;Jo4v+-GVC-t!i*_tCH5 z%GB|AC%6Ji_g{ie*<;{VEI}S!1prEFhG=9BWJ3XOy2af!nN#{Pcr(;|ld46LCZkRfbcP9+x5}pp+ zsw{@Dx~9tC9+T!Pa2c0sz7Ku&M&S&fWSlV02p78dV$0N1N!E+*a=DI72zrx)(R;g* z5JkS_`bYc_F_P_CJRi3#n!sn5uHj1k%y|1s4^A?C>9`u8lU;UGL;Klu&J0z$?Xwj< z@_8}+u47BP=BCnKY1??9+c&IR62&Vw#PYQ7gg-NShd)BDVX>w0Y?ZGJKWZNZ`s=mf?YiDP zMpK9HZ!*WNsx$EOv1n2*uEOzse~^Rw-lKtf4&wb%^z8mf7T)fKzxHbJmzSSo)4Yo~ z>iiJC?$JB;T|>rhQqu;EJ6nv#dP7jSF%O}e~m8!0=}1i#yDg^&$9 z$i9qF%#}5w_f`W{&e#E_Cwg=fHJU<_1l8Zt3IjJz#=p_=V#aSSiVv_2v@s zKDTQmvY?esi=GC{7Ho$x=4Y9dRVZ|f48<&!iD)z;35pC3f%5New&>?=JdkhzmqzvE zwt`Dgsd)mP3y+z>+yj_?rDwi7Rfh$Pk0aA{713#T9x@$c%2&sL^n#s`cWNQb$ov2q zSGUm3NwetHW;MDteH=|T(&lOA&fIlb8V@-Y$xnqvanD41ZWfWw2Ogir4JFoa*U{0m z;adnjwtXITj_gmjZm5DMQe)`d*HPg8>jKz6>w;5h+sGRpg$cDu9HrI z$(2Q*Gc^hVNA8Doi3IwlQbgBub4b*924k&eQmH@IRGS;pdpFhT;gpXMB=Y80(*k(V zmC-ydL5>%wCgZXU7e4Xn1nynXhObJ7aeIexRAG=i^_yZt2YsIi*Qcmcm(M@p*Ng<( za3!8vn>o_&lg8n*9d6ulqYs{{?uV5Ia$IT3N^V*(mj~T&@g z4SoF(yUxDB#wq^fv|B5gJbD+(-Z@R*I;PN97BlH~t26ZelCxBA_g$LeB&M@IMo^EM z5*k}JjNa60CF^I-p&Gu0w649J4%nVd-#6Ze5%F;%_t{Y*P37&PwSvFMJU2qL&ox|R zP`pX>V)HuDF~wNHOy1jn{wF_QH&r=yHEZ;%92kKJCS#xOMC@Ils;ziBgjH6)&3#iX_8+tnD0 zi4y6F-CG8l zP6ol7f&|=8i}8}|SQ>f6o_cKE&W>!-;@`EuvTgIG@Vt7d9-n^3ACGnE>mwui^rX}H zsHRz*v~M;~D7`9nx~qs{Gd-N@Elop5rGZ;O1U)uVnod9PiJY)ag`(<<{KVPotXBCT zo}U@SQ~3?_w7h`MkvUvmWQ1R|_tQTT$G|i0Gu$?Ni2LSmhh9H&@k#D0jJadSV}I>O zuj{tB!BvM__lymH#m$t)s?h2?gkOnu=SMH!N5=vFTu@0tqEioNqdt(QUtZ#6!*1xl zJcg&(*z@O^jD-GD;?GhQ$>e;r*;NB|cY?E7w$oioXmF#cF_wnwn^=aUzOtO+}aQo8R1lblB8WM5gF9P40Gw7A-O5Hqk`iN2#JD&Q7$u*pq&1i2%*jBCZ}y$ShlJT>35% zpIul&J=$k+tES9~iAE9hzQt*LGy4_P4V=US+&uA#jupSN@fH3FcnP}(DsY37R-CQ1 zql(>9^t+lH`Ej~0SAJrI{>8xw-F3Tze3{GZISja01@((FAjmli+r#22R*YU+?)hvuSK1;h0A_w1L8G6=(_w>Zp|)p!F!_8TcTQ`7Sq6kw`1hgX;tjFu zmkbYAaHsdaZ{w?jr05>Zub8+$U0hbkG0Z@pMYZ=O)%uD2SSv_a$; z;s~SI20r5PP>ju-&kM|rxo+pSvd1pEZ0vyvY<|Uadcblv+Q;^+`E!uD8m$z*$}a%R zvG&w5O-g*`%Qb9QU&mbxqv`fcH)yQ%B2L;`)MkIkblmc>ZC1@f!N9(WtVp$#O{Ql84%BCF=bk-N(w_Td_hwC*|7E!P`a1ut(b*CGK zdGmJbZ51nKne)q$e{dZcO!Y5r#F*?axY8$vn1Aqt=@tX|_Z%(GYJIqeS_o{uGMG0n z+eI1zX5gT@AzY-f5Bt|!a23f;#8OM{JhhAq7nbq&gWcdn-xhdtIh%iP=|oAD66&~f zT18D?itjwX;v#DsJaG9EnOU%nAFO^%9h9S>g+KisgFb_d3aJ{bAas>MTSp8@m9nWrb|0%0VzGFGCKt47m;S^np>l|438-PPPzmhAb--BNBWZt^Tg${*zEW{4;E2>y?RXk`2aRdZbJR6u~>T}fhHtsUTA$cT{s5CNVEUW)9#NM*Ys%#uRC%NyT(1X8}zG?j10Ylm*NW{+DL=vnNQ>;m!y$;|HkME zDzvRqn-_FtGdsUZu;`x3lU1###Dzv2*C!lp+HY{Rz$(5ANAUxXqWCa9BWh<}%v*FE z>2iw_I#-aT9V67?+nflH?U{ddq$k0n7kl^sI4<@wJ`$oh2`acRV7^tq!#KWq2m{vR$NrW(;JD{q6&+CMn_Y$EKw--{ad zap7evhl6(T65clE6sE2uG%-G%-`*`r=bA?FYYuz4)Q~mw_tr-8S|g4%oi?O1;>*$h z{v-5Gn@;cKo~QPn-*BgOE={TKEt=nZ5^ZnzgHwqwH3;{lo4m`}xyx5+`*kHQObz9& zAIzwcSV1_Y;s=W@0cJ`tyGU7OzN0XN4-JTgiLE1f->oqfRf*$hLLX<6G59*JUGage z$a&CnJ~5Lg?w6#pn=avj&G9^QgC$MBeF>~Oqd@g*6$UzIqH~{P#3`qP{8moI=`9{~ zV|F-g$Xra*I}G{qFQeedC|zE8^CgvgVb9&3dsTQH9$MiO{~jjg%fZ8%LJVBqTAuON zgmTUAbsNA_*PI?u1#~)il;jre0J-d7-lU<1#=j#(ZLJzSZ2SSx&*_G&RypFRWIdr;4vff9Rv1#WwQ?)pK zy3hV7JUkIWPlb8Xa?ib_XmJ{6A3ne<&vqK~?kHC>PURO~>hdRRo%qC>d)(m`<(?Z- z_@ljUq8*xIUr(V8QPo21jr@3y4l zZnY8!I_Uy2?iQH1J(sPXEsIOeR$=9yFmS6b!wc6QVWmin4|{JGAMxuFB6a$L+~7M< zTxCWIFK;B;FRjVpGo~c&=326RY@|r$=@!w;>Xo8z4x2P8<31w*-% zlmBRTbdV-7m>VxH&9E-dx7o~cVprO!ZQmv4#U{-8zAU`^I)_<=x})2#BI5q*A$foy z%y6W6*_eJ8$&R;KMC;83yVvcV;$AnFLx9<5oc&slUC>b=EIWt=r5|AZ`{^-leLIXi zu}=7<5JaR+L&UjjHV>u_=GIYMHk zq1a2nF4CXZ9&tdA8V$bVz*u;zR7i$R--1=j7Ua-#4ft$%g2}~SCvV*BAk7-Up*`6G>OTAG6M3tTX3`yV0xhBw$9uOyOrsq9PM8FpjfSMkrD+}l;D!5w0D z^Tb!Z=;LHt8aK5a21bQ(TTdmvV(DZ^SZDxIW>qAzduZ3hg&olrk13U2$)#?jMu zvKhW-gt`;?FmJF1w0s@}=Soar$e~LlvX4hL(a{nlH`{kfxJP&Bpbn zx-{nTQ{38MP4ir*^HpIH{Oq16E;mxd*Y+8UGwWTrtF$X$WIv5ZPWIz-w(7ufmzL9+83^prZSvRq?DjxwR=+O*^WLr&Yu-NybyH`sW;12L#t7KEJsygd48`_CYjMw; zujJx$NxQu-KcjLVCwe$gk{>%VPkdd|h&X=rmu0@B@%X!bO6|P^T&(Aq3 z@Yz!$c%NhL{Fz}S_ps^5WdrhsJ7@B+)>55&&8K5Tz`kSB+1~<^fhEi zPY(IKZ7Yem6+vWYzY+>==LmzsB*oun#*(W8F0qdN*5IX`X_GbM5y=^=NhI5NCk|&Vm&L-~-%7Y7s*Nl;f07NEC`$sZ z7nDaY4KDwYa+W#te#Iu$u7KKo>DVlF3iHfA13kA_ye`~@H@oGMGMUpj{pl&%wem21 z@l~0Y*F{qDbs1GpaD~~=KTyet659HzkU9>ssW6$WSE1!+UQst$t-{aaE%%5TTA?^c zsp4{Q?}}%cS{1rsU3RaA9Uzqj+rfX~8$mbqA&U%9h2*?sQlWQD-1*}+`I?wZEY=+r znk>@GFD`v3xV~%_{8R5*x2#YRO09L+hw5kI%YIMUmZMGOag+5}`|Lcnz}lHL_(w2R zk21lxzP(%~O_SZZ7GB;y{0Y8UPvOIIeR1GA7YP2O$d+wt5-(K#BzAi-n7vba%Um8E zBG*SLkZaWwS>B2}Ec(E2X7tGt->-WhR?T*V*JU@z08A4U&q?4}Jp+6}Clb{!N15By z0pw%IOzgj79yxtA6z5v3MWghIh_m{EnZ#`Lm~))0)x1Oojd(=t26nLTBY(0{)8CWS z;3PQUkc5L{h@JbAzVOC+vGCP03}U7FGeeUc;mqmWtd^}SW{T#q~zwo}I;s*_UG;ZXh+RFJ-wKf`etXGBfna})9?FViu_l%Fk z%UtTfZ$38r9$)k91wZ_yPsPtiv)I-(AIXEa$Jy`!sbV?*FmZ?;?E`u3_)?ToUBht_EkVy~JmB zJc%$Jft`kh%w72!xmTl$cT^=<@A1eYe-+`ovOy%Ncmauez5-r)II%7Ef$U|u2iYB@ z&Q|PSg$t#3koCS2FhswBXsl0R=Qk&_x@9-TVdG25;5H9J-1L~CdM>lS7hRsOQ!1qI z2q1c%ve>=u5&JpkB|#z!tG$P?KfhlSS6w~gu*{8v?YCo!qCZI#`0G>WW4yR-_ zlk-a{oL!SZ5@b_wgb+>Se;g%2Z3RTn{iXQ%hZ|zAmVV^D&S7}1xfWkL2H`Hbv5^04 z4$N88AHU9ZLJyheWYelLXuj@fmu|lb@B~K#=$9lo1slP-tUZu3tAR~PnFG5IPJ`?DZlK(%2wrhRpkz-5d-C`!G~XJ6j&tOB zu5}X0b5Genr2)qGI`kcig|l&VijZ6g`l*tBtU0@P}+Hqm?k(e(XP$RcqqlmkG)BM ziRt1O5_3u3tQg^$?gc^HPeCYCcpx?j-Uxe>n##9#4nv4tDqO9ZhKuM@VVLF)@@z~r zJ{cmzP3k;R+Ucq-OmxMoE2>$&;Sm@>_JDf7ROtU?9L`wrfHaky#JN`l_B}dUT=np> zt){rTDP56`6Ht3~VpeN7wHLkaEcgE29RGZ?2n}jq`8} zP(Fqy-a0{GaRi>0ZeUs4#*@fs1NKKPhm6U0WueCnd;VeOY@V}%xOjcNSmXFYIP$4k zta;uDjP492_U~U}aK~BZ_q`G$LdU}P76Tq$;?F1Z<2cz6ag~E9Tdgq`4`&_1_%}7A z!+#fipB)ZAo}3UeW{qIz+ho^zbUXxDO~!$1B3SV_C_a4e1G7jjU@C)8mU+0u z!b4;Go;!a(SiwgVY-<%`^ngNaoWBQ0{EoI0ntOJatoNW&;A*G|-w02-U4_!siGqZR zK8y6&!R!n*@Q}l1-tOLuZ@xGlJ!W1g*dKD*T*4pU~kbzHz(crmf5m>2a z!4`EN_&HGosw+DoP|XTr)a9sreFuV4Eml7wxGy~$_s2iMNsR&6{w)`V>bt<_1f z?{vX`h$JizeBE=akUTxMP5&aaQuqCtK^tSyoe8P^l+xsyD6Q>cZ} zGqvFloeG@_&cc|7v*F(AZenY>1C4{Kk$um?B9EiOg@BJl&Ok)oj~Whw*BaKE*N@Z; znSw{1-BHo?g?O+yLvXF%Cq5Uti`5@jWMUCvF8%s}g3cP!wxC70xUY=mc5f3KzwTMT z_j^zF56)tzOrHu5+k=SyE=BAtHz&c>(yY%%O<`16lXyc#Cfh4H8`tmzNc4%o-F=_2 zcFhmA14oR*1?)C^aCrukdkf6#((kf^DS2c}_vv!uLCeTk>Lqq6F~Gc3!k-T*#gX$} zac5!|&bsm&c~%?x_2a^<^g(zH?vt7FuFzo=2G>%i!RQAk#0Sqru|)Oj;-+0r9G>f= zVedke4-6;sZ?-W(BbB*cmmmvI8iI7-C-%HGpG=O|KtX~F@3jPKxaR)bjEc-JZC4kv+?B--D~mO1smx6FkN`1&`GZRP+;?g<*;{23(Lvx z$+Ov#j23$1*(E5fc|RAICJQvI^&!?OoFLt~b9hPn4e{ul+hF6_1RMKH(xGPkVCs}t zLQv9DkRE&rKA*^eq;5$ti`<4u-k#wJ+_ef_+gnmKXxY`7e$$i zS2vpxg-gA`Yn48HH_Rs<3!afxk2aG~jZQKyL=xwyOQ2NM6j)$%m>Adh#p7|3xOvw( z@#~aB;^kdoWaZm9lFOrDw8vq%790&T_6!B}`F-*2y@^z_W-L~1%z}G5;h?%|C7Hay z6S9UIL*ki5q@gWDyhUvYDnz~}OLr9z)7@qe@@6%QQrN^?57iUtcgezuol$H`TV&bD z0g+5~%zJUU%0ki*K8(os912YN`j8AgvzGNXOA}u1?GWcKtzySrOPR>Ol6iC&kb#K{ zSmsGjB7asvP_@byY6AZ-jpeJL<4q3vHadzJW=u!51*1U!i$9*OTqe%h4tR0NU2;c0 z3S^67*d)t*SRm>ne!aP8Z!;>5bq-F&kmhPKrNb3B>hA*AtQn;8p(|8Nl=1taUU>0f zBj2Ms1Xq;J?72go%eChCf#dBoc8(M1-LuYpD7lc-{w`&Ht46b*m-n$!%S`OQu8X8k zv4?AU0c_`p6YwS~l=T07nKX!P*>N9hL>onPsre{AHlsm2{ns*P_PrYSIRd;msm^O} z%Yo~IM7nqXUL2S{iyKdyMOMZhL!Ib)@SpgD4AJ%^>01gJ`efj&E>Cv&j2vMl@wh{~ zKlXW8Nww<;_6ht=_vf0DZkbHlFY_kd?sFH_E6354V~=rXXa=0qngr>=@x Vec { + fn search(&mut self, time_left: f32) -> PyResult> { let mut local_board = self.0.clone(); let mut py_actions: Vec = vec![]; - let evaluator = Evaluator::new(&local_board); + let evaluator = Evaluator::new(&local_board).map_err(PyRuntimeError::new_err)?; // Time control // For reference, on my laptop we can hit around 20k iterations per second (with incremental logic) @@ -215,7 +215,7 @@ impl PyBoard { let mut tree = GameSearchTree::new(&local_board, &evaluator); loop { - tree.run_descent_for_iter(15_000, Duration::from_secs_f32(max_duration)); + tree.run_descent_for_iter(iterations.max(50), Duration::from_secs_f32(max_duration)); let (action_id, action) = tree.get_best_action_and_index(); let player_coord = local_board.current_player_coord(); @@ -225,16 +225,16 @@ impl PyBoard { ApplyActionOutcome::Ongoing => { py_actions.push(action.to_python_primitives(&player_coord)); if action.is_final() { - return py_actions; + return Ok(py_actions); } } ApplyActionOutcome::Terminal { .. } => { py_actions.push(action.to_python_primitives(&player_coord)); - return py_actions; + return Ok(py_actions); } ApplyActionOutcome::PlayInstead { play_instead, .. } => { py_actions.push(play_instead.to_python_primitives(&player_coord)); - return py_actions; + return Ok(py_actions); } ApplyActionOutcome::Killshot { terminal: _, moves } => { py_actions.push(action.to_python_primitives(&player_coord)); @@ -251,7 +251,7 @@ impl PyBoard { coord = mv.target; } } - return py_actions; + return Ok(py_actions); } }; @@ -260,6 +260,20 @@ impl PyBoard { } } + fn bench_eval(&self, reps: usize) -> PyResult { + let reps = reps.max(1); + let evaluator = Evaluator::new(&self.0).map_err(PyRuntimeError::new_err)?; + let started_at = std::time::Instant::now(); + let mut sink = 0.0f32; + for _ in 0..reps { + sink += evaluator.evaluate(&self.0); + } + let elapsed = started_at.elapsed().as_secs_f64(); + let avg_micros = elapsed * 1_000_000.0 / reps as f64; + let _ = sink; + Ok(avg_micros) + } + fn apply_turn(&mut self, turn: &Turn) { let len = turn.actions.len(); for (i, (action, _)) in turn.actions.iter().enumerate() { diff --git a/alpha_paint/src/evaluation.rs b/alpha_paint/src/evaluation.rs index 5e71b67..771b600 100644 --- a/alpha_paint/src/evaluation.rs +++ b/alpha_paint/src/evaluation.rs @@ -1,107 +1,42 @@ +use std::sync::{Arc, OnceLock}; + use crate::board::Board; -use crate::board::board_structs::Player; -use std::cmp::min; +use crate::nn::model::ValueModel; + +static SHARED_MODEL: OnceLock, String>> = OnceLock::new(); -pub struct Evaluator; +#[derive(Clone)] +pub struct Evaluator { + model: Arc, +} impl Evaluator { - pub fn new(_board: &Board) -> Evaluator { - Evaluator + pub fn new(_board: &Board) -> Result { + let model = load_shared_model()?; + Ok(Evaluator { model }) } - pub fn evaluate(&self, board: &Board) -> i32 { - let mut white_hills = 0; - let mut black_hills = 0; - for md in board.tiles.hill_metadata().iter() { - match md.owner { - Some(Player::White) => white_hills += 1, - Some(Player::Black) => black_hills += 1, - None => (), - }; + pub fn from_model(model: ValueModel) -> Evaluator { + Evaluator { + model: Arc::new(model), } - if white_hills * 4 >= board.hills.len() * 3 { - return 2_000_000_000 - 5 * (board.turn_count as i32); - } else if black_hills * 4 >= board.hills.len() * 3 { - return -2_000_000_000 + 5 * (board.turn_count as i32); - } - - // Deci-tile scale (10 eval = 1 painted tile) - let mut evaluation = 0; - - evaluation += Self::hill_paint_eval(board); // Hill paint gets a large point bonus - evaluation += Self::paint_coverage_eval(board); - - evaluation += board.white_stamina as i32 * 5; - evaluation -= board.black_stamina as i32 * 5; - - evaluation += Self::find_target_hill_eval(board); - - evaluation } - fn find_target_hill_eval(board: &Board) -> i32 { - let mut white_best_score = usize::MAX; - let mut black_best_score = usize::MAX; - - let mut total_eval: i32 = 0; - - for (hill_tiles, hill_data) in board.hills.iter().zip(board.tiles.hill_metadata().iter()) { - let mut white_dist = usize::MAX; - let mut black_dist = usize::MAX; - - for &tile in hill_tiles.iter() { - if !board.tiles[tile].is_owned_by::() { - white_dist = min(white_dist, board.dist[(board.white_coord, tile)] as usize); - } - if !board.tiles[tile].is_owned_by::() { - black_dist = min(black_dist, board.dist[(board.black_coord, tile)] as usize); - } - } - - match hill_data.owner { - Some(Player::White) => { - total_eval += 200_000; - black_best_score = min(black_dist, black_best_score); - } - Some(Player::Black) => { - total_eval -= 200_000; - white_best_score = min(white_dist, white_best_score); - } - None => { - white_best_score = min(white_dist, white_best_score); - black_best_score = min(black_dist, black_best_score); - } - } - } - - if white_best_score != usize::MAX { - total_eval += -20 * (white_best_score as i32); - } - - if black_best_score != usize::MAX { - total_eval -= -20 * (black_best_score as i32); - } - - total_eval + pub fn benchmark() -> Evaluator { + Self::from_model(ValueModel::benchmark_model()) } - fn hill_paint_eval(board: &Board) -> i32 { - let mut evaluation = 0; - - for hill in board.hills.iter() { - for tile in hill.iter() { - if board.tiles[*tile].is_owned_by::() { - evaluation += 25_000; - } else if board.tiles[*tile].is_owned_by::() { - evaluation -= 25_000; - } - } - } - - evaluation + pub fn evaluate(&self, board: &Board) -> f32 { + self.model.evaluate(&crate::nn::features::extract(board)) } +} - fn paint_coverage_eval(board: &Board) -> i32 { - board.tiles.coverage_eval() * 5 - } +fn load_shared_model() -> Result, String> { + SHARED_MODEL + .get_or_init(|| { + ValueModel::from_embedded() + .map(Arc::new) + .map_err(|err| format!("failed to load embedded model weights: {err}")) + }) + .clone() } diff --git a/alpha_paint/src/lib.rs b/alpha_paint/src/lib.rs index e7c2fec..40dd0a6 100644 --- a/alpha_paint/src/lib.rs +++ b/alpha_paint/src/lib.rs @@ -1,11 +1,11 @@ use crate::board::{ApplyActionOutcome, Board}; use pyo3::prelude::*; -use std::time::Instant; mod bindings; pub mod board; -mod evaluation; -mod search; +pub mod evaluation; +pub mod nn; +pub mod search; fn perft(board: &mut Board, depth: usize) -> usize { if depth == 1 { diff --git a/alpha_paint/src/nn/features.rs b/alpha_paint/src/nn/features.rs new file mode 100644 index 0000000..2f73455 --- /dev/null +++ b/alpha_paint/src/nn/features.rs @@ -0,0 +1,250 @@ +use crate::board::board_structs::Player; +use crate::board::structs::Coordinate; +use crate::board::Board; + +pub const BOARD_SIDE: usize = 32; +pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; +pub const BOARD_PLANES: usize = 17; +pub const INTRINSIC_COUNT: usize = 10; +pub const INTRINSIC_SCALE: [f32; INTRINSIC_COUNT] = [ + 420.0, 420.0, 8.0, 8.0, 2000.0, 1024.0, 1024.0, 1024.0, 1024.0, 64.0, +]; + +pub const CURRENT_PAINT_L1: usize = 0; +pub const CURRENT_PAINT_L2: usize = 1; +pub const CURRENT_PAINT_L3: usize = 2; +pub const CURRENT_PAINT_L4: usize = 3; +pub const OPPONENT_PAINT_L1: usize = 4; +pub const OPPONENT_PAINT_L2: usize = 5; +pub const OPPONENT_PAINT_L3: usize = 6; +pub const OPPONENT_PAINT_L4: usize = 7; +pub const WALL: usize = 8; +pub const POWERUP: usize = 9; +pub const BEACON_CURRENT: usize = 10; +pub const BEACON_OPPONENT: usize = 11; +pub const HILL_NEUTRAL: usize = 12; +pub const HILL_CURRENT: usize = 13; +pub const HILL_OPPONENT: usize = 14; +pub const CURRENT_PLAYER: usize = 15; +pub const OPPONENT_PLAYER: usize = 16; + +pub const INTRINSIC_CURRENT_STAMINA: usize = 0; +pub const INTRINSIC_OPPONENT_STAMINA: usize = 1; +pub const INTRINSIC_CURRENT_HILLS: usize = 2; +pub const INTRINSIC_OPPONENT_HILLS: usize = 3; +pub const INTRINSIC_TURN_COUNT: usize = 4; +pub const INTRINSIC_CURRENT_TERRITORY: usize = 5; +pub const INTRINSIC_OPPONENT_TERRITORY: usize = 6; +pub const INTRINSIC_CURRENT_BEACONS: usize = 7; +pub const INTRINSIC_OPPONENT_BEACONS: usize = 8; +pub const INTRINSIC_CONSECUTIVE_MOVES: usize = 9; + +#[derive(Clone, Debug)] +pub struct Features { + pub board: Vec, + pub intrinsics: [f32; INTRINSIC_COUNT], +} + +impl Features { + pub fn zeros() -> Self { + let mut board = vec![0.0; BOARD_PLANES * BOARD_CELLS]; + for y in 0..BOARD_SIDE { + for x in 0..BOARD_SIDE { + board[plane_index(WALL, x as u8, y as u8)] = 1.0; + } + } + Self { + board, + intrinsics: [0.0; INTRINSIC_COUNT], + } + } +} + +pub fn extract(board: &Board) -> Features { + let mut features = Features::zeros(); + fill_board_planes(board, &mut features.board); + features.intrinsics = extract_intrinsics(board); + features +} + +pub fn fill_board_planes(board: &Board, out: &mut [f32]) { + assert_eq!( + out.len(), + BOARD_PLANES * BOARD_CELLS, + "feature board buffer size mismatch" + ); + + out.fill(0.0); + for y in 0..BOARD_SIDE { + for x in 0..BOARD_SIDE { + out[plane_index(WALL, x as u8, y as u8)] = 1.0; + } + } + + let current_is_white = board.is_white_turn(); + let current_coord = board.current_player_coord(); + let opponent_coord = if current_is_white { + board.black_coord + } else { + board.white_coord + }; + + for y in 0..board.rows { + for x in 0..board.cols { + let coord = Coordinate::new(x, y); + let tile = board.tiles[coord]; + let paint = tile.paint_value(); + let strength = paint.unsigned_abs() as usize; + let is_enemy_paint = if current_is_white { + paint < 0 + } else { + paint > 0 + }; + + out[plane_index(WALL, x, y)] = if tile.is_wall() { 1.0 } else { 0.0 }; + out[plane_index(POWERUP, x, y)] = if board.powerups[coord] { 1.0 } else { 0.0 }; + + if strength > 0 { + let offset = if is_enemy_paint { + OPPONENT_PAINT_L1 + } else { + CURRENT_PAINT_L1 + }; + for level in 0..strength.min(4) { + out[plane_index(offset + level, x, y)] = 1.0; + } + } + + match tile.beacon_owner() { + Some(Player::White) if current_is_white => { + out[plane_index(BEACON_CURRENT, x, y)] = 1.0; + } + Some(Player::White) => { + out[plane_index(BEACON_OPPONENT, x, y)] = 1.0; + } + Some(Player::Black) if current_is_white => { + out[plane_index(BEACON_OPPONENT, x, y)] = 1.0; + } + Some(Player::Black) => { + out[plane_index(BEACON_CURRENT, x, y)] = 1.0; + } + None => {} + } + + match board.hill_id[coord] { + u16::MAX => {} + hill_id => match board.tiles.hill_metadata()[hill_id as usize].owner { + None => out[plane_index(HILL_NEUTRAL, x, y)] = 1.0, + Some(Player::White) if current_is_white => { + out[plane_index(HILL_CURRENT, x, y)] = 1.0 + } + Some(Player::White) => out[plane_index(HILL_OPPONENT, x, y)] = 1.0, + Some(Player::Black) if current_is_white => { + out[plane_index(HILL_OPPONENT, x, y)] = 1.0 + } + Some(Player::Black) => out[plane_index(HILL_CURRENT, x, y)] = 1.0, + }, + } + + if coord == current_coord { + out[plane_index(CURRENT_PLAYER, x, y)] = 1.0; + } + if coord == opponent_coord { + out[plane_index(OPPONENT_PLAYER, x, y)] = 1.0; + } + } + } +} + +pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { + let current_is_white = board.is_white_turn(); + let mut out = [0.0; INTRINSIC_COUNT]; + + if current_is_white { + out[INTRINSIC_CURRENT_STAMINA] = board.white_stamina as f32; + out[INTRINSIC_OPPONENT_STAMINA] = board.black_stamina as f32; + out[INTRINSIC_CURRENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = + board.tiles.get_beacon_iterator::().count() as f32; + } else { + out[INTRINSIC_CURRENT_STAMINA] = board.black_stamina as f32; + out[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina as f32; + out[INTRINSIC_CURRENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_CURRENT_BEACONS] = + board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + } + + out[INTRINSIC_TURN_COUNT] = board.turn_count as f32; + out[INTRINSIC_CONSECUTIVE_MOVES] = board.consecutives_moves_so_far as f32; + + for (value, scale) in out.iter_mut().zip(INTRINSIC_SCALE) { + *value /= scale; + } + + out +} + +#[inline] +pub const fn plane_index(plane: usize, x: u8, y: u8) -> usize { + plane * BOARD_CELLS + y as usize * BOARD_SIDE + x as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_features_for_white_turn() { + let board = Board::from_fen( + "ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:1B1/d1O/3", + ) + .unwrap(); + + let features = extract(&board); + + assert_eq!(features.board[plane_index(CURRENT_PAINT_L1, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PAINT_L2, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L4, 0, 1)], 1.0); + assert_eq!(features.board[plane_index(POWERUP, 0, 1)], 1.0); + assert_eq!(features.board[plane_index(HILL_NEUTRAL, 1, 1)], 1.0); + assert_eq!(features.board[plane_index(BEACON_CURRENT, 2, 1)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PLAYER, 0, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 2, 2)], 1.0); + assert_eq!(features.board[plane_index(WALL, 31, 31)], 1.0); + + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 99.0 / 420.0); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 88.0 / 420.0); + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_TERRITORY], 1.0 / 1024.0); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], 1.0 / 1024.0); + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_BEACONS], 1.0 / 1024.0); + assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 2.0 / 64.0); + } + + #[test] + fn extracts_features_for_black_turn() { + let board = Board::from_fen( + "ap2|3x3|tc:1|cm:5|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/2o", + ) + .unwrap(); + + let features = extract(&board); + + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L1, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L2, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(HILL_CURRENT, 1, 1)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PLAYER, 2, 2)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 0, 0)], 1.0); + + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 88.0 / 420.0); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 99.0 / 420.0); + assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 5.0 / 64.0); + } +} diff --git a/alpha_paint/src/nn/mod.rs b/alpha_paint/src/nn/mod.rs new file mode 100644 index 0000000..c572841 --- /dev/null +++ b/alpha_paint/src/nn/mod.rs @@ -0,0 +1,2 @@ +pub mod features; +pub mod model; diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs new file mode 100644 index 0000000..846be9f --- /dev/null +++ b/alpha_paint/src/nn/model.rs @@ -0,0 +1,457 @@ +use ndarray::{Array1, Array2}; +use safetensors::{Dtype, SafeTensors}; +use safetensors::tensor::Metadata; + +use crate::nn::features::{BOARD_CELLS, BOARD_PLANES, BOARD_SIDE, Features, INTRINSIC_COUNT}; + +const DEFAULT_EPS: f32 = 1.0e-5; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ModelConfig { + pub width: usize, + pub num_blocks: usize, + pub hidden_dim: usize, +} + +impl Default for ModelConfig { + fn default() -> Self { + Self { + width: 8, + num_blocks: 1, + hidden_dim: 32, + } + } +} + +#[derive(Clone, Debug)] +pub struct Conv2d { + out_channels: usize, + in_channels: usize, + weights: Vec, +} + +#[derive(Clone, Debug)] +pub struct BatchNorm2d { + weight: Vec, + bias: Vec, + running_mean: Vec, + running_var: Vec, + eps: f32, +} + +#[derive(Clone, Debug)] +pub struct Linear { + weight: Array2, + bias: Array1, +} + +#[derive(Clone, Debug)] +pub struct ResidualBlock { + norm1: BatchNorm2d, + conv1: Conv2d, + norm2: BatchNorm2d, + conv2: Conv2d, +} + +#[derive(Clone, Debug)] +pub struct ValueModel { + pub config: ModelConfig, + stem_conv: Conv2d, + stem_norm: BatchNorm2d, + blocks: Vec, + head_linear1: Linear, + head_linear2: Linear, +} + +impl ValueModel { + pub fn zeroed(config: ModelConfig) -> Self { + let stem_conv = Conv2d::zeros(config.width, BOARD_PLANES); + let stem_norm = BatchNorm2d::identity(config.width); + let blocks = (0..config.num_blocks) + .map(|_| ResidualBlock::zeroed(config.width)) + .collect(); + let head_linear1 = Linear::zeros(config.hidden_dim, config.width + INTRINSIC_COUNT); + let head_linear2 = Linear::zeros(1, config.hidden_dim); + Self { + config, + stem_conv, + stem_norm, + blocks, + head_linear1, + head_linear2, + } + } + + pub fn benchmark_model() -> Self { + Self::zeroed(ModelConfig::default()) + } + + pub fn from_embedded() -> Result { + Self::from_bytes(include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/assets/value_net.safetensors" + ))) + } + + pub fn from_bytes(bytes: &[u8]) -> Result { + if bytes == b"placeholder\n" || bytes == b"placeholder" { + return Err( + "embedded weights are still the placeholder file; run python/scripts/export_value_net.py to overwrite alpha_paint/assets/value_net.safetensors".to_string(), + ); + } + let (_, metadata) = SafeTensors::read_metadata(bytes) + .map_err(|err| format!("failed to read safetensors metadata: {err}"))?; + let config = load_config(&metadata)?; + let tensors = SafeTensors::deserialize(bytes) + .map_err(|err| format!("failed to deserialize safetensors: {err}"))?; + Self::from_safetensors(&tensors, config) + } + + pub fn from_safetensors( + tensors: &SafeTensors<'_>, + config: ModelConfig, + ) -> Result { + let stem_conv = Conv2d::from_tensor(tensors, "stem.conv.weight", config.width, BOARD_PLANES)?; + let stem_norm = BatchNorm2d::from_prefix(tensors, "stem.bn", config.width)?; + + let mut blocks = Vec::with_capacity(config.num_blocks); + for idx in 0..config.num_blocks { + let prefix = format!("blocks.{idx}"); + blocks.push(ResidualBlock { + norm1: BatchNorm2d::from_prefix( + tensors, + &format!("{prefix}.norm1"), + config.width, + )?, + conv1: Conv2d::from_tensor( + tensors, + &format!("{prefix}.conv1.weight"), + config.width, + config.width, + )?, + norm2: BatchNorm2d::from_prefix( + tensors, + &format!("{prefix}.norm2"), + config.width, + )?, + conv2: Conv2d::from_tensor( + tensors, + &format!("{prefix}.conv2.weight"), + config.width, + config.width, + )?, + }); + } + + let head_linear1 = Linear::from_prefix( + tensors, + "head.fc1", + config.hidden_dim, + config.width + INTRINSIC_COUNT, + )?; + let head_linear2 = Linear::from_prefix(tensors, "head.fc2", 1, config.hidden_dim)?; + + Ok(Self { + config, + stem_conv, + stem_norm, + blocks, + head_linear1, + head_linear2, + }) + } + + pub fn evaluate(&self, features: &Features) -> f32 { + let width = self.config.width; + let plane_span = width * BOARD_CELLS; + let mut activations = vec![0.0; plane_span]; + let mut scratch_a = vec![0.0; plane_span]; + let mut scratch_b = vec![0.0; plane_span]; + + self.stem_conv.forward(&features.board, &mut activations); + self.stem_norm.relu_inplace(&mut activations); + + for block in &self.blocks { + scratch_a.copy_from_slice(&activations); + block.norm1.relu_inplace(&mut scratch_a); + block.conv1.forward(&scratch_a, &mut scratch_b); + block.norm2.relu_inplace(&mut scratch_b); + block.conv2.forward(&scratch_b, &mut scratch_a); + for (dst, residual) in activations.iter_mut().zip(&scratch_a) { + *dst += *residual; + } + } + + let pooled = average_pool_channels(&activations, width); + let mut head_input = Array1::zeros(width + INTRINSIC_COUNT); + for (dst, src) in head_input.iter_mut().take(width).zip(pooled) { + *dst = src; + } + for (dst, src) in head_input + .iter_mut() + .skip(width) + .zip(features.intrinsics.iter().copied()) + { + *dst = src; + } + + let mut hidden = self.head_linear1.forward(&head_input); + hidden.mapv_inplace(|value: f32| value.max(0.0)); + let output = self.head_linear2.forward(&hidden); + output[0].tanh() * 7.0 + } +} + +impl Conv2d { + fn zeros(out_channels: usize, in_channels: usize) -> Self { + Self { + out_channels, + in_channels, + weights: vec![0.0; out_channels * in_channels * 3 * 3], + } + } + + fn from_tensor( + tensors: &SafeTensors<'_>, + name: &str, + out_channels: usize, + in_channels: usize, + ) -> Result { + let data = tensor_f32(tensors, name, &[out_channels, in_channels, 3, 3])?; + Ok(Self { + out_channels, + in_channels, + weights: data, + }) + } + + fn forward(&self, input: &[f32], out: &mut [f32]) { + debug_assert_eq!(input.len(), self.in_channels * BOARD_CELLS); + debug_assert_eq!(out.len(), self.out_channels * BOARD_CELLS); + out.fill(0.0); + + for out_ch in 0..self.out_channels { + for y in 0..BOARD_SIDE { + for x in 0..BOARD_SIDE { + let mut sum = 0.0; + for in_ch in 0..self.in_channels { + for ky in 0..3 { + let iy = y as isize + ky as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&iy) { + continue; + } + for kx in 0..3 { + let ix = x as isize + kx as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&ix) { + continue; + } + let input_idx = in_ch * BOARD_CELLS + + iy as usize * BOARD_SIDE + + ix as usize; + let weight_idx = + (((out_ch * self.in_channels + in_ch) * 3 + ky) * 3) + kx; + sum += input[input_idx] * self.weights[weight_idx]; + } + } + } + out[out_ch * BOARD_CELLS + y * BOARD_SIDE + x] = sum; + } + } + } + } +} + +impl BatchNorm2d { + fn identity(width: usize) -> Self { + Self { + weight: vec![1.0; width], + bias: vec![0.0; width], + running_mean: vec![0.0; width], + running_var: vec![1.0; width], + eps: DEFAULT_EPS, + } + } + + fn from_prefix(tensors: &SafeTensors<'_>, prefix: &str, width: usize) -> Result { + let weight = tensor_f32(tensors, &format!("{prefix}.weight"), &[width])?; + let bias = tensor_f32(tensors, &format!("{prefix}.bias"), &[width])?; + let running_mean = tensor_f32(tensors, &format!("{prefix}.running_mean"), &[width])?; + let running_var = tensor_f32(tensors, &format!("{prefix}.running_var"), &[width])?; + Ok(Self { + weight, + bias, + running_mean, + running_var, + eps: DEFAULT_EPS, + }) + } + + fn relu_inplace(&self, values: &mut [f32]) { + debug_assert_eq!(values.len(), self.weight.len() * BOARD_CELLS); + for channel in 0..self.weight.len() { + let scale = self.weight[channel] / (self.running_var[channel] + self.eps).sqrt(); + let offset = self.bias[channel] - self.running_mean[channel] * scale; + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + for value in &mut values[start..end] { + *value = (*value * scale + offset).max(0.0); + } + } + } +} + +impl Linear { + fn zeros(out_features: usize, in_features: usize) -> Self { + Self { + weight: Array2::zeros((out_features, in_features)), + bias: Array1::zeros(out_features), + } + } + + fn from_prefix( + tensors: &SafeTensors<'_>, + prefix: &str, + out_features: usize, + in_features: usize, + ) -> Result { + let weight = tensor_f32( + tensors, + &format!("{prefix}.weight"), + &[out_features, in_features], + )?; + let bias = tensor_f32(tensors, &format!("{prefix}.bias"), &[out_features])?; + let weight = Array2::from_shape_vec((out_features, in_features), weight) + .map_err(|err| format!("failed to shape {prefix}.weight: {err}"))?; + let bias = Array1::from_vec(bias); + Ok(Self { weight, bias }) + } + + fn forward(&self, input: &Array1) -> Array1 { + let mut out = self.weight.dot(input); + out += &self.bias; + out + } +} + +impl ResidualBlock { + fn zeroed(width: usize) -> Self { + Self { + norm1: BatchNorm2d::identity(width), + conv1: Conv2d::zeros(width, width), + norm2: BatchNorm2d::identity(width), + conv2: Conv2d::zeros(width, width), + } + } +} + +fn load_config(metadata: &Metadata) -> Result { + let metadata = metadata + .metadata() + .as_ref() + .ok_or_else(|| "missing safetensors metadata".to_string())?; + let parse_usize = |key: &str| -> Result { + metadata + .get(key) + .ok_or_else(|| format!("missing metadata key {key}"))? + .parse::() + .map_err(|err| format!("invalid metadata {key}: {err}")) + }; + Ok(ModelConfig { + width: parse_usize("width")?, + num_blocks: parse_usize("num_blocks")?, + hidden_dim: parse_usize("hidden_dim")?, + }) +} + +fn tensor_f32( + tensors: &SafeTensors<'_>, + name: &str, + expected_shape: &[usize], +) -> Result, String> { + let tensor = tensors + .tensor(name) + .map_err(|err| format!("missing tensor {name}: {err}"))?; + if tensor.dtype() != Dtype::F32 { + return Err(format!( + "tensor {name} has dtype {:?}, expected F32", + tensor.dtype() + )); + } + if tensor.shape() != expected_shape { + return Err(format!( + "tensor {name} has shape {:?}, expected {:?}", + tensor.shape(), + expected_shape + )); + } + let data = tensor.data(); + if data.len() % 4 != 0 { + return Err(format!("tensor {name} has invalid byte length {}", data.len())); + } + let mut out = Vec::with_capacity(data.len() / 4); + for chunk in data.chunks_exact(4) { + out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); + } + Ok(out) +} + +fn average_pool_channels(values: &[f32], channels: usize) -> Vec { + let mut pooled = vec![0.0; channels]; + let scale = 1.0 / BOARD_CELLS as f32; + for channel in 0..channels { + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + pooled[channel] = values[start..end].iter().copied().sum::() * scale; + } + pooled +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::nn::features::{Features, plane_index}; + + #[test] + fn forward_matches_simple_hand_computation() { + let mut model = ValueModel::zeroed(ModelConfig { + width: 1, + num_blocks: 0, + hidden_dim: 1, + }); + model.stem_conv.weights[4] = 1.0; + model.head_linear1.weight[[0, 0]] = 1.0; + model.head_linear2.weight[[0, 0]] = 0.5; + + let mut features = Features::zeros(); + for y in 0..BOARD_SIDE as u8 { + for x in 0..BOARD_SIDE as u8 { + features.board[plane_index(0, x, y)] = 1.0; + } + } + + let expected = (0.5f32).tanh() * 7.0; + assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); + } + + #[test] + fn residual_block_preserves_input_when_convs_are_zero() { + let mut model = ValueModel::zeroed(ModelConfig { + width: 1, + num_blocks: 1, + hidden_dim: 1, + }); + model.stem_conv.weights[4] = 1.0; + model.head_linear1.weight[[0, 0]] = 1.0; + model.head_linear2.weight[[0, 0]] = 1.0; + + let mut features = Features::zeros(); + for y in 0..BOARD_SIDE as u8 { + for x in 0..BOARD_SIDE as u8 { + features.board[plane_index(0, x, y)] = 1.0; + } + } + + let expected = 1.0f32.tanh() * 7.0; + assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); + } +} diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index 8e317f5..e4015e4 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -1,14 +1,19 @@ +use std::cmp::Ordering; +use std::time::Duration; + +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; + use crate::board::actions::Move; use crate::board::{Action, ApplyActionOutcome, Board, TerminalState}; use crate::evaluation::Evaluator; -use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; -use std::time::Duration; + +const AVG_GAME_LENGTH: f32 = 500.0; #[derive(Debug)] pub struct ChildData { action: Action, - child_value: i32, + child_value: f32, entrance_count: usize, pub node: Option>, } @@ -16,7 +21,7 @@ pub struct ChildData { impl ChildData { fn completion_value(&self) -> i32 { match &self.node { - Some(c) => c.completion_value, + Some(child) => child.completion_value, None => 0, } } @@ -24,14 +29,14 @@ impl ChildData { #[derive(Debug)] pub struct SearchNode { - pub value: i32, + pub value: f32, pub completion_value: i32, resolved: bool, pub children: Vec, } impl SearchNode { - fn new(value: i32, completion_value: i32, is_resolved: bool) -> SearchNode { + fn new(value: f32, completion_value: i32, is_resolved: bool) -> SearchNode { SearchNode { value, completion_value, @@ -40,120 +45,120 @@ impl SearchNode { } } - /// Prefers moves with low entrance counts fn completed_best_action_dual( &self, is_max_player: bool, rng: &mut SmallRng, ) -> (usize, Action) { - let res = if is_max_player { + let result = if is_max_player { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { + .max_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action_dual without any valid actions!") + .expect("called completed_best_action_dual without any valid actions") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .min_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action_dual without any valid actions!") + .expect("called completed_best_action_dual without any valid actions") }; - - (res.0, res.1.action) + (result.0, result.1.action) } - /// Prefers moves with high entrance counts fn completed_best_action(&self, is_max_player: bool, rng: &mut SmallRng) -> (usize, Action) { - let res = if is_max_player { + let result = if is_max_player { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action without any valid actions!") + .expect("called completed_best_action without any valid actions") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { + .min_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action without any valid actions!") + .expect("called completed_best_action without any valid actions") }; - - (res.0, res.1.action) + (result.0, result.1.action) } fn backup_resolution(&self) -> bool { if self.completion_value.abs() == 1 { true } else { - self.children.iter().all(|child| { - child - .node - .as_ref() - .and_then(|c| Some(c.resolved)) - .unwrap_or(false) - }) + self.children + .iter() + .all(|child| child.node.as_ref().map(|node| node.resolved).unwrap_or(false)) } } - /// Build a chain of resolved SearchNodes for a killshot move sequence. - /// All nodes are resolved. Only the final node is terminal (no children). fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { let term_value = Self::value_from_term(board, terminal); - let comp_value = terminal.value(); + let completion_value = terminal.value(); + let mut node = SearchNode::new(term_value, completion_value, true); - // Start with the terminal leaf (the collision result) - let mut node = SearchNode::new(term_value, comp_value, true); - - // Build chain from last move to first - for i in (0..moves.len()).rev() { - let action = if i == moves.len() - 1 { - Action::FinalMove(moves[i]) + for index in (0..moves.len()).rev() { + let action = if index == moves.len() - 1 { + Action::FinalMove(moves[index]) } else { - Action::Move(moves[i]) + Action::Move(moves[index]) }; - - let parent = SearchNode { + node = SearchNode { value: term_value, - completion_value: comp_value, + completion_value, resolved: true, children: vec![ChildData { action, - child_value: comp_value, + child_value: term_value, entrance_count: 0, node: Some(Box::new(node)), }], }; - - node = parent; } node @@ -167,7 +172,7 @@ impl SearchNode { ) -> SearchNode { match outcome { ApplyActionOutcome::Ongoing => { - let mut new_node = SearchNode::new(0, 0, false); + let mut new_node = SearchNode::new(0.0, 0, false); let actions = board.get_valid_actions(); new_node.children.reserve(actions.len()); @@ -178,8 +183,8 @@ impl SearchNode { for action in actions.into_iter().copied() { let mut local_board = board.clone(); - let (outcome, _) = local_board.apply_action(action); - match outcome { + let (child_outcome, _) = local_board.apply_action(action); + match child_outcome { ApplyActionOutcome::Ongoing => { new_node.children.push(ChildData { action, @@ -188,19 +193,15 @@ impl SearchNode { node: None, }); } - // we can only guarantee that the play instead action is a valid action if - // its a Move (e.g. collision) or a Paint after a move - // Therefore, we just do the same thing as we did above, where we just - // treat this move as a terminal move too, and handle it when we consume - // the SearchTree (e.g. in bindings.rs) ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => { + let child_value = Self::value_from_term(&local_board, terminal); new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value, entrance_count: 0, node: Some(Box::new(SearchNode::new( - Self::value_from_term(&local_board, terminal), + child_value, terminal.value(), true, ))), @@ -210,7 +211,7 @@ impl SearchNode { let chain = Self::build_killshot_chain(&local_board, terminal, &moves); new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value: chain.value, entrance_count: 0, node: Some(Box::new(chain)), }); @@ -225,10 +226,6 @@ impl SearchNode { new_node.resolved = new_node.backup_resolution(); new_node } - // in search, we don't really CARE about the distinction between these two. - // e.g. the search tree doesn't _really_ care that we need to play a `Final` variant of - // the action instead. It just cares that this is a terminal state. - // We can just fix this in the consumers of the search tree, e.g. in bindings.rs ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => SearchNode::new( Self::value_from_term(board, terminal), @@ -241,8 +238,10 @@ impl SearchNode { } } - fn value_from_term(board: &Board, term: TerminalState) -> i32 { - term.value() * (2_000_000_000 - 5 * (board.turn_count as i32)) + fn value_from_term(board: &Board, term: TerminalState) -> f32 { + let sign = term.value() as f32; + let progress = board.turn_count.max(1) as f32; + sign * (AVG_GAME_LENGTH / progress).ln_1p() } fn create_child( @@ -251,20 +250,13 @@ impl SearchNode { action: Action, evaluator: &Evaluator, rng: &mut SmallRng, - ) -> i32 { + ) -> f32 { let (outcome, _) = state.apply_action(action); - let node = Box::new(SearchNode::build_self(&state, outcome, evaluator, rng)); let value = node.value; - - if let Some(id) = self - .children - .iter() - .position(|child| child.action == action) - { - self.children[id].node = Some(node); + if let Some(index) = self.children.iter().position(|child| child.action == action) { + self.children[index].node = Some(node); } - value } @@ -274,12 +266,11 @@ impl SearchNode { outcome: ApplyActionOutcome, evaluator: &Evaluator, rng: &mut SmallRng, - ) -> i32 { + ) -> f32 { let white_turn = state.is_white_turn(); - match outcome { ApplyActionOutcome::Ongoing => { - if self.children.len() == 0 { + if self.children.is_empty() { let loss = TerminalState::loss_for(white_turn); self.resolved = true; self.completion_value = loss.value(); @@ -290,12 +281,11 @@ impl SearchNode { if !self.resolved { let (best_action_id, best_action) = self.completed_best_action_dual(white_turn, rng); - self.children[best_action_id].entrance_count += 1; - if let Some(child_val) = self.children[best_action_id].node.as_mut() { - let (outcome, _) = state.apply_action(best_action); - child_val.ubfms_iteration(state, outcome, evaluator, rng); + if let Some(child) = self.children[best_action_id].node.as_mut() { + let (next_outcome, _) = state.apply_action(best_action); + child.ubfms_iteration(state, next_outcome, evaluator, rng); } else { self.children[best_action_id].child_value = self.create_child(state, best_action, evaluator, rng); @@ -307,7 +297,6 @@ impl SearchNode { self.resolved = self.backup_resolution(); } } - ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } | ApplyActionOutcome::Killshot { terminal, .. } => { @@ -316,7 +305,6 @@ impl SearchNode { self.value = Self::value_from_term(&state, terminal); } } - self.value } } @@ -330,46 +318,55 @@ pub struct GameSearchTree<'a> { impl GameSearchTree<'_> { fn safest_action(&mut self) -> (usize, Action) { - let val = if self.root_state.is_white_turn() { + let result = if self.root_state.is_white_turn() { self.root_node .children .iter() .enumerate() - .max_by_key(|(_, child)| { + .max_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - child.entrance_count, - child.child_value, - self.rng.next_u32(), + a.completion_value(), + a.entrance_count as i32, + a.child_value, ) + .partial_cmp(&( + b.completion_value(), + b.entrance_count as i32, + b.child_value, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) - .expect("No valid action at root state!") + .expect("no valid action at root state") } else { self.root_node .children .iter() .enumerate() - .min_by_key(|(_, child)| { + .min_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - -(child.entrance_count as isize), - child.child_value, - self.rng.next_u32(), + a.completion_value(), + -(a.entrance_count as i32), + a.child_value, ) + .partial_cmp(&( + b.completion_value(), + -(b.entrance_count as i32), + b.child_value, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) - .expect("No valid action at root state!") + .expect("no valid action at root state") }; - - (val.0, val.1.action) + (result.0, result.1.action) } - /// This has the invariant that the board is NOT in a terminal state. pub fn new<'a>(board: &Board, evaluator: &'a Evaluator) -> GameSearchTree<'a> { - let mut cpy = board.clone(); let mut rng = SmallRng::seed_from_u64(123312); GameSearchTree { root_node: Box::new(SearchNode::build_self( - &mut cpy, + board, ApplyActionOutcome::Ongoing, evaluator, &mut rng, @@ -382,47 +379,42 @@ impl GameSearchTree<'_> { pub fn step_tree(&mut self, new_board: &Board, action_id: usize, outcome: ApplyActionOutcome) { self.root_state = new_board.clone(); - if self.root_node.children[action_id].node.is_some() { self.root_node = self.root_node.children[action_id].node.take().unwrap(); } else { self.root_node = Box::new(SearchNode::build_self( - &mut self.root_state.clone(), + &self.root_state, outcome, self.evaluator, &mut self.rng, - )) + )); } } pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let dcv = if self.root_state.is_white_turn() { - 1 - } else { - -1 - }; - - let start_time = std::time::Instant::now(); + let desired_completion = if self.root_state.is_white_turn() { 1 } else { -1 }; + let started_at = std::time::Instant::now(); if self.root_node.children.len() <= 1 { - return; // No need to compute with only one option. + return; } + for epoch in 0..iterations { - if start_time.elapsed() > max_duration && epoch >= 50 { + if started_at.elapsed() > max_duration && epoch >= 50 { break; } - let ba = self.get_best_action_index(); - if self.root_node.children[ba].entrance_count >= 6000 - && self.root_node.children[ba].completion_value() != -dcv + let best_action = self.get_best_action_index(); + if self.root_node.children[best_action].entrance_count >= 6000 + && self.root_node.children[best_action].completion_value() != -desired_completion { break; } - if self.root_node.children[ba].completion_value() == dcv { + if self.root_node.children[best_action].completion_value() == desired_completion { break; } self.root_node.ubfms_iteration( self.root_state.clone(), ApplyActionOutcome::Ongoing, - &self.evaluator, + self.evaluator, &mut self.rng, ); } @@ -440,3 +432,49 @@ impl GameSearchTree<'_> { self.safest_action() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::board::board_structs::Player; + + fn sample_board() -> Board { + Board::from_fen(crate::TRAINING_START_FENS[0]).unwrap() + } + + #[test] + fn value_from_term_compresses_terminal_scale() { + let mut board = sample_board(); + board.turn_count = 10; + let win = SearchNode::value_from_term(&board, TerminalState::Win(Player::White)); + let loss = SearchNode::value_from_term(&board, TerminalState::Win(Player::Black)); + assert!(win > 0.0); + assert!(loss < 0.0); + assert!(win.abs() < 10.0); + assert!(loss.abs() < 10.0); + } + + #[test] + fn faster_wins_score_higher() { + let mut fast = sample_board(); + let mut slow = sample_board(); + fast.turn_count = 20; + slow.turn_count = 120; + assert!( + SearchNode::value_from_term(&fast, TerminalState::Win(Player::White)) + > SearchNode::value_from_term(&slow, TerminalState::Win(Player::White)) + ); + } + + #[test] + fn faster_losses_score_lower() { + let mut fast = sample_board(); + let mut slow = sample_board(); + fast.turn_count = 20; + slow.turn_count = 120; + assert!( + SearchNode::value_from_term(&fast, TerminalState::Win(Player::Black)) + < SearchNode::value_from_term(&slow, TerminalState::Win(Player::Black)) + ); + } +} diff --git a/pyproject.toml b/pyproject.toml index 37e8b75..94f9b9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ "cython==3.0.11", "py-cpuinfo", "torch==2.10.0", - "triton>=3.6.0", + "triton>=3.6.0; sys_platform == 'linux'", "wandb>=0.25.1", ] classifiers = [ diff --git a/python/scripts/export_value_net.py b/python/scripts/export_value_net.py new file mode 100644 index 0000000..eea0cf7 --- /dev/null +++ b/python/scripts/export_value_net.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import save_file + + +def _rename_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + renamed: dict[str, torch.Tensor] = {} + + def copy(src: str, dst: str) -> None: + tensor = state_dict[src].detach().to(dtype=torch.float32, device="cpu").contiguous() + renamed[dst] = tensor + + copy("value_net.stem.0.weight", "stem.conv.weight") + copy("value_net.stem.1.weight", "stem.bn.weight") + copy("value_net.stem.1.bias", "stem.bn.bias") + copy("value_net.stem.1.running_mean", "stem.bn.running_mean") + copy("value_net.stem.1.running_var", "stem.bn.running_var") + + width = state_dict["value_net.stem.0.weight"].shape[0] + head_linear1 = state_dict["value_net.head.0.weight"] + hidden_dim = head_linear1.shape[0] + + block_ids = sorted( + { + int(key.split(".")[2]) + for key in state_dict + if key.startswith("value_net.blocks.") and key.endswith("conv1.weight") + } + ) + + for block_id in block_ids: + prefix = f"value_net.blocks.{block_id}" + copy(f"{prefix}.norm1.weight", f"blocks.{block_id}.norm1.weight") + copy(f"{prefix}.norm1.bias", f"blocks.{block_id}.norm1.bias") + copy(f"{prefix}.norm1.running_mean", f"blocks.{block_id}.norm1.running_mean") + copy(f"{prefix}.norm1.running_var", f"blocks.{block_id}.norm1.running_var") + copy(f"{prefix}.conv1.weight", f"blocks.{block_id}.conv1.weight") + copy(f"{prefix}.norm2.weight", f"blocks.{block_id}.norm2.weight") + copy(f"{prefix}.norm2.bias", f"blocks.{block_id}.norm2.bias") + copy(f"{prefix}.norm2.running_mean", f"blocks.{block_id}.norm2.running_mean") + copy(f"{prefix}.norm2.running_var", f"blocks.{block_id}.norm2.running_var") + copy(f"{prefix}.conv2.weight", f"blocks.{block_id}.conv2.weight") + + copy("value_net.head.0.weight", "head.fc1.weight") + copy("value_net.head.0.bias", "head.fc1.bias") + copy("value_net.head.2.weight", "head.fc2.weight") + copy("value_net.head.2.bias", "head.fc2.bias") + + renamed["__meta.width"] = torch.tensor([width], dtype=torch.float32) + renamed["__meta.hidden_dim"] = torch.tensor([hidden_dim], dtype=torch.float32) + renamed["__meta.num_blocks"] = torch.tensor([len(block_ids)], dtype=torch.float32) + return renamed + + +def _extract_state_dict(checkpoint: dict[str, Any]) -> dict[str, torch.Tensor]: + maybe_model = checkpoint.get("model") + if isinstance(maybe_model, dict): + return maybe_model + return checkpoint + + +def _build_fixtures( + checkpoint_path: Path, + fixture_path: Path, + fixture_count: int, + seed: int, +) -> None: + if fixture_count <= 0: + return + + from alphapaint_training import decode_packed_observation, sample_random_terminal_batch + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + state_dict = _extract_state_dict(checkpoint) + from alphapaint_training.model import PackedValueModel + + model = PackedValueModel(board_dtype=torch.float32) + model.load_state_dict(state_dict) + model.eval() + + packed_obs_np, _ = sample_random_terminal_batch(fixture_count, seed) + packed_obs = torch.from_numpy(packed_obs_np).to(dtype=torch.uint16) + board, intrinsics = decode_packed_observation( + packed_obs, + board_dtype=torch.float32, + intrinsic_dtype=torch.float32, + ) + with torch.no_grad(): + outputs = model(packed_obs.to(dtype=torch.uint16)).to(dtype=torch.float32) + + fixture_payload = { + "seed": seed, + "samples": [ + { + "board": board[idx].reshape(-1).tolist(), + "intrinsics": intrinsics[idx].tolist(), + "output": float(outputs[idx].item()), + } + for idx in range(fixture_count) + ], + } + fixture_path.write_text(json.dumps(fixture_payload, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Export AlphaPaint value network checkpoints for Rust CPU inference." + ) + parser.add_argument("--checkpoint", type=Path, required=True, help="Input .pt checkpoint") + parser.add_argument( + "--output", + type=Path, + default=Path("alpha_paint/assets/value_net.safetensors"), + help="Output safetensors path", + ) + parser.add_argument( + "--fixtures", + type=Path, + default=None, + help="Optional JSON fixture output with unpacked features and model outputs", + ) + parser.add_argument( + "--fixture-count", + type=int, + default=0, + help="Number of random terminal samples to export as fixtures", + ) + parser.add_argument("--seed", type=int, default=42, help="Fixture sampling seed") + args = parser.parse_args() + + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + state_dict = _extract_state_dict(checkpoint) + renamed = _rename_state_dict(state_dict) + + metadata = { + "width": str(int(renamed.pop("__meta.width").item())), + "num_blocks": str(int(renamed.pop("__meta.num_blocks").item())), + "hidden_dim": str(int(renamed.pop("__meta.hidden_dim").item())), + "source_checkpoint": str(args.checkpoint), + } + + args.output.parent.mkdir(parents=True, exist_ok=True) + save_file(renamed, str(args.output), metadata=metadata) + + if args.fixtures is not None: + args.fixtures.parent.mkdir(parents=True, exist_ok=True) + _build_fixtures(args.checkpoint, args.fixtures, args.fixture_count, args.seed) + + print(f"exported={args.output}") + if args.fixtures is not None and args.fixture_count > 0: + print(f"fixtures={args.fixtures}") + + +if __name__ == "__main__": + main() From 2a7d6438588974bbdfbadd17aeed753def59bea3 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 14:09:02 -0700 Subject: [PATCH 35/59] optimize tf out of resnet cpu eval --- Cargo.lock | 1 - alpha_paint/Cargo.toml | 1 - alpha_paint/benches/nn_latency.rs | 14 + alpha_paint/src/nn/model.rs | 645 ++++++++++++++++++------------ uv.lock | 4 +- 5 files changed, 399 insertions(+), 266 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e7f531f..8c81c59 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -31,7 +31,6 @@ name = "alpha_paint" version = "0.1.0" dependencies = [ "criterion", - "ndarray", "pyo3", "rand 0.10.0", "safetensors", diff --git a/alpha_paint/Cargo.toml b/alpha_paint/Cargo.toml index add05a4..f4da249 100644 --- a/alpha_paint/Cargo.toml +++ b/alpha_paint/Cargo.toml @@ -12,7 +12,6 @@ name = "nn_latency" harness = false [dependencies] -ndarray = "0.17.2" pyo3 = { version = "0.28.2", features = ["extension-module"] } rand = "0.10.0" safetensors = "0.7.0" diff --git a/alpha_paint/benches/nn_latency.rs b/alpha_paint/benches/nn_latency.rs index 78334ac..e3d6297 100644 --- a/alpha_paint/benches/nn_latency.rs +++ b/alpha_paint/benches/nn_latency.rs @@ -4,6 +4,7 @@ use std::hint::black_box; use alpha_paint::board::Board; use alpha_paint::evaluation::Evaluator; use alpha_paint::nn::features; +use alpha_paint::nn::model::ValueModel; use alpha_paint::search::GameSearchTree; use criterion::{Criterion, criterion_group, criterion_main}; @@ -28,6 +29,18 @@ fn bench_single_eval(c: &mut Criterion) { }); } +fn bench_forward_pass(c: &mut Criterion) { + let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[1]).unwrap(); + let features = features::extract(&board); + let model = ValueModel::benchmark_model(); + c.bench_function("value_model_forward", |b| { + b.iter(|| { + let value = black_box(&model).evaluate(black_box(&features)); + black_box(value); + }); + }); +} + fn bench_tree_step(c: &mut Criterion) { let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[2]).unwrap(); let evaluator = Evaluator::benchmark(); @@ -44,6 +57,7 @@ criterion_group!( benches, bench_feature_extraction, bench_single_eval, + bench_forward_pass, bench_tree_step ); criterion_main!(benches); diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs index 846be9f..bb28625 100644 --- a/alpha_paint/src/nn/model.rs +++ b/alpha_paint/src/nn/model.rs @@ -1,10 +1,24 @@ -use ndarray::{Array1, Array2}; -use safetensors::{Dtype, SafeTensors}; +use std::cell::UnsafeCell; + use safetensors::tensor::Metadata; +use safetensors::{Dtype, SafeTensors}; use crate::nn::features::{BOARD_CELLS, BOARD_PLANES, BOARD_SIDE, Features, INTRINSIC_COUNT}; const DEFAULT_EPS: f32 = 1.0e-5; +const MODEL_WIDTH: usize = 8; +const MODEL_BLOCKS: usize = 1; +const MODEL_HIDDEN_DIM: usize = 32; +const KERNEL_SIZE: usize = 9; +const ACTIVATION_LEN: usize = MODEL_WIDTH * BOARD_CELLS; +const STEM_WEIGHT_LEN: usize = MODEL_WIDTH * BOARD_PLANES * KERNEL_SIZE; +const BLOCK_WEIGHT_LEN: usize = MODEL_WIDTH * MODEL_WIDTH * KERNEL_SIZE; +const HEAD0_IN_DIM: usize = MODEL_WIDTH + INTRINSIC_COUNT; +const HEAD0_WEIGHT_LEN: usize = MODEL_HIDDEN_DIM * HEAD0_IN_DIM; + +thread_local! { + static SCRATCH: UnsafeCell = const { UnsafeCell::new(Scratch::new()) }; +} #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ModelConfig { @@ -16,69 +30,74 @@ pub struct ModelConfig { impl Default for ModelConfig { fn default() -> Self { Self { - width: 8, - num_blocks: 1, - hidden_dim: 32, + width: MODEL_WIDTH, + num_blocks: MODEL_BLOCKS, + hidden_dim: MODEL_HIDDEN_DIM, } } } #[derive(Clone, Debug)] -pub struct Conv2d { - out_channels: usize, - in_channels: usize, - weights: Vec, -} - -#[derive(Clone, Debug)] -pub struct BatchNorm2d { - weight: Vec, - bias: Vec, - running_mean: Vec, - running_var: Vec, - eps: f32, +pub struct ValueModel { + pub config: ModelConfig, + stem_weight: [f32; STEM_WEIGHT_LEN], + stem_bias: [f32; MODEL_WIDTH], + block0_norm1_scale: [f32; MODEL_WIDTH], + block0_norm1_bias: [f32; MODEL_WIDTH], + block0_conv1_weight: [f32; BLOCK_WEIGHT_LEN], + block0_conv1_bias: [f32; MODEL_WIDTH], + block0_conv2_weight: [f32; BLOCK_WEIGHT_LEN], + head0_weight: [f32; HEAD0_WEIGHT_LEN], + head0_bias: [f32; MODEL_HIDDEN_DIM], + head1_weight: [f32; MODEL_HIDDEN_DIM], + head1_bias: f32, } -#[derive(Clone, Debug)] -pub struct Linear { - weight: Array2, - bias: Array1, +struct Scratch { + a: [f32; ACTIVATION_LEN], + b: [f32; ACTIVATION_LEN], + residual: [f32; ACTIVATION_LEN], + pooled: [f32; MODEL_WIDTH], + head_input: [f32; HEAD0_IN_DIM], + hidden: [f32; MODEL_HIDDEN_DIM], } -#[derive(Clone, Debug)] -pub struct ResidualBlock { - norm1: BatchNorm2d, - conv1: Conv2d, - norm2: BatchNorm2d, - conv2: Conv2d, +impl Scratch { + const fn new() -> Self { + Self { + a: [0.0; ACTIVATION_LEN], + b: [0.0; ACTIVATION_LEN], + residual: [0.0; ACTIVATION_LEN], + pooled: [0.0; MODEL_WIDTH], + head_input: [0.0; HEAD0_IN_DIM], + hidden: [0.0; MODEL_HIDDEN_DIM], + } + } } -#[derive(Clone, Debug)] -pub struct ValueModel { - pub config: ModelConfig, - stem_conv: Conv2d, - stem_norm: BatchNorm2d, - blocks: Vec, - head_linear1: Linear, - head_linear2: Linear, +struct BatchNormStats { + weight: [f32; C], + bias: [f32; C], + running_mean: [f32; C], + running_var: [f32; C], } impl ValueModel { pub fn zeroed(config: ModelConfig) -> Self { - let stem_conv = Conv2d::zeros(config.width, BOARD_PLANES); - let stem_norm = BatchNorm2d::identity(config.width); - let blocks = (0..config.num_blocks) - .map(|_| ResidualBlock::zeroed(config.width)) - .collect(); - let head_linear1 = Linear::zeros(config.hidden_dim, config.width + INTRINSIC_COUNT); - let head_linear2 = Linear::zeros(1, config.hidden_dim); + validate_config(config).unwrap(); Self { config, - stem_conv, - stem_norm, - blocks, - head_linear1, - head_linear2, + stem_weight: [0.0; STEM_WEIGHT_LEN], + stem_bias: [0.0; MODEL_WIDTH], + block0_norm1_scale: [1.0; MODEL_WIDTH], + block0_norm1_bias: [0.0; MODEL_WIDTH], + block0_conv1_weight: [0.0; BLOCK_WEIGHT_LEN], + block0_conv1_bias: [0.0; MODEL_WIDTH], + block0_conv2_weight: [0.0; BLOCK_WEIGHT_LEN], + head0_weight: [0.0; HEAD0_WEIGHT_LEN], + head0_bias: [0.0; MODEL_HIDDEN_DIM], + head1_weight: [0.0; MODEL_HIDDEN_DIM], + head1_bias: 0.0, } } @@ -111,236 +130,346 @@ impl ValueModel { tensors: &SafeTensors<'_>, config: ModelConfig, ) -> Result { - let stem_conv = Conv2d::from_tensor(tensors, "stem.conv.weight", config.width, BOARD_PLANES)?; - let stem_norm = BatchNorm2d::from_prefix(tensors, "stem.bn", config.width)?; - - let mut blocks = Vec::with_capacity(config.num_blocks); - for idx in 0..config.num_blocks { - let prefix = format!("blocks.{idx}"); - blocks.push(ResidualBlock { - norm1: BatchNorm2d::from_prefix( - tensors, - &format!("{prefix}.norm1"), - config.width, - )?, - conv1: Conv2d::from_tensor( - tensors, - &format!("{prefix}.conv1.weight"), - config.width, - config.width, - )?, - norm2: BatchNorm2d::from_prefix( - tensors, - &format!("{prefix}.norm2"), - config.width, - )?, - conv2: Conv2d::from_tensor( - tensors, - &format!("{prefix}.conv2.weight"), - config.width, - config.width, - )?, - }); - } + validate_config(config)?; + + let stem_weight = tensor_array::( + tensors, + "stem.conv.weight", + &[MODEL_WIDTH, BOARD_PLANES, 3, 3], + )?; + let stem_bn = batch_norm_from_prefix::(tensors, "stem.bn")?; + let (stem_weight, stem_bias) = fold_conv_bn(stem_weight, stem_bn, BOARD_PLANES); + + let block0_norm1 = batch_norm_from_prefix::(tensors, "blocks.0.norm1")?; + let (block0_norm1_scale, block0_norm1_bias) = bn_to_affine(block0_norm1); - let head_linear1 = Linear::from_prefix( + let block0_conv1_weight = tensor_array::( tensors, - "head.fc1", - config.hidden_dim, - config.width + INTRINSIC_COUNT, + "blocks.0.conv1.weight", + &[MODEL_WIDTH, MODEL_WIDTH, 3, 3], )?; - let head_linear2 = Linear::from_prefix(tensors, "head.fc2", 1, config.hidden_dim)?; + let block0_norm2 = batch_norm_from_prefix::(tensors, "blocks.0.norm2")?; + let (block0_conv1_weight, block0_conv1_bias) = + fold_conv_bn(block0_conv1_weight, block0_norm2, MODEL_WIDTH); + + let block0_conv2_weight = tensor_array::( + tensors, + "blocks.0.conv2.weight", + &[MODEL_WIDTH, MODEL_WIDTH, 3, 3], + )?; + + let head0_weight = tensor_array::( + tensors, + "head.fc1.weight", + &[MODEL_HIDDEN_DIM, HEAD0_IN_DIM], + )?; + let head0_bias = tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; + let head1_weight = + tensor_array::(tensors, "head.fc2.weight", &[1, MODEL_HIDDEN_DIM])?; + let head1_bias = tensor_array::<1>(tensors, "head.fc2.bias", &[1])?[0]; Ok(Self { config, - stem_conv, - stem_norm, - blocks, - head_linear1, - head_linear2, + stem_weight, + stem_bias, + block0_norm1_scale, + block0_norm1_bias, + block0_conv1_weight, + block0_conv1_bias, + block0_conv2_weight, + head0_weight, + head0_bias, + head1_weight, + head1_bias, }) } + #[inline(always)] pub fn evaluate(&self, features: &Features) -> f32 { - let width = self.config.width; - let plane_span = width * BOARD_CELLS; - let mut activations = vec![0.0; plane_span]; - let mut scratch_a = vec![0.0; plane_span]; - let mut scratch_b = vec![0.0; plane_span]; - - self.stem_conv.forward(&features.board, &mut activations); - self.stem_norm.relu_inplace(&mut activations); - - for block in &self.blocks { - scratch_a.copy_from_slice(&activations); - block.norm1.relu_inplace(&mut scratch_a); - block.conv1.forward(&scratch_a, &mut scratch_b); - block.norm2.relu_inplace(&mut scratch_b); - block.conv2.forward(&scratch_b, &mut scratch_a); - for (dst, residual) in activations.iter_mut().zip(&scratch_a) { - *dst += *residual; - } - } + SCRATCH.with(|scratch| { + // SAFETY: each thread has its own scratch buffer, and evaluate does not retain + // references beyond this call. + let scratch = unsafe { &mut *scratch.get() }; + self.forward(&features.board, &features.intrinsics, scratch) + }) + } - let pooled = average_pool_channels(&activations, width); - let mut head_input = Array1::zeros(width + INTRINSIC_COUNT); - for (dst, src) in head_input.iter_mut().take(width).zip(pooled) { - *dst = src; - } - for (dst, src) in head_input - .iter_mut() - .skip(width) - .zip(features.intrinsics.iter().copied()) - { - *dst = src; + #[inline(always)] + fn forward( + &self, + board: &[f32], + intrinsics: &[f32; INTRINSIC_COUNT], + scratch: &mut Scratch, + ) -> f32 { + debug_assert_eq!(board.len(), BOARD_PLANES * BOARD_CELLS); + + conv3x3_bias::( + &self.stem_weight, + &self.stem_bias, + board, + &mut scratch.a, + ); + relu_inplace(&mut scratch.a); + + scratch.residual.copy_from_slice(&scratch.a); + affine_relu_inplace::( + &mut scratch.a, + &self.block0_norm1_scale, + &self.block0_norm1_bias, + ); + conv3x3_bias::( + &self.block0_conv1_weight, + &self.block0_conv1_bias, + &scratch.a, + &mut scratch.b, + ); + relu_inplace(&mut scratch.b); + conv3x3_no_bias::( + &self.block0_conv2_weight, + &scratch.b, + &mut scratch.a, + ); + for idx in 0..ACTIVATION_LEN { + scratch.a[idx] += scratch.residual[idx]; } - let mut hidden = self.head_linear1.forward(&head_input); - hidden.mapv_inplace(|value: f32| value.max(0.0)); - let output = self.head_linear2.forward(&hidden); - output[0].tanh() * 7.0 + global_avg_pool::(&scratch.a, &mut scratch.pooled); + scratch.head_input[..MODEL_WIDTH].copy_from_slice(&scratch.pooled); + scratch.head_input[MODEL_WIDTH..].copy_from_slice(intrinsics); + linear_relu::( + &self.head0_weight, + &self.head0_bias, + &scratch.head_input, + &mut scratch.hidden, + ); + (linear_scalar::( + &self.head1_weight, + self.head1_bias, + &scratch.hidden, + )) + .tanh() + * 7.0 } } -impl Conv2d { - fn zeros(out_channels: usize, in_channels: usize) -> Self { - Self { - out_channels, - in_channels, - weights: vec![0.0; out_channels * in_channels * 3 * 3], +#[inline(always)] +fn batch_norm_from_prefix( + tensors: &SafeTensors<'_>, + prefix: &str, +) -> Result, String> { + Ok(BatchNormStats { + weight: tensor_array::(tensors, &format!("{prefix}.weight"), &[C])?, + bias: tensor_array::(tensors, &format!("{prefix}.bias"), &[C])?, + running_mean: tensor_array::(tensors, &format!("{prefix}.running_mean"), &[C])?, + running_var: tensor_array::(tensors, &format!("{prefix}.running_var"), &[C])?, + }) +} + +#[inline(always)] +fn bn_to_affine(bn: BatchNormStats) -> ([f32; C], [f32; C]) { + let mut scale = [0.0; C]; + let mut bias = [0.0; C]; + for channel in 0..C { + let channel_scale = bn.weight[channel] / (bn.running_var[channel] + DEFAULT_EPS).sqrt(); + scale[channel] = channel_scale; + bias[channel] = bn.bias[channel] - bn.running_mean[channel] * channel_scale; + } + (scale, bias) +} + +#[inline(always)] +fn fold_conv_bn( + mut conv_weight: [f32; W], + bn: BatchNormStats, + in_channels: usize, +) -> ([f32; W], [f32; C]) { + let (scale, bias) = bn_to_affine(bn); + let channel_stride = in_channels * KERNEL_SIZE; + for channel in 0..C { + let start = channel * channel_stride; + let end = start + channel_stride; + for value in &mut conv_weight[start..end] { + *value *= scale[channel]; } } + (conv_weight, bias) +} - fn from_tensor( - tensors: &SafeTensors<'_>, - name: &str, - out_channels: usize, - in_channels: usize, - ) -> Result { - let data = tensor_f32(tensors, name, &[out_channels, in_channels, 3, 3])?; - Ok(Self { - out_channels, - in_channels, - weights: data, - }) +#[inline(always)] +fn conv3x3_bias( + weight: &[f32], + bias: &[f32; OUT_C], + input: &[f32], + output: &mut [f32; ACTIVATION_LEN], +) { + for out_channel in 0..OUT_C { + let out_plane = &mut output[out_channel * BOARD_CELLS..(out_channel + 1) * BOARD_CELLS]; + out_plane.fill(bias[out_channel]); + accumulate_conv3x3::(weight, out_channel, input, out_plane); + } +} + +#[inline(always)] +fn conv3x3_no_bias( + weight: &[f32], + input: &[f32; ACTIVATION_LEN], + output: &mut [f32; ACTIVATION_LEN], +) { + for out_channel in 0..OUT_C { + let out_plane = &mut output[out_channel * BOARD_CELLS..(out_channel + 1) * BOARD_CELLS]; + out_plane.fill(0.0); + accumulate_conv3x3::(weight, out_channel, input, out_plane); } +} - fn forward(&self, input: &[f32], out: &mut [f32]) { - debug_assert_eq!(input.len(), self.in_channels * BOARD_CELLS); - debug_assert_eq!(out.len(), self.out_channels * BOARD_CELLS); - out.fill(0.0); - - for out_ch in 0..self.out_channels { - for y in 0..BOARD_SIDE { - for x in 0..BOARD_SIDE { - let mut sum = 0.0; - for in_ch in 0..self.in_channels { - for ky in 0..3 { - let iy = y as isize + ky as isize - 1; - if !(0..BOARD_SIDE as isize).contains(&iy) { - continue; - } - for kx in 0..3 { - let ix = x as isize + kx as isize - 1; - if !(0..BOARD_SIDE as isize).contains(&ix) { - continue; - } - let input_idx = in_ch * BOARD_CELLS - + iy as usize * BOARD_SIDE - + ix as usize; - let weight_idx = - (((out_ch * self.in_channels + in_ch) * 3 + ky) * 3) + kx; - sum += input[input_idx] * self.weights[weight_idx]; - } - } - } - out[out_ch * BOARD_CELLS + y * BOARD_SIDE + x] = sum; +#[inline(always)] +fn accumulate_conv3x3( + weight: &[f32], + out_channel: usize, + input: &[f32], + out_plane: &mut [f32], +) { + let out_ptr = out_plane.as_mut_ptr(); + for in_channel in 0..IN_C { + let in_plane = &input[in_channel * BOARD_CELLS..(in_channel + 1) * BOARD_CELLS]; + let kernel = &weight[(out_channel * IN_C + in_channel) * KERNEL_SIZE..][..KERNEL_SIZE]; + let k00 = kernel[0]; + let k01 = kernel[1]; + let k02 = kernel[2]; + let k10 = kernel[3]; + let k11 = kernel[4]; + let k12 = kernel[5]; + let k20 = kernel[6]; + let k21 = kernel[7]; + let k22 = kernel[8]; + let in_ptr = in_plane.as_ptr(); + + unsafe { + for y in 1..BOARD_SIDE - 1 { + let row = y * BOARD_SIDE; + let row_up = row - BOARD_SIDE; + let row_down = row + BOARD_SIDE; + for x in 1..BOARD_SIDE - 1 { + let idx = row + x; + *out_ptr.add(idx) += *in_ptr.add(row_up + x - 1) * k00 + + *in_ptr.add(row_up + x) * k01 + + *in_ptr.add(row_up + x + 1) * k02 + + *in_ptr.add(row + x - 1) * k10 + + *in_ptr.add(row + x) * k11 + + *in_ptr.add(row + x + 1) * k12 + + *in_ptr.add(row_down + x - 1) * k20 + + *in_ptr.add(row_down + x) * k21 + + *in_ptr.add(row_down + x + 1) * k22; } } } + + for x in 0..BOARD_SIDE { + out_plane[x] += conv3x3_border(in_plane, kernel, x, 0); + out_plane[(BOARD_SIDE - 1) * BOARD_SIDE + x] += + conv3x3_border(in_plane, kernel, x, BOARD_SIDE - 1); + } + for y in 1..BOARD_SIDE - 1 { + let row = y * BOARD_SIDE; + out_plane[row] += conv3x3_border(in_plane, kernel, 0, y); + out_plane[row + BOARD_SIDE - 1] += + conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); + } } } -impl BatchNorm2d { - fn identity(width: usize) -> Self { - Self { - weight: vec![1.0; width], - bias: vec![0.0; width], - running_mean: vec![0.0; width], - running_var: vec![1.0; width], - eps: DEFAULT_EPS, +#[inline(always)] +fn conv3x3_border(in_plane: &[f32], kernel: &[f32], x: usize, y: usize) -> f32 { + let mut sum = 0.0; + for ky in 0..3 { + let input_y = y as isize + ky as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&input_y) { + continue; + } + let row = input_y as usize * BOARD_SIDE; + for kx in 0..3 { + let input_x = x as isize + kx as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&input_x) { + continue; + } + sum += in_plane[row + input_x as usize] * kernel[ky * 3 + kx]; } } + sum +} - fn from_prefix(tensors: &SafeTensors<'_>, prefix: &str, width: usize) -> Result { - let weight = tensor_f32(tensors, &format!("{prefix}.weight"), &[width])?; - let bias = tensor_f32(tensors, &format!("{prefix}.bias"), &[width])?; - let running_mean = tensor_f32(tensors, &format!("{prefix}.running_mean"), &[width])?; - let running_var = tensor_f32(tensors, &format!("{prefix}.running_var"), &[width])?; - Ok(Self { - weight, - bias, - running_mean, - running_var, - eps: DEFAULT_EPS, - }) +#[inline(always)] +fn affine_relu_inplace( + values: &mut [f32; ACTIVATION_LEN], + scale: &[f32; C], + bias: &[f32; C], +) { + for channel in 0..C { + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + let channel_scale = scale[channel]; + let channel_bias = bias[channel]; + for value in &mut values[start..end] { + *value = (*value * channel_scale + channel_bias).max(0.0); + } } +} - fn relu_inplace(&self, values: &mut [f32]) { - debug_assert_eq!(values.len(), self.weight.len() * BOARD_CELLS); - for channel in 0..self.weight.len() { - let scale = self.weight[channel] / (self.running_var[channel] + self.eps).sqrt(); - let offset = self.bias[channel] - self.running_mean[channel] * scale; - let start = channel * BOARD_CELLS; - let end = start + BOARD_CELLS; - for value in &mut values[start..end] { - *value = (*value * scale + offset).max(0.0); - } - } +#[inline(always)] +fn relu_inplace(values: &mut [f32]) { + for value in values { + *value = value.max(0.0); } } -impl Linear { - fn zeros(out_features: usize, in_features: usize) -> Self { - Self { - weight: Array2::zeros((out_features, in_features)), - bias: Array1::zeros(out_features), +#[inline(always)] +fn global_avg_pool(input: &[f32; ACTIVATION_LEN], output: &mut [f32; C]) { + let scale = 1.0 / BOARD_CELLS as f32; + for channel in 0..C { + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + let mut sum = 0.0; + for value in &input[start..end] { + sum += *value; } + output[channel] = sum * scale; } +} - fn from_prefix( - tensors: &SafeTensors<'_>, - prefix: &str, - out_features: usize, - in_features: usize, - ) -> Result { - let weight = tensor_f32( - tensors, - &format!("{prefix}.weight"), - &[out_features, in_features], - )?; - let bias = tensor_f32(tensors, &format!("{prefix}.bias"), &[out_features])?; - let weight = Array2::from_shape_vec((out_features, in_features), weight) - .map_err(|err| format!("failed to shape {prefix}.weight: {err}"))?; - let bias = Array1::from_vec(bias); - Ok(Self { weight, bias }) +#[inline(always)] +fn linear_relu( + weight: &[f32], + bias: &[f32; OUT], + input: &[f32; IN], + output: &mut [f32; OUT], +) { + for out_idx in 0..OUT { + let mut sum = bias[out_idx]; + let row = &weight[out_idx * IN..(out_idx + 1) * IN]; + for in_idx in 0..IN { + sum += row[in_idx] * input[in_idx]; + } + output[out_idx] = sum.max(0.0); } +} - fn forward(&self, input: &Array1) -> Array1 { - let mut out = self.weight.dot(input); - out += &self.bias; - out +#[inline(always)] +fn linear_scalar(weight: &[f32; IN], bias: f32, input: &[f32; IN]) -> f32 { + let mut sum = bias; + for idx in 0..IN { + sum += weight[idx] * input[idx]; } + sum } -impl ResidualBlock { - fn zeroed(width: usize) -> Self { - Self { - norm1: BatchNorm2d::identity(width), - conv1: Conv2d::zeros(width, width), - norm2: BatchNorm2d::identity(width), - conv2: Conv2d::zeros(width, width), - } +fn validate_config(config: ModelConfig) -> Result<(), String> { + if config == ModelConfig::default() { + Ok(()) + } else { + Err(format!( + "only the default value model is supported (got width={}, num_blocks={}, hidden_dim={})", + config.width, config.num_blocks, config.hidden_dim + )) } } @@ -363,6 +492,17 @@ fn load_config(metadata: &Metadata) -> Result { }) } +fn tensor_array( + tensors: &SafeTensors<'_>, + name: &str, + expected_shape: &[usize], +) -> Result<[f32; N], String> { + let values = tensor_f32(tensors, name, expected_shape)?; + values + .try_into() + .map_err(|got: Vec| format!("tensor {name} has {} values, expected {N}", got.len())) +} + fn tensor_f32( tensors: &SafeTensors<'_>, name: &str, @@ -395,17 +535,6 @@ fn tensor_f32( Ok(out) } -fn average_pool_channels(values: &[f32], channels: usize) -> Vec { - let mut pooled = vec![0.0; channels]; - let scale = 1.0 / BOARD_CELLS as f32; - for channel in 0..channels { - let start = channel * BOARD_CELLS; - let end = start + BOARD_CELLS; - pooled[channel] = values[start..end].iter().copied().sum::() * scale; - } - pooled -} - #[cfg(test)] mod tests { use super::*; @@ -413,14 +542,10 @@ mod tests { #[test] fn forward_matches_simple_hand_computation() { - let mut model = ValueModel::zeroed(ModelConfig { - width: 1, - num_blocks: 0, - hidden_dim: 1, - }); - model.stem_conv.weights[4] = 1.0; - model.head_linear1.weight[[0, 0]] = 1.0; - model.head_linear2.weight[[0, 0]] = 0.5; + let mut model = ValueModel::zeroed(ModelConfig::default()); + model.stem_weight[4] = 1.0; + model.head0_weight[0] = 1.0; + model.head1_weight[0] = 0.5; let mut features = Features::zeros(); for y in 0..BOARD_SIDE as u8 { @@ -435,14 +560,10 @@ mod tests { #[test] fn residual_block_preserves_input_when_convs_are_zero() { - let mut model = ValueModel::zeroed(ModelConfig { - width: 1, - num_blocks: 1, - hidden_dim: 1, - }); - model.stem_conv.weights[4] = 1.0; - model.head_linear1.weight[[0, 0]] = 1.0; - model.head_linear2.weight[[0, 0]] = 1.0; + let mut model = ValueModel::zeroed(ModelConfig::default()); + model.stem_weight[4] = 1.0; + model.head0_weight[0] = 1.0; + model.head1_weight[0] = 1.0; let mut features = Features::zeros(); for y in 0..BOARD_SIDE as u8 { diff --git a/uv.lock b/uv.lock index 53dfac2..749aad0 100644 --- a/uv.lock +++ b/uv.lock @@ -12,7 +12,7 @@ dependencies = [ { name = "psutil" }, { name = "py-cpuinfo" }, { name = "torch" }, - { name = "triton" }, + { name = "triton", marker = "sys_platform == 'linux'" }, { name = "wandb" }, ] @@ -30,7 +30,7 @@ requires-dist = [ { name = "psutil", specifier = "==5.9.0" }, { name = "py-cpuinfo" }, { name = "torch", specifier = "==2.10.0" }, - { name = "triton", specifier = ">=3.6.0" }, + { name = "triton", marker = "sys_platform == 'linux'", specifier = ">=3.6.0" }, { name = "wandb", specifier = ">=0.25.1" }, ] From c5848e901b27cfaf3966794ee4992555c8d74ba7 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 18:02:37 -0400 Subject: [PATCH 36/59] train a larger resnet --- python/alphapaint_training/train.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 970b25f..f914a67 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -26,14 +26,17 @@ class TrainConfig: rounds: int = 1 samples_per_round: int = 1_048_576 train_steps_per_round: int = 128 - batch_size: int = 24_576 + batch_size: int = 8_192 replay_capacity: int = 16_000_000 num_threads: int = 32 workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 4 * 1024 + max_gpu_evals_per_move: int = 1_536 lr: float = 3e-4 + width: int = 20 + num_blocks: int = 1 + hidden_dim: int = 64 pretrain_terminal_samples: int = 5_000_000 - pretrain_batch_size: int = 24_576 + pretrain_batch_size: int = 8_192 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" From 3d163dd047f535b606b38890296bf809ecbb3527 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 18:06:27 -0400 Subject: [PATCH 37/59] wip: multigpu --- alpha_paint/benches/nn_latency.rs | 2 +- alpha_paint/src/nn/features.rs | 28 ++-- alpha_paint/src/nn/model.rs | 19 ++- alpha_paint/src/search.rs | 28 ++-- .../alphapaint_training/cudagraph_backend.py | 128 ++++++++++++------ python/alphapaint_training/packed_obs.py | 25 ++-- python/alphapaint_training/train.py | 45 +++++- training/src/cudagraph.rs | 51 ++++++- training/src/executor.rs | 28 ++++ training/src/lib.rs | 125 ++++++++++++++--- 10 files changed, 367 insertions(+), 112 deletions(-) diff --git a/alpha_paint/benches/nn_latency.rs b/alpha_paint/benches/nn_latency.rs index e3d6297..2ecd24f 100644 --- a/alpha_paint/benches/nn_latency.rs +++ b/alpha_paint/benches/nn_latency.rs @@ -1,5 +1,5 @@ -use std::time::Duration; use std::hint::black_box; +use std::time::Duration; use alpha_paint::board::Board; use alpha_paint::evaluation::Evaluator; diff --git a/alpha_paint/src/nn/features.rs b/alpha_paint/src/nn/features.rs index 2f73455..1a09a65 100644 --- a/alpha_paint/src/nn/features.rs +++ b/alpha_paint/src/nn/features.rs @@ -1,6 +1,6 @@ +use crate::board::Board; use crate::board::board_structs::Player; use crate::board::structs::Coordinate; -use crate::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; @@ -168,8 +168,7 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; - out[INTRINSIC_OPPONENT_BEACONS] = - board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; } else { out[INTRINSIC_CURRENT_STAMINA] = board.black_stamina as f32; out[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina as f32; @@ -177,8 +176,7 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; - out[INTRINSIC_CURRENT_BEACONS] = - board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; } @@ -221,9 +219,18 @@ mod tests { assert_eq!(features.board[plane_index(WALL, 31, 31)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 99.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 88.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_CURRENT_TERRITORY], 1.0 / 1024.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], 1.0 / 1024.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 88.0 / 420.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_CURRENT_TERRITORY], + 1.0 / 1024.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], + 1.0 / 1024.0 + ); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_BEACONS], 1.0 / 1024.0); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 2.0 / 64.0); } @@ -244,7 +251,10 @@ mod tests { assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 0, 0)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 88.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 99.0 / 420.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 99.0 / 420.0 + ); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 5.0 / 64.0); } } diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs index bb28625..823da96 100644 --- a/alpha_paint/src/nn/model.rs +++ b/alpha_paint/src/nn/model.rs @@ -163,7 +163,8 @@ impl ValueModel { "head.fc1.weight", &[MODEL_HIDDEN_DIM, HEAD0_IN_DIM], )?; - let head0_bias = tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; + let head0_bias = + tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; let head1_weight = tensor_array::(tensors, "head.fc2.weight", &[1, MODEL_HIDDEN_DIM])?; let head1_bias = tensor_array::<1>(tensors, "head.fc2.bias", &[1])?[0]; @@ -242,12 +243,8 @@ impl ValueModel { &scratch.head_input, &mut scratch.hidden, ); - (linear_scalar::( - &self.head1_weight, - self.head1_bias, - &scratch.hidden, - )) - .tanh() + (linear_scalar::(&self.head1_weight, self.head1_bias, &scratch.hidden)) + .tanh() * 7.0 } } @@ -372,8 +369,7 @@ fn accumulate_conv3x3( for y in 1..BOARD_SIDE - 1 { let row = y * BOARD_SIDE; out_plane[row] += conv3x3_border(in_plane, kernel, 0, y); - out_plane[row + BOARD_SIDE - 1] += - conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); + out_plane[row + BOARD_SIDE - 1] += conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); } } } @@ -526,7 +522,10 @@ fn tensor_f32( } let data = tensor.data(); if data.len() % 4 != 0 { - return Err(format!("tensor {name} has invalid byte length {}", data.len())); + return Err(format!( + "tensor {name} has invalid byte length {}", + data.len() + )); } let mut out = Vec::with_capacity(data.len() / 4); for chunk in data.chunks_exact(4) { diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index e4015e4..9afe8a1 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -131,9 +131,13 @@ impl SearchNode { if self.completion_value.abs() == 1 { true } else { - self.children - .iter() - .all(|child| child.node.as_ref().map(|node| node.resolved).unwrap_or(false)) + self.children.iter().all(|child| { + child + .node + .as_ref() + .map(|node| node.resolved) + .unwrap_or(false) + }) } } @@ -254,7 +258,11 @@ impl SearchNode { let (outcome, _) = state.apply_action(action); let node = Box::new(SearchNode::build_self(&state, outcome, evaluator, rng)); let value = node.value; - if let Some(index) = self.children.iter().position(|child| child.action == action) { + if let Some(index) = self + .children + .iter() + .position(|child| child.action == action) + { self.children[index].node = Some(node); } value @@ -324,11 +332,7 @@ impl GameSearchTree<'_> { .iter() .enumerate() .max_by(|&(_, a), &(_, b)| { - ( - a.completion_value(), - a.entrance_count as i32, - a.child_value, - ) + (a.completion_value(), a.entrance_count as i32, a.child_value) .partial_cmp(&( b.completion_value(), b.entrance_count as i32, @@ -392,7 +396,11 @@ impl GameSearchTree<'_> { } pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let desired_completion = if self.root_state.is_white_turn() { 1 } else { -1 }; + let desired_completion = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; let started_at = std::time::Instant::now(); if self.root_node.children.len() <= 1 { return; diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index dd62b0d..fffdc1c 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -10,6 +10,48 @@ import torch.utils.dlpack as dlpack +def _validate_lane_tensors( + obs_host: torch.Tensor, + obs_device: torch.Tensor, + value_host: torch.Tensor, + value_device: torch.Tensor, + gpu_id: int, +) -> None: + if obs_host.device.type != "cpu": + raise ValueError("obs_host must be a CPU tensor") + if value_host.device.type != "cpu": + raise ValueError("value_host must be a CPU tensor") + if obs_device.device.type != "cuda": + raise ValueError("obs_device must be a CUDA tensor") + if value_device.device.type != "cuda": + raise ValueError("value_device must be a CUDA tensor") + + if obs_host.dtype != torch.uint16: + raise ValueError(f"obs_host must be uint16, got {obs_host.dtype}") + if obs_device.dtype != torch.uint16: + raise ValueError(f"obs_device must be uint16, got {obs_device.dtype}") + if value_host.dtype != torch.float32: + raise ValueError(f"value_host must be float32, got {value_host.dtype}") + if value_device.dtype != torch.float32: + raise ValueError(f"value_device must be float32, got {value_device.dtype}") + + if obs_host.shape != obs_device.shape: + raise ValueError( + f"obs_host and obs_device shape mismatch: {obs_host.shape} vs {obs_device.shape}" + ) + batch = obs_host.shape[0] + if value_host.shape != (batch,): + raise ValueError(f"value_host must be shape ({batch},), got {value_host.shape}") + if value_device.shape != (batch,): + raise ValueError( + f"value_device must be shape ({batch},), got {value_device.shape}" + ) + + for name, tensor in (("obs_device", obs_device), ("value_device", value_device)): + if tensor.device.index != gpu_id: + raise ValueError(f"{name} must be on cuda:{gpu_id}, got {tensor.device}") + + def _autocast_dtype(precision: str) -> Optional[torch.dtype]: if precision == "fp32": return None @@ -30,6 +72,7 @@ def capture_lane_graph( value_device_dlpack, stream_handle: int, precision: str = "bf16", + gpu_id: int = 0, ) -> tuple[int, object]: """Capture a CUDA graph for value-only inference. @@ -48,49 +91,56 @@ def capture_lane_graph( Returns: Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). """ + if gpu_id < 0: + raise ValueError(f"gpu_id must be >= 0, got {gpu_id}") + obs_host = dlpack.from_dlpack(obs_host_dlpack) obs_device = dlpack.from_dlpack(obs_device_dlpack) value_host = dlpack.from_dlpack(value_host_dlpack) value_device = dlpack.from_dlpack(value_device_dlpack) - model = model.cuda() - model = model.to(memory_format=torch.channels_last) - model.eval() - torch.backends.cudnn.benchmark = True - stream = torch.cuda.ExternalStream(stream_handle) - graph = torch.cuda.CUDAGraph(keep_graph=True) - dtype = _autocast_dtype(precision) - - def run_step() -> None: - obs_device.copy_(obs_host, non_blocking=True) - if dtype is None: - value = model(obs_device) - else: - with torch.autocast(device_type="cuda", dtype=dtype): + _validate_lane_tensors(obs_host, obs_device, value_host, value_device, gpu_id) + + with torch.cuda.device(gpu_id): + model = model.to(f"cuda:{gpu_id}") + model = model.to(memory_format=torch.channels_last) + model.eval() + torch.backends.cudnn.benchmark = True + stream = torch.cuda.ExternalStream(stream_handle) + graph = torch.cuda.CUDAGraph(keep_graph=True) + dtype = _autocast_dtype(precision) + + def run_step() -> None: + obs_device.copy_(obs_host, non_blocking=True) + if dtype is None: value = model(obs_device) - # Model returns value tensor, shape (B,) or (B, 1) - if value.ndim == 2: - value = value.squeeze(-1) - value_device.copy_(value, non_blocking=True) - value_host.copy_(value_device, non_blocking=True) - - with torch.inference_mode(): - with torch.cuda.stream(stream): - for _ in range(3): + else: + with torch.autocast(device_type="cuda", dtype=dtype): + value = model(obs_device) + if value.ndim == 2: + value = value.squeeze(-1) + value_device.copy_(value, non_blocking=True) + value_host.copy_(value_device, non_blocking=True) + + with torch.inference_mode(): + with torch.cuda.stream(stream): + for _ in range(3): + run_step() + torch.cuda.synchronize(device=gpu_id) + + with torch.cuda.graph( + graph, stream=stream, capture_error_mode="thread_local" + ): run_step() - torch.cuda.synchronize() - - with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"): - run_step() - - graph.instantiate() - owner = ( - graph, - model, - obs_host, - obs_device, - value_host, - value_device, - stream, - ) - return int(graph.raw_cuda_graph_exec()), owner + + graph.instantiate() + owner = ( + graph, + model, + obs_host, + obs_device, + value_host, + value_device, + stream, + ) + return int(graph.raw_cuda_graph_exec()), owner diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index a7d9003..50f85fe 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -229,18 +229,19 @@ def decode_packed_board( raise ValueError(f"out must have dtype {dtype}, got {out.dtype}") total_cells = board_words.shape[0] * BOARD_CELLS grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) - _decode_board_kernel[grid]( - board_words, - out, - total_cells, - board_words.stride(0), - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - BLOCK=256, - num_warps=4, - ) + with torch.cuda.device(board_words.device): + _decode_board_kernel[grid]( + board_words, + out, + total_cells, + board_words.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + BLOCK=256, + num_warps=4, + ) return out diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index f914a67..22a79cb 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -26,20 +26,18 @@ class TrainConfig: rounds: int = 1 samples_per_round: int = 1_048_576 train_steps_per_round: int = 128 - batch_size: int = 8_192 + batch_size: int = 24_576 replay_capacity: int = 16_000_000 num_threads: int = 32 workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 1_536 + max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 - width: int = 20 - num_blocks: int = 1 - hidden_dim: int = 64 pretrain_terminal_samples: int = 5_000_000 - pretrain_batch_size: int = 8_192 + pretrain_batch_size: int = 24_576 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" + num_gpus: int = 0 device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" @@ -79,6 +77,10 @@ def _to_channels_last(module: torch.nn.Module) -> torch.nn.Module: return module +def _parameter_count(module: torch.nn.Module) -> int: + return sum(parameter.numel() for parameter in module.parameters()) + + def _numpy_batch_to_device( obs_np: np.ndarray, values_np: np.ndarray, @@ -307,8 +309,19 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: model = cast( PackedValueModel, - _to_channels_last(PackedValueModel().to(device)), + _to_channels_last( + PackedValueModel( + width=config.width, + num_blocks=config.num_blocks, + hidden_dim=config.hidden_dim, + ).to(device) + ), ) + model_params = _parameter_count(model) + print( + f"model width={config.width} blocks={config.num_blocks} hidden={config.hidden_dim} params={model_params}" + ) + wandb.config.update({"model_params": model_params}, allow_val_change=True) model.eval() optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) @@ -372,6 +385,14 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: ) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) + available_gpus = torch.cuda.device_count() + if available_gpus < 1: + raise RuntimeError("AlphaPaint self-play requires at least one CUDA GPU") + num_gpus = config.num_gpus if config.num_gpus > 0 else available_gpus + if num_gpus > available_gpus: + raise ValueError( + f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" + ) selfplay = SelfPlay( replay_buffer, config.num_threads, @@ -380,6 +401,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: max_gpu_evals_per_move=config.max_gpu_evals_per_move, model=model, selfplay_precision=config.selfplay_precision, + num_gpus=num_gpus, ) losses: list[float] = [] @@ -795,6 +817,9 @@ def _parse_args() -> TrainConfig: default=defaults.max_gpu_evals_per_move, ) parser.add_argument("--lr", type=float, default=defaults.lr) + parser.add_argument("--width", type=int, default=defaults.width) + parser.add_argument("--num-blocks", type=int, default=defaults.num_blocks) + parser.add_argument("--hidden-dim", type=int, default=defaults.hidden_dim) parser.add_argument( "--pretrain-terminal-samples", type=int, @@ -816,6 +841,12 @@ def _parse_args() -> TrainConfig: default=defaults.selfplay_precision, choices=["bf16", "fp16", "fp32"], ) + parser.add_argument( + "--num-gpus", + type=int, + default=defaults.num_gpus, + help="Number of GPUs for self-play (0 uses all visible CUDA GPUs)", + ) parser.add_argument("--device", default=defaults.device) parser.add_argument( "--checkpoint-interval", type=int, default=defaults.checkpoint_interval diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index 730f6fe..edbc4f5 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -178,6 +178,34 @@ fn check_cuda_or_panic(code: CudaError, context: &str) { } } +fn validate_cuda_device(gpu_id: usize) -> PyResult { + let gpu_device_id = i32::try_from(gpu_id) + .map_err(|_| PyErr::new::("gpu_id must fit in i32"))?; + + let mut device_count = 0i32; + unsafe { + check_cuda( + cuda::cudaGetDeviceCount(&mut device_count as *mut i32), + "cudaGetDeviceCount", + )?; + } + + if device_count <= 0 { + return Err(PyErr::new::( + "no CUDA devices available for AlphaPaint CUDA graph runner", + )); + } + + if gpu_device_id >= device_count { + return Err(PyErr::new::(format!( + "gpu_id {} out of range for {} CUDA devices", + gpu_id, device_count + ))); + } + + Ok(gpu_device_id) +} + fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { let mut ptr: *mut c_void = std::ptr::null_mut(); let bytes = count @@ -235,6 +263,7 @@ fn cuda_malloc_device_u16(count: usize, context: &str) -> PyResult<*mut c_void> } struct CudaGraphLane { + gpu_device_id: i32, stream: cudaStream_t, graph_exec: cudaGraphExec_t, /// Owns Python-side graph/tensor objects for this lane. @@ -271,6 +300,7 @@ unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { impl Drop for CudaGraphLane { fn drop(&mut self) { unsafe { + let _ = cuda::cudaSetDevice(self.gpu_device_id); let _ = cuda::cudaFree(self.obs_dev); let _ = cuda::cudaFree(self.value_dev); let _ = cuda::cudaFreeHost(self.obs_host.cast::()); @@ -282,6 +312,7 @@ impl Drop for CudaGraphLane { /// Per-lane CUDA graph executor for AlphaPaint value-only inference. pub struct CudaGraphRunner { + gpu_device_id: i32, batch_size: usize, lanes: Vec, dispatched_batches: AtomicU64, @@ -297,6 +328,7 @@ impl CudaGraphRunner { pub fn new( py: Python<'_>, model: Py, + gpu_id: usize, num_lanes: usize, batch_size: usize, precision: &str, @@ -308,6 +340,14 @@ impl CudaGraphRunner { return Err(PyErr::new::("batch_size must be > 0")); } + let gpu_device_id = validate_cuda_device(gpu_id)?; + unsafe { + check_cuda( + cuda::cudaSetDevice(gpu_device_id), + "cudaSetDevice in CudaGraphRunner::new", + )?; + } + let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; let capture_fn = module.getattr("capture_lane_graph")?; @@ -356,7 +396,7 @@ impl CudaGraphRunner { obs_dev, &obs_shape, DL_DEVICE_CUDA, - 0, + gpu_device_id, DL_DTYPE_UINT, 16, )?; @@ -374,7 +414,7 @@ impl CudaGraphRunner { value_dev, &value_shape, DL_DEVICE_CUDA, - 0, + gpu_device_id, DL_DTYPE_FLOAT, 32, )?; @@ -388,10 +428,12 @@ impl CudaGraphRunner { value_dev_capsule, stream as u64, precision, + gpu_id, ))? .extract()?; let lane = CudaGraphLane { + gpu_device_id, stream, graph_exec: exec_handle as cudaGraphExec_t, _py_owner: py_owner, @@ -404,6 +446,7 @@ impl CudaGraphRunner { } Ok(Self { + gpu_device_id, batch_size, lanes, dispatched_batches: AtomicU64::new(0), @@ -440,6 +483,10 @@ impl CudaGraphRunner { .fetch_add(self.batch_size as u64, Ordering::Relaxed); unsafe { + check_cuda_or_panic( + cuda::cudaSetDevice(self.gpu_device_id), + "cudaSetDevice before cudaGraphLaunch", + ); check_cuda_or_panic( cuda::cudaGraphLaunch(lane.graph_exec, lane.stream), "cudaGraphLaunch", diff --git a/training/src/executor.rs b/training/src/executor.rs index 5f29537..6af96c7 100644 --- a/training/src/executor.rs +++ b/training/src/executor.rs @@ -324,6 +324,34 @@ mod tests { assert!(completed.get()); } + #[test] + fn test_executor_checks_cancel_during_progress_loop() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + let polls = Arc::new(AtomicUsize::new(0)); + let polls_for_future = polls.clone(); + + let fut = std::future::poll_fn(move |_cx| { + let n = polls_for_future.fetch_add(1, Ordering::Relaxed) + 1; + if n < 1000 { + signal_progress(); + } + Poll::<()>::Pending + }); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + let polls_for_cancel = polls.clone(); + executor.run(&mut vec![Box::pin(fut)], &mut || { + polls_for_cancel.load(Ordering::Relaxed) >= 10 + }); + + assert!( + polls.load(Ordering::Relaxed) < 100, + "cancel should stop polling quickly" + ); + } + #[test] fn test_run_preserves_futures() { // Test that run does NOT drop pending futures. diff --git a/training/src/lib.rs b/training/src/lib.rs index 2808d15..df66e31 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,5 +1,6 @@ use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use alpha_paint::TRAINING_START_FENS; +use std::collections::HashMap; use std::sync::Arc; use std::sync::{Mutex, OnceLock}; @@ -86,7 +87,9 @@ struct GraphCacheEntry { model_ptr: usize, num_batches: usize, precision: String, - runner: Arc, + num_gpus: usize, + runners: HashMap>, + models: HashMap>, } static GRAPH_CACHE: OnceLock>> = OnceLock::new(); @@ -253,7 +256,9 @@ impl EphemeralReplayBuffer { #[pyclass] struct SelfPlay { session: Option, - runner: Arc, + runners: HashMap>, + source_model: Py, + source_model_ptr: usize, } impl SelfPlay { @@ -275,7 +280,8 @@ impl SelfPlay { *, max_gpu_evals_per_move = 4096, model, - selfplay_precision = "bf16" + selfplay_precision = "bf16", + num_gpus = 1 ))] fn new( py: Python<'_>, @@ -286,7 +292,14 @@ impl SelfPlay { max_gpu_evals_per_move: u64, model: Py, selfplay_precision: &str, + num_gpus: usize, ) -> PyResult { + if num_gpus == 0 { + return Err(PyErr::new::( + "num_gpus must be >= 1", + )); + } + let config = SessionConfig { num_threads, workers_per_thread, @@ -304,8 +317,8 @@ impl SelfPlay { let (num_batches, _total_slots) = queue_shape_for_workers(total_workers); let model_ptr = model.bind(py).as_ptr() as usize; - // Build or reuse the CUDA graph runner. - let runner = { + // Build or reuse the CUDA graph runners. + let runners = { let cache = graph_cache(); let mut guard = cache.lock().expect("graph cache mutex poisoned"); @@ -314,57 +327,75 @@ impl SelfPlay { entry.model_ptr != model_ptr || entry.num_batches != num_batches || entry.precision != selfplay_precision + || entry.num_gpus != num_gpus } None => true, }; if needs_rebuild { - let runner = Arc::new(CudaGraphRunner::new( - py, - model.clone_ref(py), - num_batches, - BATCH_SIZE, - selfplay_precision, - )?); + let copy = PyModule::import(py, "copy")?; + let deepcopy = copy.getattr("deepcopy")?; + let mut runners = HashMap::new(); + let mut models = HashMap::new(); + for gpu_id in 0..num_gpus { + let model_copy: Py = deepcopy.call1((model.clone_ref(py),))?.into(); + let runner = Arc::new(CudaGraphRunner::new( + py, + model_copy.clone_ref(py), + gpu_id, + num_batches, + BATCH_SIZE, + selfplay_precision, + )?); + runners.insert(gpu_id, runner); + models.insert(gpu_id, model_copy); + } *guard = Some(GraphCacheEntry { model_ptr, num_batches, precision: selfplay_precision.to_string(), - runner: runner.clone(), + num_gpus, + runners: runners.clone(), + models, }); - runner + runners } else { guard .as_ref() - .expect("cached runner should exist") - .runner + .expect("cached runners should exist") + .runners .clone() } }; - let runner_for_dispatch = runner.clone(); + let runners_for_dispatch = runners.clone(); let dispatch = move |batch_idx: usize, obs_view: ArrayView, completion: queue::BatchCompletion| { - runner_for_dispatch.dispatch_async(batch_idx, obs_view, completion); + let gpu_id = batch_idx % num_gpus; + runners_for_dispatch[&gpu_id].dispatch_async(batch_idx, obs_view, completion); }; let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); Ok(Self { session: Some(session), - runner, + runners, + source_model: model, + source_model_ptr: model_ptr, }) } /// Start self-play with no sample limit. - fn start(&self) -> PyResult<()> { + fn start(&self, py: Python<'_>) -> PyResult<()> { + self.sync_model_replicas(py)?; self.session()?.start(); Ok(()) } /// Block until absolute target_samples is reached, then pause and quiesce. fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { + self.sync_model_replicas(py)?; let session = self.session()?; let result = py.detach(|| session.wait_for(target_samples)); Ok(result) @@ -482,12 +513,18 @@ impl SelfPlay { /// Return the total number of CUDA graph launches completed so far. fn gpu_batches(&self) -> u64 { - self.runner.dispatched_batches() + self.runners + .values() + .map(|runner| runner.dispatched_batches()) + .sum() } /// Return the total number of packed observations sent to GPU so far. fn gpu_evals(&self) -> u64 { - self.runner.dispatched_evals() + self.runners + .values() + .map(|runner| runner.dispatched_evals()) + .sum() } /// Shut down the session. Idempotent. @@ -499,6 +536,50 @@ impl SelfPlay { } } +impl SelfPlay { + fn sync_model_replicas(&self, py: Python<'_>) -> PyResult<()> { + let replicas = { + let cache = graph_cache(); + let guard = cache.lock().expect("graph cache mutex poisoned"); + let entry = guard.as_ref().ok_or_else(|| { + PyErr::new::( + "graph cache missing while syncing model replicas", + ) + })?; + + if entry.model_ptr != self.source_model_ptr { + return Err(PyErr::new::( + "graph cache model mismatch while syncing model replicas", + )); + } + + let mut models: Vec<(usize, Py)> = entry + .models + .iter() + .map(|(gpu_id, model)| (*gpu_id, model.clone_ref(py))) + .collect(); + models.sort_by_key(|(gpu_id, _)| *gpu_id); + models + .into_iter() + .map(|(_, model)| model) + .collect::>() + }; + + let state_dict: Py = self + .source_model + .bind(py) + .call_method0("state_dict")? + .into(); + for replica in replicas { + let _ = replica + .bind(py) + .call_method1("load_state_dict", (state_dict.clone_ref(py),))?; + } + + Ok(()) + } +} + impl Drop for SelfPlay { fn drop(&mut self) { if let Some(mut session) = self.session.take() { From 6b3d24a3eb76f14cbdeb18c23b9ea56d8cf6410d Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 18:12:34 -0400 Subject: [PATCH 38/59] config for a multigpu --- .../alphapaint_training.pyi | 1 + python/alphapaint_training/train.py | 65 +++++++++++++++++-- 2 files changed, 61 insertions(+), 5 deletions(-) diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi index 94213d7..0a8cd21 100644 --- a/python/alphapaint_training/alphapaint_training.pyi +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -23,6 +23,7 @@ class SelfPlay: max_gpu_evals_per_move: int = 4096, model: object, selfplay_precision: str = "bf16", + num_gpus: int = 1, ) -> None: ... def start(self) -> None: ... def wait_for(self, target_samples: int) -> int: ... diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 22a79cb..41bc580 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -2,9 +2,10 @@ import argparse import json +import os import time from contextlib import nullcontext -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path from typing import cast @@ -21,6 +22,47 @@ from alphapaint_training.model import PackedValueModel +def _int_env(name: str) -> int | None: + value = os.environ.get(name) + if not value: + return None + try: + return int(value) + except ValueError: + return None + + +def _slurm_gpu_count() -> int | None: + for name in ("SLURM_GPUS_ON_NODE", "SLURM_GPUS"): + value = _int_env(name) + if value is not None and value > 0: + return value + + job_gpus = os.environ.get("SLURM_JOB_GPUS") + if not job_gpus: + return None + + gpu_ids = [gpu_id.strip() for gpu_id in job_gpus.split(",") if gpu_id.strip()] + if not gpu_ids: + return None + return len(gpu_ids) + + +def _default_num_threads() -> int: + cpus_per_gpu = _int_env("SLURM_CPUS_PER_GPU") + slurm_gpu_count = _slurm_gpu_count() + if cpus_per_gpu is not None and cpus_per_gpu > 0 and slurm_gpu_count is not None: + return cpus_per_gpu * slurm_gpu_count + return 32 + + +def _default_num_gpus() -> int: + slurm_gpu_count = _slurm_gpu_count() + if slurm_gpu_count is None: + return 0 + return slurm_gpu_count + + @dataclass(slots=True) class TrainConfig: rounds: int = 1 @@ -28,16 +70,19 @@ class TrainConfig: train_steps_per_round: int = 128 batch_size: int = 24_576 replay_capacity: int = 16_000_000 - num_threads: int = 32 + num_threads: int = field(default_factory=_default_num_threads) workers_per_thread: int = 8 max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 + width: int = 24 + num_blocks: int = 2 + hidden_dim: int = 128 pretrain_terminal_samples: int = 5_000_000 pretrain_batch_size: int = 24_576 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" - num_gpus: int = 0 + num_gpus: int = field(default_factory=_default_num_gpus) device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" @@ -393,6 +438,11 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: raise ValueError( f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" ) + print( + "selfplay " + f"threads={config.num_threads} workers_per_thread={config.workers_per_thread} " + f"gpus={num_gpus} max_gpu_evals_per_move={config.max_gpu_evals_per_move}" + ) selfplay = SelfPlay( replay_buffer, config.num_threads, @@ -807,7 +857,12 @@ def _parse_args() -> TrainConfig: ) parser.add_argument("--batch-size", type=int, default=defaults.batch_size) parser.add_argument("--replay-capacity", type=int, default=defaults.replay_capacity) - parser.add_argument("--num-threads", type=int, default=defaults.num_threads) + parser.add_argument( + "--num-threads", + type=int, + default=defaults.num_threads, + help="Self-play executor threads (defaults to Slurm CPU allocation when available)", + ) parser.add_argument( "--workers-per-thread", type=int, default=defaults.workers_per_thread ) @@ -845,7 +900,7 @@ def _parse_args() -> TrainConfig: "--num-gpus", type=int, default=defaults.num_gpus, - help="Number of GPUs for self-play (0 uses all visible CUDA GPUs)", + help="Number of GPUs for self-play (0 uses all visible CUDA GPUs; defaults to Slurm allocation when available)", ) parser.add_argument("--device", default=defaults.device) parser.add_argument( From 9936dbd8332c779aa764991d3b13c3c1df7de711 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 19:29:45 -0400 Subject: [PATCH 39/59] Revert "config for a multigpu" This reverts commit 6b3d24a3eb76f14cbdeb18c23b9ea56d8cf6410d. --- .../alphapaint_training.pyi | 1 - python/alphapaint_training/train.py | 65 ++----------------- 2 files changed, 5 insertions(+), 61 deletions(-) diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi index 0a8cd21..94213d7 100644 --- a/python/alphapaint_training/alphapaint_training.pyi +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -23,7 +23,6 @@ class SelfPlay: max_gpu_evals_per_move: int = 4096, model: object, selfplay_precision: str = "bf16", - num_gpus: int = 1, ) -> None: ... def start(self) -> None: ... def wait_for(self, target_samples: int) -> int: ... diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 41bc580..22a79cb 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -2,10 +2,9 @@ import argparse import json -import os import time from contextlib import nullcontext -from dataclasses import asdict, dataclass, field +from dataclasses import asdict, dataclass from pathlib import Path from typing import cast @@ -22,47 +21,6 @@ from alphapaint_training.model import PackedValueModel -def _int_env(name: str) -> int | None: - value = os.environ.get(name) - if not value: - return None - try: - return int(value) - except ValueError: - return None - - -def _slurm_gpu_count() -> int | None: - for name in ("SLURM_GPUS_ON_NODE", "SLURM_GPUS"): - value = _int_env(name) - if value is not None and value > 0: - return value - - job_gpus = os.environ.get("SLURM_JOB_GPUS") - if not job_gpus: - return None - - gpu_ids = [gpu_id.strip() for gpu_id in job_gpus.split(",") if gpu_id.strip()] - if not gpu_ids: - return None - return len(gpu_ids) - - -def _default_num_threads() -> int: - cpus_per_gpu = _int_env("SLURM_CPUS_PER_GPU") - slurm_gpu_count = _slurm_gpu_count() - if cpus_per_gpu is not None and cpus_per_gpu > 0 and slurm_gpu_count is not None: - return cpus_per_gpu * slurm_gpu_count - return 32 - - -def _default_num_gpus() -> int: - slurm_gpu_count = _slurm_gpu_count() - if slurm_gpu_count is None: - return 0 - return slurm_gpu_count - - @dataclass(slots=True) class TrainConfig: rounds: int = 1 @@ -70,19 +28,16 @@ class TrainConfig: train_steps_per_round: int = 128 batch_size: int = 24_576 replay_capacity: int = 16_000_000 - num_threads: int = field(default_factory=_default_num_threads) + num_threads: int = 32 workers_per_thread: int = 8 max_gpu_evals_per_move: int = 4 * 1024 lr: float = 3e-4 - width: int = 24 - num_blocks: int = 2 - hidden_dim: int = 128 pretrain_terminal_samples: int = 5_000_000 pretrain_batch_size: int = 24_576 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" - num_gpus: int = field(default_factory=_default_num_gpus) + num_gpus: int = 0 device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" @@ -438,11 +393,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: raise ValueError( f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" ) - print( - "selfplay " - f"threads={config.num_threads} workers_per_thread={config.workers_per_thread} " - f"gpus={num_gpus} max_gpu_evals_per_move={config.max_gpu_evals_per_move}" - ) selfplay = SelfPlay( replay_buffer, config.num_threads, @@ -857,12 +807,7 @@ def _parse_args() -> TrainConfig: ) parser.add_argument("--batch-size", type=int, default=defaults.batch_size) parser.add_argument("--replay-capacity", type=int, default=defaults.replay_capacity) - parser.add_argument( - "--num-threads", - type=int, - default=defaults.num_threads, - help="Self-play executor threads (defaults to Slurm CPU allocation when available)", - ) + parser.add_argument("--num-threads", type=int, default=defaults.num_threads) parser.add_argument( "--workers-per-thread", type=int, default=defaults.workers_per_thread ) @@ -900,7 +845,7 @@ def _parse_args() -> TrainConfig: "--num-gpus", type=int, default=defaults.num_gpus, - help="Number of GPUs for self-play (0 uses all visible CUDA GPUs; defaults to Slurm allocation when available)", + help="Number of GPUs for self-play (0 uses all visible CUDA GPUs)", ) parser.add_argument("--device", default=defaults.device) parser.add_argument( From 08ef34f08b3d11879d1f48c9ccefe84d8fd8b36e Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 19:29:45 -0400 Subject: [PATCH 40/59] Revert "wip: multigpu" This reverts commit 3d163dd047f535b606b38890296bf809ecbb3527. --- alpha_paint/benches/nn_latency.rs | 2 +- alpha_paint/src/nn/features.rs | 28 ++-- alpha_paint/src/nn/model.rs | 19 +-- alpha_paint/src/search.rs | 28 ++-- .../alphapaint_training/cudagraph_backend.py | 128 ++++++------------ python/alphapaint_training/packed_obs.py | 25 ++-- python/alphapaint_training/train.py | 45 +----- training/src/cudagraph.rs | 51 +------ training/src/executor.rs | 28 ---- training/src/lib.rs | 125 +++-------------- 10 files changed, 112 insertions(+), 367 deletions(-) diff --git a/alpha_paint/benches/nn_latency.rs b/alpha_paint/benches/nn_latency.rs index 2ecd24f..e3d6297 100644 --- a/alpha_paint/benches/nn_latency.rs +++ b/alpha_paint/benches/nn_latency.rs @@ -1,5 +1,5 @@ -use std::hint::black_box; use std::time::Duration; +use std::hint::black_box; use alpha_paint::board::Board; use alpha_paint::evaluation::Evaluator; diff --git a/alpha_paint/src/nn/features.rs b/alpha_paint/src/nn/features.rs index 1a09a65..2f73455 100644 --- a/alpha_paint/src/nn/features.rs +++ b/alpha_paint/src/nn/features.rs @@ -1,6 +1,6 @@ -use crate::board::Board; use crate::board::board_structs::Player; use crate::board::structs::Coordinate; +use crate::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; @@ -168,7 +168,8 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; - out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = + board.tiles.get_beacon_iterator::().count() as f32; } else { out[INTRINSIC_CURRENT_STAMINA] = board.black_stamina as f32; out[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina as f32; @@ -176,7 +177,8 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; - out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_CURRENT_BEACONS] = + board.tiles.get_beacon_iterator::().count() as f32; out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; } @@ -219,18 +221,9 @@ mod tests { assert_eq!(features.board[plane_index(WALL, 31, 31)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 99.0 / 420.0); - assert_eq!( - features.intrinsics[INTRINSIC_OPPONENT_STAMINA], - 88.0 / 420.0 - ); - assert_eq!( - features.intrinsics[INTRINSIC_CURRENT_TERRITORY], - 1.0 / 1024.0 - ); - assert_eq!( - features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], - 1.0 / 1024.0 - ); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 88.0 / 420.0); + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_TERRITORY], 1.0 / 1024.0); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], 1.0 / 1024.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_BEACONS], 1.0 / 1024.0); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 2.0 / 64.0); } @@ -251,10 +244,7 @@ mod tests { assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 0, 0)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 88.0 / 420.0); - assert_eq!( - features.intrinsics[INTRINSIC_OPPONENT_STAMINA], - 99.0 / 420.0 - ); + assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 99.0 / 420.0); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 5.0 / 64.0); } } diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs index 823da96..bb28625 100644 --- a/alpha_paint/src/nn/model.rs +++ b/alpha_paint/src/nn/model.rs @@ -163,8 +163,7 @@ impl ValueModel { "head.fc1.weight", &[MODEL_HIDDEN_DIM, HEAD0_IN_DIM], )?; - let head0_bias = - tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; + let head0_bias = tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; let head1_weight = tensor_array::(tensors, "head.fc2.weight", &[1, MODEL_HIDDEN_DIM])?; let head1_bias = tensor_array::<1>(tensors, "head.fc2.bias", &[1])?[0]; @@ -243,8 +242,12 @@ impl ValueModel { &scratch.head_input, &mut scratch.hidden, ); - (linear_scalar::(&self.head1_weight, self.head1_bias, &scratch.hidden)) - .tanh() + (linear_scalar::( + &self.head1_weight, + self.head1_bias, + &scratch.hidden, + )) + .tanh() * 7.0 } } @@ -369,7 +372,8 @@ fn accumulate_conv3x3( for y in 1..BOARD_SIDE - 1 { let row = y * BOARD_SIDE; out_plane[row] += conv3x3_border(in_plane, kernel, 0, y); - out_plane[row + BOARD_SIDE - 1] += conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); + out_plane[row + BOARD_SIDE - 1] += + conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); } } } @@ -522,10 +526,7 @@ fn tensor_f32( } let data = tensor.data(); if data.len() % 4 != 0 { - return Err(format!( - "tensor {name} has invalid byte length {}", - data.len() - )); + return Err(format!("tensor {name} has invalid byte length {}", data.len())); } let mut out = Vec::with_capacity(data.len() / 4); for chunk in data.chunks_exact(4) { diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index 9afe8a1..e4015e4 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -131,13 +131,9 @@ impl SearchNode { if self.completion_value.abs() == 1 { true } else { - self.children.iter().all(|child| { - child - .node - .as_ref() - .map(|node| node.resolved) - .unwrap_or(false) - }) + self.children + .iter() + .all(|child| child.node.as_ref().map(|node| node.resolved).unwrap_or(false)) } } @@ -258,11 +254,7 @@ impl SearchNode { let (outcome, _) = state.apply_action(action); let node = Box::new(SearchNode::build_self(&state, outcome, evaluator, rng)); let value = node.value; - if let Some(index) = self - .children - .iter() - .position(|child| child.action == action) - { + if let Some(index) = self.children.iter().position(|child| child.action == action) { self.children[index].node = Some(node); } value @@ -332,7 +324,11 @@ impl GameSearchTree<'_> { .iter() .enumerate() .max_by(|&(_, a), &(_, b)| { - (a.completion_value(), a.entrance_count as i32, a.child_value) + ( + a.completion_value(), + a.entrance_count as i32, + a.child_value, + ) .partial_cmp(&( b.completion_value(), b.entrance_count as i32, @@ -396,11 +392,7 @@ impl GameSearchTree<'_> { } pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let desired_completion = if self.root_state.is_white_turn() { - 1 - } else { - -1 - }; + let desired_completion = if self.root_state.is_white_turn() { 1 } else { -1 }; let started_at = std::time::Instant::now(); if self.root_node.children.len() <= 1 { return; diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index fffdc1c..dd62b0d 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -10,48 +10,6 @@ import torch.utils.dlpack as dlpack -def _validate_lane_tensors( - obs_host: torch.Tensor, - obs_device: torch.Tensor, - value_host: torch.Tensor, - value_device: torch.Tensor, - gpu_id: int, -) -> None: - if obs_host.device.type != "cpu": - raise ValueError("obs_host must be a CPU tensor") - if value_host.device.type != "cpu": - raise ValueError("value_host must be a CPU tensor") - if obs_device.device.type != "cuda": - raise ValueError("obs_device must be a CUDA tensor") - if value_device.device.type != "cuda": - raise ValueError("value_device must be a CUDA tensor") - - if obs_host.dtype != torch.uint16: - raise ValueError(f"obs_host must be uint16, got {obs_host.dtype}") - if obs_device.dtype != torch.uint16: - raise ValueError(f"obs_device must be uint16, got {obs_device.dtype}") - if value_host.dtype != torch.float32: - raise ValueError(f"value_host must be float32, got {value_host.dtype}") - if value_device.dtype != torch.float32: - raise ValueError(f"value_device must be float32, got {value_device.dtype}") - - if obs_host.shape != obs_device.shape: - raise ValueError( - f"obs_host and obs_device shape mismatch: {obs_host.shape} vs {obs_device.shape}" - ) - batch = obs_host.shape[0] - if value_host.shape != (batch,): - raise ValueError(f"value_host must be shape ({batch},), got {value_host.shape}") - if value_device.shape != (batch,): - raise ValueError( - f"value_device must be shape ({batch},), got {value_device.shape}" - ) - - for name, tensor in (("obs_device", obs_device), ("value_device", value_device)): - if tensor.device.index != gpu_id: - raise ValueError(f"{name} must be on cuda:{gpu_id}, got {tensor.device}") - - def _autocast_dtype(precision: str) -> Optional[torch.dtype]: if precision == "fp32": return None @@ -72,7 +30,6 @@ def capture_lane_graph( value_device_dlpack, stream_handle: int, precision: str = "bf16", - gpu_id: int = 0, ) -> tuple[int, object]: """Capture a CUDA graph for value-only inference. @@ -91,56 +48,49 @@ def capture_lane_graph( Returns: Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). """ - if gpu_id < 0: - raise ValueError(f"gpu_id must be >= 0, got {gpu_id}") - obs_host = dlpack.from_dlpack(obs_host_dlpack) obs_device = dlpack.from_dlpack(obs_device_dlpack) value_host = dlpack.from_dlpack(value_host_dlpack) value_device = dlpack.from_dlpack(value_device_dlpack) - _validate_lane_tensors(obs_host, obs_device, value_host, value_device, gpu_id) - - with torch.cuda.device(gpu_id): - model = model.to(f"cuda:{gpu_id}") - model = model.to(memory_format=torch.channels_last) - model.eval() - torch.backends.cudnn.benchmark = True - stream = torch.cuda.ExternalStream(stream_handle) - graph = torch.cuda.CUDAGraph(keep_graph=True) - dtype = _autocast_dtype(precision) - - def run_step() -> None: - obs_device.copy_(obs_host, non_blocking=True) - if dtype is None: + model = model.cuda() + model = model.to(memory_format=torch.channels_last) + model.eval() + torch.backends.cudnn.benchmark = True + stream = torch.cuda.ExternalStream(stream_handle) + graph = torch.cuda.CUDAGraph(keep_graph=True) + dtype = _autocast_dtype(precision) + + def run_step() -> None: + obs_device.copy_(obs_host, non_blocking=True) + if dtype is None: + value = model(obs_device) + else: + with torch.autocast(device_type="cuda", dtype=dtype): value = model(obs_device) - else: - with torch.autocast(device_type="cuda", dtype=dtype): - value = model(obs_device) - if value.ndim == 2: - value = value.squeeze(-1) - value_device.copy_(value, non_blocking=True) - value_host.copy_(value_device, non_blocking=True) - - with torch.inference_mode(): - with torch.cuda.stream(stream): - for _ in range(3): - run_step() - torch.cuda.synchronize(device=gpu_id) - - with torch.cuda.graph( - graph, stream=stream, capture_error_mode="thread_local" - ): + # Model returns value tensor, shape (B,) or (B, 1) + if value.ndim == 2: + value = value.squeeze(-1) + value_device.copy_(value, non_blocking=True) + value_host.copy_(value_device, non_blocking=True) + + with torch.inference_mode(): + with torch.cuda.stream(stream): + for _ in range(3): run_step() - - graph.instantiate() - owner = ( - graph, - model, - obs_host, - obs_device, - value_host, - value_device, - stream, - ) - return int(graph.raw_cuda_graph_exec()), owner + torch.cuda.synchronize() + + with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"): + run_step() + + graph.instantiate() + owner = ( + graph, + model, + obs_host, + obs_device, + value_host, + value_device, + stream, + ) + return int(graph.raw_cuda_graph_exec()), owner diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index 50f85fe..a7d9003 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -229,19 +229,18 @@ def decode_packed_board( raise ValueError(f"out must have dtype {dtype}, got {out.dtype}") total_cells = board_words.shape[0] * BOARD_CELLS grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) - with torch.cuda.device(board_words.device): - _decode_board_kernel[grid]( - board_words, - out, - total_cells, - board_words.stride(0), - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - BLOCK=256, - num_warps=4, - ) + _decode_board_kernel[grid]( + board_words, + out, + total_cells, + board_words.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + BLOCK=256, + num_warps=4, + ) return out diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 22a79cb..f914a67 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -26,18 +26,20 @@ class TrainConfig: rounds: int = 1 samples_per_round: int = 1_048_576 train_steps_per_round: int = 128 - batch_size: int = 24_576 + batch_size: int = 8_192 replay_capacity: int = 16_000_000 num_threads: int = 32 workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 4 * 1024 + max_gpu_evals_per_move: int = 1_536 lr: float = 3e-4 + width: int = 20 + num_blocks: int = 1 + hidden_dim: int = 64 pretrain_terminal_samples: int = 5_000_000 - pretrain_batch_size: int = 24_576 + pretrain_batch_size: int = 8_192 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" - num_gpus: int = 0 device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" @@ -77,10 +79,6 @@ def _to_channels_last(module: torch.nn.Module) -> torch.nn.Module: return module -def _parameter_count(module: torch.nn.Module) -> int: - return sum(parameter.numel() for parameter in module.parameters()) - - def _numpy_batch_to_device( obs_np: np.ndarray, values_np: np.ndarray, @@ -309,19 +307,8 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: model = cast( PackedValueModel, - _to_channels_last( - PackedValueModel( - width=config.width, - num_blocks=config.num_blocks, - hidden_dim=config.hidden_dim, - ).to(device) - ), + _to_channels_last(PackedValueModel().to(device)), ) - model_params = _parameter_count(model) - print( - f"model width={config.width} blocks={config.num_blocks} hidden={config.hidden_dim} params={model_params}" - ) - wandb.config.update({"model_params": model_params}, allow_val_change=True) model.eval() optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) @@ -385,14 +372,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: ) replay_buffer = EphemeralReplayBuffer(config.replay_capacity) - available_gpus = torch.cuda.device_count() - if available_gpus < 1: - raise RuntimeError("AlphaPaint self-play requires at least one CUDA GPU") - num_gpus = config.num_gpus if config.num_gpus > 0 else available_gpus - if num_gpus > available_gpus: - raise ValueError( - f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" - ) selfplay = SelfPlay( replay_buffer, config.num_threads, @@ -401,7 +380,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: max_gpu_evals_per_move=config.max_gpu_evals_per_move, model=model, selfplay_precision=config.selfplay_precision, - num_gpus=num_gpus, ) losses: list[float] = [] @@ -817,9 +795,6 @@ def _parse_args() -> TrainConfig: default=defaults.max_gpu_evals_per_move, ) parser.add_argument("--lr", type=float, default=defaults.lr) - parser.add_argument("--width", type=int, default=defaults.width) - parser.add_argument("--num-blocks", type=int, default=defaults.num_blocks) - parser.add_argument("--hidden-dim", type=int, default=defaults.hidden_dim) parser.add_argument( "--pretrain-terminal-samples", type=int, @@ -841,12 +816,6 @@ def _parse_args() -> TrainConfig: default=defaults.selfplay_precision, choices=["bf16", "fp16", "fp32"], ) - parser.add_argument( - "--num-gpus", - type=int, - default=defaults.num_gpus, - help="Number of GPUs for self-play (0 uses all visible CUDA GPUs)", - ) parser.add_argument("--device", default=defaults.device) parser.add_argument( "--checkpoint-interval", type=int, default=defaults.checkpoint_interval diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index edbc4f5..730f6fe 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -178,34 +178,6 @@ fn check_cuda_or_panic(code: CudaError, context: &str) { } } -fn validate_cuda_device(gpu_id: usize) -> PyResult { - let gpu_device_id = i32::try_from(gpu_id) - .map_err(|_| PyErr::new::("gpu_id must fit in i32"))?; - - let mut device_count = 0i32; - unsafe { - check_cuda( - cuda::cudaGetDeviceCount(&mut device_count as *mut i32), - "cudaGetDeviceCount", - )?; - } - - if device_count <= 0 { - return Err(PyErr::new::( - "no CUDA devices available for AlphaPaint CUDA graph runner", - )); - } - - if gpu_device_id >= device_count { - return Err(PyErr::new::(format!( - "gpu_id {} out of range for {} CUDA devices", - gpu_id, device_count - ))); - } - - Ok(gpu_device_id) -} - fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { let mut ptr: *mut c_void = std::ptr::null_mut(); let bytes = count @@ -263,7 +235,6 @@ fn cuda_malloc_device_u16(count: usize, context: &str) -> PyResult<*mut c_void> } struct CudaGraphLane { - gpu_device_id: i32, stream: cudaStream_t, graph_exec: cudaGraphExec_t, /// Owns Python-side graph/tensor objects for this lane. @@ -300,7 +271,6 @@ unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { impl Drop for CudaGraphLane { fn drop(&mut self) { unsafe { - let _ = cuda::cudaSetDevice(self.gpu_device_id); let _ = cuda::cudaFree(self.obs_dev); let _ = cuda::cudaFree(self.value_dev); let _ = cuda::cudaFreeHost(self.obs_host.cast::()); @@ -312,7 +282,6 @@ impl Drop for CudaGraphLane { /// Per-lane CUDA graph executor for AlphaPaint value-only inference. pub struct CudaGraphRunner { - gpu_device_id: i32, batch_size: usize, lanes: Vec, dispatched_batches: AtomicU64, @@ -328,7 +297,6 @@ impl CudaGraphRunner { pub fn new( py: Python<'_>, model: Py, - gpu_id: usize, num_lanes: usize, batch_size: usize, precision: &str, @@ -340,14 +308,6 @@ impl CudaGraphRunner { return Err(PyErr::new::("batch_size must be > 0")); } - let gpu_device_id = validate_cuda_device(gpu_id)?; - unsafe { - check_cuda( - cuda::cudaSetDevice(gpu_device_id), - "cudaSetDevice in CudaGraphRunner::new", - )?; - } - let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; let capture_fn = module.getattr("capture_lane_graph")?; @@ -396,7 +356,7 @@ impl CudaGraphRunner { obs_dev, &obs_shape, DL_DEVICE_CUDA, - gpu_device_id, + 0, DL_DTYPE_UINT, 16, )?; @@ -414,7 +374,7 @@ impl CudaGraphRunner { value_dev, &value_shape, DL_DEVICE_CUDA, - gpu_device_id, + 0, DL_DTYPE_FLOAT, 32, )?; @@ -428,12 +388,10 @@ impl CudaGraphRunner { value_dev_capsule, stream as u64, precision, - gpu_id, ))? .extract()?; let lane = CudaGraphLane { - gpu_device_id, stream, graph_exec: exec_handle as cudaGraphExec_t, _py_owner: py_owner, @@ -446,7 +404,6 @@ impl CudaGraphRunner { } Ok(Self { - gpu_device_id, batch_size, lanes, dispatched_batches: AtomicU64::new(0), @@ -483,10 +440,6 @@ impl CudaGraphRunner { .fetch_add(self.batch_size as u64, Ordering::Relaxed); unsafe { - check_cuda_or_panic( - cuda::cudaSetDevice(self.gpu_device_id), - "cudaSetDevice before cudaGraphLaunch", - ); check_cuda_or_panic( cuda::cudaGraphLaunch(lane.graph_exec, lane.stream), "cudaGraphLaunch", diff --git a/training/src/executor.rs b/training/src/executor.rs index 6af96c7..5f29537 100644 --- a/training/src/executor.rs +++ b/training/src/executor.rs @@ -324,34 +324,6 @@ mod tests { assert!(completed.get()); } - #[test] - fn test_executor_checks_cancel_during_progress_loop() { - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - let polls = Arc::new(AtomicUsize::new(0)); - let polls_for_future = polls.clone(); - - let fut = std::future::poll_fn(move |_cx| { - let n = polls_for_future.fetch_add(1, Ordering::Relaxed) + 1; - if n < 1000 { - signal_progress(); - } - Poll::<()>::Pending - }); - - let executor = Executor::new(|| event_listener::Event::new().listen()); - let polls_for_cancel = polls.clone(); - executor.run(&mut vec![Box::pin(fut)], &mut || { - polls_for_cancel.load(Ordering::Relaxed) >= 10 - }); - - assert!( - polls.load(Ordering::Relaxed) < 100, - "cancel should stop polling quickly" - ); - } - #[test] fn test_run_preserves_futures() { // Test that run does NOT drop pending futures. diff --git a/training/src/lib.rs b/training/src/lib.rs index df66e31..2808d15 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,6 +1,5 @@ use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use alpha_paint::TRAINING_START_FENS; -use std::collections::HashMap; use std::sync::Arc; use std::sync::{Mutex, OnceLock}; @@ -87,9 +86,7 @@ struct GraphCacheEntry { model_ptr: usize, num_batches: usize, precision: String, - num_gpus: usize, - runners: HashMap>, - models: HashMap>, + runner: Arc, } static GRAPH_CACHE: OnceLock>> = OnceLock::new(); @@ -256,9 +253,7 @@ impl EphemeralReplayBuffer { #[pyclass] struct SelfPlay { session: Option, - runners: HashMap>, - source_model: Py, - source_model_ptr: usize, + runner: Arc, } impl SelfPlay { @@ -280,8 +275,7 @@ impl SelfPlay { *, max_gpu_evals_per_move = 4096, model, - selfplay_precision = "bf16", - num_gpus = 1 + selfplay_precision = "bf16" ))] fn new( py: Python<'_>, @@ -292,14 +286,7 @@ impl SelfPlay { max_gpu_evals_per_move: u64, model: Py, selfplay_precision: &str, - num_gpus: usize, ) -> PyResult { - if num_gpus == 0 { - return Err(PyErr::new::( - "num_gpus must be >= 1", - )); - } - let config = SessionConfig { num_threads, workers_per_thread, @@ -317,8 +304,8 @@ impl SelfPlay { let (num_batches, _total_slots) = queue_shape_for_workers(total_workers); let model_ptr = model.bind(py).as_ptr() as usize; - // Build or reuse the CUDA graph runners. - let runners = { + // Build or reuse the CUDA graph runner. + let runner = { let cache = graph_cache(); let mut guard = cache.lock().expect("graph cache mutex poisoned"); @@ -327,75 +314,57 @@ impl SelfPlay { entry.model_ptr != model_ptr || entry.num_batches != num_batches || entry.precision != selfplay_precision - || entry.num_gpus != num_gpus } None => true, }; if needs_rebuild { - let copy = PyModule::import(py, "copy")?; - let deepcopy = copy.getattr("deepcopy")?; - let mut runners = HashMap::new(); - let mut models = HashMap::new(); - for gpu_id in 0..num_gpus { - let model_copy: Py = deepcopy.call1((model.clone_ref(py),))?.into(); - let runner = Arc::new(CudaGraphRunner::new( - py, - model_copy.clone_ref(py), - gpu_id, - num_batches, - BATCH_SIZE, - selfplay_precision, - )?); - runners.insert(gpu_id, runner); - models.insert(gpu_id, model_copy); - } + let runner = Arc::new(CudaGraphRunner::new( + py, + model.clone_ref(py), + num_batches, + BATCH_SIZE, + selfplay_precision, + )?); *guard = Some(GraphCacheEntry { model_ptr, num_batches, precision: selfplay_precision.to_string(), - num_gpus, - runners: runners.clone(), - models, + runner: runner.clone(), }); - runners + runner } else { guard .as_ref() - .expect("cached runners should exist") - .runners + .expect("cached runner should exist") + .runner .clone() } }; - let runners_for_dispatch = runners.clone(); + let runner_for_dispatch = runner.clone(); let dispatch = move |batch_idx: usize, obs_view: ArrayView, completion: queue::BatchCompletion| { - let gpu_id = batch_idx % num_gpus; - runners_for_dispatch[&gpu_id].dispatch_async(batch_idx, obs_view, completion); + runner_for_dispatch.dispatch_async(batch_idx, obs_view, completion); }; let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); Ok(Self { session: Some(session), - runners, - source_model: model, - source_model_ptr: model_ptr, + runner, }) } /// Start self-play with no sample limit. - fn start(&self, py: Python<'_>) -> PyResult<()> { - self.sync_model_replicas(py)?; + fn start(&self) -> PyResult<()> { self.session()?.start(); Ok(()) } /// Block until absolute target_samples is reached, then pause and quiesce. fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { - self.sync_model_replicas(py)?; let session = self.session()?; let result = py.detach(|| session.wait_for(target_samples)); Ok(result) @@ -513,18 +482,12 @@ impl SelfPlay { /// Return the total number of CUDA graph launches completed so far. fn gpu_batches(&self) -> u64 { - self.runners - .values() - .map(|runner| runner.dispatched_batches()) - .sum() + self.runner.dispatched_batches() } /// Return the total number of packed observations sent to GPU so far. fn gpu_evals(&self) -> u64 { - self.runners - .values() - .map(|runner| runner.dispatched_evals()) - .sum() + self.runner.dispatched_evals() } /// Shut down the session. Idempotent. @@ -536,50 +499,6 @@ impl SelfPlay { } } -impl SelfPlay { - fn sync_model_replicas(&self, py: Python<'_>) -> PyResult<()> { - let replicas = { - let cache = graph_cache(); - let guard = cache.lock().expect("graph cache mutex poisoned"); - let entry = guard.as_ref().ok_or_else(|| { - PyErr::new::( - "graph cache missing while syncing model replicas", - ) - })?; - - if entry.model_ptr != self.source_model_ptr { - return Err(PyErr::new::( - "graph cache model mismatch while syncing model replicas", - )); - } - - let mut models: Vec<(usize, Py)> = entry - .models - .iter() - .map(|(gpu_id, model)| (*gpu_id, model.clone_ref(py))) - .collect(); - models.sort_by_key(|(gpu_id, _)| *gpu_id); - models - .into_iter() - .map(|(_, model)| model) - .collect::>() - }; - - let state_dict: Py = self - .source_model - .bind(py) - .call_method0("state_dict")? - .into(); - for replica in replicas { - let _ = replica - .bind(py) - .call_method1("load_state_dict", (state_dict.clone_ref(py),))?; - } - - Ok(()) - } -} - impl Drop for SelfPlay { fn drop(&mut self) { if let Some(mut session) = self.session.take() { From 12a8c219672270e58551c5fb539a6c5aa8117427 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 21:07:16 -0400 Subject: [PATCH 41/59] stabilize selfplay value targets Align role-relative model outputs with white-centric search targets, fix terminal child labeling, and honor training model/LR config so self-play stops learning from contradictory values. --- alpha_paint/src/evaluation.rs | 7 ++- alpha_paint/src/nn/model.rs | 23 +++---- python/alphapaint_training/model.py | 32 +++++++--- python/alphapaint_training/packed_obs.py | 2 + python/alphapaint_training/train.py | 78 +++++++++++++++++++----- python/scripts/export_value_net.py | 49 +++++++++++---- training/src/descent.rs | 71 ++++++++++++++++++--- training/src/eval.rs | 2 +- 8 files changed, 206 insertions(+), 58 deletions(-) diff --git a/alpha_paint/src/evaluation.rs b/alpha_paint/src/evaluation.rs index 771b600..9cfb51c 100644 --- a/alpha_paint/src/evaluation.rs +++ b/alpha_paint/src/evaluation.rs @@ -27,7 +27,12 @@ impl Evaluator { } pub fn evaluate(&self, board: &Board) -> f32 { - self.model.evaluate(&crate::nn::features::extract(board)) + let value = self.model.evaluate(&crate::nn::features::extract(board)); + if board.is_white_turn() { + value + } else { + -value + } } } diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs index bb28625..db52097 100644 --- a/alpha_paint/src/nn/model.rs +++ b/alpha_paint/src/nn/model.rs @@ -163,7 +163,8 @@ impl ValueModel { "head.fc1.weight", &[MODEL_HIDDEN_DIM, HEAD0_IN_DIM], )?; - let head0_bias = tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; + let head0_bias = + tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; let head1_weight = tensor_array::(tensors, "head.fc2.weight", &[1, MODEL_HIDDEN_DIM])?; let head1_bias = tensor_array::<1>(tensors, "head.fc2.bias", &[1])?[0]; @@ -242,13 +243,7 @@ impl ValueModel { &scratch.head_input, &mut scratch.hidden, ); - (linear_scalar::( - &self.head1_weight, - self.head1_bias, - &scratch.hidden, - )) - .tanh() - * 7.0 + linear_scalar::(&self.head1_weight, self.head1_bias, &scratch.hidden) } } @@ -372,8 +367,7 @@ fn accumulate_conv3x3( for y in 1..BOARD_SIDE - 1 { let row = y * BOARD_SIDE; out_plane[row] += conv3x3_border(in_plane, kernel, 0, y); - out_plane[row + BOARD_SIDE - 1] += - conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); + out_plane[row + BOARD_SIDE - 1] += conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); } } } @@ -526,7 +520,10 @@ fn tensor_f32( } let data = tensor.data(); if data.len() % 4 != 0 { - return Err(format!("tensor {name} has invalid byte length {}", data.len())); + return Err(format!( + "tensor {name} has invalid byte length {}", + data.len() + )); } let mut out = Vec::with_capacity(data.len() / 4); for chunk in data.chunks_exact(4) { @@ -554,7 +551,7 @@ mod tests { } } - let expected = (0.5f32).tanh() * 7.0; + let expected = 0.5f32; assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); } @@ -572,7 +569,7 @@ mod tests { } } - let expected = 1.0f32.tanh() * 7.0; + let expected = 1.0f32; assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); } } diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 3ae2aee..1fcbe33 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections import OrderedDict from typing import cast import torch @@ -8,6 +9,7 @@ from .packed_obs import ( BOARD_CELLS, BOARD_PLANES, + INTRINSIC_TURN_COUNT, INTRINSIC_COUNT, INTRINSIC_SCALE, decode_packed_board, @@ -78,14 +80,14 @@ def __init__( super().__init__() self.board_dtype = board_dtype self._board_buffers: dict[ - tuple[str, int | None, int, torch.dtype], torch.Tensor - ] = {} + tuple[str, int | None, int, torch.dtype, int], torch.Tensor + ] = OrderedDict() self._intrinsic_fp32_buffers: dict[ - tuple[str, int | None, int], torch.Tensor - ] = {} + tuple[str, int | None, int, int], torch.Tensor + ] = OrderedDict() self._intrinsic_buffers: dict[ - tuple[str, int | None, int, torch.dtype], torch.Tensor - ] = {} + tuple[str, int | None, int, torch.dtype, int], torch.Tensor + ] = OrderedDict() self.register_buffer( "intrinsic_scales", torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), @@ -97,14 +99,20 @@ def __init__( hidden_dim=hidden_dim, ) + @staticmethod + def _trim_cache[K](cache: dict[K, torch.Tensor], max_entries: int = 64) -> None: + while len(cache) > max_entries: + cache.pop(next(iter(cache))) + def _buffer_key( self, packed_obs: torch.Tensor, dtype: torch.dtype - ) -> tuple[str, int | None, int, torch.dtype]: + ) -> tuple[str, int | None, int, torch.dtype, int]: return ( packed_obs.device.type, packed_obs.device.index, packed_obs.shape[0], dtype, + packed_obs.data_ptr(), ) def _ensure_decode_buffers( @@ -125,11 +133,13 @@ def _ensure_decode_buffers( memory_format=memory_format, ) self._board_buffers[board_key] = board + self._trim_cache(self._board_buffers) fp32_key = ( packed_obs.device.type, packed_obs.device.index, packed_obs.shape[0], + packed_obs.data_ptr(), ) intrinsics_fp32 = self._intrinsic_fp32_buffers.get(fp32_key) if intrinsics_fp32 is None: @@ -139,6 +149,7 @@ def _ensure_decode_buffers( dtype=torch.float32, ) self._intrinsic_fp32_buffers[fp32_key] = intrinsics_fp32 + self._trim_cache(self._intrinsic_fp32_buffers) intrinsic_key = self._buffer_key(packed_obs, self.board_dtype) intrinsics = self._intrinsic_buffers.get(intrinsic_key) @@ -149,6 +160,7 @@ def _ensure_decode_buffers( dtype=self.board_dtype, ) self._intrinsic_buffers[intrinsic_key] = intrinsics + self._trim_cache(self._intrinsic_buffers) return board, intrinsics_fp32, intrinsics @@ -172,7 +184,11 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: board = board.to(dtype=param_dtype) intrinsics = intrinsics.to(dtype=param_dtype) value = self.value_net(board, intrinsics) - return value.squeeze(-1).tanh() * 7.0 + value = value.squeeze(-1) + turn_count = packed_obs[:, BOARD_CELLS + INTRINSIC_TURN_COUNT].to(torch.int32) + white_to_move = (turn_count & 1) == 0 + sign = torch.where(white_to_move, 1.0, -1.0).to(dtype=value.dtype) + return value * sign __all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index a7d9003..b967867 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -9,6 +9,7 @@ BOARD_PLANES = 17 INTRINSIC_COUNT = 10 OBS_WORDS = BOARD_CELLS + INTRINSIC_COUNT +INTRINSIC_TURN_COUNT = 4 PAINT_STRENGTH_MASK = 0b111 PAINT_IS_ENEMY_BIT = 1 << 3 @@ -279,6 +280,7 @@ def decode_packed_observation( "BOARD_SIDE", "INTRINSIC_SCALE", "INTRINSIC_COUNT", + "INTRINSIC_TURN_COUNT", "OBS_WORDS", "decode_intrinsics", "decode_packed_board", diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index f914a67..dcdb6ae 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -32,6 +32,9 @@ class TrainConfig: workers_per_thread: int = 8 max_gpu_evals_per_move: int = 1_536 lr: float = 3e-4 + pretrain_lr: float | None = None + selfplay_lr: float | None = None + weight_decay: float = 1e-4 width: int = 20 num_blocks: int = 1 hidden_dim: int = 64 @@ -43,7 +46,7 @@ class TrainConfig: device: str = "cuda" checkpoint_interval: int = 3 run_dir: str = "runs/latest" - wandb: bool = False + wandb: bool = True @dataclass(slots=True) @@ -263,6 +266,11 @@ def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: return run_dir, checkpoint_dir +def _set_optimizer_lr(optimizer: torch.optim.Optimizer, lr: float) -> None: + for group in optimizer.param_groups: + group["lr"] = lr + + def _save_checkpoint( *, checkpoint_dir: Path, @@ -294,6 +302,8 @@ def _save_checkpoint( def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) + pretrain_lr = config.lr if config.pretrain_lr is None else config.pretrain_lr + selfplay_lr = config.lr if config.selfplay_lr is None else config.selfplay_lr if device.type == "cuda": torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True @@ -301,16 +311,35 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - wandb.init( - project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) - ) + run = None + if config.wandb: + run = wandb.init( + project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) + ) model = cast( PackedValueModel, - _to_channels_last(PackedValueModel().to(device)), + _to_channels_last( + PackedValueModel( + width=config.width, + num_blocks=config.num_blocks, + hidden_dim=config.hidden_dim, + ).to(device) + ), ) model.eval() - optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) + param_count = sum(param.numel() for param in model.parameters()) + print( + f"model width={config.width} blocks={config.num_blocks} hidden={config.hidden_dim} params={param_count}" + ) + if run is not None: + run.summary["model_params"] = param_count + initial_lr = pretrain_lr if config.pretrain_terminal_samples > 0 else selfplay_lr + optimizer = torch.optim.AdamW( + model.parameters(), + lr=initial_lr, + weight_decay=config.weight_decay, + ) if config.pretrain_terminal_samples > 0: pretrain_steps = ( @@ -354,14 +383,16 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"pretrain step={step_number}/{pretrain_steps} samples={samples_done}/{config.pretrain_terminal_samples} " f"loss={window_mean_loss:.6f} elapsed={elapsed:.2f}s" ) - wandb.log( - { - "pretrain_step": step_number, - "pretrain_samples_total": samples_done, - "pretrain_loss_mean": window_mean_loss, - "pretrain_seconds": elapsed, - } - ) + if run is not None: + wandb.log( + { + "pretrain_step": step_number, + "pretrain_samples_total": samples_done, + "pretrain_loss_mean": window_mean_loss, + "pretrain_seconds": elapsed, + "learning_rate": optimizer.param_groups[0]["lr"], + } + ) if device.type == "cuda": torch.cuda.synchronize(device) @@ -370,6 +401,9 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print( f"pretrain_complete samples={config.pretrain_terminal_samples} steps={pretrain_steps} elapsed={pretrain_seconds:.2f}s" ) + if selfplay_lr != pretrain_lr: + _set_optimizer_lr(optimizer, selfplay_lr) + print(f"selfplay_lr={selfplay_lr:.6g}") replay_buffer = EphemeralReplayBuffer(config.replay_capacity) selfplay = SelfPlay( @@ -713,6 +747,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "train_backward_seconds": train_backward_seconds, "train_optimizer_seconds": train_optimizer_seconds, "loss_mean": mean_loss, + "learning_rate": optimizer.param_groups[0]["lr"], "timestamp": time.time(), } print( @@ -728,7 +763,8 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " f"loss={mean_loss:.6f}" ) - wandb.log(record) + if run is not None: + wandb.log(record) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -769,6 +805,8 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: previous_descent_backup_nanos = descent_backup_nanos finally: selfplay.drop() + if run is not None: + run.finish() return model, losses @@ -795,6 +833,12 @@ def _parse_args() -> TrainConfig: default=defaults.max_gpu_evals_per_move, ) parser.add_argument("--lr", type=float, default=defaults.lr) + parser.add_argument("--pretrain-lr", type=float, default=defaults.pretrain_lr) + parser.add_argument("--selfplay-lr", type=float, default=defaults.selfplay_lr) + parser.add_argument("--weight-decay", type=float, default=defaults.weight_decay) + parser.add_argument("--width", type=int, default=defaults.width) + parser.add_argument("--num-blocks", type=int, default=defaults.num_blocks) + parser.add_argument("--hidden-dim", type=int, default=defaults.hidden_dim) parser.add_argument( "--pretrain-terminal-samples", type=int, @@ -821,7 +865,9 @@ def _parse_args() -> TrainConfig: "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) parser.add_argument("--run-dir", default=defaults.run_dir) - parser.add_argument("--wandb", action="store_true", default=True) + parser.add_argument("--wandb", dest="wandb", action="store_true") + parser.add_argument("--no-wandb", dest="wandb", action="store_false") + parser.set_defaults(wandb=defaults.wandb) args = parser.parse_args() return TrainConfig(**vars(args)) diff --git a/python/scripts/export_value_net.py b/python/scripts/export_value_net.py index eea0cf7..21ae238 100644 --- a/python/scripts/export_value_net.py +++ b/python/scripts/export_value_net.py @@ -9,11 +9,28 @@ from safetensors.torch import save_file +def _infer_model_config(state_dict: dict[str, torch.Tensor]) -> dict[str, int]: + width = int(state_dict["value_net.stem.0.weight"].shape[0]) + hidden_dim = int(state_dict["value_net.head.0.weight"].shape[0]) + block_ids = { + int(key.split(".")[2]) + for key in state_dict + if key.startswith("value_net.blocks.") and key.endswith("conv1.weight") + } + return { + "width": width, + "hidden_dim": hidden_dim, + "num_blocks": len(block_ids), + } + + def _rename_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: renamed: dict[str, torch.Tensor] = {} def copy(src: str, dst: str) -> None: - tensor = state_dict[src].detach().to(dtype=torch.float32, device="cpu").contiguous() + tensor = ( + state_dict[src].detach().to(dtype=torch.float32, device="cpu").contiguous() + ) renamed[dst] = tensor copy("value_net.stem.0.weight", "stem.conv.weight") @@ -22,10 +39,7 @@ def copy(src: str, dst: str) -> None: copy("value_net.stem.1.running_mean", "stem.bn.running_mean") copy("value_net.stem.1.running_var", "stem.bn.running_var") - width = state_dict["value_net.stem.0.weight"].shape[0] - head_linear1 = state_dict["value_net.head.0.weight"] - hidden_dim = head_linear1.shape[0] - + model_config = _infer_model_config(state_dict) block_ids = sorted( { int(key.split(".")[2]) @@ -52,9 +66,13 @@ def copy(src: str, dst: str) -> None: copy("value_net.head.2.weight", "head.fc2.weight") copy("value_net.head.2.bias", "head.fc2.bias") - renamed["__meta.width"] = torch.tensor([width], dtype=torch.float32) - renamed["__meta.hidden_dim"] = torch.tensor([hidden_dim], dtype=torch.float32) - renamed["__meta.num_blocks"] = torch.tensor([len(block_ids)], dtype=torch.float32) + renamed["__meta.width"] = torch.tensor([model_config["width"]], dtype=torch.float32) + renamed["__meta.hidden_dim"] = torch.tensor( + [model_config["hidden_dim"]], dtype=torch.float32 + ) + renamed["__meta.num_blocks"] = torch.tensor( + [model_config["num_blocks"]], dtype=torch.float32 + ) return renamed @@ -74,12 +92,19 @@ def _build_fixtures( if fixture_count <= 0: return - from alphapaint_training import decode_packed_observation, sample_random_terminal_batch + from alphapaint_training import ( + decode_packed_observation, + sample_random_terminal_batch, + ) + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) state_dict = _extract_state_dict(checkpoint) from alphapaint_training.model import PackedValueModel - model = PackedValueModel(board_dtype=torch.float32) + model = PackedValueModel( + board_dtype=torch.float32, + **_infer_model_config(state_dict), + ) model.load_state_dict(state_dict) model.eval() @@ -111,7 +136,9 @@ def main() -> None: parser = argparse.ArgumentParser( description="Export AlphaPaint value network checkpoints for Rust CPU inference." ) - parser.add_argument("--checkpoint", type=Path, required=True, help="Input .pt checkpoint") + parser.add_argument( + "--checkpoint", type=Path, required=True, help="Input .pt checkpoint" + ) parser.add_argument( "--output", type=Path, diff --git a/training/src/descent.rs b/training/src/descent.rs index 0fa3da2..99c4c7d 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -5,7 +5,7 @@ //! and ordinal distribution for action selection. use alpha_paint::board::actions::Move; -use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; +use alpha_paint::board::{Action, ApplyActionOutcome, Board, Rollback, TerminalState}; use rand::rngs::SmallRng; use rand::{Rng, RngExt}; use std::sync::atomic::{AtomicU64, Ordering}; @@ -222,6 +222,36 @@ impl SearchNode { node } + fn apply_play_instead( + mut board: Board, + action: Action, + rollback: Rollback, + play_instead: Action, + ) -> Board { + board.rollback(action, rollback); + let (outcome, _) = board.apply_action(play_instead); + debug_assert!(matches!(outcome, ApplyActionOutcome::Terminal { .. })); + board + } + + fn apply_killshot_moves(mut board: Board, moves: &[Move]) -> Board { + let last_idx = moves.len().saturating_sub(1); + for (idx, mv) in moves.iter().copied().enumerate() { + let action = if idx == last_idx { + Action::FinalMove(mv) + } else { + Action::Move(mv) + }; + let (outcome, _) = board.apply_action(action); + if idx == last_idx { + debug_assert!(matches!(outcome, ApplyActionOutcome::Terminal { .. })); + } else { + debug_assert!(matches!(outcome, ApplyActionOutcome::Ongoing)); + } + } + board + } + /// Expand a node: evaluate all children with the neural network. async fn build_self( board: &Board, @@ -272,7 +302,7 @@ impl SearchNode { }) .collect(); - for (&action, local_board, child_outcome, _, value) in batch { + for (&action, local_board, child_outcome, rollback, value) in batch { match child_outcome { ApplyActionOutcome::Ongoing => { // Evaluate this child with neural net @@ -281,24 +311,49 @@ impl SearchNode { timing.record_eval_await(await_started_at); eval_results.push(ChildEvalResult { action, value }); } - ApplyActionOutcome::Terminal { terminal } - | ApplyActionOutcome::PlayInstead { terminal, .. } => { + ApplyActionOutcome::Terminal { terminal } => { + let term_value = Self::value_from_term(&local_board, terminal); + new_node.children.push(ChildData { + action, + child_value: term_value, + entrance_count: 0, + node: Some(Box::new(SearchNode::new( + term_value, + terminal.value(), + true, + ))), + }); + } + ApplyActionOutcome::PlayInstead { + terminal, + play_instead, + } => { + let terminal_board = Self::apply_play_instead( + local_board, + action, + rollback, + play_instead, + ); + let term_value = Self::value_from_term(&terminal_board, terminal); new_node.children.push(ChildData { action, - child_value: Self::value_from_term(&local_board, terminal), + child_value: term_value, entrance_count: 0, node: Some(Box::new(SearchNode::new( - Self::value_from_term(&local_board, terminal), + term_value, terminal.value(), true, ))), }); } ApplyActionOutcome::Killshot { terminal, moves } => { - let chain = Self::build_killshot_chain(&local_board, terminal, &moves); + let terminal_board = Self::apply_killshot_moves(local_board, &moves); + let term_value = Self::value_from_term(&terminal_board, terminal); + let chain = + Self::build_killshot_chain(&terminal_board, terminal, &moves); new_node.children.push(ChildData { action, - child_value: Self::value_from_term(&local_board, terminal), + child_value: term_value, entrance_count: 0, node: Some(Box::new(chain)), }); diff --git a/training/src/eval.rs b/training/src/eval.rs index 626d93f..f3e124e 100644 --- a/training/src/eval.rs +++ b/training/src/eval.rs @@ -12,7 +12,7 @@ use crate::queue::GpuJobQueue; /// Async evaluator trait for neural network inference. /// /// Returns a scalar value estimate for the given board position. -/// Value is from the current player's perspective. +/// Value is from White's perspective. pub trait Evaluator { /// Evaluate the board and return a value estimate. fn evaluate(&self, board: Board) -> impl Future; From f518c83ed5e89dab8247d6f2bb3e76e20a73a469 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sat, 28 Mar 2026 21:34:22 -0400 Subject: [PATCH 42/59] cap search by total turn time Track a single wall-clock budget across the whole turn and let elapsed time, not an old heuristic iteration estimate, stop Descent so model-backed search stays within the intended move budget. --- alpha_paint/src/bindings.rs | 39 +++++++++++++------------------------ alpha_paint/src/search.rs | 33 ++++++++++++++++++++----------- 2 files changed, 36 insertions(+), 36 deletions(-) diff --git a/alpha_paint/src/bindings.rs b/alpha_paint/src/bindings.rs index 894cb72..1dbb5aa 100644 --- a/alpha_paint/src/bindings.rs +++ b/alpha_paint/src/bindings.rs @@ -1,7 +1,7 @@ -use pyo3::{PyResult, exceptions::PyRuntimeError, pyclass, pymethods}; -use rand::{RngExt, SeedableRng, rngs::StdRng, seq::SliceRandom}; +use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, PyResult}; +use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use crate::board::{board_structs::*, structs::*, *}; use crate::evaluation::Evaluator; @@ -186,36 +186,25 @@ impl PyBoard { let mut py_actions: Vec = vec![]; let evaluator = Evaluator::new(&local_board).map_err(PyRuntimeError::new_err)?; - // Time control - // For reference, on my laptop we can hit around 20k iterations per second (with incremental logic) - let max_moves_left = ((2000 - local_board.turn_count) / 2) as u32; - let (max_duration, iterations) = if time_left < 5.0 { - // At 100 iterations per move, we could finish the entire game in ~5 seconds - (time_left / (max_moves_left as f32 * 5.0), 100) + let turn_budget_secs = if time_left < 5.0 { + time_left / (max_moves_left.max(1) as f32 * 5.0) } else if time_left < 30.0 { - // Speed mode (try to leave 2.5s buffer while still finishing all moves) - let iter_budget = ((time_left - 2.5) * 18000.0) as u32; - ( - time_left / ((max_moves_left as f32 - 2.5) * 5.0), - iter_budget / (max_moves_left * 5), - ) + (time_left - 2.5).max(0.0) / (max_moves_left.max(1) as f32 * 5.0) } else if time_left < 90.0 { - // Late game mode (try to leave 20s buffer) - let iter_budget = ((time_left - 20.0) * 18000.0) as u32; - ( - time_left / ((max_moves_left as f32 - 20.0) * 5.0), - iter_budget / (max_moves_left * 5), - ) + (time_left - 20.0).max(0.0) / (max_moves_left.max(1) as f32 * 5.0) } else { - // For the first half of our time, let's use it more liberally. We can - // finish most games before we get under 90s anyway - (0.150, 15_000) + 0.150 }; + let turn_budget = Duration::from_secs_f32(turn_budget_secs.max(0.005)); + let turn_started_at = Instant::now(); let mut tree = GameSearchTree::new(&local_board, &evaluator); loop { - tree.run_descent_for_iter(iterations.max(50), Duration::from_secs_f32(max_duration)); + let remaining_budget = turn_budget.saturating_sub(turn_started_at.elapsed()); + if !remaining_budget.is_zero() { + tree.run_descent_for_iter(u32::MAX, remaining_budget); + } let (action_id, action) = tree.get_best_action_and_index(); let player_coord = local_board.current_player_coord(); diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index e4015e4..280d59d 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -131,9 +131,13 @@ impl SearchNode { if self.completion_value.abs() == 1 { true } else { - self.children - .iter() - .all(|child| child.node.as_ref().map(|node| node.resolved).unwrap_or(false)) + self.children.iter().all(|child| { + child + .node + .as_ref() + .map(|node| node.resolved) + .unwrap_or(false) + }) } } @@ -254,7 +258,11 @@ impl SearchNode { let (outcome, _) = state.apply_action(action); let node = Box::new(SearchNode::build_self(&state, outcome, evaluator, rng)); let value = node.value; - if let Some(index) = self.children.iter().position(|child| child.action == action) { + if let Some(index) = self + .children + .iter() + .position(|child| child.action == action) + { self.children[index].node = Some(node); } value @@ -324,11 +332,7 @@ impl GameSearchTree<'_> { .iter() .enumerate() .max_by(|&(_, a), &(_, b)| { - ( - a.completion_value(), - a.entrance_count as i32, - a.child_value, - ) + (a.completion_value(), a.entrance_count as i32, a.child_value) .partial_cmp(&( b.completion_value(), b.entrance_count as i32, @@ -392,14 +396,21 @@ impl GameSearchTree<'_> { } pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let desired_completion = if self.root_state.is_white_turn() { 1 } else { -1 }; + let desired_completion = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; + if max_duration.is_zero() { + return; + } let started_at = std::time::Instant::now(); if self.root_node.children.len() <= 1 { return; } for epoch in 0..iterations { - if started_at.elapsed() > max_duration && epoch >= 50 { + if epoch > 0 && started_at.elapsed() >= max_duration { break; } let best_action = self.get_best_action_index(); From 9014f504a25a99aafe62f33eeaf316fc03372277 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 01:06:01 -0400 Subject: [PATCH 43/59] retune default 90k training config Switch the default trainer to the 90k model and the faster 16x8/2048 setup, and shrink the production GPU queue batch to 128 so CUDA graph lane memory stays under control during longer runs. --- python/alphapaint_training/train.py | 16 ++++++++-------- training/src/queue.rs | 4 ++-- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index dcdb6ae..823b9b4 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -25,19 +25,19 @@ class TrainConfig: rounds: int = 1 samples_per_round: int = 1_048_576 - train_steps_per_round: int = 128 - batch_size: int = 8_192 + train_steps_per_round: int = 256 + batch_size: int = 4_096 replay_capacity: int = 16_000_000 - num_threads: int = 32 + num_threads: int = 16 workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 1_536 + max_gpu_evals_per_move: int = 2_048 lr: float = 3e-4 pretrain_lr: float | None = None - selfplay_lr: float | None = None + selfplay_lr: float | None = 5e-5 weight_decay: float = 1e-4 - width: int = 20 - num_blocks: int = 1 - hidden_dim: int = 64 + width: int = 32 + num_blocks: int = 4 + hidden_dim: int = 256 pretrain_terminal_samples: int = 5_000_000 pretrain_batch_size: int = 8_192 pretrain_log_interval: int = 10 diff --git a/training/src/queue.rs b/training/src/queue.rs index 6ad0fac..1b2e685 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -16,12 +16,12 @@ use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; use crate::BatchDim; -/// Number of jobs per batch. In production this would be 256. +/// Number of jobs per batch. In production this would be 128. /// Using a smaller value for tests to avoid deadlock with few workers. #[cfg(test)] pub const BATCH_SIZE: usize = 16; #[cfg(not(test))] -pub const BATCH_SIZE: usize = 256; +pub const BATCH_SIZE: usize = 128; const SLOT_MULTIPLIER: usize = 32; From 86d2389f6a5ce4f0019e079b97dccdf24a3e3a98 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 16:09:22 -0400 Subject: [PATCH 44/59] some debug stuff --- alpha_paint/benches/nn_latency.rs | 2 +- alpha_paint/src/bindings.rs | 4 +- alpha_paint/src/evaluation.rs | 6 +-- alpha_paint/src/nn/features.rs | 28 +++++++---- python/alphapaint_training/train.py | 2 +- training/src/lib.rs | 1 + training/src/queue.rs | 55 ++++++++++++++++++++- training/src/training.rs | 75 ++++++++++++++++++++++++++--- 8 files changed, 146 insertions(+), 27 deletions(-) diff --git a/alpha_paint/benches/nn_latency.rs b/alpha_paint/benches/nn_latency.rs index e3d6297..2ecd24f 100644 --- a/alpha_paint/benches/nn_latency.rs +++ b/alpha_paint/benches/nn_latency.rs @@ -1,5 +1,5 @@ -use std::time::Duration; use std::hint::black_box; +use std::time::Duration; use alpha_paint::board::Board; use alpha_paint::evaluation::Evaluator; diff --git a/alpha_paint/src/bindings.rs b/alpha_paint/src/bindings.rs index 1dbb5aa..5314fa1 100644 --- a/alpha_paint/src/bindings.rs +++ b/alpha_paint/src/bindings.rs @@ -1,5 +1,5 @@ -use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, PyResult}; -use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng}; +use pyo3::{PyResult, exceptions::PyRuntimeError, pyclass, pymethods}; +use rand::{RngExt, SeedableRng, rngs::StdRng, seq::SliceRandom}; use std::sync::Arc; use std::time::{Duration, Instant}; diff --git a/alpha_paint/src/evaluation.rs b/alpha_paint/src/evaluation.rs index 9cfb51c..297b703 100644 --- a/alpha_paint/src/evaluation.rs +++ b/alpha_paint/src/evaluation.rs @@ -28,11 +28,7 @@ impl Evaluator { pub fn evaluate(&self, board: &Board) -> f32 { let value = self.model.evaluate(&crate::nn::features::extract(board)); - if board.is_white_turn() { - value - } else { - -value - } + if board.is_white_turn() { value } else { -value } } } diff --git a/alpha_paint/src/nn/features.rs b/alpha_paint/src/nn/features.rs index 2f73455..1a09a65 100644 --- a/alpha_paint/src/nn/features.rs +++ b/alpha_paint/src/nn/features.rs @@ -1,6 +1,6 @@ +use crate::board::Board; use crate::board::board_structs::Player; use crate::board::structs::Coordinate; -use crate::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; @@ -168,8 +168,7 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; - out[INTRINSIC_OPPONENT_BEACONS] = - board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; } else { out[INTRINSIC_CURRENT_STAMINA] = board.black_stamina as f32; out[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina as f32; @@ -177,8 +176,7 @@ pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; - out[INTRINSIC_CURRENT_BEACONS] = - board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; } @@ -221,9 +219,18 @@ mod tests { assert_eq!(features.board[plane_index(WALL, 31, 31)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 99.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 88.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_CURRENT_TERRITORY], 1.0 / 1024.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], 1.0 / 1024.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 88.0 / 420.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_CURRENT_TERRITORY], + 1.0 / 1024.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], + 1.0 / 1024.0 + ); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_BEACONS], 1.0 / 1024.0); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 2.0 / 64.0); } @@ -244,7 +251,10 @@ mod tests { assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 0, 0)], 1.0); assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 88.0 / 420.0); - assert_eq!(features.intrinsics[INTRINSIC_OPPONENT_STAMINA], 99.0 / 420.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 99.0 / 420.0 + ); assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 5.0 / 64.0); } } diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 823b9b4..c55b908 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -23,7 +23,7 @@ @dataclass(slots=True) class TrainConfig: - rounds: int = 1 + rounds: int = 200 samples_per_round: int = 1_048_576 train_steps_per_round: int = 256 batch_size: int = 4_096 diff --git a/training/src/lib.rs b/training/src/lib.rs index 2808d15..cafa906 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -291,6 +291,7 @@ impl SelfPlay { num_threads, workers_per_thread, seed, + thread_stack_size_bytes: SessionConfig::default().thread_stack_size_bytes, worker: WorkerConfig { max_gpu_evals_per_move, }, diff --git a/training/src/queue.rs b/training/src/queue.rs index 1b2e685..e1ff52a 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -25,6 +25,16 @@ pub const BATCH_SIZE: usize = 128; const SLOT_MULTIPLIER: usize = 32; +#[derive(Debug, Clone, Copy)] +pub struct QueueDebugSnapshot { + pub next_ticket: u64, + pub inflight_batches: u64, + pub remainder: usize, + pub open_batch_writes: usize, + pub num_batches: usize, + pub total_slots: usize, +} + /// Compute queue shape for a given worker count. /// /// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a @@ -289,6 +299,27 @@ where } } + pub fn debug_snapshot(&self) -> QueueDebugSnapshot { + let next_ticket = self.write_ticket.load(Ordering::Acquire); + let remainder = (next_ticket % BATCH_SIZE as u64) as usize; + let open_batch_writes = if remainder == 0 { + 0 + } else { + let batch_number = next_ticket / BATCH_SIZE as u64; + let batch_idx = (batch_number as usize) % self.num_batches; + self.state.batch_writes[batch_idx].load(Ordering::Acquire) as usize + }; + + QueueDebugSnapshot { + next_ticket, + inflight_batches: self.state.inflight_batches.load(Ordering::Acquire), + remainder, + open_batch_writes, + num_batches: self.num_batches, + total_slots: self.total_slots, + } + } + /// Flush the open partial batch, if any, and wait for all in-flight GPU work. /// /// # Safety @@ -298,8 +329,19 @@ where /// The intended callsite is after all session worker threads have left the /// executor polling loop at a pause boundary. pub unsafe fn quiesce_exclusive(&self) -> bool { - let next_ticket = self.write_ticket.load(Ordering::Acquire); - let remainder = (next_ticket % BATCH_SIZE as u64) as usize; + let before = self.debug_snapshot(); + eprintln!( + "queue_quiesce start next_ticket={} inflight={} remainder={} open_writes={} num_batches={} total_slots={}", + before.next_ticket, + before.inflight_batches, + before.remainder, + before.open_batch_writes, + before.num_batches, + before.total_slots, + ); + + let next_ticket = before.next_ticket; + let remainder = before.remainder; let mut flushed = false; if remainder != 0 { @@ -322,6 +364,15 @@ where } self.wait_until_idle(); + let after = self.debug_snapshot(); + eprintln!( + "queue_quiesce done flushed={} next_ticket={} inflight={} remainder={} open_writes={}", + flushed, + after.next_ticket, + after.inflight_batches, + after.remainder, + after.open_batch_writes, + ); flushed } diff --git a/training/src/training.rs b/training/src/training.rs index 66692b4..22545fd 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -14,7 +14,7 @@ use rand_chacha::ChaCha8Rng; use crate::eval::GpuEvaluator; use crate::executor::Executor; -use crate::queue::{BatchCompletion, GpuJobQueue}; +use crate::queue::{BatchCompletion, GpuJobQueue, QueueDebugSnapshot}; use crate::replay_buffer::ReplayBuffer; use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; use crate::BatchDim; @@ -139,6 +139,8 @@ pub struct SessionConfig { pub worker: WorkerConfig, /// Random seed for reproducibility. pub seed: u64, + /// Stack size for each self-play OS thread. + pub thread_stack_size_bytes: usize, } impl Default for SessionConfig { @@ -148,6 +150,7 @@ impl Default for SessionConfig { workers_per_thread: 8, worker: WorkerConfig::default(), seed: 42, + thread_stack_size_bytes: 128 * 1024 * 1024, } } } @@ -156,6 +159,7 @@ impl Default for SessionConfig { trait QueueNotify: Send + Sync { unsafe fn quiesce_exclusive(&self) -> bool; fn notify_all(&self); + fn debug_snapshot(&self) -> QueueDebugSnapshot; } impl QueueNotify for GpuJobQueue @@ -171,6 +175,10 @@ where fn notify_all(&self) { GpuJobQueue::notify_all(self); } + + fn debug_snapshot(&self) -> QueueDebugSnapshot { + GpuJobQueue::debug_snapshot(self) + } } /// A persistent self-play session that owns worker threads and preserves @@ -210,9 +218,15 @@ impl SelfPlaySession { let config = config.clone(); let replay_buffer = replay_buffer.clone(); - let handle = thread::spawn(move || { - session_thread_main(thread_id, queue, config, control, &replay_buffer); - }); + let thread_name = format!("selfplay-{thread_id}"); + let stack_size = config.thread_stack_size_bytes; + let handle = thread::Builder::new() + .name(thread_name) + .stack_size(stack_size) + .spawn(move || { + session_thread_main(thread_id, queue, config, control, &replay_buffer); + }) + .expect("failed to spawn self-play thread"); threads.push(handle); } @@ -236,6 +250,17 @@ impl SelfPlaySession { /// Block until at least `target_samples` absolute samples have been /// collected, then pause and quiesce all workers. pub fn wait_for(&self, target_samples: usize) -> usize { + let start_snapshot = self.queue_notify.debug_snapshot(); + eprintln!( + "selfplay_wait start target_samples={} current_samples={} active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", + target_samples, + self.control.samples_collected.load(Ordering::Acquire), + self.control.active_pollers.load(Ordering::Acquire), + start_snapshot.next_ticket, + start_snapshot.inflight_batches, + start_snapshot.remainder, + start_snapshot.open_batch_writes, + ); self.control .target_samples .store(target_samples, Ordering::Release); @@ -260,10 +285,27 @@ impl SelfPlaySession { } } + let reached_snapshot = self.queue_notify.debug_snapshot(); + eprintln!( + "selfplay_wait target_reached samples={} active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", + self.control.samples_collected.load(Ordering::Acquire), + self.control.active_pollers.load(Ordering::Acquire), + reached_snapshot.next_ticket, + reached_snapshot.inflight_batches, + reached_snapshot.remainder, + reached_snapshot.open_batch_writes, + ); + self.control.running.store(false, Ordering::Release); self.queue_notify.notify_all(); self.control.wake_all(); + eprintln!( + "selfplay_wait pausing target_samples={} active_pollers={}", + target_samples, + self.control.active_pollers.load(Ordering::Acquire), + ); + { let mut guard = self .control @@ -281,12 +323,31 @@ impl SelfPlaySession { } } + let drained_snapshot = self.queue_notify.debug_snapshot(); + eprintln!( + "selfplay_wait pollers_drained active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", + self.control.active_pollers.load(Ordering::Acquire), + drained_snapshot.next_ticket, + drained_snapshot.inflight_batches, + drained_snapshot.remainder, + drained_snapshot.open_batch_writes, + ); + // SAFETY: `wait_for` has already stopped the session and waited until // every session thread has left the executor loop (`active_pollers == 0`), // so no thread can submit new queue writes or race another quiesce call. - unsafe { - self.queue_notify.quiesce_exclusive(); - } + let flushed = unsafe { self.queue_notify.quiesce_exclusive() }; + + let end_snapshot = self.queue_notify.debug_snapshot(); + eprintln!( + "selfplay_wait quiesce_done flushed={} samples={} queue_next_ticket={} inflight={} remainder={} open_writes={}", + flushed, + self.control.samples_collected.load(Ordering::Acquire), + end_snapshot.next_ticket, + end_snapshot.inflight_batches, + end_snapshot.remainder, + end_snapshot.open_batch_writes, + ); self.control.samples_collected.load(Ordering::Acquire) } From 851374a9eed7d1f652e0d84158c4ee03e12a63b4 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 19:14:57 -0400 Subject: [PATCH 45/59] cleanup logging --- python/alphapaint_training/train.py | 76 ++++++++++++++++++++++++----- training/src/queue.rs | 55 +-------------------- training/src/training.rs | 60 ++--------------------- 3 files changed, 70 insertions(+), 121 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index c55b908..34d992c 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -299,6 +299,60 @@ def _save_checkpoint( return checkpoint_path +def _tensor_summary(prefix: str, values: torch.Tensor) -> dict[str, float]: + values = values.detach().float().cpu() + return { + f"{prefix}_mean": float(values.mean().item()), + f"{prefix}_std": float(values.std(unbiased=False).item()), + f"{prefix}_min": float(values.min().item()), + f"{prefix}_max": float(values.max().item()), + } + + +def _collect_round_diagnostics( + model: PackedValueModel, + replay_buffer: EphemeralReplayBuffer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[dict[str, float], dict[str, wandb.Histogram]]: + obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( + replay_buffer, batch_size, seed, device + ) + autocast = ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if device.type == "cuda" + else nullcontext() + ) + model.eval() + with torch.no_grad(): + with autocast: + pred = model(obs) + if device.type == "cuda": + torch.cuda.synchronize(device) + + pred = pred.float() + residual = pred - target + stats = { + "diag_sample_seconds": sample_seconds, + "diag_h2d_seconds": h2d_seconds, + **_tensor_summary("diag_target", target), + **_tensor_summary("diag_pred", pred), + **_tensor_summary("diag_residual", residual), + } + histograms = { + "diag_target_hist": wandb.Histogram( + target.detach().float().cpu().numpy().tolist() + ), + "diag_pred_hist": wandb.Histogram(pred.detach().float().cpu().numpy().tolist()), + "diag_residual_hist": wandb.Histogram( + residual.detach().float().cpu().numpy().tolist() + ), + } + return stats, histograms + + def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device = _device(config.device) torch.manual_seed(config.seed) @@ -677,6 +731,13 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: mean_loss = float(round_loss_values.mean().item()) else: mean_loss = 0.0 + diag_stats, diag_histograms = _collect_round_diagnostics( + model, + replay_buffer, + batch_size=config.batch_size, + seed=config.seed + round_idx * 10_000 + config.train_steps_per_round, + device=device, + ) record = { "round": round_number, "samples_total": collected, @@ -749,22 +810,13 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "loss_mean": mean_loss, "learning_rate": optimizer.param_groups[0]["lr"], "timestamp": time.time(), + **diag_stats, } print( - f"round={round_number} samples={collected} (+{samples_added}) games={games} (+{games_added}) " - f"gpu_evals={gpu_evals_added} ({gpu_evals_added / max(collect_seconds, 1e-9):.1f}/s) " - f"batches={gpu_batches_added} replay={replay_size} act={action_steps_added} final={final_actions_added} " - f"nonfinal={nonfinal_actions_added} act/turn={actions_per_turn_display} act/game={actions_per_game_display} " - f"avg_tc={avg_turn_count:.1f} max_tc={max_turn_count_seen} cg_tc={completed_game_turn_count_display} " - f"collect={collect_seconds:.2f}s train={train_seconds:.2f}s " - f"sp_ms/act=build:{build_ms_per_action:.1f} desc:{descent_ms_per_action:.1f} coll:{collect_ms_per_action:.1f} push:{push_ms_per_action:.1f} " - f"desc_ms/act=exp:{descent_expand_ms_per_action:.1f} app:{descent_apply_ms_per_action:.1f} sub:{descent_submit_ms_per_action:.1f} wait:{descent_wait_ms_per_action:.1f} bk:{descent_backup_ms_per_action:.1f} other:{descent_other_ms_per_action:.1f} " - f"desc_us/eval=sub:{descent_submit_us_per_eval:.2f} wait:{descent_wait_us_per_eval:.2f} " - f"tr_s=sample:{train_sample_seconds:.2f} h2d:{train_h2d_seconds:.2f} fwd:{train_forward_seconds:.2f} bwd:{train_backward_seconds:.2f} opt:{train_optimizer_seconds:.2f} " - f"loss={mean_loss:.6f}" + f"round={round_number} samples={collected} games={games} gpu_evals/s={gpu_evals_added / max(collect_seconds, 1e-9):.1f} replay={replay_size} loss={mean_loss:.6f}" ) if run is not None: - wandb.log(record) + wandb.log({**record, **diag_histograms}) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 diff --git a/training/src/queue.rs b/training/src/queue.rs index e1ff52a..1b2e685 100644 --- a/training/src/queue.rs +++ b/training/src/queue.rs @@ -25,16 +25,6 @@ pub const BATCH_SIZE: usize = 128; const SLOT_MULTIPLIER: usize = 32; -#[derive(Debug, Clone, Copy)] -pub struct QueueDebugSnapshot { - pub next_ticket: u64, - pub inflight_batches: u64, - pub remainder: usize, - pub open_batch_writes: usize, - pub num_batches: usize, - pub total_slots: usize, -} - /// Compute queue shape for a given worker count. /// /// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a @@ -299,27 +289,6 @@ where } } - pub fn debug_snapshot(&self) -> QueueDebugSnapshot { - let next_ticket = self.write_ticket.load(Ordering::Acquire); - let remainder = (next_ticket % BATCH_SIZE as u64) as usize; - let open_batch_writes = if remainder == 0 { - 0 - } else { - let batch_number = next_ticket / BATCH_SIZE as u64; - let batch_idx = (batch_number as usize) % self.num_batches; - self.state.batch_writes[batch_idx].load(Ordering::Acquire) as usize - }; - - QueueDebugSnapshot { - next_ticket, - inflight_batches: self.state.inflight_batches.load(Ordering::Acquire), - remainder, - open_batch_writes, - num_batches: self.num_batches, - total_slots: self.total_slots, - } - } - /// Flush the open partial batch, if any, and wait for all in-flight GPU work. /// /// # Safety @@ -329,19 +298,8 @@ where /// The intended callsite is after all session worker threads have left the /// executor polling loop at a pause boundary. pub unsafe fn quiesce_exclusive(&self) -> bool { - let before = self.debug_snapshot(); - eprintln!( - "queue_quiesce start next_ticket={} inflight={} remainder={} open_writes={} num_batches={} total_slots={}", - before.next_ticket, - before.inflight_batches, - before.remainder, - before.open_batch_writes, - before.num_batches, - before.total_slots, - ); - - let next_ticket = before.next_ticket; - let remainder = before.remainder; + let next_ticket = self.write_ticket.load(Ordering::Acquire); + let remainder = (next_ticket % BATCH_SIZE as u64) as usize; let mut flushed = false; if remainder != 0 { @@ -364,15 +322,6 @@ where } self.wait_until_idle(); - let after = self.debug_snapshot(); - eprintln!( - "queue_quiesce done flushed={} next_ticket={} inflight={} remainder={} open_writes={}", - flushed, - after.next_ticket, - after.inflight_batches, - after.remainder, - after.open_batch_writes, - ); flushed } diff --git a/training/src/training.rs b/training/src/training.rs index 22545fd..8e3c472 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -14,7 +14,7 @@ use rand_chacha::ChaCha8Rng; use crate::eval::GpuEvaluator; use crate::executor::Executor; -use crate::queue::{BatchCompletion, GpuJobQueue, QueueDebugSnapshot}; +use crate::queue::{BatchCompletion, GpuJobQueue}; use crate::replay_buffer::ReplayBuffer; use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; use crate::BatchDim; @@ -159,7 +159,6 @@ impl Default for SessionConfig { trait QueueNotify: Send + Sync { unsafe fn quiesce_exclusive(&self) -> bool; fn notify_all(&self); - fn debug_snapshot(&self) -> QueueDebugSnapshot; } impl QueueNotify for GpuJobQueue @@ -175,10 +174,6 @@ where fn notify_all(&self) { GpuJobQueue::notify_all(self); } - - fn debug_snapshot(&self) -> QueueDebugSnapshot { - GpuJobQueue::debug_snapshot(self) - } } /// A persistent self-play session that owns worker threads and preserves @@ -250,17 +245,6 @@ impl SelfPlaySession { /// Block until at least `target_samples` absolute samples have been /// collected, then pause and quiesce all workers. pub fn wait_for(&self, target_samples: usize) -> usize { - let start_snapshot = self.queue_notify.debug_snapshot(); - eprintln!( - "selfplay_wait start target_samples={} current_samples={} active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", - target_samples, - self.control.samples_collected.load(Ordering::Acquire), - self.control.active_pollers.load(Ordering::Acquire), - start_snapshot.next_ticket, - start_snapshot.inflight_batches, - start_snapshot.remainder, - start_snapshot.open_batch_writes, - ); self.control .target_samples .store(target_samples, Ordering::Release); @@ -285,27 +269,10 @@ impl SelfPlaySession { } } - let reached_snapshot = self.queue_notify.debug_snapshot(); - eprintln!( - "selfplay_wait target_reached samples={} active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", - self.control.samples_collected.load(Ordering::Acquire), - self.control.active_pollers.load(Ordering::Acquire), - reached_snapshot.next_ticket, - reached_snapshot.inflight_batches, - reached_snapshot.remainder, - reached_snapshot.open_batch_writes, - ); - self.control.running.store(false, Ordering::Release); self.queue_notify.notify_all(); self.control.wake_all(); - eprintln!( - "selfplay_wait pausing target_samples={} active_pollers={}", - target_samples, - self.control.active_pollers.load(Ordering::Acquire), - ); - { let mut guard = self .control @@ -323,31 +290,12 @@ impl SelfPlaySession { } } - let drained_snapshot = self.queue_notify.debug_snapshot(); - eprintln!( - "selfplay_wait pollers_drained active_pollers={} queue_next_ticket={} inflight={} remainder={} open_writes={}", - self.control.active_pollers.load(Ordering::Acquire), - drained_snapshot.next_ticket, - drained_snapshot.inflight_batches, - drained_snapshot.remainder, - drained_snapshot.open_batch_writes, - ); - // SAFETY: `wait_for` has already stopped the session and waited until // every session thread has left the executor loop (`active_pollers == 0`), // so no thread can submit new queue writes or race another quiesce call. - let flushed = unsafe { self.queue_notify.quiesce_exclusive() }; - - let end_snapshot = self.queue_notify.debug_snapshot(); - eprintln!( - "selfplay_wait quiesce_done flushed={} samples={} queue_next_ticket={} inflight={} remainder={} open_writes={}", - flushed, - self.control.samples_collected.load(Ordering::Acquire), - end_snapshot.next_ticket, - end_snapshot.inflight_batches, - end_snapshot.remainder, - end_snapshot.open_batch_writes, - ); + unsafe { + self.queue_notify.quiesce_exclusive(); + } self.control.samples_collected.load(Ordering::Acquire) } From 49abb9a47a4bb5aa3829fc24ec77b6b4b22a4587 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 20:21:02 -0400 Subject: [PATCH 46/59] refine terminal value heuristic Blend hill control, territory margin, and depth into terminal targets and normalize territory by playable cells so training and search get a less collapse-prone value signal on wall-heavy maps. --- alpha_paint/src/board/tile_map.rs | 12 ++++++++++++ alpha_paint/src/search.rs | 25 ++++++++++++++++++++++--- training/src/descent.rs | 25 ++++++++++++++++++++++--- training/src/lib.rs | 21 ++++++++++++++++++++- 4 files changed, 76 insertions(+), 7 deletions(-) diff --git a/alpha_paint/src/board/tile_map.rs b/alpha_paint/src/board/tile_map.rs index eacdf50..8e8378a 100644 --- a/alpha_paint/src/board/tile_map.rs +++ b/alpha_paint/src/board/tile_map.rs @@ -21,6 +21,7 @@ pub struct TileMap { coverage_eval: i32, white_tiles: usize, black_tiles: usize, + wall_tiles: usize, hill_id: Arc>, hills: Arc>>, hill_metadata: Vec, @@ -37,6 +38,7 @@ impl Default for TileMap { coverage_eval: 0, white_tiles: 0, black_tiles: 0, + wall_tiles: 0, hill_id: Arc::new(Array32x32 { 0: [[u16::MAX; 32]; 32], }), @@ -126,6 +128,9 @@ impl TileMap { if old.is_owned_by::() { self.black_tiles -= 1; } + if old.is_wall() { + self.wall_tiles -= 1; + } if tile.is_owned_by::() { self.white_tiles += 1; @@ -133,6 +138,9 @@ impl TileMap { if tile.is_owned_by::() { self.black_tiles += 1; } + if tile.is_wall() { + self.wall_tiles += 1; + } if old.is_beacon() || tile.is_beacon() { if old.is_beacon_of::() && !tile.is_beacon_of::() { @@ -198,6 +206,10 @@ impl TileMap { self.coverage_eval } + pub fn wall_count(&self) -> usize { + self.wall_tiles + } + pub fn hill_metadata(&self) -> &[HillData] { &self.hill_metadata } diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index 280d59d..42f89b7 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -244,8 +244,27 @@ impl SearchNode { fn value_from_term(board: &Board, term: TerminalState) -> f32 { let sign = term.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + let progress = board.turn_count.max(1) as f32; - sign * (AVG_GAME_LENGTH / progress).ln_1p() + let depth = (AVG_GAME_LENGTH / progress).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) } fn create_child( @@ -461,8 +480,8 @@ mod tests { let loss = SearchNode::value_from_term(&board, TerminalState::Win(Player::Black)); assert!(win > 0.0); assert!(loss < 0.0); - assert!(win.abs() < 10.0); - assert!(loss.abs() < 10.0); + assert!(win.abs() <= 1.0); + assert!(loss.abs() <= 1.0); } #[test] diff --git a/training/src/descent.rs b/training/src/descent.rs index 99c4c7d..fcc01d3 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -398,8 +398,27 @@ impl SearchNode { /// Reinforcement heuristic: preserve faster-win ordering without exploding. fn value_from_term(board: &Board, term: TerminalState) -> f32 { let sign = term.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + let p = board.turn_count.max(1) as f32; - sign * (AVG_GAME_LENGTH / p).ln_1p() + let depth = (AVG_GAME_LENGTH / p).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) } async fn ubfms_iteration( @@ -593,8 +612,8 @@ mod tests { let board = board_with_turn_count(1); let value = SearchNode::value_from_term(&board, TerminalState::Win(Player::White)); - assert!(value > 6.0); - assert!(value < 7.0); + assert!(value > 0.3); + assert!(value <= 1.0); } #[test] diff --git a/training/src/lib.rs b/training/src/lib.rs index cafa906..f67ddcc 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -100,8 +100,27 @@ const PRETRAIN_SAMPLE_CHUNK: usize = 256; fn terminal_value_from_term(board: &Board, terminal: TerminalState) -> f32 { let sign = terminal.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + let p = board.turn_count.max(1) as f32; - sign * (AVG_GAME_LENGTH / p).ln_1p() + let depth = (AVG_GAME_LENGTH / p).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) } fn sample_random_action(board: &Board, rng: &mut R) -> Option { From cd9c93c329021bc7f21698d7d7763f34756a0063 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 21:12:08 -0400 Subject: [PATCH 47/59] mix in some terminal states into replay buffer --- python/alphapaint_training/train.py | 8 ++ training/src/descent.rs | 159 ++++++++++++++++++++++++++-- training/src/lib.rs | 5 + training/src/training.rs | 11 ++ training/src/worker.rs | 35 ++++-- 5 files changed, 201 insertions(+), 17 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 34d992c..e6dfa7b 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -473,6 +473,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: losses: list[float] = [] target_samples = 0 previous_samples = 0 + previous_terminal_mix_samples = 0 previous_games = 0 previous_gpu_batches = 0 previous_gpu_evals = 0 @@ -498,6 +499,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: collect_started_at = time.perf_counter() target_samples += config.samples_per_round collected = selfplay.wait_for(target_samples) + terminal_mix_samples = selfplay.terminal_mix_samples_added() games = selfplay.games() gpu_batches = selfplay.gpu_batches() gpu_evals = selfplay.gpu_evals() @@ -525,6 +527,9 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: descent_backup_nanos = selfplay.descent_backup_nanos() collect_seconds = time.perf_counter() - collect_started_at samples_added = collected - previous_samples + terminal_mix_samples_added = ( + terminal_mix_samples - previous_terminal_mix_samples + ) games_added = games - previous_games gpu_batches_added = gpu_batches - previous_gpu_batches gpu_evals_added = gpu_evals - previous_gpu_evals @@ -742,6 +747,8 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: "round": round_number, "samples_total": collected, "samples_added": samples_added, + "terminal_mix_samples_total": terminal_mix_samples, + "terminal_mix_samples_added": terminal_mix_samples_added, "games_total": games, "games_added": games_added, "gpu_batches_total": gpu_batches, @@ -836,6 +843,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print(f"checkpoint={checkpoint_path}") previous_samples = collected + previous_terminal_mix_samples = terminal_mix_samples previous_games = games previous_gpu_batches = gpu_batches previous_gpu_evals = gpu_evals diff --git a/training/src/descent.rs b/training/src/descent.rs index fcc01d3..2696b2d 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -7,6 +7,8 @@ use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, Rollback, TerminalState}; use rand::rngs::SmallRng; +#[cfg(test)] +use rand::SeedableRng; use rand::{Rng, RngExt}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; @@ -19,6 +21,7 @@ const AVG_GAME_LENGTH: f32 = 500.0; pub struct TreeLearningSample { pub board: Board, pub value: f32, + pub terminal_mix: bool, } #[derive(Clone)] @@ -252,6 +255,26 @@ impl SearchNode { board } + fn terminal_leaf_sample_board( + state: &Board, + action: Action, + rollback: Rollback, + outcome: ApplyActionOutcome, + ) -> Board { + match outcome { + ApplyActionOutcome::Terminal { .. } => state.clone(), + ApplyActionOutcome::PlayInstead { play_instead, .. } => { + Self::apply_play_instead(state.clone(), action, rollback, play_instead) + } + ApplyActionOutcome::Killshot { moves, .. } => { + Self::apply_killshot_moves(state.clone(), &moves) + } + ApplyActionOutcome::Ongoing => { + unreachable!("terminal leaf sampling only applies to resolved terminal outcomes") + } + } + } + /// Expand a node: evaluate all children with the neural network. async fn build_self( board: &Board, @@ -514,35 +537,84 @@ impl SearchNode { /// An internal node is one that has children and at least one expanded child. /// Non-terminal leaf nodes (where the network estimate was used without /// minimax backing) are excluded per Athénan's tree learning rules. - fn collect_samples(&self, state: &mut Board, out: &mut Vec) { + fn collect_samples_with_terminal_mix( + &self, + state: &mut Board, + out: &mut Vec, + rng: &mut R, + terminal_mix_prob: f32, + ) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); if has_expanded_child { out.push(TreeLearningSample { board: state.clone(), value: self.value, + terminal_mix: false, + }); + } else if self.resolved + && self.children.is_empty() + && rng.random::() <= terminal_mix_prob + { + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + terminal_mix: true, }); } for child in &self.children { if let Some(node) = child.node.as_ref() { - let (_, rollback) = state.apply_action(child.action); - node.collect_samples(state, out); + let (outcome, rollback) = state.apply_action(child.action); + if node.children.is_empty() && node.resolved { + if rng.random::() <= terminal_mix_prob { + out.push(TreeLearningSample { + board: Self::terminal_leaf_sample_board( + state, + child.action, + rollback, + outcome, + ), + value: node.value, + terminal_mix: true, + }); + } + } else { + node.collect_samples_with_terminal_mix(state, out, rng, terminal_mix_prob); + } state.rollback(child.action, rollback); } } } - fn collect_samples_excluding_child( + #[cfg(test)] + fn collect_samples(&self, state: &mut Board, out: &mut Vec) { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_samples_with_terminal_mix(state, out, &mut rng, 0.0); + } + + fn collect_samples_excluding_child_with_terminal_mix( &self, state: &mut Board, excluded_child: usize, out: &mut Vec, + rng: &mut R, + terminal_mix_prob: f32, ) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); if has_expanded_child { out.push(TreeLearningSample { board: state.clone(), value: self.value, + terminal_mix: false, + }); + } else if self.resolved + && self.children.is_empty() + && rng.random::() <= terminal_mix_prob + { + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + terminal_mix: true, }); } @@ -551,12 +623,44 @@ impl SearchNode { continue; } if let Some(node) = child.node.as_ref() { - let (_, rollback) = state.apply_action(child.action); - node.collect_samples(state, out); + let (outcome, rollback) = state.apply_action(child.action); + if node.children.is_empty() && node.resolved { + if rng.random::() <= terminal_mix_prob { + out.push(TreeLearningSample { + board: Self::terminal_leaf_sample_board( + state, + child.action, + rollback, + outcome, + ), + value: node.value, + terminal_mix: true, + }); + } + } else { + node.collect_samples_with_terminal_mix(state, out, rng, terminal_mix_prob); + } state.rollback(child.action, rollback); } } } + + #[cfg(test)] + fn collect_samples_excluding_child( + &self, + state: &mut Board, + excluded_child: usize, + out: &mut Vec, + ) { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_samples_excluding_child_with_terminal_mix( + state, + excluded_child, + out, + &mut rng, + 0.0, + ); + } } #[cfg(test)] @@ -827,22 +931,45 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { } } - pub fn collect_tree_learning_samples(&self) -> Vec { + pub fn collect_tree_learning_samples_with_terminal_mix( + &self, + rng: &mut R, + terminal_mix_prob: f32, + ) -> Vec { let mut state = self.root_state.clone(); let mut samples = Vec::new(); - self.root_node.collect_samples(&mut state, &mut samples); + self.root_node.collect_samples_with_terminal_mix( + &mut state, + &mut samples, + rng, + terminal_mix_prob, + ); samples } - pub fn step_tree_and_collect_dropped_samples( + #[cfg(test)] + pub fn collect_tree_learning_samples(&self) -> Vec { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_tree_learning_samples_with_terminal_mix(&mut rng, 0.0) + } + + pub fn step_tree_and_collect_dropped_samples_with_terminal_mix( &mut self, new_board: &Board, action_id: usize, + rng: &mut R, + terminal_mix_prob: f32, ) -> Vec { let mut state = self.root_state.clone(); let mut samples = Vec::new(); self.root_node - .collect_samples_excluding_child(&mut state, action_id, &mut samples); + .collect_samples_excluding_child_with_terminal_mix( + &mut state, + action_id, + &mut samples, + rng, + terminal_mix_prob, + ); self.root_state = new_board.clone(); @@ -857,6 +984,18 @@ impl<'a, E: Evaluator> GameSearchTree<'a, E> { samples } + #[cfg(test)] + pub fn step_tree_and_collect_dropped_samples( + &mut self, + new_board: &Board, + action_id: usize, + ) -> Vec { + let mut rng = SmallRng::seed_from_u64(0); + self.step_tree_and_collect_dropped_samples_with_terminal_mix( + new_board, action_id, &mut rng, 0.0, + ) + } + pub async fn run_descent_for_iter(&mut self, iterations: u32) { let dcv: i32 = if self.root_state.is_white_turn() { 1 diff --git a/training/src/lib.rs b/training/src/lib.rs index f67ddcc..86b8dd9 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -400,6 +400,11 @@ impl SelfPlay { Ok(self.session()?.games()) } + /// Return the total number of terminal leaf samples injected into replay. + fn terminal_mix_samples_added(&self) -> PyResult { + Ok(self.session()?.terminal_mix_samples_added()) + } + /// Return the total number of selected actions. fn action_steps(&self) -> PyResult { Ok(self.session()?.action_steps()) diff --git a/training/src/training.rs b/training/src/training.rs index 8e3c472..09c0262 100644 --- a/training/src/training.rs +++ b/training/src/training.rs @@ -29,6 +29,8 @@ struct SessionControl { target_samples: AtomicUsize, /// Total samples collected (monotonic across the session lifetime). samples_collected: Arc, + /// Total terminal leaf samples injected into replay via terminal mixing. + terminal_mix_samples_added: Arc, /// Total games completed. games_completed: Arc, /// Total actions selected across all workers. @@ -85,6 +87,7 @@ impl SessionControl { shutdown: AtomicBool::new(false), target_samples: AtomicUsize::new(0), samples_collected: Arc::new(AtomicUsize::new(0)), + terminal_mix_samples_added: Arc::new(AtomicU64::new(0)), games_completed: Arc::new(AtomicUsize::new(0)), action_steps: Arc::new(AtomicU64::new(0)), final_actions: Arc::new(AtomicU64::new(0)), @@ -310,6 +313,13 @@ impl SelfPlaySession { self.control.games_completed.load(Ordering::Acquire) } + /// Return the total number of terminal leaf samples injected into replay. + pub fn terminal_mix_samples_added(&self) -> u64 { + self.control + .terminal_mix_samples_added + .load(Ordering::Acquire) + } + /// Return the total number of selected actions. pub fn action_steps(&self) -> u64 { self.control.action_steps.load(Ordering::Acquire) @@ -465,6 +475,7 @@ fn session_thread_main( let metrics = SelfPlayMetrics { samples_collected: control.samples_collected.clone(), + terminal_mix_samples_added: control.terminal_mix_samples_added.clone(), games_completed: control.games_completed.clone(), action_steps: control.action_steps.clone(), final_actions: control.final_actions.clone(), diff --git a/training/src/worker.rs b/training/src/worker.rs index 92293cb..e91184b 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -19,6 +19,8 @@ use crate::eval::{CountingEvaluator, Evaluator}; use crate::observation; use crate::replay_buffer::ReplayBuffer; +const RANDOM_TERMINAL_MIX_PROB: f32 = 0.2; + fn update_max_usize(max_value: &AtomicUsize, candidate: usize) { let mut current = max_value.load(Ordering::Acquire); while candidate > current { @@ -56,6 +58,7 @@ fn elapsed_nanos(started_at: Instant) -> u64 { #[derive(Clone)] pub struct SelfPlayMetrics { pub samples_collected: Arc, + pub terminal_mix_samples_added: Arc, pub games_completed: Arc, pub action_steps: Arc, pub final_actions: Arc, @@ -104,6 +107,8 @@ fn push_samples_to_replay( return; } + let terminal_mix_count = samples.iter().filter(|sample| sample.terminal_mix).count(); + let started_at = Instant::now(); let mut guard = replay_buffer.reserve(num_samples); @@ -119,31 +124,42 @@ fn push_samples_to_replay( metrics .samples_collected .fetch_add(num_samples, Ordering::AcqRel); + metrics + .terminal_mix_samples_added + .fetch_add(terminal_mix_count as u64, Ordering::AcqRel); metrics .replay_push_nanos .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); } -fn collect_tree_learning_samples_timed( +fn collect_tree_learning_samples_timed( tree: &GameSearchTree<'_, E>, metrics: &SelfPlayMetrics, + rng: &mut R, ) -> Vec { let started_at = Instant::now(); - let samples = tree.collect_tree_learning_samples(); + let samples = + tree.collect_tree_learning_samples_with_terminal_mix(rng, RANDOM_TERMINAL_MIX_PROB); metrics .sample_collect_nanos .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); samples } -fn step_tree_and_collect_dropped_samples_timed( +fn step_tree_and_collect_dropped_samples_timed( tree: &mut GameSearchTree<'_, E>, new_board: &Board, action_id: usize, metrics: &SelfPlayMetrics, + rng: &mut R, ) -> Vec { let started_at = Instant::now(); - let samples = tree.step_tree_and_collect_dropped_samples(new_board, action_id); + let samples = tree.step_tree_and_collect_dropped_samples_with_terminal_mix( + new_board, + action_id, + rng, + RANDOM_TERMINAL_MIX_PROB, + ); metrics .sample_collect_nanos .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); @@ -154,13 +170,14 @@ fn finish_game_with_current_tree( tree: &GameSearchTree<'_, E>, replay_buffer: &ReplayBuffer, metrics: &SelfPlayMetrics, + rng: &mut impl Rng, action_steps_in_game: u64, completed_turn_count_in_game: usize, ) { push_samples_to_replay( replay_buffer, metrics, - collect_tree_learning_samples_timed(tree, metrics), + collect_tree_learning_samples_timed(tree, metrics, rng), ); metrics .completed_game_actions_total @@ -248,6 +265,7 @@ async fn play_game( &tree, replay_buffer, metrics, + rng, action_steps_in_game, tree.root_state.turn_count, ); @@ -277,6 +295,7 @@ async fn play_game( &tree, replay_buffer, metrics, + rng, action_steps_in_game, new_board.turn_count, ); @@ -287,6 +306,7 @@ async fn play_game( &tree, replay_buffer, metrics, + rng, action_steps_in_game, tree.root_state.turn_count + 1, ); @@ -297,6 +317,7 @@ async fn play_game( &tree, replay_buffer, metrics, + rng, action_steps_in_game, tree.root_state.turn_count + 1, ); @@ -305,14 +326,14 @@ async fn play_game( alpha_paint::board::ApplyActionOutcome::Ongoing => { if tree.root_node.children[action_id].node.is_some() { let samples = step_tree_and_collect_dropped_samples_timed( - &mut tree, &new_board, action_id, metrics, + &mut tree, &new_board, action_id, metrics, rng, ); push_samples_to_replay(replay_buffer, metrics, samples); } else { push_samples_to_replay( replay_buffer, metrics, - collect_tree_learning_samples_timed(&tree, metrics), + collect_tree_learning_samples_timed(&tree, metrics, rng), ); let next_rng = SmallRng::from_rng(rng); let build_started_at = Instant::now(); From 5f6d7fc1c23dd54ae8eda6bb3e542f1304772fac Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Sun, 29 Mar 2026 21:47:20 -0400 Subject: [PATCH 48/59] fix bug in terminal state mixing --- training/src/descent.rs | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/training/src/descent.rs b/training/src/descent.rs index 2696b2d..d752342 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -553,6 +553,7 @@ impl SearchNode { }); } else if self.resolved && self.children.is_empty() + && state.get_valid_actions().len() == 0 && rng.random::() <= terminal_mix_prob { out.push(TreeLearningSample { @@ -565,7 +566,13 @@ impl SearchNode { for child in &self.children { if let Some(node) = child.node.as_ref() { let (outcome, rollback) = state.apply_action(child.action); - if node.children.is_empty() && node.resolved { + if node.children.is_empty() + && matches!( + outcome, + ApplyActionOutcome::Terminal { .. } + | ApplyActionOutcome::PlayInstead { .. } + ) + { if rng.random::() <= terminal_mix_prob { out.push(TreeLearningSample { board: Self::terminal_leaf_sample_board( @@ -609,6 +616,7 @@ impl SearchNode { }); } else if self.resolved && self.children.is_empty() + && state.get_valid_actions().len() == 0 && rng.random::() <= terminal_mix_prob { out.push(TreeLearningSample { @@ -624,7 +632,13 @@ impl SearchNode { } if let Some(node) = child.node.as_ref() { let (outcome, rollback) = state.apply_action(child.action); - if node.children.is_empty() && node.resolved { + if node.children.is_empty() + && matches!( + outcome, + ApplyActionOutcome::Terminal { .. } + | ApplyActionOutcome::PlayInstead { .. } + ) + { if rng.random::() <= terminal_mix_prob { out.push(TreeLearningSample { board: Self::terminal_leaf_sample_board( From daf39f5c90a132ba4bb8702f03eeca90fce633cf Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 00:15:50 -0400 Subject: [PATCH 49/59] simplify logging --- python/alphapaint_training/__init__.py | 2 + python/alphapaint_training/logger.py | 494 +++++++++++++++++++++++++ python/alphapaint_training/train.py | 398 +++----------------- training/src/worker.rs | 4 - 4 files changed, 538 insertions(+), 360 deletions(-) create mode 100644 python/alphapaint_training/logger.py diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 6715edf..81d638c 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -9,6 +9,7 @@ SelfPlay, sample_random_terminal_batch, ) +from .logger import TrainingLogger from .model import PackedValueModel, ResidualBlock, TinyValueNet from .packed_obs import ( BOARD_CELLS, @@ -33,6 +34,7 @@ "ResidualBlock", "SelfPlay", "TinyValueNet", + "TrainingLogger", "decode_intrinsics", "decode_packed_board", "decode_packed_board_reference", diff --git a/python/alphapaint_training/logger.py b/python/alphapaint_training/logger.py new file mode 100644 index 0000000..2ae46b6 --- /dev/null +++ b/python/alphapaint_training/logger.py @@ -0,0 +1,494 @@ +"""Training logger for AlphaPaint - handles wandb and console output.""" + +from __future__ import annotations + +import time +from dataclasses import asdict +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +import wandb + +if TYPE_CHECKING: + from alphapaint_training.train import TrainConfig + from alphapaint_training.model import PackedValueModel + + +class TrainingLogger: + """Handles all logging for training runs including wandb and console output.""" + + def __init__( + self, + config: TrainConfig, + run_dir: Path, + checkpoint_dir: Path, + model: PackedValueModel, + param_count: int, + enable_wandb: bool = True, + ) -> None: + self.config = config + self.run_dir = run_dir + self.checkpoint_dir = checkpoint_dir + self.enable_wandb = enable_wandb + self.run = None + + if enable_wandb: + self.run = wandb.init( + project="alphapaint", + name=run_dir.name, + dir=run_dir, + config=asdict(config), + ) + self.run.summary["model_params"] = param_count + + self._pretrain_loss_window: list[float] = [] + self._pretrain_started_at: float | None = None + self._previous_samples = 0 + self._previous_terminal_mix_samples = 0 + self._previous_games = 0 + self._previous_gpu_batches = 0 + self._previous_gpu_evals = 0 + self._previous_action_steps = 0 + self._previous_final_actions = 0 + self._previous_nonfinal_actions = 0 + self._previous_completed_turns = 0 + self._previous_action_turn_count_total = 0 + self._previous_completed_game_actions_total = 0 + self._previous_completed_game_turn_count_total = 0 + self._previous_tree_build_nanos = 0 + self._previous_descent_nanos = 0 + self._previous_sample_collect_nanos = 0 + self._previous_replay_push_nanos = 0 + self._previous_descent_expand_cpu_nanos = 0 + self._previous_descent_apply_action_nanos = 0 + self._previous_descent_eval_submit_nanos = 0 + self._previous_descent_eval_await_nanos = 0 + self._previous_descent_backup_nanos = 0 + + def start_pretrain(self) -> None: + """Call at the start of pretraining.""" + self._pretrain_started_at = time.perf_counter() + self._pretrain_loss_window.clear() + print( + f"pretrain samples={self.config.pretrain_terminal_samples} " + f"batch={self.config.pretrain_batch_size}" + ) + + def log_pretrain_step( + self, + step_idx: int, + step_batch_size: int, + loss_value: float, + pretrain_steps: int, + device: torch.device, + ) -> None: + """Log a single pretrain step.""" + self._pretrain_loss_window.append(loss_value) + pretrain_log_interval = max(1, self.config.pretrain_log_interval) + if len(self._pretrain_loss_window) > pretrain_log_interval: + self._pretrain_loss_window.pop(0) + + step_number = step_idx + 1 + if step_number % pretrain_log_interval == 0 or step_number == pretrain_steps: + if device.type == "cuda": + torch.cuda.synchronize(device) + samples_done = step_idx * self.config.pretrain_batch_size + step_batch_size + elapsed = ( + time.perf_counter() - self._pretrain_started_at + if self._pretrain_started_at + else 0.0 + ) + window_mean_loss = float(np.mean(self._pretrain_loss_window)) + print( + f"pretrain step={step_number}/{pretrain_steps} " + f"samples={samples_done}/{self.config.pretrain_terminal_samples} " + f"loss={window_mean_loss:.6f} elapsed={elapsed:.2f}s" + ) + if self.run is not None: + wandb.log( + { + "pretrain_step": step_number, + "pretrain_samples_total": samples_done, + "pretrain_loss_mean": window_mean_loss, + "pretrain_seconds": elapsed, + "learning_rate": self._get_current_lr(), + } + ) + + def finish_pretrain(self, pretrain_steps: int, device: torch.device) -> None: + """Call at the end of pretraining.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + elapsed = ( + time.perf_counter() - self._pretrain_started_at + if self._pretrain_started_at + else 0.0 + ) + print( + f"pretrain_complete samples={self.config.pretrain_terminal_samples} " + f"steps={pretrain_steps} elapsed={elapsed:.2f}s" + ) + + def _get_current_lr(self) -> float: + """Get current learning rate from optimizer (stored externally).""" + return 0.0 + + def set_lr_getter(self, getter: Any) -> None: + """Set a function to get the current learning rate.""" + self._get_current_lr = getter + + def update( + self, + selfplay: Any, + round_number: int, + collected: int, + replay_size: int, + mean_loss: float, + collect_seconds: float, + train_seconds: float, + train_sample_seconds: float, + train_h2d_seconds: float, + train_forward_seconds: float, + train_backward_seconds: float, + train_optimizer_seconds: float, + diag_stats: dict[str, float], + diag_histograms: dict[str, wandb.Histogram], + ) -> dict[str, object]: + """Update logs for a completed round. Returns the record dict for checkpointing.""" + # Extract all metrics from selfplay + terminal_mix_samples = selfplay.terminal_mix_samples_added() + games = selfplay.games() + gpu_batches = selfplay.gpu_batches() + gpu_evals = selfplay.gpu_evals() + action_steps = selfplay.action_steps() + final_actions = selfplay.final_actions() + nonfinal_actions = selfplay.nonfinal_actions() + completed_turns = selfplay.completed_turns() + action_turn_count_total = selfplay.action_turn_count_total() + max_turn_count_seen = selfplay.max_turn_count_seen() + completed_game_actions_total = selfplay.completed_game_actions_total() + max_turn_count_in_completed_game = selfplay.max_turn_count_in_completed_game() + completed_game_turn_counts = selfplay.take_completed_game_turn_counts() + completed_game_turn_count_total = selfplay.completed_game_turn_count_total() + max_actions_in_completed_game = selfplay.max_actions_in_completed_game() + tree_build_nanos = selfplay.tree_build_nanos() + descent_nanos = selfplay.descent_nanos() + sample_collect_nanos = selfplay.sample_collect_nanos() + replay_push_nanos = selfplay.replay_push_nanos() + descent_expand_cpu_nanos = selfplay.descent_expand_cpu_nanos() + descent_apply_action_nanos = selfplay.descent_apply_action_nanos() + descent_eval_submit_nanos = selfplay.descent_eval_submit_nanos() + descent_eval_await_nanos = selfplay.descent_eval_await_nanos() + descent_backup_nanos = selfplay.descent_backup_nanos() + + # Calculate deltas + samples_added = collected - self._previous_samples + terminal_mix_samples_added = ( + terminal_mix_samples - self._previous_terminal_mix_samples + ) + games_added = games - self._previous_games + gpu_batches_added = gpu_batches - self._previous_gpu_batches + gpu_evals_added = gpu_evals - self._previous_gpu_evals + action_steps_added = action_steps - self._previous_action_steps + final_actions_added = final_actions - self._previous_final_actions + nonfinal_actions_added = nonfinal_actions - self._previous_nonfinal_actions + completed_turns_added = completed_turns - self._previous_completed_turns + action_turn_count_added = ( + action_turn_count_total - self._previous_action_turn_count_total + ) + completed_game_actions_added = ( + completed_game_actions_total - self._previous_completed_game_actions_total + ) + completed_game_turn_count_added = ( + completed_game_turn_count_total + - self._previous_completed_game_turn_count_total + ) + + tree_build_seconds = (tree_build_nanos - self._previous_tree_build_nanos) / 1e9 + descent_seconds = (descent_nanos - self._previous_descent_nanos) / 1e9 + sample_collect_seconds = ( + sample_collect_nanos - self._previous_sample_collect_nanos + ) / 1e9 + replay_push_seconds = ( + replay_push_nanos - self._previous_replay_push_nanos + ) / 1e9 + descent_expand_cpu_seconds = ( + descent_expand_cpu_nanos - self._previous_descent_expand_cpu_nanos + ) / 1e9 + descent_apply_action_seconds = ( + descent_apply_action_nanos - self._previous_descent_apply_action_nanos + ) / 1e9 + descent_eval_submit_seconds = ( + descent_eval_submit_nanos - self._previous_descent_eval_submit_nanos + ) / 1e9 + descent_eval_await_seconds = ( + descent_eval_await_nanos - self._previous_descent_eval_await_nanos + ) / 1e9 + descent_backup_seconds = ( + descent_backup_nanos - self._previous_descent_backup_nanos + ) / 1e9 + descent_other_seconds = max( + 0.0, + descent_seconds + - descent_expand_cpu_seconds + - descent_apply_action_seconds + - descent_eval_submit_seconds + - descent_eval_await_seconds + - descent_backup_seconds, + ) + + avg_turn_count = ( + action_turn_count_added / action_steps_added if action_steps_added else 0.0 + ) + actions_per_turn = ( + action_steps_added / completed_turns_added if completed_turns_added else 0.0 + ) + actions_per_game = ( + completed_game_actions_added / games_added if games_added else 0.0 + ) + + if completed_game_turn_counts: + completed_game_turn_counts_arr = np.asarray( + completed_game_turn_counts, dtype=np.float32 + ) + completed_game_turn_count_mean = float( + completed_game_turn_counts_arr.mean() + ) + completed_game_turn_count_p50 = float( + np.percentile(completed_game_turn_counts_arr, 50) + ) + completed_game_turn_count_p90 = float( + np.percentile(completed_game_turn_counts_arr, 90) + ) + else: + completed_game_turn_count_mean = float("nan") + completed_game_turn_count_p50 = float("nan") + completed_game_turn_count_p90 = float("nan") + + build_ms_per_action = ( + tree_build_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_ms_per_action = ( + descent_seconds * 1000.0 / action_steps_added if action_steps_added else 0.0 + ) + collect_ms_per_action = ( + sample_collect_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + push_ms_per_action = ( + replay_push_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_expand_ms_per_action = ( + descent_expand_cpu_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_apply_ms_per_action = ( + descent_apply_action_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_ms_per_action = ( + descent_eval_submit_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_wait_ms_per_action = ( + descent_eval_await_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_backup_ms_per_action = ( + descent_backup_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_other_ms_per_action = ( + descent_other_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_us_per_eval = ( + descent_eval_submit_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) + descent_wait_us_per_eval = ( + descent_eval_await_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) + + record: dict[str, Any] = { + "round": round_number, + "samples_total": collected, + "samples_added": samples_added, + "terminal_mix_samples_total": terminal_mix_samples, + "terminal_mix_samples_added": terminal_mix_samples_added, + "games_total": games, + "games_added": games_added, + "gpu_batches_total": gpu_batches, + "gpu_batches_added": gpu_batches_added, + "gpu_batches_per_second": gpu_batches_added / max(collect_seconds, 1e-9), + "gpu_evals_total": gpu_evals, + "gpu_evals_added": gpu_evals_added, + "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), + "action_steps_total": action_steps, + "action_steps_added": action_steps_added, + "final_actions_total": final_actions, + "final_actions_added": final_actions_added, + "nonfinal_actions_total": nonfinal_actions, + "nonfinal_actions_added": nonfinal_actions_added, + "completed_turns_total": completed_turns, + "completed_turns_added": completed_turns_added, + "action_turn_count_total": action_turn_count_total, + "action_turn_count_added": action_turn_count_added, + "avg_turn_count": avg_turn_count, + "max_turn_count_seen": max_turn_count_seen, + "actions_per_turn": actions_per_turn, + "completed_game_actions_total": completed_game_actions_total, + "completed_game_actions_added": completed_game_actions_added, + "actions_per_completed_game": actions_per_game, + "max_actions_in_completed_game": max_actions_in_completed_game, + "completed_game_turn_count_total": completed_game_turn_count_total, + "completed_game_turn_count_added": completed_game_turn_count_added, + "completed_game_turn_count_mean": completed_game_turn_count_mean, + "completed_game_turn_count_p50": completed_game_turn_count_p50, + "completed_game_turn_count_p90": completed_game_turn_count_p90, + "max_turn_count_in_completed_game": max_turn_count_in_completed_game, + "replay_size": replay_size, + "collection_seconds": collect_seconds, + "training_seconds": train_seconds, + "samples_per_second": samples_added / max(collect_seconds, 1e-9), + "train_steps_per_second": self.config.train_steps_per_round + / max(train_seconds, 1e-9), + "selfplay_tree_build_seconds": tree_build_seconds, + "selfplay_descent_seconds": descent_seconds, + "selfplay_sample_collect_seconds": sample_collect_seconds, + "selfplay_replay_push_seconds": replay_push_seconds, + "selfplay_tree_build_ms_per_action": build_ms_per_action, + "selfplay_descent_ms_per_action": descent_ms_per_action, + "selfplay_sample_collect_ms_per_action": collect_ms_per_action, + "selfplay_replay_push_ms_per_action": push_ms_per_action, + "descent_expand_cpu_seconds": descent_expand_cpu_seconds, + "descent_apply_action_seconds": descent_apply_action_seconds, + "descent_eval_submit_seconds": descent_eval_submit_seconds, + "descent_eval_await_seconds": descent_eval_await_seconds, + "descent_backup_seconds": descent_backup_seconds, + "descent_other_seconds": descent_other_seconds, + "descent_expand_cpu_ms_per_action": descent_expand_ms_per_action, + "descent_apply_action_ms_per_action": descent_apply_ms_per_action, + "descent_eval_submit_ms_per_action": descent_submit_ms_per_action, + "descent_eval_await_ms_per_action": descent_wait_ms_per_action, + "descent_backup_ms_per_action": descent_backup_ms_per_action, + "descent_other_ms_per_action": descent_other_ms_per_action, + "descent_eval_submit_us_per_eval": descent_submit_us_per_eval, + "descent_eval_await_us_per_eval": descent_wait_us_per_eval, + "train_sample_seconds": train_sample_seconds, + "train_h2d_seconds": train_h2d_seconds, + "train_forward_seconds": train_forward_seconds, + "train_backward_seconds": train_backward_seconds, + "train_optimizer_seconds": train_optimizer_seconds, + "loss_mean": mean_loss, + "learning_rate": self._get_current_lr(), + "timestamp": time.time(), + **diag_stats, + } + + print( + f"round={round_number} samples={collected} games={games} " + f"gpu_evals/s={gpu_evals_added / max(collect_seconds, 1e-9):.1f} " + f"replay={replay_size} loss={mean_loss:.6f}" + ) + + if self.run is not None: + wandb.log({**record, **diag_histograms}) + + self._update_previous_values( + collected, + terminal_mix_samples, + games, + gpu_batches, + gpu_evals, + action_steps, + final_actions, + nonfinal_actions, + completed_turns, + action_turn_count_total, + completed_game_actions_total, + completed_game_turn_count_total, + tree_build_nanos, + descent_nanos, + sample_collect_nanos, + replay_push_nanos, + descent_expand_cpu_nanos, + descent_apply_action_nanos, + descent_eval_submit_nanos, + descent_eval_await_nanos, + descent_backup_nanos, + ) + + return record + + def _update_previous_values( + self, + collected: int, + terminal_mix_samples: int, + games: int, + gpu_batches: int, + gpu_evals: int, + action_steps: int, + final_actions: int, + nonfinal_actions: int, + completed_turns: int, + action_turn_count_total: int, + completed_game_actions_total: int, + completed_game_turn_count_total: int, + tree_build_nanos: int, + descent_nanos: int, + sample_collect_nanos: int, + replay_push_nanos: int, + descent_expand_cpu_nanos: int, + descent_apply_action_nanos: int, + descent_eval_submit_nanos: int, + descent_eval_await_nanos: int, + descent_backup_nanos: int, + ) -> None: + """Update previous values for delta calculations.""" + self._previous_samples = collected + self._previous_terminal_mix_samples = terminal_mix_samples + self._previous_games = games + self._previous_gpu_batches = gpu_batches + self._previous_gpu_evals = gpu_evals + self._previous_action_steps = action_steps + self._previous_final_actions = final_actions + self._previous_nonfinal_actions = nonfinal_actions + self._previous_completed_turns = completed_turns + self._previous_action_turn_count_total = action_turn_count_total + self._previous_completed_game_actions_total = completed_game_actions_total + self._previous_completed_game_turn_count_total = completed_game_turn_count_total + self._previous_tree_build_nanos = tree_build_nanos + self._previous_descent_nanos = descent_nanos + self._previous_sample_collect_nanos = sample_collect_nanos + self._previous_replay_push_nanos = replay_push_nanos + self._previous_descent_expand_cpu_nanos = descent_expand_cpu_nanos + self._previous_descent_apply_action_nanos = descent_apply_action_nanos + self._previous_descent_eval_submit_nanos = descent_eval_submit_nanos + self._previous_descent_eval_await_nanos = descent_eval_await_nanos + self._previous_descent_backup_nanos = descent_backup_nanos + + def log_checkpoint(self, checkpoint_path: Path) -> None: + """Log a checkpoint save.""" + print(f"checkpoint={checkpoint_path}") + + def finish(self) -> None: + """Clean up logging.""" + if self.run is not None: + self.run.finish() diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index e6dfa7b..324dfdb 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -18,6 +18,7 @@ SelfPlay, sample_random_terminal_batch, ) +from alphapaint_training.logger import TrainingLogger from alphapaint_training.model import PackedValueModel @@ -365,12 +366,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: run_dir, checkpoint_dir = _prepare_run_dir(config) - run = None - if config.wandb: - run = wandb.init( - project="alphapaint", name=run_dir.name, dir=run_dir, config=asdict(config) - ) - model = cast( PackedValueModel, _to_channels_last( @@ -386,8 +381,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print( f"model width={config.width} blocks={config.num_blocks} hidden={config.hidden_dim} params={param_count}" ) - if run is not None: - run.summary["model_params"] = param_count + initial_lr = pretrain_lr if config.pretrain_terminal_samples > 0 else selfplay_lr optimizer = torch.optim.AdamW( model.parameters(), @@ -395,16 +389,22 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: weight_decay=config.weight_decay, ) + logger = TrainingLogger( + config=config, + run_dir=run_dir, + checkpoint_dir=checkpoint_dir, + model=model, + param_count=param_count, + enable_wandb=config.wandb, + ) + logger.set_lr_getter(lambda: optimizer.param_groups[0]["lr"]) + if config.pretrain_terminal_samples > 0: pretrain_steps = ( config.pretrain_terminal_samples + config.pretrain_batch_size - 1 ) // config.pretrain_batch_size pretrain_log_interval = max(1, config.pretrain_log_interval) - pretrain_started_at = time.perf_counter() - pretrain_loss_window: list[float] = [] - print( - f"pretrain samples={config.pretrain_terminal_samples} batch={config.pretrain_batch_size} steps={pretrain_steps}" - ) + logger.start_pretrain() for step_idx in range(pretrain_steps): samples_done = step_idx * config.pretrain_batch_size step_batch_size = min( @@ -419,42 +419,14 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device=device, ) loss_value = float(result.loss.float().cpu().item()) - pretrain_loss_window.append(loss_value) - if len(pretrain_loss_window) > pretrain_log_interval: - pretrain_loss_window.pop(0) - - step_number = step_idx + 1 - if ( - step_number % pretrain_log_interval == 0 - or step_number == pretrain_steps - ): - if device.type == "cuda": - torch.cuda.synchronize(device) - samples_done += step_batch_size - elapsed = time.perf_counter() - pretrain_started_at - window_mean_loss = float(np.mean(pretrain_loss_window)) - print( - f"pretrain step={step_number}/{pretrain_steps} samples={samples_done}/{config.pretrain_terminal_samples} " - f"loss={window_mean_loss:.6f} elapsed={elapsed:.2f}s" - ) - if run is not None: - wandb.log( - { - "pretrain_step": step_number, - "pretrain_samples_total": samples_done, - "pretrain_loss_mean": window_mean_loss, - "pretrain_seconds": elapsed, - "learning_rate": optimizer.param_groups[0]["lr"], - } - ) - - if device.type == "cuda": - torch.cuda.synchronize(device) - pretrain_seconds = time.perf_counter() - pretrain_started_at - model.eval() - print( - f"pretrain_complete samples={config.pretrain_terminal_samples} steps={pretrain_steps} elapsed={pretrain_seconds:.2f}s" - ) + logger.log_pretrain_step( + step_idx=step_idx, + step_batch_size=step_batch_size, + loss_value=loss_value, + pretrain_steps=pretrain_steps, + device=device, + ) + logger.finish_pretrain(pretrain_steps=pretrain_steps, device=device) if selfplay_lr != pretrain_lr: _set_optimizer_lr(optimizer, selfplay_lr) print(f"selfplay_lr={selfplay_lr:.6g}") @@ -472,212 +444,13 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: losses: list[float] = [] target_samples = 0 - previous_samples = 0 - previous_terminal_mix_samples = 0 - previous_games = 0 - previous_gpu_batches = 0 - previous_gpu_evals = 0 - previous_action_steps = 0 - previous_final_actions = 0 - previous_nonfinal_actions = 0 - previous_completed_turns = 0 - previous_action_turn_count_total = 0 - previous_completed_game_actions_total = 0 - previous_completed_game_turn_count_total = 0 - previous_tree_build_nanos = 0 - previous_descent_nanos = 0 - previous_sample_collect_nanos = 0 - previous_replay_push_nanos = 0 - previous_descent_expand_cpu_nanos = 0 - previous_descent_apply_action_nanos = 0 - previous_descent_eval_submit_nanos = 0 - previous_descent_eval_await_nanos = 0 - previous_descent_backup_nanos = 0 try: for round_idx in range(config.rounds): round_number = round_idx + 1 collect_started_at = time.perf_counter() target_samples += config.samples_per_round collected = selfplay.wait_for(target_samples) - terminal_mix_samples = selfplay.terminal_mix_samples_added() - games = selfplay.games() - gpu_batches = selfplay.gpu_batches() - gpu_evals = selfplay.gpu_evals() - action_steps = selfplay.action_steps() - final_actions = selfplay.final_actions() - nonfinal_actions = selfplay.nonfinal_actions() - completed_turns = selfplay.completed_turns() - action_turn_count_total = selfplay.action_turn_count_total() - max_turn_count_seen = selfplay.max_turn_count_seen() - completed_game_actions_total = selfplay.completed_game_actions_total() - completed_game_turn_count_total = selfplay.completed_game_turn_count_total() - max_turn_count_in_completed_game = ( - selfplay.max_turn_count_in_completed_game() - ) - completed_game_turn_counts = selfplay.take_completed_game_turn_counts() - max_actions_in_completed_game = selfplay.max_actions_in_completed_game() - tree_build_nanos = selfplay.tree_build_nanos() - descent_nanos = selfplay.descent_nanos() - sample_collect_nanos = selfplay.sample_collect_nanos() - replay_push_nanos = selfplay.replay_push_nanos() - descent_expand_cpu_nanos = selfplay.descent_expand_cpu_nanos() - descent_apply_action_nanos = selfplay.descent_apply_action_nanos() - descent_eval_submit_nanos = selfplay.descent_eval_submit_nanos() - descent_eval_await_nanos = selfplay.descent_eval_await_nanos() - descent_backup_nanos = selfplay.descent_backup_nanos() collect_seconds = time.perf_counter() - collect_started_at - samples_added = collected - previous_samples - terminal_mix_samples_added = ( - terminal_mix_samples - previous_terminal_mix_samples - ) - games_added = games - previous_games - gpu_batches_added = gpu_batches - previous_gpu_batches - gpu_evals_added = gpu_evals - previous_gpu_evals - action_steps_added = action_steps - previous_action_steps - final_actions_added = final_actions - previous_final_actions - nonfinal_actions_added = nonfinal_actions - previous_nonfinal_actions - completed_turns_added = completed_turns - previous_completed_turns - action_turn_count_added = ( - action_turn_count_total - previous_action_turn_count_total - ) - completed_game_actions_added = ( - completed_game_actions_total - previous_completed_game_actions_total - ) - completed_game_turn_count_added = ( - completed_game_turn_count_total - - previous_completed_game_turn_count_total - ) - tree_build_seconds = (tree_build_nanos - previous_tree_build_nanos) / 1e9 - descent_seconds = (descent_nanos - previous_descent_nanos) / 1e9 - sample_collect_seconds = ( - sample_collect_nanos - previous_sample_collect_nanos - ) / 1e9 - replay_push_seconds = (replay_push_nanos - previous_replay_push_nanos) / 1e9 - descent_expand_cpu_seconds = ( - descent_expand_cpu_nanos - previous_descent_expand_cpu_nanos - ) / 1e9 - descent_apply_action_seconds = ( - descent_apply_action_nanos - previous_descent_apply_action_nanos - ) / 1e9 - descent_eval_submit_seconds = ( - descent_eval_submit_nanos - previous_descent_eval_submit_nanos - ) / 1e9 - descent_eval_await_seconds = ( - descent_eval_await_nanos - previous_descent_eval_await_nanos - ) / 1e9 - descent_backup_seconds = ( - descent_backup_nanos - previous_descent_backup_nanos - ) / 1e9 - descent_other_seconds = max( - 0.0, - descent_seconds - - descent_expand_cpu_seconds - - descent_apply_action_seconds - - descent_eval_submit_seconds - - descent_eval_await_seconds - - descent_backup_seconds, - ) - avg_turn_count = ( - action_turn_count_added / action_steps_added - if action_steps_added - else 0.0 - ) - actions_per_turn = ( - action_steps_added / completed_turns_added - if completed_turns_added - else 0.0 - ) - actions_per_game = ( - completed_game_actions_added / games_added if games_added else 0.0 - ) - actions_per_turn_display = ( - f"{actions_per_turn:.2f}" if completed_turns_added else "-" - ) - actions_per_game_display = f"{actions_per_game:.1f}" if games_added else "-" - if completed_game_turn_counts: - completed_game_turn_counts_arr = np.asarray( - completed_game_turn_counts, dtype=np.float32 - ) - completed_game_turn_count_mean = float( - completed_game_turn_counts_arr.mean() - ) - completed_game_turn_count_p50 = float( - np.percentile(completed_game_turn_counts_arr, 50) - ) - completed_game_turn_count_p90 = float( - np.percentile(completed_game_turn_counts_arr, 90) - ) - completed_game_turn_count_display = ( - f"mean:{completed_game_turn_count_mean:.1f} " - f"p50:{completed_game_turn_count_p50:.1f} " - f"p90:{completed_game_turn_count_p90:.1f} " - f"max:{max_turn_count_in_completed_game}" - ) - else: - completed_game_turn_count_mean = float("nan") - completed_game_turn_count_p50 = float("nan") - completed_game_turn_count_p90 = float("nan") - completed_game_turn_count_display = "-" - build_ms_per_action = ( - tree_build_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_ms_per_action = ( - descent_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - collect_ms_per_action = ( - sample_collect_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - push_ms_per_action = ( - replay_push_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_expand_ms_per_action = ( - descent_expand_cpu_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_apply_ms_per_action = ( - descent_apply_action_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_submit_ms_per_action = ( - descent_eval_submit_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_wait_ms_per_action = ( - descent_eval_await_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_backup_ms_per_action = ( - descent_backup_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_other_ms_per_action = ( - descent_other_seconds * 1000.0 / action_steps_added - if action_steps_added - else 0.0 - ) - descent_submit_us_per_eval = ( - descent_eval_submit_seconds * 1e6 / gpu_evals_added - if gpu_evals_added - else 0.0 - ) - descent_wait_us_per_eval = ( - descent_eval_await_seconds * 1e6 / gpu_evals_added - if gpu_evals_added - else 0.0 - ) round_results: list[TrainStepResult] = [] train_started_at = time.perf_counter() @@ -696,7 +469,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: model.eval() train_seconds = time.perf_counter() - train_started_at - replay_size = len(replay_buffer) train_sample_seconds = sum( result.sample_seconds for result in round_results ) @@ -736,6 +508,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: mean_loss = float(round_loss_values.mean().item()) else: mean_loss = 0.0 + diag_stats, diag_histograms = _collect_round_diagnostics( model, replay_buffer, @@ -743,87 +516,23 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: seed=config.seed + round_idx * 10_000 + config.train_steps_per_round, device=device, ) - record = { - "round": round_number, - "samples_total": collected, - "samples_added": samples_added, - "terminal_mix_samples_total": terminal_mix_samples, - "terminal_mix_samples_added": terminal_mix_samples_added, - "games_total": games, - "games_added": games_added, - "gpu_batches_total": gpu_batches, - "gpu_batches_added": gpu_batches_added, - "gpu_batches_per_second": gpu_batches_added - / max(collect_seconds, 1e-9), - "gpu_evals_total": gpu_evals, - "gpu_evals_added": gpu_evals_added, - "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), - "action_steps_total": action_steps, - "action_steps_added": action_steps_added, - "final_actions_total": final_actions, - "final_actions_added": final_actions_added, - "nonfinal_actions_total": nonfinal_actions, - "nonfinal_actions_added": nonfinal_actions_added, - "completed_turns_total": completed_turns, - "completed_turns_added": completed_turns_added, - "action_turn_count_total": action_turn_count_total, - "action_turn_count_added": action_turn_count_added, - "avg_turn_count": avg_turn_count, - "max_turn_count_seen": max_turn_count_seen, - "actions_per_turn": actions_per_turn, - "completed_game_actions_total": completed_game_actions_total, - "completed_game_actions_added": completed_game_actions_added, - "actions_per_completed_game": actions_per_game, - "max_actions_in_completed_game": max_actions_in_completed_game, - "completed_game_turn_count_total": completed_game_turn_count_total, - "completed_game_turn_count_added": completed_game_turn_count_added, - "completed_game_turn_count_mean": completed_game_turn_count_mean, - "completed_game_turn_count_p50": completed_game_turn_count_p50, - "completed_game_turn_count_p90": completed_game_turn_count_p90, - "max_turn_count_in_completed_game": max_turn_count_in_completed_game, - "replay_size": replay_size, - "collection_seconds": collect_seconds, - "training_seconds": train_seconds, - "samples_per_second": samples_added / max(collect_seconds, 1e-9), - "train_steps_per_second": config.train_steps_per_round - / max(train_seconds, 1e-9), - "selfplay_tree_build_seconds": tree_build_seconds, - "selfplay_descent_seconds": descent_seconds, - "selfplay_sample_collect_seconds": sample_collect_seconds, - "selfplay_replay_push_seconds": replay_push_seconds, - "selfplay_tree_build_ms_per_action": build_ms_per_action, - "selfplay_descent_ms_per_action": descent_ms_per_action, - "selfplay_sample_collect_ms_per_action": collect_ms_per_action, - "selfplay_replay_push_ms_per_action": push_ms_per_action, - "descent_expand_cpu_seconds": descent_expand_cpu_seconds, - "descent_apply_action_seconds": descent_apply_action_seconds, - "descent_eval_submit_seconds": descent_eval_submit_seconds, - "descent_eval_await_seconds": descent_eval_await_seconds, - "descent_backup_seconds": descent_backup_seconds, - "descent_other_seconds": descent_other_seconds, - "descent_expand_cpu_ms_per_action": descent_expand_ms_per_action, - "descent_apply_action_ms_per_action": descent_apply_ms_per_action, - "descent_eval_submit_ms_per_action": descent_submit_ms_per_action, - "descent_eval_await_ms_per_action": descent_wait_ms_per_action, - "descent_backup_ms_per_action": descent_backup_ms_per_action, - "descent_other_ms_per_action": descent_other_ms_per_action, - "descent_eval_submit_us_per_eval": descent_submit_us_per_eval, - "descent_eval_await_us_per_eval": descent_wait_us_per_eval, - "train_sample_seconds": train_sample_seconds, - "train_h2d_seconds": train_h2d_seconds, - "train_forward_seconds": train_forward_seconds, - "train_backward_seconds": train_backward_seconds, - "train_optimizer_seconds": train_optimizer_seconds, - "loss_mean": mean_loss, - "learning_rate": optimizer.param_groups[0]["lr"], - "timestamp": time.time(), - **diag_stats, - } - print( - f"round={round_number} samples={collected} games={games} gpu_evals/s={gpu_evals_added / max(collect_seconds, 1e-9):.1f} replay={replay_size} loss={mean_loss:.6f}" + + record = logger.update( + selfplay=selfplay, + round_number=round_number, + collected=collected, + replay_size=len(replay_buffer), + mean_loss=mean_loss, + collect_seconds=collect_seconds, + train_seconds=train_seconds, + train_sample_seconds=train_sample_seconds, + train_h2d_seconds=train_h2d_seconds, + train_forward_seconds=train_forward_seconds, + train_backward_seconds=train_backward_seconds, + train_optimizer_seconds=train_optimizer_seconds, + diag_stats=diag_stats, + diag_histograms=diag_histograms, ) - if run is not None: - wandb.log({**record, **diag_histograms}) if config.checkpoint_interval > 0 and ( round_number % config.checkpoint_interval == 0 @@ -836,37 +545,14 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: config=config, round_idx=round_number, samples=collected, - games=games, - replay_size=replay_size, + games=selfplay.games(), + replay_size=len(replay_buffer), record=record, ) - print(f"checkpoint={checkpoint_path}") - - previous_samples = collected - previous_terminal_mix_samples = terminal_mix_samples - previous_games = games - previous_gpu_batches = gpu_batches - previous_gpu_evals = gpu_evals - previous_action_steps = action_steps - previous_final_actions = final_actions - previous_nonfinal_actions = nonfinal_actions - previous_completed_turns = completed_turns - previous_action_turn_count_total = action_turn_count_total - previous_completed_game_actions_total = completed_game_actions_total - previous_completed_game_turn_count_total = completed_game_turn_count_total - previous_tree_build_nanos = tree_build_nanos - previous_descent_nanos = descent_nanos - previous_sample_collect_nanos = sample_collect_nanos - previous_replay_push_nanos = replay_push_nanos - previous_descent_expand_cpu_nanos = descent_expand_cpu_nanos - previous_descent_apply_action_nanos = descent_apply_action_nanos - previous_descent_eval_submit_nanos = descent_eval_submit_nanos - previous_descent_eval_await_nanos = descent_eval_await_nanos - previous_descent_backup_nanos = descent_backup_nanos + logger.log_checkpoint(checkpoint_path) finally: selfplay.drop() - if run is not None: - run.finish() + logger.finish() return model, losses diff --git a/training/src/worker.rs b/training/src/worker.rs index e91184b..ffb16fb 100644 --- a/training/src/worker.rs +++ b/training/src/worker.rs @@ -257,10 +257,6 @@ async fn play_game( // Select action via ordinal distribution if tree.root_node.children.is_empty() { - eprintln!( - "how the fuck did we get here?? board has no children. FEN={}", - tree.root_state - ); finish_game_with_current_tree( &tree, replay_buffer, From 0f3ca5477587f968a153fd51bbf2edcee293e775 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 00:24:38 -0400 Subject: [PATCH 50/59] even larger resnet --- python/alphapaint_training/train.py | 37 +++++++++++++++-------------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 324dfdb..add6d1c 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -31,21 +31,21 @@ class TrainConfig: replay_capacity: int = 16_000_000 num_threads: int = 16 workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 2_048 + max_gpu_evals_per_move: int = 4096 lr: float = 3e-4 pretrain_lr: float | None = None selfplay_lr: float | None = 5e-5 weight_decay: float = 1e-4 - width: int = 32 - num_blocks: int = 4 + width: int = 48 + num_blocks: int = 6 hidden_dim: int = 256 - pretrain_terminal_samples: int = 5_000_000 - pretrain_batch_size: int = 8_192 + pretrain_terminal_samples: int = 7_500_000 + pretrain_batch_size: int = 16_384 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" device: str = "cuda" - checkpoint_interval: int = 3 + checkpoint_interval: int = 2 run_dir: str = "runs/latest" wandb: bool = True @@ -454,18 +454,19 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: round_results: list[TrainStepResult] = [] train_started_at = time.perf_counter() - for step_idx in range(config.train_steps_per_round): - result = train_step( - model, - replay_buffer, - optimizer, - batch_size=config.batch_size, - seed=config.seed + round_idx * 10_000 + step_idx, - device=device, - ) - round_results.append(result) - if device.type == "cuda" and round_results: - torch.cuda.synchronize(device) + if collected > config.replay_capacity // 2: + for step_idx in range(config.train_steps_per_round): + result = train_step( + model, + replay_buffer, + optimizer, + batch_size=config.batch_size, + seed=config.seed + round_idx * 10_000 + step_idx, + device=device, + ) + round_results.append(result) + if device.type == "cuda" and round_results: + torch.cuda.synchronize(device) model.eval() train_seconds = time.perf_counter() - train_started_at From db517fa70ed453f92ebaacb6a7f2bbfdbc140a85 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 00:48:36 -0400 Subject: [PATCH 51/59] multigpu!! --- .../alphapaint_training.pyi | 1 + .../alphapaint_training/cudagraph_backend.py | 129 ++++++++++++------ python/alphapaint_training/packed_obs.py | 25 ++-- python/alphapaint_training/train.py | 85 +++++++++++- training/src/cudagraph.rs | 51 ++++++- training/src/lib.rs | 125 ++++++++++++++--- 6 files changed, 335 insertions(+), 81 deletions(-) diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi index 94213d7..0a8cd21 100644 --- a/python/alphapaint_training/alphapaint_training.pyi +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -23,6 +23,7 @@ class SelfPlay: max_gpu_evals_per_move: int = 4096, model: object, selfplay_precision: str = "bf16", + num_gpus: int = 1, ) -> None: ... def start(self) -> None: ... def wait_for(self, target_samples: int) -> int: ... diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py index dd62b0d..c6e3e19 100644 --- a/python/alphapaint_training/cudagraph_backend.py +++ b/python/alphapaint_training/cudagraph_backend.py @@ -10,6 +10,48 @@ import torch.utils.dlpack as dlpack +def _validate_lane_tensors( + obs_host: torch.Tensor, + obs_device: torch.Tensor, + value_host: torch.Tensor, + value_device: torch.Tensor, + gpu_id: int, +) -> None: + if obs_host.device.type != "cpu": + raise ValueError("obs_host must be a CPU tensor") + if value_host.device.type != "cpu": + raise ValueError("value_host must be a CPU tensor") + if obs_device.device.type != "cuda": + raise ValueError("obs_device must be a CUDA tensor") + if value_device.device.type != "cuda": + raise ValueError("value_device must be a CUDA tensor") + + if obs_host.dtype != torch.uint16: + raise ValueError(f"obs_host must be uint16, got {obs_host.dtype}") + if obs_device.dtype != torch.uint16: + raise ValueError(f"obs_device must be uint16, got {obs_device.dtype}") + if value_host.dtype != torch.float32: + raise ValueError(f"value_host must be float32, got {value_host.dtype}") + if value_device.dtype != torch.float32: + raise ValueError(f"value_device must be float32, got {value_device.dtype}") + + if obs_host.shape != obs_device.shape: + raise ValueError( + f"obs_host and obs_device shape mismatch: {obs_host.shape} vs {obs_device.shape}" + ) + batch = obs_host.shape[0] + if value_host.shape != (batch,): + raise ValueError(f"value_host must be shape ({batch},), got {value_host.shape}") + if value_device.shape != (batch,): + raise ValueError( + f"value_device must be shape ({batch},), got {value_device.shape}" + ) + + for name, tensor in (("obs_device", obs_device), ("value_device", value_device)): + if tensor.device.index != gpu_id: + raise ValueError(f"{name} must be on cuda:{gpu_id}, got {tensor.device}") + + def _autocast_dtype(precision: str) -> Optional[torch.dtype]: if precision == "fp32": return None @@ -30,6 +72,7 @@ def capture_lane_graph( value_device_dlpack, stream_handle: int, precision: str = "bf16", + gpu_id: int = 0, ) -> tuple[int, object]: """Capture a CUDA graph for value-only inference. @@ -44,53 +87,61 @@ def capture_lane_graph( value_device_dlpack: DLPack capsule for device value buffer. stream_handle: Raw CUDA stream handle. precision: "bf16", "fp16", or "fp32". + gpu_id: CUDA device index to capture the graph on. Returns: Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). """ + if gpu_id < 0: + raise ValueError(f"gpu_id must be >= 0, got {gpu_id}") + obs_host = dlpack.from_dlpack(obs_host_dlpack) obs_device = dlpack.from_dlpack(obs_device_dlpack) value_host = dlpack.from_dlpack(value_host_dlpack) value_device = dlpack.from_dlpack(value_device_dlpack) - model = model.cuda() - model = model.to(memory_format=torch.channels_last) - model.eval() - torch.backends.cudnn.benchmark = True - stream = torch.cuda.ExternalStream(stream_handle) - graph = torch.cuda.CUDAGraph(keep_graph=True) - dtype = _autocast_dtype(precision) - - def run_step() -> None: - obs_device.copy_(obs_host, non_blocking=True) - if dtype is None: - value = model(obs_device) - else: - with torch.autocast(device_type="cuda", dtype=dtype): + _validate_lane_tensors(obs_host, obs_device, value_host, value_device, gpu_id) + + with torch.cuda.device(gpu_id): + model = model.to(f"cuda:{gpu_id}") + model = model.to(memory_format=torch.channels_last) + model.eval() + torch.backends.cudnn.benchmark = True + stream = torch.cuda.ExternalStream(stream_handle) + graph = torch.cuda.CUDAGraph(keep_graph=True) + dtype = _autocast_dtype(precision) + + def run_step() -> None: + obs_device.copy_(obs_host, non_blocking=True) + if dtype is None: value = model(obs_device) - # Model returns value tensor, shape (B,) or (B, 1) - if value.ndim == 2: - value = value.squeeze(-1) - value_device.copy_(value, non_blocking=True) - value_host.copy_(value_device, non_blocking=True) - - with torch.inference_mode(): - with torch.cuda.stream(stream): - for _ in range(3): + else: + with torch.autocast(device_type="cuda", dtype=dtype): + value = model(obs_device) + if value.ndim == 2: + value = value.squeeze(-1) + value_device.copy_(value, non_blocking=True) + value_host.copy_(value_device, non_blocking=True) + + with torch.inference_mode(): + with torch.cuda.stream(stream): + for _ in range(3): + run_step() + torch.cuda.synchronize(device=gpu_id) + + with torch.cuda.graph( + graph, stream=stream, capture_error_mode="thread_local" + ): run_step() - torch.cuda.synchronize() - - with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"): - run_step() - - graph.instantiate() - owner = ( - graph, - model, - obs_host, - obs_device, - value_host, - value_device, - stream, - ) - return int(graph.raw_cuda_graph_exec()), owner + + graph.instantiate() + owner = ( + graph, + model, + obs_host, + obs_device, + value_host, + value_device, + stream, + ) + return int(graph.raw_cuda_graph_exec()), owner diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index b967867..2433747 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -230,18 +230,19 @@ def decode_packed_board( raise ValueError(f"out must have dtype {dtype}, got {out.dtype}") total_cells = board_words.shape[0] * BOARD_CELLS grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) - _decode_board_kernel[grid]( - board_words, - out, - total_cells, - board_words.stride(0), - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - BLOCK=256, - num_warps=4, - ) + with torch.cuda.device(board_words.device): + _decode_board_kernel[grid]( + board_words, + out, + total_cells, + board_words.stride(0), + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + BLOCK=256, + num_warps=4, + ) return out diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index add6d1c..8c31fce 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -2,9 +2,10 @@ import argparse import json +import os import time from contextlib import nullcontext -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, field from pathlib import Path from typing import cast @@ -22,16 +23,57 @@ from alphapaint_training.model import PackedValueModel +def _int_env(name: str) -> int | None: + value = os.environ.get(name) + if not value: + return None + try: + return int(value) + except ValueError: + return None + + +def _slurm_gpu_count() -> int | None: + for name in ("SLURM_GPUS_ON_NODE", "SLURM_GPUS"): + value = _int_env(name) + if value is not None and value > 0: + return value + + job_gpus = os.environ.get("SLURM_JOB_GPUS") + if not job_gpus: + return None + + gpu_ids = [gpu_id.strip() for gpu_id in job_gpus.split(",") if gpu_id.strip()] + if not gpu_ids: + return None + return len(gpu_ids) + + +def _default_num_threads() -> int: + cpus_per_gpu = _int_env("SLURM_CPUS_PER_GPU") + slurm_gpu_count = _slurm_gpu_count() + if cpus_per_gpu is not None and cpus_per_gpu > 0 and slurm_gpu_count is not None: + return cpus_per_gpu * slurm_gpu_count + return 32 + + +def _default_num_gpus() -> int: + slurm_gpu_count = _slurm_gpu_count() + if slurm_gpu_count is None: + return 0 + return slurm_gpu_count + + @dataclass(slots=True) class TrainConfig: rounds: int = 200 samples_per_round: int = 1_048_576 train_steps_per_round: int = 256 - batch_size: int = 4_096 + batch_size: int = 8_192 replay_capacity: int = 16_000_000 - num_threads: int = 16 + num_threads: int = field(default_factory=_default_num_threads) workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 4096 + max_gpu_evals_per_move: int = 4 * 4096 lr: float = 3e-4 pretrain_lr: float | None = None selfplay_lr: float | None = 5e-5 @@ -40,10 +82,11 @@ class TrainConfig: num_blocks: int = 6 hidden_dim: int = 256 pretrain_terminal_samples: int = 7_500_000 - pretrain_batch_size: int = 16_384 + pretrain_batch_size: int = 8_192 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" + num_gpus: int = field(default_factory=_default_num_gpus) device: str = "cuda" checkpoint_interval: int = 2 run_dir: str = "runs/latest" @@ -75,6 +118,8 @@ class TrainStepResult: def _device(device: str) -> torch.device: if device == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA is required for the training runner") + if device == "cuda": + return torch.device("cuda:0") return torch.device(device) @@ -432,6 +477,19 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: print(f"selfplay_lr={selfplay_lr:.6g}") replay_buffer = EphemeralReplayBuffer(config.replay_capacity) + available_gpus = torch.cuda.device_count() + if available_gpus < 1: + raise RuntimeError("AlphaPaint self-play requires at least one CUDA GPU") + num_gpus = config.num_gpus if config.num_gpus > 0 else available_gpus + if num_gpus > available_gpus: + raise ValueError( + f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" + ) + print( + "selfplay " + f"threads={config.num_threads} workers_per_thread={config.workers_per_thread} " + f"gpus={num_gpus} max_gpu_evals_per_move={config.max_gpu_evals_per_move}" + ) selfplay = SelfPlay( replay_buffer, config.num_threads, @@ -440,7 +498,11 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: max_gpu_evals_per_move=config.max_gpu_evals_per_move, model=model, selfplay_precision=config.selfplay_precision, + num_gpus=num_gpus, ) + # Graph capture via cudaSetDevice may leave the CUDA default on another GPU; + # pin it back to the training device so H2D transfers / diagnostics land on cuda:0. + torch.cuda.set_device(device) losses: list[float] = [] target_samples = 0 @@ -570,7 +632,12 @@ def _parse_args() -> TrainConfig: ) parser.add_argument("--batch-size", type=int, default=defaults.batch_size) parser.add_argument("--replay-capacity", type=int, default=defaults.replay_capacity) - parser.add_argument("--num-threads", type=int, default=defaults.num_threads) + parser.add_argument( + "--num-threads", + type=int, + default=defaults.num_threads, + help="Self-play executor threads (defaults to Slurm CPU allocation when available)", + ) parser.add_argument( "--workers-per-thread", type=int, default=defaults.workers_per_thread ) @@ -607,6 +674,12 @@ def _parse_args() -> TrainConfig: default=defaults.selfplay_precision, choices=["bf16", "fp16", "fp32"], ) + parser.add_argument( + "--num-gpus", + type=int, + default=defaults.num_gpus, + help="Number of GPUs for self-play (0 uses all visible CUDA GPUs; defaults to Slurm allocation when available)", + ) parser.add_argument("--device", default=defaults.device) parser.add_argument( "--checkpoint-interval", type=int, default=defaults.checkpoint_interval diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs index 730f6fe..edbc4f5 100644 --- a/training/src/cudagraph.rs +++ b/training/src/cudagraph.rs @@ -178,6 +178,34 @@ fn check_cuda_or_panic(code: CudaError, context: &str) { } } +fn validate_cuda_device(gpu_id: usize) -> PyResult { + let gpu_device_id = i32::try_from(gpu_id) + .map_err(|_| PyErr::new::("gpu_id must fit in i32"))?; + + let mut device_count = 0i32; + unsafe { + check_cuda( + cuda::cudaGetDeviceCount(&mut device_count as *mut i32), + "cudaGetDeviceCount", + )?; + } + + if device_count <= 0 { + return Err(PyErr::new::( + "no CUDA devices available for AlphaPaint CUDA graph runner", + )); + } + + if gpu_device_id >= device_count { + return Err(PyErr::new::(format!( + "gpu_id {} out of range for {} CUDA devices", + gpu_id, device_count + ))); + } + + Ok(gpu_device_id) +} + fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { let mut ptr: *mut c_void = std::ptr::null_mut(); let bytes = count @@ -235,6 +263,7 @@ fn cuda_malloc_device_u16(count: usize, context: &str) -> PyResult<*mut c_void> } struct CudaGraphLane { + gpu_device_id: i32, stream: cudaStream_t, graph_exec: cudaGraphExec_t, /// Owns Python-side graph/tensor objects for this lane. @@ -271,6 +300,7 @@ unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { impl Drop for CudaGraphLane { fn drop(&mut self) { unsafe { + let _ = cuda::cudaSetDevice(self.gpu_device_id); let _ = cuda::cudaFree(self.obs_dev); let _ = cuda::cudaFree(self.value_dev); let _ = cuda::cudaFreeHost(self.obs_host.cast::()); @@ -282,6 +312,7 @@ impl Drop for CudaGraphLane { /// Per-lane CUDA graph executor for AlphaPaint value-only inference. pub struct CudaGraphRunner { + gpu_device_id: i32, batch_size: usize, lanes: Vec, dispatched_batches: AtomicU64, @@ -297,6 +328,7 @@ impl CudaGraphRunner { pub fn new( py: Python<'_>, model: Py, + gpu_id: usize, num_lanes: usize, batch_size: usize, precision: &str, @@ -308,6 +340,14 @@ impl CudaGraphRunner { return Err(PyErr::new::("batch_size must be > 0")); } + let gpu_device_id = validate_cuda_device(gpu_id)?; + unsafe { + check_cuda( + cuda::cudaSetDevice(gpu_device_id), + "cudaSetDevice in CudaGraphRunner::new", + )?; + } + let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; let capture_fn = module.getattr("capture_lane_graph")?; @@ -356,7 +396,7 @@ impl CudaGraphRunner { obs_dev, &obs_shape, DL_DEVICE_CUDA, - 0, + gpu_device_id, DL_DTYPE_UINT, 16, )?; @@ -374,7 +414,7 @@ impl CudaGraphRunner { value_dev, &value_shape, DL_DEVICE_CUDA, - 0, + gpu_device_id, DL_DTYPE_FLOAT, 32, )?; @@ -388,10 +428,12 @@ impl CudaGraphRunner { value_dev_capsule, stream as u64, precision, + gpu_id, ))? .extract()?; let lane = CudaGraphLane { + gpu_device_id, stream, graph_exec: exec_handle as cudaGraphExec_t, _py_owner: py_owner, @@ -404,6 +446,7 @@ impl CudaGraphRunner { } Ok(Self { + gpu_device_id, batch_size, lanes, dispatched_batches: AtomicU64::new(0), @@ -440,6 +483,10 @@ impl CudaGraphRunner { .fetch_add(self.batch_size as u64, Ordering::Relaxed); unsafe { + check_cuda_or_panic( + cuda::cudaSetDevice(self.gpu_device_id), + "cudaSetDevice before cudaGraphLaunch", + ); check_cuda_or_panic( cuda::cudaGraphLaunch(lane.graph_exec, lane.stream), "cudaGraphLaunch", diff --git a/training/src/lib.rs b/training/src/lib.rs index 86b8dd9..1d22009 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -1,5 +1,6 @@ use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; use alpha_paint::TRAINING_START_FENS; +use std::collections::HashMap; use std::sync::Arc; use std::sync::{Mutex, OnceLock}; @@ -86,7 +87,9 @@ struct GraphCacheEntry { model_ptr: usize, num_batches: usize, precision: String, - runner: Arc, + num_gpus: usize, + runners: HashMap>, + models: HashMap>, } static GRAPH_CACHE: OnceLock>> = OnceLock::new(); @@ -272,7 +275,9 @@ impl EphemeralReplayBuffer { #[pyclass] struct SelfPlay { session: Option, - runner: Arc, + runners: HashMap>, + source_model: Py, + source_model_ptr: usize, } impl SelfPlay { @@ -294,7 +299,8 @@ impl SelfPlay { *, max_gpu_evals_per_move = 4096, model, - selfplay_precision = "bf16" + selfplay_precision = "bf16", + num_gpus = 1 ))] fn new( py: Python<'_>, @@ -305,7 +311,14 @@ impl SelfPlay { max_gpu_evals_per_move: u64, model: Py, selfplay_precision: &str, + num_gpus: usize, ) -> PyResult { + if num_gpus == 0 { + return Err(PyErr::new::( + "num_gpus must be >= 1", + )); + } + let config = SessionConfig { num_threads, workers_per_thread, @@ -324,8 +337,8 @@ impl SelfPlay { let (num_batches, _total_slots) = queue_shape_for_workers(total_workers); let model_ptr = model.bind(py).as_ptr() as usize; - // Build or reuse the CUDA graph runner. - let runner = { + // Build or reuse the CUDA graph runners. + let runners = { let cache = graph_cache(); let mut guard = cache.lock().expect("graph cache mutex poisoned"); @@ -334,57 +347,75 @@ impl SelfPlay { entry.model_ptr != model_ptr || entry.num_batches != num_batches || entry.precision != selfplay_precision + || entry.num_gpus != num_gpus } None => true, }; if needs_rebuild { - let runner = Arc::new(CudaGraphRunner::new( - py, - model.clone_ref(py), - num_batches, - BATCH_SIZE, - selfplay_precision, - )?); + let copy = PyModule::import(py, "copy")?; + let deepcopy = copy.getattr("deepcopy")?; + let mut runners = HashMap::new(); + let mut models = HashMap::new(); + for gpu_id in 0..num_gpus { + let model_copy: Py = deepcopy.call1((model.clone_ref(py),))?.into(); + let runner = Arc::new(CudaGraphRunner::new( + py, + model_copy.clone_ref(py), + gpu_id, + num_batches, + BATCH_SIZE, + selfplay_precision, + )?); + runners.insert(gpu_id, runner); + models.insert(gpu_id, model_copy); + } *guard = Some(GraphCacheEntry { model_ptr, num_batches, precision: selfplay_precision.to_string(), - runner: runner.clone(), + num_gpus, + runners: runners.clone(), + models, }); - runner + runners } else { guard .as_ref() - .expect("cached runner should exist") - .runner + .expect("cached runners should exist") + .runners .clone() } }; - let runner_for_dispatch = runner.clone(); + let runners_for_dispatch = runners.clone(); let dispatch = move |batch_idx: usize, obs_view: ArrayView, completion: queue::BatchCompletion| { - runner_for_dispatch.dispatch_async(batch_idx, obs_view, completion); + let gpu_id = batch_idx % num_gpus; + runners_for_dispatch[&gpu_id].dispatch_async(batch_idx, obs_view, completion); }; let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); Ok(Self { session: Some(session), - runner, + runners, + source_model: model, + source_model_ptr: model_ptr, }) } /// Start self-play with no sample limit. - fn start(&self) -> PyResult<()> { + fn start(&self, py: Python<'_>) -> PyResult<()> { + self.sync_model_replicas(py)?; self.session()?.start(); Ok(()) } /// Block until absolute target_samples is reached, then pause and quiesce. fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { + self.sync_model_replicas(py)?; let session = self.session()?; let result = py.detach(|| session.wait_for(target_samples)); Ok(result) @@ -507,12 +538,18 @@ impl SelfPlay { /// Return the total number of CUDA graph launches completed so far. fn gpu_batches(&self) -> u64 { - self.runner.dispatched_batches() + self.runners + .values() + .map(|runner| runner.dispatched_batches()) + .sum() } /// Return the total number of packed observations sent to GPU so far. fn gpu_evals(&self) -> u64 { - self.runner.dispatched_evals() + self.runners + .values() + .map(|runner| runner.dispatched_evals()) + .sum() } /// Shut down the session. Idempotent. @@ -532,6 +569,50 @@ impl Drop for SelfPlay { } } +impl SelfPlay { + fn sync_model_replicas(&self, py: Python<'_>) -> PyResult<()> { + let replicas = { + let cache = graph_cache(); + let guard = cache.lock().expect("graph cache mutex poisoned"); + let entry = guard.as_ref().ok_or_else(|| { + PyErr::new::( + "graph cache missing while syncing model replicas", + ) + })?; + + if entry.model_ptr != self.source_model_ptr { + return Err(PyErr::new::( + "graph cache model mismatch while syncing model replicas", + )); + } + + let mut models: Vec<(usize, Py)> = entry + .models + .iter() + .map(|(gpu_id, model)| (*gpu_id, model.clone_ref(py))) + .collect(); + models.sort_by_key(|(gpu_id, _)| *gpu_id); + models + .into_iter() + .map(|(_, model)| model) + .collect::>() + }; + + let state_dict: Py = self + .source_model + .bind(py) + .call_method0("state_dict")? + .into(); + for replica in replicas { + let _ = replica + .bind(py) + .call_method1("load_state_dict", (state_dict.clone_ref(py),))?; + } + + Ok(()) + } +} + #[pymodule] fn alphapaint_training(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; From a06e645939d3b9d0ef1c3c91afc2da9269682ce8 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 02:44:34 -0400 Subject: [PATCH 52/59] scuffed nnue training but like we r evaling nnue on gpu for some reason lol --- python/alphapaint_training/__init__.py | 31 +- python/alphapaint_training/model.py | 284 ++++++------ python/alphapaint_training/packed_obs.py | 394 ++++++---------- python/alphapaint_training/train.py | 184 ++------ training/src/observation.rs | 546 ++++++++++++++--------- 5 files changed, 668 insertions(+), 771 deletions(-) diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 81d638c..d9e640d 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -10,34 +10,35 @@ sample_random_terminal_batch, ) from .logger import TrainingLogger -from .model import PackedValueModel, ResidualBlock, TinyValueNet +from .model import NnueValueNet, PackedValueModel from .packed_obs import ( BOARD_CELLS, - BOARD_PLANES, BOARD_SIDE, - INTRINSIC_COUNT, + GLOBAL_FEATURES, + GLOBAL_SCALES, + LOCAL_FEATURES, + LOCAL_WINDOW_TILES, OBS_WORDS, - decode_intrinsics, - decode_packed_board, - decode_packed_board_reference, - decode_packed_observation, + TILE_PLANES, + TOTAL_TILE_FEATURES, + decode_nnue_obs, ) __all__ = [ "BOARD_CELLS", - "BOARD_PLANES", "BOARD_SIDE", "EphemeralReplayBuffer", - "INTRINSIC_COUNT", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", + "NnueValueNet", "OBS_WORDS", "PackedValueModel", - "ResidualBlock", "SelfPlay", - "TinyValueNet", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", "TrainingLogger", - "decode_intrinsics", - "decode_packed_board", - "decode_packed_board_reference", - "decode_packed_observation", + "decode_nnue_obs", "sample_random_terminal_batch", ] diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 1fcbe33..8227224 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -1,194 +1,186 @@ +"""NNUE value model for GPU training. + +Architecture mirrors the CPU NNUE evaluator: + - Accumulator: shared linear 15360 -> acc_dim for both perspectives + - Local MLP: shared linear 375 -> local_hidden_dim (ReLU) for both players + - Global features: 20 normalized scalars + - Head: concat all -> fc1 (ReLU) -> fc2 (ReLU) -> fc3 -> scalar + +The same weight matrix is applied to both the current-player and +opponent perspectives, matching the CPU NNUE design where one set of +feature weights serves both the white and black accumulators. +""" + from __future__ import annotations -from collections import OrderedDict from typing import cast import torch from torch import nn from .packed_obs import ( - BOARD_CELLS, - BOARD_PLANES, - INTRINSIC_TURN_COUNT, - INTRINSIC_COUNT, - INTRINSIC_SCALE, - decode_packed_board, + GLOBAL_FEATURES, + GLOBAL_SCALES, + GLOBAL_TURN_COUNT_IDX, + LOCAL_FEATURES, + OFFSET_GLOBALS, + OPP_PLANE_PERM, + TILE_PLANES, + TOTAL_TILE_FEATURES, + decode_nnue_obs, ) -class ResidualBlock(nn.Module): - def __init__(self, width: int): - super().__init__() - self.norm1 = nn.BatchNorm2d(width) - self.conv1 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.norm2 = nn.BatchNorm2d(width) - self.conv2 = nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False) - self.act = nn.ReLU(inplace=True) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - residual = x - x = self.norm1(x) - x = self.act(x) - x = self.conv1(x) - x = self.norm2(x) - x = self.act(x) - x = self.conv2(x) - return x + residual - - -class TinyValueNet(nn.Module): +class NnueValueNet(nn.Module): + """NNUE-style value network for GPU training. + + Input components: + acc_mine: [B, acc_dim] from accumulator (my perspective) + acc_opp: [B, acc_dim] from accumulator (opp perspective) + local_my: [B, local_hidden] from local MLP (my position) + local_opp: [B, local_hidden] from local MLP (opp position) + globals: [B, 20] normalized global features + + Architecture: + concat -> fc1 (ReLU) -> fc2 (ReLU) -> fc3 -> scalar + """ + def __init__( self, *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, + acc_dim: int = 128, + local_hidden_dim: int = 64, + fc1_dim: int = 128, + fc2_dim: int = 32, ): super().__init__() - self.stem = nn.Sequential( - nn.Conv2d(BOARD_PLANES, width, kernel_size=3, padding=1, bias=False), - nn.BatchNorm2d(width), - nn.ReLU(inplace=True), - ) - self.blocks = nn.Sequential(*(ResidualBlock(width) for _ in range(num_blocks))) - self.pool = nn.AdaptiveAvgPool2d(1) - self.flatten = nn.Flatten() - self.head = nn.Sequential( - nn.Linear(width + INTRINSIC_COUNT, hidden_dim), + self.acc_dim = acc_dim + self.local_hidden_dim = local_hidden_dim + + # Accumulator: shared for both perspectives + self.acc_linear = nn.Linear(TOTAL_TILE_FEATURES, acc_dim) + + # Local MLP: shared for both player windows + self.local_mlp = nn.Sequential( + nn.Linear(LOCAL_FEATURES, local_hidden_dim), nn.ReLU(inplace=True), - nn.Linear(hidden_dim, 1), ) - def forward(self, board: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor: - x = self.stem(board) - x = self.blocks(x) - x = self.pool(x) - x = self.flatten(x) - x = torch.cat((x, intrinsics), dim=1) - return self.head(x) + # Final head + fc1_in = acc_dim * 2 + local_hidden_dim * 2 + GLOBAL_FEATURES + self.fc1 = nn.Linear(fc1_in, fc1_dim) + self.fc2 = nn.Linear(fc1_dim, fc2_dim) + self.fc3 = nn.Linear(fc2_dim, 1) + + def forward( + self, + my_features: torch.Tensor, + opp_features: torch.Tensor, + local_my: torch.Tensor, + local_opp: torch.Tensor, + globals_norm: torch.Tensor, + ) -> torch.Tensor: + # Shared accumulator + acc_mine = self.acc_linear(my_features) + acc_opp = self.acc_linear(opp_features) + + # Shared local MLP + local_mine_out = self.local_mlp(local_my) + local_opp_out = self.local_mlp(local_opp) + + # Concat and head + x = torch.cat( + [acc_mine, acc_opp, local_mine_out, local_opp_out, globals_norm], dim=1 + ) + x = torch.relu(self.fc1(x)) + x = torch.relu(self.fc2(x)) + return self.fc3(x) class PackedValueModel(nn.Module): + """Wraps NnueValueNet with packed observation decoding. + + Takes [B, OBS_WORDS] uint16 packed observations, decodes them on-device, + runs the NNUE value net, and returns [B] scalar values from White's + perspective. + """ + def __init__( self, *, - width: int = 8, - num_blocks: int = 1, - hidden_dim: int = 32, + acc_dim: int = 128, + local_hidden_dim: int = 64, + fc1_dim: int = 128, + fc2_dim: int = 32, board_dtype: torch.dtype = torch.bfloat16, ): super().__init__() self.board_dtype = board_dtype - self._board_buffers: dict[ - tuple[str, int | None, int, torch.dtype, int], torch.Tensor - ] = OrderedDict() - self._intrinsic_fp32_buffers: dict[ - tuple[str, int | None, int, int], torch.Tensor - ] = OrderedDict() - self._intrinsic_buffers: dict[ - tuple[str, int | None, int, torch.dtype, int], torch.Tensor - ] = OrderedDict() + self.value_net = NnueValueNet( + acc_dim=acc_dim, + local_hidden_dim=local_hidden_dim, + fc1_dim=fc1_dim, + fc2_dim=fc2_dim, + ) + + # Decode buffers — registered so they follow .to(device) and are + # captured in CUDA graphs. self.register_buffer( - "intrinsic_scales", - torch.tensor(INTRINSIC_SCALE, dtype=torch.float32), + "_plane_shift", + torch.arange(TILE_PLANES, dtype=torch.int32), persistent=False, ) - self.value_net = TinyValueNet( - width=width, - num_blocks=num_blocks, - hidden_dim=hidden_dim, + self.register_buffer( + "_opp_perm", + torch.tensor(OPP_PLANE_PERM, dtype=torch.long), + persistent=False, ) - - @staticmethod - def _trim_cache[K](cache: dict[K, torch.Tensor], max_entries: int = 64) -> None: - while len(cache) > max_entries: - cache.pop(next(iter(cache))) - - def _buffer_key( - self, packed_obs: torch.Tensor, dtype: torch.dtype - ) -> tuple[str, int | None, int, torch.dtype, int]: - return ( - packed_obs.device.type, - packed_obs.device.index, - packed_obs.shape[0], - dtype, - packed_obs.data_ptr(), + self.register_buffer( + "_global_scales", + torch.tensor(GLOBAL_SCALES, dtype=torch.float32), + persistent=False, ) - def _ensure_decode_buffers( + def decode( self, packed_obs: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - board_key = self._buffer_key(packed_obs, self.board_dtype) - board = self._board_buffers.get(board_key) - if board is None: - memory_format = ( - torch.channels_last - if packed_obs.device.type == "cuda" - else torch.contiguous_format - ) - board = torch.empty( - (packed_obs.shape[0], BOARD_PLANES, 32, 32), - device=packed_obs.device, - dtype=self.board_dtype, - memory_format=memory_format, - ) - self._board_buffers[board_key] = board - self._trim_cache(self._board_buffers) - - fp32_key = ( - packed_obs.device.type, - packed_obs.device.index, - packed_obs.shape[0], - packed_obs.data_ptr(), - ) - intrinsics_fp32 = self._intrinsic_fp32_buffers.get(fp32_key) - if intrinsics_fp32 is None: - intrinsics_fp32 = torch.empty( - (packed_obs.shape[0], INTRINSIC_COUNT), - device=packed_obs.device, - dtype=torch.float32, - ) - self._intrinsic_fp32_buffers[fp32_key] = intrinsics_fp32 - self._trim_cache(self._intrinsic_fp32_buffers) - - intrinsic_key = self._buffer_key(packed_obs, self.board_dtype) - intrinsics = self._intrinsic_buffers.get(intrinsic_key) - if intrinsics is None: - intrinsics = torch.empty( - (packed_obs.shape[0], INTRINSIC_COUNT), - device=packed_obs.device, - dtype=self.board_dtype, - ) - self._intrinsic_buffers[intrinsic_key] = intrinsics - self._trim_cache(self._intrinsic_buffers) - - return board, intrinsics_fp32, intrinsics - - def decode(self, packed_obs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - board, intrinsics_fp32, intrinsics = self._ensure_decode_buffers(packed_obs) - board = decode_packed_board( - packed_obs[:, :BOARD_CELLS], + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + return decode_nnue_obs( + packed_obs, + plane_shift=cast(torch.Tensor, self._plane_shift), + opp_perm=cast(torch.Tensor, self._opp_perm), + global_scales=cast(torch.Tensor, self._global_scales), dtype=self.board_dtype, - out=board, ) - intrinsic_scales = cast(torch.Tensor, self.intrinsic_scales) - intrinsics_fp32.copy_(packed_obs[:, BOARD_CELLS:]) - intrinsics_fp32.div_(intrinsic_scales) - intrinsics.copy_(intrinsics_fp32) - return board, intrinsics def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: - board, intrinsics = self.decode(packed_obs) + my_features, opp_features, local_my, local_opp, globals_norm, turn_count = ( + self.decode(packed_obs) + ) + if not torch.is_autocast_enabled(): param_dtype = next(self.value_net.parameters()).dtype - board = board.to(dtype=param_dtype) - intrinsics = intrinsics.to(dtype=param_dtype) - value = self.value_net(board, intrinsics) + my_features = my_features.to(dtype=param_dtype) + opp_features = opp_features.to(dtype=param_dtype) + local_my = local_my.to(dtype=param_dtype) + local_opp = local_opp.to(dtype=param_dtype) + globals_norm = globals_norm.to(dtype=param_dtype) + + value = self.value_net( + my_features, opp_features, local_my, local_opp, globals_norm + ) value = value.squeeze(-1) - turn_count = packed_obs[:, BOARD_CELLS + INTRINSIC_TURN_COUNT].to(torch.int32) + + # Convert from current-player perspective to White perspective. white_to_move = (turn_count & 1) == 0 sign = torch.where(white_to_move, 1.0, -1.0).to(dtype=value.dtype) return value * sign -__all__ = ["PackedValueModel", "ResidualBlock", "TinyValueNet"] +__all__ = ["NnueValueNet", "PackedValueModel"] diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index 2433747..aa2634c 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -1,33 +1,76 @@ +"""Packed observation decoding for NNUE training. + +The Rust encoder packs board state into a flat u16 array with this layout: + [0, 1024) — per-tile 15-bit bitmask (32×32 board) + [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) + [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) + [1078, 1098) — 20 global scalar features + +Tile bitmask planes (from current player's perspective): + bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) + bits 4-7: opponent paint thermometer + bit 8: wall + bit 9: powerup + bit 10: current player beacon + bit 11: opponent beacon + bit 12: hill neutral + bit 13: hill current player + bit 14: hill opponent +""" + from __future__ import annotations import torch -import triton -import triton.language as tl - -BOARD_SIDE = 32 -BOARD_CELLS = BOARD_SIDE * BOARD_SIDE -BOARD_PLANES = 17 -INTRINSIC_COUNT = 10 -OBS_WORDS = BOARD_CELLS + INTRINSIC_COUNT -INTRINSIC_TURN_COUNT = 4 - -PAINT_STRENGTH_MASK = 0b111 -PAINT_IS_ENEMY_BIT = 1 << 3 -WALL_BIT = 1 << 4 -POWERUP_BIT = 1 << 5 -BEACON_SHIFT = 6 -HILL_SHIFT = 8 -CURRENT_PLAYER_BIT = 1 << 10 -OPPONENT_PLAYER_BIT = 1 << 11 -BEACON_CURRENT = 1 -BEACON_OPPONENT = 2 +# Layout constants ---------------------------------------------------------- -HILL_NEUTRAL = 1 -HILL_CURRENT = 2 -HILL_OPPONENT = 3 - -INTRINSIC_SCALE = (420.0, 420.0, 8.0, 8.0, 2000.0, 1024.0, 1024.0, 1024.0, 1024.0, 64.0) +BOARD_SIDE = 32 +BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 +TILE_PLANES = 15 +LOCAL_WINDOW_TILES = 25 +GLOBAL_FEATURES = 20 +TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 15360 +LOCAL_FEATURES = LOCAL_WINDOW_TILES * TILE_PLANES # 375 + +OFFSET_LOCALS = BOARD_CELLS # 1024 +OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 +OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 +OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 + +# Index of turn_count within the global section +GLOBAL_TURN_COUNT_IDX = 17 + +# Plane permutation for opponent perspective: +# swap my_paint[0:4] <-> opp_paint[4:8] +# swap my_beacon[10] <-> opp_beacon[11] +# swap my_hill[13] <-> opp_hill[14] +# wall[8], powerup[9], hill_neutral[12] stay +OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13] + +# Normalization divisors for the 20 global features. +# Brings raw u16 values into roughly [0, 1] range. +GLOBAL_SCALES: tuple[float, ...] = ( + 420.0, # 0 my_stamina + 420.0, # 1 opp_stamina + 420.0, # 2 my_max_stamina + 420.0, # 3 opp_max_stamina + 1024.0, # 4 my_territory + 1024.0, # 5 opp_territory + 8.0, # 6 my_hills + 8.0, # 7 opp_hills + 1024.0, # 8 my_hill_tiles + 1024.0, # 9 opp_hill_tiles + 8.0, # 10 contested_hills + 1024.0, # 11 my_beacons + 1024.0, # 12 opp_beacons + 64.0, # 13 player_dist + 64.0, # 14 my_hill_dist + 64.0, # 15 opp_hill_dist + 32.0, # 16 consecutive_moves + 2000.0, # 17 turn_count + 32.0, # 18 rows + 32.0, # 19 cols +) def _check_packed_obs(packed_obs: torch.Tensor) -> None: @@ -41,250 +84,85 @@ def _check_packed_obs(packed_obs: torch.Tensor) -> None: ) -@triton.jit -def _decode_board_kernel( - packed_ptr, - out_ptr, - total_cells, - packed_stride0, - out_stride0, - out_stride1, - out_stride2, - out_stride3, - BLOCK: tl.constexpr, -): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < total_cells - - batch = offs // 1024 - cell = offs % 1024 - x = cell % 32 - y = cell // 32 - - words = tl.load(packed_ptr + batch * packed_stride0 + cell, mask=mask, other=0).to( - tl.uint16 - ) - - strength = words & 0b111 - enemy = (words >> 3) & 1 - wall = (words >> 4) & 1 - powerup = (words >> 5) & 1 - beacon = (words >> 6) & 0b11 - hill = (words >> 8) & 0b11 - current_player = (words >> 10) & 1 - opponent_player = (words >> 11) & 1 - - current_paint = (strength != 0) & (enemy == 0) - opponent_paint = (strength != 0) & (enemy != 0) - - out_base = batch * out_stride0 + y * out_stride2 + x * out_stride3 - - tl.store( - out_ptr + out_base + 0 * out_stride1, current_paint.to(tl.float32), mask=mask - ) - tl.store( - out_ptr + out_base + 1 * out_stride1, - (current_paint & (strength >= 2)).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 2 * out_stride1, - (current_paint & (strength >= 3)).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 3 * out_stride1, - (current_paint & (strength >= 4)).to(tl.float32), - mask=mask, - ) - - tl.store( - out_ptr + out_base + 4 * out_stride1, opponent_paint.to(tl.float32), mask=mask - ) - tl.store( - out_ptr + out_base + 5 * out_stride1, - (opponent_paint & (strength >= 2)).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 6 * out_stride1, - (opponent_paint & (strength >= 3)).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 7 * out_stride1, - (opponent_paint & (strength >= 4)).to(tl.float32), - mask=mask, - ) - - tl.store(out_ptr + out_base + 8 * out_stride1, wall.to(tl.float32), mask=mask) - tl.store(out_ptr + out_base + 9 * out_stride1, powerup.to(tl.float32), mask=mask) - tl.store( - out_ptr + out_base + 10 * out_stride1, - (beacon == 1).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 11 * out_stride1, - (beacon == 2).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 12 * out_stride1, - (hill == 1).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 13 * out_stride1, - (hill == 2).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 14 * out_stride1, - (hill == 3).to(tl.float32), - mask=mask, - ) - tl.store( - out_ptr + out_base + 15 * out_stride1, current_player.to(tl.float32), mask=mask - ) - tl.store( - out_ptr + out_base + 16 * out_stride1, opponent_player.to(tl.float32), mask=mask - ) - - -def decode_packed_board_reference(board_words: torch.Tensor) -> torch.Tensor: - if board_words.dtype != torch.uint16: - raise ValueError(f"board_words must be uint16, got {board_words.dtype}") - if board_words.ndim != 2 or board_words.shape[1] != BOARD_CELLS: - raise ValueError( - f"board_words must have shape (B, {BOARD_CELLS}), got {tuple(board_words.shape)}" - ) - - words = board_words.to(torch.int32) - strength = words & PAINT_STRENGTH_MASK - enemy = (words & PAINT_IS_ENEMY_BIT) != 0 - - current_paint = (strength != 0) & (~enemy) - opponent_paint = (strength != 0) & enemy - beacon = (words >> BEACON_SHIFT) & 0b11 - hill = (words >> HILL_SHIFT) & 0b11 - - planes = torch.stack( - [ - current_paint, - current_paint & (strength >= 2), - current_paint & (strength >= 3), - current_paint & (strength >= 4), - opponent_paint, - opponent_paint & (strength >= 2), - opponent_paint & (strength >= 3), - opponent_paint & (strength >= 4), - (words & WALL_BIT) != 0, - (words & POWERUP_BIT) != 0, - beacon == BEACON_CURRENT, - beacon == BEACON_OPPONENT, - hill == HILL_NEUTRAL, - hill == HILL_CURRENT, - hill == HILL_OPPONENT, - (words & CURRENT_PLAYER_BIT) != 0, - (words & OPPONENT_PLAYER_BIT) != 0, - ], - dim=1, - ) - return planes.to(torch.float32).reshape(-1, BOARD_PLANES, BOARD_SIDE, BOARD_SIDE) - - -def decode_packed_board( - board_words: torch.Tensor, - *, - dtype: torch.dtype = torch.bfloat16, - out: torch.Tensor | None = None, -) -> torch.Tensor: - if board_words.dtype != torch.uint16: - raise ValueError(f"board_words must be uint16, got {board_words.dtype}") - if board_words.ndim != 2 or board_words.shape[1] != BOARD_CELLS: - raise ValueError( - f"board_words must have shape (B, {BOARD_CELLS}), got {tuple(board_words.shape)}" - ) +def _unpack_bits(words: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + """Extract TILE_PLANES binary planes from u16 words. - if board_words.device.type != "cuda": - decoded = decode_packed_board_reference(board_words).to(dtype=dtype) - if out is not None: - out.copy_(decoded) - return out - return decoded + Args: + words: [...] int32 tensor of packed bitmasks + shift: [TILE_PLANES] int32 tensor = arange(TILE_PLANES) - board_words = board_words.contiguous() - expected_shape = (board_words.shape[0], BOARD_PLANES, BOARD_SIDE, BOARD_SIDE) - if out is None: - out = torch.empty(expected_shape, device=board_words.device, dtype=dtype) - else: - if out.shape != expected_shape: - raise ValueError( - f"out must have shape {expected_shape}, got {tuple(out.shape)}" - ) - if out.device != board_words.device: - raise ValueError("out must be on the same device as board_words") - if out.dtype != dtype: - raise ValueError(f"out must have dtype {dtype}, got {out.dtype}") - total_cells = board_words.shape[0] * BOARD_CELLS - grid = lambda meta: (triton.cdiv(total_cells, meta["BLOCK"]),) - with torch.cuda.device(board_words.device): - _decode_board_kernel[grid]( - board_words, - out, - total_cells, - board_words.stride(0), - out.stride(0), - out.stride(1), - out.stride(2), - out.stride(3), - BLOCK=256, - num_warps=4, - ) - return out + Returns: + [..., TILE_PLANES] float tensor of 0/1 values + """ + return ((words.unsqueeze(-1) >> shift) & 1).float() -def decode_intrinsics( +def decode_nnue_obs( packed_obs: torch.Tensor, *, + plane_shift: torch.Tensor, + opp_perm: torch.Tensor, + global_scales: torch.Tensor, dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - _check_packed_obs(packed_obs) - scales = torch.tensor( - INTRINSIC_SCALE, device=packed_obs.device, dtype=torch.float32 - ) - intrinsics = packed_obs[:, BOARD_CELLS:].to(torch.float32) - intrinsics = intrinsics / scales - return intrinsics.to(dtype=dtype) - - -def decode_packed_observation( - packed_obs: torch.Tensor, - *, - board_dtype: torch.dtype = torch.bfloat16, - intrinsic_dtype: torch.dtype | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - _check_packed_obs(packed_obs) - board = decode_packed_board(packed_obs[:, :BOARD_CELLS], dtype=board_dtype) - intrinsics = decode_intrinsics( - packed_obs, - dtype=board_dtype if intrinsic_dtype is None else intrinsic_dtype, - ) - return board, intrinsics +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Decode packed NNUE observation into model inputs. + + All buffers (plane_shift, opp_perm, global_scales) should be registered + as model buffers so they live on the right device and are captured by + CUDA graphs. + + Returns: + my_features: [B, 15360] dtype — accumulator input (my perspective) + opp_features: [B, 15360] dtype — accumulator input (opponent perspective) + local_my: [B, 375] dtype — local window around my position + local_opp: [B, 375] dtype — local window around opponent position + globals_norm: [B, 20] dtype — normalized global scalars + turn_count: [B] int32 — for white-to-move sign + """ + B = packed_obs.shape[0] + + # --- Tile features: [B, 1024] u16 -> [B, 15360] ---------------------- + tiles = packed_obs[:, :BOARD_CELLS].to(torch.int32) + bits = _unpack_bits(tiles, plane_shift) # [B, 1024, 15] + my_features = bits.reshape(B, -1).to(dtype) + + # Opponent perspective: permute planes + opp_features = bits[:, :, opp_perm].reshape(B, -1).to(dtype) + + # --- Local windows: [B, 50] u16 -> [B, 375] each --------------------- + locals_raw = packed_obs[:, OFFSET_LOCALS:OFFSET_POSITIONS].to(torch.int32) + local_bits = _unpack_bits(locals_raw, plane_shift) # [B, 50, 15] + local_my = local_bits[:, :LOCAL_WINDOW_TILES, :].reshape(B, -1).to(dtype) + local_opp = local_bits[:, LOCAL_WINDOW_TILES:, :].reshape(B, -1).to(dtype) + + # --- Global features: [B, 20] u16 -> [B, 20] float normalized -------- + globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) + globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) + + # --- Turn count for white-to-move sign -------------------------------- + turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to(torch.int32) + + return my_features, opp_features, local_my, local_opp, globals_norm, turn_count __all__ = [ "BOARD_CELLS", - "BOARD_PLANES", "BOARD_SIDE", - "INTRINSIC_SCALE", - "INTRINSIC_COUNT", - "INTRINSIC_TURN_COUNT", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "GLOBAL_TURN_COUNT_IDX", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", "OBS_WORDS", - "decode_intrinsics", - "decode_packed_board", - "decode_packed_board_reference", - "decode_packed_observation", + "OPP_PLANE_PERM", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", + "decode_nnue_obs", ] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 8c31fce..41a6168 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -4,7 +4,6 @@ import json import os import time -from contextlib import nullcontext from dataclasses import asdict, dataclass, field from pathlib import Path from typing import cast @@ -78,16 +77,12 @@ class TrainConfig: pretrain_lr: float | None = None selfplay_lr: float | None = 5e-5 weight_decay: float = 1e-4 - width: int = 48 - num_blocks: int = 6 - hidden_dim: int = 256 - pretrain_terminal_samples: int = 7_500_000 + pretrain_terminal_samples: int = 4_000_000 pretrain_batch_size: int = 8_192 pretrain_log_interval: int = 10 seed: int = 42 selfplay_precision: str = "bf16" num_gpus: int = field(default_factory=_default_num_gpus) - device: str = "cuda" checkpoint_interval: int = 2 run_dir: str = "runs/latest" wandb: bool = True @@ -107,25 +102,9 @@ class TrainStepResult: loss: torch.Tensor sample_seconds: float h2d_seconds: float - forward_seconds: float = 0.0 - backward_seconds: float = 0.0 - optimizer_seconds: float = 0.0 - forward_timing: CudaSectionTiming | None = None - backward_timing: CudaSectionTiming | None = None - optimizer_timing: CudaSectionTiming | None = None - - -def _device(device: str) -> torch.device: - if device == "cuda" and not torch.cuda.is_available(): - raise RuntimeError("CUDA is required for the training runner") - if device == "cuda": - return torch.device("cuda:0") - return torch.device(device) - - -def _to_channels_last(module: torch.nn.Module) -> torch.nn.Module: - module = module.to(memory_format=torch.channels_last) # type: ignore[call-overload] - return module + forward_timing: CudaSectionTiming + backward_timing: CudaSectionTiming + optimizer_timing: CudaSectionTiming def _numpy_batch_to_device( @@ -186,75 +165,44 @@ def _train_step_from_batch( target: torch.Tensor, sample_seconds: float, h2d_seconds: float, - *, - device: torch.device, ) -> TrainStepResult: model.train() optimizer.zero_grad(set_to_none=True) - autocast = ( - torch.autocast(device_type="cuda", dtype=torch.bfloat16) - if device.type == "cuda" - else nullcontext() - ) - if device.type == "cuda": - forward_timing = CudaSectionTiming( - start=torch.cuda.Event(enable_timing=True), - end=torch.cuda.Event(enable_timing=True), - ) - backward_timing = CudaSectionTiming( - start=torch.cuda.Event(enable_timing=True), - end=torch.cuda.Event(enable_timing=True), - ) - optimizer_timing = CudaSectionTiming( - start=torch.cuda.Event(enable_timing=True), - end=torch.cuda.Event(enable_timing=True), - ) - - forward_timing.start.record() - with autocast: - pred = model(obs) - loss = F.mse_loss(pred.float(), target) - forward_timing.end.record() - - backward_timing.start.record() - loss.backward() - backward_timing.end.record() - - optimizer_timing.start.record() - optimizer.step() - optimizer_timing.end.record() - - return TrainStepResult( - loss=loss.detach(), - sample_seconds=sample_seconds, - h2d_seconds=h2d_seconds, - forward_timing=forward_timing, - backward_timing=backward_timing, - optimizer_timing=optimizer_timing, - ) + forward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + backward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + optimizer_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) - forward_started_at = time.perf_counter() - with autocast: + forward_timing.start.record() + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pred = model(obs) - loss = F.mse_loss(pred.float(), target) - forward_seconds = time.perf_counter() - forward_started_at + loss = F.huber_loss(pred.float(), target) + forward_timing.end.record() - backward_started_at = time.perf_counter() + backward_timing.start.record() loss.backward() - backward_seconds = time.perf_counter() - backward_started_at + backward_timing.end.record() - optimizer_started_at = time.perf_counter() + optimizer_timing.start.record() optimizer.step() - optimizer_seconds = time.perf_counter() - optimizer_started_at + optimizer_timing.end.record() return TrainStepResult( loss=loss.detach(), sample_seconds=sample_seconds, h2d_seconds=h2d_seconds, - forward_seconds=forward_seconds, - backward_seconds=backward_seconds, - optimizer_seconds=optimizer_seconds, + forward_timing=forward_timing, + backward_timing=backward_timing, + optimizer_timing=optimizer_timing, ) @@ -271,13 +219,7 @@ def train_step( replay_buffer, batch_size, seed, device ) return _train_step_from_batch( - model, - optimizer, - obs, - target, - sample_seconds, - h2d_seconds, - device=device, + model, optimizer, obs, target, sample_seconds, h2d_seconds ) @@ -293,13 +235,7 @@ def pretrain_step( batch_size, seed, device ) return _train_step_from_batch( - model, - optimizer, - obs, - target, - sample_seconds, - h2d_seconds, - device=device, + model, optimizer, obs, target, sample_seconds, h2d_seconds ) @@ -366,17 +302,11 @@ def _collect_round_diagnostics( obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( replay_buffer, batch_size, seed, device ) - autocast = ( - torch.autocast(device_type="cuda", dtype=torch.bfloat16) - if device.type == "cuda" - else nullcontext() - ) model.eval() with torch.no_grad(): - with autocast: + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): pred = model(obs) - if device.type == "cuda": - torch.cuda.synchronize(device) + torch.cuda.synchronize(device) pred = pred.float() residual = pred - target @@ -400,32 +330,25 @@ def _collect_round_diagnostics( def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: - device = _device(config.device) + device = torch.device("cuda:0") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the training runner") torch.manual_seed(config.seed) pretrain_lr = config.lr if config.pretrain_lr is None else config.pretrain_lr selfplay_lr = config.lr if config.selfplay_lr is None else config.selfplay_lr - if device.type == "cuda": - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True run_dir, checkpoint_dir = _prepare_run_dir(config) model = cast( PackedValueModel, - _to_channels_last( - PackedValueModel( - width=config.width, - num_blocks=config.num_blocks, - hidden_dim=config.hidden_dim, - ).to(device) - ), + PackedValueModel().to(device), ) model.eval() param_count = sum(param.numel() for param in model.parameters()) - print( - f"model width={config.width} blocks={config.num_blocks} hidden={config.hidden_dim} params={param_count}" - ) + print(f"model (nnue) params={param_count}") initial_lr = pretrain_lr if config.pretrain_terminal_samples > 0 else selfplay_lr optimizer = torch.optim.AdamW( @@ -448,7 +371,6 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: pretrain_steps = ( config.pretrain_terminal_samples + config.pretrain_batch_size - 1 ) // config.pretrain_batch_size - pretrain_log_interval = max(1, config.pretrain_log_interval) logger.start_pretrain() for step_idx in range(pretrain_steps): samples_done = step_idx * config.pretrain_batch_size @@ -527,7 +449,7 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: device=device, ) round_results.append(result) - if device.type == "cuda" and round_results: + if round_results: torch.cuda.synchronize(device) model.eval() train_seconds = time.perf_counter() - train_started_at @@ -537,30 +459,14 @@ def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: ) train_h2d_seconds = sum(result.h2d_seconds for result in round_results) train_forward_seconds = sum( - result.forward_seconds for result in round_results + result.forward_timing.seconds() for result in round_results ) train_backward_seconds = sum( - result.backward_seconds for result in round_results + result.backward_timing.seconds() for result in round_results ) train_optimizer_seconds = sum( - result.optimizer_seconds for result in round_results + result.optimizer_timing.seconds() for result in round_results ) - if device.type == "cuda": - train_forward_seconds += sum( - result.forward_timing.seconds() - for result in round_results - if result.forward_timing is not None - ) - train_backward_seconds += sum( - result.backward_timing.seconds() - for result in round_results - if result.backward_timing is not None - ) - train_optimizer_seconds += sum( - result.optimizer_timing.seconds() - for result in round_results - if result.optimizer_timing is not None - ) if round_results: round_loss_values = ( @@ -650,9 +556,6 @@ def _parse_args() -> TrainConfig: parser.add_argument("--pretrain-lr", type=float, default=defaults.pretrain_lr) parser.add_argument("--selfplay-lr", type=float, default=defaults.selfplay_lr) parser.add_argument("--weight-decay", type=float, default=defaults.weight_decay) - parser.add_argument("--width", type=int, default=defaults.width) - parser.add_argument("--num-blocks", type=int, default=defaults.num_blocks) - parser.add_argument("--hidden-dim", type=int, default=defaults.hidden_dim) parser.add_argument( "--pretrain-terminal-samples", type=int, @@ -680,7 +583,6 @@ def _parse_args() -> TrainConfig: default=defaults.num_gpus, help="Number of GPUs for self-play (0 uses all visible CUDA GPUs; defaults to Slurm allocation when available)", ) - parser.add_argument("--device", default=defaults.device) parser.add_argument( "--checkpoint-interval", type=int, default=defaults.checkpoint_interval ) diff --git a/training/src/observation.rs b/training/src/observation.rs index fb25fd4..22ec701 100644 --- a/training/src/observation.rs +++ b/training/src/observation.rs @@ -1,188 +1,285 @@ -use alpha_paint::board::board_structs::Player; +use alpha_paint::board::board_structs::{HillData, Player}; +use alpha_paint::board::consts::{BASE_MAX_STAMINA, HILL_MAX_STAMINA_BONUS}; use alpha_paint::board::structs::Coordinate; use alpha_paint::board::Board; pub const BOARD_SIDE: usize = 32; pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; -pub const INTRINSIC_COUNT: usize = 10; -pub const OBS_WORDS: usize = BOARD_CELLS + INTRINSIC_COUNT; -pub const BOARD_PLANES: usize = 17; - -pub const PAINT_STRENGTH_MASK: u16 = 0b111; -pub const PAINT_IS_ENEMY_BIT: u16 = 1 << 3; -pub const WALL_BIT: u16 = 1 << 4; -pub const POWERUP_BIT: u16 = 1 << 5; -pub const BEACON_SHIFT: u16 = 6; -pub const BEACON_MASK: u16 = 0b11 << BEACON_SHIFT; -pub const HILL_SHIFT: u16 = 8; -pub const HILL_MASK: u16 = 0b11 << HILL_SHIFT; -pub const CURRENT_PLAYER_BIT: u16 = 1 << 10; -pub const OPPONENT_PLAYER_BIT: u16 = 1 << 11; - -pub const BEACON_NONE: u16 = 0; -pub const BEACON_CURRENT: u16 = 1; -pub const BEACON_OPPONENT: u16 = 2; - -pub const HILL_NONE: u16 = 0; -pub const HILL_NEUTRAL: u16 = 1; -pub const HILL_CURRENT: u16 = 2; -pub const HILL_OPPONENT: u16 = 3; - -pub const INTRINSIC_CURRENT_STAMINA: usize = 0; -pub const INTRINSIC_OPPONENT_STAMINA: usize = 1; -pub const INTRINSIC_CURRENT_HILLS: usize = 2; -pub const INTRINSIC_OPPONENT_HILLS: usize = 3; -pub const INTRINSIC_TURN_COUNT: usize = 4; -pub const INTRINSIC_CURRENT_TERRITORY: usize = 5; -pub const INTRINSIC_OPPONENT_TERRITORY: usize = 6; -pub const INTRINSIC_CURRENT_BEACONS: usize = 7; -pub const INTRINSIC_OPPONENT_BEACONS: usize = 8; -pub const INTRINSIC_CONSECUTIVE_MOVES: usize = 9; - -#[inline] +pub const TILE_PLANES: usize = 15; +pub const LOCAL_WINDOW_TILES: usize = 25; +pub const GLOBAL_FEATURES: usize = 20; + +/// 32×32 tiles + 2×25 local windows + 4 positions + 20 globals +pub const OBS_WORDS: usize = BOARD_CELLS + (LOCAL_WINDOW_TILES * 2) + 4 + GLOBAL_FEATURES; + +const OFFSET_LOCALS: usize = BOARD_CELLS; +const OFFSET_POSITIONS: usize = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2; +const OFFSET_GLOBALS: usize = OFFSET_POSITIONS + 4; + +/// Index of turn_count within the global features section. +pub const GLOBAL_TURN_COUNT_OFFSET: usize = 17; + +/// Encodes the observation to ship off to the GPU for training. +/// +/// Layout (all from current player's perspective): +/// [0, 1024) — per-tile 15-bit bitmask (32×32 board) +/// [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) +/// [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) +/// [1078, 1098) — 20 global scalar features +/// +/// Tile bitmask planes: +/// bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) +/// bits 4-7: opponent paint thermometer +/// bit 8: wall +/// bit 9: powerup +/// bit 10: current player beacon +/// bit 11: opponent beacon +/// bit 12: hill neutral +/// bit 13: hill current player +/// bit 14: hill opponent pub fn encode_into_slice(board: &Board, out: &mut [u16]) { - assert_eq!( - out.len(), - OBS_WORDS, - "packed observation buffer size mismatch" - ); + debug_assert!(out.len() >= OBS_WORDS); - out.fill(0); - out[..BOARD_CELLS].fill(WALL_BIT); - - let current_is_white = board.is_white_turn(); - let current_coord = board.current_player_coord(); - let opponent_coord = if current_is_white { - board.black_coord + let is_white = board.is_white_turn(); + let (my_coord, opp_coord) = if is_white { + (board.white_coord, board.black_coord) + } else { + (board.black_coord, board.white_coord) + }; + let (my_stamina, opp_stamina) = if is_white { + (board.white_stamina, board.black_stamina) } else { - board.white_coord + (board.black_stamina, board.white_stamina) }; - for y in 0..board.rows { - for x in 0..board.cols { + let hill_meta = board.tiles.hill_metadata(); + + // 1. Tile bitmasks [0, 1024) + for x in 0..32u8 { + for y in 0..32u8 { let coord = Coordinate::new(x, y); - let tile = board.tiles[coord]; - - let mut word = 0u16; - let paint = tile.paint_value(); - if paint != 0 { - word |= (paint.unsigned_abs() as u16) & PAINT_STRENGTH_MASK; - let is_enemy_paint = if current_is_white { - paint < 0 - } else { - paint > 0 - }; - if is_enemy_paint { - word |= PAINT_IS_ENEMY_BIT; - } - } + out[(x as usize) * 32 + (y as usize)] = encode_tile(board, is_white, hill_meta, coord); + } + } - if tile.is_wall() { - word |= WALL_BIT; - } - if board.powerups[coord] { - word |= POWERUP_BIT; - } + // 2. Local windows [1024, 1074) + encode_local_window( + board, + is_white, + hill_meta, + my_coord, + &mut out[OFFSET_LOCALS..], + ); + encode_local_window( + board, + is_white, + hill_meta, + opp_coord, + &mut out[OFFSET_LOCALS + 25..], + ); - let beacon_state = match tile.beacon_owner() { - Some(Player::White) if current_is_white => BEACON_CURRENT, - Some(Player::White) => BEACON_OPPONENT, - Some(Player::Black) if current_is_white => BEACON_OPPONENT, - Some(Player::Black) => BEACON_CURRENT, - None => BEACON_NONE, - }; - word |= beacon_state << BEACON_SHIFT; - - let hill_state = match board.hill_id[coord] { - u16::MAX => HILL_NONE, - hill_id => match board.tiles.hill_metadata()[hill_id as usize].owner { - None => HILL_NEUTRAL, - Some(Player::White) if current_is_white => HILL_CURRENT, - Some(Player::White) => HILL_OPPONENT, - Some(Player::Black) if current_is_white => HILL_OPPONENT, - Some(Player::Black) => HILL_CURRENT, - }, - }; - word |= hill_state << HILL_SHIFT; + // 3. Player positions [1074, 1078) + out[OFFSET_POSITIONS] = my_coord.x as u16; + out[OFFSET_POSITIONS + 1] = my_coord.y as u16; + out[OFFSET_POSITIONS + 2] = opp_coord.x as u16; + out[OFFSET_POSITIONS + 3] = opp_coord.y as u16; - if coord == current_coord { - word |= CURRENT_PLAYER_BIT; - } - if coord == opponent_coord { - word |= OPPONENT_PLAYER_BIT; - } + // 4. Global features [1078, 1098) + let my_hills = if is_white { + board.tiles.controlled_hill_count::() + } else { + board.tiles.controlled_hill_count::() + }; + let opp_hills = if is_white { + board.tiles.controlled_hill_count::() + } else { + board.tiles.controlled_hill_count::() + }; + + let my_territory = if is_white { + board.tiles.territory_count::() + } else { + board.tiles.territory_count::() + }; + let opp_territory = if is_white { + board.tiles.territory_count::() + } else { + board.tiles.territory_count::() + }; - out[cell_index(x, y)] = word; + let mut my_hill_tiles: usize = 0; + let mut opp_hill_tiles: usize = 0; + let mut contested_hills: usize = 0; + for hd in hill_meta { + let (mine, theirs) = if is_white { + (hd.white_count, hd.black_count) + } else { + (hd.black_count, hd.white_count) + }; + my_hill_tiles += mine; + opp_hill_tiles += theirs; + if mine > 0 && theirs > 0 && hd.owner.is_none() { + contested_hills += 1; } } - let tail = &mut out[BOARD_CELLS..]; - if current_is_white { - tail[INTRINSIC_CURRENT_STAMINA] = board.white_stamina.try_into().unwrap(); - tail[INTRINSIC_OPPONENT_STAMINA] = board.black_stamina.try_into().unwrap(); - tail[INTRINSIC_CURRENT_HILLS] = board - .tiles - .controlled_hill_count::() - .try_into() - .unwrap(); - tail[INTRINSIC_OPPONENT_HILLS] = board - .tiles - .controlled_hill_count::() - .try_into() - .unwrap(); - tail[INTRINSIC_CURRENT_TERRITORY] = - board.tiles.territory_count::().try_into().unwrap(); - tail[INTRINSIC_OPPONENT_TERRITORY] = - board.tiles.territory_count::().try_into().unwrap(); - tail[INTRINSIC_CURRENT_BEACONS] = board - .tiles - .get_beacon_iterator::() - .count() - .try_into() - .unwrap(); - tail[INTRINSIC_OPPONENT_BEACONS] = board - .tiles - .get_beacon_iterator::() - .count() - .try_into() - .unwrap(); + let my_beacons = if is_white { + board.tiles.get_beacon_iterator::().count() } else { - tail[INTRINSIC_CURRENT_STAMINA] = board.black_stamina.try_into().unwrap(); - tail[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina.try_into().unwrap(); - tail[INTRINSIC_CURRENT_HILLS] = board - .tiles - .controlled_hill_count::() - .try_into() - .unwrap(); - tail[INTRINSIC_OPPONENT_HILLS] = board - .tiles - .controlled_hill_count::() - .try_into() - .unwrap(); - tail[INTRINSIC_CURRENT_TERRITORY] = - board.tiles.territory_count::().try_into().unwrap(); - tail[INTRINSIC_OPPONENT_TERRITORY] = - board.tiles.territory_count::().try_into().unwrap(); - tail[INTRINSIC_CURRENT_BEACONS] = board - .tiles - .get_beacon_iterator::() - .count() - .try_into() - .unwrap(); - tail[INTRINSIC_OPPONENT_BEACONS] = board - .tiles - .get_beacon_iterator::() - .count() - .try_into() - .unwrap(); - } - tail[INTRINSIC_TURN_COUNT] = board.turn_count.try_into().unwrap(); - tail[INTRINSIC_CONSECUTIVE_MOVES] = board.consecutives_moves_so_far.try_into().unwrap(); + board.tiles.get_beacon_iterator::().count() + }; + let opp_beacons = if is_white { + board.tiles.get_beacon_iterator::().count() + } else { + board.tiles.get_beacon_iterator::().count() + }; + + let player_dist = board.dist[(my_coord, opp_coord)]; + let my_hill_dist = nearest_non_controlled_hill_dist(board, my_coord, is_white, hill_meta); + let opp_hill_dist = nearest_non_controlled_hill_dist(board, opp_coord, !is_white, hill_meta); + + let my_max_stamina = BASE_MAX_STAMINA + my_hills * HILL_MAX_STAMINA_BONUS; + let opp_max_stamina = BASE_MAX_STAMINA + opp_hills * HILL_MAX_STAMINA_BONUS; + + let g = &mut out[OFFSET_GLOBALS..]; + g[0] = my_stamina as u16; + g[1] = opp_stamina as u16; + g[2] = my_max_stamina as u16; + g[3] = opp_max_stamina as u16; + g[4] = my_territory as u16; + g[5] = opp_territory as u16; + g[6] = my_hills as u16; + g[7] = opp_hills as u16; + g[8] = my_hill_tiles as u16; + g[9] = opp_hill_tiles as u16; + g[10] = contested_hills as u16; + g[11] = my_beacons as u16; + g[12] = opp_beacons as u16; + g[13] = player_dist; + g[14] = my_hill_dist; + g[15] = opp_hill_dist; + g[16] = board.consecutives_moves_so_far as u16; + g[17] = board.turn_count as u16; + g[18] = board.rows as u16; + g[19] = board.cols as u16; } #[inline] -pub const fn cell_index(x: u8, y: u8) -> usize { - y as usize * BOARD_SIDE + x as usize +fn encode_tile(board: &Board, is_white: bool, hill_meta: &[HillData], coord: Coordinate) -> u16 { + let tile = board.tiles[coord]; + let mut bits: u16 = 0; + + let paint = tile.paint_value(); + let (my_paint, opp_paint) = if is_white { + (paint.max(0) as u16, (-paint).max(0) as u16) + } else { + ((-paint).max(0) as u16, paint.max(0) as u16) + }; + // Thermometer encoding: bits 0-3 for my paint, bits 4-7 for opponent paint + for t in 0..my_paint { + bits |= 1 << t; + } + for t in 0..opp_paint { + bits |= 1 << (4 + t); + } + + if tile.is_wall() { + bits |= 1 << 8; + } + if board.powerups[coord] { + bits |= 1 << 9; + } + + match tile.beacon_owner() { + Some(Player::White) if is_white => bits |= 1 << 10, + Some(Player::Black) if !is_white => bits |= 1 << 10, + Some(Player::White) => bits |= 1 << 11, + Some(Player::Black) => bits |= 1 << 11, + None => {} + } + + let hill_id = board.hill_id[coord]; + if hill_id != u16::MAX { + match hill_meta[hill_id as usize].owner { + None => bits |= 1 << 12, + Some(Player::White) if is_white => bits |= 1 << 13, + Some(Player::Black) if !is_white => bits |= 1 << 13, + _ => bits |= 1 << 14, + } + } + + bits +} + +const WALL_BITS: u16 = 1 << 8; + +/// Write 25 u16 bitmasks for the manhattan-distance-3 diamond around `center`. +/// Out-of-bounds tiles default to the wall bit. +/// +/// Enumeration order: dx from -3..=3, dy from -3..=3, keeping |dx|+|dy| <= 3. +fn encode_local_window( + board: &Board, + is_white: bool, + hill_meta: &[HillData], + center: Coordinate, + out: &mut [u16], +) { + let cx = center.x as i16; + let cy = center.y as i16; + let mut i = 0; + for dx in -3i16..=3 { + for dy in -3i16..=3 { + if dx.abs() + dy.abs() > 3 { + continue; + } + let nx = cx + dx; + let ny = cy + dy; + out[i] = if (0..32).contains(&nx) && (0..32).contains(&ny) { + encode_tile( + board, + is_white, + hill_meta, + Coordinate::new(nx as u8, ny as u8), + ) + } else { + WALL_BITS + }; + i += 1; + } + } + debug_assert_eq!(i, 25); +} + +/// Find the minimum distance from `player_coord` to any tile in a hill +/// that is NOT controlled by the player identified by `is_white`. +fn nearest_non_controlled_hill_dist( + board: &Board, + player_coord: Coordinate, + is_white: bool, + hill_meta: &[HillData], +) -> u16 { + let my_owner = if is_white { + Some(Player::White) + } else { + Some(Player::Black) + }; + + let mut best = u16::MAX; + + for (i, hd) in hill_meta.iter().enumerate() { + // Skip hills we already control + if hd.owner == my_owner { + continue; + } + for &tile_coord in &board.hills[i] { + let d = board.dist[(player_coord, tile_coord)]; + if d < best { + best = d; + } + } + } + + best } #[cfg(test)] @@ -197,61 +294,88 @@ mod tests { } #[test] - fn packs_board_bits_and_intrinsics_for_white_turn() { - let words = encode_words( - "ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:1B1/d1O/3", - ); - - assert_eq!(words[cell_index(1, 0)] & PAINT_STRENGTH_MASK, 2); - assert_eq!(words[cell_index(1, 0)] & PAINT_IS_ENEMY_BIT, 0); - - let enemy_paint = words[cell_index(0, 1)]; - assert_eq!(enemy_paint & PAINT_STRENGTH_MASK, 4); - assert_ne!(enemy_paint & PAINT_IS_ENEMY_BIT, 0); - assert_ne!(enemy_paint & POWERUP_BIT, 0); - - let hill_word = words[cell_index(1, 1)]; - assert_eq!((hill_word & HILL_MASK) >> HILL_SHIFT, HILL_NEUTRAL); - - let beacon_word = words[cell_index(2, 1)]; - assert_eq!((beacon_word & BEACON_MASK) >> BEACON_SHIFT, BEACON_CURRENT); - - assert_ne!(words[cell_index(0, 0)] & CURRENT_PLAYER_BIT, 0); - assert_ne!(words[cell_index(2, 2)] & OPPONENT_PLAYER_BIT, 0); - - assert_ne!(words[cell_index(31, 31)] & WALL_BIT, 0); - - let tail = &words[BOARD_CELLS..]; - assert_eq!(tail[INTRINSIC_CURRENT_STAMINA], 99); - assert_eq!(tail[INTRINSIC_OPPONENT_STAMINA], 88); - assert_eq!(tail[INTRINSIC_CURRENT_HILLS], 0); - assert_eq!(tail[INTRINSIC_OPPONENT_HILLS], 0); - assert_eq!(tail[INTRINSIC_TURN_COUNT], 0); - assert_eq!(tail[INTRINSIC_CURRENT_TERRITORY], 1); - assert_eq!(tail[INTRINSIC_OPPONENT_TERRITORY], 1); - assert_eq!(tail[INTRINSIC_CURRENT_BEACONS], 1); - assert_eq!(tail[INTRINSIC_OPPONENT_BEACONS], 0); - assert_eq!(tail[INTRINSIC_CONSECUTIVE_MOVES], 2); + fn obs_words_matches_expected_layout() { + assert_eq!(OBS_WORDS, 1024 + 50 + 4 + 20); + assert_eq!(OBS_WORDS, 1098); } #[test] - fn keeps_absolute_locations_on_black_turn() { + fn tile_thermometer_paint_encoding() { + // White turn, white paint of strength 3 at (1,0) via 'B' = 2 strength let words = - encode_words("ap2|3x3|tc:1|cm:5|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/2o"); + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:1B1/3/3"); + // White paint strength 2 at (1,0) -> bits 0,1 set (my paint >=1, >=2) + let tile_word = words[1 * 32 + 0]; + assert_eq!(tile_word & 0b1111, 0b0011); // bits 0,1 = my paint thermometer + assert_eq!((tile_word >> 4) & 0b1111, 0); // no opponent paint + } - let white_paint = words[cell_index(1, 0)]; - assert_eq!(white_paint & PAINT_STRENGTH_MASK, 2); - assert_ne!(white_paint & PAINT_IS_ENEMY_BIT, 0); + #[test] + fn tile_wall_and_powerup_bits() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:3/3/3"); + // (0,1) has powerup + let pu_word = words[0 * 32 + 1]; + assert_ne!(pu_word & (1 << 9), 0); + + // out-of-bounds cell should be wall + let oob_word = words[31 * 32 + 31]; + assert_ne!(oob_word & (1 << 8), 0); + } - let hill_word = words[cell_index(1, 1)]; - assert_eq!((hill_word & HILL_MASK) >> HILL_SHIFT, HILL_CURRENT); + #[test] + fn global_features_in_correct_positions() { + let words = + encode_words("ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + let g = &words[OFFSET_GLOBALS..]; + assert_eq!(g[0], 99); // my_stamina (white turn) + assert_eq!(g[1], 88); // opp_stamina + assert_eq!(g[16], 2); // consecutive_moves + assert_eq!(g[17], 0); // turn_count + assert_eq!(g[18], 3); // rows + assert_eq!(g[19], 3); // cols + } - assert_ne!(words[cell_index(2, 2)] & CURRENT_PLAYER_BIT, 0); - assert_ne!(words[cell_index(0, 0)] & OPPONENT_PLAYER_BIT, 0); + #[test] + fn local_window_has_25_tiles() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + // Local windows are at OFFSET_LOCALS..OFFSET_LOCALS+50 + let my_local = &words[OFFSET_LOCALS..OFFSET_LOCALS + 25]; + let opp_local = &words[OFFSET_LOCALS + 25..OFFSET_LOCALS + 50]; + // Both should have 25 entries, out-of-bounds tiles = WALL_BITS + assert_eq!(my_local.len(), 25); + assert_eq!(opp_local.len(), 25); + } - let tail = &words[BOARD_CELLS..]; - assert_eq!(tail[INTRINSIC_CURRENT_STAMINA], 88); - assert_eq!(tail[INTRINSIC_OPPONENT_STAMINA], 99); - assert_eq!(tail[INTRINSIC_CONSECUTIVE_MOVES], 5); + #[test] + fn player_positions_encoded() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + // White turn: my=(0,0), opp=(2,2) + assert_eq!(words[OFFSET_POSITIONS], 0); // my_x + assert_eq!(words[OFFSET_POSITIONS + 1], 0); // my_y + assert_eq!(words[OFFSET_POSITIONS + 2], 2); // opp_x + assert_eq!(words[OFFSET_POSITIONS + 3], 2); // opp_y + } + + #[test] + fn black_turn_flips_perspective() { + let words = + encode_words("ap2|3x3|tc:1|cm:0|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/3"); + // Black turn: my=(2,2), opp=(0,0) + assert_eq!(words[OFFSET_POSITIONS], 2); // my_x (black) + assert_eq!(words[OFFSET_POSITIONS + 1], 2); // my_y + assert_eq!(words[OFFSET_POSITIONS + 2], 0); // opp_x (white) + assert_eq!(words[OFFSET_POSITIONS + 3], 0); // opp_y + + // White paint at (1,0) should be opponent paint from black's perspective + let tile_word = words[1 * 32 + 0]; + assert_eq!(tile_word & 0b1111, 0); // no my paint (black has no paint here) + assert_ne!((tile_word >> 4) & 0b1111, 0); // opponent paint (white's paint) + + // Hill owned by black should be "my hill" from black's perspective + let hill_word = words[1 * 32 + 1]; + assert_ne!(hill_word & (1 << 13), 0); // hill current player } } From a06371ca7362053f75293de35dce35f8641069c9 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 03:30:27 -0400 Subject: [PATCH 53/59] try tanh nnue --- python/alphapaint_training/model.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 8227224..175d7b6 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -21,9 +21,7 @@ from .packed_obs import ( GLOBAL_FEATURES, GLOBAL_SCALES, - GLOBAL_TURN_COUNT_IDX, LOCAL_FEATURES, - OFFSET_GLOBALS, OPP_PLANE_PERM, TILE_PLANES, TOTAL_TILE_FEATURES, @@ -180,7 +178,7 @@ def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: # Convert from current-player perspective to White perspective. white_to_move = (turn_count & 1) == 0 sign = torch.where(white_to_move, 1.0, -1.0).to(dtype=value.dtype) - return value * sign + return (value * sign).tanh() __all__ = ["NnueValueNet", "PackedValueModel"] From 91f33a734ad036666dd17a7a7c7c48bf82cc8a2d Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Mon, 30 Mar 2026 09:37:32 -0400 Subject: [PATCH 54/59] fix weird bugs with training data not being bounded --- AGENTS.md | 2 ++ python/alphapaint_training/train.py | 4 ++-- training/src/descent.rs | 31 ++++++++++++++++++++++++++--- training/src/replay_buffer.rs | 2 +- 4 files changed, 33 insertions(+), 6 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ac2ff11..d364e10 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -11,6 +11,8 @@ check out /mnt/github/TheConverseEngineer/AlphaSnake/ for last years impl `just test` runs tests. u should use this rather than doign it yourself bc this auto installs the package +- this just tests the rust board impl. do not use it for eg training code changes + `nix develop` GETS YOU CUDA `maturin develop --release` to bring python bindings in. you can `--manifest-path training/Cargo.toml` etc to choose the package to reinstall `PYTHONPATH=python` is uesful sometimes diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py index 41a6168..0d55710 100644 --- a/python/alphapaint_training/train.py +++ b/python/alphapaint_training/train.py @@ -66,13 +66,13 @@ def _default_num_gpus() -> int: @dataclass(slots=True) class TrainConfig: rounds: int = 200 - samples_per_round: int = 1_048_576 + samples_per_round: int = 1024 * 1024 * 2 train_steps_per_round: int = 256 batch_size: int = 8_192 replay_capacity: int = 16_000_000 num_threads: int = field(default_factory=_default_num_threads) workers_per_thread: int = 8 - max_gpu_evals_per_move: int = 4 * 4096 + max_gpu_evals_per_move: int = 16 * 4096 lr: float = 3e-4 pretrain_lr: float | None = None selfplay_lr: float | None = 5e-5 diff --git a/training/src/descent.rs b/training/src/descent.rs index d752342..327d1b0 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -441,7 +441,7 @@ impl SearchNode { let p = board.turn_count.max(1) as f32; let depth = (AVG_GAME_LENGTH / p).ln_1p() / AVG_GAME_LENGTH.ln_1p(); - sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) + (sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth)).clamp(-1.0, 1.0) } async fn ubfms_iteration( @@ -532,6 +532,27 @@ impl SearchNode { self.value } + /// Compute the minimax value considering only expanded children. + /// This is used for training targets to avoid including unexpanded + /// children's raw network predictions in the value signal. + fn compute_backed_value(&self, is_white_turn: bool) -> f32 { + if is_white_turn { + self.children + .iter() + .filter(|c| c.node.is_some()) + .map(|c| c.child_value) + .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .expect("compute_backed_value called without expanded children") + } else { + self.children + .iter() + .filter(|c| c.node.is_some()) + .map(|c| c.child_value) + .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .expect("compute_backed_value called without expanded children") + } + } + /// Collect tree learning samples from all internal nodes. /// /// An internal node is one that has children and at least one expanded child. @@ -546,9 +567,11 @@ impl SearchNode { ) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); if has_expanded_child { + // Use backed value (only expanded children) for training targets + let backed_value = self.compute_backed_value(state.is_white_turn()); out.push(TreeLearningSample { board: state.clone(), - value: self.value, + value: backed_value, terminal_mix: false, }); } else if self.resolved @@ -609,9 +632,11 @@ impl SearchNode { ) { let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); if has_expanded_child { + // Use backed value (only expanded children) for training targets + let backed_value = self.compute_backed_value(state.is_white_turn()); out.push(TreeLearningSample { board: state.clone(), - value: self.value, + value: backed_value, terminal_mix: false, }); } else if self.resolved diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs index 88ac817..9637f1c 100644 --- a/training/src/replay_buffer.rs +++ b/training/src/replay_buffer.rs @@ -97,7 +97,7 @@ where "observation length must match env observation size" ); - self.push_with_observation(value, |mut out| { + self.push_with_observation(value.clamp(-1., 1.), |mut out| { for (dst, src) in out.iter_mut().zip(observation.iter()) { *dst = src.clone(); } From 616397f8bd8f374e19c46e1338185098c8d781f7 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Tue, 31 Mar 2026 03:30:06 -0400 Subject: [PATCH 55/59] port over killshots --- alpha_paint/src/bindings.rs | 57 +++--- alpha_paint/src/board/bitboard.rs | 127 +++++++++++++ alpha_paint/src/board/board_impl.rs | 274 +++++++++++++++++++++++----- alpha_paint/src/board/fen.rs | 7 +- alpha_paint/src/board/mod.rs | 1 + alpha_paint/src/board/structs.rs | 29 ++- alpha_paint/src/board/tile_map.rs | 171 ++++++++++++++++- alpha_paint/src/search.rs | 27 +-- training/src/descent.rs | 38 ++-- training/src/lib.rs | 12 +- 10 files changed, 617 insertions(+), 126 deletions(-) create mode 100644 alpha_paint/src/board/bitboard.rs diff --git a/alpha_paint/src/bindings.rs b/alpha_paint/src/bindings.rs index 5314fa1..7d2f8f8 100644 --- a/alpha_paint/src/bindings.rs +++ b/alpha_paint/src/bindings.rs @@ -1,5 +1,5 @@ -use pyo3::{PyResult, exceptions::PyRuntimeError, pyclass, pymethods}; -use rand::{RngExt, SeedableRng, rngs::StdRng, seq::SliceRandom}; +use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, PyResult}; +use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -67,7 +67,7 @@ impl PyBoard { 0: [[u16::MAX; 32]; 32], }), hills: Arc::new(vec![Vec::new(); hill_owners.len()]), - dist: Arc::new(DoubleArray32x32(vec![])), // computed after tiles are set + dist: Arc::new(DoubleArray32x32::splat(2000)), // computed after tiles are set tiles: TileMap::default(), powerups: Array32x32 { @@ -139,7 +139,11 @@ impl PyBoard { }) .collect(), ); - board.dist = Arc::new(compute_distances(&board.tiles)); + let dist = Arc::new(compute_distances(&board.tiles)); + board + .tiles + .attach_dist(dist.clone(), board.rows, board.cols); + board.dist = dist; PyBoard(board) } @@ -225,19 +229,17 @@ impl PyBoard { py_actions.push(play_instead.to_python_primitives(&player_coord)); return Ok(py_actions); } - ApplyActionOutcome::Killshot { terminal: _, moves } => { + ApplyActionOutcome::Killshot { + terminal: _, + actions: ks_actions, + } => { py_actions.push(action.to_python_primitives(&player_coord)); if !action.is_final() { // Killshot for us - add all killshot moves let mut coord = local_board.current_player_coord(); - for (i, mv) in moves.iter().enumerate() { - let a = if i == moves.len() - 1 { - Action::FinalMove(*mv) - } else { - Action::Move(*mv) - }; - py_actions.push(a.to_python_primitives(&coord)); - coord = mv.target; + for ks_action in ks_actions.iter() { + py_actions.push(ks_action.to_python_primitives(&coord)); + coord = ks_action.target(); } } return Ok(py_actions); @@ -341,28 +343,31 @@ fn enumerate_single_turns( }); move_stack.pop(); } + // if we played a final action, this killshot is for the OPPONENT, so we dont include + // it in the single turn perft. + ApplyActionOutcome::Killshot { .. } if action.is_final() => { + move_stack.push((action, player_coord)); + save_turn(Turn { + actions: move_stack.clone(), + is_terminal: None, + }); + move_stack.pop(); + } ApplyActionOutcome::Killshot { - terminal, moves, .. + terminal, + actions: ks_actions, } => { - debug_assert!(!action.is_final()); move_stack.push((action, player_coord)); - // Add all killshot intermediate moves let mut coord = board.current_player_coord(); - for (i, mv) in moves.iter().enumerate() { - let a = if i == moves.len() - 1 { - Action::FinalMove(*mv) - } else { - Action::Move(*mv) - }; - move_stack.push((a, coord)); - coord = mv.target; + for &ks_action in ks_actions.iter() { + move_stack.push((ks_action, coord)); + coord = ks_action.target(); } save_turn(Turn { actions: move_stack.clone(), is_terminal: Some(terminal), }); - // Pop all the killshot moves + the original action - for _ in 0..moves.len() + 1 { + for _ in 0..ks_actions.len() + 1 { move_stack.pop(); } } diff --git a/alpha_paint/src/board/bitboard.rs b/alpha_paint/src/board/bitboard.rs new file mode 100644 index 0000000..f8af802 --- /dev/null +++ b/alpha_paint/src/board/bitboard.rs @@ -0,0 +1,127 @@ +use crate::board::structs::{Array32x32, Coordinate}; + +/// 32×32 bitboard: row `y` holds bits for `x = 0..31`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Bitboard { + pub rows: [u32; 32], +} + +impl Default for Bitboard { + fn default() -> Self { + Self::empty() + } +} + +impl Bitboard { + pub const fn empty() -> Self { + Self { rows: [0; 32] } + } + + pub fn from_array32x32(arr: &Array32x32, conv: impl Fn(&T) -> bool) -> Self { + let mut rows = [0u32; 32]; + // TODO: we can simd this if compiler doesnt do it alr + for y in 0..32 { + for x in 0..32 { + if conv(&arr[Coordinate::new(x, y)]) { + rows[y as usize] |= 1u32 << x; + } + } + } + Self { rows } + } + + #[inline] + pub const fn get(self, c: Coordinate) -> bool { + (self.rows[c.y as usize] >> c.x) & 1 != 0 + } + + #[inline] + pub fn set_bit(&mut self, c: Coordinate) { + self.rows[c.y as usize] |= 1u32 << c.x; + } + + #[inline] + pub fn clear_bit(&mut self, c: Coordinate) { + self.rows[c.y as usize] &= !(1u32 << c.x); + } + + pub fn and(self, other: Self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = self.rows[i] & other.rows[i]; + } + Self { rows } + } + + pub fn or(self, other: Self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = self.rows[i] | other.rows[i]; + } + Self { rows } + } + + pub fn not(self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = !self.rows[i]; + } + Self { rows } + } + + pub fn any(self) -> bool { + self.rows.iter().any(|&r| r != 0) + } + + pub fn count_ones(self) -> u32 { + self.rows.iter().map(|r| r.count_ones()).sum() + } + + /// First set bit in row-major order (y then x), or `None` if empty. + pub fn first_set(self) -> Option { + for y in 0u8..32 { + let r = self.rows[y as usize]; + if r != 0 { + let x = r.trailing_zeros() as u8; + return Some(Coordinate::new(x, y)); + } + } + None + } + + pub fn iter_set(self) -> IterSet { + IterSet { + bb: self, + y: 0, + cur_row: self.rows[0], + } + } +} + +pub struct IterSet { + bb: Bitboard, + y: u8, + cur_row: u32, +} + +impl Iterator for IterSet { + type Item = Coordinate; + + fn next(&mut self) -> Option { + loop { + if self.y >= 32 { + return None; + } + if self.cur_row != 0 { + let x = self.cur_row.trailing_zeros() as u8; + let c = Coordinate::new(x, self.y); + self.cur_row &= self.cur_row - 1; + return Some(c); + } + self.y += 1; + if self.y < 32 { + self.cur_row = self.bb.rows[self.y as usize]; + } + } + } +} diff --git a/alpha_paint/src/board/board_impl.rs b/alpha_paint/src/board/board_impl.rs index bd0a716..2ed9594 100644 --- a/alpha_paint/src/board/board_impl.rs +++ b/alpha_paint/src/board/board_impl.rs @@ -2,11 +2,11 @@ use crate::board::actions::{Action, Move, MoveKind, Paint}; use crate::board::board_structs::{Player, Powerup, TerminalState}; use crate::board::consts::*; use crate::board::structs::{ - Array32x32, Coordinate, DoubleArray32x32, killshot_stamina_cost, min_path_to, + cost_for_movement, min_path_to, Array32x32, Coordinate, DoubleArray32x32, }; use crate::board::tile::Tile; use crate::board::tile_map::TileMap; -use std::cmp::{Ordering, min}; +use std::cmp::{min, Ordering}; use std::sync::Arc; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -60,12 +60,13 @@ pub enum ApplyActionOutcome { terminal: TerminalState, play_instead: Action, }, - /// The current player can win by collision via a sequence of moves. - /// The moves lead the current player to the opponent's position for a kill. - /// Only the last move is terminal; the rest are intermediate steps. + /// The player whose turn it is to move (which is evaluated AFTER a potential `Final` move + /// variant) can force a collision win using `actions`. NOTE: this does not necessarily mean + /// that the player who played the most recent action is WINNING and it is important to check + /// the `terminal` value. Killshot { terminal: TerminalState, - moves: Vec, + actions: Vec, }, } @@ -253,16 +254,10 @@ impl Board { (None, Some(_)) => unreachable!("we never have a finalizer for a non terminal"), }; - // If ongoing, check if current player can win by collision (killshot) - if matches!(outcome, ApplyActionOutcome::Ongoing) && !action.is_final() { - if let Some((ks_terminal, ks_moves)) = self.check_killshot() { - return ( - ApplyActionOutcome::Killshot { - terminal: ks_terminal, - moves: ks_moves, - }, - rollback, - ); + // NOTE: this can be the NEXT player's killshot, because we run this AFTER `end_turn` + if matches!(outcome, ApplyActionOutcome::Ongoing) { + if let Some((terminal, actions)) = self.check_killshot() { + return (ApplyActionOutcome::Killshot { terminal, actions }, rollback); } } @@ -639,54 +634,245 @@ impl Board { regen } - /// Check if the current player can win by collision (killshot). - /// Returns the terminal state and the sequence of moves to reach the opponent. - fn check_killshot(&self) -> Option<(TerminalState, Vec)> { - let is_white_turn = self.is_white_turn(); - let (player_coord, player_stamina, opponent_coord) = if is_white_turn { - (self.white_coord, self.white_stamina, self.black_coord) + + /// Returns if the player to move (e.g. based on is_white_turn) can win via collision + fn check_killshot(&self) -> Option<(TerminalState, Vec)> { + if self.is_white_turn() { + self.check_killshot_generic::() } else { - (self.black_coord, self.black_stamina, self.white_coord) + self.check_killshot_generic::() + } + } + + /// Three ways this player (per `IS_WHITE`) can execute a killshot: + /// 1. walk straight into the opponent + /// 2. walk to a beacon, teleport, then walk into the opponent + /// 3. walk, place a beacon, teleport through it, then walk into the opponent + fn check_killshot_generic(&self) -> Option<(TerminalState, Vec)> { + use crate::board::bitboard::Bitboard; + + let moves_so_far = self.consecutives_moves_so_far; + let (player_coord, stamina, opponent_coord, beacons) = if IS_WHITE { + ( + self.white_coord, + self.white_stamina, + self.black_coord, + &self.tiles.white_beacons, + ) + } else { + ( + self.black_coord, + self.black_stamina, + self.white_coord, + &self.tiles.black_beacons, + ) }; let pdist = self.dist[(player_coord, opponent_coord)] as usize; + // we r already colliding, handled by terminal() if pdist == 0 { - return None; // already colliding, handled by terminal() + return None; } - // Check if the opponent's tile is favorable for collision - // (if opponent stands on their own paint, they win the collision) - if is_white_turn { + // make sure we'd acc win the collision before doing more work + if IS_WHITE { if self.tiles[opponent_coord].is_owned_by::() { return None; } - } else { - if self.tiles[opponent_coord].is_owned_by::() { - return None; + } else if self.tiles[opponent_coord].is_owned_by::() { + return None; + } + + let win = TerminalState::win_for(IS_WHITE); + + // walk only + let cost_walk = cost_for_movement(pdist, moves_so_far); + if cost_walk <= stamina { + let path = min_path_to(&self.dist, player_coord, opponent_coord); + if path.is_empty() { + return None; // Unreachable but defensive. } + + let moves: Vec = path + .iter() + .map(|&target| Move { + target, + kind: MoveKind::Regular, + place_beacon: false, + }) + .collect(); + + return Some((win, moves_to_killshot_actions(moves))); } - let stamina_cost = killshot_stamina_cost(pdist, self.consecutives_moves_so_far); - if stamina_cost > player_stamina { + // both case 2 and 3 require already placed beacons + if beacons.is_empty() { return None; } - // We can reach the opponent with enough stamina - compute the path - let path = min_path_to(&self.dist, player_coord, opponent_coord); - if path.is_empty() { - return None; // unreachable but defensive + // both beacon_dist lookups use OUR beacons — we want "how far are + // we from our nearest beacon" and "how far is our nearest beacon from + // the opponent" + let (us_to_beacon, opp_to_beacon) = if IS_WHITE { + ( + self.tiles.white_beacon_dist[player_coord], + self.tiles.white_beacon_dist[opponent_coord], + ) + } else { + ( + self.tiles.black_beacon_dist[player_coord], + self.tiles.black_beacon_dist[opponent_coord], + ) + }; + debug_assert!(us_to_beacon != u16::MAX); + debug_assert!(opp_to_beacon != u16::MAX); + + // case 2: walk to beacon + teleport + walk + let consecutive_at_beacon = moves_so_far + us_to_beacon as usize; + let stamina_required = cost_for_movement(us_to_beacon as usize, moves_so_far) + + EXTRA_MOVE_COST * consecutive_at_beacon // teleport cost + + cost_for_movement(opp_to_beacon as usize, 1); + if stamina_required <= stamina { + let (b1, b2) = ( + self.find_beacon_by_dist::(player_coord, us_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"), + self.find_beacon_by_dist::(opponent_coord, opp_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"), + ); + // walk to b1 (may be empty if already on it) + let to_b1 = min_path_to(&self.dist, player_coord, b1); + let mut moves: Vec<_> = to_b1 + .into_iter() + .map(|coord| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: coord, + }) + .collect(); + // teleport from b1 to b2 (appended, not replacing) + moves.push(Move { + kind: MoveKind::BeaconTravel, + place_beacon: false, + target: b2, + }); + // walk from b2 to opponent + let to_opp = min_path_to(&self.dist, b2, opponent_coord); + moves.extend(to_opp.into_iter().map(|coord| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: coord, + })); + + return Some((win, moves_to_killshot_actions(moves))); } - let terminal = TerminalState::win_for(is_white_turn); - let moves: Vec = path - .iter() - .map(|&target| Move { - target, + // case 3: walk to X, place beacon, teleport to existing beacon near opp, walk to opp + // budget = stamina minus (teleport from X + walk from target_beacon to opp) + // the teleport cost depends on how far we walk to X, so we use a conservative + // lower bound: teleport costs at least EXTRA_MOVE_COST * (moves_so_far + 1) + let min_teleport_cost = EXTRA_MOVE_COST * (moves_so_far + 1); + let stamina_budget = stamina + .saturating_sub(cost_for_movement(opp_to_beacon as usize, 1)) + .saturating_sub(min_teleport_cost); + if stamina_budget < EXTRA_MOVE_COST * moves_so_far { + return None; + } + let moves_we_can_afford = max_possible_move_amt(stamina_budget, moves_so_far) as u16; + if moves_we_can_afford == 0 { + return None; + } + let reach = Bitboard::from_array32x32(self.dist.get(player_coord), |&d| { + 0 < d && d <= moves_we_can_afford + }); + + let placement = if IS_WHITE { + self.tiles.white_placement_valid + } else { + self.tiles.black_placement_valid + }; + let candidates = reach.and(placement); + if !candidates.any() { + return None; + } + let target_beacon = self + .find_beacon_by_dist::(opponent_coord, opp_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"); + let target_to_opp = min_path_to(&self.dist, target_beacon, opponent_coord); + + for x in candidates.iter_set() { + if x == player_coord || self.tiles[x].is_beacon() { + continue; + } + let dx = self.dist[(player_coord, x)] as usize; + let consecutive_at_x = moves_so_far + dx; + let total_cost = cost_for_movement(dx, moves_so_far) + + EXTRA_MOVE_COST * consecutive_at_x + + cost_for_movement(opp_to_beacon as usize, 1); + if total_cost > stamina { + continue; + } + + let path_to_x = min_path_to(&self.dist, player_coord, x); + let mut moves: Vec<_> = path_to_x + .into_iter() + .map(|t| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: t, + }) + .collect(); + // place beacon on the last step (landing on X) + moves + .last_mut() + .expect("path_to_x should not be empty") + .place_beacon = true; + // teleport from X (our new beacon) to target_beacon (existing beacon near opp) + moves.push(Move { + kind: MoveKind::BeaconTravel, + place_beacon: false, + target: target_beacon, + }); + // walk from target_beacon to opponent + moves.extend(target_to_opp.iter().map(|&t| Move { kind: MoveKind::Regular, place_beacon: false, - }) - .collect(); + target: t, + })); + return Some((win, moves_to_killshot_actions(moves))); + } - Some((terminal, moves)) + None + } + + fn find_beacon_by_dist( + &self, + coord: Coordinate, + dist: u16, + ) -> Option { + self.tiles + .get_beacon_iterator::() + .find(|&&b| self.dist[(coord, b)] == dist) + .copied() } } + +/// Walking steps as [`Action::Move`], collision as [`Action::FinalMove`]. +fn moves_to_killshot_actions(mut moves: Vec) -> Vec { + if moves.is_empty() { + return Vec::new(); + } + let last = moves.pop().expect("non-empty"); + let mut out: Vec = moves.into_iter().map(Action::Move).collect(); + out.push(Action::FinalMove(last)); + out +} + +/// Cost is 5n(2c+n-1) <= B. we use quadratic equation to get the following equation we use +/// where n=moves we can afford, B=stamina, c=moves so far +fn max_possible_move_amt(stamina: usize, moves_so_far: usize) -> usize { + let c = moves_so_far as isize; + let b = stamina as isize; + + let n = (5 - 10 * c + ((10 * c - 5).pow(2) + 20 * b).isqrt()) / 10; + n as usize +} diff --git a/alpha_paint/src/board/fen.rs b/alpha_paint/src/board/fen.rs index 1bdbd32..512e7e7 100644 --- a/alpha_paint/src/board/fen.rs +++ b/alpha_paint/src/board/fen.rs @@ -1,6 +1,6 @@ use crate::board::board_impl::Board; use crate::board::board_structs::{Player, Powerup}; -use crate::board::structs::{Array32x32, Coordinate, compute_distances}; +use crate::board::structs::{compute_distances, Array32x32, Coordinate}; use crate::board::tile::Tile; use crate::board::tile_map::TileMap; use std::fmt::{self, Display, Write}; @@ -237,6 +237,9 @@ impl Board { let hills = Arc::new(hills_vec); tiles.attach_hills(hill_id.clone(), hills.clone(), hill_owners); + let dist = Arc::new(compute_distances(&tiles)); + tiles.attach_dist(dist.clone(), rows, cols); + Ok(Board { rows, cols, @@ -244,7 +247,7 @@ impl Board { event_pointer, hill_id, hills, - dist: Arc::new(compute_distances(&tiles)), + dist, tiles, powerups, white_coord, diff --git a/alpha_paint/src/board/mod.rs b/alpha_paint/src/board/mod.rs index 0b5a107..9566d30 100644 --- a/alpha_paint/src/board/mod.rs +++ b/alpha_paint/src/board/mod.rs @@ -1,6 +1,7 @@ mod action_generation; #[allow(dead_code)] pub mod actions; +pub mod bitboard; pub mod board_impl; pub mod board_structs; pub mod consts; diff --git a/alpha_paint/src/board/structs.rs b/alpha_paint/src/board/structs.rs index 71f27a2..97c869d 100644 --- a/alpha_paint/src/board/structs.rs +++ b/alpha_paint/src/board/structs.rs @@ -1,3 +1,4 @@ +use std::array; use std::collections::VecDeque; use std::ops::{Index, IndexMut}; @@ -111,6 +112,12 @@ impl Array32x32 { return default; } } + pub fn splat(value: T) -> Self + where + T: Copy, + { + Array32x32([[value; 32]; 32]) + } } impl Index for Array32x32 { @@ -128,27 +135,37 @@ impl IndexMut for Array32x32 { } #[derive(Debug, Clone)] -pub struct DoubleArray32x32(pub Vec); +pub struct DoubleArray32x32(pub Box>>); impl DoubleArray32x32 { + #[allow(dead_code)] fn flat_index(index: (Coordinate, Coordinate)) -> usize { (((index.0.x as usize) * 32 + index.0.y as usize) * 32 + index.1.x as usize) * 32 + index.1.y as usize } + pub fn get(&self, index: Coordinate) -> &Array32x32 { + &self.0[index] + } + pub fn splat(value: T) -> Self + where + T: Copy, + { + let arr = array::from_fn(|_| array::from_fn(|_| Array32x32::splat(value))); + DoubleArray32x32(Box::new(Array32x32(arr))) + } } impl Index<(Coordinate, Coordinate)> for DoubleArray32x32 { type Output = T; fn index(&self, index: (Coordinate, Coordinate)) -> &Self::Output { - &self.0[Self::flat_index(index)] + &self.0[index.0][index.1] } } impl IndexMut<(Coordinate, Coordinate)> for DoubleArray32x32 { fn index_mut(&mut self, index: (Coordinate, Coordinate)) -> &mut Self::Output { - let idx = Self::flat_index(index); - &mut self.0[idx] + &mut self.0[index.0][index.1] } } @@ -188,7 +205,7 @@ pub fn min_path_to( } /// Compute the stamina cost for `pdist` consecutive moves starting from `consecutive_so_far`. -pub fn killshot_stamina_cost(pdist: usize, consecutive_so_far: usize) -> usize { +pub fn cost_for_movement(pdist: usize, consecutive_so_far: usize) -> usize { // Each move i (0-indexed) costs EXTRA_MOVE_COST * (consecutive_so_far + i) // Total = EXTRA_MOVE_COST * sum(consecutive_so_far + i for i in 0..pdist) // = EXTRA_MOVE_COST * (pdist * consecutive_so_far + pdist*(pdist-1)/2) @@ -201,7 +218,7 @@ pub fn compute_distances(tiles: &T) -> DoubleArray32x32 where T: Index, { - let mut dist = DoubleArray32x32(vec![2000; 32 * 32 * 32 * 32]); + let mut dist = DoubleArray32x32::splat(2000); for i in 0..32u8 { for j in 0..32u8 { diff --git a/alpha_paint/src/board/tile_map.rs b/alpha_paint/src/board/tile_map.rs index 8e8378a..e2c482d 100644 --- a/alpha_paint/src/board/tile_map.rs +++ b/alpha_paint/src/board/tile_map.rs @@ -2,9 +2,10 @@ use std::ops::Index; use std::slice::Iter; use std::sync::Arc; +use crate::board::bitboard::Bitboard; use crate::board::board_structs::{HillData, Player}; use crate::board::consts::HILL_CONTROL_THRESHOLD; -use crate::board::structs::{Array32x32, Coordinate}; +use crate::board::structs::{Array32x32, Coordinate, DoubleArray32x32}; use crate::board::tile::Tile; /// The idea of the TileMap is to provide a useful abstraction layer between the game @@ -25,8 +26,18 @@ pub struct TileMap { hill_id: Arc>, hills: Arc>>, hill_metadata: Vec, - white_beacons: Vec, - black_beacons: Vec, + pub white_beacons: Vec, + pub black_beacons: Vec, + + /// All-pairs walking distances (walls only). Set via [`Self::attach_dist`]. + pub dist: Option>>, + pub rows: u8, + pub cols: u8, + /// Min walking distance from any friendly beacon to each tile (`u16::MAX` if no beacons). + pub white_beacon_dist: Array32x32, + pub black_beacon_dist: Array32x32, + pub white_placement_valid: Bitboard, + pub black_placement_valid: Bitboard, } impl Default for TileMap { @@ -46,6 +57,17 @@ impl Default for TileMap { hill_metadata: Vec::new(), white_beacons: Vec::new(), black_beacons: Vec::new(), + dist: None, + rows: 32, + cols: 32, + white_beacon_dist: Array32x32 { + 0: [[u16::MAX; 32]; 32], + }, + black_beacon_dist: Array32x32 { + 0: [[u16::MAX; 32]; 32], + }, + white_placement_valid: Bitboard::empty(), + black_placement_valid: Bitboard::empty(), } } } @@ -108,6 +130,145 @@ impl TileMap { } } + pub fn on_edge(&self, coord: Coordinate) -> bool { + coord.x == 0 || coord.x + 1 == self.cols || coord.y == 0 || coord.y + 1 == self.rows + } + + /// Must be called after the tile map is fully built; wires walking distances and + /// initializes beacon distance and beacon placement bitboards. + pub fn attach_dist(&mut self, dist: Arc>, rows: u8, cols: u8) { + self.dist = Some(dist); + self.rows = rows; + self.cols = cols; + let d = self.dist.as_ref().unwrap().clone(); + self.recompute_beacon_dist::(d.as_ref()); + self.recompute_beacon_dist::(d.as_ref()); + self.recompute_all_placement_valid(); + } + + fn recompute_beacon_dist(&mut self, dist: &DoubleArray32x32) { + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + let mut m = u16::MAX; + for &b in self.get_beacon_iterator::() { + let dbt = dist[(b, c)]; + if dbt < m { + m = dbt; + } + } + if IS_WHITE { + self.white_beacon_dist[c] = m; + } else { + self.black_beacon_dist[c] = m; + } + } + } + } + + fn apply_beacon_add( + &mut self, + dist: &DoubleArray32x32, + b: Coordinate, + ) { + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + let d = dist[(b, c)]; + if IS_WHITE { + if d < self.white_beacon_dist[c] { + self.white_beacon_dist[c] = d; + } + } else if d < self.black_beacon_dist[c] { + self.black_beacon_dist[c] = d; + } + } + } + } + + fn update_beacon_dist_after_tile_change(&mut self, old: Tile, new: Tile, coord: Coordinate) { + let Some(dist) = self.dist.as_ref().cloned() else { + return; + }; + let d = dist.as_ref(); + if old.is_beacon_of::() && !new.is_beacon_of::() { + self.recompute_beacon_dist::(d); + } else if !old.is_beacon_of::() && new.is_beacon_of::() { + self.apply_beacon_add::(d, coord); + } + if old.is_beacon_of::() && !new.is_beacon_of::() { + self.recompute_beacon_dist::(d); + } else if !old.is_beacon_of::() && new.is_beacon_of::() { + self.apply_beacon_add::(d, coord); + } + } + + fn recompute_all_placement_valid(&mut self) { + self.white_placement_valid = Bitboard::empty(); + self.black_placement_valid = Bitboard::empty(); + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + if self.placement_valid_tile::(c) { + self.white_placement_valid.set_bit(c); + } + if self.placement_valid_tile::(c) { + self.black_placement_valid.set_bit(c); + } + } + } + } + + fn placement_valid_tile(&self, target: Coordinate) -> bool { + if !target.in_bounds() { + return false; + } + let tile = self[target]; + if tile.is_wall() { + return false; + } + if self.on_edge(target) { + return false; + } + if tile.is_beacon_of::() { + return false; + } + let (valid_cells, controlled_cells) = self.beacon_window_counts::(target); + if controlled_cells * 3 < valid_cells * 2 { + return false; + } + let mut post = tile; + if tile.is_erasable::() { + post.erase::(); + } else { + post.maybe_erase1::(); + } + let opponent = if IS_WHITE { + Player::Black + } else { + Player::White + }; + post.paint_owner() != Some(opponent) + } + + fn update_placement_neighborhood(&mut self, coord: Coordinate) { + for center in coord.region_3x3_with_self() { + if !center.in_bounds() { + continue; + } + if self.placement_valid_tile::(center) { + self.white_placement_valid.set_bit(center); + } else { + self.white_placement_valid.clear_bit(center); + } + if self.placement_valid_tile::(center) { + self.black_placement_valid.set_bit(center); + } else { + self.black_placement_valid.clear_bit(center); + } + } + } + /// Set the value of a specified tile, and update hill ownership accordingly /// Returns the original tile at that location pub fn set(&mut self, coord: Coordinate, tile: Tile) -> Tile { @@ -161,6 +322,10 @@ impl TileMap { self.tiles[coord] = tile; self.apply_owner_transition(coord, before_owner, after_owner); + if self.dist.is_some() { + self.update_beacon_dist_after_tile_change(old, tile, coord); + self.update_placement_neighborhood(coord); + } old } diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index 42f89b7..ec06da5 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -4,7 +4,6 @@ use std::time::Duration; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; -use crate::board::actions::Move; use crate::board::{Action, ApplyActionOutcome, Board, TerminalState}; use crate::evaluation::Evaluator; @@ -141,17 +140,17 @@ impl SearchNode { } } - fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { + fn build_killshot_chain( + board: &Board, + terminal: TerminalState, + actions: &[Action], + ) -> SearchNode { let term_value = Self::value_from_term(board, terminal); let completion_value = terminal.value(); let mut node = SearchNode::new(term_value, completion_value, true); - for index in (0..moves.len()).rev() { - let action = if index == moves.len() - 1 { - Action::FinalMove(moves[index]) - } else { - Action::Move(moves[index]) - }; + for index in (0..actions.len()).rev() { + let action = actions[index]; node = SearchNode { value: term_value, completion_value, @@ -211,8 +210,12 @@ impl SearchNode { ))), }); } - ApplyActionOutcome::Killshot { terminal, moves } => { - let chain = Self::build_killshot_chain(&local_board, terminal, &moves); + ApplyActionOutcome::Killshot { + terminal, + actions: ks_actions, + } => { + let chain = + Self::build_killshot_chain(&local_board, terminal, &ks_actions); new_node.children.push(ChildData { action, child_value: chain.value, @@ -236,8 +239,8 @@ impl SearchNode { terminal.value(), true, ), - ApplyActionOutcome::Killshot { terminal, moves } => { - Self::build_killshot_chain(board, terminal, &moves) + ApplyActionOutcome::Killshot { terminal, actions } => { + Self::build_killshot_chain(board, terminal, &actions) } } } diff --git a/training/src/descent.rs b/training/src/descent.rs index 327d1b0..21fe59b 100644 --- a/training/src/descent.rs +++ b/training/src/descent.rs @@ -4,7 +4,6 @@ //! Implements tree learning (collect training samples from internal nodes) //! and ordinal distribution for action selection. -use alpha_paint::board::actions::Move; use alpha_paint::board::{Action, ApplyActionOutcome, Board, Rollback, TerminalState}; use rand::rngs::SmallRng; #[cfg(test)] @@ -193,19 +192,15 @@ impl SearchNode { } } - /// Build a chain of resolved SearchNodes for a killshot move sequence. - fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { + /// Build a chain of resolved SearchNodes for a killshot action sequence. + fn build_killshot_chain(board: &Board, terminal: TerminalState, actions: &[Action]) -> SearchNode { let term_value = Self::value_from_term(board, terminal); let comp_value = terminal.value(); let mut node = SearchNode::new(term_value, comp_value, true); - for i in (0..moves.len()).rev() { - let action = if i == moves.len() - 1 { - Action::FinalMove(moves[i]) - } else { - Action::Move(moves[i]) - }; + for i in (0..actions.len()).rev() { + let action = actions[i]; let parent = SearchNode { value: term_value, @@ -237,14 +232,9 @@ impl SearchNode { board } - fn apply_killshot_moves(mut board: Board, moves: &[Move]) -> Board { - let last_idx = moves.len().saturating_sub(1); - for (idx, mv) in moves.iter().copied().enumerate() { - let action = if idx == last_idx { - Action::FinalMove(mv) - } else { - Action::Move(mv) - }; + fn apply_killshot_actions(mut board: Board, actions: &[Action]) -> Board { + let last_idx = actions.len().saturating_sub(1); + for (idx, action) in actions.iter().copied().enumerate() { let (outcome, _) = board.apply_action(action); if idx == last_idx { debug_assert!(matches!(outcome, ApplyActionOutcome::Terminal { .. })); @@ -266,8 +256,8 @@ impl SearchNode { ApplyActionOutcome::PlayInstead { play_instead, .. } => { Self::apply_play_instead(state.clone(), action, rollback, play_instead) } - ApplyActionOutcome::Killshot { moves, .. } => { - Self::apply_killshot_moves(state.clone(), &moves) + ApplyActionOutcome::Killshot { actions, .. } => { + Self::apply_killshot_actions(state.clone(), &actions) } ApplyActionOutcome::Ongoing => { unreachable!("terminal leaf sampling only applies to resolved terminal outcomes") @@ -369,11 +359,11 @@ impl SearchNode { ))), }); } - ApplyActionOutcome::Killshot { terminal, moves } => { - let terminal_board = Self::apply_killshot_moves(local_board, &moves); + ApplyActionOutcome::Killshot { terminal, actions } => { + let terminal_board = Self::apply_killshot_actions(local_board, &actions); let term_value = Self::value_from_term(&terminal_board, terminal); let chain = - Self::build_killshot_chain(&terminal_board, terminal, &moves); + Self::build_killshot_chain(&terminal_board, terminal, &actions); new_node.children.push(ChildData { action, child_value: term_value, @@ -412,8 +402,8 @@ impl SearchNode { terminal.value(), true, ), - ApplyActionOutcome::Killshot { terminal, moves } => { - Self::build_killshot_chain(board, terminal, &moves) + ApplyActionOutcome::Killshot { terminal, actions } => { + Self::build_killshot_chain(board, terminal, &actions) } } } diff --git a/training/src/lib.rs b/training/src/lib.rs index 1d22009..0e720a2 100644 --- a/training/src/lib.rs +++ b/training/src/lib.rs @@ -156,16 +156,10 @@ fn random_terminal_board(rng: &mut R) -> (Board, TerminalState) { debug_assert!(matches!(play_outcome, ApplyActionOutcome::Terminal { .. })); return (board, terminal); } - ApplyActionOutcome::Killshot { terminal, moves } => { - let last_idx = moves.len().saturating_sub(1); - for (idx, mv) in moves.into_iter().enumerate() { - let action = if idx == last_idx { - Action::FinalMove(mv) - } else { - Action::Move(mv) - }; + ApplyActionOutcome::Killshot { terminal, actions } => { + for (idx, action) in actions.iter().copied().enumerate() { let (killshot_outcome, _) = board.apply_action(action); - if idx == last_idx { + if idx == actions.len() - 1 { debug_assert!(matches!( killshot_outcome, ApplyActionOutcome::Terminal { .. } From 72a959b0f4e7ca2111ea9b851b549a140c2ddd4e Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Tue, 31 Mar 2026 04:07:40 -0400 Subject: [PATCH 56/59] include player location in nnue --- python/alphapaint_training/model.py | 6 +- python/alphapaint_training/packed_obs.py | 72 ++++++++++++++++++------ 2 files changed, 57 insertions(+), 21 deletions(-) diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 175d7b6..cf9f015 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -1,7 +1,7 @@ """NNUE value model for GPU training. Architecture mirrors the CPU NNUE evaluator: - - Accumulator: shared linear 15360 -> acc_dim for both perspectives + - Accumulator: shared linear 17408 -> acc_dim for both perspectives - Local MLP: shared linear 375 -> local_hidden_dim (ReLU) for both players - Global features: 20 normalized scalars - Head: concat all -> fc1 (ReLU) -> fc2 (ReLU) -> fc3 -> scalar @@ -23,7 +23,7 @@ GLOBAL_SCALES, LOCAL_FEATURES, OPP_PLANE_PERM, - TILE_PLANES, + TILE_BITMASK_PLANES, TOTAL_TILE_FEATURES, decode_nnue_obs, ) @@ -125,7 +125,7 @@ def __init__( # captured in CUDA graphs. self.register_buffer( "_plane_shift", - torch.arange(TILE_PLANES, dtype=torch.int32), + torch.arange(TILE_BITMASK_PLANES, dtype=torch.int32), persistent=False, ) self.register_buffer( diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py index aa2634c..6089baf 100644 --- a/python/alphapaint_training/packed_obs.py +++ b/python/alphapaint_training/packed_obs.py @@ -16,6 +16,10 @@ bit 12: hill neutral bit 13: hill current player bit 14: hill opponent + +Additional planes added during decoding: + plane 15: one-hot my player location (derived from positions) + plane 16: one-hot opponent player location (derived from positions) """ from __future__ import annotations @@ -26,11 +30,14 @@ BOARD_SIDE = 32 BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 -TILE_PLANES = 15 +# Bits 0–14 from each u16 tile word (Rust encoder); +2 derived position planes → accumulator +TILE_BITMASK_PLANES = 15 +TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 LOCAL_WINDOW_TILES = 25 GLOBAL_FEATURES = 20 -TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 15360 -LOCAL_FEATURES = LOCAL_WINDOW_TILES * TILE_PLANES # 375 +TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 +LOCAL_TILE_PLANES = TILE_BITMASK_PLANES # Local windows don't have player position planes +LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 OFFSET_LOCALS = BOARD_CELLS # 1024 OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 @@ -44,8 +51,9 @@ # swap my_paint[0:4] <-> opp_paint[4:8] # swap my_beacon[10] <-> opp_beacon[11] # swap my_hill[13] <-> opp_hill[14] +# swap my_pos[15] <-> opp_pos[16] # wall[8], powerup[9], hill_neutral[12] stay -OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13] +OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] # Normalization divisors for the 20 global features. # Brings raw u16 values into roughly [0, 1] range. @@ -85,14 +93,14 @@ def _check_packed_obs(packed_obs: torch.Tensor) -> None: def _unpack_bits(words: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: - """Extract TILE_PLANES binary planes from u16 words. + """Extract binary planes from u16 words (one channel per shift index). Args: words: [...] int32 tensor of packed bitmasks - shift: [TILE_PLANES] int32 tensor = arange(TILE_PLANES) + shift: [K] int32 tensor of bit indices (e.g. arange(15) for tile bitmasks) Returns: - [..., TILE_PLANES] float tensor of 0/1 values + [..., K] float tensor of 0/1 values """ return ((words.unsqueeze(-1) >> shift) & 1).float() @@ -119,26 +127,53 @@ def decode_nnue_obs( CUDA graphs. Returns: - my_features: [B, 15360] dtype — accumulator input (my perspective) - opp_features: [B, 15360] dtype — accumulator input (opponent perspective) - local_my: [B, 375] dtype — local window around my position - local_opp: [B, 375] dtype — local window around opponent position - globals_norm: [B, 20] dtype — normalized global scalars - turn_count: [B] int32 — for white-to-move sign + my_features: [B, 17408] dtype — accumulator input (my perspective) + opp_features: [B, 17408] dtype — accumulator input (opponent perspective) + local_my: [B, 375] dtype — local window around my position + local_opp: [B, 375] dtype — local window around opponent position + globals_norm: [B, 20] dtype — normalized global scalars + turn_count: [B] int32 — for white-to-move sign """ B = packed_obs.shape[0] - # --- Tile features: [B, 1024] u16 -> [B, 15360] ---------------------- + # --- Tile features: [B, 1024] u16 -> [B, 17408] ---------------------- tiles = packed_obs[:, :BOARD_CELLS].to(torch.int32) - bits = _unpack_bits(tiles, plane_shift) # [B, 1024, 15] - my_features = bits.reshape(B, -1).to(dtype) + # Only bits 0–14 are defined on tile words; planes 15–16 come from positions below. + bits = _unpack_bits(tiles, plane_shift[:TILE_BITMASK_PLANES]) # [B, 1024, 15] + + # --- Player positions -> one-hot planes [B, 1024, 2] ----------------- + my_x = packed_obs[:, OFFSET_POSITIONS].long() # [B] + my_y = packed_obs[:, OFFSET_POSITIONS + 1].long() # [B] + opp_x = packed_obs[:, OFFSET_POSITIONS + 2].long() # [B] + opp_y = packed_obs[:, OFFSET_POSITIONS + 3].long() # [B] + + my_pos_flat = my_x * BOARD_SIDE + my_y # [B] + opp_pos_flat = opp_x * BOARD_SIDE + opp_y # [B] + + my_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + opp_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + my_pos_plane.scatter_(1, my_pos_flat.unsqueeze(1), 1.0) + opp_pos_plane.scatter_(1, opp_pos_flat.unsqueeze(1), 1.0) + + # Concatenate: [B, 1024, 15] + [B, 1024, 1] + [B, 1024, 1] = [B, 1024, 17] + bits_with_pos = torch.cat( + [bits, my_pos_plane.unsqueeze(-1), opp_pos_plane.unsqueeze(-1)], dim=-1 + ) # [B, 1024, 17] + + my_features = bits_with_pos.reshape(B, -1).to(dtype) # Opponent perspective: permute planes - opp_features = bits[:, :, opp_perm].reshape(B, -1).to(dtype) + opp_features = bits_with_pos[:, :, opp_perm].reshape(B, -1).to(dtype) # --- Local windows: [B, 50] u16 -> [B, 375] each --------------------- + # Local windows only have 15 planes (no player positions) locals_raw = packed_obs[:, OFFSET_LOCALS:OFFSET_POSITIONS].to(torch.int32) - local_bits = _unpack_bits(locals_raw, plane_shift) # [B, 50, 15] + local_shift = plane_shift[:TILE_BITMASK_PLANES] + local_bits = _unpack_bits(locals_raw, local_shift) # [B, 50, 15] local_my = local_bits[:, :LOCAL_WINDOW_TILES, :].reshape(B, -1).to(dtype) local_opp = local_bits[:, LOCAL_WINDOW_TILES:, :].reshape(B, -1).to(dtype) @@ -162,6 +197,7 @@ def decode_nnue_obs( "LOCAL_WINDOW_TILES", "OBS_WORDS", "OPP_PLANE_PERM", + "TILE_BITMASK_PLANES", "TILE_PLANES", "TOTAL_TILE_FEATURES", "decode_nnue_obs", From 7d5fd63938e1993015c275119d8259bf9766f3b8 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Tue, 31 Mar 2026 04:20:48 -0400 Subject: [PATCH 57/59] switch packed format decode to custom triton kernel --- python/alphapaint_training/__init__.py | 2 +- python/alphapaint_training/decode_triton.py | 257 ++++++++++++++++++++ python/alphapaint_training/model.py | 42 +--- python/alphapaint_training/packed_obs.py | 204 ---------------- python/scripts/benchmark_decode_forward.py | 82 +++++++ python/scripts/export_value_net.py | 5 +- 6 files changed, 351 insertions(+), 241 deletions(-) create mode 100644 python/alphapaint_training/decode_triton.py delete mode 100644 python/alphapaint_training/packed_obs.py create mode 100644 python/scripts/benchmark_decode_forward.py diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index d9e640d..9f544d2 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -11,7 +11,7 @@ ) from .logger import TrainingLogger from .model import NnueValueNet, PackedValueModel -from .packed_obs import ( +from .decode_triton import ( BOARD_CELLS, BOARD_SIDE, GLOBAL_FEATURES, diff --git a/python/alphapaint_training/decode_triton.py b/python/alphapaint_training/decode_triton.py new file mode 100644 index 0000000..200b20a --- /dev/null +++ b/python/alphapaint_training/decode_triton.py @@ -0,0 +1,257 @@ +"""Packed NNUE observation layout + decode (CUDA + Triton only). + +The Rust encoder packs board state into a flat u16 array with this layout: + [0, 1024) — per-tile 15-bit bitmask (32×32 board) + [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) + [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) + [1078, 1098) — 20 global scalar features + +Decode fuses tile bitmask unpack, position one-hots, opponent plane permutation (accum), +and local-window unpack. Globals + turn_count use PyTorch on the result. +""" + +from __future__ import annotations + +import torch + +BOARD_SIDE = 32 +BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 +TILE_BITMASK_PLANES = 15 +TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 +LOCAL_WINDOW_TILES = 25 +GLOBAL_FEATURES = 20 +TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 +LOCAL_TILE_PLANES = TILE_BITMASK_PLANES +LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 + +OFFSET_LOCALS = BOARD_CELLS # 1024 +OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 +OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 +OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 + +GLOBAL_TURN_COUNT_IDX = 17 + +OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] + +GLOBAL_SCALES: tuple[float, ...] = ( + 420.0, + 420.0, + 420.0, + 420.0, + 1024.0, + 1024.0, + 8.0, + 8.0, + 1024.0, + 1024.0, + 8.0, + 1024.0, + 1024.0, + 64.0, + 64.0, + 64.0, + 32.0, + 2000.0, + 32.0, + 32.0, +) + + +def _check_packed_obs(packed_obs: torch.Tensor) -> None: + if packed_obs.dtype != torch.uint16: + raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") + if packed_obs.ndim != 2: + raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") + if packed_obs.shape[1] != OBS_WORDS: + raise ValueError( + f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" + ) + + +try: + import triton + import triton.language as tl + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + triton = None # type: ignore[assignment] + tl = None # type: ignore[assignment] + + +if _TRITON_AVAILABLE: + _OPP_PERM17 = tl.constexpr(tuple(OPP_PLANE_PERM)) + + @triton.jit + def _decode_accum_kernel( + packed_ptr, + my_out_ptr, + opp_out_ptr, + row_stride: tl.constexpr, + off_positions: tl.constexpr, + board_side: tl.constexpr, + board_cells: tl.constexpr, + tile_planes: tl.constexpr, + ) -> None: + pid = tl.program_id(0) + b = pid // board_cells + cell = pid - b * board_cells + + base_tile = b * row_stride + cell + w = tl.load(packed_ptr + base_tile).to(tl.int32) + + pos0 = b * row_stride + off_positions + my_x = tl.load(packed_ptr + pos0 + 0).to(tl.int32) + my_y = tl.load(packed_ptr + pos0 + 1).to(tl.int32) + opp_x = tl.load(packed_ptr + pos0 + 2).to(tl.int32) + opp_y = tl.load(packed_ptr + pos0 + 3).to(tl.int32) + my_flat = my_x * board_side + my_y + opp_flat = opp_x * board_side + opp_y + cell_i = cell.to(tl.int32) + my_pos = tl.where(cell_i == my_flat, 1.0, 0.0) + opp_pos = tl.where(cell_i == opp_flat, 1.0, 0.0) + + out_base = b * (board_cells * tile_planes) + cell * tile_planes + + for p in tl.static_range(15): + v = ((w >> p) & 1).to(tl.float32) + tl.store(my_out_ptr + out_base + p, v.to(tl.bfloat16)) + tl.store(my_out_ptr + out_base + 15, my_pos.to(tl.bfloat16)) + tl.store(my_out_ptr + out_base + 16, opp_pos.to(tl.bfloat16)) + + for out_p in tl.static_range(17): + src = _OPP_PERM17[out_p] + if src < 15: + v = ((w >> src) & 1).to(tl.float32) + elif src == 15: + v = my_pos + else: + v = opp_pos + tl.store(opp_out_ptr + out_base + out_p, v.to(tl.bfloat16)) + + @triton.jit + def _decode_locals_kernel( + packed_ptr, + local_my_ptr, + local_opp_ptr, + row_stride: tl.constexpr, + off_locals: tl.constexpr, + local_tiles: tl.constexpr, + bitmask_planes: tl.constexpr, + ) -> None: + pid = tl.program_id(0) + nloc = 2 * local_tiles + b = pid // nloc + li = pid - b * nloc + + idx = b * row_stride + off_locals + li + w = tl.load(packed_ptr + idx).to(tl.int32) + + for p in tl.static_range(15): + v = ((w >> p) & 1).to(tl.bfloat16) + dst_my = b * (local_tiles * bitmask_planes) + li * bitmask_planes + p + li2 = li - local_tiles + dst_opp = b * (local_tiles * bitmask_planes) + li2 * bitmask_planes + p + tl.store(local_my_ptr + dst_my, v, mask=li < local_tiles) + tl.store(local_opp_ptr + dst_opp, v, mask=li >= local_tiles) + + def decode_nnue_obs( + packed_obs: torch.Tensor, + *, + global_scales: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + """Decode packed observation to bf16 features (CUDA + Triton only).""" + if not torch.cuda.is_available() or not packed_obs.is_cuda: + raise RuntimeError("decode_nnue_obs requires CUDA device tensors") + _check_packed_obs(packed_obs) + + packed_obs = packed_obs.contiguous() + device = packed_obs.device + B = packed_obs.shape[0] + row_stride = packed_obs.stride(0) + if packed_obs.stride(1) != 1: + raise ValueError("packed_obs must be contiguous in last dimension") + + dtype = torch.bfloat16 + my_features = torch.empty(B, TOTAL_TILE_FEATURES, dtype=dtype, device=device) + opp_features = torch.empty(B, TOTAL_TILE_FEATURES, dtype=dtype, device=device) + local_my = torch.empty( + B, LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES, dtype=dtype, device=device + ) + local_opp = torch.empty( + B, LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES, dtype=dtype, device=device + ) + + grid_acc = (B * BOARD_CELLS,) + _decode_accum_kernel[grid_acc]( + packed_obs, + my_features, + opp_features, + row_stride=row_stride, + off_positions=OFFSET_POSITIONS, + board_side=BOARD_SIDE, + board_cells=BOARD_CELLS, + tile_planes=TILE_PLANES, + ) + + grid_loc = (B * 2 * LOCAL_WINDOW_TILES,) + _decode_locals_kernel[grid_loc]( + packed_obs, + local_my, + local_opp, + row_stride=row_stride, + off_locals=OFFSET_LOCALS, + local_tiles=LOCAL_WINDOW_TILES, + bitmask_planes=TILE_BITMASK_PLANES, + ) + + globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) + globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) + turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to( + torch.int32 + ) + + return my_features, opp_features, local_my, local_opp, globals_norm, turn_count + +else: + + def decode_nnue_obs( + packed_obs: torch.Tensor, + *, + global_scales: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + raise RuntimeError("decode_nnue_obs requires Triton (linux dependency)") + + +__all__ = [ + "BOARD_CELLS", + "BOARD_SIDE", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "GLOBAL_TURN_COUNT_IDX", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", + "OBS_WORDS", + "OFFSET_GLOBALS", + "OFFSET_LOCALS", + "OFFSET_POSITIONS", + "OPP_PLANE_PERM", + "TILE_BITMASK_PLANES", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", + "decode_nnue_obs", +] diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index cf9f015..5b0c67a 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -13,17 +13,13 @@ from __future__ import annotations -from typing import cast - import torch from torch import nn -from .packed_obs import ( +from .decode_triton import ( GLOBAL_FEATURES, GLOBAL_SCALES, LOCAL_FEATURES, - OPP_PLANE_PERM, - TILE_BITMASK_PLANES, TOTAL_TILE_FEATURES, decode_nnue_obs, ) @@ -110,10 +106,8 @@ def __init__( local_hidden_dim: int = 64, fc1_dim: int = 128, fc2_dim: int = 32, - board_dtype: torch.dtype = torch.bfloat16, ): super().__init__() - self.board_dtype = board_dtype self.value_net = NnueValueNet( acc_dim=acc_dim, local_hidden_dim=local_hidden_dim, @@ -121,18 +115,6 @@ def __init__( fc2_dim=fc2_dim, ) - # Decode buffers — registered so they follow .to(device) and are - # captured in CUDA graphs. - self.register_buffer( - "_plane_shift", - torch.arange(TILE_BITMASK_PLANES, dtype=torch.int32), - persistent=False, - ) - self.register_buffer( - "_opp_perm", - torch.tensor(OPP_PLANE_PERM, dtype=torch.long), - persistent=False, - ) self.register_buffer( "_global_scales", torch.tensor(GLOBAL_SCALES, dtype=torch.float32), @@ -150,25 +132,21 @@ def decode( torch.Tensor, ]: return decode_nnue_obs( - packed_obs, - plane_shift=cast(torch.Tensor, self._plane_shift), - opp_perm=cast(torch.Tensor, self._opp_perm), - global_scales=cast(torch.Tensor, self._global_scales), - dtype=self.board_dtype, + packed_obs, global_scales=self._global_scales ) def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: my_features, opp_features, local_my, local_opp, globals_norm, turn_count = ( self.decode(packed_obs) ) - - if not torch.is_autocast_enabled(): - param_dtype = next(self.value_net.parameters()).dtype - my_features = my_features.to(dtype=param_dtype) - opp_features = opp_features.to(dtype=param_dtype) - local_my = local_my.to(dtype=param_dtype) - local_opp = local_opp.to(dtype=param_dtype) - globals_norm = globals_norm.to(dtype=param_dtype) + # Decode is always bf16; align to parameter dtype when not under autocast. + wdt = next(self.value_net.parameters()).dtype + if my_features.dtype != wdt: + my_features = my_features.to(wdt) + opp_features = opp_features.to(wdt) + local_my = local_my.to(wdt) + local_opp = local_opp.to(wdt) + globals_norm = globals_norm.to(wdt) value = self.value_net( my_features, opp_features, local_my, local_opp, globals_norm diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py deleted file mode 100644 index 6089baf..0000000 --- a/python/alphapaint_training/packed_obs.py +++ /dev/null @@ -1,204 +0,0 @@ -"""Packed observation decoding for NNUE training. - -The Rust encoder packs board state into a flat u16 array with this layout: - [0, 1024) — per-tile 15-bit bitmask (32×32 board) - [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) - [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) - [1078, 1098) — 20 global scalar features - -Tile bitmask planes (from current player's perspective): - bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) - bits 4-7: opponent paint thermometer - bit 8: wall - bit 9: powerup - bit 10: current player beacon - bit 11: opponent beacon - bit 12: hill neutral - bit 13: hill current player - bit 14: hill opponent - -Additional planes added during decoding: - plane 15: one-hot my player location (derived from positions) - plane 16: one-hot opponent player location (derived from positions) -""" - -from __future__ import annotations - -import torch - -# Layout constants ---------------------------------------------------------- - -BOARD_SIDE = 32 -BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 -# Bits 0–14 from each u16 tile word (Rust encoder); +2 derived position planes → accumulator -TILE_BITMASK_PLANES = 15 -TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 -LOCAL_WINDOW_TILES = 25 -GLOBAL_FEATURES = 20 -TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 -LOCAL_TILE_PLANES = TILE_BITMASK_PLANES # Local windows don't have player position planes -LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 - -OFFSET_LOCALS = BOARD_CELLS # 1024 -OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 -OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 -OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 - -# Index of turn_count within the global section -GLOBAL_TURN_COUNT_IDX = 17 - -# Plane permutation for opponent perspective: -# swap my_paint[0:4] <-> opp_paint[4:8] -# swap my_beacon[10] <-> opp_beacon[11] -# swap my_hill[13] <-> opp_hill[14] -# swap my_pos[15] <-> opp_pos[16] -# wall[8], powerup[9], hill_neutral[12] stay -OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] - -# Normalization divisors for the 20 global features. -# Brings raw u16 values into roughly [0, 1] range. -GLOBAL_SCALES: tuple[float, ...] = ( - 420.0, # 0 my_stamina - 420.0, # 1 opp_stamina - 420.0, # 2 my_max_stamina - 420.0, # 3 opp_max_stamina - 1024.0, # 4 my_territory - 1024.0, # 5 opp_territory - 8.0, # 6 my_hills - 8.0, # 7 opp_hills - 1024.0, # 8 my_hill_tiles - 1024.0, # 9 opp_hill_tiles - 8.0, # 10 contested_hills - 1024.0, # 11 my_beacons - 1024.0, # 12 opp_beacons - 64.0, # 13 player_dist - 64.0, # 14 my_hill_dist - 64.0, # 15 opp_hill_dist - 32.0, # 16 consecutive_moves - 2000.0, # 17 turn_count - 32.0, # 18 rows - 32.0, # 19 cols -) - - -def _check_packed_obs(packed_obs: torch.Tensor) -> None: - if packed_obs.dtype != torch.uint16: - raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") - if packed_obs.ndim != 2: - raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") - if packed_obs.shape[1] != OBS_WORDS: - raise ValueError( - f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" - ) - - -def _unpack_bits(words: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: - """Extract binary planes from u16 words (one channel per shift index). - - Args: - words: [...] int32 tensor of packed bitmasks - shift: [K] int32 tensor of bit indices (e.g. arange(15) for tile bitmasks) - - Returns: - [..., K] float tensor of 0/1 values - """ - return ((words.unsqueeze(-1) >> shift) & 1).float() - - -def decode_nnue_obs( - packed_obs: torch.Tensor, - *, - plane_shift: torch.Tensor, - opp_perm: torch.Tensor, - global_scales: torch.Tensor, - dtype: torch.dtype = torch.bfloat16, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Decode packed NNUE observation into model inputs. - - All buffers (plane_shift, opp_perm, global_scales) should be registered - as model buffers so they live on the right device and are captured by - CUDA graphs. - - Returns: - my_features: [B, 17408] dtype — accumulator input (my perspective) - opp_features: [B, 17408] dtype — accumulator input (opponent perspective) - local_my: [B, 375] dtype — local window around my position - local_opp: [B, 375] dtype — local window around opponent position - globals_norm: [B, 20] dtype — normalized global scalars - turn_count: [B] int32 — for white-to-move sign - """ - B = packed_obs.shape[0] - - # --- Tile features: [B, 1024] u16 -> [B, 17408] ---------------------- - tiles = packed_obs[:, :BOARD_CELLS].to(torch.int32) - # Only bits 0–14 are defined on tile words; planes 15–16 come from positions below. - bits = _unpack_bits(tiles, plane_shift[:TILE_BITMASK_PLANES]) # [B, 1024, 15] - - # --- Player positions -> one-hot planes [B, 1024, 2] ----------------- - my_x = packed_obs[:, OFFSET_POSITIONS].long() # [B] - my_y = packed_obs[:, OFFSET_POSITIONS + 1].long() # [B] - opp_x = packed_obs[:, OFFSET_POSITIONS + 2].long() # [B] - opp_y = packed_obs[:, OFFSET_POSITIONS + 3].long() # [B] - - my_pos_flat = my_x * BOARD_SIDE + my_y # [B] - opp_pos_flat = opp_x * BOARD_SIDE + opp_y # [B] - - my_pos_plane = torch.zeros( - B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device - ) - opp_pos_plane = torch.zeros( - B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device - ) - my_pos_plane.scatter_(1, my_pos_flat.unsqueeze(1), 1.0) - opp_pos_plane.scatter_(1, opp_pos_flat.unsqueeze(1), 1.0) - - # Concatenate: [B, 1024, 15] + [B, 1024, 1] + [B, 1024, 1] = [B, 1024, 17] - bits_with_pos = torch.cat( - [bits, my_pos_plane.unsqueeze(-1), opp_pos_plane.unsqueeze(-1)], dim=-1 - ) # [B, 1024, 17] - - my_features = bits_with_pos.reshape(B, -1).to(dtype) - - # Opponent perspective: permute planes - opp_features = bits_with_pos[:, :, opp_perm].reshape(B, -1).to(dtype) - - # --- Local windows: [B, 50] u16 -> [B, 375] each --------------------- - # Local windows only have 15 planes (no player positions) - locals_raw = packed_obs[:, OFFSET_LOCALS:OFFSET_POSITIONS].to(torch.int32) - local_shift = plane_shift[:TILE_BITMASK_PLANES] - local_bits = _unpack_bits(locals_raw, local_shift) # [B, 50, 15] - local_my = local_bits[:, :LOCAL_WINDOW_TILES, :].reshape(B, -1).to(dtype) - local_opp = local_bits[:, LOCAL_WINDOW_TILES:, :].reshape(B, -1).to(dtype) - - # --- Global features: [B, 20] u16 -> [B, 20] float normalized -------- - globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) - globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) - - # --- Turn count for white-to-move sign -------------------------------- - turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to(torch.int32) - - return my_features, opp_features, local_my, local_opp, globals_norm, turn_count - - -__all__ = [ - "BOARD_CELLS", - "BOARD_SIDE", - "GLOBAL_FEATURES", - "GLOBAL_SCALES", - "GLOBAL_TURN_COUNT_IDX", - "LOCAL_FEATURES", - "LOCAL_WINDOW_TILES", - "OBS_WORDS", - "OPP_PLANE_PERM", - "TILE_BITMASK_PLANES", - "TILE_PLANES", - "TOTAL_TILE_FEATURES", - "decode_nnue_obs", -] diff --git a/python/scripts/benchmark_decode_forward.py b/python/scripts/benchmark_decode_forward.py new file mode 100644 index 0000000..912f292 --- /dev/null +++ b/python/scripts/benchmark_decode_forward.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Benchmark NNUE decode vs full `PackedValueModel` forward on GPU. + +Run on the training machine (e.g. 4×L40): + + PYTHONPATH=python python python/scripts/benchmark_decode_forward.py --batch 256 + +Prints median milliseconds per step and ``decode_time / forward_time`` so you +can validate the decode fraction (often ~15–35% depending on batch and head). +""" + +from __future__ import annotations + +import argparse +import time +from collections.abc import Callable + +import torch + +from alphapaint_training.model import PackedValueModel +from alphapaint_training.decode_triton import OBS_WORDS, OFFSET_GLOBALS, OFFSET_POSITIONS + + +def _random_packed(batch: int, device: torch.device) -> torch.Tensor: + g = torch.Generator(device=device).manual_seed(0) + x = torch.randint(0, 65535, (batch, OBS_WORDS), dtype=torch.uint16, device=device, generator=g) + for o in range(4): + x[:, OFFSET_POSITIONS + o] = torch.randint(0, 32, (batch,), device=device, dtype=torch.uint16) + x[:, OFFSET_GLOBALS + 17] = torch.randint(0, 2000, (batch,), device=device, dtype=torch.uint16) + return x + + +def _median_ms( + fn: Callable[[], None], + *, + warmup: int, + iters: int, +) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters * 1000.0 + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--batch", type=int, default=256) + p.add_argument("--warmup", type=int, default=20) + p.add_argument("--iters", type=int, default=100) + args = p.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA required") + + device = torch.device("cuda") + x = _random_packed(args.batch, device) + m = PackedValueModel().to(device) + + def decode_only() -> None: + m.decode(x) + + def forward_only() -> None: + m(x) + + _ = m(x) # alloc buffers + + d_ms = _median_ms(decode_only, warmup=args.warmup, iters=args.iters) + f_ms = _median_ms(forward_only, warmup=args.warmup, iters=args.iters) + ratio = d_ms / f_ms if f_ms > 0 else float("nan") + + print(f"batch={args.batch}") + print(f" decode median: {d_ms:.3f} ms") + print(f" forward median: {f_ms:.3f} ms") + print(f" decode / forward: {ratio:.3f} ({100.0 * ratio:.1f}%)") + + +if __name__ == "__main__": + main() diff --git a/python/scripts/export_value_net.py b/python/scripts/export_value_net.py index 21ae238..62542d6 100644 --- a/python/scripts/export_value_net.py +++ b/python/scripts/export_value_net.py @@ -101,10 +101,7 @@ def _build_fixtures( state_dict = _extract_state_dict(checkpoint) from alphapaint_training.model import PackedValueModel - model = PackedValueModel( - board_dtype=torch.float32, - **_infer_model_config(state_dict), - ) + model = PackedValueModel(**_infer_model_config(state_dict)) model.load_state_dict(state_dict) model.eval() From 702691a8e338b508605c18c52108e7b61e9d8b0f Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Tue, 31 Mar 2026 04:33:06 -0400 Subject: [PATCH 58/59] Revert "switch packed format decode to custom triton kernel" This reverts commit 7d5fd63938e1993015c275119d8259bf9766f3b8. --- python/alphapaint_training/__init__.py | 2 +- python/alphapaint_training/decode_triton.py | 257 -------------------- python/alphapaint_training/model.py | 42 +++- python/alphapaint_training/packed_obs.py | 204 ++++++++++++++++ python/scripts/benchmark_decode_forward.py | 82 ------- python/scripts/export_value_net.py | 5 +- 6 files changed, 241 insertions(+), 351 deletions(-) delete mode 100644 python/alphapaint_training/decode_triton.py create mode 100644 python/alphapaint_training/packed_obs.py delete mode 100644 python/scripts/benchmark_decode_forward.py diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py index 9f544d2..d9e640d 100644 --- a/python/alphapaint_training/__init__.py +++ b/python/alphapaint_training/__init__.py @@ -11,7 +11,7 @@ ) from .logger import TrainingLogger from .model import NnueValueNet, PackedValueModel -from .decode_triton import ( +from .packed_obs import ( BOARD_CELLS, BOARD_SIDE, GLOBAL_FEATURES, diff --git a/python/alphapaint_training/decode_triton.py b/python/alphapaint_training/decode_triton.py deleted file mode 100644 index 200b20a..0000000 --- a/python/alphapaint_training/decode_triton.py +++ /dev/null @@ -1,257 +0,0 @@ -"""Packed NNUE observation layout + decode (CUDA + Triton only). - -The Rust encoder packs board state into a flat u16 array with this layout: - [0, 1024) — per-tile 15-bit bitmask (32×32 board) - [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) - [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) - [1078, 1098) — 20 global scalar features - -Decode fuses tile bitmask unpack, position one-hots, opponent plane permutation (accum), -and local-window unpack. Globals + turn_count use PyTorch on the result. -""" - -from __future__ import annotations - -import torch - -BOARD_SIDE = 32 -BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 -TILE_BITMASK_PLANES = 15 -TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 -LOCAL_WINDOW_TILES = 25 -GLOBAL_FEATURES = 20 -TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 -LOCAL_TILE_PLANES = TILE_BITMASK_PLANES -LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 - -OFFSET_LOCALS = BOARD_CELLS # 1024 -OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 -OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 -OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 - -GLOBAL_TURN_COUNT_IDX = 17 - -OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] - -GLOBAL_SCALES: tuple[float, ...] = ( - 420.0, - 420.0, - 420.0, - 420.0, - 1024.0, - 1024.0, - 8.0, - 8.0, - 1024.0, - 1024.0, - 8.0, - 1024.0, - 1024.0, - 64.0, - 64.0, - 64.0, - 32.0, - 2000.0, - 32.0, - 32.0, -) - - -def _check_packed_obs(packed_obs: torch.Tensor) -> None: - if packed_obs.dtype != torch.uint16: - raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") - if packed_obs.ndim != 2: - raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") - if packed_obs.shape[1] != OBS_WORDS: - raise ValueError( - f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" - ) - - -try: - import triton - import triton.language as tl - - _TRITON_AVAILABLE = True -except ImportError: - _TRITON_AVAILABLE = False - triton = None # type: ignore[assignment] - tl = None # type: ignore[assignment] - - -if _TRITON_AVAILABLE: - _OPP_PERM17 = tl.constexpr(tuple(OPP_PLANE_PERM)) - - @triton.jit - def _decode_accum_kernel( - packed_ptr, - my_out_ptr, - opp_out_ptr, - row_stride: tl.constexpr, - off_positions: tl.constexpr, - board_side: tl.constexpr, - board_cells: tl.constexpr, - tile_planes: tl.constexpr, - ) -> None: - pid = tl.program_id(0) - b = pid // board_cells - cell = pid - b * board_cells - - base_tile = b * row_stride + cell - w = tl.load(packed_ptr + base_tile).to(tl.int32) - - pos0 = b * row_stride + off_positions - my_x = tl.load(packed_ptr + pos0 + 0).to(tl.int32) - my_y = tl.load(packed_ptr + pos0 + 1).to(tl.int32) - opp_x = tl.load(packed_ptr + pos0 + 2).to(tl.int32) - opp_y = tl.load(packed_ptr + pos0 + 3).to(tl.int32) - my_flat = my_x * board_side + my_y - opp_flat = opp_x * board_side + opp_y - cell_i = cell.to(tl.int32) - my_pos = tl.where(cell_i == my_flat, 1.0, 0.0) - opp_pos = tl.where(cell_i == opp_flat, 1.0, 0.0) - - out_base = b * (board_cells * tile_planes) + cell * tile_planes - - for p in tl.static_range(15): - v = ((w >> p) & 1).to(tl.float32) - tl.store(my_out_ptr + out_base + p, v.to(tl.bfloat16)) - tl.store(my_out_ptr + out_base + 15, my_pos.to(tl.bfloat16)) - tl.store(my_out_ptr + out_base + 16, opp_pos.to(tl.bfloat16)) - - for out_p in tl.static_range(17): - src = _OPP_PERM17[out_p] - if src < 15: - v = ((w >> src) & 1).to(tl.float32) - elif src == 15: - v = my_pos - else: - v = opp_pos - tl.store(opp_out_ptr + out_base + out_p, v.to(tl.bfloat16)) - - @triton.jit - def _decode_locals_kernel( - packed_ptr, - local_my_ptr, - local_opp_ptr, - row_stride: tl.constexpr, - off_locals: tl.constexpr, - local_tiles: tl.constexpr, - bitmask_planes: tl.constexpr, - ) -> None: - pid = tl.program_id(0) - nloc = 2 * local_tiles - b = pid // nloc - li = pid - b * nloc - - idx = b * row_stride + off_locals + li - w = tl.load(packed_ptr + idx).to(tl.int32) - - for p in tl.static_range(15): - v = ((w >> p) & 1).to(tl.bfloat16) - dst_my = b * (local_tiles * bitmask_planes) + li * bitmask_planes + p - li2 = li - local_tiles - dst_opp = b * (local_tiles * bitmask_planes) + li2 * bitmask_planes + p - tl.store(local_my_ptr + dst_my, v, mask=li < local_tiles) - tl.store(local_opp_ptr + dst_opp, v, mask=li >= local_tiles) - - def decode_nnue_obs( - packed_obs: torch.Tensor, - *, - global_scales: torch.Tensor, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: - """Decode packed observation to bf16 features (CUDA + Triton only).""" - if not torch.cuda.is_available() or not packed_obs.is_cuda: - raise RuntimeError("decode_nnue_obs requires CUDA device tensors") - _check_packed_obs(packed_obs) - - packed_obs = packed_obs.contiguous() - device = packed_obs.device - B = packed_obs.shape[0] - row_stride = packed_obs.stride(0) - if packed_obs.stride(1) != 1: - raise ValueError("packed_obs must be contiguous in last dimension") - - dtype = torch.bfloat16 - my_features = torch.empty(B, TOTAL_TILE_FEATURES, dtype=dtype, device=device) - opp_features = torch.empty(B, TOTAL_TILE_FEATURES, dtype=dtype, device=device) - local_my = torch.empty( - B, LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES, dtype=dtype, device=device - ) - local_opp = torch.empty( - B, LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES, dtype=dtype, device=device - ) - - grid_acc = (B * BOARD_CELLS,) - _decode_accum_kernel[grid_acc]( - packed_obs, - my_features, - opp_features, - row_stride=row_stride, - off_positions=OFFSET_POSITIONS, - board_side=BOARD_SIDE, - board_cells=BOARD_CELLS, - tile_planes=TILE_PLANES, - ) - - grid_loc = (B * 2 * LOCAL_WINDOW_TILES,) - _decode_locals_kernel[grid_loc]( - packed_obs, - local_my, - local_opp, - row_stride=row_stride, - off_locals=OFFSET_LOCALS, - local_tiles=LOCAL_WINDOW_TILES, - bitmask_planes=TILE_BITMASK_PLANES, - ) - - globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) - globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) - turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to( - torch.int32 - ) - - return my_features, opp_features, local_my, local_opp, globals_norm, turn_count - -else: - - def decode_nnue_obs( - packed_obs: torch.Tensor, - *, - global_scales: torch.Tensor, - ) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ]: - raise RuntimeError("decode_nnue_obs requires Triton (linux dependency)") - - -__all__ = [ - "BOARD_CELLS", - "BOARD_SIDE", - "GLOBAL_FEATURES", - "GLOBAL_SCALES", - "GLOBAL_TURN_COUNT_IDX", - "LOCAL_FEATURES", - "LOCAL_WINDOW_TILES", - "OBS_WORDS", - "OFFSET_GLOBALS", - "OFFSET_LOCALS", - "OFFSET_POSITIONS", - "OPP_PLANE_PERM", - "TILE_BITMASK_PLANES", - "TILE_PLANES", - "TOTAL_TILE_FEATURES", - "decode_nnue_obs", -] diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py index 5b0c67a..cf9f015 100644 --- a/python/alphapaint_training/model.py +++ b/python/alphapaint_training/model.py @@ -13,13 +13,17 @@ from __future__ import annotations +from typing import cast + import torch from torch import nn -from .decode_triton import ( +from .packed_obs import ( GLOBAL_FEATURES, GLOBAL_SCALES, LOCAL_FEATURES, + OPP_PLANE_PERM, + TILE_BITMASK_PLANES, TOTAL_TILE_FEATURES, decode_nnue_obs, ) @@ -106,8 +110,10 @@ def __init__( local_hidden_dim: int = 64, fc1_dim: int = 128, fc2_dim: int = 32, + board_dtype: torch.dtype = torch.bfloat16, ): super().__init__() + self.board_dtype = board_dtype self.value_net = NnueValueNet( acc_dim=acc_dim, local_hidden_dim=local_hidden_dim, @@ -115,6 +121,18 @@ def __init__( fc2_dim=fc2_dim, ) + # Decode buffers — registered so they follow .to(device) and are + # captured in CUDA graphs. + self.register_buffer( + "_plane_shift", + torch.arange(TILE_BITMASK_PLANES, dtype=torch.int32), + persistent=False, + ) + self.register_buffer( + "_opp_perm", + torch.tensor(OPP_PLANE_PERM, dtype=torch.long), + persistent=False, + ) self.register_buffer( "_global_scales", torch.tensor(GLOBAL_SCALES, dtype=torch.float32), @@ -132,21 +150,25 @@ def decode( torch.Tensor, ]: return decode_nnue_obs( - packed_obs, global_scales=self._global_scales + packed_obs, + plane_shift=cast(torch.Tensor, self._plane_shift), + opp_perm=cast(torch.Tensor, self._opp_perm), + global_scales=cast(torch.Tensor, self._global_scales), + dtype=self.board_dtype, ) def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: my_features, opp_features, local_my, local_opp, globals_norm, turn_count = ( self.decode(packed_obs) ) - # Decode is always bf16; align to parameter dtype when not under autocast. - wdt = next(self.value_net.parameters()).dtype - if my_features.dtype != wdt: - my_features = my_features.to(wdt) - opp_features = opp_features.to(wdt) - local_my = local_my.to(wdt) - local_opp = local_opp.to(wdt) - globals_norm = globals_norm.to(wdt) + + if not torch.is_autocast_enabled(): + param_dtype = next(self.value_net.parameters()).dtype + my_features = my_features.to(dtype=param_dtype) + opp_features = opp_features.to(dtype=param_dtype) + local_my = local_my.to(dtype=param_dtype) + local_opp = local_opp.to(dtype=param_dtype) + globals_norm = globals_norm.to(dtype=param_dtype) value = self.value_net( my_features, opp_features, local_my, local_opp, globals_norm diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py new file mode 100644 index 0000000..6089baf --- /dev/null +++ b/python/alphapaint_training/packed_obs.py @@ -0,0 +1,204 @@ +"""Packed observation decoding for NNUE training. + +The Rust encoder packs board state into a flat u16 array with this layout: + [0, 1024) — per-tile 15-bit bitmask (32×32 board) + [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) + [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) + [1078, 1098) — 20 global scalar features + +Tile bitmask planes (from current player's perspective): + bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) + bits 4-7: opponent paint thermometer + bit 8: wall + bit 9: powerup + bit 10: current player beacon + bit 11: opponent beacon + bit 12: hill neutral + bit 13: hill current player + bit 14: hill opponent + +Additional planes added during decoding: + plane 15: one-hot my player location (derived from positions) + plane 16: one-hot opponent player location (derived from positions) +""" + +from __future__ import annotations + +import torch + +# Layout constants ---------------------------------------------------------- + +BOARD_SIDE = 32 +BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 +# Bits 0–14 from each u16 tile word (Rust encoder); +2 derived position planes → accumulator +TILE_BITMASK_PLANES = 15 +TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 +LOCAL_WINDOW_TILES = 25 +GLOBAL_FEATURES = 20 +TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 +LOCAL_TILE_PLANES = TILE_BITMASK_PLANES # Local windows don't have player position planes +LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 + +OFFSET_LOCALS = BOARD_CELLS # 1024 +OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 +OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 +OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 + +# Index of turn_count within the global section +GLOBAL_TURN_COUNT_IDX = 17 + +# Plane permutation for opponent perspective: +# swap my_paint[0:4] <-> opp_paint[4:8] +# swap my_beacon[10] <-> opp_beacon[11] +# swap my_hill[13] <-> opp_hill[14] +# swap my_pos[15] <-> opp_pos[16] +# wall[8], powerup[9], hill_neutral[12] stay +OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] + +# Normalization divisors for the 20 global features. +# Brings raw u16 values into roughly [0, 1] range. +GLOBAL_SCALES: tuple[float, ...] = ( + 420.0, # 0 my_stamina + 420.0, # 1 opp_stamina + 420.0, # 2 my_max_stamina + 420.0, # 3 opp_max_stamina + 1024.0, # 4 my_territory + 1024.0, # 5 opp_territory + 8.0, # 6 my_hills + 8.0, # 7 opp_hills + 1024.0, # 8 my_hill_tiles + 1024.0, # 9 opp_hill_tiles + 8.0, # 10 contested_hills + 1024.0, # 11 my_beacons + 1024.0, # 12 opp_beacons + 64.0, # 13 player_dist + 64.0, # 14 my_hill_dist + 64.0, # 15 opp_hill_dist + 32.0, # 16 consecutive_moves + 2000.0, # 17 turn_count + 32.0, # 18 rows + 32.0, # 19 cols +) + + +def _check_packed_obs(packed_obs: torch.Tensor) -> None: + if packed_obs.dtype != torch.uint16: + raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") + if packed_obs.ndim != 2: + raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") + if packed_obs.shape[1] != OBS_WORDS: + raise ValueError( + f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" + ) + + +def _unpack_bits(words: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + """Extract binary planes from u16 words (one channel per shift index). + + Args: + words: [...] int32 tensor of packed bitmasks + shift: [K] int32 tensor of bit indices (e.g. arange(15) for tile bitmasks) + + Returns: + [..., K] float tensor of 0/1 values + """ + return ((words.unsqueeze(-1) >> shift) & 1).float() + + +def decode_nnue_obs( + packed_obs: torch.Tensor, + *, + plane_shift: torch.Tensor, + opp_perm: torch.Tensor, + global_scales: torch.Tensor, + dtype: torch.dtype = torch.bfloat16, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Decode packed NNUE observation into model inputs. + + All buffers (plane_shift, opp_perm, global_scales) should be registered + as model buffers so they live on the right device and are captured by + CUDA graphs. + + Returns: + my_features: [B, 17408] dtype — accumulator input (my perspective) + opp_features: [B, 17408] dtype — accumulator input (opponent perspective) + local_my: [B, 375] dtype — local window around my position + local_opp: [B, 375] dtype — local window around opponent position + globals_norm: [B, 20] dtype — normalized global scalars + turn_count: [B] int32 — for white-to-move sign + """ + B = packed_obs.shape[0] + + # --- Tile features: [B, 1024] u16 -> [B, 17408] ---------------------- + tiles = packed_obs[:, :BOARD_CELLS].to(torch.int32) + # Only bits 0–14 are defined on tile words; planes 15–16 come from positions below. + bits = _unpack_bits(tiles, plane_shift[:TILE_BITMASK_PLANES]) # [B, 1024, 15] + + # --- Player positions -> one-hot planes [B, 1024, 2] ----------------- + my_x = packed_obs[:, OFFSET_POSITIONS].long() # [B] + my_y = packed_obs[:, OFFSET_POSITIONS + 1].long() # [B] + opp_x = packed_obs[:, OFFSET_POSITIONS + 2].long() # [B] + opp_y = packed_obs[:, OFFSET_POSITIONS + 3].long() # [B] + + my_pos_flat = my_x * BOARD_SIDE + my_y # [B] + opp_pos_flat = opp_x * BOARD_SIDE + opp_y # [B] + + my_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + opp_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + my_pos_plane.scatter_(1, my_pos_flat.unsqueeze(1), 1.0) + opp_pos_plane.scatter_(1, opp_pos_flat.unsqueeze(1), 1.0) + + # Concatenate: [B, 1024, 15] + [B, 1024, 1] + [B, 1024, 1] = [B, 1024, 17] + bits_with_pos = torch.cat( + [bits, my_pos_plane.unsqueeze(-1), opp_pos_plane.unsqueeze(-1)], dim=-1 + ) # [B, 1024, 17] + + my_features = bits_with_pos.reshape(B, -1).to(dtype) + + # Opponent perspective: permute planes + opp_features = bits_with_pos[:, :, opp_perm].reshape(B, -1).to(dtype) + + # --- Local windows: [B, 50] u16 -> [B, 375] each --------------------- + # Local windows only have 15 planes (no player positions) + locals_raw = packed_obs[:, OFFSET_LOCALS:OFFSET_POSITIONS].to(torch.int32) + local_shift = plane_shift[:TILE_BITMASK_PLANES] + local_bits = _unpack_bits(locals_raw, local_shift) # [B, 50, 15] + local_my = local_bits[:, :LOCAL_WINDOW_TILES, :].reshape(B, -1).to(dtype) + local_opp = local_bits[:, LOCAL_WINDOW_TILES:, :].reshape(B, -1).to(dtype) + + # --- Global features: [B, 20] u16 -> [B, 20] float normalized -------- + globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) + globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) + + # --- Turn count for white-to-move sign -------------------------------- + turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to(torch.int32) + + return my_features, opp_features, local_my, local_opp, globals_norm, turn_count + + +__all__ = [ + "BOARD_CELLS", + "BOARD_SIDE", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "GLOBAL_TURN_COUNT_IDX", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", + "OBS_WORDS", + "OPP_PLANE_PERM", + "TILE_BITMASK_PLANES", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", + "decode_nnue_obs", +] diff --git a/python/scripts/benchmark_decode_forward.py b/python/scripts/benchmark_decode_forward.py deleted file mode 100644 index 912f292..0000000 --- a/python/scripts/benchmark_decode_forward.py +++ /dev/null @@ -1,82 +0,0 @@ -#!/usr/bin/env python3 -"""Benchmark NNUE decode vs full `PackedValueModel` forward on GPU. - -Run on the training machine (e.g. 4×L40): - - PYTHONPATH=python python python/scripts/benchmark_decode_forward.py --batch 256 - -Prints median milliseconds per step and ``decode_time / forward_time`` so you -can validate the decode fraction (often ~15–35% depending on batch and head). -""" - -from __future__ import annotations - -import argparse -import time -from collections.abc import Callable - -import torch - -from alphapaint_training.model import PackedValueModel -from alphapaint_training.decode_triton import OBS_WORDS, OFFSET_GLOBALS, OFFSET_POSITIONS - - -def _random_packed(batch: int, device: torch.device) -> torch.Tensor: - g = torch.Generator(device=device).manual_seed(0) - x = torch.randint(0, 65535, (batch, OBS_WORDS), dtype=torch.uint16, device=device, generator=g) - for o in range(4): - x[:, OFFSET_POSITIONS + o] = torch.randint(0, 32, (batch,), device=device, dtype=torch.uint16) - x[:, OFFSET_GLOBALS + 17] = torch.randint(0, 2000, (batch,), device=device, dtype=torch.uint16) - return x - - -def _median_ms( - fn: Callable[[], None], - *, - warmup: int, - iters: int, -) -> float: - for _ in range(warmup): - fn() - torch.cuda.synchronize() - t0 = time.perf_counter() - for _ in range(iters): - fn() - torch.cuda.synchronize() - return (time.perf_counter() - t0) / iters * 1000.0 - - -def main() -> None: - p = argparse.ArgumentParser() - p.add_argument("--batch", type=int, default=256) - p.add_argument("--warmup", type=int, default=20) - p.add_argument("--iters", type=int, default=100) - args = p.parse_args() - - if not torch.cuda.is_available(): - raise SystemExit("CUDA required") - - device = torch.device("cuda") - x = _random_packed(args.batch, device) - m = PackedValueModel().to(device) - - def decode_only() -> None: - m.decode(x) - - def forward_only() -> None: - m(x) - - _ = m(x) # alloc buffers - - d_ms = _median_ms(decode_only, warmup=args.warmup, iters=args.iters) - f_ms = _median_ms(forward_only, warmup=args.warmup, iters=args.iters) - ratio = d_ms / f_ms if f_ms > 0 else float("nan") - - print(f"batch={args.batch}") - print(f" decode median: {d_ms:.3f} ms") - print(f" forward median: {f_ms:.3f} ms") - print(f" decode / forward: {ratio:.3f} ({100.0 * ratio:.1f}%)") - - -if __name__ == "__main__": - main() diff --git a/python/scripts/export_value_net.py b/python/scripts/export_value_net.py index 62542d6..21ae238 100644 --- a/python/scripts/export_value_net.py +++ b/python/scripts/export_value_net.py @@ -101,7 +101,10 @@ def _build_fixtures( state_dict = _extract_state_dict(checkpoint) from alphapaint_training.model import PackedValueModel - model = PackedValueModel(**_infer_model_config(state_dict)) + model = PackedValueModel( + board_dtype=torch.float32, + **_infer_model_config(state_dict), + ) model.load_state_dict(state_dict) model.eval() From 8a4b6dd66b193eafa2dd1428fa4d3cefa5399e50 Mon Sep 17 00:00:00 2001 From: Rohan Godha Date: Tue, 31 Mar 2026 04:58:26 -0400 Subject: [PATCH 59/59] fix stack overflow issue for DoubleArray::splat --- alpha_paint/src/board/structs.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/alpha_paint/src/board/structs.rs b/alpha_paint/src/board/structs.rs index 97c869d..935bbe0 100644 --- a/alpha_paint/src/board/structs.rs +++ b/alpha_paint/src/board/structs.rs @@ -150,8 +150,15 @@ impl DoubleArray32x32 { where T: Copy, { - let arr = array::from_fn(|_| array::from_fn(|_| Array32x32::splat(value))); - DoubleArray32x32(Box::new(Array32x32(arr))) + // Allocate on heap to avoid stack overflow (~2MB array) + let vec: Vec<[Array32x32; 32]> = (0..32) + .map(|_| array::from_fn(|_| Array32x32::splat(value))) + .collect(); + let boxed_slice: Box<[[Array32x32; 32]]> = vec.into_boxed_slice(); + assert_eq!(boxed_slice.len(), 32); + let ptr = Box::into_raw(boxed_slice); + let boxed_arr: Box<[[Array32x32; 32]; 32]> = unsafe { Box::from_raw(ptr as *mut _) }; + DoubleArray32x32(Box::new(Array32x32(*boxed_arr))) } }