1"""Action catalog and ephemeral workspace helpers.23Action catalog (M = 16 slots):4 SINGLE_j (j = 0..M-1) -> edit one slot : actions 0 .. M-15 PAIR_{2i,2i+1} (i = 0..M/2-1) -> edit two fixed slots : actions M .. M+M/2-16 STOP -> no edit : action M+M/27Exactly M + M/2 nonterminal actions + STOP (= 25 for M = 16). Arbitrary pairs are NOT available.8"""9from __future__ import annotations1011import torch1213from .config import PrismeConfig141516def action_table(cfg: PrismeConfig, device=None) -> torch.Tensor:17 """[n_actions + 1, 2] slot indices per action; -1 marks an unused query lane. Last row = STOP."""18 M = cfg.n_slots19 rows = [[j, -1] for j in range(M)]20 if getattr(cfg, "pair_scheme", "adjacent") == "offset":21 rows += [[i, i + M // 2] for i in range(M // 2)]22 else:23 rows += [[2 * i, 2 * i + 1] for i in range(M // 2)]24 rows += [[-1, -1]]25 return torch.tensor(rows, dtype=torch.long, device=device)262728def action_names(cfg: PrismeConfig) -> list[str]:29 tbl = action_table(cfg).tolist()30 return [f"SINGLE_{a}" if b < 0 else f"PAIR_{a}_{b}" for a, b in tbl[:-1]] + ["STOP"]313233def action_group_size(cfg: PrismeConfig, device=None) -> torch.Tensor:34 """k(a): number of edited slots per action (0 for STOP)."""35 M = cfg.n_slots36 return torch.tensor([1] * M + [2] * (M // 2) + [0], dtype=torch.long, device=device)373839def init_workspace(slot_embed: torch.Tensor, z: torch.Tensor) -> torch.Tensor:40 """S0[j] = e[j] + z_t. slot_embed: [M, w]; z: [..., w] -> [..., M, w]."""41 return slot_embed + z.unsqueeze(-2)424344def rms(x: torch.Tensor) -> torch.Tensor:45 return x.float().pow(2).mean(-1).sqrt()46