| """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 |
| |
| 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 |
| |
| 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() |
|
|
