summaryrefslogtreecommitdiff
path: root/annotate_pgns/sample_egtb_positions.py
blob: f4cfa45508a6d902acb952b4a6424585b9604e89 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
#!/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()