scripts/train.py

#!/usr/bin/env python3
"""Train a Jevon decision model on procedurally generated maze and Snake states."""
from __future__ import annotations

import argparse
import json
import math
import random
import sys
import time
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from jevon.data import SamplerConfig, TaskSampler, build_corpus, collate, load_frozen
from jevon.losses import LossWeights, decision_loss, field_loss
from jevon.modeling import JevonConfig, JevonModel
from jevon.planner import default_iterations
from jevon.tokenizer import JevonTokenizer


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out", default="runs/jevon-s")
    p.add_argument("--steps", type=int, default=6000)
    p.add_argument("--states-per-batch", type=int, default=12)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--warmup", type=int, default=300)
    p.add_argument("--weight-decay", type=float, default=0.01)
    p.add_argument("--grad-clip", type=float, default=1.0)
    p.add_argument("--seed", type=int, default=17)
    p.add_argument("--device", default="mps")
    p.add_argument("--d-model", type=int, default=384)
    p.add_argument("--planner-dim", type=int, default=96)
    p.add_argument("--planner-grad-steps", type=int, default=8)
    p.add_argument("--planner-lr-mult", type=float, default=10.0,
                   help="learning-rate multiplier for the planner and field head")
    p.add_argument("--field-target", choices=("log", "linear"), default="log",
                   help="shape of the distance target; linear keeps neighbour "
                        "contrast uniform, which is what greedy descent needs")
    p.add_argument("--descent-weight", type=float, default=0.0,
                   help="hinge the field's 4-neighbour ordering toward the true "
                        "downhill step; this is what greedy descent reads")
    p.add_argument("--descent-margin", type=float, default=1.0,
                   help="required neighbour separation, in units of one true step")
    p.add_argument("--score-smoothing", type=float, default=0.0,
                   help="spread score targets over neighbouring levels (0 = one-hot)")
    p.add_argument("--min-plus-field", action="store_true",
                   help="read the distance off a min-plus recurrence, making the "
                        "field traversable by construction")
    p.add_argument("--expose-field", action="store_true",
                   help="let the text side read the planner's own distance read-out")
    p.add_argument("--field-weight", type=float, default=1.0,
                   help="weight on the dense planner distance-field loss")
    p.add_argument("--iteration-jitter", type=float, default=0.5,
                   help="planner budget is scaled by U(1-j, 1+j) during training")
    p.add_argument("--snake-fraction", type=float, default=0.45)
    p.add_argument("--eval-every", type=int, default=250)
    p.add_argument("--eval-file", default="data/frozen/val.jsonl")
    p.add_argument("--log-every", type=int, default=25)
    p.add_argument("--resume", default=None)
    p.add_argument("--no-planner", action="store_true")
    return p.parse_args()


def lr_at(step: int, args) -> float:
    if step < args.warmup:
        return args.lr * (step + 1) / args.warmup
    progress = (step - args.warmup) / max(1, args.steps - args.warmup)
    return args.lr * (0.05 + 0.95 * 0.5 * (1 + math.cos(math.pi * min(progress, 1.0))))


def bucket_batch(sampler: TaskSampler, states: int, rng: random.Random) -> list:
    """One batch holds a single game and board size: exact planner budget, no padding."""
    cfg = sampler.cfg
    snake = rng.random() < cfg.snake_fraction
    size = rng.choice(list(cfg.snake_sizes if snake else cfg.maze_sizes))
    topology = None if snake else rng.choice(list(cfg.maze_topologies))
    out = []
    guard = 0
    while len(out) < states and guard < states * 40:
        guard += 1
        seed = sampler._draw_seed()
        if snake:
            from jevon.data import sample_snake_state
            sample = sample_snake_state(seed, size, sampler.rng, cfg.augment)
        else:
            from jevon.data import sample_maze_state
            sample = sample_maze_state(seed, size, topology, sampler.rng, cfg.augment)
        if sample is not None:
            out.append(sample)
    return out


@torch.no_grad()
def evaluate(model, states, tokenizer, device, args) -> dict:
    model.eval()
    totals: dict[str, list[float]] = {}
    by_game: dict[str, list[float]] = {}
    for start in range(0, len(states), args.states_per_batch):
        group = states[start:start + args.states_per_batch]
        batch = collate(group, tokenizer, device=device)
        steps = default_iterations(int(batch["board_size"].max()))
        cells = None
        if model.planner is not None and "boards" in batch:
            cells = model.encode_boards(batch["boards"], batch["passable"], steps, 0)
        logits = model(batch, iterations=steps, grad_steps=0, cells=cells)
        _, stats = decision_loss(logits, batch)
        if cells is not None and "dist_field" in batch:
            # Scored against the model's own target shape, so field_r2 means
            # "how well did it fit what it was asked to fit". It is therefore
            # NOT comparable across --field-target settings; descent accuracy
            # from probe.py is the shape-independent metric for that.
            _, f_stats = field_loss(model.supervised_field(cells),
                                    batch["dist_field"],
                                    batch["passable"], batch["board_size"],
                                    args.field_target,
                                    args.descent_weight, args.descent_margin)
            stats.update(f_stats)
        for key, value in stats.items():
            v = float(value)
            if v == v:
                totals.setdefault(key, []).append(v)
        target, valid = batch["target_probs"], batch["candidate_valid"]
        picked = logits.argmax(-1)
        hit = target.gather(1, picked.unsqueeze(1)).squeeze(1)
        correct = (hit >= target.max(-1).values - 1e-6).float()
        for meta, ok, qtype in zip(batch["meta"], correct.tolist(),
                                   batch["question_type"].tolist()):
            if qtype == 0:                       # choice == the gameplay decision
                by_game.setdefault(f"{meta['game']}_action", []).append(ok)
        del valid
    model.train()
    out = {k: sum(v) / len(v) for k, v in totals.items() if not k.startswith("_")}
    if "_f_n" in totals:                       # pool R^2 instead of averaging it
        n = sum(totals["_f_n"])
        sse, tsum, tsq = (sum(totals[k]) for k in ("_f_sse", "_f_sum", "_f_sumsq"))
        centred = max(tsq - tsum * tsum / max(n, 1.0), 1e-9)
        out["field_r2"] = 1.0 - sse / centred
        out["field_mae"] = sum(totals["_f_abs"]) / max(n, 1.0)
    out.update({k: sum(v) / len(v) for k, v in by_game.items() if v})
    return out


def balanced_score(metrics: dict) -> float:
    """How well the checkpoint plays *both* games, for model selection.

    Eval loss is dominated by the ~970 boolean questions in the frozen split,
    so the ~170 action questions every headline number comes from barely move
    it -- a checkpoint that stopped playing snake can still own the lowest
    loss, and one did. ``min`` rather than the mean because the failure being
    guarded against is exactly the trade of one game for the other, and a mean
    would average it away.

    A missing metric scores zero rather than being skipped: a run that reports
    no snake_action has not earned a balanced checkpoint.
    """
    return min(metrics.get("maze_action", 0.0), metrics.get("snake_action", 0.0))


def main() -> None:
    args = parse_args()
    torch.manual_seed(args.seed)
    rng = random.Random(args.seed)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    device = torch.device(args.device if torch.backends.mps.is_available()
                          or args.device != "mps" else "cpu")

    sampler_cfg = SamplerConfig(snake_fraction=args.snake_fraction)
    train_sampler = TaskSampler(sampler_cfg, "train", args.seed)

    tok_path = out_dir / "tokenizer.json"
    if tok_path.exists():
        tokenizer = JevonTokenizer.load(tok_path)
    else:
        print("building tokenizer vocabulary ...", flush=True)
        tokenizer = JevonTokenizer.build(build_corpus(train_sampler, 300),
                                         max_vocab=4096, min_count=1)
        tokenizer.save(tok_path)
    print(f"vocab={len(tokenizer)} device={device}")

    cfg = JevonConfig(vocab_size=len(tokenizer), d_model=args.d_model,
                      planner_dim=args.planner_dim,
                      planner_grad_steps=args.planner_grad_steps,
                      expose_field=args.expose_field,
                     min_plus_field=args.min_plus_field,
                      use_planner=not args.no_planner)
    model = JevonModel(cfg).to(device)
    if args.resume:
        model.load_state_dict(torch.load(args.resume, map_location=device))
    print(f"parameters: {model.num_parameters()/1e6:.2f}M")
    (out_dir / "config.json").write_text(json.dumps(cfg.to_dict(), indent=2))
    # config.json is the *architecture*; this is the *recipe*. Without it a run
    # directory records what was built but not how it was trained, and the
    # flags that distinguish these runs from each other -- --min-plus-field,
    # --planner-lr-mult, --field-weight, --seed -- live only in the shell
    # history of whoever launched them. argv is recorded verbatim alongside the
    # parsed values because the two disagree in the way that matters: argv says
    # what was asked for, the parsed dict says what the defaults turned it into.
    (out_dir / "args.json").write_text(json.dumps(
        {"argv": sys.argv, "args": vars(args)}, indent=2, default=str))

    # The planner is one small block applied several hundred times, and only the
    # last `grad_steps` applications carry gradient, so its updates arrive at a
    # completely different scale from the transformer's. Measured on a fixed set
    # of 11x11 mazes, 200 steps of field-only training: lr 3e-4 reaches R^2 0.24,
    # lr 3e-3 reaches 0.64. It needs its own learning rate.
    planner_names = {n for n, _ in model.named_parameters()
                     if n.startswith("planner.") or n.startswith("field_head.")}
    groups = []
    for is_planner in (False, True):
        for wants_decay in (True, False):
            params = [p for n, p in model.named_parameters()
                      if (n in planner_names) == is_planner
                      and (p.dim() >= 2) == wants_decay]
            if params:
                groups.append({
                    "params": params,
                    "weight_decay": args.weight_decay if wants_decay else 0.0,
                    "lr_mult": args.planner_lr_mult if is_planner else 1.0})
    opt = torch.optim.AdamW(groups, lr=args.lr, betas=(0.9, 0.95))

    eval_states = load_frozen(args.eval_file) if Path(args.eval_file).exists() else []
    print(f"frozen eval states: {len(eval_states)}")

    weights = LossWeights(score_smoothing=args.score_smoothing)
    history, best = [], float("inf")
    # Eval loss is dominated by the ~970 boolean questions in the split; the
    # 170 action questions that every headline number is computed from barely
    # move it. Selecting on loss alone therefore picked a checkpoint that had
    # traded one game away -- a step reached maze_action 0.979 while snake_action
    # fell to 0.297, below its own uniform control, and still took best.pt
    # because the aggregate loss was the lowest so far. `min` rather than the
    # mean because the failure being guarded against is exactly that trade.
    best_balanced = -float("inf")
    running: dict[str, float] = {}
    start_time = time.time()

    for step in range(args.steps):
        for group in opt.param_groups:
            group["lr"] = lr_at(step, args) * group["lr_mult"]
        states = bucket_batch(train_sampler, args.states_per_batch, rng)
        batch = collate(states, tokenizer, device=device)
        size = int(batch["board_size"].max())
        # Jitter skews upward: too few iterations teaches the model to decide
        # before the distance field has arrived, which is the failure we are fixing.
        jitter = 1.0 + rng.uniform(-0.25 * args.iteration_jitter, args.iteration_jitter)
        steps_t = max(8, int(default_iterations(size) * jitter))

        cells = None
        if model.planner is not None and "boards" in batch:
            cells = model.encode_boards(batch["boards"], batch["passable"],
                                        steps_t, args.planner_grad_steps)
        logits = model(batch, iterations=steps_t, cells=cells)
        loss, stats = decision_loss(logits, batch, weights)
        # Dense planner supervision. Without it the recurrence gets gradient
        # from one distance question per state, which measurably fails to
        # teach it distances (linear probe R^2 ~= 0.21).
        if cells is not None and args.field_weight > 0 and "dist_field" in batch:
            f_loss, f_stats = field_loss(model.supervised_field(cells),
                                         batch["dist_field"], batch["passable"],
                                         batch["board_size"], args.field_target,
                                         args.descent_weight, args.descent_margin)
            loss = loss + args.field_weight * f_loss
            stats.update(f_stats)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
        opt.step()

        for key, value in stats.items():
            v = float(value)
            if v == v:
                running[key] = 0.9 * running.get(key, v) + 0.1 * v

        if (step + 1) % args.log_every == 0:
            rate = (step + 1) / (time.time() - start_time)
            print(f"step {step+1:5d}/{args.steps} loss {running['loss']:.4f} "
                  f"acc {running['acc']:.3f} choice {running.get('acc_choice', float('nan')):.3f} "
                  f"bool {running.get('acc_boolean', float('nan')):.3f} "
                  f"tvd {running['tvd']:.3f} T={steps_t} {rate:.2f} it/s", flush=True)

        if eval_states and ((step + 1) % args.eval_every == 0 or step + 1 == args.steps):
            metrics = evaluate(model, eval_states, tokenizer, device, args)
            metrics["step"] = step + 1
            history.append(metrics)
            print("  eval " + " ".join(
                f"{k}={v:.4f}" for k, v in metrics.items() if k != "step"), flush=True)
            (out_dir / "history.json").write_text(json.dumps(history, indent=2))
            if metrics["loss"] < best:
                best = metrics["loss"]
                torch.save(model.state_dict(), out_dir / "best.pt")
                (out_dir / "best.json").write_text(json.dumps(metrics, indent=2))
            balanced = balanced_score(metrics)
            if balanced > best_balanced:
                best_balanced = balanced
                torch.save(model.state_dict(), out_dir / "balanced.pt")
                (out_dir / "balanced.json").write_text(
                    json.dumps({**metrics, "balanced": balanced}, indent=2))
        if (step + 1) % 500 == 0:
            torch.save(model.state_dict(), out_dir / "last.pt")

    torch.save(model.state_dict(), out_dir / "last.pt")
    print(f"done in {(time.time()-start_time)/60:.1f} min; best eval loss {best:.4f}; "
          f"best balanced action acc {best_balanced:.4f}")


if __name__ == "__main__":
    main()