""" ByteFight Policy Model v3: ViT board encoder + Transformer action decoder. Board encoder: 970 discrete tokens (9 scalars + 961 cells) → shared embedding → 2-layer bidirectional self-attention → mean pool → 1 vector Uses same tokenization as original alphabyte (tokenizer.py): Vocab 2462: CLS=0, stamina 1-381, position 382-412, turn 413-2413, cells 2414-2461 Actions (21): 20 action types + EOS """ import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass NUM_ACTIONS = 21 EOS_ACTION = 20 BOARD_VOCAB = 2462 # from tokenizer.py BOARD_SEQ_LEN = 970 # 9 scalars + 961 cells (no CLS) MAX_BOARD = 31 @dataclass class Config: d_model: int = 256 n_layer: int = 6 n_head: int = 8 max_seq: int = 6000 n_actions: int = NUM_ACTIONS board_vocab: int = BOARD_VOCAB board_seq_len: int = BOARD_SEQ_LEN dropout: float = 0.1 class BoardAttnBlock(nn.Module): def __init__(self, d, n_head): super().__init__() self.norm1 = nn.LayerNorm(d) self.n_head = n_head self.head_dim = d // n_head self.qkv = nn.Linear(d, 3 * d, bias=False) self.out = nn.Linear(d, d, bias=False) self.norm2 = nn.LayerNorm(d) h = d * 4 self.ffn = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d)) def forward(self, x): B, T, C = x.shape h = self.norm1(x) qkv = self.qkv(h).reshape(B, T, 3, self.n_head, self.head_dim) q, k, v = qkv.unbind(2) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) attn = F.scaled_dot_product_attention(q, k, v) x = x + self.out(attn.transpose(1, 2).reshape(B, T, C)) x = x + self.ffn(self.norm2(x)) return x class BoardEncoder(nn.Module): """ViT: 970 discrete board tokens → embedding → self-attention → 1 vector.""" def __init__(self, cfg: Config): super().__init__() self.embed = nn.Embedding(cfg.board_vocab, cfg.d_model) self.pos_embed = nn.Parameter(torch.randn(1, cfg.board_seq_len, cfg.d_model) * 0.02) self.blocks = nn.ModuleList([ BoardAttnBlock(cfg.d_model, cfg.n_head) for _ in range(2) ]) self.norm = nn.LayerNorm(cfg.d_model) def forward(self, board_tokens): """board_tokens: (B, 970) int64 → (B, d_model)""" x = self.embed(board_tokens) # (B, 970, d_model) x = x + self.pos_embed[:, :x.shape[1]] for block in self.blocks: x = block(x) return self.norm(x).mean(dim=1) # (B, d_model) class RMSNorm(nn.Module): def __init__(self, d, eps=1e-6): super().__init__() self.w = nn.Parameter(torch.ones(d)) self.eps = eps def forward(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.w class Attention(nn.Module): def __init__(self, cfg): super().__init__() self.n_head = cfg.n_head self.head_dim = cfg.d_model // cfg.n_head self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=False) self.out = nn.Linear(cfg.d_model, cfg.d_model, bias=False) self.dropout = cfg.dropout def forward(self, x, kv_cache=None): B, T, C = x.shape qkv = self.qkv(x).reshape(B, T, 3, self.n_head, self.head_dim) q, k, v = qkv.unbind(2) q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) if kv_cache is not None: k_prev, v_prev = kv_cache k = torch.cat([k_prev, k], dim=2) v = torch.cat([v_prev, v], dim=2) new_cache = (k, v) x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.dropout if self.training else 0.0, is_causal=(kv_cache is None)) return self.out(x.transpose(1, 2).reshape(B, T, C)), new_cache class FFN(nn.Module): def __init__(self, cfg): super().__init__() h = (int(cfg.d_model * 8 / 3) + 15) // 16 * 16 self.w1 = nn.Linear(cfg.d_model, h, bias=False) self.w2 = nn.Linear(h, cfg.d_model, bias=False) self.w3 = nn.Linear(cfg.d_model, h, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) class Block(nn.Module): def __init__(self, cfg): super().__init__() self.norm1 = RMSNorm(cfg.d_model) self.attn = Attention(cfg) self.norm2 = RMSNorm(cfg.d_model) self.ffn = FFN(cfg) def forward(self, x, kv_cache=None): attn_out, new_cache = self.attn(self.norm1(x), kv_cache) x = x + attn_out x = x + self.ffn(self.norm2(x)) return x, new_cache class PolicyModel(nn.Module): """16-channel CNN encoder + Transformer decoder.""" def __init__(self, cfg: Config = None): super().__init__() if cfg is None: cfg = Config() self.cfg = cfg self.encoder = BoardEncoder(cfg) self.action_embed = nn.Embedding(cfg.n_actions, cfg.d_model) self.pos_embed = nn.Embedding(cfg.max_seq, cfg.d_model) self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg.n_layer)]) self.norm = RMSNorm(cfg.d_model) self.head = nn.Linear(cfg.d_model, cfg.n_actions, bias=False) def forward(self, board_tokens, seq_actions, seq_targets, seq_is_board, board_counts): """Training forward. board_tokens: (total_boards, 970) int64.""" B, T = seq_actions.shape device = seq_actions.device board_embs = self.encoder(board_tokens) # (total_boards, d_model) act_emb = self.action_embed(seq_actions) seq_emb = act_emb.clone() board_idx = 0 for b in range(B): positions = seq_is_board[b].nonzero(as_tuple=True)[0] n_boards = board_counts[b].item() for i in range(n_boards): if i < len(positions): seq_emb[b, positions[i]] = board_embs[board_idx] board_idx += 1 pos = self.pos_embed(torch.arange(T, device=device)) x = seq_emb + pos for block in self.blocks: x, _ = block(x) logits = self.head(self.norm(x)) loss = F.cross_entropy( logits.reshape(-1, self.cfg.n_actions), seq_targets.reshape(-1), ignore_index=-100) return logits, loss @torch.no_grad() def generate(self, board_tokens, kv_caches=None, seq_pos=0, max_actions=10, temperature=0.0): """Generate actions with KV cache. board_tokens: (1, 970) int64.""" self.eval() device = board_tokens.device board_emb = self.encoder(board_tokens).unsqueeze(1) pos = self.pos_embed(torch.tensor([seq_pos], device=device)).unsqueeze(0) x = board_emb + pos if kv_caches is None: kv_caches = [None] * len(self.blocks) new_caches = [] for block, cache in zip(self.blocks, kv_caches): x, new_cache = block(x, cache) new_caches.append(new_cache) kv_caches = new_caches logits = self.head(self.norm(x)) next_logits = logits[0, -1] seq_pos += 1 if temperature <= 0: action = next_logits.argmax().item() else: action = torch.multinomial( F.softmax(next_logits / temperature, dim=-1), 1).item() actions = [] if action == EOS_ACTION: return actions, kv_caches, seq_pos actions.append(action) for _ in range(max_actions - 1): act_emb = self.action_embed(torch.tensor([[action]], device=device)) pos = self.pos_embed(torch.tensor([seq_pos], device=device)).unsqueeze(0) x = act_emb + pos new_caches = [] for block, cache in zip(self.blocks, kv_caches): x, new_cache = block(x, cache) new_caches.append(new_cache) kv_caches = new_caches logits = self.head(self.norm(x)) next_logits = logits[0, -1] seq_pos += 1 if temperature <= 0: action = next_logits.argmax().item() else: action = torch.multinomial( F.softmax(next_logits / temperature, dim=-1), 1).item() if action == EOS_ACTION: break actions.append(action) return actions, kv_caches, seq_pos def count_params(self): return sum(p.numel() for p in self.parameters()) if __name__ == '__main__': cfg = Config() model = PolicyModel(cfg) print(f"Total params: {model.count_params():,}") print(f" Encoder: {sum(p.numel() for p in model.encoder.parameters()):,}") print(f" Decoder: {model.count_params() - sum(p.numel() for p in model.encoder.parameters()):,}") # Test encoder tokens = torch.randint(0, BOARD_VOCAB, (2, BOARD_SEQ_LEN)) emb = model.encoder(tokens) print(f" Board tokens {tokens.shape} -> {emb.shape}") # Test KV cache generation kv = None pos = 0 for turn in range(5): tokens = torch.randint(0, BOARD_VOCAB, (1, BOARD_SEQ_LEN)) acts, kv, pos = model.generate(tokens, kv_caches=kv, seq_pos=pos) cache_size = kv[0][0].shape[2] print(f" Turn {turn}: actions={acts}, seq_pos={pos}, cache_len={cache_size}")