89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
"""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()
|