File size: 2,198 Bytes
e791b16 | 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 | # PPO loss implementation with optional off‑policy correction (M2PO)
import torch
from torch.nn import functional as F
def compute_ppo_loss(
policy_log_probs: torch.Tensor, # (B, seq_len)
old_log_probs: torch.Tensor, # (B, seq_len)
advantages: torch.Tensor, # (B, seq_len)
clip_eps: float = 0.2,
use_m2po: bool = False,
m2_threshold: float = 2.0,
) -> torch.Tensor:
"""Compute PPO (or M2PO) loss.
Parameters
----------
policy_log_probs : torch.Tensor
Log‑probs of actions under the current policy.
old_log_probs : torch.Tensor
Log‑probs recorded during rollout.
advantages : torch.Tensor
Advantage estimates (reward‑to‑go minus baseline).
clip_eps : float, default 0.2
Standard PPO clipping epsilon.
use_m2po : bool, default False
If True, applies the second‑moment trust‑region (M2PO) instead of clipping.
m2_threshold : float, default 2.0
Maximum allowed value of (ratio ** 2) mean for M2PO.
"""
# ratio = pi(a|s) / pi_old(a|s)
ratio = torch.exp(policy_log_probs - old_log_probs)
if use_m2po:
# second‑moment constraint
m2 = (ratio ** 2).mean()
mask = ratio < m2_threshold
loss = -torch.min(
ratio * advantages,
torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * advantages,
)
loss = (loss * mask.float()).sum() / mask.float().sum().clamp(min=1.0)
else:
# standard PPO clipping
clipped_ratio = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps)
loss = -torch.min(ratio * advantages, clipped_ratio * advantages).mean()
return loss
def compute_m2po_loss(
policy_log_probs: torch.Tensor,
old_log_probs: torch.Tensor,
advantages: torch.Tensor,
clip_eps: float = 0.2,
m2_threshold: float = 2.0,
) -> torch.Tensor:
"""Convenience wrapper: compute_ppo_loss with use_m2po=True.
Returns a scalar torch.Tensor.
"""
return compute_ppo_loss(
policy_log_probs,
old_log_probs,
advantages,
clip_eps=clip_eps,
use_m2po=True,
m2_threshold=m2_threshold,
)
|