| |
|
|
| import torch |
| from torch.nn import functional as F |
|
|
|
|
| def compute_ppo_loss( |
| policy_log_probs: torch.Tensor, |
| old_log_probs: torch.Tensor, |
| advantages: torch.Tensor, |
| 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 = torch.exp(policy_log_probs - old_log_probs) |
|
|
| if use_m2po: |
| |
| 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: |
| |
| 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, |
| ) |
|
|
