checkpoint

scripts/mcts_agent.py

@@ -0,0 +1,89 @@
+"""Competition stdio agent using particle PUCT and a trained checkpoint."""
+
+import argparse
+import os
+import sys
+import time
+
+import equinox as eqx
+import jax.random as jrandom
+
+from general_bots_training.mcts import MCTSConfig, ParticlePUCT, observation_from_wire
+from general_bots_training.network import PolicyValueNetwork
+
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+
+
+def _read_grid(stream, height: int):
+    return [[int(value) for value in stream.readline().split()] for _ in range(height)]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("checkpoint", help="Equinox policy checkpoint")
+    parser.add_argument("--time-budget-ms", type=float, default=125.0)
+    parser.add_argument("--max-simulations", type=int, default=128)
+    parser.add_argument("--rollout-depth", type=int, default=2)
+    parser.add_argument("--top-k", type=int, default=20)
+    parser.add_argument("--seed", type=int, default=0)
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    handshake = sys.stdin.readline()
+    if not handshake:
+        return
+    player_index, height, width = (int(value) for value in handshake.split())
+
+    network = PolicyValueNetwork(jrandom.PRNGKey(args.seed))
+    network = eqx.tree_deserialise_leaves(args.checkpoint, network)
+    search = ParticlePUCT(
+        network,
+        player_index,
+        MCTSConfig(
+            time_budget_ms=args.time_budget_ms,
+            max_simulations=args.max_simulations,
+            rollout_depth=args.rollout_depth,
+            top_k=args.top_k,
+        ),
+        seed=args.seed,
+    )
+    search.warmup(height, width)
+    print(f"[mcts] warmup complete for {height}x{width}", file=sys.stderr, flush=True)
+
+    while True:
+        scalar_line = sys.stdin.readline()
+        if not scalar_line:
+            return
+        timestep, own_land, own_army, opponent_land, opponent_army = (
+            int(value) for value in scalar_line.split()
+        )
+        type_grid = _read_grid(sys.stdin, height)
+        owner_grid = _read_grid(sys.stdin, height)
+        army_grid = _read_grid(sys.stdin, height)
+        observation = observation_from_wire(
+            timestep,
+            own_land,
+            own_army,
+            opponent_land,
+            opponent_army,
+            type_grid,
+            owner_grid,
+            army_grid,
+        )
+
+        started_at = time.perf_counter()
+        action, stats = search.search(observation)
+        elapsed_ms = (time.perf_counter() - started_at) * 1000
+        print(
+            f"[mcts] turn={timestep} simulations={int(stats['simulations'])} "
+            f"elapsed_ms={elapsed_ms:.1f}",
+            file=sys.stderr,
+            flush=True,
+        )
+        print(" ".join(str(int(value)) for value in action), flush=True)
+
+
+if __name__ == "__main__":
+    main()