code: detached probe, so every version can be measured

code/test_policy_distill.py

@@ -121,6 +121,26 @@ def test_train_step(net):
     return model, head
 
 
+def test_probe_is_detached(net):
+    """The probe must train itself and leave the model it measures alone."""
+    class Args:
+        hidden, layers, heads = 32, 2, 2
+
+    model = D.build_policy(Args())
+    head, probe = D.StateHead(32), D.StateHead(32)
+    rng = np.random.default_rng(11)
+    states0, _ = C.random_walk(8, 8, rng)
+    batch = D.rollout_batch(net, states0, "cpu", 6, 0.35, 0.0, rng, amp=False)
+    # state_weight 0 removes the attached head, so anything reaching the trunk
+    # after that comes from the move loss alone -- never from the probe.
+    loss, _, _, _ = D.losses(model, head, batch, "cpu", 0.0, torch.bfloat16, False,
+                             probe)
+    loss.backward()
+    assert all(q.grad is not None and q.grad.abs().sum() > 0 for q in probe.parameters())
+    assert all(q.grad is None or q.grad.abs().sum() == 0 for q in head.parameters())
+    print("ok  the probe trains itself and does not reach the attached head")
+
+
 def test_eval(net, model, head):
     rates = D.student_greedy(model, "cpu", [3], 8, np.random.default_rng(4))
     assert set(rates) == {3} and 0.0 <= rates[3] <= 1.0
@@ -154,6 +174,7 @@ if __name__ == "__main__":
         test_targets(teacher)
         test_hard_targets(teacher)
         test_solved_start(teacher)
+        test_probe_is_detached(teacher)
         model, head = test_train_step(teacher)
         test_eval(teacher, model, head)
     print("\nall smoke tests passed")

code/train_policy_distill.py

@@ -296,6 +296,20 @@ class StateHead(torch.nn.Module):
     to reconstruct the cube itself and let the trunk off the hook, which would
     make the diagnostic meaningless. The question is whether the trunk carries
     the cube, so the probe must be weak.
+
+    Two of these are built, and the difference is the whole reason there are two:
+
+    - the **probe** reads `hidden.detach()`, so it is a pure measurement. It is
+      trained in every arm, including the ones with no auxiliary loss, and it
+      never changes the model it measures.
+    - the **auxiliary head** reads the live hidden state, so its loss flows into
+      the trunk. That is the intervention, and `--state-weight 0` removes it.
+
+    The first version of this trainer used one head for both jobs. That reads
+    correctly in the arms that use the auxiliary loss and reads pure chance in the
+    arms that do not -- the head is simply never trained there -- so tracking
+    could not be compared across arms at all, which was the comparison worth
+    having. Runs before 2026-09-16 have that limitation.
     """
 
     def __init__(self, hidden):
@@ -306,7 +320,7 @@ class StateHead(torch.nn.Module):
         return self.proj(h).view(*h.shape[:-1], T.STATE_LEN, len(C.COLORS))
 
 
-def losses(model, head, batch, device, state_weight, amp_dtype, use_amp):
+def losses(model, head, batch, device, state_weight, amp_dtype, use_amp, probe=None):
     ids = batch["ids"].to(device)
     attn = batch["attn"].to(device)
     posmask = batch["posmask"].to(device)
@@ -325,11 +339,23 @@ def losses(model, head, batch, device, state_weight, amp_dtype, use_amp):
     logp = torch.log_softmax(logits.float(), dim=-1)[..., ANSWER_IDS]
     move_loss = -(soft * logp).sum(-1)[posmask].mean()
 
+    target = faces[posmask].reshape(-1)
     sel = face_logits.float()[posmask]
-    face_loss = F.cross_entropy(sel.reshape(-1, len(C.COLORS)), faces[posmask].reshape(-1))
+    face_loss = F.cross_entropy(sel.reshape(-1, len(C.COLORS)), target)
+
+    total = move_loss + state_weight * face_loss
+    if probe is not None:
+        # Detached, so this term trains the probe and nothing else. Always on, at
+        # weight 1, in every arm -- it is the measurement, not the intervention.
+        with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
+            probe_logits = probe(hidden.detach())
+        psel = probe_logits.float()[posmask]
+        probe_loss = F.cross_entropy(psel.reshape(-1, len(C.COLORS)), target)
+        total = total + probe_loss
+        sel = psel
     with torch.no_grad():
         face_acc = (sel.argmax(-1) == faces[posmask]).float().mean()
-    return move_loss + state_weight * face_loss, move_loss, face_loss, face_acc
+    return total, move_loss, face_loss, face_acc
 
 
 # --------------------------------------------------------------------------
@@ -535,6 +561,7 @@ def main():
 
     model = build_policy(args).to(device)
     head = StateHead(args.hidden).to(device)
+    probe = StateHead(args.hidden).to(device)
     n_params = sum(q.numel() for q in model.parameters())
     print(f"policy {n_params/1e6:.1f}M params | teacher {args.teacher} "
           f"({getattr(net, 'value_kind', '?')}) | {device}", flush=True)
@@ -547,7 +574,8 @@ def main():
         amp_dtype = torch.float16
         print("bf16 unsupported on this GPU; using fp16", flush=True)
 
-    opt = torch.optim.AdamW(list(model.parameters()) + list(head.parameters()),
+    opt = torch.optim.AdamW(list(model.parameters()) + list(head.parameters())
+                            + list(probe.parameters()),
                             lr=args.lr, weight_decay=0.1, betas=(0.9, 0.95))
     sched = torch.optim.lr_scheduler.LambdaLR(
         opt, lambda s: min((s + 1) / max(args.warmup, 1), 1.0)
@@ -583,11 +611,12 @@ def main():
                 buffer.pop(0)
         batch = buffer[rng.integers(len(buffer))]
         loss, mloss, floss, facc = losses(model, head, batch, device,
-                                          args.state_weight, amp_dtype, amp)
+                                          args.state_weight, amp_dtype, amp, probe)
         opt.zero_grad(set_to_none=True)
         loss.backward()
         torch.nn.utils.clip_grad_norm_(
-            list(model.parameters()) + list(head.parameters()), 1.0)
+            list(model.parameters()) + list(head.parameters())
+            + list(probe.parameters()), 1.0)
         opt.step()
         sched.step()
 
@@ -601,7 +630,7 @@ def main():
             ev = np.random.default_rng(12345)
             rates = student_greedy(model, device, depths, args.eval_n, ev)
             ev = np.random.default_rng(999)
-            per_pos, exact = state_accuracy(model, head, net, device, depths,
+            per_pos, exact = state_accuracy(model, probe, net, device, depths,
                                             max(args.eval_n // 4, 8), ev,
                                             args.rollout_len, args.tau, amp)
             mean = float(np.mean(list(rates.values())))
@@ -620,13 +649,13 @@ def main():
             # Latest always, best additionally -- a saturating metric must never
             # be the only thing that chooses a checkpoint (CLAUDE.md).
             model.save_pretrained(str(out_dir))
-            torch.save({"head": head.state_dict(), "hidden": args.hidden},
-                       out_dir / "state_head.pt")
+            torch.save({"head": head.state_dict(), "probe": probe.state_dict(),
+                        "hidden": args.hidden}, out_dir / "state_head.pt")
             if mean >= best:
                 best = mean
                 model.save_pretrained(f"{args.out}-best")
-                torch.save({"head": head.state_dict(), "hidden": args.hidden},
-                           Path(f"{args.out}-best") / "state_head.pt")
+                torch.save({"head": head.state_dict(), "probe": probe.state_dict(),
+                            "hidden": args.hidden}, Path(f"{args.out}-best") / "state_head.pt")
                 if api is not None:
                     try:
                         # Keyed on the output directory, so arms run back to