diff options
Diffstat (limited to 'src/nn_train.py')
| -rwxr-xr-x | src/nn_train.py | 116 |
1 files changed, 65 insertions, 51 deletions
diff --git a/src/nn_train.py b/src/nn_train.py index 8025b68..73a5eaf 100755 --- a/src/nn_train.py +++ b/src/nn_train.py @@ -48,7 +48,7 @@ import pandas as pd import torch from tqdm import tqdm -from torch.nn import MSELoss, Linear, Module +from torch.nn import Embedding, MSELoss, Linear, Module from torch import Tensor @@ -97,38 +97,35 @@ def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str: def td_lambda_targets( - values: Tensor, + leaf_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 = torch.empty_like(leaf_values) targets[-1] = result - for ply in range(length - 2, -1, -1): - targets[ply] = (1.0 - td_lambda) * values[ply + 1] + td_lambda * targets[ + + # G[t] = (1-lambda)V[t+1] + lambda G[t+1] + for ply in range(len(leaf_values) - 2, -1, -1): + targets[ply] = (1.0 - td_lambda) * leaf_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 + positions: Tensor # [n, inp_dim] 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) + results: Tensor # [n] final played result from black's perspective + terminals: Tensor # [n] true when the recorded position is a terminal leaf + terminal_values: Tensor # [n] exact value of a terminal leaf, otherwise 0 @classmethod def load_selfplay( cls, fn: Path, size: int, device: torch.device ) -> SelfPlayRecords: - # [ net input: 0-inp_dim ] [game_id] [ply] [game result] + # [net input] [game_id] [played result] [terminal flag] [leaf result] df = pd.read_csv(fn, header=None) inp_dim = NET_DIMS[size][0] return cls( @@ -141,12 +138,18 @@ class SelfPlayRecords: df.iloc[:, inp_dim].to_numpy(dtype=np.int64), device=device, ), - plies=torch.tensor( - df.iloc[:, inp_dim + 1].to_numpy(dtype=np.int64), + results=torch.tensor( + df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32), + dtype=torch.float32, + device=device, + ), + terminals=torch.tensor( + df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_), + dtype=torch.bool, device=device, ), - results=torch.tensor( - df.iloc[:, inp_dim + 2].to_numpy(dtype=np.float32), + terminal_values=torch.tensor( + df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32), dtype=torch.float32, device=device, ), @@ -160,8 +163,9 @@ class SelfPlayRecords: return cls( positions=torch.concat([spr.positions for spr in sprs]), game_ids=torch.concat([spr.game_ids for spr in sprs]), - plies=torch.concat([spr.plies for spr in sprs]), results=torch.concat([spr.results for spr in sprs]), + terminals=torch.concat([spr.terminals for spr in sprs]), + terminal_values=torch.concat([spr.terminal_values for spr in sprs]), ) def make_training_data( @@ -175,9 +179,16 @@ class SelfPlayRecords: for game_id in torch.unique(self.game_ids): in_game = self.game_ids == game_id - game_values = values[in_game] + game_terminals = self.terminals[in_game] + game_terminal_values = self.terminal_values[in_game] + game_values = values[in_game].clone() + game_values[game_terminals] = game_terminal_values[game_terminals] game_result = self.results[in_game][0].item() - targets[in_game] = td_lambda_targets(game_values, game_result, td_lambda) + game_targets = td_lambda_targets(game_values, game_result, td_lambda) + # Terminal positions are also direct supervised examples. Do this + # after constructing the TD sequence so it does not reset the EMA. + game_targets[game_terminals] = game_terminal_values[game_terminals] + targets[in_game] = game_targets return self.positions, targets @@ -190,11 +201,18 @@ class SelfPlayRecords: class Net(Module): def __init__(self, inp_dim: int, hidden: int) -> None: super().__init__() + # learned 7-way cell-state embedding (empty, {black, white} x + # {flat, standing, road}); the slot layout must match nn.c + self.embedding = Embedding(7, 1) 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)) + # x: [n, 3 + n_cells]; the first 3 columns are the fixed head + # inputs, the rest are raw cell states (0..6) + n = x.shape[0] + inp = torch.cat([x[:, :3], self.embedding(x[:, 3:].long()).view(n, -1)], dim=1) + hidden = torch.relu(self.hidden(inp)) return torch.tanh(self.output(hidden)).squeeze(-1) def fit( @@ -272,6 +290,7 @@ class Net(Module): prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six") upper = "FIVE" if size == 5 else "SIX" + E = self.embedding.weight.detach().cpu().numpy() W = self.hidden.weight.detach().cpu().numpy() B = self.hidden.bias.detach().cpu().numpy() LW = self.output.weight.detach().cpu().numpy().ravel() @@ -283,6 +302,10 @@ class Net(Module): _ = f.write(f"/* {p} */\n") _ = f.write(f'#include "weights_{size}.h"\n\n') _ = f.write( + f"const float {prefix}_embedding[7] =\n" + f"{{\n{_fmt_flat(E, per_line=7)}\n}};\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" ) @@ -340,7 +363,6 @@ class TournamentResult: champ: int cand: int draws: int - cand_rate: float def run_matches( @@ -365,14 +387,7 @@ def run_matches( 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, - ) + return TournamentResult(champ=champ, cand=cand, draws=draws) # =================================================================== @@ -420,10 +435,10 @@ def run_step( print(f"\n== iteration {iter} ==") - # 1. self-play with the installed champion (include/weights_<size>.c) + # 1. rebuild selfplay _ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True) - # 2. generate self-play data (one game per process, in parallel) + # 2. generate self-play data eps = max(eps_min, eps_start * (eps_decay ** (iter - 1))) csv = make_csv(iter) print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs") @@ -446,7 +461,7 @@ def run_step( # 3. fine-tune the champion on TD(lambda) targets -> candidate device = next(champ.parameters()).device csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)] - print("Loading data from", csvs) + print("Loading data from", [str(c) for c in csvs]) records = SelfPlayRecords.load_merge(csvs, size, device) cand = copy.deepcopy(champ) @@ -465,7 +480,7 @@ def run_step( cand=True, ) - # 4. build the match driver (champion + candidate linked) + # 4. build the match driver (champion + candidate) _ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True) # 5. tournament + decision @@ -480,8 +495,7 @@ def run_step( 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}" + f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, winrate={rate:.3f}" ) promote = decided >= int(0.25 * tournament_games) and rate >= promote_rate @@ -489,7 +503,7 @@ def run_step( cand.to_c( size, root / "include" / f"weights_{size}.c", - (f"v{iter:03d}, val_corr={corr:.4f}",), + (f"v{iter:03d}, val_corr={corr:.4f}, rate={rate:.4f}",), ) print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==") return cand @@ -513,23 +527,23 @@ def main() -> None: ) _ = 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" + "--games", type=int, default=100, help="selfplay games per iteration" ) _ = p.add_argument( "--replay-iters", type=int, - default=4, + default=0, help="number of previous self-play instances to mix in with current dataset during training", ) _ = 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", + help="initial self-play epsilon: in asymmetric self-play one side (chosen at random per game) plays a uniform random move with this probability while the greedy side has its plies recorded", + ) + _ = p.add_argument( + "--eps-min", type=float, default=0.05, help="self-play epsilon floor" ) - _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor") _ = p.add_argument( "--eps-decay", type=float, default=0.9, help="per-iteration epsilon decay" ) @@ -548,20 +562,20 @@ def main() -> None: _ = p.add_argument( "--tournament-eps", type=float, - default=0.05, + default=0.01, help="tournament epsilon (random-move probability)", ) _ = p.add_argument( "--tournament-games", type=int, default=None, - help="tournament games (10/tour_eps if unset)", + help="tournament games (1/tournament-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)", + default=0.0, + help="candidate win rate among non-drawn tournament games needed to promote", ) _ = p.add_argument( "--procs", @@ -582,7 +596,7 @@ def main() -> None: args = p.parse_args() if args.tournament_games is None: - args.tournament_games = max(20, int(round(10 / args.tournament_eps))) + args.tournament_games = max(100, int(round(1 / args.tournament_eps))) if args.procs is None: args.procs = max(1, (os.cpu_count() or 1) // 2) |