Wire NNUE eval into search, build the training pipeline, and add a fast

move/eval-agreement benchmark against Stockfish (tests/lop)

Implements the architecture decisions in src/docs/NNUE.md: HalfKAv2_hm
feature set, int16/int8 quantization, per-ply accumulator backup for
king moves, and a lazy-eval-ready integration into search.c/eval.c.
Adds the full annotate_pgns/ training pipeline (extraction, K-fit/
labeling, quiet filtering, tensor building, training, export) plus an
end-to-end orchestrator (run_nnue_pipeline.py) and validation tooling
(sanity_check_net.py, compare_vs_stockfish_net.py). Adds tests/lop:
a 10k-position list-of-positions benchmark comparing engine eval rank
and move choice against pre-computed Stockfish/lc0 references, much
cheaper than a full self-play match for day-to-day progress checks.

First sd10-matched self-play result against classic HCE: LLR +2.28
(H1=+2.94), -72=119+169 over 360 games, confirming the net is
"smarter per node" at matched depth. A real-time (--st 4) match is
currently weaker (candidate loses ~3 plies of depth to a ~3.3-3.7x
nps gap against an unrecalibrated search), consistent with
NNUE.md's adoption-criterion section -- pruning/reduction constants
still need retuning for the new eval's score scale before a fair
real-time comparison is possible.

Co-Authored-By: Claude Sonnet 5 <[email protected]>

1 files changed, 306 insertions, 0 deletions

diff --git a/annotate_pgns/sanity_check_net.py b/annotate_pgns/sanity_check_net.py
new file mode 100755
index 0000000..fd3b240
--- /dev/null
+++ b/annotate_pgns/sanity_check_net.py
@@ -0,0 +1,306 @@
+#!/usr/bin/env python3
+"""Material-sensitivity sanity check for a trained NNUE checkpoint.
+
+Loads a raw PyTorch checkpoint (nnue_train.py's ckpt_epochNNN.pt) and
+runs it, via nnue_model.py's forward() directly -- no export/C
+round-trip needed, so this is the earliest point in the pipeline a
+checkpoint can be checked -- against a battery of hand-picked base
+positions, each with one piece removed at a time. Reports the raw
+score delta from the base position for each removal and flags
+whether the deltas are monotonic in material value (pawn < knight ~
+bishop < rook < queen), which a net with a real concept of material
+should satisfy even though nothing forces it to structurally.
+
+This exists because of a real, measured failure found 2026-10-17:
+removing a queen, a rook, a knight, and a pawn from the same base
+position produced nearly indistinguishable score drops (-0.56 to
+-0.70) -- directionally correct (material loss always hurts) but
+with no real sense of HOW MUCH. Likely root cause: training loss is
+MSE against sigmoid(raw/K), and sigmoid's gradient vanishes once a
+position is already clearly lost/won, so a decisively-bad position
+gets little pressure to become MORE decisively bad regardless of
+how much additional material disappears. This script exists to make
+that failure mode visible at a glance on any future checkpoint,
+rather than rediscovering it by hand each time.
+
+Base positions deliberately avoid the opening setup: deleting a
+queen from the starting position is itself wildly out-of-distribution
+(no real game ever reaches that exact shape), so a net's response to
+it isn't a trustworthy read -- these use realistic midgame/endgame-ish
+material counts instead.
+
+2026-10-17: added --cross-sectional, a population-level version of the
+same idea. The hand-picked BASE_POSITIONS checks above are useful as a
+fast smoke test, but a single removal tests one specific piece on one
+specific square in one specific position -- its result is confounded
+by that piece's concrete role there (an active bishop vs. a boxed-in
+rook can legitimately score differently for positional reasons having
+nothing to do with the net's general sense of material), as observed
+directly: bishop-removal outscored rook-removal in one hand-picked
+position and the reverse held in another. --cross-sectional instead
+samples many REAL positions from a holdout pool, removes one piece of
+each type per position (when present), and pools deltas by piece
+type across the whole sample -- so per-position context should average
+out, leaving a cleaner read on the net's actual learned material
+prior (mean/median/stdev per piece type, not a single number).
+
+Usage:
+    python3 sanity_check_net.py --ckpt runs/pooled_gen2/nnue_ckpt/ckpt_epoch032.pt
+    python3 sanity_check_net.py --ckpt ... --cross-sectional \
+        --sample-file runs/pooled_gen2/games.holdout.labeled.quiet.jsonl \
+        --n-positions 2000
+"""
+import argparse
+import glob
+import json
+import random
+
+import chess
+import numpy as np
+import torch
+
+from nnue_feature_spec import active_feature_indices
+from nnue_model import NnueNet
+
+# (piece_type, approximate pawn value) -- value only used to order
+# the monotonicity check, not fed to the net.
+PIECE_VALUES = [
+    (chess.PAWN, 1),
+    (chess.KNIGHT, 3),
+    (chess.BISHOP, 3),
+    (chess.ROOK, 5),
+    (chess.QUEEN, 9),
+]
+
+BASE_POSITIONS = {
+    "midgame_balanced":
+        "r2qk2r/ppp2ppp/2n1bn2/2b1p3/4P3/2N1BN2/PPP2PPP/R2QK2R w KQkq - 0 1",
+    "midgame_balanced_black_to_move":
+        "r2qk2r/ppp2ppp/2n1bn2/2b1p3/4P3/2N1BN2/PPP2PPP/R2QK2R b KQkq - 0 1",
+    "endgame_rook_pawns":
+        "8/5pk1/6p1/8/4R3/5PK1/6P1/8 w - - 0 1",
+    "queen_endgame":
+        "6k1/5ppp/8/8/8/8/4QPPP/6K1 w - - 0 1",
+}
+
+
+def load_model(ckpt_path):
+    ck = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+    model = NnueNet()
+    model.load_state_dict(ck["model"])
+    model.eval()
+    return model, ck.get("holdout_mse")
+
+
+def score(model, board):
+    w = active_feature_indices(board, chess.WHITE)
+    b = active_feature_indices(board, chess.BLACK)
+    w_idx = torch.tensor(w, dtype=torch.long)
+    w_off = torch.tensor([0], dtype=torch.long)
+    b_idx = torch.tensor(b, dtype=torch.long)
+    b_off = torch.tensor([0], dtype=torch.long)
+    stm = torch.tensor([1 if board.turn == chess.WHITE else 0], dtype=torch.long)
+    with torch.no_grad():
+        return model.forward(w_idx, w_off, b_idx, b_off, stm).item()
+
+
+def one_sided_removals(board, color):
+    """Yields (piece_type, value, square, board_with_piece_removed)
+    for every piece of `color` on the board, one at a time -- skips
+    kings (removing a king produces an illegal/meaningless position)."""
+    for piece_type, value in PIECE_VALUES:
+        squares = list(board.pieces(piece_type, color))
+        if not squares:
+            continue
+        sq = squares[0]
+        b2 = board.copy()
+        b2.remove_piece_at(sq)
+        yield piece_type, value, sq, b2
+
+
+def run_base_position(model, name, fen):
+    board = chess.Board(fen)
+    mover_color = board.turn
+    mover_name = "White" if mover_color == chess.WHITE else "Black"
+    base_score = score(model, board)
+    print(f"\n=== {name} ({mover_name} to move) ===")
+    print(f"  base score: {base_score:+.3f}")
+
+    print(f"  -- removing {mover_name}'s OWN pieces (score should drop, "
+          f"and drop MORE for more valuable pieces) --")
+    rows = []
+    for piece_type, value, sq, b2 in one_sided_removals(board, mover_color):
+        s = score(model, b2)
+        delta = s - base_score
+        rows.append((value, chess.piece_name(piece_type), delta))
+        print(f"    remove own {chess.piece_name(piece_type):8s} "
+              f"(val={value}): score={s:+.3f}  delta={delta:+.3f}")
+
+    # Monotonicity check: sort by piece value, deltas should be
+    # non-increasing (more valuable piece removed -> more negative
+    # delta), allowing ties between knight/bishop (same nominal value).
+    rows.sort(key=lambda r: r[0])
+    deltas = [d for _, _, d in rows]
+    monotonic = all(deltas[i] >= deltas[i + 1] - 1e-6 for i in range(len(deltas) - 1))
+    spread = max(deltas) - min(deltas) if deltas else 0.0
+    print(f"  monotonic (pawn..queen delta non-increasing)? {monotonic}  "
+          f"(spread between least/most valuable piece's delta: {spread:.3f})")
+    if not monotonic:
+        print("  *** FAIL: removing a MORE valuable piece did not hurt "
+              "the score MORE -- net may not have learned material. ***")
+    if deltas and spread < 0.5:
+        print(f"  *** WARNING: spread of only {spread:.3f} across pawn..queen "
+              f"removal is small -- net may be material-insensitive even if "
+              f"technically monotonic. ***")
+
+    print(f"  -- removing opponent's pieces (score should RISE) --")
+    opp_color = not mover_color
+    for piece_type, value, sq, b2 in one_sided_removals(board, opp_color):
+        s = score(model, b2)
+        delta = s - base_score
+        flag = "" if delta > 0 else "  *** FAIL: should have risen ***"
+        print(f"    remove opp {chess.piece_name(piece_type):8s} "
+              f"(val={value}): score={s:+.3f}  delta={delta:+.3f}{flag}")
+
+
+def sample_fens(sample_file, n_positions, seed):
+    """Reservoir-sample n_positions FENs out of every position in
+    every game in a games.jsonl-shaped file (quiet_filter_games.py
+    output). Loads all FENs into memory first (a string per position,
+    cheap even at a few hundred thousand rows) rather than true
+    streaming reservoir sampling -- simpler, and the holdout pools
+    this is meant to run against are a few hundred MB at most."""
+    fens = []
+    with open(sample_file) as f:
+        for line in f:
+            game = json.loads(line)
+            for pos in game["positions"]:
+                fens.append(pos["fen"])
+    rng = random.Random(seed)
+    if n_positions < len(fens):
+        fens = rng.sample(fens, n_positions)
+    return fens
+
+
+def one_random_removal_per_piece_type(board, color, rng):
+    """Yields (piece_type, value, board_with_one_piece_removed) --
+    exactly one randomly-chosen piece per type present for `color`
+    (not every piece, unlike one_sided_removals -- across thousands
+    of sampled positions this still covers many squares per type in
+    aggregate without each position contributing >1 sample per type,
+    which would overweight positions with e.g. two bishops)."""
+    for piece_type, value in PIECE_VALUES:
+        squares = list(board.pieces(piece_type, color))
+        if not squares:
+            continue
+        sq = rng.choice(squares)
+        b2 = board.copy()
+        b2.remove_piece_at(sq)
+        yield piece_type, value, b2
+
+
+def run_cross_sectional(model, sample_file, n_positions, seed):
+    print(f"\n=== cross-sectional population check ===")
+    print(f"  sampling {n_positions} positions from {sample_file} ...")
+    fens = sample_fens(sample_file, n_positions, seed)
+    print(f"  got {len(fens)} FENs")
+
+    rng = random.Random(seed)
+    # own_deltas[piece_type] = list of deltas from removing the MOVER's
+    # own piece of that type (expected negative, more negative for
+    # more valuable pieces). opp_deltas = removing the opponent's
+    # piece instead (expected positive).
+    own_deltas = {pt: [] for pt, _ in PIECE_VALUES}
+    opp_deltas = {pt: [] for pt, _ in PIECE_VALUES}
+    n_skipped = 0
+
+    for i, fen in enumerate(fens):
+        try:
+            board = chess.Board(fen)
+        except ValueError:
+            n_skipped += 1
+            continue
+        mover = board.turn
+        base_score = score(model, board)
+        for piece_type, value, b2 in one_random_removal_per_piece_type(board, mover, rng):
+            own_deltas[piece_type].append(score(model, b2) - base_score)
+        for piece_type, value, b2 in one_random_removal_per_piece_type(board, not mover, rng):
+            opp_deltas[piece_type].append(score(model, b2) - base_score)
+        if (i + 1) % 500 == 0:
+            print(f"  ...{i + 1}/{len(fens)} positions processed")
+
+    if n_skipped:
+        print(f"  ({n_skipped} FENs failed to parse, skipped)")
+
+    def report(deltas_by_type, label, expect_sign):
+        print(f"\n  -- {label} (expect sign: {'negative' if expect_sign < 0 else 'positive'}) --")
+        means = []
+        for piece_type, value in PIECE_VALUES:
+            d = np.array(deltas_by_type[piece_type])
+            if len(d) == 0:
+                print(f"    {chess.piece_name(piece_type):8s} (val={value}): no samples")
+                continue
+            means.append((value, chess.piece_name(piece_type), d.mean()))
+            print(f"    {chess.piece_name(piece_type):8s} (val={value}): "
+                  f"n={len(d):6d}  mean={d.mean():+.4f}  median={np.median(d):+.4f}  "
+                  f"std={d.std():.4f}  p25={np.percentile(d, 25):+.4f}  "
+                  f"p75={np.percentile(d, 75):+.4f}")
+        means.sort(key=lambda r: r[0])
+        mean_vals = [m for _, _, m in means]
+        if expect_sign < 0:
+            monotonic = all(mean_vals[i] >= mean_vals[i + 1] - 1e-6 for i in range(len(mean_vals) - 1))
+        else:
+            monotonic = all(mean_vals[i] <= mean_vals[i + 1] + 1e-6 for i in range(len(mean_vals) - 1))
+        print(f"    population means monotonic in piece value? {monotonic}")
+        if not monotonic:
+            order = ", ".join(f"{name}={m:+.4f}" for _, name, m in means)
+            print(f"    *** FAIL: population-level mean deltas not ordered "
+                  f"by piece value: {order} ***")
+
+    report(own_deltas, "removing MOVER's own piece", expect_sign=-1)
+    report(opp_deltas, "removing OPPONENT's piece", expect_sign=+1)
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__,
+                                  formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--ckpt", default=None,
+                     help="checkpoint path. Default: latest "
+                          "ckpt_epoch*.pt under --ckpt-glob-dir")
+    ap.add_argument("--ckpt-glob-dir", default="runs/pooled_gen2/nnue_ckpt",
+                     help="directory to find the latest checkpoint in, "
+                          "when --ckpt is not given")
+    ap.add_argument("--cross-sectional", action="store_true",
+                     help="run the population-level check (see module "
+                          "docstring's 2026-10-17 note) instead of the "
+                          "hand-picked single-position checks")
+    ap.add_argument("--sample-file",
+                     default="runs/pooled_gen2/games.holdout.labeled.quiet.jsonl",
+                     help="games.jsonl-shaped file to sample real "
+                          "positions from, for --cross-sectional")
+    ap.add_argument("--n-positions", type=int, default=2000,
+                     help="number of positions to sample, for --cross-sectional")
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args()
+
+    ckpt_path = args.ckpt
+    if ckpt_path is None:
+        candidates = sorted(glob.glob(f"{args.ckpt_glob_dir}/ckpt_epoch*.pt"))
+        if not candidates:
+            raise SystemExit(f"No checkpoints found under {args.ckpt_glob_dir}")
+        ckpt_path = candidates[-1]
+
+    print(f"Loading {ckpt_path} ...")
+    model, holdout_mse = load_model(ckpt_path)
+    if holdout_mse is not None:
+        print(f"  (checkpoint's recorded holdout_mse: {holdout_mse:.5f})")
+
+    if args.cross_sectional:
+        run_cross_sectional(model, args.sample_file, args.n_positions, args.seed)
+    else:
+        for name, fen in BASE_POSITIONS.items():
+            run_base_position(model, name, fen)
+
+
+if __name__ == "__main__":
+    main()