working towards hybrid rootstrap

2 files changed, 300 insertions, 194 deletions

diff --git a/src/nn_train.py b/src/nn_train.py
index 5332aa8..a97c964 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -3,7 +3,7 @@
 # /// script
 # requires-python = ">=3.14"
 # dependencies = [
-#   "pandas>=3.0.5",
+#   "numpy>=2.5.3",
 #   "torch>=2.14.0",
 #   "tqdm>=4.70.0",
 # ]
@@ -39,13 +39,12 @@ from math import log, ceil
 
 import argparse
 import concurrent.futures
-import copy
+import json
 import os
 from pathlib import Path
 import subprocess
 
 import numpy as np
-import pandas as pd
 import torch
 from tqdm import tqdm
 
@@ -94,28 +93,51 @@ def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
 
 
 # ===================================================================
+# Config
+# ===================================================================
+
+
+@dataclass(frozen=True)
+class DataGenConfig:
+    size: int
+    depth: int
+    min_ply: int
+    games: int
+    procs: int
+    replay_iters: int
+
+
+@dataclass(frozen=True)
+class TrainingConfig:
+    smooth_factor: float
+    epochs: int
+    batch: int
+    lr: float
+    val_frac: float
+
+
+# ===================================================================
 # Self-play records
 # ===================================================================
 
 
-def td_lambda_targets(
+def smooth(
     leaf_values: Tensor,
     result: float,
-    td_lambda: 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
 
-    # G[t] = (1-lambda)V[t+1] + lambda G[t+1]
     for ply in range(len(leaf_values) - 2, -1, -1):
-        targets[ply] = (1.0 - td_lambda) * leaf_values[ply + 1] + td_lambda * targets[
-            ply + 1
-        ]
+        targets[ply] = (1.0 - lam) * leaf_values[ply] + lam * targets[ply + 1]
 
     return targets
 
 
-SPOOKY: int = 3
+SPOOKY: int = 4
 
 
 @dataclass
@@ -133,14 +155,38 @@ class Performance:
 
 
 @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
+
+    def batches(self, batch_size: int) -> Iterator[tuple[Tensor, Tensor, Tensor]]:
+        order = torch.randperm(
+            len(self.train_positions), device=self.train_positions.device
+        )
+        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],
+            )
+
+
+@dataclass(frozen=True)
 class CandidateRecords:
-    positions: list[Tensor]  # [n, inp_dim]
-    game_ids: list[Tensor]  # [n]
-    results: list[Tensor]  # [n]
-    cand_results: list[Tensor]  # [n]
-    terminals: list[Tensor]  # [n]
-    terminal_values: list[Tensor]  # [n]
-    parity_signs: list[Tensor]  # [n]
+    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 performance(self) -> Performance:
         result = Performance(0, 0, 0)
@@ -156,107 +202,176 @@ class CandidateRecords:
                     result.draw += 1
         return result
 
-    def make_training_data(
-        self, model: Module, td_lambda: float
-    ) -> tuple[Tensor, Tensor, Tensor]:
+    def make_training_data(self, model: Module, training: TrainingConfig) -> DataSplit:
         model.eval()
         with torch.no_grad():
-            values = model(self.positions) * self.parity_signs
+            # 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
+            )
 
-        targets = torch.empty_like(values)
+        targets = torch.empty_like(leaf_values)
 
         for game_id in torch.unique(self.game_ids):
             in_game = self.game_ids == game_id
-            game_terminals = self.terminals[in_game]
-            game_terminal_values = self.terminal_values[in_game]
-            game_values = values[in_game].clone()
-            game_values[game_terminals] = game_terminal_values[game_terminals]
             game_result = self.results[in_game][0].item()
-            game_targets = td_lambda_targets(game_values, game_result, td_lambda)
-            # Terminal positions are also direct supervised examples. Do this
-            # after constructing the TD sequence so it does not reset the EMA.
-            game_targets[game_terminals] = game_terminal_values[game_terminals]
-            targets[in_game] = game_targets
+            # 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:
+                    continue
+                targets[idxs] = smooth(
+                    leaf_values[idxs], game_result, training.smooth_factor
+                )
+
+        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:]
 
-        return self.positions, targets, self.parity_signs
+        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],
+        )
 
 
 @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]
-    terminals: list[Tensor]  # [candidate, n_i]
-    terminal_values: 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]
 
     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],
-            terminals=self.terminals[index],
-            terminal_values=self.terminal_values[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_selfplay(cls, fn, size):
-        df = pd.read_csv(fn, header=None)
+    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
+
+        games: list[dict] = []
+        for file in files:
+            with open(file) as f:
+                for line in f:
+                    if line.strip():
+                        games.append(json.loads(line))
 
-        fields = {
-            "parity_signs": df.iloc[:, 0].to_numpy(dtype=np.float32),
-            "positions": df.iloc[:, 1 : inp_dim + 1].to_numpy(dtype=np.float32),
-            "results": df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32),
-            "terminals": df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_),
-            "terminal_values": df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32),
-            "game_ids": df.iloc[:, inp_dim + 4].to_numpy(dtype=np.int64),
-            "cand_results": df.iloc[:, inp_dim + 6].to_numpy(dtype=np.float32),
-        }
-        cand_ids = df.iloc[:, inp_dim + 5].to_numpy(dtype=np.int64)
+        def board(row_part: dict) -> list[float]:
+            return [*row_part["head"], *row_part["cells"]]
 
-        n_cands = int(cand_ids.max()) + 1
-        gathered = {k: [] for k in fields}
-        for cid in range(n_cands):
-            mask = cand_ids == cid
-            for k, arr in fields.items():
-                gathered[k].append(torch.from_numpy(arr[mask]))
+        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",
+        )
+        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,
+                        )
+                    )
+                )
 
-        return cls(**gathered)
+        def merged(cand: int, k: str, empty: Tensor) -> Tensor:
+            blocks = per_cand[cand][k]
+            return torch.concat(blocks, dim=0) if blocks else empty
 
-    @classmethod
-    def load_merge(cls, files: list[Path], size: int) -> PlayRecords:
-        sprs = [PlayRecords.load_selfplay(file, size) for file in files]
-        # for each field, gather [file0[i], file1[i], ...] for every index i
-        # and concat along dim 0 -> merged list indexed by cand id
         return cls(
-            positions=[
-                torch.concat(vals, dim=0) for vals in zip(*(s.positions for s in sprs))
-            ],
-            game_ids=[
-                torch.concat(vals, dim=0) for vals in zip(*(s.game_ids for s in sprs))
-            ],
-            results=[
-                torch.concat(vals, dim=0) for vals in zip(*(s.results for s in sprs))
-            ],
-            cand_results=[
-                torch.concat(vals, dim=0)
-                for vals in zip(*(s.cand_results for s in sprs))
-            ],
-            terminals=[
-                torch.concat(vals, dim=0) for vals in zip(*(s.terminals for s in sprs))
-            ],
-            terminal_values=[
-                torch.concat(vals, dim=0)
-                for vals in zip(*(s.terminal_values for s in sprs))
-            ],
-            parity_signs=[
-                torch.concat(vals, dim=0)
-                for vals in zip(*(s.parity_signs for s in sprs))
-            ],
+            **{
+                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
+            }
         )
 
 
@@ -282,55 +397,29 @@ class Net(Module):
         return tanh(self.output(hidden)).squeeze(-1)
 
     def fit(
-        self,
-        *,
-        records: CandidateRecords,
-        td_lambda: float,
-        val_frac: float,
-        epochs: int,
-        batch: int,
-        lr: float,
+        self, *, records: CandidateRecords, training: TrainingConfig
     ) -> tuple[float, float]:
         # training widgets
-        optimizer = torch.optim.AdamW(self.parameters(), lr=lr)
+        optimizer = torch.optim.AdamW(self.parameters(), lr=training.lr)
         loss_fn = MSELoss()
 
-        pbar = tqdm(range(epochs), desc="train", unit="ep")
+        pbar = tqdm(range(training.epochs), desc="train", unit="ep")
         for _ in pbar:
-            # make data
-            positions, targets, parity_signs = records.make_training_data(
-                self, td_lambda
-            )
-            n = len(positions)
-            perm = torch.randperm(n, device=positions.device)
-            n_val = max(1, int(n * val_frac))
-            val_idx = perm[:n_val]
-            train_idx = perm[n_val:]
-
-            val_positions = positions[val_idx]
-            val_targets = targets[val_idx]
-            val_parity_signs = parity_signs[val_idx]
-            train_positions = positions[train_idx]
-            train_targets = targets[train_idx]
-            train_parity_signs = parity_signs[train_idx]
+            split = records.make_training_data(self, training)
 
             self.train()
-            order = torch.randperm(len(train_positions), device=positions.device)
-            for start in range(0, len(train_positions), batch):
-                batch_idx = order[start : start + batch]
+            for positions, targets, parity_signs in split.batches(training.batch):
                 optimizer.zero_grad()
-                predictions = (
-                    self(train_positions[batch_idx]) * train_parity_signs[batch_idx]
-                )
-                loss = loss_fn(predictions, train_targets[batch_idx])
+                predictions = self(positions) * parity_signs
+                loss = loss_fn(predictions, targets)
                 loss.backward()
                 optimizer.step()
 
             with torch.no_grad():
                 self.eval()
-                predictions = self(val_positions) * val_parity_signs
-                mse = loss_fn(predictions, val_targets).item()
-                corr = torch.corrcoef(torch.stack([predictions, val_targets]))[
+                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()
 
@@ -426,39 +515,30 @@ def cand_copy_for(size: int, id: int, root: Path) -> None:
 # ===================================================================
 
 
-def setup(num_candidates: int, size: int, root: Path) -> list[Net]:
-    inp_dim, hidden = NET_DIMS[size]
+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=size,
+            size=data_gen.size,
             id=id,
-            path=root / "include" / f"cand_{id}_weights_{size}.c",
+            path=root / "include" / f"cand_{id}_weights_{data_gen.size}.c",
             note=f"{id=}, v000, random init",
         )
-        cand_copy_for(6 if size == 5 else 5, id, root)
+        cand_copy_for(6 if data_gen.size == 5 else 5, id, root)
     return candidates
 
 
 def run_step(
     *,
-    batch: int,
     candidates: list[Net],
-    depth: int,
-    epochs: int,
-    games: int,
+    data_gen: DataGenConfig,
     iter: int,
-    lr: float,
-    min_ply: int,
-    procs: int,
-    replay_iters: int,
     root: Path,
-    size: int,
-    td_lambda: float,
-    val_frac: float,
+    training: TrainingConfig,
 ):
-    def csv_path(i: int) -> Path:
-        return root / "data" / f"selfplay-{size}-v{i:03d}.csv"
+    def data_path(i: int) -> Path:
+        return root / "data" / f"selfplay-{data_gen.size}-v{i:03d}.json"
 
     print(f"\n== iteration {iter} ==")
 
@@ -474,30 +554,40 @@ def run_step(
         )
 
     # 2. generate self-play data
-    csv = csv_path(iter)
-    print(f"selfplay: {num_candidates=}, {games=}, {size=}, {depth=}, {min_ply=}")
-    out_rows: list[str] = []
+    data = data_path(iter)
+    print(
+        f"selfplay: {num_candidates=}, {data_gen.games=}, "
+        f"{data_gen.size=}, {data_gen.depth=}, {data_gen.min_ply=}"
+    )
+    out_games: list[str] = []
     for out in fan_out(
         [
-            [f"./selfplay-{i}-{j}", str(size), str(depth), str(min_ply)]
+            [
+                f"./selfplay-{i}-{j}",
+                str(data_gen.size),
+                str(data_gen.depth),
+                str(data_gen.min_ply),
+            ]
             for (i, j) in cand_pairs
         ],
-        games,
-        procs,
+        data_gen.games,
+        data_gen.procs,
         "selfplay",
         root,
     ):
-        out_rows.extend(out.splitlines())
-    print(f"writing {len(out_rows):,} rows to {csv}")
-    with open(csv, "w") as f:
-        _ = f.write("\n".join(out_rows))
-        if out_rows:
+        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. compute performance
-    csvs = [csv_path(i) for i in range(max(iter - replay_iters, 1), iter + 1)]
-    print("Loading data from", [str(c) for c in csvs], "...", end=" ")
-    records = PlayRecords.load_merge(csvs, size)
+    data_files = [
+        data_path(i) for i in range(max(iter - data_gen.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)
@@ -509,9 +599,9 @@ def run_step(
 
     # 4. write champ
     idx, champ, perf = max(results, key=lambda t: t[2].rate())
-    path = root / "include" / f"weights_{size}.c"
+    path = root / "include" / f"weights_{data_gen.size}.c"
     champ.to_c(
-        size=size,
+        size=data_gen.size,
         id=None,
         path=path,
         note=f"{iter=}, perf={str(perf)}, rate={perf.rate() * 100.0:.3f}",
@@ -520,20 +610,13 @@ def run_step(
     _ = subprocess.run(["make", "native"], cwd=root, check=True, capture_output=True)
     print(f"Wrote candidate {idx:03d} to {str(path)}")
 
-    # 5. fine-tune on TD-leaf(lambda) targets -> candidate
+    # 5. fine-tune on the hybrid targets -> candidate
     for idx, cand in enumerate(candidates):
-        corr, mse = cand.fit(
-            records=records.index(idx),
-            td_lambda=td_lambda,
-            val_frac=val_frac,
-            epochs=epochs,
-            batch=batch,
-            lr=lr,
-        )
+        corr, mse = cand.fit(records=records.index(idx), training=training)
         cand.to_c(
-            size=size,
+            size=data_gen.size,
             id=idx,
-            path=root / "include" / f"cand_{idx}_weights_{size}.c",
+            path=root / "include" / f"cand_{idx}_weights_{data_gen.size}.c",
             note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
         )
 
@@ -546,7 +629,7 @@ def run_step(
 def main() -> None:
     root = Path(__file__).resolve().parent.parent
     p = argparse.ArgumentParser(
-        description="the TD(lambda) self-play bootstrap loop",
+        description="Train the neural network evaluators of ctak",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     _ = p.add_argument(
@@ -577,19 +660,19 @@ def main() -> None:
     _ = p.add_argument(
         "--replay-iters",
         type=int,
-        default=0,
+        default=10,
         help="number of previous data sources to retain",
     )
     _ = p.add_argument(
-        "--td-lambda",
+        "--smooth-factor",
         type=float,
         default=0.7,
-        help="TD(lambda) lambda",
+        help="The EMA smoothing factor",
     )
     _ = p.add_argument(
         "--epochs",
         type=int,
-        default=20,
+        default=5,
         help="candidate training epochs",
     )
     _ = p.add_argument(
@@ -626,7 +709,7 @@ def main() -> None:
         "--min-ply",
         type=int,
         default=8,
-        help="record positions at this ply or later",
+        help="length of random opening in games",
     )
 
     args = p.parse_args()
@@ -642,28 +725,35 @@ def main() -> None:
         )
     )
 
+    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,
+    )
+    training = TrainingConfig(
+        smooth_factor=args.smooth_factor,
+        epochs=args.epochs,
+        batch=args.batch,
+        lr=args.lr,
+        val_frac=args.val_frac,
+    )
+
     (root / "data").mkdir(exist_ok=True)
-    candidates = setup(num_candidates=args.num_candidates, size=args.size, root=root)
+    candidates = setup(num_candidates=args.num_candidates, data_gen=data_gen, root=root)
 
     num_params = sum(p.numel() for p in candidates[0].parameters() if p.requires_grad)
     print(f"Number of trainable parameters: {num_params:,}")
 
     for iter in range(1, args.iters + 1):
         run_step(
-            batch=args.batch,
             candidates=candidates,
-            depth=args.depth,
-            epochs=args.epochs,
-            games=args.games,
+            data_gen=data_gen,
             iter=iter,
-            lr=args.lr,
-            min_ply=args.min_ply,
-            procs=args.procs,
-            replay_iters=args.replay_iters,
             root=root,
-            size=args.size,
-            td_lambda=args.td_lambda,
-            val_frac=args.val_frac,
+            training=training,
         )
     print(f"\n== done, champion in include/weights_{args.size}.c ==")
 
diff --git a/src/selfplay.c b/src/selfplay.c
index 1d6bd84..b2a5821 100644
--- a/src/selfplay.c
+++ b/src/selfplay.c
@@ -91,10 +91,20 @@ static int8_t win_value(const enum WIN_TYPE w) {
   }
 }
 
+static void emit_board(const float *head, const uint8_t *cells,
+                       const uint8_t n_cells) {
+  fprintf(stdout, "\"head\":[%.8f,%.8f],\"cells\":[", head[0], head[1]);
+  for (uint8_t c = 0; c < n_cells; c++)
+    fprintf(stdout, "%s%u", c ? "," : "", cells[c]);
+  fprintf(stdout, "]");
+}
+
 static int play_game(const uint8_t min_ply, const nn_function fn_1,
                      const nn_function fn_2, const uint32_t cand_id_1,
                      const uint32_t cand_id_2,
-                     const uint32_t game_id, char record[][16], uint8_t replay) {
+                     char record[][16], uint8_t replay) {
+  float root_head[MAX_PLY][2];
+  uint8_t root_cells[MAX_PLY][MAX_CELLS];
   float row_head[MAX_PLY][2];
   uint8_t row_cells[MAX_PLY][MAX_CELLS];
   uint8_t row_term[MAX_PLY];
@@ -133,6 +143,10 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
       row_parity_sign[n_rows] = leaf->parity_sign;
       row_cand_id[n_rows] = (state->current_colour&1) ? cand_id_1 : cand_id_2;
       row_mover_sign[n_rows] = (state->current_colour&1) ? 1 : -1;
+
+      // Net input of the root state s_t, before the move is played
+      nn_prepare_head(root_head[n_rows], state);
+      nn_cell_idxs(root_cells[n_rows], state);
       n_rows++;
     }
 
@@ -155,13 +169,17 @@ 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);
   for (uint32_t k = 0; k < n_rows; k++) {
-    fprintf(stdout, "%.1f,%.8f,%.8f,", row_parity_sign[k], row_head[k][0], row_head[k][1]);
-    for (uint8_t c = 0; c < n_cells; c++)
-      fprintf(stdout, "%u,", row_cells[k][c]);
-    fprintf(stdout, "%i,%u,%.0f,%u,%u,%i\n", R, row_term[k], row_terminal_value[k], game_id, row_cand_id[k],
-            row_mover_sign[k] * R);
+    fprintf(stdout, "%s{\"cand\":%u,\"root\":{\"parity\":%d,", k ? "," : "",
+            row_cand_id[k], row_mover_sign[k]);
+    emit_board(root_head[k], root_cells[k], n_cells);
+    fprintf(stdout, "},\"leaf\":{\"parity\":%.0f,", row_parity_sign[k]);
+    emit_board(row_head[k], row_cells[k], n_cells);
+    fprintf(stdout, ",\"terminal\":%u,\"terminal_value\":%.1f}}",
+            row_term[k], row_terminal_value[k]);
   }
+  fprintf(stdout, "]}\n");
 
   return (int)n_rows;
 }
@@ -169,9 +187,9 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
 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 CSV row per recorded search, from ply <min_ply> onwards:\n\n\
-\tparity_sign,mover_flats,opponent_flats,<cell ids...>,R,terminal,leaf_R,game_id,cand_id,cand_R\n\n\
-R is the outcome from black's perspective, cand_R the outcome from the recorded candidate's perspective. nn_train.py runs many of these in parallel to generate training data.\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\
 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.";
 
@@ -212,13 +230,11 @@ int main(int argc, char **argv) {
             size == 5 ? FIVE_A : SIX_A, size == 5 ? FIVE_B : SIX_B,
             CAND_A,
             CAND_B,
-            rand32(),
             record, 0);
   play_game(min_ply,
             size == 5 ? FIVE_B : SIX_B, size == 5 ? FIVE_A : SIX_A,
             CAND_B,
             CAND_A,
-            rand32(),
             record, CAND_A == CAND_B ? 0: 1);