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