Enhancing Rubric-based RL via Self-Distillation

Abstract: Rubric-based RL has recently shown promise in improving LLMs on open-ended tasks. A widely recognized limitation of rubric-based RL is limited exploration: criteria that no rollout manages to satisfy (Unexplored Criteria, UC) receive no optimization signal. Recent methods address this by incorporating rubric information as external guidance during rollout, yet they introduce a train-inference mismatch: the policy is optimized on rollouts produced under external guidance while this guidance is absent at inference time, causing error accumulation through autoregressive decoding. Moreover, these exploration-focused approaches overlook a fundamentally different failure mode that we term Suppressed Criteria (SC) -- criteria that are satisfied by some rollouts yet whose learning signals are lost during optimization because scalar reward aggregation assigns them non-positive aggregate advantages. Our analysis reveals that SC are remarkably prevalent: over 57% of samples exhibit this failure mode throughout training, with an average of 1.8 SC per sample. To simultaneously address both UC and SC without introducing training-inference mismatch, we propose Criterion-Distilled Policy Optimization (CriPO), which enhances rubric-based RL via on-policy self-distillation. For UC, CriPO constructs a criterion-injection self-teacher and computes a localized forward-KL loss to inject missing behaviors into the policy. For SC, CriPO employs a counterfactual self-teacher to locate criterion-relevant tokens in negative-advantage rollouts and flips their token-level advantages to positive values, preserving useful patterns that would otherwise be suppressed. Experiments on medicine and science benchmarks demonstrate that CriPO consistently outperforms rubric-based RL, achieving stronger final performance with approximately $2\times$ fewer optimization steps.

• The paper introduces CriPO, a criterion-distilled policy optimization technique that mitigates unexplored and suppressed criteria in rubric-based RL.
• It utilizes localized forward-KL distillation and counterfactual advantage flipping to inject precise, criterion-relevant corrections without train–inference mismatch.
• Experimental results demonstrate that CriPO significantly improves compute efficiency and reward scores, outperforming established methods on multiple benchmarks.

Criterion-Distilled Policy Optimization (CriPO) for Rubric-based RL via Self-Distillation

Motivation and Failure Modes in Rubric-based RL

Rubric-based RL leverages criterion-decomposed rewards for training LLMs in open-ended domains where traditional verifiable supervision (e.g., exact answers, test cases) is unavailable. The canonical pipeline involves prompt-conditioned on-policy rollouts scored along explicit criteria using an LLM-as-Judge, followed by scalar-aggregated reward optimization via algorithms such as GRPO. Despite progress, two critical criterion-level failure modes undermine the reliability of reward signal propagation:

1. Unexplored Criteria (UC): Criteria with no satisfying rollouts in a group receive zero optimization signal, leading to unlearned behaviors in the policy.
2. Suppressed Criteria (SC): Criteria satisfied by some rollouts but receiving non-positive aggregate advantages due to scalar reward aggregation, resulting in systematic loss of criterion-relevant learning signals.

Statistical analysis reveals persistent and pronounced prevalence of these modes (Figure 1):

Figure 1: Left: Unexplored and suppressed criteria emerge in GRPO due to rollout limitations and reward aggregation. Right: UC and SC persist across training, highlighting structural inefficiency.

Previous methods aiming to improve exploration—e.g., rubric-conditioned generation (Zhou et al., 23 Aug 2025), hindsight experience replay (Zhang et al., 20 Mar 2026)—inject privileged information during training, resulting in a policy distribution mismatch at inference, with negligible improvement to SC.

The CriPO Algorithm: Unified Correction for Exploration and Suppression

CriPO (Criterion-Distilled Policy Optimization) directly addresses both failure modes via targeted, on-policy self-distillation. It extends GRPO's stable optimization backbone by selectively introducing dense, token-level corrections using on-policy self-teaching, entirely circumventing train–inference mismatch.

UC Correction: Localized Forward-KL Distillation

For unexplored criteria, CriPO constructs a criterion-injection self-teacher by prompting the LLM to minimally revise a high-quality (maximum-advantage) rollout to satisfy the unmet criteria. The student is then distilled via a forward-KL objective, but only at token positions with high KL contribution—identified by cumulative filtering (typically covering ≈35% tokens), substantially reducing noise:

Figure 2: CriPO first identifies UC/SC from rubric-evaluated rollouts. Missing UC behaviors are injected through localized OPSD loss.

Performance gains stem from focusing distillation solely on tokens where criterion satisfaction modulates behavior (as confirmed by cumulative KL plots), ensuring efficient transfer and minimal degradation from noisy supervision.

SC Correction: Counterfactual Advantage Flipping

Suppressed criteria are corrected by leveraging a counterfactual self-teacher. Given a negative-advantage, criterion-satisfying rollout, the self-teacher is prompted to revise or excise the parts that satisfy the suppressed criteria. Tokens are selected for advantage "flipping" based on two signals:

• Substantial log-probability drop for the sampled token under the counterfactual teacher;
• High teacher certainty for an alternative token at that position.

These tokens receive a positive local advantage, ensuring preservation, rather than penalization, of useful but suppressed behaviors:

Figure 3: Counterfactual self-teacher correctly localizes criterion-relevant tokens for advantage flipping, e.g., the token "Avoid".

Both corrections are integrated as an auxiliary loss (with strength $\beta$), yielding the overall training objective:

$$\mathcal{L}_\mathrm{CriPO} = \mathcal{L}_\mathrm{GRPO} + \beta\mathcal{L}_\mathrm{OPSD}$$

Experimental Analysis and Training Dynamics

CriPO demonstrates robust and consistent improvement over GRPO, HeRL, and vanilla OPSD across Qwen3-1.7B and Qwen3-4B on RaR-Medicine/Science, HealthBench, LLMEval-Med, and ResearchQA. Notably:

• Absolute improvement: $+3.2$ (1.7B) and $+1.4$ (4B) average points over GRPO.
• Optimization efficiency: Achieves best GRPO performance with $2\times$ fewer steps (Figure 4).

  Figure 4: CriPO exhibits superior compute efficiency, requiring significantly fewer steps to surpass GRPO convergence.

  

  Figure 5: CriPO achieves higher rewards, maintains greater entropy (exploration), and produces longer, content-rich responses.

CriPO delivers both higher peak reward and improved sample efficiency due to precise criterion-level supervision. Analysis reveals substantial reduction in both UC and SC prevalence, with each submodule (CriPO-U for UC, CriPO-S for SC) demonstrating its targeted effect (Figure 6):

Figure 6: Occurrence and mean count of UC and SC—CriPO variants reduce their target failure modes most effectively.

Qualitative Case Studies

CriPO encourages LLMs to produce more complete and rubric-relevant reasoning, as confirmed by both token-selection analysis and real task cases. For instance, CriPO localizes and preserves critical clinical reasoning tokens that standard GRPO suppresses, enabling more robust, explainable medical reasoning (Figure 7):

Figure 7: CriPO responses evidence deeper, more explicit clinical rationale by leveraging criterion-guided exploration.

Theoretical and Practical Implications

CriPO aligns with recent findings that dense self-distillation as a sole supervisory signal is unstable due to privileged information leakage and entropy collapse (Hübotter et al., 28 Jan 2026, Zhu et al., 11 May 2026). Instead, it integrates localized OPSD with reward-grounded PPO-style training, preserving optimization direction and stability. This approach can be extended to other complex RL settings where scalar reward ambiguity and sparse exploration jointly afflict policy learning. Further, CriPO’s explicit localization of reward attribution lays groundwork for fine-grained behavioral interventions and more interpretable RL training paradigms.

Figure 8: Token-wise forward KL is highest near boundaries; CriPO masks these regions to avoid noise in distillation.

Figure 9: Example prompt for rubric-based judge, providing criterion-wise reward signal for training and analysis.

Conclusion

CriPO introduces a milestone solution to two core limitations in rubric-based RL—exploration incompleteness and reward ambiguity—by integrating on-policy, criterion-conditioned self-distillation with stable policy optimization. The method consistently outperforms strong baselines both in effectiveness and compute efficiency, and generalizes across evaluation metrics and judge models. Its modular design sets the foundation for future rubric-aware RL frameworks, facilitating more trustworthy and interpretable LLM alignment.

Future Directions

Continued research should address:

• Automatic selection of criterion sets for self-distillation to further enhance scalability,
• Joint optimization with diverse judge architectures to increase robustness,
• Defending against overfit or spurious attribution effects in highly multi-criteria regimes,
• Formal extension to hierarchical or structured reward decompositions.

Overall, CriPO represents a significant advancement in aligning LLMs with nuanced, criterion-rich guidance in open-ended task spaces, and is directly extensible to broader domains involving complex, decomposed reward supervision.