File size: 8,589 Bytes
026f63e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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", {})) # the checkpoint decides the choice labels
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) # TorchDecider may have moved it
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"]
|