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 rows += [[2 * i, 2 * i + 1] for i in range(M // 2)]21 rows += [[-1, -1]]22 return torch.tensor(rows, dtype=torch.long, device=device)232425def action_names(cfg: PrismeConfig) -> list[str]:26 M = cfg.n_slots27 return [f"SINGLE_{j}" for j in range(M)] + [f"PAIR_{2*i}_{2*i+1}" for i in range(M // 2)] + ["STOP"]282930def action_group_size(cfg: PrismeConfig, device=None) -> torch.Tensor:31 """k(a): number of edited slots per action (0 for STOP)."""32 M = cfg.n_slots33 return torch.tensor([1] * M + [2] * (M // 2) + [0], dtype=torch.long, device=device)343536def init_workspace(slot_embed: torch.Tensor, z: torch.Tensor) -> torch.Tensor:37 """S0[j] = e[j] + z_t. slot_embed: [M, w]; z: [..., w] -> [..., M, w]."""38 return slot_embed + z.unsqueeze(-2)394041def rms(x: torch.Tensor) -> torch.Tensor:42 return x.float().pow(2).mean(-1).sqrt()43