TD-leaf and another negamax bug

TD-lambda and TD-leaf are separated by the observation that minimax with a parameterised leaf heuristic is differentiable* in the parameters. "All" that had to be done was to extract the "principal variation": which one position survived all of the pruning to become the reported value.

This was a bit fiddly and uncovered a subtle bug (again) in negamax.c: we were not being complete in our overriding of the TT entry and by omitting key it was possible that later searches would find incorrect values.

The more interesting question to me was: how do we introduce enough variety while not corrupting the PV chain? I settled on having selfplay have one of the players be epsilon-random but the other greedy (minimax) and simply report only the greedy portion PV chain. This appears to work well enough in practice and doesn't introduce nonsense training data for the net (we don't force its actions to be nonsensical). I also deciced, perhaps in a departure from tradition, that we had good supervised labels for PVs that happened to be winning states, so the training overrides the targets in that case. Finally rolled back some of the ambition in nn_train.py default arguments.

*almost always

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)