fix(algorithms): honour the configured KL clamps in the reward-side KL

[tianyi-zhang-02]

What does this PR do ?

Makes the reward-side KL use the clamp values that are configured for it, instead of falling back to calculate_kl's defaults.

Issues

None to close.

Usage

No new config. loss_fn.kl_input_clamp_value / loss_fn.kl_output_clamp_value now reach the reward-side KL as well as the loss-side one.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests (unit only)
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

GeneralizedAdvantageEstimator and ReinforcePlusPlusAdvantageEstimator both take a ClippedPGLossConfig and read three things off it:

self.use_kl_in_reward = loss_config.use_kl_in_reward
self.kl_coef = loss_config.reference_policy_kl_penalty
self.kl_type = loss_config.reference_policy_kl_type

The same config also carries kl_input_clamp_value and kl_output_clamp_value, and ClippedPGLossFn passes both into calculate_kl. The estimators did not, so with use_kl_in_reward the configured bounds were silently replaced by the function's own defaults, 20.0 and 10.0.

Same config, same KL approximation, two different bounds depending on which side of the algorithm reads it — and no error either way.

On the tests. Getting them to assert anything took one non-obvious step worth flagging for review: both estimators normalize the advantage globally at the end, so a KL that is constant across the batch normalizes away to all-zeros no matter what the clamp does. My first version of these tests passed with the fix reverted for exactly that reason. The reference logprobs have to diverge by a different amount at each position before the clamp is observable at all.

Mutation-tested: dropping the two kwargs from either call site, and hardcoding a clamp back to its old default, each turn one of these red.

Adjacent to my #3512, which wraps these returns in an AdvantageResult — the tests accept a bare tensor, a (advantages, returns) tuple, or the dataclass, so they survive that rebase whichever lands first. :)

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.