Research map / Reinforcement learning / Policy gradient and actor-critic

Pass@K Policy Optimization: Solving Harder Reinforcement Learning Problems

NeurIPS 2025 · Reinforcement learning › Policy gradient and actor-critic · approach: Variance-reduced policy gradients · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes PKPO, unbiased low-variance estimators of pass@k and its gradient for RL with multiple samples per prompt, implemented as a batch reward transformation that lets k be annealed during training.

Problem

Standard RL with n independent samples per prompt optimizes pass@1 (per-sample expected reward), under-utilizing the sampling budget and limiting exploration on hard problems; the goal is to directly optimize pass@k (expected max reward over k samples) for arbitrary k ≤ n.

Approach

Derives unbiased estimators of pass@k (and max-reward@k for continuous rewards) and their policy gradients by averaging over all k-subsets of a batch, yielding a per-sample reward transformation (sloo / sloo_minus_one) computable in O(k+n log n) that can be dropped into any policy gradient method and tuned to any k_opt ≤ n.

Key results

• On MATH with Gemma2-9B, k_opt=8 reaches 48.37 pass@16 vs 42.59 for k_opt=1 (standard pass@1 RL), and k_opt matching k_eval gives the best pass@k in most columns.
• On MATH with Llama3.1-8B, k_opt=8 gives 61.88 pass@16 vs 56.83 for k_opt=1 and annealing k_opt from 8 to 1 keeps pass@1 at the baseline level while improving pass@k for k>1.
• On ARC-AGI-1 (easy subset), conventional pass@1 optimization stalls while PKPO unblocks learning and improves pass@k across all k including k=1.

Measured effects reported in the paper

Surprising findings

• Entropy regularization, a common exploration baseline, sacrifices pass@1 and only slightly improves pass@k, while PKPO outperforms it.
• Without basic LOO mean-centering the training diverges, so the baseline is essential even for the k_opt=1 setting.

Limitations

• Variance of the estimator increases as k_opt approaches n because fewer subsets are available, making gains at large k less reliable.
• Gradient estimator for k<n is biased in the baseline (not in the estimator itself) when using leave-one-out-1 baseline, per the derivation.
• Experiments are limited to 2B-9B models on three task sets with three seeds.

Experimental setup

Datasets: MATH MBPP ARC-AGI-1

Benchmarks: MATH test split HumanEval ARC-AGI-1 test split

Metrics: pass@k pass@1 cumulative solve rate entropy

Baselines: pass@1 RL (k_opt=1) Tang et al. (k_opt=n) max-baseline method entropy regularization PPO

Builds on: PPO pass@k estimator of Chen et al. (Codex) tangloo (Tang et al. 2025 pass@n RL with leave-one-out max baseline) REINFORCE expected-maximum objective (Chung et al.) leave-one-out baseline

Closest papers by meaning

• The Perils of Optimizing Learned Reward Functions: Low Training Error Does Not Guarantee Low Regret (ICML 2025 · similarity 0.87)
• To the Max: Reinventing Reward in Reinforcement Learning (ICML 2024 · similarity 0.87)
• Oracle-Efficient Reinforcement Learning for Max Value Ensembles (NeurIPS 2024 · similarity 0.87)
• Correlated Proxies: A New Definition and Improved Mitigation for Reward Hacking (ICLR 2025 · similarity 0.87)
• Multi-Reward Best Policy Identification (NeurIPS 2024 · similarity 0.87)
• Reward-Free Kernel-Based Reinforcement Learning (ICML 2024 · similarity 0.87)