Advantage Shaping as Surrogate Reward Maximization: Unifying Pass@K Policy Gradients

TL;DR

This note reconciles two seemingly distinct approaches to policy gradient optimization for the Pass@K objective in reinforcement learning with verifiable rewards: direct REINFORCE-style methods, and advantage-shaping techniques that directly modify GRPO.

Abstract

This note reconciles two seemingly distinct approaches to policy gradient optimization for the Pass@K objective in reinforcement learning with verifiable rewards: (1) direct REINFORCE-style methods, and (2) advantage-shaping techniques that directly modify GRPO. We show that these are two sides of the same coin. By reverse-engineering existing advantage-shaping algorithms, we reveal that they implicitly optimize surrogate rewards. We specifically interpret practical "hard-example up-weighting" modifications to GRPO as reward-level regularization. Conversely, starting from surrogate reward objectives, we provide a simple recipe for deriving both existing and new advantage-shaping methods. This perspective provides a lens for RLVR policy gradient optimization beyond our original motivation of Pass@K.

Figures (5)

• Figure 1: Comparison of effective gradient weights for the two Pass@K methods: GRPO$_K$ ($A_K^\pm$ in Eq. \ref{['eq:sadegh vs vanilla']}) and $\text{$\widetilde{\text{GRPO}}_K$}$ ($\widetilde{A}_K^\pm$ in Eq. \ref{['eq:wenlong vs vanilla']}). These weights scale the gradients of correct (solid lines) and wrong (dashed lines) responses by a factor $A^+ \cdot \hat{\rho}$ and $A^- \cdot (1-\hat{\rho})$, respectively. The scores are plotted against the empirical 0/1 rate ($\hat{\rho}$) for a fixed sample size of $N=16$ and varying values of $K$. For $K=1$, both weights coincide with the vanilla GRPO weights $\sqrt{\widehat{\rho}(1-\widehat{\rho})}$ (Eq. \ref{['eq:GRPO']}).
• Figure 2: Combined Pass@K effective gradient scores for $N=16$ and $K\in\{2,4,6,8,10\}$. Green uses Eq. \ref{['eq:combo direct']}; Blue uses Eq. \ref{['eq:combo bydance']}; the dotted gray line is the baseline $(1-\widehat{\rho})A^\pm$ (which we call skew-R). For $K=2,4$ the two combinations differ slightly at mid/high $\widehat{\rho}$ (blue is modestly larger); for $K\ge 6$ they nearly coincide across $\widehat{\rho}$. This is because, for most $\widehat{\rho}$, the Pass@K scalers vanish for $\widehat{\rho}<1-K/N$, reducing both combinations to the same envelope $(1-\widehat{\rho})\sqrt{\widehat{\rho}(1-\widehat{\rho})}$.
• Figure 3: Comparison of (regularized) surrogate objectives (GRPO, Skew-R, and Entropy-Augmented). (a) Effective Gradient Weight: Empirical gradient's scaling factor as a function of the empirical reward $\hat{\rho}$. The vanilla GRPO (blue, dashed) is symmetric. Skew-R (green) and the Entropy-augmented methods (red, purple) are asymmetric, suppressing gradients for high $\hat{\rho}$. Skew-R gradients are multiplied here by $2$, so all four curves have the same area under the curve. (b) Surrogate Reward Value: Corresponding surrogate reward functions, normalized so that they evaluate to $1$ at $\rho=1$.
• Figure 4: Comparison of absolute effective gradient weights (log scale) for Pass@K training with $N=16$. Each panel shows a different $K$ and plots the magnitude applied to correct (solid) and wrong (dashed) responses under three schemes: GRPO$_K$ (Eq. \ref{['eq:sadegh vs vanilla']}; green), $\text{$\widetilde{\text{GRPO}}_K$}$ (Eq. \ref{['eq:wenlong vs vanilla']}; blue), and the biased scaler (Eq. \ref{['eq:biased']}; red). For small $\hat{\rho}$, GRPO$_K$ assigns slightly larger weights than the biased scaler. As $\hat{\rho}$ increases, GRPO$_K$ enforces a hard zero once $\hat{\rho}>1-\frac{K-1}{N}$, whereas the biased scaler decays smoothly. $\text{$\widetilde{\text{GRPO}}_K$}$ applies a symmetric, example-level scale and often yields the largest mid-$\hat{\rho}$ weights until $\hat{\rho}_K$ saturates. Only the GRPO$_K$ exhibits asymmetric scaling between correct and wrong responses.
• Figure 5: A comparison of surrogate reward functions, plotted against the Pass@K reward $\rho_K$. The colored lines show the GRPO$_K$ surrogate, Eq. \ref{['eq:surrogate of grpok']} for various $K$ values. The black dashed line shows the (normalized) $\text{$\widetilde{\text{GRPO}}_K$}$ surrogate, $\frac{2}{\pi} \arcsin(\sqrt{\rho_K})$. (All curves are normalized to reach value $1$ at $\rho_K=1$.) The surrogates agree for $K=1$ (dashed and purple solid lines coincide), but for $K>1$ the GRPO$_K$ surrogates become concave providing stronger optimization incentive for hard examples.

Theorems & Definitions (8)

• Claim 1: Reverse-Engineering $\widetilde{\text{GRPO}}_K$
• Claim 2: $\widetilde{\text{GRPO}}_K$=RLOO on Surrogate Reward
• Corollary 1: GRPO as Surrogate Reward Optimization
• Claim 3: Skew-R$=$ Regularized Surrogate Reward Maximization
• Lemma 1: Relating $\widehat{f}_{K-1}^{ +}$, $\widehat{f}_{K-1}^{ -}$, and $\widehat{\rho}_{K}$
• proof
• Lemma 2: Conditional-gradient identity
• proof