Enhancing Rubric-based RL via Self-Distillation

Research Question and Approach

Rubric-based reinforcement learning (RL) is a powerful technique for aligning LLMs on open-ended tasks, but it suffers from two primary failure modes: Unexplored Criteria (UC), where the model never satisfies a criterion, and Suppressed Criteria (SC), where the model satisfies a criterion but the overall reward aggregation obscures the signal, leading to negative advantages. Existing methods often introduce a train-inference mismatch by using external guidance during training that is absent during inference. The authors propose Criterion-Distilled Policy Optimization (CriPO), which uses on-policy self-distillation to provide dense, token-level supervision without creating a distribution gap between training and inference.

Methodology

CriPO integrates two mechanisms into the standard Group Relative Policy Optimization (GRPO) framework:

• Behavior Injection for UC: The model constructs a 'criterion-injection' self-teacher that revises the best-performing rollout to satisfy missing criteria. It then uses a forward-KL divergence loss to distill these behaviors into the student policy, filtering for the most informative tokens to avoid noise.
• Advantage Flipping for SC: The model identifies tokens responsible for satisfying suppressed criteria by comparing the student with a 'counterfactual' teacher that is prompted to remove those specific behaviors. It then flips the advantages of these tokens from negative to positive, ensuring that useful patterns are reinforced rather than penalized by the aggregate reward.

By addressing both UC and SC, CriPO allows models to learn more effectively from multi-dimensional rubric feedback. Experiments on medical and scientific benchmarks show that CriPO consistently outperforms standard GRPO and other exploration-focused methods, achieving superior final performance while requiring approximately 2x fewer optimization steps. This demonstrates that fine-grained, on-policy supervision is a more efficient way to handle complex, open-ended alignment tasks than relying solely on scalar reward signals.