working towards hybrid rootstrap

1 files changed, 273 insertions, 183 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 ==")