fix(ppo): preserve raw KL so rollout/kl logging is correct

[EazyReal]

Problem

compute_advantages_and_returns writes per-token KL into rollout_data["kl"], which is what gets logged as the KL metric.

The PPO branch then did:

for reward, k in zip(old_rewards, kl, strict=False):
    k *= kl_coef
    if cp_rank == 0:
        k[-1] += reward
    rewards.append(k)

k is the same tensor already stored in rollout_data["kl"], so logging sees kl_coef-scaled KL with the scalar reward added on the last token. GRPO / GSPO / CISPO / R++ do not mutate kl this way.

Change

Build token-level rewards out of place:

token_level_rewards = per_token_kl * kl_coef
if cp_rank == 0:
    token_level_rewards[-1] += reward
rewards.append(token_level_rewards)

PPO advantages/returns still use the scaled+reward tensor. rollout_data["kl"] stays the raw metric. Other estimators are untouched.

Test plan

• tests/test_ppo_kl_metric.py — after PPO, rollout_data["kl"] still matches compute_approx_kl (registered in pr-test.yml.j2)

[EazyReal]

@zhuzilin could you review this now-cleaned version? PPO reward shaping was mutating the raw KL tensor before metrics, so rollout/kl logging could report shaped rewards instead of KL. The fix keeps raw KL separate and mirrors the local-k pattern used by the reinforce loss helpers.

[EazyReal]

@zhuzilin refreshed on latest main and checks are green. This keeps raw KL separate from token-level rewards so rollout/kl logging stays correct while PPO still applies the KL penalty. Could you review when you have bandwidth?