1"""Reference routing policies that share PRISME's editor (actions chosen without the learned value controller).2All are target-free and deployable; they differ only in how the action is picked at each step.3 random uniform over the 24 nonterminal actions, fixed budget b4 uniform round-robin single-slot schedule SINGLE_0, SINGLE_1, ... (deterministic, fixed budget b)5 fixed_pair PAIR_0_1 repeated (fixed budget b)6 entropy continue (random action) while the current predictive entropy exceeds a threshold, else STOP7 norm continue while the last write moved the slot by more than a threshold (RMS), else STOP8"""9from __future__ import annotations1011import torch1213from prisme.model import PrismeModel141516@torch.no_grad()17def rollout(model: PrismeModel, tokens, policy: str, budget: int, threshold: float = 0.0, seed: int = 0):18 cfg = model.cfg19 ctx = model.encode(tokens)20 S = ctx.S021 B, T = tokens.shape22 g = torch.Generator(device="cpu").manual_seed(seed)23 stopped = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)24 n_edits = torch.zeros(B, T, dtype=torch.long, device=tokens.device)25 trace = torch.full((B, T, budget), cfg.stop_action, dtype=torch.long, device=tokens.device)26 last_rms = torch.full((B, T), float("inf"), device=tokens.device)27 for r in range(budget):28 if policy == "random":29 a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)30 elif policy == "uniform":31 a = torch.full((B, T), r % cfg.n_slots, device=tokens.device)32 elif policy == "fixed_pair":33 a = torch.full((B, T), cfg.n_slots + (r % cfg.n_pairs), device=tokens.device)34 elif policy == "entropy":35 logits, _ = model.logits_from_state(S, ctx, n_edits > 0)36 p = torch.softmax(logits.float(), -1)37 ent = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)38 a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)39 a = torch.where(ent > threshold, a, torch.full_like(a, cfg.stop_action))40 elif policy == "norm":41 a = torch.randint(0, cfg.n_actions, (B, T), generator=g).to(tokens.device)42 a = torch.where(last_rms > threshold, a, torch.full_like(a, cfg.stop_action))43 else:44 raise ValueError(policy)45 a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)46 stopped |= a == cfg.stop_action47 trace[:, :, r] = a48 S_new, info = model.apply_action(S, ctx, a, r, budget - r - 1)49 moved = (info["eta"].squeeze(-1) * info["delta_rms"] * info["scale"]) # [B,T,k] realised rms step per lane50 last_rms = torch.where(a != cfg.stop_action, moved.max(-1).values, last_rms)51 S = S_new52 n_edits += (a != cfg.stop_action).long()53 logits, _ = model.logits_from_state(S, ctx, n_edits > 0)54 return {"logits": logits, "actions": trace, "n_edits": n_edits}55