336 lines
12 KiB
Python
336 lines
12 KiB
Python
"""Competition stdio agent: direct policy sampling over a trained equinox checkpoint.
|
|
|
|
Self-contained competition bot. It speaks the wire protocol in
|
|
`generals-bots/competition/protocol.py` (handshake, then one observation frame
|
|
per turn, one action line per turn, EOF on stdin = game over) and, on each
|
|
turn, samples an action from the policy network exactly as the model is used
|
|
during training/evaluation — no tree search.
|
|
|
|
Unlike `scripts/mcts_agent.py`, this module does **not** import the
|
|
`general_bots_training` training package. It inlines the policy-value network
|
|
and the observation encoding so the script can be dropped into the competition
|
|
sandbox with just the engine (`generals`), `equinox`, `jax`, and `numpy`
|
|
available.
|
|
|
|
This is the search-free counterpart to `scripts/competition_puct.py`: it shares
|
|
the same network and observation encoding, but selects an action by sampling
|
|
from the masked policy (as in `PolicyValueNetwork.__call__`) instead of running
|
|
particle PUCT.
|
|
|
|
Run it directly through the bundled matchup driver, e.g.:
|
|
|
|
PYTHONPATH=src:generals-bots .venv/bin/python \
|
|
generals-bots/competition/matchup.py \
|
|
scripts/competition.py --checkpoint ppo_model.eqx \
|
|
generals-bots/competition/agents/expander_python/run.sh \
|
|
--mode competition
|
|
|
|
or wrap it in a `run.sh` that passes the checkpoint path.
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import equinox as eqx
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import jax.random as jrandom
|
|
import numpy as np
|
|
from generals.core import game
|
|
from generals.core.action import compute_valid_move_mask
|
|
from generals.core.game import GameState
|
|
from generals.core.observation import Observation
|
|
from generals.modifiers import build_castles, deathtouch
|
|
|
|
os.environ.setdefault("JAX_PLATFORMS", "cpu")
|
|
|
|
PASS_ACTION = np.array([1, 0, 0, 0, 0], dtype=np.int32)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Policy-value network (inlined from general_bots_training.network)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def obs_to_tensor(obs: Observation) -> jnp.ndarray:
|
|
"""Encode an Observation into a (C, H, W) float32 tensor for the network.
|
|
|
|
Armies and army-counts are log-normalized; land counts are divided by the
|
|
number of cells; scalar values are broadcast to spatial planes so a plain
|
|
conv stack can consume them. Grid-size agnostic.
|
|
"""
|
|
H, W = obs.armies.shape
|
|
armies = jnp.log1p(obs.armies.astype(jnp.float32)) / jnp.log(50.0)
|
|
|
|
def bcast(scalar):
|
|
return jnp.broadcast_to(scalar.astype(jnp.float32), (H, W))
|
|
|
|
own_land = bcast(obs.owned_land_count / (H * W))
|
|
own_army = bcast(jnp.log1p(obs.owned_army_count.astype(jnp.float32)) / jnp.log(50.0))
|
|
opp_land = bcast(obs.opponent_land_count / (H * W))
|
|
opp_army = bcast(jnp.log1p(obs.opponent_army_count.astype(jnp.float32)) / jnp.log(50.0))
|
|
timestep = bcast(jnp.log1p(obs.timestep.astype(jnp.float32)) / jnp.log(1201.0))
|
|
|
|
return jnp.stack(
|
|
[
|
|
armies,
|
|
obs.generals.astype(jnp.float32),
|
|
obs.castles.astype(jnp.float32),
|
|
obs.mountains.astype(jnp.float32),
|
|
obs.neutral_cells.astype(jnp.float32),
|
|
obs.owned_cells.astype(jnp.float32),
|
|
obs.opponent_cells.astype(jnp.float32),
|
|
obs.fog_cells.astype(jnp.float32),
|
|
obs.structures_in_fog.astype(jnp.float32),
|
|
own_land,
|
|
own_army,
|
|
opp_land,
|
|
opp_army,
|
|
timestep,
|
|
],
|
|
axis=0,
|
|
)
|
|
|
|
|
|
class PolicyValueNetwork(eqx.Module):
|
|
"""Conv policy-value network.
|
|
|
|
Action layout: 4 full-move directions and 4 half-move (split) directions
|
|
per source cell, plus one global pass action. Invalid moves are masked to
|
|
-1e9; pass is always available.
|
|
|
|
The value head uses global average pooling over the spatial grid, so the
|
|
network works for any board size without reshaping linear layers.
|
|
"""
|
|
|
|
conv1: eqx.nn.Conv2d
|
|
conv2: eqx.nn.Conv2d
|
|
conv3: eqx.nn.Conv2d
|
|
conv4: eqx.nn.Conv2d
|
|
policy_conv: eqx.nn.Conv2d
|
|
value_conv: eqx.nn.Conv2d
|
|
value_linear1: eqx.nn.Linear
|
|
value_linear2: eqx.nn.Linear
|
|
|
|
def __init__(self, key, in_channels: int = 14, channels=(32, 32, 32, 16)):
|
|
keys = jrandom.split(key, 8)
|
|
self.conv1 = eqx.nn.Conv2d(in_channels, channels[0], kernel_size=3, padding=1, key=keys[0])
|
|
self.conv2 = eqx.nn.Conv2d(channels[0], channels[1], kernel_size=3, padding=1, key=keys[1])
|
|
self.conv3 = eqx.nn.Conv2d(channels[1], channels[2], kernel_size=3, padding=1, key=keys[2])
|
|
self.conv4 = eqx.nn.Conv2d(channels[2], channels[3], kernel_size=3, padding=1, key=keys[3])
|
|
# 9 = 4 dirs (full) + 4 dirs (half) + 1 pass
|
|
self.policy_conv = eqx.nn.Conv2d(channels[3], 9, kernel_size=1, key=keys[4])
|
|
self.value_conv = eqx.nn.Conv2d(channels[3], 4, kernel_size=1, key=keys[5])
|
|
self.value_linear1 = eqx.nn.Linear(4, 64, key=keys[6])
|
|
self.value_linear2 = eqx.nn.Linear(64, 1, key=keys[7])
|
|
|
|
def _features(self, obs):
|
|
x = jax.nn.relu(self.conv1(obs))
|
|
x = jax.nn.relu(self.conv2(x))
|
|
x = jax.nn.relu(self.conv3(x))
|
|
x = jax.nn.relu(self.conv4(x))
|
|
return x
|
|
|
|
def _value_from_features(self, feat):
|
|
v = jax.nn.relu(self.value_conv(feat)) # (4, H, W)
|
|
v = v.mean(axis=(1, 2)) # (4,) global average pool
|
|
v = jax.nn.relu(self.value_linear1(v))
|
|
return self.value_linear2(v)[0]
|
|
|
|
def policy_value(self, obs, mask):
|
|
"""Return deterministic masked policy logits and the critic value."""
|
|
features = self._features(obs)
|
|
return self._policy_logits(features, mask), self._value_from_features(features)
|
|
|
|
def _policy_logits(self, feat, mask):
|
|
"""Return 8*H*W move logits plus one global pass logit."""
|
|
logits = self.policy_conv(feat) # (9, H, W)
|
|
mask_t = jnp.transpose(mask, (2, 0, 1)) # (4, H, W)
|
|
penalty = (1.0 - mask_t) * -1e9
|
|
move_penalty = jnp.concatenate([penalty, penalty], axis=0)
|
|
move_logits = (logits[:8] + move_penalty).reshape(-1)
|
|
pass_logit = jnp.mean(logits[8])[None]
|
|
return jnp.concatenate([move_logits, pass_logit])
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Action selection (matches training/evaluation rollout)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@jax.jit
|
|
def competition_move_mask(obs: Observation) -> jnp.ndarray:
|
|
"""Legal move mask, identical to the one used during training/evaluation."""
|
|
return compute_valid_move_mask(obs.armies, obs.owned_cells, obs.mountains)
|
|
|
|
|
|
def _decode_policy_index(index: int, height: int, width: int) -> np.ndarray:
|
|
cells = height * width
|
|
if index == 8 * cells:
|
|
return PASS_ACTION.copy()
|
|
encoded_direction, position = divmod(index, cells)
|
|
row, col = divmod(position, width)
|
|
split = int(encoded_direction >= 4)
|
|
direction = encoded_direction - 4 if split else encoded_direction
|
|
return np.array([0, row, col, direction, split], dtype=np.int32)
|
|
|
|
|
|
def observation_from_wire(
|
|
timestep: int,
|
|
owned_land_count: int,
|
|
owned_army_count: int,
|
|
opponent_land_count: int,
|
|
opponent_army_count: int,
|
|
type_grid,
|
|
owner_grid,
|
|
army_grid,
|
|
) -> Observation:
|
|
"""Convert a competition wire frame into the network's Observation type."""
|
|
types = jnp.asarray(type_grid, dtype=jnp.int32)
|
|
owners = jnp.asarray(owner_grid, dtype=jnp.int32)
|
|
armies = jnp.asarray(army_grid, dtype=jnp.int32)
|
|
visible = (types != 0) & (types != 5)
|
|
return Observation(
|
|
armies=armies,
|
|
generals=types == 4,
|
|
castles=types == 3,
|
|
mountains=types == 2,
|
|
neutral_cells=visible & (owners == 0) & (types != 2),
|
|
owned_cells=owners == 1,
|
|
opponent_cells=owners == 2,
|
|
fog_cells=types == 0,
|
|
structures_in_fog=types == 5,
|
|
owned_land_count=jnp.int32(owned_land_count),
|
|
owned_army_count=jnp.int32(owned_army_count),
|
|
opponent_land_count=jnp.int32(opponent_land_count),
|
|
opponent_army_count=jnp.int32(opponent_army_count),
|
|
timestep=jnp.int32(timestep),
|
|
)
|
|
|
|
|
|
@jax.jit
|
|
def competition_step(state: GameState, actions: jnp.ndarray):
|
|
state, actions = build_castles.apply_build_actions(state, actions)
|
|
return deathtouch.step(state, actions, turn=800)
|
|
|
|
|
|
class PolicyAgent:
|
|
"""Direct policy agent that samples actions as in training/evaluation.
|
|
|
|
On each turn it computes the masked policy logits over all legal moves and
|
|
samples one action with `jrandom.categorical`, exactly mirroring
|
|
`PolicyValueNetwork.__call__` used during rollout collection.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
network: PolicyValueNetwork,
|
|
seed: int = 0,
|
|
):
|
|
self.network = network
|
|
self.key = jrandom.PRNGKey(seed)
|
|
self._infer = eqx.filter_jit(
|
|
lambda network, tensor, mask: network.policy_value(tensor, mask)
|
|
)
|
|
|
|
def warmup(self, height: int, width: int) -> None:
|
|
"""Compile shape-dependent network and competition transition kernels."""
|
|
grid = jnp.zeros((height, width), dtype=jnp.int32)
|
|
grid = grid.at[0, 0].set(1).at[height - 1, width - 1].set(2)
|
|
state = game.create_initial_state(grid)
|
|
state = state._replace(
|
|
armies=state.armies.at[0, 0].set(100).at[height - 1, width - 1].set(100),
|
|
time=jnp.int32(100),
|
|
)
|
|
obs = game.get_observation(state, 0)
|
|
self.policy_value(obs)
|
|
actions = jnp.stack([jnp.asarray(PASS_ACTION), jnp.asarray(PASS_ACTION)])
|
|
warmed_state, _ = competition_step(state, actions)
|
|
jax.block_until_ready(warmed_state)
|
|
|
|
def policy_value(self, obs: Observation):
|
|
mask = competition_move_mask(obs)
|
|
logits, value = self._infer(self.network, obs_to_tensor(obs), mask)
|
|
return np.asarray(logits), float(value), mask
|
|
|
|
def select(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]:
|
|
"""Sample one action from the masked policy and report statistics."""
|
|
started_at = time.perf_counter()
|
|
logits, value, _ = self.policy_value(obs)
|
|
self.key, sample_key = jrandom.split(self.key)
|
|
index = int(jrandom.categorical(sample_key, jnp.asarray(logits)))
|
|
height, width = obs.armies.shape
|
|
action = _decode_policy_index(index, height, width)
|
|
elapsed_ms = (time.perf_counter() - started_at) * 1000
|
|
return action, {
|
|
"value": value,
|
|
"elapsed_ms": elapsed_ms,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Stdio driver
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
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("--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)
|
|
agent = PolicyAgent(network, seed=args.seed)
|
|
agent.warmup(height, width)
|
|
print(f"[competition] 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 = agent.select(observation)
|
|
elapsed_ms = (time.perf_counter() - started_at) * 1000
|
|
print(
|
|
f"[competition] turn={timestep} value={stats['value']:+.3f} "
|
|
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()
|