code/train_value.py

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