#!/usr/bin/env python3 """Generate direct Syzygy-ground-truth training positions for the EGTB-mixing idea in NNUE.md: random LEGAL positions for a curated set of endgame piece configurations, each probed directly against the Syzygy tablebase for its exact WDL result and labeled with that result (1.0/0.0/0.5, mover-relative) -- NOT derived from any engine search score, so there's no sigmoid/lambda involved at all for these rows (see build_training_tensors.py's position_target: exact rows with an r_outcome use it directly). Why this exists as a separate generator rather than just trusting Stockfish self-play games to supply this signal: a self-play game's score field, even when the search is quietly influenced by a TB probe a few plies down, is NOT a reliable way to detect "this exact position is a TB-determined result" after the fact -- see parse_comment's MATE_CLAMP_THRESHOLD handling in extract_selfplay_positions.py, which can only catch the unambiguous extreme case (|score| close to the engine's own mate-score cap). An ordinary-looking, e.g., "+0.30" 6-piece position could just as easily be a genuine TB draw as a normal balanced middlegame-ish score, and there is no way to tell the two apart from the log alone. Rather than guess, this script sidesteps the ambiguity entirely: it builds positions whose piece count and side-to-move are known up front, and asks the tablebase directly, which is unambiguous by construction. Curated piece-set list: endgames that are either common in real games (KPK, KRK, KQK, basic mating/drawing technique) or specifically instructive for a net's positional judgment near the piece-trade boundary (KRPKR, KBPKB, KNPKN -- same-minor-vs-pawn fortress/win distinctions; KPPKPP -- connected/outside-passer races). Deliberately excludes same-type-heavy, essentially-never-arising-from-real-play configurations like KNNNKB (three knights) or KQQKQQ -- these exist in the 6-piece Syzygy set for completeness but add no realistic training signal, only wasted sampling effort chasing positions the net will never see at inference time. Usage: python3 sample_egtb_positions.py \\ --egtb-path /zscratch/egtb \\ --out runs/egtb_direct/positions.jsonl \\ --per-config 2000 """ import argparse import json import random import sys import chess import chess.syzygy # Each entry: (white_pieces, black_pieces), piece letters excluding the # king (every side always has exactly one king, added automatically). # Ordered white-majority-first by convention; sampling itself places # White/Black and side-to-move uniformly at random regardless. USEFUL_CONFIGS = [ # 3-piece: the universal basics. ("P", ""), # KPK ("R", ""), # KRK ("Q", ""), # KQK ("B", "B"), # KBKB (same-color-bishop draw / opposite-color technique) ("N", "N"), # KNKN # 4-piece: one side up material, or pawn-down defensive technique. ("R", "P"), # KRKP ("Q", "P"), # KQKP ("Q", "R"), # KQKR ("R", "B"), # KRKB ("R", "N"), # KRKN ("B", "N"), # KBKN ("P", "P"), # KPKP (opposition/zugzwang pawn races) # 5-piece: minor+pawn fortress/technique boundary cases, specifically # called out as instructive in the discussion that motivated this. ("RP", "R"), # KRPKR (the single most common rook-endgame skeleton) ("BP", "B"), # KBPKB (same/opposite-color bishop-vs-pawn technique) ("NP", "N"), # KNPKN ("QP", "Q"), # KQPKQ ("PP", "P"), # KPPKP # 6-piece: the specific "KPPKPP, not KNNNKB" case from the discussion. ("PP", "PP"), # KPPKPP (connected/outside passer races) ("RP", "RP"), # KRPKRP ("BP", "BP"), # KBPKBP ("RR", "R"), # KRRKR ] PIECE_SYMBOL = { "P": chess.PAWN, "N": chess.KNIGHT, "B": chess.BISHOP, "R": chess.ROOK, "Q": chess.QUEEN, } def random_board(white_pieces, black_pieces, rng): """Attempts one random legal placement for this piece configuration (kings included). Returns a chess.Board, or None if this attempt didn't produce a legal position (retried by the caller) -- most retries are from kings ending up adjacent or the side not on move being left in check, both of which are just rejection-sampled away rather than avoided by construction, since legal-position rejection sampling is simple, fast for this few-piece count, and impossible to get subtly wrong the way a hand-built adjacency/legality check could.""" squares = rng.sample(range(64), 2 + len(white_pieces) + len(black_pieces)) board = chess.Board.empty() board.set_piece_at(squares[0], chess.Piece(chess.KING, chess.WHITE)) board.set_piece_at(squares[1], chess.Piece(chess.KING, chess.BLACK)) idx = 2 for letter in white_pieces: sq = squares[idx] idx += 1 if letter == "P" and chess.square_rank(sq) in (0, 7): return None board.set_piece_at(sq, chess.Piece(PIECE_SYMBOL[letter], chess.WHITE)) for letter in black_pieces: sq = squares[idx] idx += 1 if letter == "P" and chess.square_rank(sq) in (0, 7): return None board.set_piece_at(sq, chess.Piece(PIECE_SYMBOL[letter], chess.BLACK)) board.turn = chess.WHITE if rng.random() < 0.5 else chess.BLACK board.castling_rights = chess.BB_EMPTY if not board.is_valid(): return None return board def sample_config(tablebase, white_pieces, black_pieces, n, rng, max_attempts): """Returns up to n (fen, r_outcome) pairs for one piece configuration, r_outcome mover-relative (1.0/0.5/0.0), cursed-win/blessed-loss collapsed to the practical 1.0/0.0 (not 0.5) -- under the engine's own 50-move-rule handling a cursed win is still winnable with correct play short of the counter running out, so it is a much better training label as "winning" than as "drawn"; collapsing it the other way would teach the net to misjudge a real winning chance as hopeless.""" out = [] attempts = 0 while len(out) < n and attempts < max_attempts: attempts += 1 board = random_board(white_pieces, black_pieces, rng) if board is None: continue try: wdl = tablebase.probe_wdl(board) except chess.syzygy.MissingTableError: return out # this exact configuration has no table -- give up if wdl > 0: r_outcome = 1.0 elif wdl < 0: r_outcome = 0.0 else: r_outcome = 0.5 out.append((board.fen(), r_outcome)) return out def main(): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("--egtb-path", default="/zscratch/egtb") ap.add_argument("--out", required=True) ap.add_argument("--per-config", type=int, default=2000, help="target sample count per piece configuration " "(each configuration is also sampled with both " "colors' roles swapped, so total yield per entry " "in USEFUL_CONFIGS is up to 2x this)") ap.add_argument("--max-attempts-factor", type=int, default=50, help="give up on a configuration after " "per-config * this many random placements, in " "case legal-position rejection sampling is " "pathologically slow for some piece set") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args() rng = random.Random(args.seed) tablebase = chess.syzygy.open_tablebase(args.egtb_path) total = 0 with open(args.out, "w") as fout: for white_pieces, black_pieces in USEFUL_CONFIGS: for wp, bp in ({(white_pieces, black_pieces), (black_pieces, white_pieces)}): rows = sample_config(tablebase, wp, bp, args.per_config, rng, args.per_config * args.max_attempts_factor) for fen, r_outcome in rows: game_rec = { "game_id": f"egtb_direct_K{wp}K{bp}_{total}", "result": "*", "positions": [{ "fen": fen, "exact": True, "exact_reason": "egtb_direct", "r_outcome": r_outcome, }], } fout.write(json.dumps(game_rec) + "\n") total += 1 print(f" K{wp}vK{bp}: {len(rows)}/{args.per_config} sampled", file=sys.stderr) print(f"Wrote {args.out}: {total} direct-EGTB positions", file=sys.stderr) if __name__ == "__main__": main()