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