File size: 2,554 Bytes
7e3d211 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Fully-convolutional autoencoder over OpenFront tile states.
Input encoding per tile:
- owner slot -> learned embedding (static slot per player for a whole game,
slot 0 = unowned, slots 1..MAX_SLOTS-1 assigned by smallID order at spawn)
- terrain scalars: land flag, normalized magnitude, fallout flag
The encoder downsamples by 16x into a spatial latent grid (LATENT_C channels
per 16x16 tile region), so any map size divisible by 16 works. The decoder
reconstructs per-tile owner-slot logits.
"""
import torch
import torch.nn as nn
MAX_SLOTS = 128 # owner classes: 0 = unowned, 1..127 player slots
OWNER_EMB_DIM = 8
TERRAIN_CHANNELS = 3 # land, magnitude, fallout
def conv_block(c_in: int, c_out: int, stride: int) -> nn.Sequential:
return nn.Sequential(
nn.Conv2d(c_in, c_out, kernel_size=3, stride=stride, padding=1),
nn.GroupNorm(8, c_out),
nn.SiLU(),
)
def deconv_block(c_in: int, c_out: int) -> nn.Sequential:
return nn.Sequential(
nn.ConvTranspose2d(c_in, c_out, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(8, c_out),
nn.SiLU(),
)
class TileAutoencoder(nn.Module):
def __init__(self, latent_c: int = 64):
super().__init__()
self.owner_emb = nn.Embedding(MAX_SLOTS, OWNER_EMB_DIM)
c_in = OWNER_EMB_DIM + TERRAIN_CHANNELS
self.encoder = nn.Sequential(
conv_block(c_in, 32, stride=1),
conv_block(32, 64, stride=2), # /2
conv_block(64, 96, stride=2), # /4
conv_block(96, 128, stride=2), # /8
conv_block(128, 128, stride=2), # /16
nn.Conv2d(128, latent_c, kernel_size=1),
)
self.decoder = nn.Sequential(
conv_block(latent_c, 128, stride=1),
deconv_block(128, 128), # /8
deconv_block(128, 96), # /4
deconv_block(96, 64), # /2
deconv_block(64, 32), # /1
nn.Conv2d(32, MAX_SLOTS, kernel_size=1),
)
def encode(self, owners: torch.Tensor, terrain: torch.Tensor) -> torch.Tensor:
"""owners: (B, H, W) int64 slots; terrain: (B, 3, H, W) float."""
emb = self.owner_emb(owners).permute(0, 3, 1, 2) # (B, E, H, W)
x = torch.cat([emb, terrain], dim=1)
return self.encoder(x)
def forward(
self, owners: torch.Tensor, terrain: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
z = self.encode(owners, terrain)
logits = self.decoder(z) # (B, MAX_SLOTS, H, W)
return logits, z
|