File size: 3,836 Bytes
45507bb | 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 | """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()}
|