composer_replication/trainer/tests/test_kl_in_reward.py

"""Tests for k1-in-reward KL (F5 #1 fidelity fix; Composer-2 §4.1 / verl).

The load-bearing test is `test_apply_equals_fold_then_baseline`: it proves the
advantage adjustment `adv -= coef·(KL - group_mean(KL))` is EXACTLY equal to
folding the KL penalty into the reward and re-running GRPO's group-mean
baseline (with no std-norm). That equivalence is the entire justification for
adjusting advantages post-hoc instead of forking TRL's reward→advantage code.
"""
from __future__ import annotations

import pytest
import torch

from composer_replication.trainer.kl_in_reward import (
    apply_kl_in_reward,
    k1_kl_penalty_per_sequence,
    k3_kl_penalty_per_sequence,
    kl_penalty_per_sequence,
)


# ---------------------------------------------------------------------
# Per-sequence KL estimators
# ---------------------------------------------------------------------


def test_k1_penalty_sums_masked_logp_diff():
    policy = torch.tensor([[0.0, -1.0, -2.0], [-0.5, -0.5, -0.5]])
    ref = torch.tensor([[0.0, -0.5, -1.0], [-1.0, -1.0, -1.0]])
    mask = torch.tensor([[1.0, 1.0, 0.0], [1.0, 1.0, 1.0]])  # row0 drops last token
    out = k1_kl_penalty_per_sequence(policy, ref, mask)
    # row0: (0-0) + (-1-(-0.5)) [+ masked 0] = -0.5
    # row1: (-0.5-(-1.0))*3 = +1.5
    torch.testing.assert_close(out, torch.tensor([-0.5, 1.5]))


def test_k1_can_be_negative_k3_cannot():
    """Structural difference: k1 is signed, k3 ≥ 0 (the whole reason they differ)."""
    policy = torch.tensor([[0.0, 0.0]])
    ref = torch.tensor([[1.0, 1.0]])  # ref > policy → Δ=ref-logp>0 → k1<0
    mask = torch.ones_like(policy)
    k1 = k1_kl_penalty_per_sequence(policy, ref, mask)
    k3 = k3_kl_penalty_per_sequence(policy, ref, mask)
    assert (k1 < 0).all(), "k1 = Σ(logp-ref) is negative when ref>logp"
    assert (k3 >= -1e-6).all(), "k3 (Schulman) is always non-negative"


def test_k3_leading_order_is_half_delta_squared():
    """For small Δ, k3 ≈ Δ²/2 — the minor-delta claim in make_dr_grpo_config."""
    policy = torch.tensor([[0.0, 0.0, 0.0]])
    ref = torch.tensor([[0.01, -0.02, 0.005]])
    mask = torch.ones_like(policy)
    k3 = k3_kl_penalty_per_sequence(policy, ref, mask)
    delta = ref - policy
    expected = (0.5 * delta**2).sum()
    torch.testing.assert_close(k3, expected.unsqueeze(0), atol=1e-4, rtol=1e-3)


def test_dispatch_and_unknown_estimator():
    policy = torch.zeros(1, 2)
    ref = torch.ones(1, 2)
    mask = torch.ones(1, 2)
    torch.testing.assert_close(
        kl_penalty_per_sequence(policy, ref, mask, "k1"),
        k1_kl_penalty_per_sequence(policy, ref, mask),
    )
    with pytest.raises(ValueError, match="Unknown KL estimator"):
        kl_penalty_per_sequence(policy, ref, mask, "k2")


def test_penalty_shape_validation():
    with pytest.raises(ValueError, match="identical shape"):
        k1_kl_penalty_per_sequence(torch.zeros(1, 3), torch.zeros(1, 2), torch.zeros(1, 3))
    with pytest.raises(ValueError, match="must match"):
        k1_kl_penalty_per_sequence(torch.zeros(1, 3), torch.zeros(1, 3), torch.zeros(1, 2))


# ---------------------------------------------------------------------
# apply_kl_in_reward — the advantage adjustment
# ---------------------------------------------------------------------


def test_apply_coef_zero_is_identity():
    adv = torch.tensor([1.0, -1.0, 0.5, -0.5])
    kl = torch.tensor([2.0, 3.0, 1.0, 0.0])
    out = apply_kl_in_reward(adv, kl, num_generations=2, coef=0.0)
    torch.testing.assert_close(out, adv)


def test_apply_centers_kl_within_group():
    # Two groups of 2. coef=1. adv -= (KL - group_mean(KL)).
    adv = torch.zeros(4)
    kl = torch.tensor([1.0, 3.0, 10.0, 20.0])
    out = apply_kl_in_reward(adv, kl, num_generations=2, coef=1.0)
    # group0 mean=2 → centered [-1,+1] → adv-(-1,+1)=[1,-1]
    # group1 mean=15 → centered [-5,+5] → adv-(-5,+5)=[5,-5]
    torch.testing.assert_close(out, torch.tensor([1.0, -1.0, 5.0, -5.0]))


def test_apply_divisibility_validation():
    with pytest.raises(ValueError, match="multiple of num_generations"):
        apply_kl_in_reward(torch.zeros(5), torch.zeros(5), num_generations=2, coef=1.0)
    with pytest.raises(ValueError, match="identical shape"):
        apply_kl_in_reward(torch.zeros(4), torch.zeros(2), num_generations=2, coef=1.0)


@pytest.mark.parametrize("num_generations", [2, 3, 4])
@pytest.mark.parametrize("n_groups", [1, 2, 5])
def test_apply_equals_fold_then_baseline(num_generations, n_groups):
    """THE load-bearing property: adjusting baselined advantages by
    -coef·(KL - group_mean(KL)) equals folding -coef·KL into the reward and
    re-running GRPO's group-mean baseline (scale_rewards='none').

    This proves the post-hoc advantage adjustment IS exact k1-in-reward, not an
    approximation — the justification for not forking TRL's scoring code.
    """
    torch.manual_seed(0)
    g, k = num_generations, n_groups
    b = g * k
    coef = 0.137

    rewards = torch.randn(b)
    kl = torch.randn(b).abs()  # KL ≥ 0 in spirit, though sign-agnostic here

    # GRPO baseline (no std-norm): adv = reward - group_mean(reward).
    def group_baseline(x):
        means = x.view(k, g).mean(dim=1).repeat_interleave(g)  # (b,)
        return x - means

    advantages = group_baseline(rewards)

    # Reference: fold KL into reward, THEN baseline.
    folded_reward = rewards - coef * kl
    adv_fold_then_baseline = group_baseline(folded_reward)

    # Under test: adjust the ALREADY-baselined advantages.
    adv_adjusted = apply_kl_in_reward(advantages, kl, num_generations=g, coef=coef)

    torch.testing.assert_close(adv_adjusted, adv_fold_then_baseline, atol=1e-5, rtol=1e-5)


def test_apply_does_not_mutate_input():
    adv = torch.tensor([1.0, 2.0])
    adv_copy = adv.clone()
    apply_kl_in_reward(adv, torch.tensor([0.0, 1.0]), num_generations=2, coef=1.0)
    torch.testing.assert_close(adv, adv_copy)  # functional, not in-place