code/train_policy_distill.py

"""
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()