A Negative KL Penalty Doesn't Mean Your AI Agent's Reward Calculation Is Broken

[ym p]

A code agent generates a response. You compare its token log probabilities with those from a fixed reference policy, calculate a penalty, and discover that the “KL term” is negative.

The tempting fix is one line:

penalty = max(0, penalty)

Before adding it, check what the number actually represents. An individual sampled log probability ratio can be negative even though the exact KL divergence between normalized distributions is nonnegative. Clamping the former changes the quantity you are estimating.

This distinction is easy to miss when implementing reward bookkeeping for execution-feedback code agents. We can examine it using Python's standard library and a two-outcome distribution.

Start with a distribution small enough to enumerate​

Let the current policy assign probabilities (0.1, 0.9) to two possible outcomes. Let the reference assign (0.5, 0.5).

For a sampled outcome, calculate log(policy_probability / reference_probability):

Outcome	Current policy	Reference	Sampled log ratio
first	0.1	0.5	−1.609438
second	0.9	0.5	+0.587787

The first outcome is less likely under the current policy than under the reference, so its ratio is below one and its log is negative. There is no contradiction.

Exact forward KL averages these log ratios under the current policy. The sampling distribution matters: an unweighted average of the two rows is not the same calculation.

import math

policy = (0.1, 0.9)
reference = (0.5, 0.5)
ratios = [math.log(p / q) for p, q in zip(policy, reference)]

exact_kl = math.fsum(p * r for p, r in zip(policy, ratios))
clipped_expectation = math.fsum(
    p * max(0.0, r) for p, r in zip(policy, ratios)
)

assert ratios[0] < 0.0 < ratios[1]
assert exact_kl > 0.0
assert clipped_expectation > exact_kl

print(f"first sample log ratio: {ratios[0]:.6f}")
print(f"exact forward KL:       {exact_kl:.6f}")
print(f"clipped expectation:    {clipped_expectation:.6f}")

Output:

first sample log ratio: -1.609438
exact forward KL:       0.368064
clipped expectation:    0.529008

Clipping removes the negative contribution while leaving the positive one. Even in this tiny example, it produces a different expectation. It is a new estimator choice, not a numerical repair of the original expression.

Where RLEF uses a related quantity​

RLEF: Grounding Code LLMs in Execution Feedback with Reinforcement Learning describes iterative code generation using public execution feedback, with private evaluation determining final correctness. Its reward includes a reference-policy log-ratio term. For that term, the paper uses geometric-mean response probabilities rather than products of token probabilities, so the response log ratio is a mean of token log ratios. The reference is the initial policy, distinct from PPO's changing old policy.

The helper below isolates that arithmetic. It assumes both lists score the same sampled tokens in their corresponding contexts. It neither runs code nor implements PPO.

def mean_sampled_log_ratio(policy_logps, reference_logps):
    if not policy_logps or len(policy_logps) != len(reference_logps):
        raise ValueError("expected nonempty equal-length lists")
    if not all(math.isfinite(x) and x <= 0.0
               for x in (*policy_logps, *reference_logps)):
        raise ValueError("expected finite natural log probabilities <= 0")
    return math.fsum(
        p - q for p, q in zip(policy_logps, reference_logps)
    ) / len(policy_logps)

def shaped_reward(base_reward, policy_logps, reference_logps, beta):
    if not math.isfinite(base_reward):
        raise ValueError("base reward must be finite")
    if not math.isfinite(beta) or beta < 0.0:
        raise ValueError("beta must be finite and nonnegative")
    ratio = mean_sampled_log_ratio(policy_logps, reference_logps)
    return base_reward - beta * ratio

# This sampled response is less likely under the current policy.
ratio = mean_sampled_log_ratio([-4.0], [-2.0])
reward = shaped_reward(0.0, [-4.0], [-2.0], 0.5)
assert ratio == -2.0 and reward == 1.0
print(f"sampled ratio={ratio:.1f}; shaped reward={reward:.1f}")

Here the reference term increases the sampled shaped reward. That does not mean the code passed a test. The base reward and regularization contribution are separate components and should be logged separately.

Also, the exact distribution calculation in the first section is not a claim that every mean-over-token estimator is an unbiased estimate of sequence-level KL. Token contexts, response length, stopping rules, and the chosen normalization all matter. The narrow point is that nonnegativity of exact KL does not justify clamping every observed log ratio.

Mean and sum produce different length dependence​

For token ratios (1, 1), the mean is 1 and the sum is 2. Repeating those values eight times preserves the mean but multiplies the sum by eight.

current = [-1.0, -3.0]
fixed_ref = [-2.0, -4.0]

short_reward = shaped_reward(0.0, current, fixed_ref, 0.5)
repeated_reward = shaped_reward(0.0, current * 8, fixed_ref * 8, 0.5)
assert short_reward == repeated_reward == -0.5

wrong_sum_reward = -0.5 * math.fsum(
    p-q for p, q in zip(current * 8, fixed_ref * 8)
)
assert wrong_sum_reward == -8.0
print(f"mean-based reward: {repeated_reward:.1f}")
print(f"sum-based reward:  {wrong_sum_reward:.1f}")

This is an arithmetic invariance test using repeated arrays, not a claim that repeating real text leaves its conditional token probabilities unchanged.

Summing log probabilities is appropriate when computing an ordinary sequence log probability. The bug is substituting that sum into a formula whose contract explicitly calls for the mean, without adjusting the objective or coefficient.

Tests that catch the common mistakes​

Test signs, normalization, zero coefficients, and invalid input independently. A successful end-to-end run can hide each of these errors.

# Equal token probabilities give zero reference contribution.
assert mean_sampled_log_ratio([-1.0, -2.0], [-1.0, -2.0]) == 0.0

# A positive sampled ratio reduces the reward.
assert shaped_reward(1.0, [-1.0], [-3.0], 0.5) == 0.0

# beta=0 preserves the supplied task reward.
assert shaped_reward(-1.0, [-1.0], [-9.0], 0.0) == -1.0

# A pure arithmetic swap reverses the signed ratio.
a = mean_sampled_log_ratio(current, fixed_ref)
b = mean_sampled_log_ratio(fixed_ref, current)
assert a == -b
assert current == [-1.0, -3.0] and fixed_ref == [-2.0, -4.0]

for p, q in [([], []), ([-1.0], [-1.0, -2.0]),
             ([float("nan")], [-1.0]), ([0.1], [-1.0])]:
    try:
        mean_sampled_log_ratio(p, q)
    except ValueError:
        pass
    else:
        raise AssertionError("invalid probability inputs accepted")
print("All reward arithmetic checks passed")

The swap check does not imply forward KL is symmetric: it holds the token sample fixed and tests subtraction only. Changing the distribution that supplies samples changes the expectation.

Keep the diagnostic components visible​

When reviewing reward logs, store at least the task reward, sampled mean log ratio, coefficient, and final shaped reward separately. Check that both policies scored the same tokens with the intended contexts. Check whether the implementation uses a fixed reference or a changing behavior policy. Document the length normalization.

A reward increase from this term should not be reported as improved program correctness. Evaluate correctness with its own tests.

I maintain TensorDrill, an AI coding practice site. Its RLEF turn-reward exercise combines this arithmetic with execution-reward branches in a deterministic Python task. The complete experiment above can also be run independently of the site.

Before “fixing” a negative KL value, identify whether you calculated a distribution-level divergence, a sampled log ratio, or a length-normalized response statistic. They are related quantities with different guarantees.