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
general-bots-training/scripts/competition_puct.py
Moritz Gmeiner a69868e950 checkpoint
2026-08-08 21:25:35 +02:00

552 lines
21 KiB
Python
Raw Permalink Blame History

"""Competition stdio agent: particle PUCT 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 runs
deadline-bounded root PUCT using only the perspective-relative wire
observation.
Unlike `scripts/mcts_agent.py`, this module does **not** import the
`general_bots_training` training package. It inlines the policy-value network,
the observation encoding, and the particle search so the script can be dropped
into the competition sandbox with just the engine (`generals`), `equinox`,
`jax`, and `numpy` available.
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
from dataclasses import dataclass
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)
# Search/agent defaults, previously exposed as CLI flags. Kept as constants so
# the script can be launched by a bare `run.sh` with no extra arguments.
TIME_BUDGET_MS = 125.0
MAX_SIMULATIONS = 128
ROLLOUT_DEPTH = 2
TOP_K = 20
SEED = 0
# ---------------------------------------------------------------------------
# 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])
# ---------------------------------------------------------------------------
# Particle PUCT search (inlined from general_bots_training.mcts)
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class MCTSConfig:
time_budget_ms: float = 100.0
max_simulations: int = 128
rollout_depth: int = 2
top_k: int = 20
opponent_top_k: int = 12
max_build_actions: int = 3
c_puct: float = 1.5
value_scale: float = 5.0
@jax.jit
def competition_move_mask(obs: Observation) -> jnp.ndarray:
"""Legal move mask that conservatively blocks unresolved fog structures."""
blocked = obs.mountains | obs.structures_in_fog
return compute_valid_move_mask(obs.armies, obs.owned_cells, blocked)
@jax.jit
def build_cost_grid(obs: Observation) -> jnp.ndarray:
"""Exact own-castle build costs derivable from a fog observation."""
structures = ((obs.castles | obs.generals) & obs.owned_cells).astype(jnp.int32)
height, width = structures.shape
radius = 6
padded = jnp.pad(structures, radius)
costs = jnp.full((height, width), 35, dtype=jnp.int32)
for row_offset in range(-radius, radius + 1):
for col_offset in range(-radius, radius + 1):
surcharge = 14 - 2 * (abs(row_offset) + abs(col_offset))
if surcharge > 0:
shifted = padded[
radius + row_offset : radius + row_offset + height,
radius + col_offset : radius + col_offset + width,
]
costs = costs + surcharge * shifted
return costs
@jax.jit
def valid_build_mask(obs: Observation) -> jnp.ndarray:
costs = build_cost_grid(obs)
plain_owned = obs.owned_cells & ~obs.generals & ~obs.castles
return plain_owned & (obs.armies >= costs)
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),
)
def sample_determinization(
obs: Observation, player_index: int, rng: np.random.Generator
) -> GameState:
"""Sample a conservative hidden state consistent with visible cells/totals."""
armies = np.asarray(obs.armies, dtype=np.int32).copy()
owned = np.asarray(obs.owned_cells, dtype=bool)
visible_opponent = np.asarray(obs.opponent_cells, dtype=bool)
mountains = np.asarray(obs.mountains | obs.structures_in_fog, dtype=bool)
castles = np.asarray(obs.castles, dtype=bool)
generals = np.asarray(obs.generals, dtype=bool).copy()
fog = np.asarray(obs.fog_cells, dtype=bool) & ~mountains
opponent = visible_opponent.copy()
hidden_candidates = np.argwhere(fog & ~owned)
missing_land = max(0, int(obs.opponent_land_count) - int(opponent.sum()))
if len(hidden_candidates):
chosen = hidden_candidates[
rng.choice(
len(hidden_candidates),
size=min(missing_land, len(hidden_candidates)),
replace=False,
)
]
opponent[chosen[:, 0], chosen[:, 1]] = True
visible_enemy_general = generals & opponent
if not visible_enemy_general.any():
candidates = np.argwhere((opponent | fog) & ~owned & ~mountains)
if len(candidates):
row, col = candidates[rng.integers(len(candidates))]
generals[row, col] = True
opponent[row, col] = True
hidden_opponent = opponent & ~visible_opponent
remaining_army = max(0, int(obs.opponent_army_count) - int(armies[visible_opponent].sum()))
hidden_cells = np.argwhere(hidden_opponent)
if len(hidden_cells):
base = min(remaining_army, len(hidden_cells))
armies[hidden_opponent] = 0
armies[hidden_cells[:base, 0], hidden_cells[:base, 1]] = 1
remaining_army -= base
if remaining_army:
allocations = rng.multinomial(
remaining_army, np.full(len(hidden_cells), 1 / len(hidden_cells))
)
armies[hidden_cells[:, 0], hidden_cells[:, 1]] += allocations.astype(np.int32)
ownership_relative = np.stack([owned, opponent])
if player_index == 1:
ownership = ownership_relative[::-1]
else:
ownership = ownership_relative
passable = ~mountains
neutral = passable & ~ownership[0] & ~ownership[1]
general_positions = []
for absolute_player in range(2):
positions = np.argwhere(generals & ownership[absolute_player])
if len(positions):
general_positions.append(positions[0])
else:
fallback = np.argwhere(ownership[absolute_player] & passable)
general_positions.append(fallback[0] if len(fallback) else np.array([0, 0]))
return GameState(
armies=jnp.asarray(armies),
ownership=jnp.asarray(ownership),
ownership_neutral=jnp.asarray(neutral),
generals=jnp.asarray(generals),
castles=jnp.asarray(castles),
mountains=jnp.asarray(mountains),
passable=jnp.asarray(passable),
general_positions=jnp.asarray(general_positions, dtype=jnp.int32),
time=jnp.asarray(obs.timestep, dtype=jnp.int32),
winner=jnp.int32(-1),
pool_idx=jnp.int32(0),
)
@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 ParticlePUCT:
"""Deadline-bounded root PUCT with policy-guided simultaneous rollouts."""
def __init__(
self,
network: PolicyValueNetwork,
player_index: int,
config: MCTSConfig = MCTSConfig(),
seed: int = 0,
):
self.network = network
self.player_index = player_index
self.config = config
self.rng = np.random.default_rng(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, self.player_index)
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)
candidates, _ = self.candidates(obs)
self._simulate(obs, candidates[0])
self._simulate(obs, candidates[min(1, len(candidates) - 1)])
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 candidates(self, obs: Observation, top_k: int | None = None):
logits, _, _ = self.policy_value(obs)
height, width = obs.armies.shape
valid_indices = np.flatnonzero(logits > -1e8)
count = min(top_k or self.config.top_k, len(valid_indices))
ranked = np.argsort(logits[valid_indices])[::-1][:count]
indices = valid_indices[ranked]
selected_logits = logits[indices]
selected_logits = selected_logits - selected_logits.max()
priors = np.exp(selected_logits)
actions = [_decode_policy_index(int(index), height, width) for index in indices]
build_mask = np.asarray(valid_build_mask(obs))
build_positions = np.argwhere(build_mask)
if len(build_positions):
costs = np.asarray(build_cost_grid(obs))
armies = np.asarray(obs.armies)
scores = np.array([armies[r, c] - costs[r, c] for r, c in build_positions])
order = np.argsort(scores)[::-1][: self.config.max_build_actions]
build_prior = max(float(priors.sum()) * 0.05, 1e-3)
for position_index in order:
row, col = build_positions[position_index]
actions.append(np.array([2, row, col, 0, 0], dtype=np.int32))
priors = np.append(priors, build_prior)
priors = priors / priors.sum()
return actions, priors
def _sample_policy_action(self, obs: Observation, top_k: int) -> np.ndarray:
actions, priors = self.candidates(obs, top_k)
return actions[int(self.rng.choice(len(actions), p=priors))]
def _leaf_value(self, state: GameState, info) -> float:
if bool(info.is_done):
winner = int(info.winner)
return 0.0 if winner < 0 else (1.0 if winner == self.player_index else -1.0)
obs = game.get_observation(state, self.player_index)
_, value, _ = self.policy_value(obs)
return float(np.tanh(value / self.config.value_scale))
def _simulate(self, root_obs: Observation, root_action: np.ndarray) -> float:
state = sample_determinization(root_obs, self.player_index, self.rng)
info = game.get_info(state)
our_action = root_action
for _ in range(self.config.rollout_depth):
opponent_index = 1 - self.player_index
opponent_obs = game.get_observation(state, opponent_index)
opponent_action = self._sample_policy_action(opponent_obs, self.config.opponent_top_k)
joint_actions = [None, None]
joint_actions[self.player_index] = our_action
joint_actions[opponent_index] = opponent_action
state, info = competition_step(state, jnp.asarray(joint_actions, dtype=jnp.int32))
jax.block_until_ready(state)
if bool(info.is_done):
break
our_obs = game.get_observation(state, self.player_index)
our_action = self._sample_policy_action(our_obs, self.config.top_k)
return self._leaf_value(state, info)
def search(self, obs: Observation) -> tuple[np.ndarray, dict[str, float]]:
started_at = time.perf_counter()
deadline = started_at + self.config.time_budget_ms / 1000.0
actions, priors = self.candidates(obs)
visits = np.zeros(len(actions), dtype=np.int32)
value_sums = np.zeros(len(actions), dtype=np.float64)
simulations = 0
estimated_simulation_seconds = 0.0
while (
simulations < self.config.max_simulations
and time.perf_counter() + estimated_simulation_seconds < deadline
):
total_visits = max(1, int(visits.sum()))
q_values = np.divide(
value_sums,
visits,
out=np.zeros_like(value_sums),
where=visits > 0,
)
scores = q_values + self.config.c_puct * priors * np.sqrt(total_visits) / (1 + visits)
action_index = int(np.argmax(scores))
simulation_started_at = time.perf_counter()
value = self._simulate(obs, actions[action_index])
simulation_elapsed = time.perf_counter() - simulation_started_at
estimated_simulation_seconds = max(estimated_simulation_seconds, simulation_elapsed)
visits[action_index] += 1
value_sums[action_index] += value
simulations += 1
selected = int(np.argmax(visits)) if simulations else int(np.argmax(priors))
return actions[selected], {
"simulations": float(simulations),
"selected_visits": float(visits[selected]),
"elapsed_budget_ms": self.config.time_budget_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")
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(SEED))
network = eqx.tree_deserialise_leaves(args.checkpoint, network)
search = ParticlePUCT(
network,
player_index,
MCTSConfig(
time_budget_ms=TIME_BUDGET_MS,
max_simulations=MAX_SIMULATIONS,
rollout_depth=ROLLOUT_DEPTH,
top_k=TOP_K,
),
seed=SEED,
)
search.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 = search.search(observation)
elapsed_ms = (time.perf_counter() - started_at) * 1000
print(
f"[competition] 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()
Reference in a new issue View git blame Copy permalink