156 lines
6 KiB
Python
156 lines
6 KiB
Python
"""Policy-value network and observation encoding for Generals.io PPO."""
|
|
|
|
import equinox as eqx
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import jax.random as jrandom
|
|
|
|
|
|
def obs_to_tensor(obs) -> 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 value(self, obs):
|
|
return self._value_from_features(self._features(obs))
|
|
|
|
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])
|
|
|
|
def __call__(self, obs, mask, key, action=None):
|
|
"""
|
|
Args:
|
|
obs: (C, H, W) tensor.
|
|
mask: (H, W, 4) valid-move mask.
|
|
key: PRNG key (only used when action is None).
|
|
action: if given, evaluate logprob of this [pass,row,col,dir,split];
|
|
otherwise sample an action.
|
|
|
|
Returns:
|
|
(action, value, logprob, entropy)
|
|
"""
|
|
logits, value = self.policy_value(obs, mask)
|
|
|
|
H, W = mask.shape[:2]
|
|
cells = H * W
|
|
|
|
if action is None:
|
|
idx = jrandom.categorical(key, logits)
|
|
else:
|
|
is_pass, row, col, direction, is_half = action
|
|
encoded_dir = jnp.where(is_half > 0, direction + 4, direction)
|
|
move_idx = encoded_dir * cells + row * W + col
|
|
idx = jnp.where(is_pass > 0, 8 * cells, move_idx)
|
|
|
|
log_probs = jax.nn.log_softmax(logits)
|
|
logprob = log_probs[idx]
|
|
probs = jax.nn.softmax(logits)
|
|
entropy = -jnp.sum(probs * log_probs)
|
|
|
|
if action is None:
|
|
is_pass = idx == 8 * cells
|
|
move_idx = jnp.minimum(idx, 8 * cells - 1)
|
|
direction = move_idx // cells
|
|
position = move_idx % cells
|
|
row = jnp.where(is_pass, 0, position // W)
|
|
col = jnp.where(is_pass, 0, position % W)
|
|
is_half = (~is_pass) & (direction >= 4)
|
|
actual_dir = jnp.where(is_pass, 0, jnp.where(is_half, direction - 4, direction))
|
|
action = jnp.array(
|
|
[is_pass.astype(jnp.int32), row, col, actual_dir, is_half.astype(jnp.int32)],
|
|
dtype=jnp.int32,
|
|
)
|
|
|
|
return action, value, logprob, entropy
|