"""State-once rank/history attention with role-conditioned action heads.""" from dataclasses import asdict, dataclass import torch from torch import nn from torch.nn import functional as F from .features import ACTION_DIM, HISTORY, HISTORY_DIM, SCALAR_DIM @dataclass class ModelConfig: width: int = 128 heads: int = 4 rank_layers: int = 2 history_layers: int = 2 correction_limit: float = 6. atoms: int = 51 support_limit: float = 5. teacher_temperature: float = .6 action_attention_scale: float = 0. def dict(self): return asdict(self) class AttentionBlock(nn.Module): def __init__(self, width, heads): super().__init__() self.norm1 = nn.LayerNorm(width) self.attention = nn.MultiheadAttention(width, heads, batch_first=True, dropout=0) self.norm2 = nn.LayerNorm(width) self.ff = nn.Sequential(nn.Linear(width, width * 3), nn.GELU(), nn.Linear(width * 3, width)) def forward(self, x, mask=None): z = self.norm1(x) x = x + self.attention(z, z, z, key_padding_mask=mask, need_weights=False)[0] return x + self.ff(self.norm2(x)) class QOJNet(nn.Module): def __init__(self, config=None): super().__init__() self.config = config or ModelConfig() c, w = self.config, self.config.width if min(w, c.heads, c.rank_layers, c.history_layers) < 1 or w % c.heads: raise ValueError('width must be divisible by attention heads') self.rank_input = nn.Linear(7, w) self.rank_embedding = nn.Parameter(torch.randn(1, 15, w) * .02) self.role_embedding = nn.Embedding(4, w) self.prior_gate = nn.Parameter(torch.zeros(4)) self.scalar_input = nn.Linear(SCALAR_DIM, w) self.rank_blocks = nn.ModuleList(AttentionBlock(w, c.heads) for _ in range(c.rank_layers)) self.history_input = nn.Linear(HISTORY_DIM, w) self.time_embedding = nn.Parameter(torch.randn(1, HISTORY + 1, w) * .02) self.summary_token = nn.Parameter(torch.zeros(1, 1, w)) self.history_blocks = nn.ModuleList( AttentionBlock(w, c.heads) for _ in range(c.history_layers)) self.context = nn.Sequential(nn.Linear(3 * w, w), nn.LayerNorm(w), nn.GELU()) self.belief = nn.Linear(w, 2) self.belief_input = nn.Linear(30, w) self.action_input = nn.Sequential(nn.Linear(ACTION_DIM, w), nn.GELU(), nn.Linear(w, w)) self.action_trunk = nn.Sequential(nn.Linear(4 * w, 2 * w), nn.LayerNorm(2 * w), nn.GELU(), nn.Linear(2 * w, w), nn.GELU()) self.actors = nn.ModuleList(nn.Linear(w, 1) for _ in range(4)) self.critic = nn.Sequential(nn.Linear(w, w), nn.GELU(), nn.Linear(w, 1)) self.score_head = nn.Linear(w, c.atoms) self.win_head = nn.Linear(w, 1) self.register_buffer('support', torch.linspace(-c.support_limit, c.support_limit, c.atoms)) for head in self.actors: nn.init.zeros_(head.weight) nn.init.zeros_(head.bias) def encode(self, batch): role = self.role_embedding(batch['role']) scalar = self.scalar_input(batch['scalars']) + role ranks = self.rank_input(batch['ranks']) + self.rank_embedding + scalar.unsqueeze(1) for block in self.rank_blocks: ranks = block(ranks) history = self.history_input(batch['history']) summary = self.summary_token.expand(len(role), -1, -1) + scalar.unsqueeze(1) history = torch.cat((summary, history), dim=1) + self.time_embedding # The unmasked summary token keeps empty histories finite on CPU and MPS. padding = torch.cat((torch.zeros(len(role), 1, device=role.device, dtype=torch.bool), ~batch['mask']), dim=1) for block in self.history_blocks: history = block(history, padding) belief_logits = self.belief(ranks) allocation = belief_logits.softmax(-1) expected = allocation * batch['ranks'][:, :, 1:2] * 4 # Public bottom cards belong to the landlord with certainty. known = batch['ranks'][:, :, 2] * 4 for actor in (1, 2): fixed = (batch['role'] == actor).unsqueeze(1) landlord_slot = 0 fixed_cards = torch.zeros_like(expected) fixed_cards[:, :, landlord_slot] = known constrained = allocation * (batch['ranks'][:, :, 1:2] * 4 - known.unsqueeze(-1)).clamp_min(0) + fixed_cards expected = torch.where(fixed.unsqueeze(-1), constrained, expected) context = self.context(torch.cat((ranks.mean(1), history[:, 0], scalar), dim=-1)) context = context + self.belief_input(expected.flatten(1) / 4) return {'context': context, 'ranks': ranks, 'belief_logits': belief_logits, 'belief_counts': expected, 'value': self.critic(context).squeeze(-1), 'prior_scale': 2 * self.prior_gate[batch['role']].sigmoid()} def score(self, encoded, actions, owner, roles): query = self.action_input(actions) context = encoded['context'][owner] ranks = encoded['ranks'][owner] scale = self.config.action_attention_scale or self.config.width ** .5 weights = torch.einsum('ad,ard->ar', query, ranks) / scale attended = torch.einsum('ar,ard->ad', weights.softmax(-1), ranks) z = self.action_trunk(torch.cat((query, context, query * context, attended), dim=-1)) all_heads = torch.cat([head(z) for head in self.actors], dim=-1) correction = self.config.correction_limit * torch.tanh( all_heads.gather(1, roles[owner, None]).squeeze(1)) return {'correction': correction, 'score_logits': self.score_head(z), 'win_logits': self.win_head(z).squeeze(-1)} def forward(self, batch, action_chunk=512): encoded = self.encode(batch) results = [] for start in range(0, len(batch['actions']), action_chunk): results.append(self.score(encoded, batch['actions'][start:start + action_chunk], batch['owner'][start:start + action_chunk], batch['role'])) return {**encoded, **{key: torch.cat([result[key] for result in results]) for key in results[0]}} def teacher_logits(scores, temperature): scores = torch.as_tensor(scores, dtype=torch.float32) return ((scores - scores.mean()) / scores.std(unbiased=False).clamp_min(.1)) / temperature def distribution_target(scores, support): target = scores.sign() * torch.log1p(scores.abs()) position = ((target.clamp(support[0], support[-1]) - support[0]) / (support[-1] - support[0]) * (len(support) - 1)) lo = position.floor().long() hi = (lo + 1).clamp_max(len(support) - 1) weight = position - lo result = torch.zeros(len(scores), len(support), device=scores.device) result.scatter_add_(1, lo[:, None], (1 - weight)[:, None]) result.scatter_add_(1, hi[:, None], weight[:, None]) return result def score_loss(logits, scores, support): return -(distribution_target(scores, support) * F.log_softmax(logits, -1)).sum(-1).mean()