add stuff, dunno what this was

python/requirements.txt

@@ -15,3 +15,7 @@ wandb
 tqdm
 shed
 zstandard==0.19.0
+plotly
+pandas
+jupyterlab==3.5.*
+jupyter==1.0.*

python/scripts/train_4x4.py

@@ -1,5 +1,6 @@
 from tak.alphazero import cli, hooks, trainer, schedule
 import os.path
+import yaml
 import shlex
 
 

python/tak/alphazero/trainer.py

@@ -111,6 +111,8 @@ def dedup_batch(batch):
             out[k][idx] += batch[k][i]
 
     for k in keys:
+        if not out[k].dtype.is_floating_point:
+            out[k] = out[k].float()
         out[k] /= counts.reshape((-1,) + (1,) * (len(out[k].shape) - 1))
     return {k: v[:next] for (k, v) in out.items()}
 

python/tak/self_play.py

@@ -245,11 +245,13 @@ def encode_games(logs: list[Transcript]):
     all_values = [v for tr in logs for v in tr.values]
     all_move_probs = torch.cat([tr.logits for tr in logs])
     all_results = [r for tr in logs for r in tr.results]
+    all_plies = [p.ply for tr in logs for p in tr.positions]
     encoded, mask = encoding.encode_batch(all_positions)
     return dict(
         positions=encoded,
         mask=mask,
         moves=all_move_probs,
         values=torch.tensor(all_values),
         results=torch.tensor(all_results, dtype=torch.float32),
+        plies=torch.tensor(all_plies, dtype=torch.int),
     )