1"""Paired intervention-value targets: g(s,a;y) = loss(D(s),y) - loss(D(T_a(s)),y).23An intervention on the model's own computation (same checkpoint, same target, cloned state). Positive = the4action reduced the teacher-forced next-token loss. These labels are used ONLY offline (controller training and5oracle audits); the deployed controller never sees y.6"""7from __future__ import annotations89import torch1011from .losses import token_nll12from .model import EncodedContext, PrismeModel131415@torch.no_grad()16def state_loss(model: PrismeModel, S: torch.Tensor, ctx: EncodedContext, targets: torch.Tensor,17 edited: torch.Tensor) -> torch.Tensor:18 logits, _ = model.logits_from_state(S, ctx, edited)19 return token_nll(logits, targets) # [B,T]202122@torch.no_grad()23def paired_gains(model: PrismeModel, S: torch.Tensor, ctx: EncodedContext, targets: torch.Tensor,24 edited: torch.Tensor, actions: torch.Tensor, step, remaining) -> torch.Tensor:25 """actions: [B,T,n] candidate actions evaluated from the SAME state S (cloned implicitly: apply_action is26 functional). Returns gains [B,T,n] = base_loss - loss_after_action (STOP gives exactly 0)."""27 base = state_loss(model, S, ctx, targets, edited)28 gains = []29 for i in range(actions.shape[-1]):30 a = actions[:, :, i]31 S_a, _ = model.apply_action(S, ctx, a, step, remaining)32 ed = edited | (a != model.cfg.stop_action)33 la = state_loss(model, S_a, ctx, targets, ed)34 gains.append(base - la)35 return torch.stack(gains, -1)363738@torch.no_grad()39def all_action_gains(model: PrismeModel, S, ctx, targets, edited, step, remaining) -> torch.Tensor:40 """Exhaustive audit over all nonterminal actions -> [B,T,A]."""41 A = model.cfg.n_actions42 B, T = targets.shape43 acts = torch.arange(A, device=targets.device).view(1, 1, A).expand(B, T, A)44 return paired_gains(model, S, ctx, targets, edited, acts, step, remaining)454647def sample_logged_actions(cfg, B, T, policy_pref: torch.Tensor | None, generator=None, device=None):48 """Four candidates per state: current-policy preference (or random if none), a competitive alternative49 (random single/pair of the other group), and two uniformly random actions, balanced singles/pairs.50 Returns actions [B,T,4] and log-propensities [B,T,4] under the mixture used here."""51 A, M = cfg.n_actions, cfg.n_slots52 g = generator53 def rand_single():54 return torch.randint(0, M, (B, T), generator=g, device=device)55 def rand_pair():56 return torch.randint(M, A, (B, T), generator=g, device=device)57 a0 = policy_pref if policy_pref is not None else rand_single()58 a1 = torch.where(a0 < M, rand_pair(), rand_single())59 a2, a3 = rand_single(), rand_pair()60 acts = torch.stack([a0, a1, a2, a3], -1)61 # propensity of each candidate under the logging mixture (uniform over the 4 lanes, lanes 2/3 uniform in group)62 p = torch.zeros(B, T, 4, device=device)63 p[..., 0] = 1.0 if policy_pref is not None else 1.0 / M64 p[..., 1] = torch.where(a0 < M, torch.full_like(p[..., 1], 1.0 / (A - M)), torch.full_like(p[..., 1], 1.0 / M))65 p[..., 2] = 1.0 / M66 p[..., 3] = 1.0 / (A - M)67 return acts, (p / 4).log()68