big refactor

* use elo
* optionally train on only diagonal
* train-inference nonlinearity alignment

2 files changed, 446 insertions, 346 deletions

diff --git a/src/nn_train.py b/src/nn_train.py
index 4de0a5a..f82a369 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -33,26 +33,23 @@
 # You should have received a copy of the GNU General Public License
 # along with ctak. If not, see <https://www.gnu.org/licenses/>.
 
-from collections.abc import Iterator
-from dataclasses import dataclass
-from itertools import product
-from math import log, ceil
-
 import argparse
 import concurrent.futures
 import json
+import math
 import os
-from pathlib import Path
+import random
 import subprocess
+from collections.abc import Iterator
+from dataclasses import dataclass
+from itertools import product
+from pathlib import Path
 
 import numpy as np
 import torch
-from tqdm import tqdm
-
-from torch.nn import Embedding, MSELoss, Linear, Module
-from torch.nn.functional import silu, tanh
 from torch import Tensor
-
+from torch.nn import Embedding, Linear, Module, MSELoss
+from tqdm import tqdm
 
 # ===================================================================
 # Net dimensions
@@ -106,6 +103,9 @@ class DataGenConfig:
     games: int
     procs: int
     replay_iters: int
+    root: Path
+    train_only_self_play: bool
+    k: float
 
 
 @dataclass(frozen=True)
@@ -116,292 +116,393 @@ class TrainingGridSearchConfig:
     lr: list[float]
     val_frac: list[float]
 
-    def make_configs(self) -> list[CandidateTrainingConfig]:
+    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
+        inp_dim, hidden = NET_DIMS[data_gen_config.size]
         flds = self.__dataclass_fields__.keys()
-        return [
-            CandidateTrainingConfig(**dict(zip(flds, vals)))
-            for vals in product(*(getattr(self, k) for k in flds))
+        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(frozen=True)
-class CandidateTrainingConfig:
+@dataclass
+class Candidate:
+    train_only_self_play: bool
     smooth_factor: float
     epochs: int
     batch: int
     lr: float
     val_frac: float
+    net: "Net"
+    elo: float = 1500.0
+    rated_games: int = 0
 
     def __str__(self) -> str:
-        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr}, V%={self.val_frac * 100.0:.1f})"
+        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
 
-# ===================================================================
-# Self-play records
-# ===================================================================
+    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()
 
-def smooth(
-    leaf_values: Tensor,
-    result: float,
-    lam: float,
-) -> Tensor:
-    # Smooth(leaf_values, result, lambda), rooted at the game outcome:
-    # G[-1] = result, G[t] = (1-lambda)V[t] + lambda G[t+1]
-    targets = torch.empty_like(leaf_values)
-    targets[-1] = result
+        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
+        for _ in pbar:
+            split.refresh_targets(self.net, self.smooth_factor)
 
-    for ply in range(len(leaf_values) - 2, -1, -1):
-        targets[ply] = (1.0 - lam) * leaf_values[ply] + lam * targets[ply + 1]
+            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()
 
-    return targets
+            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}")
 
-SPOOKY: int = 4
+        return corr, mse
 
 
-@dataclass
-class Performance:
-    win: int
-    loss: int
-    draw: int
+# ===================================================================
+# Self-play records: games (for ratings) and chains (for training)
+# ===================================================================
 
-    def rate(self) -> float:
-        return float(self.win) / float(max(1, self.win + self.loss + self.draw))
 
-    def __str__(self) -> str:
-        global SPOOKY
-        return f"+{str(self.win).zfill(SPOOKY)}-{str(self.loss).zfill(SPOOKY)}={str(self.draw).zfill(SPOOKY)}"
+@dataclass(frozen=True)
+class Games:
+    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}
+    newest: np.ndarray  # [G] game comes from the last loaded file
+
+    def head_to_head(self) -> Iterator[tuple[int, int, float]]:
+        for i in np.flatnonzero(self.newest & (self.black != self.white)):
+            black, white, result = (
+                int(self.black[i]),
+                int(self.white[i]),
+                float(self.result[i]),
+            )
+            # elo score: 1 win, 0.5 draw, 0 loss, from a's perspective
+            yield black, white, (result + 1.0) / 2.0
 
 
 @dataclass(frozen=True)
-class DataSplit:
-    train_positions: Tensor
-    train_targets: Tensor
-    train_parity_signs: Tensor
-    val_positions: Tensor
-    val_targets: Tensor
-    val_parity_signs: Tensor
+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)
 
-    def batches(self, batch_size: int) -> Iterator[tuple[Tensor, Tensor, Tensor]]:
-        order = torch.randperm(
-            len(self.train_positions), device=self.train_positions.device
+        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),
         )
-        for start in range(0, len(order), batch_size):
-            idx = order[start : start + batch_size]
-            yield (
-                self.train_positions[idx],
-                self.train_targets[idx],
-                self.train_parity_signs[idx],
-            )
 
+    @property
+    def n_rows(self) -> int:
+        return self.head.shape[0]
 
-@dataclass(frozen=True)
-class CandidateRecords:
-    positions: Tensor  # [n, inp_dim] net input of the root state s_t
-    leaf_positions: Tensor  # [n, inp_dim] net input of the PV leaf pv_t
-    game_ids: Tensor  # [n]
-    results: Tensor  # [n] game result R, black's perspective
-    cand_results: Tensor  # [n] R from the recorded candidate's perspective
-    parity_signs: Tensor  # [n] root parity
-    leaf_parity_signs: Tensor  # [n] leaf parity
-    leaf_terminals: Tensor  # [n] leaf is an exact terminal
-    leaf_terminal_values: Tensor  # [n] exact terminal value, black's perspective
+    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],
+        )
 
-    def performance(self) -> Performance:
-        result = Performance(0, 0, 0)
-        for game_id in torch.unique(self.game_ids):
-            in_game = self.game_ids == game_id
-            game_result = self.cand_results[in_game][0].item()
-            match game_result:
-                case 1:
-                    result.win += 1
-                case -1:
-                    result.loss += 1
-                case _:
-                    result.draw += 1
-        return result
 
-    def make_training_data(
-        self, model: Module, training: CandidateTrainingConfig
-    ) -> DataSplit:
-        model.eval()
-        with torch.no_grad():
-            # Net(pv_t), black's perspective; exact terminals override the net
-            leaf_values = model(self.leaf_positions) * self.leaf_parity_signs
-            leaf_values = torch.where(
-                self.leaf_terminals, self.leaf_terminal_values, leaf_values
+# ===================================================================
+# Reading game records
+# ===================================================================
+
+_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)
 
-        targets = torch.empty_like(leaf_values)
 
-        for game_id in torch.unique(self.game_ids):
-            in_game = self.game_ids == game_id
-            game_result = self.results[in_game][0].item()
-            # self-pair games (i,i) record the candidate on every ply,
-            # so smooth one chain per colour, not across players
-            for sign in (1.0, -1.0):
-                idxs = torch.nonzero(
-                    in_game & (self.parity_signs == sign), as_tuple=True
-                )[0]
-                if len(idxs) == 0:
+def load_window(
+    files: list[Path], size: int, n_candidates: int
+) -> tuple[list[Pool], Games]:
+    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, newest = [], [], [], []
+
+    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
-                targets[idxs] = smooth(
-                    leaf_values[idxs], game_result, training.smooth_factor
-                )
+                game = json.loads(line)
+                game_result = float(game["result"])
+                b, w = _game_colours(game)
+                black.append(b)
+                white.append(w)
+                result.append(game_result)
+                newest.append(is_newest)
+
+                by_cand: dict[int, list[dict]] = {}
+                for r in game["rows"]:
+                    by_cand.setdefault(int(r["cand"]), []).append(r)
 
-        n_val = max(1, int(len(leaf_values) * training.val_frac))
-        perm = torch.randperm(len(leaf_values), device=leaf_values.device)
-        val_idx, train_idx = perm[:n_val], perm[n_val:]
+                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(b == w)
 
-        return DataSplit(
-            train_positions=self.positions[train_idx],
-            train_targets=targets[train_idx],
-            train_parity_signs=self.parity_signs[train_idx],
-            val_positions=self.positions[val_idx],
-            val_targets=targets[val_idx],
-            val_parity_signs=self.parity_signs[val_idx],
+    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 = Games(
+        black=np.array(black, dtype=np.int16),
+        white=np.array(white, dtype=np.int16),
+        result=np.array(result, dtype=np.float32),
+        newest=np.array(newest, dtype=np.bool_),
+    )
+    return [pool_of(c) for c in range(n_candidates)], games
+
 
-@dataclass(frozen=True)
-class PlayRecords:
-    positions: list[Tensor]  # [candidate, n_i, inp_dim]
-    leaf_positions: list[Tensor]  # [candidate, n_i, inp_dim]
-    game_ids: list[Tensor]  # [candidate, n_i]
-    results: list[Tensor]  # [candidate, n_i]
-    cand_results: list[Tensor]  # [candidate, n_i]
-    parity_signs: list[Tensor]  # [candidate, n_i]
-    leaf_parity_signs: list[Tensor]  # [candidate, n_i]
-    leaf_terminals: list[Tensor]  # [candidate, n_i]
-    leaf_terminal_values: list[Tensor]  # [candidate, n_i]
+# ===================================================================
+# Training data: fixed rows, refreshed targets
+# ===================================================================
 
-    def index(self, index: int) -> CandidateRecords:
-        return CandidateRecords(
-            positions=self.positions[index],
-            leaf_positions=self.leaf_positions[index],
-            game_ids=self.game_ids[index],
-            results=self.results[index],
-            cand_results=self.cand_results[index],
-            parity_signs=self.parity_signs[index],
-            leaf_parity_signs=self.leaf_parity_signs[index],
-            leaf_terminals=self.leaf_terminals[index],
-            leaf_terminal_values=self.leaf_terminal_values[index],
-        )
 
-    @classmethod
-    def load_merge(cls, files: list[Path], size: int) -> PlayRecords:
-        # one JSON record per line:
-        # {"result": R, "rows": [{"cand", "root": {...}, "leaf": {...}}]}
-        inp_dim = NET_DIMS[size][0]
-        n_cells = inp_dim - 2
+@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
 
-        games: list[dict] = []
-        for file in files:
-            with open(file) as f:
-                for line in f:
-                    if line.strip():
-                        games.append(json.loads(line))
+    @property
+    def n(self) -> int:
+        return self.pool.n_rows
 
-        def board(row_part: dict) -> list[float]:
-            return [*row_part["head"], *row_part["cells"]]
+    @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)
 
-        n_cands = max(r["cand"] for g in games for r in g["rows"]) + 1
-        fields = (
-            "positions",
-            "leaf_positions",
-            "game_ids",
-            "results",
-            "cand_results",
-            "parity_signs",
-            "leaf_parity_signs",
-            "leaf_terminals",
-            "leaf_terminal_values",
+        # 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)
         )
-        per_cand = {c: {k: [] for k in fields} for c in range(n_cands)}
-        for game_id, game in enumerate(games):
-            result = float(game["result"])
-            for cand in range(n_cands):
-                rows = [r for r in game["rows"] if r["cand"] == cand]
-                if not rows:
-                    continue
-                for part in ("root", "leaf"):
-                    for r in rows:
-                        if len(r[part]["cells"]) != n_cells:
-                            raise ValueError(
-                                f"{files[-1]}: {len(r[part]['cells'])} cells, "
-                                f"expected {n_cells} for size {size}"
-                            )
-                parity = np.array([r["root"]["parity"] for r in rows], dtype=np.float32)
-                per_cand[cand]["positions"].append(
-                    torch.from_numpy(
-                        np.array([board(r["root"]) for r in rows], dtype=np.float32)
-                    )
-                )
-                per_cand[cand]["leaf_positions"].append(
-                    torch.from_numpy(
-                        np.array([board(r["leaf"]) for r in rows], dtype=np.float32)
-                    )
-                )
-                per_cand[cand]["game_ids"].append(
-                    torch.full((len(rows),), game_id, dtype=torch.int64)
-                )
-                per_cand[cand]["results"].append(
-                    torch.full((len(rows),), result, dtype=torch.float32)
-                )
-                per_cand[cand]["cand_results"].append(torch.from_numpy(parity * result))
-                per_cand[cand]["parity_signs"].append(torch.from_numpy(parity))
-                per_cand[cand]["leaf_parity_signs"].append(
-                    torch.from_numpy(
-                        np.array([r["leaf"]["parity"] for r in rows], dtype=np.float32)
-                    )
-                )
-                per_cand[cand]["leaf_terminals"].append(
-                    torch.from_numpy(
-                        np.array(
-                            [r["leaf"]["terminal"] == 1 for r in rows],
-                            dtype=np.bool_,
-                        )
-                    )
-                )
-                per_cand[cand]["leaf_terminal_values"].append(
-                    torch.from_numpy(
-                        np.array(
-                            [r["leaf"]["terminal_value"] for r in rows],
-                            dtype=np.float32,
-                        )
-                    )
-                )
-
-        def merged(cand: int, k: str, empty: Tensor) -> Tensor:
-            blocks = per_cand[cand][k]
-            return torch.concat(blocks, dim=0) if blocks else empty
 
+        n_val = max(1, int(n * val_frac))
+        perm = torch.randperm(n)
         return cls(
-            **{
-                k: [
-                    merged(
-                        cand,
-                        k,
-                        torch.empty(
-                            (0, inp_dim) if k.endswith("positions") else (0,),
-                            dtype=torch.float32
-                            if k not in ("game_ids", "leaf_terminals")
-                            else (torch.int64 if k == "game_ids" else torch.bool),
-                        ),
-                    )
-                    for cand in range(n_cands)
-                ]
-                for k in fields
-            }
+            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],
+            )
+
 
 # ===================================================================
 # Network
 # ===================================================================
 
 
+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__()
@@ -410,56 +511,24 @@ class Net(Module):
         self.hidden = Linear(inp_dim, hidden, bias=True)
         self.output = Linear(hidden, 1, bias=True)
 
-    def forward(self, x: Tensor) -> Tensor:
-        # x: [n, 2 + n_cells]; the first 2 columns are the fixed head
-        # inputs, the rest are raw cell states (0..6)
-        n = x.shape[0]
-        inp = torch.cat([x[:, :2], self.embedding(x[:, 2:].long()).view(n, -1)], dim=1)
-        hidden = silu(self.hidden(inp))
-        return tanh(self.output(hidden)).squeeze(-1)
-
-    def fit(
-        self, *, records: CandidateRecords, config: CandidateTrainingConfig
-    ) -> tuple[float, float]:
-        # training widgets
-        optimizer = torch.optim.AdamW(self.parameters(), lr=config.lr)
-        loss_fn = MSELoss()
-
-        pbar = tqdm(range(config.epochs), desc="train", unit="ep")
-        for _ in pbar:
-            split = records.make_training_data(self, config)
-
-            self.train()
-            for positions, targets, parity_signs in split.batches(config.batch):
-                optimizer.zero_grad()
-                predictions = self(positions) * parity_signs
-                loss = loss_fn(predictions, targets)
-                loss.backward()
-                optimizer.step()
-
-            with torch.no_grad():
-                self.eval()
-                predictions = self(split.val_positions) * split.val_parity_signs
-                mse = loss_fn(predictions, split.val_targets).item()
-                corr = torch.corrcoef(torch.stack([predictions, split.val_targets]))[
-                    0, 1
-                ].item()
-
-            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
-
-        return corr, mse
+    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,
-        id: int | None,
+        ident: int | None,
         path: Path,
         note: str,
     ) -> None:
         upper = "FIVE" if size == 5 else "SIX"
-        prefix = (f"cand_{id}_" if id is not None else "") + upper
-        guard = f"WEIGHTS_{upper}" + (f"_{id}" if id is not None else "")
+        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()
@@ -478,7 +547,7 @@ class Net(Module):
                 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",
-                f"#endif",
+                "#endif",
             ]
             for line in lines:
                 _ = f.write(line)
@@ -537,30 +606,26 @@ def cand_copy_for(size: int, id: int, root: Path) -> None:
 # ===================================================================
 
 
-def setup(num_candidates: int, data_gen: DataGenConfig, root: Path) -> list[Net]:
-    inp_dim, hidden = NET_DIMS[data_gen.size]
-    candidates = [Net(inp_dim, hidden) for _ in range(num_candidates)]
-    for id, cand in enumerate(candidates):
-        cand.to_c(
-            size=data_gen.size,
-            id=id,
-            path=root / "include" / f"cand_{id}_weights_{data_gen.size}.c",
-            note=f"{id=}, v000, random init",
-        )
-        cand_copy_for(6 if data_gen.size == 5 else 5, id, root)
-    return candidates
+def update_elo(a: Candidate, b: Candidate, score: float, k: float) -> None:
+    # score is a's result as an elo score: 1.0 win, 0.5 draw, 0.0 loss
+    ea = 1.0 / (1.0 + 10.0 ** ((b.elo - a.elo) / 400.0))
+    delta = k * (score - ea)
+    a.elo += delta
+    b.elo -= delta
 
 
 def run_step(
     *,
-    candidates: list[Net],
-    data_gen: DataGenConfig,
+    candidates: list[Candidate],
+    data_gen_config: DataGenConfig,
     iter: int,
-    root: Path,
-    training_configs: list[CandidateTrainingConfig],
 ):
     def data_path(i: int) -> Path:
-        return root / "data" / f"selfplay-{data_gen.size}-v{i:03d}.json"
+        return (
+            data_gen_config.root
+            / "data"
+            / f"selfplay-{data_gen_config.size}-v{i:03d}.json"
+        )
 
     print(f"\n== iteration {iter} ==")
 
@@ -572,30 +637,34 @@ def run_step(
     # 1. rebuild
     for i, j in cand_pairs:
         _ = subprocess.run(
-            ["make", f"selfplay-{i}-{j}"], cwd=root, check=True, capture_output=True
+            ["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.games=}, "
-        f"{data_gen.size=}, {data_gen.depth=}, {data_gen.min_ply=}"
+        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.size),
-                str(data_gen.depth),
-                str(data_gen.min_ply),
+                str(data_gen_config.size),
+                str(data_gen_config.depth),
+                str(data_gen_config.min_ply),
             ]
             for (i, j) in cand_pairs
         ],
-        data_gen.games,
-        data_gen.procs,
+        data_gen_config.games,
+        data_gen_config.procs,
         "selfplay",
-        root,
+        data_gen_config.root,
     ):
         out_games.extend(out.splitlines())
     print(f"writing {len(out_games):,} games to {data}")
@@ -604,43 +673,57 @@ def run_step(
         if out_games:
             _ = f.write("\n")
 
-    # 3. compute performance
+    # 3. rate every head-to-head game played in this iteration
     data_files = [
-        data_path(i) for i in range(max(iter - data_gen.replay_iters, 1), iter + 1)
+        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=" ")
-    records = PlayRecords.load_merge(data_files, data_gen.size)
-    print("candidates have", [f"{r.shape[0]:,}" for r in records.results], "rows.")
-    results = [
-        (i, cand, records.index(i).performance()) for i, cand in enumerate(candidates)
-    ]
+    pools, games = load_window(data_files, data_gen_config.size, num_candidates)
+    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")
 
-    print("Tournament results:")
-    for idx, _, perf in sorted(results, key=lambda t: t[2].rate(), reverse=True):
-        print(
-            f"\tcandidate {idx:03d} {perf} {perf.rate() * 100.0:.2f}% {training_configs[idx]}"
-        )
+    fresh = list(games.head_to_head())
+    random.shuffle(fresh)
+    for a, b, score in fresh:
+        update_elo(candidates[a], candidates[b], score, data_gen_config.k)
+        candidates[a].rated_games += 1
+        candidates[b].rated_games += 1
 
-    # 4. write champ
-    idx, champ, perf = max(results, key=lambda t: t[2].rate())
-    path = root / "include" / f"weights_{data_gen.size}.c"
-    champ.to_c(
-        size=data_gen.size,
-        id=None,
+    total_rated_games = sum(c.rated_games for c in candidates)
+    print(
+        f"Tournament standings ({total_rated_games:,} rated games, {total_rated_games / len(candidates):,} games/candidate):"
+    )
+    for idx, cand in sorted(
+        enumerate(candidates), key=lambda t: t[1].elo, reverse=True
+    ):
+        print(f"\tcandidate {idx:03d}: Elo {cand.elo:4.1f} Config {cand}")
+
+    # 4. write champ, unconditionally
+    idx, champ = max(enumerate(candidates), key=lambda t: t[1].elo)
+    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
+    champ.net.to_c(
+        size=data_gen_config.size,
+        ident=None,
         path=path,
-        note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
+        note=f"{iter=}, elo={champ.elo:.1f}, rated_games={champ.rated_games:,}",
     )
-    _ = subprocess.run(["make", "clean"], cwd=root, check=True, capture_output=True)
-    _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True)
-    print(f"Wrote candidate {idx:03d} to {str(path)}")
+    _ = 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}")
 
     # 5. fine-tune on the hybrid targets -> candidate
     for idx, cand in enumerate(candidates):
-        corr, mse = cand.fit(records=records.index(idx), config=training_configs[idx])
-        cand.to_c(
-            size=data_gen.size,
-            id=idx,
-            path=root / "include" / f"cand_{idx}_weights_{data_gen.size}.c",
+        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}",
         )
 
@@ -734,9 +817,25 @@ def main() -> None:
         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)",
+    )
+    _ = p.add_argument(
+        "--k-factor",
+        type=float,
+        default=32.0,
+        help="elo k-factor",
+    )
 
     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))
 
@@ -747,6 +846,9 @@ def main() -> None:
         games=args.games,
         procs=args.procs,
         replay_iters=args.replay_iters,
+        root=root,
+        train_only_self_play=not args.train_all_games,
+        k=args.k_factor,
     )
     tgsc = TrainingGridSearchConfig(
         smooth_factor=args.smooth_factor,
@@ -756,23 +858,20 @@ def main() -> None:
         val_frac=args.val_frac,
     )
 
-    training_configs = tgsc.make_configs()
+    candidates = tgsc.make_configs(data_gen)
 
     (root / "data").mkdir(exist_ok=True)
-    candidates = setup(
-        num_candidates=len(training_configs), data_gen=data_gen, root=root
-    )
 
-    num_params = sum(p.numel() for p in candidates[0].parameters() if p.requires_grad)
+    num_params = sum(
+        p.numel() for p in candidates[0].net.parameters() if p.requires_grad
+    )
     print(f"Number of trainable parameters: {num_params:,}")
 
     for iter in range(1, args.iters + 1):
         run_step(
             candidates=candidates,
-            data_gen=data_gen,
+            data_gen_config=data_gen,
             iter=iter,
-            root=root,
-            training_configs=training_configs,
         )
     print(f"\n== done, champion in include/weights_{args.size}.c ==")
 
diff --git a/src/selfplay.c b/src/selfplay.c
index b2a5821..37f94fa 100644
--- a/src/selfplay.c
+++ b/src/selfplay.c
@@ -169,7 +169,8 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
   // final outcome of the played game, from black's perspective
   const int8_t R = win_value(state->won);
 
-  fprintf(stdout, "{\"result\":%d,\"rows\":[", R);
+  fprintf(stdout, "{\"result\":%d,\"cand_A\":%u,\"cand_B\":%u,\"rows\":[", R,
+          CAND_A, CAND_B);
   for (uint32_t k = 0; k < n_rows; k++) {
     fprintf(stdout, "%s{\"cand\":%u,\"root\":{\"parity\":%d,", k ? "," : "",
             row_cand_id[k], row_mover_sign[k]);
@@ -188,8 +189,8 @@ const char *usage =
   "selfplay, one game of training data between candidates.\n\n\
 Usage: selfplay <size> <depth> <min_ply>\n\n\
 Plays a single game of <size>x<size>, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one JSON record per game, from ply <min_ply> onwards:\n\n\
-\t{\"result\": R, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\
-R is the outcome of the game from black's perspective. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\
+\t{\"result\": R, \"cand_A\": i, \"cand_B\": j, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\
+R is the outcome of the game from black's perspective; cand_A and cand_B are the ids of the two candidates that played the game. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\
 Copyright (C) 2026, tslil clingman\n\n\
 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.";