| """Task-agnostic PPO update for MiniJEV (actor = choice logits, critic = CLS value head). |
| |
| Why PPO (and not DQN) for MiniJEV: |
| * the model already outputs a distribution over *dynamic* choices; PPO trains |
| exactly that distribution, DQN would turn scores into Q-values and lose the |
| probabilistic reading that calibration relies on; |
| * it starts cleanly from the supervised policy (same head), with a KL anchor to |
| the supervised reference so early noisy advantages do not wreck it; |
| * on-policy + vectorised envs keeps it simple and stable at small scale. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def gae(rewards: np.ndarray, values: np.ndarray, dones: np.ndarray, last_value: np.ndarray, |
| gamma: float, lam: float) -> tuple[np.ndarray, np.ndarray]: |
| """rewards/values/dones: (T, N). Returns (advantages, returns).""" |
| T = rewards.shape[0] |
| adv = np.zeros_like(rewards) |
| last = np.zeros_like(last_value) |
| for t in reversed(range(T)): |
| nv = last_value if t == T - 1 else values[t + 1] |
| nonterminal = 1.0 - dones[t] |
| delta = rewards[t] + gamma * nv * nonterminal - values[t] |
| last = delta + gamma * lam * nonterminal * last |
| adv[t] = last |
| return adv, adv + values |
|
|
|
|
| def ppo_update(model, opt, data: dict, cfg: dict, ref_model=None, policy_coef: float = 1.0) -> dict: |
| """data: state (dict of tensors), question_ids, choice_ids, actions, old_logp, adv, ret (flattened).""" |
| n = len(data["actions"]) |
| mb = int(cfg.get("minibatch", 1024)) |
| clip = cfg.get("clip", 0.2) |
| stats = {"pg": 0.0, "vf": 0.0, "entropy": 0.0, "kl_ref": 0.0, "approx_kl": 0.0, "clipfrac": 0.0} |
| count = 0 |
| model.train() |
| for _ in range(int(cfg.get("epochs", 2))): |
| perm = torch.randperm(n, device=data["actions"].device) |
| for i in range(0, n, mb): |
| idx = perm[i:i + mb] |
| state = {k: v[idx] for k, v in data["state"].items()} |
| out = model(state, data["question_ids"][idx], data["choice_ids"][idx]) |
| logp_all = F.log_softmax(out["logits"], -1) |
| logp = logp_all.gather(1, data["actions"][idx, None]).squeeze(1) |
| ratio = (logp - data["old_logp"][idx]).exp() |
| adv = data["adv"][idx] |
| pg = -torch.min(ratio * adv, ratio.clamp(1 - clip, 1 + clip) * adv).mean() |
| vf = 0.5 * (out["value"] - data["ret"][idx]).pow(2).mean() |
| p = logp_all.exp() |
| ent = -(p * logp_all.masked_fill(~out["choice_mask"], 0.0)).sum(-1).mean() |
| loss = policy_coef * pg + cfg.get("vf_coef", 0.5) * vf - cfg.get("ent_coef", 0.001) * ent |
| kl_ref = torch.zeros(()) |
| if ref_model is not None and cfg.get("kl_ref_coef", 0.0) > 0: |
| with torch.no_grad(): |
| ref_logp = F.log_softmax(ref_model(state, data["question_ids"][idx], |
| data["choice_ids"][idx])["logits"], -1) |
| kl_ref = (ref_logp.exp() * (ref_logp - logp_all)).masked_fill(~out["choice_mask"], 0.0).sum(-1).mean() |
| loss = loss + cfg["kl_ref_coef"] * kl_ref |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.get("grad_clip", 0.5)) |
| opt.step() |
| with torch.no_grad(): |
| stats["pg"] += float(pg) |
| stats["vf"] += float(vf) |
| stats["entropy"] += float(ent) |
| stats["kl_ref"] += float(kl_ref) |
| stats["approx_kl"] += float((data["old_logp"][idx] - logp).mean()) |
| stats["clipfrac"] += float(((ratio - 1).abs() > clip).float().mean()) |
| count += 1 |
| return {k: v / max(1, count) for k, v in stats.items()} |
|
|
