1"""Diagnostic oracles (use the target; NEVER deployable):2 * greedy one-step oracle: at each step pick argmax_a g(s,a), STOP when max gain <= 0;3 * two-action lookahead oracle on audit subsets: best pair (a1, a2) by total gain;4 * uniform / random / fixed-budget reference policies for the opportunity analysis.5"""6from __future__ import annotations78import torch910from .counterfactual import all_action_gains, state_loss11from .model import PrismeModel121314@torch.no_grad()15def greedy_oracle_rollout(model: PrismeModel, tokens, targets, max_budget: int, ctx=None, threshold: float = 0.0):16 """Returns dict(loss_per_step [R+1,B,T], actions [B,T,R], n_edits [B,T])."""17 cfg = model.cfg18 ctx = model.encode(tokens) if ctx is None else ctx19 S = ctx.S020 B, T = targets.shape21 edited = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)22 stopped = torch.zeros_like(edited)23 losses = [state_loss(model, S, ctx, targets, edited)]24 trace = torch.full((B, T, max_budget), cfg.stop_action, dtype=torch.long, device=tokens.device)25 for r in range(max_budget):26 gains = all_action_gains(model, S, ctx, targets, edited, r, max_budget - r - 1)27 best_g, best_a = gains.max(-1)28 a = torch.where((best_g > threshold) & ~stopped, best_a, torch.full_like(best_a, cfg.stop_action))29 stopped = stopped | (a == cfg.stop_action)30 trace[:, :, r] = a31 S, _ = model.apply_action(S, ctx, a, r, max_budget - r - 1)32 edited = edited | (a != cfg.stop_action)33 losses.append(state_loss(model, S, ctx, targets, edited))34 return {"loss_per_step": torch.stack(losses), "actions": trace, "n_edits": (trace != cfg.stop_action).sum(-1)}353637@torch.no_grad()38def two_step_oracle(model: PrismeModel, tokens, targets, ctx=None):39 """Exhaustive best pair of actions (a1 then a2, a2 may be STOP) on a (small) batch -> total gain [B,T],40 and the best greedy-then-greedy total for comparison. Cost: A + A*(A+1) editor passes."""41 cfg = model.cfg42 A = cfg.n_actions43 ctx = model.encode(tokens) if ctx is None else ctx44 S0 = ctx.S045 B, T = targets.shape46 no = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)47 yes = ~no48 base = state_loss(model, S0, ctx, targets, no)49 g1 = all_action_gains(model, S0, ctx, targets, no, 0, 1) # [B,T,A]50 best_total = g1.clone() # a2 = STOP51 best_pair_a2 = torch.full((B, T, A), cfg.stop_action, dtype=torch.long, device=tokens.device)52 for a1 in range(A):53 act1 = torch.full((B, T), a1, device=tokens.device)54 S1, _ = model.apply_action(S0, ctx, act1, 0, 1)55 g2 = all_action_gains(model, S1, ctx, targets, yes, 1, 0) # [B,T,A]56 tot = g1[..., a1:a1 + 1] + g257 m, am = tot.max(-1)58 better = m > best_total[..., a1]59 best_total[..., a1] = torch.where(better, m, best_total[..., a1])60 best_pair_a2[..., a1] = torch.where(better, am, best_pair_a2[..., a1])61 two_step_best, a1_star = best_total.max(-1)62 # greedy baseline: a1 = argmax g1, then best a2 given a1 (already in best_total[a1_greedy])63 g1_best, a1_greedy = g1.max(-1)64 greedy_then_best = best_total.gather(-1, a1_greedy.unsqueeze(-1)).squeeze(-1)65 return {"base_loss": base, "one_step_best": g1_best.clamp(min=0), "two_step_best": two_step_best.clamp(min=0),66 "greedy_then_best": greedy_then_best.clamp(min=0), "a1_star": a1_star, "a1_greedy": a1_greedy}676869@torch.no_grad()70def fixed_policy_rollout(model: PrismeModel, tokens, targets, actions: torch.Tensor):71 """Loss trajectory under a forced action tensor [B,T,R]."""72 cfg = model.cfg73 ctx = model.encode(tokens)74 S = ctx.S075 B, T, R = actions.shape76 edited = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)77 losses = [state_loss(model, S, ctx, targets, edited)]78 for r in range(R):79 a = actions[:, :, r]80 S, _ = model.apply_action(S, ctx, a, r, R - r - 1)81 edited = edited | (a != cfg.stop_action)82 losses.append(state_loss(model, S, ctx, targets, edited))83 return torch.stack(losses)84