candypunk/NanoJev-Web

1"""Decision network extracted from NanoJev; no trainer or game imports."""2import torch3from torch import nn4import torch.nn.functional as F5 6class DecisionModel(nn.Module):7    def __init__(self, backbone, set_head):8        super().__init__()9        self.backbone = backbone10        hidden = backbone.config.hidden_size11        self.norm = nn.LayerNorm(hidden)12        self.scalar = nn.Linear(hidden, 1)  # Nonzero random initialization avoids a dead first step.13        nn.init.normal_(self.scalar.weight, std=0.02)14        nn.init.zeros_(self.scalar.bias)15        self.set_head = set_head16        if set_head == 'attention':17            self.set_project = nn.Linear(hidden + 1, 128)18            self.set_attention = nn.MultiheadAttention(128, 4, dropout=0.0, batch_first=True)19            self.set_output = nn.Linear(128, 1)20            # Only the final residual projection starts at zero; its upstream layers are nonzero.21            nn.init.zeros_(self.set_output.weight)22            nn.init.zeros_(self.set_output.bias)23 24    def forward(self, examples, pad_token):25        paths = [ids for ex in examples for ids in ex['leaf_tokens']]26        device = self.scalar.weight.device27        lengths = torch.tensor([len(ids) for ids in paths], device=device)28        width = int(lengths.max())29        tokens = torch.full((len(paths), width), pad_token, dtype=torch.long, device=device)30        for i, ids in enumerate(paths):31            tokens[i, :len(ids)] = torch.tensor(ids, device=device)32        attention = torch.arange(width, device=device)[None, :] < lengths[:, None]33        hidden = self.backbone(input_ids=tokens, attention_mask=attention,34                               use_cache=False).last_hidden_state35        leaves = hidden[torch.arange(len(paths), device=device), lengths-1]36        kmax = max(len(ex['candidate_ids']) for ex in examples)37        h = leaves.new_zeros((len(examples), kmax, leaves.shape[-1]))38        valid = torch.zeros((len(examples), kmax), dtype=torch.bool, device=device)39        offset = 040        for i, ex in enumerate(examples):41            n = len(ex['leaf_tokens'])42            h[i, :n] = leaves[offset:offset+n]43            valid[i, :len(ex['candidate_ids'])] = True44            offset += n45        h = self.norm(h)46        z = self.scalar(h).squeeze(-1).float()47        choice = torch.tensor([i for i, ex in enumerate(examples) if ex['type'] == 'choice'], device=device)48        if self.set_head == 'attention' and len(choice):49            log_k = valid[choice].sum(-1).float().log()[:, None, None].expand(-1, kmax, 1)50            u = self.set_project(torch.cat([h[choice], log_k.to(h.dtype)], dim=-1))51            mixed, _ = self.set_attention(u, u, u, key_padding_mask=~valid[choice], need_weights=False)52            delta = self.set_output(torch.tanh(u + mixed)).squeeze(-1).float()53            z = z.index_add(0, choice, delta)54        # Boolean has one semantic path and one scalar, representing logits [0,z].55        out = []56        for i, ex in enumerate(examples):57            if ex['type'] == 'boolean':58                out.append(F.pad(torch.stack([z[i, 0] * 0, z[i, 0]]), (0, kmax-2)))59            else:60                out.append(z[i])61        return torch.stack(out).masked_fill(~valid, -1e9), valid62