path: root/annotate_pgns/sanity_check_net.py

#!/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()