working towards hybrid rootstrap

1 files changed, 27 insertions, 11 deletions

diff --git a/src/selfplay.c b/src/selfplay.c
index 1d6bd84..b2a5821 100644
--- a/src/selfplay.c
+++ b/src/selfplay.c
@@ -91,10 +91,20 @@ static int8_t win_value(const enum WIN_TYPE w) {
   }
 }
 
+static void emit_board(const float *head, const uint8_t *cells,
+                       const uint8_t n_cells) {
+  fprintf(stdout, "\"head\":[%.8f,%.8f],\"cells\":[", head[0], head[1]);
+  for (uint8_t c = 0; c < n_cells; c++)
+    fprintf(stdout, "%s%u", c ? "," : "", cells[c]);
+  fprintf(stdout, "]");
+}
+
 static int play_game(const uint8_t min_ply, const nn_function fn_1,
                      const nn_function fn_2, const uint32_t cand_id_1,
                      const uint32_t cand_id_2,
-                     const uint32_t game_id, char record[][16], uint8_t replay) {
+                     char record[][16], uint8_t replay) {
+  float root_head[MAX_PLY][2];
+  uint8_t root_cells[MAX_PLY][MAX_CELLS];
   float row_head[MAX_PLY][2];
   uint8_t row_cells[MAX_PLY][MAX_CELLS];
   uint8_t row_term[MAX_PLY];
@@ -133,6 +143,10 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
       row_parity_sign[n_rows] = leaf->parity_sign;
       row_cand_id[n_rows] = (state->current_colour&1) ? cand_id_1 : cand_id_2;
       row_mover_sign[n_rows] = (state->current_colour&1) ? 1 : -1;
+
+      // Net input of the root state s_t, before the move is played
+      nn_prepare_head(root_head[n_rows], state);
+      nn_cell_idxs(root_cells[n_rows], state);
       n_rows++;
     }
 
@@ -155,13 +169,17 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
   // final outcome of the played game, from black's perspective
   const int8_t R = win_value(state->won);
 
+  fprintf(stdout, "{\"result\":%d,\"rows\":[", R);
   for (uint32_t k = 0; k < n_rows; k++) {
-    fprintf(stdout, "%.1f,%.8f,%.8f,", row_parity_sign[k], row_head[k][0], row_head[k][1]);
-    for (uint8_t c = 0; c < n_cells; c++)
-      fprintf(stdout, "%u,", row_cells[k][c]);
-    fprintf(stdout, "%i,%u,%.0f,%u,%u,%i\n", R, row_term[k], row_terminal_value[k], game_id, row_cand_id[k],
-            row_mover_sign[k] * R);
+    fprintf(stdout, "%s{\"cand\":%u,\"root\":{\"parity\":%d,", k ? "," : "",
+            row_cand_id[k], row_mover_sign[k]);
+    emit_board(root_head[k], root_cells[k], n_cells);
+    fprintf(stdout, "},\"leaf\":{\"parity\":%.0f,", row_parity_sign[k]);
+    emit_board(row_head[k], row_cells[k], n_cells);
+    fprintf(stdout, ",\"terminal\":%u,\"terminal_value\":%.1f}}",
+            row_term[k], row_terminal_value[k]);
   }
+  fprintf(stdout, "]}\n");
 
   return (int)n_rows;
 }
@@ -169,9 +187,9 @@ static int play_game(const uint8_t min_ply, const nn_function fn_1,
 const char *usage =
   "selfplay, one game of training data between candidates.\n\n\
 Usage: selfplay <size> <depth> <min_ply>\n\n\
-Plays a single game of <size>x<size>, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one CSV row per recorded search, from ply <min_ply> onwards:\n\n\
-\tparity_sign,mover_flats,opponent_flats,<cell ids...>,R,terminal,leaf_R,game_id,cand_id,cand_R\n\n\
-R is the outcome from black's perspective, cand_R the outcome from the recorded candidate's perspective. nn_train.py runs many of these in parallel to generate training data.\n\n\
+Plays a single game of <size>x<size>, with <min_ply> moves made randomly. Only the deterministic plies are written to stdout, one JSON record per game, from ply <min_ply> onwards:\n\n\
+\t{\"result\": R, \"rows\": [{\"cand\": id, \"root\": {\"parity\", \"head\", \"cells\"}, \"leaf\": {\"parity\", \"head\", \"cells\", \"terminal\", \"terminal_value\"}}]}\n\n\
+R is the outcome of the game from black's perspective. Each row records the net input of the root state s_t and the PV leaf selected at that ply; the leaf value Net(pv) is taken from the trained net. nn_train.py runs many of these in parallel to generate training data.\n\n\
 Copyright (C) 2026, tslil clingman\n\n\
 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 the Free Software Foundation, either version 3 of the License, or (at your option) any later version.";
 
@@ -212,13 +230,11 @@ int main(int argc, char **argv) {
             size == 5 ? FIVE_A : SIX_A, size == 5 ? FIVE_B : SIX_B,
             CAND_A,
             CAND_B,
-            rand32(),
             record, 0);
   play_game(min_ply,
             size == 5 ? FIVE_B : SIX_B, size == 5 ? FIVE_A : SIX_A,
             CAND_B,
             CAND_A,
-            rand32(),
             record, CAND_A == CAND_B ? 0: 1);