Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity

TL;DR

This paper introduces Uniform-Correct Policy Optimization (UCPO), a method that enhances diversity and multi-sample coverage in reinforcement learning with verifiable rewards by promoting uniform probability distribution over correct solutions.

Contribution

The paper formalizes the cause of diversity collapse in RLVR and proposes UCPO, a novel optimization technique that improves diversity and coverage without sacrificing accuracy.

Findings

UCPO improves Pass@K and diversity across multiple models and benchmarks.

UCPO achieves up to +10% on AIME24 Pass@64.

UCPO increases equation-level diversity by up to 45%.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has achieved substantial gains in single-attempt accuracy (Pass@1) on reasoning tasks, yet often suffers from reduced multi-sample coverage (Pass@K), indicating diversity collapse. We identify a structural cause for this degradation: common RLVR objectives, such as GRPO, are indifferent to how probability mass is distributed among correct solutions. Combined with stochastic training dynamics, this indifference induces a self-reinforcing collapse, in which probability mass concentrates on a narrow subset of correct outputs while alternative valid solutions are suppressed. We formalize this collapse mechanism and further characterize the optimal policy structure under two complementary criteria: robustness and entropy-regularized optimality, which identify the Uniform-Correct Policy as uniquely optimal. Motivated by this analysis, we propose Uniform-Correct Policy Optimization (UCPO), a modification to GRPO that adds a conditional uniformity penalty on the policy's distribution over correct solutions. The penalty redistributes gradient signal toward underrepresented correct responses, encouraging uniform allocation of probability mass within the correct set. Across three models (1.5B-7B parameters) and five mathematical reasoning benchmarks, UCPO improves Pass@K and diversity while maintaining competitive Pass@1, achieving up to +10\% absolute improvement on AIME24 at Pass@64 and up to 45\% higher equation-level diversity within the correct set. The code is available at https://github.com/AnamikaLochab/UCPO.