""" Trains the one-shot policy against the value model, with explicit state tracking. The policy model (`train.py`) hits a hard reach wall at ~8 moves: 69% at depth 6, 1.5% at depth 10, zero beyond. Four interventions moved it none of the way (see `CLAUDE.md`). Two suspects survive, and this trainer attacks both at once. ## Suspect 1 -- the labels punish correct moves `train.py` trains on one Kociemba solution per state with hard cross-entropy. A cube state usually has several optimal first moves, and Kociemba emits one path. Every other optimal move is scored as an error, so the loss actively pushes down moves that are perfectly good. The fix here replaces the single hard label with a soft distribution supplied by the value model. For a state s, score all 18 children with the teacher, take `q(a) = 1 + J(s')`, and train against `softmax(-q / tau)`. Two optimal moves get equal mass instead of one getting all of it. No search is needed to build this target: it is 18 forward passes of a 33.6M net, the same expansion DAVI already does every step in `train_value_iteration.py`. ## Suspect 2 -- the decoder loses the cube A one-shot policy must carry the cube in its activations. It emits move 1, then must predict move 2 from a state it was never shown. Nothing in `train.py` ever checks that the internal state is right, so a model can learn plausible-looking move preferences while its internal cube drifts. The fix is an auxiliary head. At every output position the decoder must also predict the 54 facelets of the state it is currently in. The engine knows that state exactly, so the supervision costs nothing. `state accuracy by position` in the eval output is the direct read: if it falls away with position, the decoder is losing the cube, and that is the wall. ## What the numbers mean The eval prints the student's greedy solve rate and the teacher's side by side. This is the point of the experiment, not decoration: - **teacher greedy** is the ceiling. The student is trained to imitate the teacher's per-state preference, and the teacher descends with the real state supplied by the engine at every step. A perfect student matches it and cannot beat it. Measured on the published value-iteration weights that is ~34% at depth 10 -- far below the 100% that beam 1024 reaches, because a beam is a search and this is not. - **student greedy** is the capability, since a one-shot policy is what the LLMs on the board are scored on. - **the gap between them** is the cost of internal state tracking, because the two differ in nothing else. So a student that lands near the teacher says the labels were the problem and state tracking is fine. A student that stays far below it says the opposite. The old run gives neither reading, which is why it stalled. Three arms come out of one script: --state-weight 0 soft targets only --state-weight 1 soft targets + state tracking (the full proposal) --hard-targets teacher argmin as a one-hot label, for the ablation that separates "soft vs hard" from "value teacher vs Kociemba" The saved checkpoint is a plain `LlamaForCausalLM`, so `serve.py`, `eval_canonical.py` and `run_benchmark.sh` all read it unchanged. The auxiliary head is a diagnostic and is saved beside it, never needed at inference. """ import argparse, json, math, os, sys, time from pathlib import Path import numpy as np import torch import torch.nn.functional as F sys.path.insert(0, str(Path(__file__).parent)) import cube_ops as C import cube_tokenizer as T from train_value_iteration import encode_batch # The 19 tokens the policy may emit on the answer span: 18 moves, then EOS. # Derived from the tokenizer rather than by arithmetic on ids, because bare-face # move tokens share their ids with colour tokens (see cube_tokenizer). Sharing is # harmless here -- colours never appear as a target on the answer span -- but the # ids are not contiguous, so they have to be looked up. MOVE_IDS = [T.STOI[m] for m in T.MOVES] ANSWER_IDS = MOVE_IDS + [T.EOS] EOS_SLOT = len(MOVE_IDS) # index of EOS inside a 19-wide target vector PROMPT_LEN = 1 + T.STATE_LEN + 1 # [BOS] + 54 facelets + [SEP] assert len(set(ANSWER_IDS)) == len(ANSWER_IDS), "answer token ids must be distinct" assert T.MOVES == list(C.MOVES), "tokenizer and cube_ops must agree on move order" # -------------------------------------------------------------------------- # teacher # -------------------------------------------------------------------------- def load_teacher(path, device): """Loads either value-model flavour through search.load_value_model. That function is the one place that tells a scalar head from a 27-way distribution head, and reading a scalar head as a distribution silently returns a constant. Do not bypass it (recorded trap in CLAUDE.md). """ from search import load_value_model net = load_value_model(path, device) net.eval() for q in net.parameters(): q.requires_grad_(False) return net @torch.no_grad() def teacher_values(net, states: np.ndarray, device, batch_size=8192, amp=True): """Estimated distance-to-solved for each state in an (N, 54) array.""" kind = getattr(net, "value_kind", "scalar") out = np.empty(len(states), dtype=np.float32) for i in range(0, len(states), batch_size): chunk = states[i:i + batch_size] ids = encode_batch(chunk, device) with torch.autocast("cuda", dtype=torch.bfloat16, enabled=amp and device == "cuda"): raw = net(ids).float() if kind == "scalar": v = raw.clamp(min=0.0) else: probs = torch.softmax(raw, dim=-1) grid = torch.arange(probs.shape[-1], device=probs.device, dtype=torch.float32) v = (probs * grid).sum(-1) out[i:i + len(chunk)] = v.cpu().numpy() # A solved state is worth 0 by definition, never by prediction. out[C.is_solved(states)] = 0.0 return out def child_costs(net, states: np.ndarray, device, amp=True): """(N, 54) -> (N, 18) cost `1 + J(child)` for every move, and the children.""" kids = C.all_children(states) flat = kids.reshape(-1, 54) v = teacher_values(net, flat, device, amp=amp).reshape(len(states), C.NUM_MOVES) # A move onto solved costs exactly 1, whatever the net thinks of that state. solved_kid = C.is_solved(flat).reshape(len(states), C.NUM_MOVES) return np.where(solved_kid, 1.0, 1.0 + v).astype(np.float32), kids # -------------------------------------------------------------------------- # rollouts -> training batch # -------------------------------------------------------------------------- def rollout_batch(net, states0: np.ndarray, device, rollout_len, tau, explore_frac, rng, hard=False, amp=True): """Rolls the teacher out and returns one padded training batch. Position t of the answer span sits at input index `PROMPT_LEN - 1 + t` and predicts the move to make from state s_t. Its target is the teacher's preference over the 18 children of s_t -- or all mass on EOS when s_t is already solved. Trajectories that never reach solved are kept, not discarded. Dropping them would bias the whole dataset toward states the teacher can already descend greedily, which are exactly the shallow ones, and would rebuild the reach wall out of the training distribution. Every position is correct supervision on its own -- "here is the best move from here" -- whether or not the trajectory that produced it ever terminated. A truncated trajectory simply carries no EOS target, which teaches nothing false. """ b = len(states0) cap = rollout_len + 1 # +1 for the position that emits EOS tgt = np.zeros((b, cap, len(ANSWER_IDS)), dtype=np.float32) seen = np.zeros((b, cap, T.STATE_LEN), dtype=np.uint8) moves = np.full((b, cap), -1, dtype=np.int64) npos = np.zeros(b, dtype=np.int64) cur = states0.copy() alive = np.ones(b, dtype=bool) for t in range(cap): idx = np.where(alive)[0] if not len(idx): break seen[idx, t] = cur[idx] npos[idx] = t + 1 solved_now = C.is_solved(cur[idx]) stop = idx[solved_now] tgt[stop, t, EOS_SLOT] = 1.0 alive[stop] = False go = idx[~solved_now] if not len(go): continue if t == cap - 1: # No room left to emit a move from here, so this position carries no # target and must not be trained on. npos[go] = t alive[go] = False continue cost, kids = child_costs(net, cur[go], device, amp=amp) best = cost.argmin(axis=1) if hard: tgt[go, t, best] = 1.0 else: p = torch.softmax(torch.from_numpy(-cost / tau), dim=1).numpy() tgt[go, t, :C.NUM_MOVES] = p chosen = best if explore_frac > 0: take = rng.random(len(go)) < explore_frac if take.any(): w = torch.softmax(torch.from_numpy(-cost[take] / tau), dim=1).numpy() u = rng.random((int(take.sum()), 1)) chosen = chosen.copy() chosen[take] = (w.cumsum(axis=1) < u).sum(axis=1).clip(0, C.NUM_MOVES - 1) moves[go, t] = chosen cur[go] = kids[np.arange(len(go)), chosen] return pack(states0, tgt, seen, moves, npos) def pack(states0, tgt, seen, moves, npos): """Pads the rollouts into tensors aligned to the language-model sequence.""" b = len(states0) keep = npos > 0 length = int(npos.max()) seq = PROMPT_LEN + max(length - 1, 0) ids = np.full((b, seq), T.PAD, dtype=np.int64) ids[:, 0] = T.BOS ids[:, 1:1 + T.STATE_LEN] = np.asarray(T.COLOR_IDS, dtype=np.int64)[states0.astype(np.int64)] ids[:, T.STATE_LEN + 1] = T.SEP attn = np.zeros((b, seq), dtype=np.int64) posmask = np.zeros((b, seq), dtype=bool) # float16 / uint8 rather than float32 / int64: a packed batch is kept in a # replay buffer, and the int64 facelets alone would be 9MB a batch. soft = np.zeros((b, seq, tgt.shape[2]), dtype=np.float16) faces = np.zeros((b, seq, T.STATE_LEN), dtype=np.uint8) order = np.zeros((b, seq), dtype=np.int16) for i in range(b): n = int(npos[i]) if n <= 0: continue # The prefix that reaches position t is the prompt plus moves 1..t, so a # trajectory with n scored positions contributes n-1 move tokens. mv = moves[i, : n - 1] ids[i, PROMPT_LEN:PROMPT_LEN + len(mv)] = np.asarray(MOVE_IDS, dtype=np.int64)[mv] attn[i, : PROMPT_LEN + len(mv)] = 1 lo = PROMPT_LEN - 1 posmask[i, lo:lo + n] = True soft[i, lo:lo + n] = tgt[i, :n] faces[i, lo:lo + n] = seen[i, :n] order[i, lo:lo + n] = np.arange(n) return { "ids": torch.from_numpy(ids), "attn": torch.from_numpy(attn), "posmask": torch.from_numpy(posmask), "soft": torch.from_numpy(soft), "faces": torch.from_numpy(faces), "order": torch.from_numpy(order), "solved": torch.from_numpy(keep & (tgt[:, :, EOS_SLOT].sum(axis=1) > 0)), } # -------------------------------------------------------------------------- # model # -------------------------------------------------------------------------- def build_policy(args): from transformers import LlamaConfig, LlamaForCausalLM cfg = LlamaConfig( vocab_size=T.VOCAB_SIZE, hidden_size=args.hidden, intermediate_size=args.hidden * 4, num_hidden_layers=args.layers, num_attention_heads=args.heads, num_key_value_heads=args.heads, max_position_embeddings=T.MAX_SEQ, rms_norm_eps=1e-5, pad_token_id=T.PAD, bos_token_id=T.BOS, eos_token_id=T.EOS, tie_word_embeddings=True, ) return LlamaForCausalLM(cfg) class StateHead(torch.nn.Module): """Predicts the 54 facelets of the state the decoder is currently in. One linear layer on the hidden state, deliberately: a deep head could learn 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): super().__init__() self.proj = torch.nn.Linear(hidden, T.STATE_LEN * len(C.COLORS)) def forward(self, h): 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, probe=None): ids = batch["ids"].to(device) attn = batch["attn"].to(device) posmask = batch["posmask"].to(device) soft = batch["soft"].to(device).float() faces = batch["faces"].to(device).long() with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp): hidden = model.model(input_ids=ids, attention_mask=attn).last_hidden_state logits = model.lm_head(hidden) face_logits = head(hidden) # Cross-entropy over the full vocabulary against a soft target that is zero # off the answer tokens. Normalising over the whole vocabulary rather than # over the 19 answer tokens keeps the remaining logits pushed down, so greedy # decoding cannot wander into a colour token. 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)), 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 total, move_loss, face_loss, face_acc # -------------------------------------------------------------------------- # evaluation # -------------------------------------------------------------------------- def _walk_exact(n, depth, rng): """n states, each from a walk of exactly `depth` moves off solved.""" states = np.repeat(C.SOLVED_ARR[None, :], n, axis=0) prev = np.full(n, -1) for _ in range(depth): mv = rng.integers(0, C.NUM_MOVES, size=n) clash = (mv // 3) == prev while clash.any(): mv[clash] = rng.integers(0, C.NUM_MOVES, size=int(clash.sum())) clash = (mv // 3) == prev states = states[np.arange(n)[:, None], C.PERMS[mv]] prev = mv // 3 return states @torch.no_grad() def teacher_greedy(net, device, depths, n, rng, max_steps=30, amp=True): """The ceiling: descend to the lowest-cost child with the engine supplying the state at every step.""" out = {} for d in depths: states = _walk_exact(n, d, rng) done = C.is_solved(states) for _ in range(max_steps): live = np.where(~done)[0] if not len(live): break cost, kids = child_costs(net, states[live], device, amp=amp) best = cost.argmin(axis=1) states[live] = kids[np.arange(len(live)), best] done[live] = C.is_solved(states[live]) out[d] = float(done.mean()) return out @torch.no_grad() def student_greedy(model, device, depths, n, rng, max_new=T.MAX_SOLUTION + 1, batch_size=256): """The capability: one prompt, one greedy sequence, replayed by the engine. Greedy is the honest number for a policy, unlike for the value model, where greedy is a training signal and search is the system. Here the sequence is the whole answer, and it is what every LLM on the board is scored on. """ model.eval() out = {} for d in depths: states = _walk_exact(n, d, rng) solved = 0 for start in range(0, n, batch_size): chunk = states[start:start + batch_size] cur = chunk.copy() ids = encode_batch(chunk, device) # Decoding stops at EOS and only at EOS, then the whole emitted # sequence is replayed. Stopping early because the cube happens to be # solved mid-sequence would score a model the harness would mark # wrong: the harness applies every move the model wrote, and moves # after a solve unsolve it. live = np.ones(len(chunk), dtype=bool) for _ in range(max_new): logits = model(input_ids=ids).logits[:, -1, :] nxt = logits[:, ANSWER_IDS].argmax(-1).cpu().numpy() live = live & (nxt != EOS_SLOT) step = np.where(live)[0] if not len(step): break cur[step] = C.apply_moves(cur[step], nxt[step]) tok = torch.full((len(chunk), 1), T.PAD, dtype=torch.long, device=device) tok[step, 0] = torch.tensor( [MOVE_IDS[m] for m in nxt[step]], dtype=torch.long, device=device) ids = torch.cat([ids, tok], dim=1) if ids.shape[1] >= T.MAX_SEQ: break solved += int(C.is_solved(cur).sum()) out[d] = solved / max(n, 1) model.train() return out @torch.no_grad() def state_accuracy(model, head, net, device, depths, n, rng, rollout_len, tau, amp): """Facelet accuracy of the auxiliary probe, bucketed by output position. The direct test of suspect 2. Teacher-forced, so the prefix is a real trajectory and the true state at every position is known exactly. """ model.eval() head.eval() bins = {"0": [0, 0], "1-4": [0, 0], "5-8": [0, 0], "9-12": [0, 0], "13+": [0, 0]} exact = [0, 0] for d in depths: batch = rollout_batch(net, _walk_exact(n, d, rng), device, rollout_len, tau, 0.0, rng, amp=amp) hidden = model.model(input_ids=batch["ids"].to(device), attention_mask=batch["attn"].to(device)).last_hidden_state pred = head(hidden.float()).argmax(-1).cpu().to(torch.uint8) ok = (pred == batch["faces"]) mask = batch["posmask"] order = batch["order"] for key, lo, hi in (("0", 0, 0), ("1-4", 1, 4), ("5-8", 5, 8), ("9-12", 9, 12), ("13+", 13, 10**6)): sel = mask & (order >= lo) & (order <= hi) if sel.any(): bins[key][0] += int(ok[sel].sum()) bins[key][1] += int(sel.sum()) * T.STATE_LEN exact[0] += int(ok.all(-1)[mask].sum()) exact[1] += int(mask.sum()) model.train() head.train() per_pos = " ".join(f"{k}:{v[0]/v[1]:.1%}" for k, v in bins.items() if v[1]) return per_pos, exact[0] / max(exact[1], 1) # -------------------------------------------------------------------------- def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--teacher", required=True, help="Value-model checkpoint (.pt) from train_value_iteration.py.") p.add_argument("--hidden", type=int, default=512) p.add_argument("--layers", type=int, default=10) p.add_argument("--heads", type=int, default=8) p.add_argument("--batch-size", type=int, default=128, help="Start states per step. Each contributes up to " "--rollout-len scored positions, and costs 18 teacher " "forward passes per position.") p.add_argument("--rollout-len", type=int, default=16, help="Cap on trajectory length. Must stay at or under 25, which " "is where the sequence reaches cube_tokenizer.MAX_SEQ.") p.add_argument("--scramble-depth", type=int, default=12, help="States come from walks of 1..this many moves. Deliberately " "NOT the value trainer's 30: the teacher's predicted distance " "saturates near 12 moves, so past that every child scores the " "same and the soft target is noise with a shape. Sampling " "there would train the policy on the teacher's ignorance. " "Raise it only with a teacher whose frontier reaches further.") p.add_argument("--buffer", type=int, default=32, help="Packed rollout batches kept for reuse. The rollouts come " "from the teacher, not from the student, so the data " "distribution does not move as the student learns and reuse " "is plain caching rather than off-policy training.") p.add_argument("--updates-per-rollout", type=int, default=4, help="Gradient steps taken per batch of rollouts generated, " "sampled from the buffer. Building a target costs 18 teacher " "forward passes per position and a gradient step costs one, " "so without reuse the teacher is ~95%% of the run and the " "student sees too few steps to converge.") p.add_argument("--tau", type=float, default=0.35, metavar="MOVES", help="Softmax temperature on child costs, in moves. Costs " "differ by whole moves, so tau=1 spreads ~30%% of the mass " "onto the best child and the rest onto plainly worse ones; " "0.35 keeps ties sharing equally while a one-move-worse " "child takes ~5%%. Ties are what this is for, not noise.") p.add_argument("--explore-frac", type=float, default=0.2, help="Fraction of steps that sample from the target instead of " "taking the argmin, to widen the visited states.") p.add_argument("--state-weight", type=float, default=1.0, help="Weight on the auxiliary facelet loss. 0 disables it, " "which is the soft-targets-only arm.") p.add_argument("--hard-targets", action="store_true", help="One-hot on the teacher's argmin instead of a soft " "distribution. The ablation for suspect 1.") p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--warmup", type=int, default=200) p.add_argument("--max-steps", type=int, default=20000) p.add_argument("--eval-every", type=int, default=500) p.add_argument("--eval-n", type=int, default=100) p.add_argument("--eval-depths", default="3,6,10") p.add_argument("--ceiling-only", action="store_true", help="Measure the teacher's greedy solve rate per depth and " "exit, training nothing. This is the upper limit on what " "the student can score, so it is worth knowing before a " "training run is paid for. Inference only.") p.add_argument("--out", default="checkpoints/policy_distill") p.add_argument("--hub-repo", default="") p.add_argument("--seed", type=int, default=0) args = p.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" amp = device == "cuda" torch.manual_seed(args.seed) rng = np.random.default_rng(args.seed) net = load_teacher(args.teacher, device) if args.ceiling_only: depths = [int(d) for d in args.eval_depths.split(",")] t = time.time() rates = teacher_greedy(net, device, depths, args.eval_n, np.random.default_rng(12345), amp=amp) print(f"teacher greedy, n={args.eval_n} per depth, {time.time()-t:.0f}s") for d in depths: print(f" depth {d:>3}: {rates[d]:.0%}") print(json.dumps({"teacher_greedy": rates, "n": args.eval_n})) return 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) print(f"targets {'hard argmin' if args.hard_targets else f'soft tau={args.tau}'} | " f"state weight {args.state_weight} | batch {args.batch_size} x " f"rollout {args.rollout_len}", flush=True) amp_dtype = torch.bfloat16 if amp and not torch.cuda.is_bf16_supported(): amp_dtype = torch.float16 print("bf16 unsupported on this GPU; using fp16", flush=True) 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) * 0.5 * (1 + math.cos(math.pi * min(s / max(args.max_steps, 1), 1.0)))) depths = [int(d) for d in args.eval_depths.split(",")] out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) # The ceiling, measured once before training. Everything the student does is # read against this, so it is printed first and stored with the history. ceiling = teacher_greedy(net, device, depths, args.eval_n, np.random.default_rng(12345), amp=amp) print("teacher greedy (the ceiling): " + " ".join(f"d{d}:{r:.0%}" for d, r in ceiling.items()), flush=True) api = None if args.hub_repo: from huggingface_hub import HfApi api = HfApi(token=os.environ.get("HF_TOKEN")) run, best, t0 = [], -1.0, time.time() buffer, solved_frac = [], 0.0 for step in range(1, args.max_steps + 1): if step % args.updates_per_rollout == 1 or args.updates_per_rollout == 1: states0, _ = C.random_walk(args.batch_size, args.scramble_depth, rng) fresh = rollout_batch(net, states0, device, args.rollout_len, args.tau, args.explore_frac, rng, hard=args.hard_targets, amp=amp) solved_frac = float(fresh["solved"].float().mean()) buffer.append(fresh) if len(buffer) > args.buffer: buffer.pop(0) batch = buffer[rng.integers(len(buffer))] loss, mloss, floss, facc = losses(model, head, batch, device, 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()) + list(probe.parameters()), 1.0) opt.step() sched.step() if step % 50 == 0: print(f"step {step}/{args.max_steps} move {mloss.item():.4f} " f"face {floss.item():.4f} face_acc {facc.item():.1%} " f"solved_rollouts {solved_frac:.0%} " f"{step/(time.time()-t0):.2f} it/s", flush=True) if step % args.eval_every == 0 or step == args.max_steps: 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, probe, net, device, depths, max(args.eval_n // 4, 8), ev, args.rollout_len, args.tau, amp) mean = float(np.mean(list(rates.values()))) gap = {d: ceiling[d] - rates[d] for d in depths} print(f" step {step} student greedy " + " ".join(f"d{d}:{r:.0%}" for d, r in rates.items()) + f" | mean {mean:.1%}", flush=True) print(f" gap to teacher " + " ".join(f"d{d}:{g:+.0%}" for d, g in gap.items()), flush=True) print(f" state accuracy by position {per_pos} " f"| whole cube exact {exact:.1%}", flush=True) run.append({"step": step, "student": rates, "teacher": ceiling, "gap": gap, "mean": mean, "state_by_pos": per_pos, "state_exact": exact}) (out_dir / "history.json").write_text(json.dumps(run, indent=2)) # 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(), "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(), "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 # back on one box do not overwrite each other's weights. api.upload_folder(repo_id=args.hub_repo, folder_path=f"{args.out}-best", path_in_repo=f"checkpoints/{out_dir.name}") except Exception as e: print(f" upload failed: {type(e).__name__}: {e}", flush=True) print(f"done. best mean student greedy {best:.1%}. weights in {out_dir}", flush=True) if __name__ == "__main__": main()