On Advantage Estimates for Max@K Policy Gradients

2026-06-04

AI Summary

This work addresses the challenge of biased and ambiguous advantage estimation in sparse-reward reinforcement learning under the max@K setting, where inconsistencies among reward signals, baselines, and normalization lead to suboptimal policy updates. The authors propose a unified max@K policy gradient framework that introduces a novel Leave-Two-Out (L2O) baseline, enabling— for the first time—unbiased and precisely centered batch advantage estimation. They derive a standard advantage formulation valid under finite batch sizes, grounded in policy gradient theory and implemented via a quadratic-time algorithm suitable for group-wise post-training of large language models. Experimental results demonstrate that the L2O baseline substantially reduces gradient variance, outperforms non-centered alternatives in max@K tasks, and enhances both training stability and overall performance.

Abstract

Reinforcement learning with verifiable rewards is widely used for post-training reasoning models, but sparse outcome rewards make exploration difficult. A complementary approach is to optimize inference-time objectives such as pass@K and max@K directly, yet existing policy-gradient estimators for these objectives use different signals, baselines, and normalizations, making their relationships unclear. We study this issue through baseline design and advantage centering. Starting from the advantage estimator of a leading method in the field, we show that it is policy-gradient unbiased but yields a non-centered advantage. We then introduce a Leave-Two-Out baseline that preserves policy-gradient unbiasedness while making realized batch advantages exactly centered. The resulting method, MaxPO, has an efficient quadratic-time implementation and integrates naturally into group-based RL for LLM post-training. We further derive the canonical finite-batch advantage for max@K, providing a unified view of existing advantage estimators. Empirically, we verify that the L2O baseline reduces gradient variance and outperforms non-centered alternatives.

Problem

• reinforcement learning
• sparse rewards
• policy gradients
• max@K
• advantage estimation

Innovation

• max@K policy gradient
• advantage estimation
• Leave-Two-Out baseline
• gradient variance reduction
• LLM post-training