UCPO: Uniform Correct Policy Optimization

Official implementation of Uniform Correct Policy Optimization.

UCPO is a reinforcement learning framework for training LLMs on mathematical reasoning. It introduces an inverse Q-learning (IQ) based advantage estimator that promotes uniform coverage over correct responses, encouraging the model to explore diverse solution paths rather than collapsing to a single high-reward mode.

Built on top of verl and E3 .

Standard GRPO-style algorithms weight updates by reward, causing the policy to concentrate probability mass on a few high-reward responses. UCPO instead redistributes gradient such that spreads mass uniformly across all correct responses.

Key hyperparameters:

• tau — diversity scaling; controls the uniformity of the correct-response distribution

# Install verl and dependencies
cd verl
pip install -e .

DeepSeek-R1-Distill-Qwen-1.5B (4 GPUs):

bash scripts_c/run_ucpo_1_5b.sh

DeepSeek-R1-Distill-Qwen-7B (8 GPUs):

bash scripts_c/run_ucpo_7b.sh

Key training settings:

Checkpoints are saved under ../checkpoint_ds/.

The evaluation pipeline has three stages:

1. Inference — generate rollouts from a checkpoint:

bash scripts_c/infer.sh

2. Scoring — compute per-response accuracy:

bash scripts_c/eval_all.sh

3. Metrics — aggregate Pass@K results:

bash scripts_c/metric_all.sh

Training uses a 10K math reasoning dataset (dataset/train_data_10k.parquet) with a held-out validation split (dataset/valid_data.parquet). Each example contains a prompt and a ground-truth answer in \boxed{} format.

@article{lochab2026uniformcorrectpolicyoptimizationbreaking,
      title={Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity}, 
      author={Anamika Lochab and Bolian Li and Ruqi Zhang},
      year={2026},
      eprint={2605.00365},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2605.00365}, 
}