Function compute_reference_kl

rl_engine/testing/reference_ops.py:95–106

Compute the common GRPO/PPO reference KL approximation.

(
    current_logps: torch.Tensor,
    ref_logps: torch.Tensor,
    mask: torch.Tensor | None = None,
)

def compute_reference_kl(
    current_logps: torch.Tensor,
    ref_logps: torch.Tensor,
    mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Compute the common GRPO/PPO reference KL approximation."""

    diff = ref_logps.float() - current_logps.float()
    kl = torch.exp(diff) - diff - 1.0
    if mask is not None:
        kl = kl.masked_fill(~_bool_mask(mask, device=kl.device), 0.0)
    return kl


def summarize_kernel_drift(