mg/general-bots-training
Watch
1
0
Fork
You've already forked general-bots-training
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

checkpoint

Browse source
This commit is contained in:
Moritz Gmeiner 2026-08-08 21:25:23 +02:00
commit a69868e950
23 changed files with 1269 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

336
competition/competition.py Normal file
View file

@ -0,0 +1,336 @@
"""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()