| """ |
| Trains a value function: cube state -> distance to solved. |
| |
| ## Why this and not another policy |
| |
| Four attempts to make a sequence-predicting policy reach past ~8 moves all |
| failed (capacity, depth, chunked inference, training on its own rollouts), and |
| beam search over that policy raised depth 6 to 94% while leaving depth 15+ at |
| exactly zero. Beam search explores only what the policy already ranks highly; it |
| has no independent notion of which states are closer to solved, so where the |
| policy is wrong there is nothing better to find. |
| |
| A value function supplies that missing compass. This is DeepCubeA's shape |
| (Agostinelli et al., Nature MI 2019), which solves the full cube with a network |
| of roughly this size by learning distance-to-solved and searching over it. |
| |
| ## Why this is cheaper here than in the paper |
| |
| DeepCubeA has no expert, so it bootstraps the value function by approximate |
| value iteration -- expensive, and the bulk of its compute. We *do* have an |
| expert: Kociemba's solution length is the distance, and every training row |
| already carries it. So this is plain supervised learning on labels we generate |
| for free. |
| |
| The labels are Kociemba two-phase lengths, which are upper bounds rather than |
| true optimal distances. That makes the heuristic inadmissible (A* over it is not |
| guaranteed optimal), which does not matter here: the goal is to find *a* |
| solution, not the shortest. The lengths were separately confirmed monotonic along |
| solution paths, so the signal is consistent. |
| |
| ## Classification, not regression |
| |
| Distance is predicted as a distribution over 0..MAX_DIST rather than a scalar. |
| Regression to a mean is actively harmful on this task: the state space is |
| overwhelmingly distance 18-21, so a squared-error model collapses toward that |
| mode and loses exactly the near-solved discrimination the search depends on. |
| |
| **The number to watch is mean absolute error.** Search needs the heuristic to |
| rank neighbouring states correctly; at MAE around 1 move that works, and at MAE |
| of several moves it cannot, whatever the loss curve says. |
| """ |
| import argparse, json, math, random, sys, time |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch.utils.data import Dataset, DataLoader |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| import cube_tokenizer as T |
|
|
| MAX_DIST = 26 |
|
|
|
|
| class ValueDataset(Dataset): |
| def __init__(self, rows): |
| self.rows = rows |
|
|
| def __len__(self): |
| return len(self.rows) |
|
|
| def __getitem__(self, i): |
| state, dist = self.rows[i] |
| return torch.tensor(T.encode_state(state)), torch.tensor(dist) |
|
|
|
|
| def load_value_rows(path, limit=0): |
| """Label is the solution length -- the distance -- regardless of how the row |
| was generated, so augmented and canonical rows are equally usable.""" |
| rows = [] |
| with open(path) as fh: |
| for line in fh: |
| if limit and len(rows) >= limit: |
| break |
| r = json.loads(line) |
| d = len(r["solution"].split()) |
| if 0 < d <= MAX_DIST: |
| rows.append((r["state"], d)) |
| return rows |
|
|
|
|
| class ValueNet(torch.nn.Module): |
| """Llama encoder over the 56-token state, mean-pooled, then a distance head.""" |
|
|
| def __init__(self, hidden, layers, heads): |
| super().__init__() |
| from transformers import LlamaConfig, LlamaModel |
| cfg = LlamaConfig( |
| vocab_size=T.VOCAB_SIZE, hidden_size=hidden, intermediate_size=hidden * 4, |
| num_hidden_layers=layers, num_attention_heads=heads, num_key_value_heads=heads, |
| max_position_embeddings=T.MAX_SEQ, pad_token_id=T.PAD, |
| ) |
| self.encoder = LlamaModel(cfg) |
| self.head = torch.nn.Linear(hidden, MAX_DIST + 1) |
|
|
| def forward(self, ids): |
| h = self.encoder(input_ids=ids).last_hidden_state.mean(dim=1) |
| return self.head(h) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, rows, device, batch_size=512): |
| """Reports MAE and exact-match, plus MAE split by distance band. |
| |
| The split matters: an aggregate MAE can look fine while the near-solved band |
| -- the only place the heuristic has to be sharp for search to make progress |
| -- is useless. |
| """ |
| model.eval() |
| abs_err, exact, n = 0.0, 0, 0 |
| bands = defaultdict(lambda: [0.0, 0]) |
| for i in range(0, len(rows), batch_size): |
| chunk = rows[i:i + batch_size] |
| ids = torch.tensor([T.encode_state(s) for s, _ in chunk], device=device) |
| true = torch.tensor([d for _, d in chunk], device=device) |
| pred = model(ids).argmax(-1) |
| err = (pred - true).abs().float() |
| abs_err += err.sum().item() |
| exact += (pred == true).sum().item() |
| n += len(chunk) |
| for e, t in zip(err.tolist(), true.tolist()): |
| key = "1-8" if t <= 8 else ("9-14" if t <= 14 else "15+") |
| bands[key][0] += e |
| bands[key][1] += 1 |
| model.train() |
| parts = " ".join(f"{k}:{bands[k][0]/bands[k][1]:.2f}" for k in ("1-8", "9-14", "15+") if bands[k][1]) |
| return abs_err / n, exact / n, parts |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--data", required=True) |
| p.add_argument("--val-data", default="") |
| p.add_argument("--limit", type=int, default=0) |
| p.add_argument("--hidden", type=int, default=512) |
| p.add_argument("--layers", type=int, default=8) |
| p.add_argument("--heads", type=int, default=8) |
| p.add_argument("--batch-size", type=int, default=1024) |
| p.add_argument("--lr", type=float, default=6e-4) |
| p.add_argument("--max-steps", type=int, default=30000) |
| p.add_argument("--warmup", type=int, default=200) |
| p.add_argument("--eval-every", type=int, default=1000) |
| p.add_argument("--eval-n", type=int, default=4096) |
| p.add_argument("--out", default="checkpoints/value") |
| p.add_argument("--hub-repo", default="") |
| p.add_argument("--seed", type=int, default=0) |
| args = p.parse_args() |
|
|
| torch.manual_seed(args.seed) |
| random.seed(args.seed) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| rows = load_value_rows(args.data, args.limit) |
| val = load_value_rows(args.val_data, args.eval_n) if args.val_data else rows[-args.eval_n:] |
| if not args.val_data: |
| rows = rows[:-args.eval_n] |
| print(f"train {len(rows)} | val {len(val)} | device {device}", flush=True) |
|
|
| model = ValueNet(args.hidden, args.layers, args.heads).to(device) |
| print(f"params {sum(q.numel() for q in model.parameters())/1e6:.1f}M", flush=True) |
|
|
| loader = DataLoader(ValueDataset(rows), batch_size=args.batch_size, shuffle=True, |
| drop_last=True, num_workers=2) |
| opt = torch.optim.AdamW(model.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)))) |
| use_amp = device == "cuda" |
| amp_dtype = torch.bfloat16 if (use_amp and torch.cuda.is_bf16_supported()) else torch.float16 |
| scaler = torch.amp.GradScaler("cuda", enabled=use_amp and amp_dtype is torch.float16) |
|
|
| Path(args.out).mkdir(parents=True, exist_ok=True) |
| step, t0, best = 0, time.time(), 1e9 |
| while step < args.max_steps: |
| for ids, dist in loader: |
| ids, dist = ids.to(device), dist.to(device) |
| with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp): |
| loss = F.cross_entropy(model(ids), dist) |
| scaler.scale(loss).backward() |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| scaler.step(opt) |
| scaler.update() |
| sched.step() |
| opt.zero_grad(set_to_none=True) |
| step += 1 |
| if step % 100 == 0: |
| print(f"step {step}/{args.max_steps} loss {loss.item():.4f} " |
| f"{step/(time.time()-t0):.1f} it/s", flush=True) |
| if step % args.eval_every == 0 or step >= args.max_steps: |
| mae, exact, parts = evaluate(model, val, device) |
| print(f" step {step} MAE {mae:.3f} moves exact {exact:.1%} " |
| f"by true distance: {parts}", flush=True) |
| torch.save({"state_dict": model.state_dict(), |
| "hidden": args.hidden, "layers": args.layers, |
| "heads": args.heads}, f"{args.out}/value.pt") |
| if mae < best: |
| best = mae |
| torch.save({"state_dict": model.state_dict(), |
| "hidden": args.hidden, "layers": args.layers, |
| "heads": args.heads}, f"{args.out}/value-best.pt") |
| if args.hub_repo: |
| try: |
| from huggingface_hub import HfApi |
| HfApi().upload_file(path_or_fileobj=f"{args.out}/value.pt", |
| path_in_repo="value.pt", repo_id=args.hub_repo) |
| except Exception as e: |
| print(f" hub push failed (continuing): {e}", flush=True) |
| if step >= args.max_steps: |
| break |
| print(f"done. best MAE {best:.3f} moves. weights in {args.out}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
