default engines to depth 4, add a rolling window to nn_train.py

I've noticed high variability in training efficacy, part of that i think is due to the fact that the dataset per iteration is not really that large. Sure, it has many rows, but it's fundamentally given by the same engine playing itself with some entropy injection. To remedy this i've expanded the input data for any given iteration to contain --replay-iters many previous datasets as input too. Alongside that decreased the epsilon decay rate for slower "cooling", and included much more of the early game in the training set (min ply 8 -> 4).

3 files changed, 72 insertions, 42 deletions

diff --git a/src/ctlm.c b/src/ctlm.c
index c9844c1..d2b59c9 100644
--- a/src/ctlm.c
+++ b/src/ctlm.c
@@ -487,7 +487,7 @@ int main(int argc, char **argv) {
 
   puts(license);
 
-  negamax_search_depth = 5;
+  negamax_search_depth = 4;
   new_game(5);
   negamax_init(5);
 
diff --git a/src/geminict.c b/src/geminict.c
index 2e94e7d..48d4380 100644
--- a/src/geminict.c
+++ b/src/geminict.c
@@ -302,18 +302,22 @@ static enum TURN_RESULT negamax_turn(void) {
 }
 
 int main(int argc, char **argv) {
-  if (argc != 2) {
-    puts("Usage: PTN1.[PTN2.][PTN3.] etc");
+  if (argc != 3) {
+    puts("Usage: <depth> PTN1.[PTN2.][PTN3.] etc");
     return EXIT_FAILURE;
   };
 
-  negamax_search_depth = 5;
+  negamax_search_depth = atoi(argv[1]);
+  if (negamax_search_depth == 0 || negamax_search_depth > 6) {
+    puts("Depth must be an integer in [1, 6]");
+    return EXIT_FAILURE;
+  };
   new_game(5);
   negamax_init(5);
 
   enum TURN_RESULT tr;
 
-  char *line = argv[1];
+  char *line = argv[2];
   // Some maximum length we're willing to parse
   uint32_t len = strnlen(line, 65535);
   if (!line || len < 2 || len == 65535) {
@@ -365,7 +369,7 @@ int main(int argc, char **argv) {
   tr = negamax_turn();
 
   if (tr == T_ERR) {
-    puts("This shouldn't happen, ct1986 encountered an error.");
+    puts("This shouldn't happen, geminict encountered an error.");
     print_everything();
     return EXIT_FAILURE;
   } else if (tr == T_WIN)
diff --git a/src/nn_train.py b/src/nn_train.py
index 52549d8..8025b68 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -152,6 +152,18 @@ class SelfPlayRecords:
             ),
         )
 
+    @classmethod
+    def load_merge(
+        cls, files: list[Path], size: int, device: torch.device
+    ) -> SelfPlayRecords:
+        sprs = [SelfPlayRecords.load_selfplay(file, size, device) for file in files]
+        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]),
+        )
+
     def make_training_data(
         self, model: Module, td_lambda: float
     ) -> tuple[Tensor, Tensor]:
@@ -382,26 +394,30 @@ def setup(size: int, root: Path) -> Net:
 
 def run_step(
     *,
-    iter: int,
-    size: int,
-    eps_min: float,
+    batch: int,
     champ: Net,
-    eps_start: float,
-    eps_decay: float,
     depth: int,
-    min_ply: int,
+    epochs: int,
+    eps_decay: float,
+    eps_min: float,
+    eps_start: float,
     games: int,
+    iter: int,
+    lr: float,
+    min_ply: int,
     procs: int,
+    promote_rate: float,
+    replay_iters: int,
+    root: Path,
+    size: int,
     td_lambda: float,
-    val_frac: float,
-    epochs: int,
-    batch: int,
-    lr: float,
     tournament_eps: float,
     tournament_games: int,
-    promote_rate: float,
-    root: Path,
+    val_frac: float,
 ) -> Net:
+    def make_csv(i: int) -> Path:
+        return root / "data" / f"selfplay-{size}-v{i:03d}.csv"
+
     print(f"\n== iteration {iter} ==")
 
     # 1. self-play with the installed champion (include/weights_<size>.c)
@@ -409,7 +425,7 @@ def run_step(
 
     # 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"
+    csv = make_csv(iter)
     print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs")
     out_rows: list[str] = []
     for out in fan_out(
@@ -429,7 +445,9 @@ def run_step(
 
     # 3. fine-tune the champion on TD(lambda) targets -> candidate
     device = next(champ.parameters()).device
-    records = SelfPlayRecords.load_selfplay(csv, size, device)
+    csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)]
+    print("Loading data from", csvs)
+    records = SelfPlayRecords.load_merge(csvs, size, device)
     cand = copy.deepcopy(champ)
 
     corr, mse = cand.fit(
@@ -498,6 +516,12 @@ def main() -> None:
         "--games", type=int, default=2000, help="selfplay games per iteration"
     )
     _ = p.add_argument(
+        "--replay-iters",
+        type=int,
+        default=4,
+        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,
@@ -507,16 +531,7 @@ def main() -> None:
     )
     _ = 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"
+        "--eps-decay", type=float, default=0.9, help="per-iteration epsilon decay"
     )
     _ = p.add_argument("--td-lambda", type=float, default=0.6, help="TD(lambda) lambda")
     _ = p.add_argument(
@@ -554,6 +569,16 @@ def main() -> None:
         default=None,
         help="parallel game processes (half the CPUs if unset)",
     )
+    _ = p.add_argument(
+        "--depth",
+        type=int,
+        default=4,
+        help="search depth for selfplay and the tournament",
+    )
+    _ = p.add_argument(
+        "--min-ply", type=int, default=4, help="record positions at this ply or later"
+    )
+
     args = p.parse_args()
 
     if args.tournament_games is None:
@@ -565,25 +590,26 @@ def main() -> None:
     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,
+            batch=args.batch,
             champ=champ,
-            eps_start=args.eps_start,
-            eps_decay=args.eps_decay,
             depth=args.depth,
-            min_ply=args.min_ply,
+            epochs=args.epochs,
+            eps_decay=args.eps_decay,
+            eps_min=args.eps_min,
+            eps_start=args.eps_start,
             games=args.games,
+            iter=iter,
+            lr=args.lr,
+            min_ply=args.min_ply,
             procs=args.procs,
+            promote_rate=args.promote,
+            replay_iters=args.replay_iters,
+            root=root,
+            size=args.size,
             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,
+            val_frac=args.val_frac,
         )
     print(f"\n== done, champion in include/weights_{args.size}.c ==")