factor out nn_train.py into a separate python project to make it more manageable, add plotting and offline strength fit

6 files changed, 971 insertions, 0 deletions

diff --git a/train/src/ctak_train/__init__.py b/train/src/ctak_train/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/train/src/ctak_train/__init__.py
diff --git a/train/src/ctak_train/candidate.py b/train/src/ctak_train/candidate.py
new file mode 100644
index 0000000..aa2a6a6
--- /dev/null
+++ b/train/src/ctak_train/candidate.py
@@ -0,0 +1,125 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import math
+from dataclasses import dataclass
+from itertools import product
+from pathlib import Path
+
+import torch
+from torch.nn import MSELoss
+from tqdm import tqdm
+
+from ctak_train.data import DataSplit, Pool
+from ctak_train.net import NET_DIMS, Net
+from ctak_train.selfplay import cand_copy_for
+
+
+@dataclass(frozen=True)
+class DataGenConfig:
+    size: int
+    depth: int
+    min_ply: int
+    games: int
+    procs: int
+    replay_iters: int
+    root: Path
+    train_only_self_play: bool
+
+
+@dataclass(frozen=True)
+class TrainingGridSearchConfig:
+    smooth_factor: list[float]
+    epochs: list[int]
+    batch: list[int]
+    lr: list[float]
+    val_frac: list[float]
+
+    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
+        inp_dim, hidden = NET_DIMS[data_gen_config.size]
+        flds = self.__dataclass_fields__.keys()
+        almost_cands = [
+            dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
+        ]
+        for ident, almost_cand in enumerate(almost_cands):
+            net = Net(inp_dim, hidden)
+            net.to_c(
+                size=data_gen_config.size,
+                ident=ident,
+                path=data_gen_config.root
+                / "include"
+                / f"cand_{ident}_weights_{data_gen_config.size}.c",
+                note=f"{ident=}, v000, random init",
+            )
+            cand_copy_for(
+                6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
+            )
+            almost_cand["net"] = net
+            almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
+        return [Candidate(**d) for d in almost_cands]
+
+
+@dataclass
+class Candidate:
+    train_only_self_play: bool
+    smooth_factor: float
+    epochs: int
+    batch: int
+    lr: float
+    val_frac: float
+    net: "Net"
+
+    def __str__(self) -> str:
+        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"
+
+    def training_pool(self, pool: Pool) -> Pool:
+        if self.train_only_self_play:
+            return pool.select_chains(pool.chain_self_play)
+        return pool
+
+    def fit(self, pool: Pool) -> tuple[float, float]:
+        # The train/val split is fixed but targets are refreshed
+        split = DataSplit.build(self.training_pool(pool), self.val_frac)
+        if split.n == 0:
+            return math.nan, math.nan
+
+        optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
+        loss_fn = MSELoss()
+
+        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
+        for _ in pbar:
+            split.refresh_targets(self.net, self.smooth_factor)
+
+            self.net.train()
+            for head, cells, parity, target in split.batches(self.batch):
+                optimizer.zero_grad()
+                predictions = self.net(head, cells) * parity
+                loss = loss_fn(predictions, target)
+                loss.backward()
+                optimizer.step()
+
+            with torch.no_grad():
+                self.net.eval()
+                val = split.val_idx
+                predictions = self.net(split.pool.head[val], split.pool.cells[val])
+                predictions = predictions * split.pool.parity[val]
+                targets = split.targets[val]
+                mse = loss_fn(predictions, targets).item()
+                corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()
+
+            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
+
+        return corr, mse
diff --git a/train/src/ctak_train/data.py b/train/src/ctak_train/data.py
new file mode 100644
index 0000000..6dc9ec0
--- /dev/null
+++ b/train/src/ctak_train/data.py
@@ -0,0 +1,300 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import json
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from ctak_train.net import NET_DIMS
+from torch import Tensor
+from torch.nn import Module
+
+
+@dataclass
+class GameLog:
+    iteration: np.ndarray  # [G] iteration the game was played in
+    black: np.ndarray  # [G] candidate id playing black
+    white: np.ndarray  # [G] candidate id playing white
+    result: np.ndarray  # [G] outcome from black's perspective, {-1, 0, 1}
+
+    @classmethod
+    def empty(cls) -> GameLog:
+        return cls(
+            iteration=np.empty(0, dtype=np.int32),
+            black=np.empty(0, dtype=np.int16),
+            white=np.empty(0, dtype=np.int16),
+            result=np.empty(0, dtype=np.float32),
+        )
+
+    def append(self, newer: GameLog) -> None:
+        self.iteration = np.concatenate([self.iteration, newer.iteration])
+        self.black = np.concatenate([self.black, newer.black])
+        self.white = np.concatenate([self.white, newer.white])
+        self.result = np.concatenate([self.result, newer.result])
+
+
+@dataclass(frozen=True)
+class Pool:
+    head: Tensor  # [n, 2] fixed head inputs of the root state s_t
+    cells: Tensor  # [n, n_cells] cell states 0..6 of s_t
+    leaf_head: Tensor  # [n, 2] head inputs of the PV leaf pv_t
+    leaf_cells: Tensor  # [n, n_cells] cell states of pv_t
+    parity: Tensor  # [n] +1 when black is to move at s_t
+    leaf_parity: Tensor  # [n] +1 when black is to move at pv_t
+    terminal: Tensor  # [n] leaf is an exact terminal
+    terminal_value: Tensor  # [n] exact terminal value, black's perspective
+    chain_lens: Tensor  # [C] rows per chain
+    chain_result: Tensor  # [C] game result, black's perspective
+    chain_self_play: Tensor  # [C] chain comes from a self vs self game
+
+    @classmethod
+    def empty(cls, n_cells: int) -> Pool:
+        def rows(dtype: torch.dtype, width: int = 0) -> Tensor:
+            return torch.empty((0, width) if width else (0,), dtype=dtype)
+
+        return cls(
+            head=rows(torch.float32, 2),
+            cells=rows(torch.uint8, n_cells),
+            leaf_head=rows(torch.float32, 2),
+            leaf_cells=rows(torch.uint8, n_cells),
+            parity=rows(torch.float32),
+            leaf_parity=rows(torch.float32),
+            terminal=rows(torch.bool),
+            terminal_value=rows(torch.float32),
+            chain_lens=rows(torch.int64),
+            chain_result=rows(torch.float32),
+            chain_self_play=rows(torch.bool),
+        )
+
+    @property
+    def n_rows(self) -> int:
+        return self.head.shape[0]
+
+    def select_chains(self, keep: Tensor) -> Pool:
+        keep = keep.bool()
+        rows = torch.repeat_interleave(
+            torch.arange(self.chain_lens.shape[0], dtype=torch.int64), self.chain_lens
+        )
+        idx = torch.nonzero(keep[rows]).flatten()
+        return Pool(
+            head=self.head[idx],
+            cells=self.cells[idx],
+            leaf_head=self.leaf_head[idx],
+            leaf_cells=self.leaf_cells[idx],
+            parity=self.parity[idx],
+            leaf_parity=self.leaf_parity[idx],
+            terminal=self.terminal[idx],
+            terminal_value=self.terminal_value[idx],
+            chain_lens=self.chain_lens[keep],
+            chain_result=self.chain_result[keep],
+            chain_self_play=self.chain_self_play[keep],
+        )
+
+
+_ROW_FIELDS = (
+    "head",
+    "cells",
+    "leaf_head",
+    "leaf_cells",
+    "parity",
+    "leaf_parity",
+    "terminal",
+    "terminal_value",
+)
+
+
+def _chain_arrays(
+    rows: list[dict], n_cells: int, source: Path
+) -> dict[str, np.ndarray]:
+
+    def column(part: str, key: str, dtype: np.dtype) -> np.ndarray:
+        return np.array([r[part][key] for r in rows], dtype=dtype)
+
+    out = {
+        "head": column("root", "head", np.float32),
+        "cells": column("root", "cells", np.uint8),
+        "leaf_head": column("leaf", "head", np.float32),
+        "leaf_cells": column("leaf", "cells", np.uint8),
+        "parity": column("root", "parity", np.float32),
+        "leaf_parity": column("leaf", "parity", np.float32),
+        "terminal": column("leaf", "terminal", np.bool_),
+        "terminal_value": column("leaf", "terminal_value", np.float32),
+    }
+    for name in ("cells", "leaf_cells"):
+        if out[name].shape[1] != n_cells:
+            raise ValueError(
+                f"{source}: {name} width {out[name].shape[1]}, expected {n_cells}"
+            )
+    return out
+
+
+def _game_colours(game: dict) -> tuple[int, int]:
+    a, b = int(game["cand_A"]), int(game["cand_B"])
+    for r in game["rows"]:
+        if r["root"]["parity"] > 0:
+            cand = int(r["cand"])
+            return cand, b if cand == a else a
+    recorded = {int(r["cand"]) for r in game["rows"]}
+    return (a, b) if a not in recorded else (b, a)
+
+
+def load_window(
+    files: list[Path], size: int, n_candidates: int, iteration: int
+) -> tuple[list[Pool], GameLog]:
+    """Training pools for every file in the window, plus the game records of
+    the newest file only - the caller's GameLog already holds everything older."""
+    n_cells = NET_DIMS[size][0] - 2
+    blocks = [{f: [] for f in _ROW_FIELDS} for _ in range(n_candidates)]
+    chain_lens: list[list[int]] = [[] for _ in range(n_candidates)]
+    chain_result: list[list[float]] = [[] for _ in range(n_candidates)]
+    chain_self_play: list[list[bool]] = [[] for _ in range(n_candidates)]
+    black, white, result = [], [], []
+
+    for file_index, path in enumerate(files):
+        is_newest = file_index == len(files) - 1
+        with open(path) as f:
+            for line in f:
+                if not line.strip():
+                    continue
+                game = json.loads(line)
+                game_result = float(game["result"])
+                self_play = int(game["cand_A"]) == int(game["cand_B"])
+                if is_newest:
+                    b, w = _game_colours(game)
+                    black.append(b)
+                    white.append(w)
+                    result.append(game_result)
+
+                by_cand: dict[int, list[dict]] = {}
+                for r in game["rows"]:
+                    by_cand.setdefault(int(r["cand"]), []).append(r)
+
+                for cand, rows in by_cand.items():
+                    parity = [r["root"]["parity"] for r in rows]
+                    for sign in (1, -1):
+                        chain = [
+                            r for r, p in zip(rows, parity) if (p > 0) == (sign > 0)
+                        ]
+                        if not chain:
+                            continue
+                        for name, array in _chain_arrays(chain, n_cells, path).items():
+                            blocks[cand][name].append(array)
+                        chain_lens[cand].append(len(chain))
+                        chain_result[cand].append(game_result)
+                        chain_self_play[cand].append(self_play)
+
+    def pool_of(cand: int) -> Pool:
+        if not chain_lens[cand]:
+            return Pool.empty(n_cells)
+        arrays = {
+            name: torch.from_numpy(np.concatenate(blocks[cand][name]))
+            for name in _ROW_FIELDS
+        }
+        return Pool(
+            **arrays,
+            chain_lens=torch.tensor(chain_lens[cand], dtype=torch.int64),
+            chain_result=torch.tensor(chain_result[cand], dtype=torch.float32),
+            chain_self_play=torch.tensor(chain_self_play[cand], dtype=torch.bool),
+        )
+
+    games = GameLog(
+        iteration=np.full(len(black), iteration, dtype=np.int32),
+        black=np.array(black, dtype=np.int16),
+        white=np.array(white, dtype=np.int16),
+        result=np.array(result, dtype=np.float32),
+    )
+    return [pool_of(c) for c in range(n_candidates)], games
+
+
+@dataclass
+class DataSplit:
+    pool: Pool
+    train_idx: Tensor
+    val_idx: Tensor
+    targets: Tensor  # [n] aligned with pool rows, refreshed per epoch
+    chain_end: Tensor  # [C] last row index of each chain
+    wave_order: Tensor  # [n] rows ordered by distance from their chain's end
+    wave_bounds: Tensor  # [d] start of distance d in wave_order
+
+    @property
+    def n(self) -> int:
+        return self.pool.n_rows
+
+    @classmethod
+    def build(cls, pool: Pool, val_frac: float) -> DataSplit:
+        n = pool.n_rows
+        if n == 0:
+            empty = torch.empty(0, dtype=torch.int64)
+            return cls(pool, empty, empty, torch.empty(0), empty, empty, empty)
+
+        # Smoothing is G[t] = (1-λ)V[t] + λG[t+1] with G at the chain end.
+        # Chains are contiguous, so every row that is d plies from its chain's
+        # end can be updated in one vectorised wavefront step per d.
+        lens = pool.chain_lens
+        chain_end = torch.cumsum(lens, dim=0) - 1
+        chain_of_row = torch.repeat_interleave(
+            torch.arange(lens.shape[0], dtype=torch.int64), lens
+        )
+        dist = chain_end[chain_of_row] - torch.arange(n, dtype=torch.int64)
+        wave_order = torch.argsort(dist, stable=True)
+        wave_bounds = torch.searchsorted(
+            dist[wave_order], torch.arange(int(dist.max()) + 2, dtype=torch.int64)
+        )
+
+        n_val = max(1, int(n * val_frac))
+        perm = torch.randperm(n)
+        return cls(
+            pool=pool,
+            train_idx=perm[n_val:],
+            val_idx=perm[:n_val],
+            targets=torch.zeros(n),
+            chain_end=chain_end,
+            wave_order=wave_order,
+            wave_bounds=wave_bounds,
+        )
+
+    def refresh_targets(self, net: Module, lam: float) -> None:
+        pool = self.pool
+        net.eval()
+        with torch.no_grad():
+            # Net(pv_t), black's perspective; exact terminals override the net
+            leaf = net(pool.leaf_head, pool.leaf_cells) * pool.leaf_parity
+            leaf = torch.where(pool.terminal, pool.terminal_value, leaf)
+
+        targets = torch.zeros_like(leaf)
+        targets[self.chain_end] = pool.chain_result
+        for lo, hi in zip(self.wave_bounds[1:-1], self.wave_bounds[2:]):
+            idx = self.wave_order[lo:hi]
+            targets[idx] = (1.0 - lam) * leaf[idx] + lam * targets[idx + 1]
+        self.targets = targets
+
+    def batches(
+        self, batch_size: int
+    ) -> Iterator[tuple[Tensor, Tensor, Tensor, Tensor]]:
+        order = torch.randperm(self.train_idx.shape[0])
+        train_targets = self.targets[self.train_idx]
+        for start in range(0, order.shape[0], batch_size):
+            idx = order[start : start + batch_size]
+            rows = self.train_idx[idx]
+            yield (
+                self.pool.head[rows],
+                self.pool.cells[rows],
+                self.pool.parity[rows],
+                train_targets[idx],
+            )
diff --git a/train/src/ctak_train/net.py b/train/src/ctak_train/net.py
new file mode 100644
index 0000000..cc8bf5c
--- /dev/null
+++ b/train/src/ctak_train/net.py
@@ -0,0 +1,129 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch import Tensor
+from torch.nn import Embedding, Linear, Module
+
+# (input dim, hidden dim) per board size; must match include/weights_<size>.h
+NET_DIMS: dict[int, tuple[int, int]] = {
+    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
+    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
+}
+
+
+def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
+    a = np.ravel(np.asarray(a, dtype=np.float32))
+    lines: list[str] = []
+    for i in range(0, len(a), per_line):
+        chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
+        if i + per_line < len(a):
+            chunk += ","
+        lines.append("    " + chunk)
+    return "\n".join(lines)
+
+
+def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
+    mat = np.asarray(mat, dtype=np.float32)
+    rows: list[str] = []
+    for r in mat:
+        chunks = [
+            ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
+            for i in range(0, len(r), per_line)
+        ]
+        rows.append("    {" + ",\n     ".join(chunks) + "},")
+    return "\n".join(rows)
+
+
+def fast_tanh(x: Tensor) -> Tensor:
+    # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3],
+    # saturating to +-1 outside.
+    x2 = x * x
+    poly = x * (
+        0.997136032
+        + x2
+        * (
+            -0.305215124
+            + x2
+            * (
+                0.0832466371
+                + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))
+            )
+        )
+    )
+    return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly))
+
+
+def fast_silu(x: Tensor) -> Tensor:
+    # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff.
+    # The cutoff below is evaluated at x = 0 (instead of x) because torch
+    # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10.
+    xs = torch.where(x < -4.6, 0.0, x)
+    poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs)
+    return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly))
+
+
+class Net(Module):
+    def __init__(self, inp_dim: int, hidden: int) -> None:
+        super().__init__()
+        # learned cell-state embedding, see nn.c
+        self.embedding = Embedding(7, 1)
+        self.hidden = Linear(inp_dim, hidden, bias=True)
+        self.output = Linear(hidden, 1, bias=True)
+
+    def forward(self, head: Tensor, cells: Tensor) -> Tensor:
+        # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states
+        # (0..6) - the only place they become anything other than integers
+        inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1)
+        hidden = fast_silu(self.hidden(inp))
+        return fast_tanh(self.output(hidden)).squeeze(-1)
+
+    def to_c(
+        self,
+        *,
+        size: int,
+        ident: int | None,
+        path: Path,
+        note: str,
+    ) -> None:
+        upper = "FIVE" if size == 5 else "SIX"
+        prefix = (f"cand_{ident}_" if ident is not None else "") + upper
+        guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "")
+
+        E = self.embedding.weight.detach().cpu().numpy()
+        W = self.hidden.weight.detach().cpu().numpy()
+        B = self.hidden.bias.detach().cpu().numpy()
+        LW = self.output.weight.detach().cpu().numpy().ravel()
+        LB = float(self.output.bias.detach().cpu().item())
+
+        with open(path, "w") as f:
+            lines = [
+                "/* generated by nn_train.py, do not edit by hand */\n",
+                f"/* {note} */\n",
+                f"#ifndef {guard}\n#define {guard}\n",
+                f'#include "weights_{size}.h"\n\n',
+                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
+                f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n",
+                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
+                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
+                f"const float {prefix}_logit_bias = {LB:.9g};\n",
+                "#endif",
+            ]
+            for line in lines:
+                _ = f.write(line)
diff --git a/train/src/ctak_train/nn_train.py b/train/src/ctak_train/nn_train.py
new file mode 100644
index 0000000..d2c2990
--- /dev/null
+++ b/train/src/ctak_train/nn_train.py
@@ -0,0 +1,299 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import argparse
+import os
+import subprocess
+from pathlib import Path
+
+import numpy as np
+
+from ctak_train.candidate import Candidate, DataGenConfig, TrainingGridSearchConfig
+from ctak_train.data import GameLog, load_window
+from ctak_train.net import NET_DIMS
+from ctak_train.selfplay import fan_out
+from ctak_train.strength import fit_strengths, plot_strengths
+
+
+def run_step(
+    *,
+    candidates: list[Candidate],
+    data_gen_config: DataGenConfig,
+    log: GameLog,
+    strength_history: list[np.ndarray],
+    iter: int,
+):
+    def data_path(i: int) -> Path:
+        return (
+            data_gen_config.root
+            / "data"
+            / f"selfplay-{data_gen_config.size}-v{i:03d}.json"
+        )
+
+    print(f"\n== iteration {iter} ==")
+
+    num_candidates = len(candidates)
+    cand_pairs = [
+        (i, j) for i in range(num_candidates) for j in range(i, num_candidates)
+    ]
+
+    # 1. rebuild
+    for i, j in cand_pairs:
+        _ = subprocess.run(
+            ["make", f"selfplay-{i}-{j}"],
+            cwd=data_gen_config.root,
+            check=True,
+            capture_output=True,
+        )
+
+    # 2. generate self-play data
+    data = data_path(iter)
+    print(
+        f"selfplay: {num_candidates=}, {data_gen_config.games=}, "
+        f"{data_gen_config.size=}, {data_gen_config.depth=}, {data_gen_config.min_ply=}, "
+        f"{data_gen_config.train_only_self_play=}"
+    )
+    out_games: list[str] = []
+    for out in fan_out(
+        [
+            [
+                f"./selfplay-{i}-{j}",
+                str(data_gen_config.size),
+                str(data_gen_config.depth),
+                str(data_gen_config.min_ply),
+            ]
+            for (i, j) in cand_pairs
+        ],
+        data_gen_config.games,
+        data_gen_config.procs,
+        "selfplay",
+        data_gen_config.root,
+    ):
+        out_games.extend(out.splitlines())
+    print(f"writing {len(out_games):,} games to {data}")
+    with open(data, "w") as f:
+        _ = f.write("\n".join(out_games))
+        if out_games:
+            _ = f.write("\n")
+
+    # 3. load the training window and append this iteration's games to the log
+    data_files = [
+        data_path(i)
+        for i in range(max(iter - data_gen_config.replay_iters, 1), iter + 1)
+    ]
+    print("Loading data from", [str(c) for c in data_files], "...", end=" ")
+    pools, new_games = load_window(
+        data_files, data_gen_config.size, num_candidates, iter
+    )
+    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")
+    log.append(new_games)
+
+    # 4. fit strengths over the whole history, and plot the trajectory
+    strengths, se = fit_strengths(log, num_candidates)
+    strength_history.append(strengths)
+    plot_strengths(
+        strength_history,
+        candidates,
+        data_gen_config.root / "data" / "candidate_strength.png",
+    )
+
+    rated_games = int(np.count_nonzero(log.black != log.white))
+    ranked = sorted(range(num_candidates), key=lambda i: strengths[i], reverse=True)
+    print(
+        f"Strength fit ({rated_games:,} head-to-head games, "
+        f"iterations 1..{iter}, {rated_games // num_candidates:,} games/candidate):"
+    )
+    for idx in ranked:
+        print(
+            f"\tcandidate {idx:03d}: strength {strengths[idx]:6.1f} "
+            f"± {1.96 * se[idx]:5.1f} Config {candidates[idx]}"
+        )
+
+    # 5. write champ, unconditionally
+    idx = ranked[0]
+    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
+    candidates[idx].net.to_c(
+        size=data_gen_config.size,
+        ident=None,
+        path=path,
+        note=f"{iter=}, strength={strengths[idx]:.1f}",
+    )
+    _ = subprocess.run(
+        ["make", "clean"], cwd=data_gen_config.root, check=True, capture_output=True
+    )
+    _ = subprocess.run(
+        ["make", "native"], cwd=data_gen_config.root, check=True, capture_output=True
+    )
+    print(f"Wrote candidate {idx:03d} to {path!s}")
+
+    # 6. fine-tune on the hybrid targets -> candidate
+    for idx, cand in enumerate(candidates):
+        corr, mse = cand.fit(pools[idx])
+        cand.net.to_c(
+            size=data_gen_config.size,
+            ident=idx,
+            path=data_gen_config.root
+            / "include"
+            / f"cand_{idx}_weights_{data_gen_config.size}.c",
+            note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
+        )
+
+
+def main() -> None:
+    # repo root, relative to this file at <root>/train/src/ctak_train/nn_train.py
+    root = Path(__file__).resolve().parents[3]
+    p = argparse.ArgumentParser(
+        description="Train the neural network evaluators of ctak",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    _ = p.add_argument(
+        "--size",
+        type=int,
+        choices=sorted(NET_DIMS),
+        default=5,
+        help="board size",
+    )
+    _ = p.add_argument(
+        "--iters",
+        type=int,
+        default=100,
+        help="max loop iterations",
+    )
+    _ = p.add_argument(
+        "--games",
+        type=int,
+        default=200,
+        help="games per iteration",
+    )
+    _ = p.add_argument(
+        "--replay-iters",
+        type=int,
+        default=10,
+        help="number of previous data files to keep in the training pools "
+        "(every game record is kept for the strength fit)",
+    )
+    _ = p.add_argument(
+        "--smooth-factor",
+        type=float,
+        nargs="+",
+        default=[0.3, 0.5, 0.7, 0.9],
+        help="grid points for the EMA smoothing factor",
+    )
+    _ = p.add_argument(
+        "--epochs",
+        type=int,
+        nargs="+",
+        default=[5],
+        help="grid points for candidate training epochs",
+    )
+    _ = p.add_argument(
+        "--lr",
+        type=float,
+        nargs="+",
+        default=[1e-4],
+        help="grid points for learning rate",
+    )
+    _ = p.add_argument(
+        "--batch",
+        type=int,
+        nargs="+",
+        default=[128],
+        help="grid points for training batch size",
+    )
+    _ = p.add_argument(
+        "--val-frac",
+        type=float,
+        nargs="+",
+        default=[0.05],
+        help="grid points for fraction of rows held out for validation",
+    )
+    _ = p.add_argument(
+        "--procs",
+        type=int,
+        default=None,
+        help="parallel game processes (number of CPUs if unset)",
+    )
+    _ = p.add_argument(
+        "--depth",
+        type=int,
+        default=4,
+        help="search depth for selfplay and the matches",
+    )
+    _ = p.add_argument(
+        "--min-ply",
+        type=int,
+        default=8,
+        help="length of random opening in games",
+    )
+    _ = p.add_argument(
+        "--train-all-games",
+        action="store_true",
+        help="train on all selfplay games, including matches between "
+        "different candidates (default: self vs self games only)",
+    )
+
+    args = p.parse_args()
+
+    # avoid counter-play start
+    if args.min_ply < 2:
+        _ = p.error("--min-ply must be >= 2")
+
+    if args.procs is None:
+        args.procs = max(1, (os.cpu_count() or 1))
+
+    data_gen = DataGenConfig(
+        size=args.size,
+        depth=args.depth,
+        min_ply=args.min_ply,
+        games=args.games,
+        procs=args.procs,
+        replay_iters=args.replay_iters,
+        root=root,
+        train_only_self_play=not args.train_all_games,
+    )
+    tgsc = TrainingGridSearchConfig(
+        smooth_factor=args.smooth_factor,
+        epochs=args.epochs,
+        batch=args.batch,
+        lr=args.lr,
+        val_frac=args.val_frac,
+    )
+
+    candidates = tgsc.make_configs(data_gen)
+
+    (root / "data").mkdir(exist_ok=True)
+
+    num_params = sum(
+        p.numel() for p in candidates[0].net.parameters() if p.requires_grad
+    )
+    print(f"Number of trainable parameters: {num_params:,}")
+
+    log = GameLog.empty()
+    strength_history: list[np.ndarray] = []
+
+    for iter in range(1, args.iters + 1):
+        run_step(
+            candidates=candidates,
+            data_gen_config=data_gen,
+            log=log,
+            strength_history=strength_history,
+            iter=iter,
+        )
+    print(f"\n== done, champion in include/weights_{args.size}.c ==")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/train/src/ctak_train/strength.py b/train/src/ctak_train/strength.py
new file mode 100644
index 0000000..2302606
--- /dev/null
+++ b/train/src/ctak_train/strength.py
@@ -0,0 +1,118 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import math
+from pathlib import Path
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+from ctak_train.candidate import Candidate
+from ctak_train.data import GameLog
+from scipy.optimize import minimize
+
+matplotlib.use("Agg")
+
+# Elo points per unit of logit strength
+ELO_PER_LOGIT = 400.0 / math.log(10.0)
+
+# Spread of the zero-mean prior on the ratings, in logit units. Head-to-head
+# results alone leave the likelihood unbounded while the data is thin (a pair
+# that has never split draws runs off to +-inf).
+PRIOR_SD_LOGIT = 1.0
+
+
+def fit_strengths(log: GameLog, n_candidates: int) -> tuple[np.ndarray, np.ndarray]:
+    """Bradley-Terry logistic model (draws scored 0.5) fit in natural-logit
+    units, mean anchored to 0, reported as Elo points around 1500. Returns
+    (strengths, +-1 sd).
+    """
+    rated = log.black != log.white
+    black = log.black[rated].astype(np.int64)
+    white = log.white[rated].astype(np.int64)
+    if n_candidates < 2 or black.size == 0:
+        return np.full(n_candidates, 1500.0), np.zeros(n_candidates)
+    score = (log.result[rated].astype(np.float64) + 1.0) / 2.0
+
+    # n_candidates - 1 free parameters; the last candidate anchors the mean to 0
+    design = np.eye(n_candidates, n_candidates - 1)
+    design[n_candidates - 1] = -1.0
+
+    def nll_grad(theta: np.ndarray) -> tuple[float, np.ndarray]:
+        rating = design @ theta
+        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
+        p = 1.0 / (1.0 + np.exp(-x))
+        sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x)))  # softplus(x)
+        sp_neg = np.maximum(-x, 0.0) + np.log1p(np.exp(-np.abs(x)))
+        # nll = score*softplus(-x) + (1-score)*softplus(x): finite for all x
+        nll = float(np.sum(score * sp_neg + (1.0 - score) * sp))
+        nll += 0.5 * float(rating @ rating) / PRIOR_SD_LOGIT**2
+
+        grad_rating = rating / PRIOR_SD_LOGIT**2
+        np.add.at(grad_rating, black, p - score)
+        np.add.at(grad_rating, white, score - p)
+        return nll, design.T @ grad_rating
+
+    def hessian(theta: np.ndarray) -> np.ndarray:
+        rating = design @ theta
+        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
+        p = 1.0 / (1.0 + np.exp(-x))
+        w = p * (1.0 - p)
+
+        diag = np.bincount(black, weights=w, minlength=n_candidates) + np.bincount(
+            white, weights=w, minlength=n_candidates
+        )
+        off = np.bincount(
+            black * n_candidates + white, weights=-w, minlength=n_candidates**2
+        ).reshape(n_candidates, n_candidates)
+        hess_rating = (
+            off
+            + np.diag(diag - np.diag(off))
+            + np.eye(n_candidates) / PRIOR_SD_LOGIT**2
+        )
+        return design.T @ hess_rating @ design
+
+    res = minimize(nll_grad, np.zeros(n_candidates - 1), jac=True, method="L-BFGS-B")
+    if not res.success:
+        print(f"strength fit did not converge: {res.message}")
+    cov = design @ np.linalg.pinv(hessian(res.x)) @ design.T
+    se = np.sqrt(np.maximum(np.diag(cov), 0.0)) * ELO_PER_LOGIT
+    return 1500.0 + (design @ res.x) * ELO_PER_LOGIT, se
+
+
+def plot_strengths(
+    history: list[np.ndarray], candidates: list[Candidate], path: Path
+) -> None:
+    fig, ax = plt.subplots(figsize=(10, 6))
+    iterations = range(1, len(history) + 1)
+    for idx, cand in enumerate(candidates):
+        ax.plot(
+            iterations,
+            [row[idx] for row in history],
+            marker=".",
+            ms=3,
+            lw=1.2,
+            label=f"cand {idx:03d} {cand}",
+        )
+    ax.axhline(1500.0, color="gray", lw=0.5)
+    ax.set_xscale("log")
+    ax.set_xlabel("Iteration")
+    ax.set_ylabel("Strength (fit on all games so far, mean 1500)")
+    ax.set_title("Candidate strength, refitted on the full history each iteration")
+    ax.legend(fontsize=8)
+    fig.tight_layout()
+    _ = fig.savefig(path, dpi=150)
+    plt.close(fig)