1"""Value controller: scores every nonterminal action from target-free features.23Per state: a shared state representation u = f(z_t, step, remaining) is computed once (one SwiGLU MLP);4each action a then costs only one w x w projection of its chosen-slot summary s_a (mean of the slots in I(a)),5a nonlinearity and a 4-way linear head. Outputs per action: g_hat (predicted loss gain), c_hat (predicted cost6proxy), log_sigma (uncertainty scale), harm_logit (P[loss increases]).7Action choice: Q(s,a) = g_hat - lambda*c(a) - kappa*sigma_hat, Q(STOP) = 0.8The controller never sees the target token and never executes candidate edits.9"""10from __future__ import annotations1112import torch13import torch.nn as nn14import torch.nn.functional as F1516from .attention import RMSNorm, SwiGLU17from .config import PrismeConfig18from .workspace import action_table, action_group_size192021class ValueController(nn.Module):22 N_CTX = 2 # step/R, remaining/R2324 def __init__(self, cfg: PrismeConfig):25 super().__init__()26 self.cfg = cfg27 hidden = cfg.controller_hidden or cfg.ffn_w28 w = cfg.w29 self.z_proj = nn.Linear(cfg.d_model if cfg.ctrl_input == "h" else w, w, bias=False)30 self.ent_proj = nn.Linear(2, w, bias=False) if cfg.ctrl_use_entropy else None31 self.ctx_proj = nn.Linear(self.N_CTX, w, bias=True)32 self.norm = RMSNorm(w, cfg.norm_eps)33 self.mlp = SwiGLU(w, hidden)34 self.s_proj = nn.Linear(w, w, bias=False)35 self.k_embed = nn.Parameter(torch.zeros(2, w)) # group size 1 or 236 self.heads = nn.Linear(w, 4, bias=True)37 self.register_buffer("actions", action_table(cfg), persistent=False)38 self.register_buffer("group_size", action_group_size(cfg), persistent=False)3940 def state_repr(self, z: torch.Tensor, step: torch.Tensor, remaining: torch.Tensor, entropy=None) -> torch.Tensor:41 """z: [B,Tq,w] summary (or [B,Tq,d] normalised residual when ctrl_input == "h"); entropy: [B,Tq] optional."""42 ctx = torch.stack([step.to(z.dtype) / self.cfg.r_max, remaining.to(z.dtype) / self.cfg.r_max], -1)43 u = self.z_proj(z) + self.ctx_proj(ctx)44 if self.ent_proj is not None:45 assert entropy is not None, "ctrl_use_entropy requires the backbone entropy feature"46 e = entropy.to(z.dtype).unsqueeze(-1)47 u = u + self.ent_proj(torch.cat([e, torch.log(e + 1e-4)], -1))48 return u + self.mlp(self.norm(u)) # [B,Tq,w]4950 def slot_summaries(self, S: torch.Tensor) -> torch.Tensor:51 A = self.cfg.n_actions52 acts = self.actions[:A]53 amask = (acts >= 0).to(S.dtype)54 sel = S[..., acts.clamp(min=0), :] # [B,Tq,A,2,w]55 return (sel * amask.view(1, 1, A, 2, 1)).sum(-2) / amask.sum(-1).view(1, 1, A, 1)5657 def forward(self, z, S, step, remaining, entropy=None):58 """z: controller input ([B,Tq,w] summary or [B,Tq,d] residual); S: [B,Tq,M,w]; step/remaining: [B,Tq]."""59 A = self.cfg.n_actions60 u = self.state_repr(z, step, remaining, entropy).unsqueeze(2) # [B,Tq,1,w]61 s = self.s_proj(self.slot_summaries(S)) # [B,Tq,A,w]62 kidx = (self.group_size[:A] - 1).clamp(min=0)63 h = F.silu(u + s + self.k_embed[kidx].view(1, 1, A, -1))64 out = self.heads(h)65 return {"g_hat": out[..., 0], "c_hat": out[..., 1], "log_sigma": out[..., 2], "harm_logit": out[..., 3]}6667 def q_values(self, preds: dict, cost: torch.Tensor, lambda_cost: float | None = None, kappa_risk: float | None = None):68 lam = self.cfg.lambda_cost if lambda_cost is None else lambda_cost69 kap = self.cfg.kappa_risk if kappa_risk is None else kappa_risk70 sigma = preds["log_sigma"].clamp(-7, 3).exp()71 return preds["g_hat"] - lam * cost - kap * sigma7273 def select(self, preds: dict, cost: torch.Tensor, feasible: torch.Tensor | None = None, **kw) -> torch.Tensor:74 """Greedy action with STOP (index A) when max Q <= 0. feasible: optional [B,Tq,A] bool."""75 q = self.q_values(preds, cost, **kw)76 if feasible is not None:77 q = q.masked_fill(~feasible, float("-inf"))78 best_q, best_a = q.max(-1)79 return torch.where(best_q > 0, best_a, torch.full_like(best_a, self.cfg.stop_action))80