| """PPO fine-tuning of a supervised MiniJEV Snake policy (Phase 2). |
| |
| Rollouts: `n_envs` games stepped together, batched inference, sampled actions. |
| Reward comes from `env.reward` in the config (food +, death -, tiny/zero per |
| step). The value head is warmed up first (it was never trained by behaviour |
| cloning), and a KL penalty to the frozen supervised policy keeps early updates |
| from destroying what imitation learned. Resumable from `rl_last.pt`. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from minijev.calibration.metrics import calibration_report, fit_temperature |
| from minijev.datasets.batching import BatchEncoder |
| from minijev.datasets.decision_data import DecisionData |
| from minijev.inference.runtime import TorchDecider |
| from minijev.runs import RunLogger, environment_info |
| from minijev.tokenizer import encode_batch |
| from minijev.training.ppo import gae, ppo_update |
| from minijev.training.supervised import make_model, predict_logits |
|
|
| from .env import SnakeConfig, VectorSnakeEnv |
| from .evaluate import model_policy, play_games |
| from .task import SnakeTask |
|
|
| QUESTION = SnakeTask.default_question |
|
|
|
|
| def train_ppo(cfg: dict, checkpoint: str, run_dir: str | None = None, resume: bool = False) -> dict: |
| rl = cfg.get("rl", {}) |
| seed = int(cfg.get("seed", 0)) |
| torch.manual_seed(seed) |
| torch.set_num_threads(int(rl.get("threads", 8))) |
| ck = torch.load(checkpoint, map_location="cpu", weights_only=False) |
| base_cfg = ck["config"] |
| task = SnakeTask(base_cfg.get("task", {})) |
| model = make_model(base_cfg, task) |
| model.load_state_dict(ck["model"]) |
| model.log_temperature.zero_() |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| model.to(device) |
| ref = copy.deepcopy(model).eval() |
| for p in ref.parameters(): |
| p.requires_grad_(False) |
|
|
| run_dir = Path(run_dir or f"runs/{cfg.get('name', 'snake')}-ppo") |
| log = RunLogger(run_dir, {**cfg, "init_checkpoint": checkpoint}, resume=resume) |
| opt = torch.optim.AdamW(model.parameters(), lr=rl.get("lr", 1e-4), weight_decay=0.0) |
| env_cfg = SnakeConfig.from_dict(cfg.get("env", {})) |
| n_envs, T = int(rl.get("n_envs", 64)), int(rl.get("rollout", 128)) |
| total_updates = int(rl.get("total_env_steps", 1_000_000)) // (n_envs * T) |
| start, best = 0, -1.0 |
| last_ck = run_dir / "checkpoints" / "rl_last.pt" |
| if resume and last_ck.exists(): |
| st = torch.load(last_ck, weights_only=False) |
| model.load_state_dict(st["model"]) |
| opt.load_state_dict(st["optimizer"]) |
| start, best = st["update"], st["best"] |
| print(f"[rl] resumed at update {start}") |
| envs = VectorSnakeEnv(n_envs, env_cfg, seed=seed * 7919 + start * n_envs) |
| mcfg = base_cfg["model"] |
| q = torch.as_tensor(encode_batch([QUESTION], mcfg.get("max_question_len", 32)), device=device).repeat(n_envs, 1) |
| c = torch.as_tensor(encode_batch(list(task.default_choices), mcfg.get("max_choice_len", 16)), |
| device=device)[None].repeat(n_envs, 1, 1) |
| gamma, lam = rl.get("gamma", 0.99), rl.get("gae_lambda", 0.95) |
| warm = int(rl.get("value_warmup_updates", 5)) |
| eval_every = int(rl.get("eval_every", 10)) |
| log.write_json("env.json", {"env": environment_info(), "updates": total_updates, "n_envs": n_envs, "rollout": T}) |
| ep_scores: list[int] = [] |
| t0 = time.time() |
| print(f"[rl] PPO: {total_updates} updates x {n_envs * T} steps, reward={cfg.get('env', {}).get('reward')}") |
| for upd in range(start, total_updates): |
| obs, acts, logps, vals, rews, dns = [], [], [], [], [], [] |
| model.eval() |
| for _ in range(T): |
| x = torch.as_tensor(task.encode_states(envs.states())["state_grid"], device=device) |
| with torch.no_grad(): |
| out = model({"state_grid": x}, q, c) |
| dist = torch.distributions.Categorical(logits=out["logits"]) |
| a = dist.sample() |
| r, d, infos = envs.step(a.cpu().tolist()) |
| ep_scores += [i["score"] for i in infos if i] |
| obs.append(x) |
| acts.append(a) |
| logps.append(dist.log_prob(a)) |
| vals.append(out["value"].cpu()) |
| rews.append(r) |
| dns.append(d.astype(np.float32)) |
| with torch.no_grad(): |
| x = torch.as_tensor(task.encode_states(envs.states())["state_grid"], device=device) |
| last_v = model({"state_grid": x}, q, c)["value"].cpu().numpy() |
| V = torch.stack(vals).numpy() |
| adv, ret = gae(np.stack(rews), V, np.stack(dns), last_v, gamma, lam) |
| adv_t = torch.as_tensor(adv.reshape(-1), device=device) |
| data = {"state": {"state_grid": torch.cat(obs)}, "question_ids": q.repeat(T, 1), |
| "choice_ids": c.repeat(T, 1, 1), "actions": torch.cat(acts), "old_logp": torch.cat(logps), |
| "adv": (adv_t - adv_t.mean()) / (adv_t.std() + 1e-8), "ret": torch.as_tensor(ret.reshape(-1), device=device)} |
| stats = ppo_update(model, opt, data, rl, ref_model=ref, policy_coef=0.0 if upd < warm else 1.0) |
| recent = ep_scores[-200:] |
| log.log(upd + 1, env_steps=(upd + 1) * n_envs * T, train_avg_score=float(np.mean(recent)) if recent else 0.0, |
| episodes=len(ep_scores), **{k: round(v, 5) for k, v in stats.items()}) |
| print(f" update {upd + 1}/{total_updates} | sampled-policy avg score {np.mean(recent) if recent else 0:.1f} " |
| f"| entropy {stats['entropy']:.3f} kl_ref {stats['kl_ref']:.4f} vf {stats['vf']:.4f} " |
| f"| {time.time() - t0:.0f}s", flush=True) |
| if (upd + 1) % eval_every == 0 or upd + 1 == total_updates: |
| dec = TorchDecider(model=model, cfg=base_cfg, device=str(device)) |
| g = play_games(model_policy(dec), int(rl.get("eval_games", 100)), cfg.get("env", {}), seed=777) |
| log.log(upd + 1, eval_avg_score=g["avg_score"], eval_fatal=g["fatal_move_rate"]) |
| print(f" [eval] greedy avg score {g['avg_score']:.2f} median {g['median_score']} " |
| f"fatal-move {g['fatal_move_rate']:.4%}", flush=True) |
| if g["avg_score"] > best: |
| best = g["avg_score"] |
| model.to(device) |
| torch.save({"model": model.state_dict(), "config": base_cfg, "update": upd + 1, "eval": g}, |
| run_dir / "checkpoints" / "rl_best.pt") |
| torch.save({"model": model.state_dict(), "optimizer": opt.state_dict(), "update": upd + 1, |
| "best": best}, last_ck) |
| return finalize_rl(cfg, base_cfg, run_dir, log) |
|
|
|
|
| def finalize_rl(cfg, base_cfg, run_dir: Path, log) -> dict: |
| """Re-calibrate the best RL policy on the validation split and evaluate it.""" |
| task = SnakeTask(base_cfg.get("task", {})) |
| best = torch.load(run_dir / "checkpoints" / "rl_best.pt", weights_only=False) |
| model = make_model(base_cfg, task) |
| model.load_state_dict(best["model"]) |
| data_dir = Path(base_cfg["data"]["dir"]) |
| mc = base_cfg["model"] |
| kw = dict(max_question_len=mc.get("max_question_len", 32), max_choice_len=mc.get("max_choice_len", 16)) |
| val = BatchEncoder(DecisionData.load(data_dir / "val.npz"), task, **kw) |
| test = BatchEncoder(DecisionData.load(data_dir / "test.npz"), task, **kw) |
| vl, vt = predict_logits(model, val, "cpu") |
| tl, tt = predict_logits(model, test, "cpu") |
| T = fit_temperature(vl, vt) |
| model.log_temperature.fill_(float(np.log(T))) |
| acc_ok = test.data.soft > 0 if test.data.soft is not None else None |
| rep = calibration_report(tl, tt, T, acceptable=acc_ok) |
| torch.save({"model": model.state_dict(), "config": base_cfg, "temperature": T}, run_dir / "checkpoints" / "final.pt") |
| games = play_games(model_policy(TorchDecider(run_dir / "checkpoints" / "final.pt", threads=4)), |
| cfg.get("eval", {}).get("games", 200), cfg.get("env", {})) |
| (run_dir / "games.json").write_text(json.dumps(games, indent=2)) |
| summary = {"params": model.num_parameters(), "temperature": T, "test_calibrated": rep, |
| "test_raw": calibration_report(tl, tt, 1.0, acceptable=acc_ok), "games": games, |
| "final_checkpoint": str(run_dir / "checkpoints" / "final.pt")} |
| log.write_json("summary.json", summary) |
| log.close() |
| print(f"[rl] final: avg score {games['avg_score']:.2f} | oracle-agreement {rep['accuracy']:.4f} " |
| f"| ECE vs oracle {rep['ece']:.4f} (T={T:.3f})") |
| return summary |
|
|
|
|
| __all__ = ["train_ppo"] |
|
|
