Skip to main content

compute_token_kl_schulman

rlox_rl_ops::kl::f32_ops

Function compute_token_kl_schulman 

Source 
pub fn compute_token_kl_schulman(
    log_probs_policy: &[f32],
    log_probs_ref: &[f32],
) -> Result<f32, RlOpsError>
Expand description

Token-level KL divergence using the Schulman (2020) estimator: sum(exp(log_p - log_q) - (log_p - log_q) - 1).

This is the estimator used by TRL (HuggingFace). It is unbiased and numerically more stable than the exact exp(log_p) * (log_p - log_q).