From 241156da6f35cb888253e6b8814e7b61a76b4fc7 Mon Sep 17 00:00:00 2001 From: tslil Date: Sun, 30 Aug 2026 19:46:07 +0100 Subject: TD(lambda), and a big rework NOTE: 6x6 weights are as-yet untrained On the performance side: * switch from linked-list chaining hash table to direct mapping * cache check_win by hash * persist transposition tables across plies (this was a big deal and difficult to puzzle out) * add align directives to the network floats * prepare_nn_input gained a small LUT to avoid branching (i suppose you might call this 1D embedding) * apply mtune=native for the native builds On the correctness side: * Zobrist has now includes the ply parity * the network now predicts something in [-1, 1] using tanh (approximated by polynomials) * the training object is now MSE because we're learning a value estimate * use .clangd instead of compile_commands.json On the training side: * completely forgo human-derived data set * implement td(lambda), self-play, tournaments --- src/nn_train.py | 665 ++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 575 insertions(+), 90 deletions(-) mode change 100644 => 100755 src/nn_train.py (limited to 'src/nn_train.py') diff --git a/src/nn_train.py b/src/nn_train.py old mode 100644 new mode 100755 index 03852f5..52549d8 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -1,6 +1,24 @@ -# cnn_train.py, train a small CNN to recognise winning Tak positions +#!/usr/bin/env -S uv run --script # -# Copyright (C) 2021, tslil clingman +# /// script +# requires-python = ">=3.14" +# dependencies = [ +# "pandas>=3.0.5", +# "torch>=2.14.0", +# "tqdm>=4.70.0", +# ] +# +# [tool.uv.sources] +# torch = { index = "pytorch-cpu" } +# +# [[tool.uv.index]] +# name = "pytorch-cpu" +# url = "https://download.pytorch.org/whl/cpu" +# explicit = true +# /// + +# +# This file is part of ctak. # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by @@ -13,95 +31,562 @@ # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License -# along with this program. If not, see . +# along with ctak. If not, see . +from collections.abc import Iterator +from dataclasses import dataclass + +import argparse +import concurrent.futures +import copy +import os +from pathlib import Path +import subprocess -# Loading data -import pandas as pd import numpy as np +import pandas as pd +import torch +from tqdm import tqdm + +from torch.nn import MSELoss, Linear, Module +from torch import Tensor + + +# =================================================================== +# Net dimensions +# =================================================================== + +# (input dim, hidden dim) per board size; must match include/weights_.h +NET_DIMS: dict[int, tuple[int, int]] = { + 5: (28, 64), # FIVE_INP_NUM, FIVE_DENSE_NUM + 6: (39, 64), # SIX_INP_NUM, SIX_DENSE_NUM +} + + +# =================================================================== +# C weight formatting +# =================================================================== + + +def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str: + a = np.ravel(np.asarray(a, dtype=np.float32)) + lines: list[str] = [] + for i in range(0, len(a), per_line): + chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line]) + if i + per_line < len(a): + chunk += "," + lines.append(" " + chunk) + return "\n".join(lines) + + +def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str: + mat = np.asarray(mat, dtype=np.float32) + rows: list[str] = [] + for r in mat: + chunks = [ + ", ".join(f"{v:.9g}" for v in r[i : i + per_line]) + for i in range(0, len(r), per_line) + ] + rows.append(" {" + ",\n ".join(chunks) + "},") + return "\n".join(rows) + + +# =================================================================== +# Self-play records +# =================================================================== + + +def td_lambda_targets( + values: Tensor, + result: float, + td_lambda: float, +) -> Tensor: + # values: V(s_t) for t = 0..T-1, shape [T] + # result: final result R of the game + + # This is ema, computed backwards + # target[T-1] = R + # target[t] = (1 - lambda) * V(s_{t+1}) + lambda * target[t+1] + length = len(values) + targets = torch.empty_like(values) + targets[-1] = result + for ply in range(length - 2, -1, -1): + targets[ply] = (1.0 - td_lambda) * values[ply + 1] + td_lambda * targets[ + ply + 1 + ] + return targets + + +@dataclass(frozen=True) +class SelfPlayRecords: + positions: Tensor # [n, inp_dim] feature vector of each position + game_ids: Tensor # [n] id of the game the position belongs to + plies: Tensor # [n] ply number (position is before this move) + results: Tensor # [n] final result from black's perspective (+1 black won, -1 white won, 0 draw) + + @classmethod + def load_selfplay( + cls, fn: Path, size: int, device: torch.device + ) -> SelfPlayRecords: + # [ net input: 0-inp_dim ] [game_id] [ply] [game result] + df = pd.read_csv(fn, header=None) + inp_dim = NET_DIMS[size][0] + return cls( + positions=torch.tensor( + df.iloc[:, 0:inp_dim].to_numpy(dtype=np.float32), + dtype=torch.float32, + device=device, + ), + game_ids=torch.tensor( + df.iloc[:, inp_dim].to_numpy(dtype=np.int64), + device=device, + ), + plies=torch.tensor( + df.iloc[:, inp_dim + 1].to_numpy(dtype=np.int64), + device=device, + ), + results=torch.tensor( + df.iloc[:, inp_dim + 2].to_numpy(dtype=np.float32), + dtype=torch.float32, + device=device, + ), + ) + + def make_training_data( + self, model: Module, td_lambda: float + ) -> tuple[Tensor, Tensor]: + model.eval() + with torch.no_grad(): + values = model(self.positions) + + targets = torch.empty_like(values) + + for game_id in torch.unique(self.game_ids): + in_game = self.game_ids == game_id + game_values = values[in_game] + game_result = self.results[in_game][0].item() + targets[in_game] = td_lambda_targets(game_values, game_result, td_lambda) + + return self.positions, targets + + +# =================================================================== +# Network +# =================================================================== + + +class Net(Module): + def __init__(self, inp_dim: int, hidden: int) -> None: + super().__init__() + self.hidden = Linear(inp_dim, hidden, bias=True) + self.output = Linear(hidden, 1, bias=True) + + def forward(self, x: Tensor) -> Tensor: + hidden = torch.relu(self.hidden(x)) + return torch.tanh(self.output(hidden)).squeeze(-1) + + def fit( + self, + *, + self_play_records: SelfPlayRecords, + td_lambda: float, + val_frac: float, + epochs: int, + batch: int, + lr: float, + ) -> tuple[float, float]: + # data + positions, targets = self_play_records.make_training_data(self, td_lambda) + n = len(positions) + perm = torch.randperm(n, device=positions.device) + n_val = max(1, int(n * val_frac)) + val_idx = perm[:n_val] + train_idx = perm[n_val:] + + val_positions = positions[val_idx] + val_targets = targets[val_idx] + train_positions = positions[train_idx] + train_targets = targets[train_idx] + + # training widgets + optimizer = torch.optim.Adam(self.parameters(), lr=lr) + loss_fn = MSELoss() + + best_corr = -1.0 + best_state: dict[str, Tensor] | None = None + best_mse = float("nan") + + pbar = tqdm(range(epochs), desc="train", unit="ep") + for epoch in pbar: + self.train() + order = torch.randperm(len(train_positions), device=positions.device) + for start in range(0, len(train_positions), batch): + batch_idx = order[start : start + batch] + optimizer.zero_grad() + predictions = self(train_positions[batch_idx]) + loss = loss_fn(predictions, train_targets[batch_idx]) + loss.backward() + optimizer.step() + + with torch.no_grad(): + self.eval() + predictions = self(val_positions) + mse = loss_fn(predictions, val_targets).item() + corr = torch.corrcoef(torch.stack([predictions, val_targets]))[ + 0, 1 + ].item() + + pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}") + + if corr > best_corr: + best_corr = corr + best_state = { + name: param.clone() for name, param in self.state_dict().items() + } + best_mse = mse + + if best_state is not None: + self.load_state_dict(best_state) + + return best_corr, best_mse + + def to_c( + self, + size: int, + path: Path, + performance: tuple[str, ...] = (), + cand: bool = False, + ) -> None: + prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six") + upper = "FIVE" if size == 5 else "SIX" + + W = self.hidden.weight.detach().cpu().numpy() + B = self.hidden.bias.detach().cpu().numpy() + LW = self.output.weight.detach().cpu().numpy().ravel() + LB = float(self.output.bias.detach().cpu().item()) + + with open(path, "w") as f: + _ = f.write("/* generated by nn_train.py, do not edit by hand */\n") + for p in performance: + _ = f.write(f"/* {p} */\n") + _ = f.write(f'#include "weights_{size}.h"\n\n') + _ = f.write( + f"const float {prefix}_dense_weights[{upper}_DENSE_NUM]" + f"[{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n" + ) + _ = f.write( + f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n" + f"{{\n{_fmt_flat(B)}\n}};\n\n" + ) + _ = f.write( + f"const float {prefix}_logit_weights[{upper}_DENSE_NUM]" + f" NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n" + ) + _ = f.write(f"const float {prefix}_logit_bias = {LB:.9g};\n") + + +# =================================================================== +# Process helpers +# =================================================================== + + +def fan_out(cmd: list[str], n: int, procs: int, tag: str, root: Path) -> Iterator[str]: + with concurrent.futures.ThreadPoolExecutor(max_workers=procs) as pool: + futs = [ + pool.submit( + subprocess.run, + cmd, + cwd=root, + check=True, + capture_output=True, + text=True, + ) + for _ in range(n) + ] + for fut in tqdm( + concurrent.futures.as_completed(futs), total=n, desc=tag, unit="game" + ): + yield fut.result().stdout + + +def cand_copy_for(size: int, root: Path) -> None: + src = root / "include" / f"weights_{size}.c" + if not src.is_file(): + raise SystemExit(f"{src} not found; build a size-{size} champion first") + prefix = "five" if size == 5 else "six" + txt = src.read_text().replace(f"{prefix}_", f"cand_{prefix}_") + _ = (root / "include" / f"cand_weights_{size}.c").write_text(txt) + + +# =================================================================== +# Tournament +# =================================================================== + + +@dataclass(frozen=True) +class TournamentResult: + champ: int + cand: int + draws: int + cand_rate: float + + +def run_matches( + *, + size: int, + depth: int, + tournament_eps: float, + games: int, + procs: int, + root: Path, +) -> TournamentResult: + outs = list( + fan_out( + ["./match", str(size), f"{tournament_eps:.4f}", str(depth)], + games, + procs, + "match", + root, + ) + ) + winners = [line.split()[0] for line in outs if line.strip()] + champ = winners.count("current") + cand = winners.count("candidate") + draws = winners.count("draw") + aborts = len(winners) - champ - cand - draws + played = len(winners) - aborts + return TournamentResult( + champ=champ, + cand=cand, + draws=draws, + cand_rate=(cand + 0.5 * draws) / played if played else 0.0, + ) + + +# =================================================================== +# The bootstrap loop +# =================================================================== + + +def setup(size: int, root: Path) -> Net: + inp_dim, hidden = NET_DIMS[size] + champ = Net(inp_dim, hidden) + champ.to_c( + size, + root / "include" / f"weights_{size}.c", + ("v000, random init",), + ) + cand_copy_for(6 if size == 5 else 5, root) + return champ + + +def run_step( + *, + iter: int, + size: int, + eps_min: float, + champ: Net, + eps_start: float, + eps_decay: float, + depth: int, + min_ply: int, + games: int, + procs: int, + td_lambda: float, + val_frac: float, + epochs: int, + batch: int, + lr: float, + tournament_eps: float, + tournament_games: int, + promote_rate: float, + root: Path, +) -> Net: + print(f"\n== iteration {iter} ==") + + # 1. self-play with the installed champion (include/weights_.c) + _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True) + + # 2. generate self-play data (one game per process, in parallel) + eps = max(eps_min, eps_start * (eps_decay ** (iter - 1))) + csv = root / "data" / f"selfplay-{size}-v{iter:03d}.csv" + print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs") + out_rows: list[str] = [] + for out in fan_out( + ["./selfplay", str(size), f"{eps:.4f}", str(depth), str(min_ply)], + games, + procs, + "selfplay", + root, + ): + out_rows.extend(out.splitlines()) + with open(csv, "w") as f: + _ = f.write("\n".join(out_rows)) + if out_rows: + _ = f.write("\n") + rows = len(out_rows) + print(f" {rows} rows") + + # 3. fine-tune the champion on TD(lambda) targets -> candidate + device = next(champ.parameters()).device + records = SelfPlayRecords.load_selfplay(csv, size, device) + cand = copy.deepcopy(champ) + + corr, mse = cand.fit( + self_play_records=records, + td_lambda=td_lambda, + val_frac=val_frac, + epochs=epochs, + batch=batch, + lr=lr, + ) + cand.to_c( + size, + root / "include" / f"cand_weights_{size}.c", + (f"v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",), + cand=True, + ) + + # 4. build the match driver (champion + candidate linked) + _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True) + + # 5. tournament + decision + tournament_result = run_matches( + root=root, + size=size, + depth=depth, + tournament_eps=tournament_eps, + games=tournament_games, + procs=procs, + ) + decided = tournament_result.champ + tournament_result.cand + rate = tournament_result.cand / decided if decided else 0.0 + print( + f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, " + + f"winrate={rate:.3f}, cand_rate={tournament_result.cand_rate:.3f}" + ) + promote = decided >= int(0.25 * tournament_games) and rate >= promote_rate + + if promote: + cand.to_c( + size, + root / "include" / f"weights_{size}.c", + (f"v{iter:03d}, val_corr={corr:.4f}",), + ) + print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==") + return cand + print(f"== candidate v{iter:03d} rejected ==") + return champ + + +# =================================================================== +# Main +# =================================================================== + + +def main() -> None: + root = Path(__file__).resolve().parent.parent + p = argparse.ArgumentParser( + description="the TD(lambda) self-play bootstrap loop", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + _ = p.add_argument( + "--size", type=int, choices=sorted(NET_DIMS), default=5, help="board size" + ) + _ = p.add_argument("--iters", type=int, default=10, help="max loop iterations") + _ = p.add_argument( + "--games", type=int, default=2000, help="selfplay games per iteration" + ) + _ = p.add_argument( + "--eps-start", + type=float, + default=0.15, + help="initial self-play epsilon (random-move probability); kept close " + "to --tour-eps so the value net trains on the policy it is deployed " + "under", + ) + _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor") + _ = p.add_argument( + "--eps-decay", type=float, default=0.8, help="per-iteration epsilon decay" + ) + _ = p.add_argument( + "--depth", + type=int, + default=4, + help="search depth for selfplay AND the tournament", + ) + _ = p.add_argument( + "--min-ply", type=int, default=8, help="record positions at this ply or later" + ) + _ = p.add_argument("--td-lambda", type=float, default=0.6, help="TD(lambda) lambda") + _ = p.add_argument( + "--epochs", type=int, default=20, help="candidate training epochs" + ) + _ = p.add_argument("--lr", type=float, default=1e-3, help="learning rate") + _ = p.add_argument("--batch", type=int, default=128, help="training batch size") + _ = p.add_argument( + "--val-frac", + type=float, + default=0.05, + help="fraction of rows held out for validation", + ) + _ = p.add_argument( + "--tournament-eps", + type=float, + default=0.05, + help="tournament epsilon (random-move probability)", + ) + _ = p.add_argument( + "--tournament-games", + type=int, + default=None, + help="tournament games (10/tour_eps if unset)", + ) + _ = p.add_argument( + "--promote", + type=float, + default=0.57, + help="candidate win rate among non-drawn tournament games needed to promote (~p<0.03 one-sided binomial at the default 200 games)", + ) + _ = p.add_argument( + "--procs", + type=int, + default=None, + help="parallel game processes (half the CPUs if unset)", + ) + args = p.parse_args() + + if args.tournament_games is None: + args.tournament_games = max(20, int(round(10 / args.tournament_eps))) + if args.procs is None: + args.procs = max(1, (os.cpu_count() or 1) // 2) + + (root / "data").mkdir(exist_ok=True) + champ = setup(size=args.size, root=root) + for iter in range(1, args.iters + 1): + champ = run_step( + iter=iter, + size=args.size, + eps_min=args.eps_min, + champ=champ, + eps_start=args.eps_start, + eps_decay=args.eps_decay, + depth=args.depth, + min_ply=args.min_ply, + games=args.games, + procs=args.procs, + td_lambda=args.td_lambda, + val_frac=args.val_frac, + epochs=args.epochs, + batch=args.batch, + lr=args.lr, + tournament_eps=args.tournament_eps, + tournament_games=args.tournament_games, + promote_rate=args.promote, + root=root, + ) + print(f"\n== done, champion in include/weights_{args.size}.c ==") + -# Output of weights -from tensorflow import transpose - -# Building models -from tensorflow.keras.models import Model -from tensorflow.keras.layers import Input, Dense - - -def load_data(size): - output_len = 2 - # Load training data - tr_fn = "data/training-"+str(size)+".csv" - training_csv = pd.read_csv(tr_fn) - training_data = training_csv - training_input = np.array(training_data.iloc[:, 0:-output_len]) - training_outcome = np.array(training_data.iloc[:, -output_len:]) - # Load validation data - val_fn = "data/validation-"+str(size)+".csv" - val_data = pd.read_csv(val_fn) - val_input = np.array(val_data.iloc[:, 0:-output_len]) - val_outcome = np.array(val_data.iloc[:, -output_len:]) - return ((training_input, training_outcome), (val_input, val_outcome)) - - -def train(size, model, data, iterations=1, epochs=10, batch=None): - (tra_input, tra_outcome), (val_input, val_outcome) = data - results = [] - for i in range(0, iterations): - print("Iteration {0}/{1}".format(i+1, iterations)) - model.fit(tra_input, tra_outcome, epochs=epochs, - validation_data=(val_input, val_outcome), - verbose=True, batch_size=batch) - val_res = model.evaluate(val_input, val_outcome, verbose=False) - tra_res = model.evaluate(tra_input, tra_outcome, verbose=False) - results.append((tra_res, val_res)) - print(val_res) - write_weights(size, model, str(i+1).zfill(len(str(iterations))), (tra_res, val_res)) - print("\nScores") - for i, data in enumerate(results): - print(f"Iteration {i+1}: {data}") - return results - - -def make_model(size, magic): - inputs = Input(shape=(size * size + 3,)) - model = inputs - model = Dense(magic, activation="relu", use_bias=True)(model) - model = Dense(2, activation="relu", use_bias=True)(model) - model = Model(inputs=inputs, outputs=model) - model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy']) - model.summary() - return model - - -def write_weights(size, model, iteration, performance): - def fix(val): - string = str(np.array(val).tolist()) - string = string.replace("[", "{").replace("]", "}") - return string - size_str = "five" if size==5 else "six" - dense_weights = transpose(model.trainable_variables[0], perm=[1, 0]) - dense_biases = model.trainable_variables[1] - output_weights = transpose(model.trainable_variables[2], perm=[1, 0]) - output_bias = model.trainable_variables[3] - # Prepare output - to_output = [(f"{size_str}_dense_weights[{size_str.upper()}_DENSE_NUM][{size_str.upper()}_INP_NUM]", dense_weights), - (f"{size_str}_dense_biases[{size_str.upper()}_DENSE_NUM]", dense_biases), - (f"{size_str}_output_weights[2][{size_str.upper()}_DENSE_NUM]", output_weights), - (f"{size_str}_output_bias[2]", output_bias)] - # Write to file - f = open(f"{size_str}_weights-"+iteration+".txt", "w") - f.write("/*\n") - model.summary(print_fn=lambda l: f.write(" * "+l+"\n")) - f.write(" * "+str(performance)+"\n*/\n\n") - f.write(f"#include \"weights_{size}.h\"\n\n") - for (name, val) in to_output: - f.write("const float "+name+" =\n"+fix(val)+";\n\n") - f.close() - - -for size in (5,6): - data = load_data(size) - model = make_model(size, 64) - print("Before training", model.evaluate(data[1][0], data[1][1], verbose=False, batch_size=16)) - results = train(size, model, data, iterations=5, epochs=10, batch=128) +if __name__ == "__main__": + main() -- cgit v1.3.1