E1c: k*=0 routing + pre-loop halt target (item 19 scored, E1c pre-registered)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

scripts/train_gate_head.py

@@ -38,7 +38,7 @@ ap.add_argument("--seed", type=int, default=0)
 ap.add_argument("--lr", type=float, default=3e-3)
 ARGS = ap.parse_args()
 BATCH = 4
-TARGET_K = {"easy": 1, "hard": ARGS.kmax}
+TARGET_K = {"easy": 0, "hard": ARGS.kmax}  # easy: halt BEFORE looping
 
 
 def main():
@@ -81,7 +81,9 @@ def main():
             calls, _ = looper.capture(ids, msk, logits_to_keep=1)
             e = looper._hin[looper.l0].detach()
             s = looper.band(e, calls)
-        loss = 0.0
+        # pre-loop halt decision on s_0: target 1 for easy (k*=0)
+        p0 = adapter.halt_prob(e[bidx, pl], s[bidx, pl])
+        loss = F.binary_cross_entropy(p0.float(), (tk == 0).float())
         for i in range(ARGS.kmax):
             with torch.no_grad():
                 x = adapter(e, s)
@@ -90,7 +92,7 @@ def main():
             p = adapter.halt_prob(e[bidx, pl], s[bidx, pl])
             tgt = ((i + 1) >= tk).float()
             loss = loss + F.binary_cross_entropy(p.float(), tgt)
-        loss = loss / ARGS.kmax
+        loss = loss / (ARGS.kmax + 1)
         opt.zero_grad(set_to_none=True)
         loss.backward()
         opt.step()