code: policy distillation trainer, tests and bootstrap mode

code/test_policy_distill.py

@@ -0,0 +1,159 @@
+"""
+Smoke tests for train_policy_distill.py, runnable on CPU with no checkpoint.
+
+The point is not coverage. It is that a rented GPU box must not be the place
+where a shape error or a misaligned target index shows up: that costs metered
+time, and a *silently* misaligned target costs a whole run that looks healthy.
+So the alignment between the language-model sequence and the rollout is checked
+against the string engine in gen_data, which is a completely separate
+implementation from the numpy one in cube_ops.
+
+    python test_policy_distill.py
+"""
+import sys, tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+sys.path.insert(0, str(Path(__file__).parent))
+import cube_ops as C
+import cube_tokenizer as T
+import train_policy_distill as D
+from gen_data import SOLVED, apply_sequence
+from train_value_iteration import ValueNet
+
+
+def make_teacher(tmp):
+    """A randomly initialised value net, saved in the value-iteration format."""
+    net = ValueNet(32, 2, 2)
+    path = Path(tmp) / "teacher.pt"
+    torch.save({"model": net.state_dict(), "hidden": 32, "layers": 2, "heads": 2,
+                "step": 0}, path)
+    return D.load_teacher(str(path), "cpu")
+
+
+def test_token_tables():
+    assert len(D.ANSWER_IDS) == 19
+    assert D.MOVE_IDS == [T.STOI[m] for m in T.MOVES]
+    assert D.ANSWER_IDS[D.EOS_SLOT] == T.EOS
+    # The move a target index names must be the move cube_ops applies for it.
+    s = C.encode(SOLVED)
+    for i, mv in enumerate(T.MOVES):
+        arr = C.decode(C.apply_moves(s, [i]))[0]
+        assert arr == apply_sequence(SOLVED, [mv]), mv
+    print("ok  token tables agree with both engines")
+
+
+def test_pack_alignment(net):
+    """Every masked position must carry the state its prefix actually reaches."""
+    rng = np.random.default_rng(0)
+    states0, _ = C.random_walk(6, 12, rng)
+    batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.0, rng, amp=False)
+
+    ids = batch["ids"].numpy()
+    faces = batch["faces"].numpy()
+    mask = batch["posmask"].numpy()
+    order = batch["order"].numpy()
+
+    for b in range(len(states0)):
+        start = C.decode(states0[b:b + 1])[0]
+        assert list(ids[b, :T.STATE_LEN + 2]) == T.encode_state(start)
+        idx = np.where(mask[b])[0]
+        assert list(order[b, idx]) == list(range(len(idx)))
+        for t, pos in enumerate(idx):
+            # The prefix that reaches this position is the prompt plus t moves.
+            prefix = ids[b, D.PROMPT_LEN:D.PROMPT_LEN + t]
+            moves = [T.ITOS[int(i)] for i in prefix]
+            assert apply_sequence(start, moves) == C.decode(faces[b, pos:pos + 1])[0], (b, t)
+    print(f"ok  pack alignment on {len(states0)} rollouts, string engine agrees")
+
+
+def test_targets(net):
+    rng = np.random.default_rng(1)
+    states0, _ = C.random_walk(32, 6, rng)
+    batch = D.rollout_batch(net, states0, "cpu", 10, 0.35, 0.0, rng, amp=False)
+    soft = batch["soft"].numpy()
+    mask = batch["posmask"].numpy()
+    faces = batch["faces"].numpy()
+
+    # float32 before summing: the packed target is float16, so a float16
+    # accumulation over 19 entries has more error than the assertion allows.
+    rows = soft[mask].astype(np.float32)
+    assert np.allclose(rows.sum(axis=1), 1.0, atol=1e-3), rows.sum(axis=1)[:5]
+    # EOS carries all the mass exactly where the state is solved, and none anywhere else.
+    solved = C.is_solved(faces[mask])
+    assert np.allclose(rows[solved][:, D.EOS_SLOT], 1.0)
+    assert np.allclose(rows[~solved][:, D.EOS_SLOT], 0.0)
+    assert (soft[~mask] == 0).all()
+    print(f"ok  targets normalised; EOS mass matches solved on {mask.sum()} positions")
+
+
+def test_hard_targets(net):
+    rng = np.random.default_rng(2)
+    states0, _ = C.random_walk(8, 6, rng)
+    batch = D.rollout_batch(net, states0, "cpu", 6, 0.35, 0.0, rng, hard=True, amp=False)
+    rows = batch["soft"].numpy()[batch["posmask"].numpy()]
+    assert np.allclose(rows.max(axis=1), 1.0)
+    print("ok  --hard-targets produces one-hot rows")
+
+
+def test_train_step(net):
+    class Args:
+        hidden, layers, heads = 32, 2, 2
+
+    model = D.build_policy(Args())
+    head = D.StateHead(32)
+    rng = np.random.default_rng(3)
+    states0, _ = C.random_walk(8, 10, rng)
+    batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.2, rng, amp=False)
+    loss, mloss, floss, facc = D.losses(model, head, batch, "cpu", 1.0,
+                                        torch.bfloat16, False)
+    assert torch.isfinite(loss), loss
+    loss.backward()
+    grads = [q.grad for q in head.parameters() if q.grad is not None]
+    assert grads and any(g.abs().sum() > 0 for g in grads)
+    # The facelet probe must reach the trunk, not just its own linear layer.
+    trunk = model.model.layers[0].self_attn.q_proj.weight.grad
+    assert trunk is not None and trunk.abs().sum() > 0
+    print(f"ok  one training step: move {mloss.item():.3f} face {floss.item():.3f} "
+          f"acc {facc.item():.1%}, gradients reach the trunk")
+    return model, 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
+    ceil = D.teacher_greedy(net, "cpu", [3], 8, np.random.default_rng(5), max_steps=4,
+                            amp=False)
+    assert 0.0 <= ceil[3] <= 1.0
+    per_pos, exact = D.state_accuracy(model, head, net, "cpu", [3], 8,
+                                      np.random.default_rng(6), 6, 0.35, False)
+    assert "0:" in per_pos and 0.0 <= exact <= 1.0
+    print(f"ok  eval paths run: student {rates[3]:.0%}, teacher {ceil[3]:.0%}, "
+          f"state {per_pos}")
+
+
+def test_solved_start(net):
+    """A start state that is already solved must produce EOS and nothing else."""
+    states0 = C.SOLVED_ARR[None, :].copy()
+    batch = D.rollout_batch(net, states0, "cpu", 8, 0.35, 0.0,
+                            np.random.default_rng(7), amp=False)
+    mask = batch["posmask"].numpy()[0]
+    assert mask.sum() == 1
+    assert batch["soft"].numpy()[0, mask][0, D.EOS_SLOT] == 1.0
+    print("ok  a solved start state yields exactly one EOS position")
+
+
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    with tempfile.TemporaryDirectory() as tmp:
+        teacher = make_teacher(tmp)
+        test_token_tables()
+        test_pack_alignment(teacher)
+        test_targets(teacher)
+        test_hard_targets(teacher)
+        test_solved_start(teacher)
+        model, head = test_train_step(teacher)
+        test_eval(teacher, model, head)
+    print("\nall smoke tests passed")

code/train_policy_distill.py

@@ -0,0 +1,644 @@
+"""
+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.
+    """
+
+    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):
+    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()
+
+    sel = face_logits.float()[posmask]
+    face_loss = F.cross_entropy(sel.reshape(-1, len(C.COLORS)), faces[posmask].reshape(-1))
+    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
+
+
+# --------------------------------------------------------------------------
+# 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)
+    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()),
+                            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)
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(
+            list(model.parameters()) + list(head.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, head, 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(), "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")
+                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()

code/vast_bootstrap.sh

@@ -49,7 +49,7 @@ echo "building venv from $PYBIN ($($PYBIN -V 2>&1))"
 # because a failed C-extension build falls back to an upstream-git rebuild --
 # observed on a run that then never imported the module at all.
 NEEDS_SOLVER=1
-if [ -n "${EVAL_VALUE_ONLY:-}${TRAIN_VALUE_ITER:-}" ]; then NEEDS_SOLVER=""; fi
+if [ -n "${EVAL_VALUE_ONLY:-}${TRAIN_VALUE_ITER:-}${TRAIN_POLICY_DISTILL:-}" ]; then NEEDS_SOLVER=""; fi
 
 # kociemba's setup.py predates modern setuptools' removal of install_layout.
 /opt/venv/bin/pip install -q "setuptools<60" wheel
@@ -97,6 +97,93 @@ if [ -n "${EVAL_VALUE_ONLY:-}" ]; then
   exit 0
 fi
 
+# Policy distillation from a trained value model. Needs no dataset and no solver:
+# the targets are the teacher's scores for the 18 children of each visited state,
+# and the auxiliary facelet targets come from applying moves. So, like value
+# iteration, this goes straight from a GPU check to training.
+#
+# DISTILL_ARMS runs several arms back to back on one box, which is the point of
+# the experiment rather than a convenience -- the arms differ only in the loss, so
+# comparing them across two rentals would add a hardware variable to a difference
+# of a few points:
+#   full  soft targets + the state-tracking head (the whole proposal)
+#   soft  soft targets only            -- isolates what the head contributes
+#   hard  teacher argmin, no soft mass -- isolates soft-vs-hard from the teacher
+if [ -n "${TRAIN_POLICY_DISTILL:-}" ]; then
+  cd "${CODE_DIR:-/workspace/code}"
+  mkdir -p checkpoints
+  echo "=== GPU check ==="
+  POLICY_HIDDEN="${HIDDEN:-512}" POLICY_LAYERS="${LAYERS:-10}" POLICY_BATCH="${BATCH:-256}" \
+  /opt/venv/bin/python - <<'PYPD'
+import os, sys, torch
+if not torch.cuda.is_available():
+    sys.exit("ABORT: no CUDA device")
+h, l, b = (int(os.environ[k]) for k in ("POLICY_HIDDEN", "POLICY_LAYERS", "POLICY_BATCH"))
+vram = torch.cuda.get_device_properties(0).total_memory / 1e9
+print(f"gpu: {torch.cuda.get_device_name(0)} "
+      f"sm_{''.join(map(str, torch.cuda.get_device_capability(0)))} {vram:.0f}GB "
+      f"| torch {torch.__version__}")
+# The real config at the real batch, for the reason recorded in the main block: a
+# scaled-down probe passes on a GPU that cannot hold the model, so the failure
+# surfaces only once the run has been paid for.
+try:
+    from transformers import LlamaConfig, LlamaForCausalLM
+    m = LlamaForCausalLM(LlamaConfig(
+        vocab_size=28, hidden_size=h, intermediate_size=h * 4, num_hidden_layers=l,
+        num_attention_heads=8, num_key_value_heads=8, max_position_embeddings=83,
+        tie_word_embeddings=True)).cuda()
+    opt = torch.optim.AdamW(m.parameters(), lr=1e-4)
+    ids = torch.randint(0, 28, (b, 83), device="cuda")
+    with torch.autocast("cuda", dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16):
+        m(input_ids=ids, labels=ids).loss.backward()
+    opt.step()
+    torch.cuda.synchronize()
+    print(f"gpu smoke test: full step at batch {b} OK, "
+          f"peak {torch.cuda.max_memory_allocated()/1e9:.1f}GB of {vram:.0f}GB")
+except torch.OutOfMemoryError:
+    sys.exit(f"ABORT: batch {b} of a {l}-layer/{h}-wide model does not fit in {vram:.0f}GB")
+except Exception as e:
+    sys.exit(f"ABORT: GPU cannot run the model ({type(e).__name__}: {e})")
+PYPD
+
+  echo "=== fetching teacher ${TEACHER_REPO:-briscoooe/tiny-cube-value-iter} ==="
+  TEACHER="$(TR="${TEACHER_REPO:-briscoooe/tiny-cube-value-iter}" \
+             TF="${TEACHER_FILE:-checkpoints/value_iter/value_iter_best.pt}" \
+             /opt/venv/bin/python -c 'import os; from huggingface_hub import hf_hub_download; print(hf_hub_download(os.environ["TR"], os.environ["TF"], token=os.environ.get("HF_TOKEN")))')"
+  echo "teacher: $TEACHER"
+
+  # The highest score any arm can reach, printed before a single gradient step.
+  # The student copies the teacher's one-step choice while the teacher descends
+  # with the true state supplied by the engine, so the teacher's greedy rate caps
+  # the student. Cheap, and it makes an undertrained result impossible to confuse
+  # with a saturated one.
+  /opt/venv/bin/python train_policy_distill.py --teacher "$TEACHER" --ceiling-only \
+    --eval-depths "${EVAL_DEPTHS:-3,6,10}" --eval-n "${EVAL_N:-100}"
+
+  for ARM in $(echo "${DISTILL_ARMS:-full}" | tr ',' ' '); do
+    case "$ARM" in
+      full) ARM_FLAGS="--state-weight ${STATE_WEIGHT:-1.0}" ;;
+      soft) ARM_FLAGS="--state-weight 0" ;;
+      hard) ARM_FLAGS="--state-weight ${STATE_WEIGHT:-1.0} --hard-targets" ;;
+      *) echo "unknown arm: $ARM" >&2; exit 1 ;;
+    esac
+    echo "=== arm $ARM ($ARM_FLAGS) ==="
+    /opt/venv/bin/python train_policy_distill.py \
+      --teacher "$TEACHER" $ARM_FLAGS \
+      --hidden "${HIDDEN:-512}" --layers "${LAYERS:-10}" --heads 8 \
+      --batch-size "${BATCH:-256}" --lr "${LR:-3e-4}" --max-steps "${STEPS:-12000}" \
+      --rollout-len "${ROLLOUT_LEN:-16}" --scramble-depth "${WALK_DEPTH:-12}" \
+      --tau "${TAU:-0.35}" --explore-frac "${EXPLORE_FRAC:-0.2}" \
+      --buffer "${BUFFER:-32}" --updates-per-rollout "${UPDATES_PER_ROLLOUT:-4}" \
+      --eval-every "${EVAL_EVERY:-1000}" --eval-n "${EVAL_N:-100}" \
+      --eval-depths "${EVAL_DEPTHS:-3,6,10}" \
+      --out "checkpoints/policy_distill_$ARM" \
+      --hub-repo "${HUB_REPO:-}"
+  done
+  echo "=== done ==="
+  exit 0
+fi
+
 # Value iteration needs no dataset and no solver: it derives its own targets by
 # bootstrapping from the solved state. So this mode skips both the kociemba gate
 # and the multi-million-pair generation that dominates every other run's startup