#!/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()