Skip to main content

compute_token_kl_schulman

rlox_rl_ops::kl::f64_ops

Function compute_token_kl_schulman 

Source 
pub fn compute_token_kl_schulman(
    log_probs_policy: &[f64],
    log_probs_ref: &[f64],
) -> Result<f64, 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).