Function k3_kl

src/post_training/grpo.py:31–34

Per-token unbiased, non-negative KL estimator (Schulman's k3) for KL(policy||ref).

(new_logp: torch.Tensor, ref_logp: torch.Tensor)

def k3_kl(new_logp: torch.Tensor, ref_logp: torch.Tensor) -> torch.Tensor:
    """Per-token unbiased, non-negative KL estimator (Schulman's k3) for KL(policy||ref)."""
    diff = ref_logp - new_logp
    return torch.exp(diff) - diff - 1.0


def grpo_loss(