Enhancing Rubric-based RL via Self-Distillation

The Problem: Two Hidden Failures in Rubric-Based RL

Rubric-based reinforcement learning trains language models on open-ended tasks by scoring outputs against structured criteria—think grading an essay against a rubric. The approach has shown promise, but the authors identify two critical failure modes that existing methods miss or create.

The first is Unexplored Criteria (UC): some criteria never get satisfied in any rollout during training, so the model never learns to optimize for them. Prior work tries to fix this by injecting criterion information as guidance during rollout generation—but this creates a mismatch. The policy trains on guided rollouts that don't reflect what happens at inference, where guidance is absent. This mismatch causes error accumulation as the model decodes autoregressively.

The second failure, Suppressed Criteria (SC), is more subtle and overlooked. Some criteria do get satisfied in certain rollouts, but the learning signal disappears during optimization. Here's why: rubric-based RL typically converts criterion satisfaction into a scalar reward, then aggregates it into a single advantage value for the entire rollout. If a rollout satisfies criterion A but fails criterion B, the aggregate advantage can become non-positive, and the whole rollout gets treated as unhelpful—even though the model did learn something valuable about criterion A. The authors' analysis shows this is pervasive: over 57% of training samples suffer from SC, averaging 1.8 suppressed criteria per sample.

The Solution: Self-Distillation Without Train-Inference Mismatch

The authors propose Criterion-Distilled Policy Optimization (CriPO), which uses on-policy self-distillation to address both failure modes while keeping training and inference aligned.

For UC, CriPO creates a "criterion-injection self-teacher"—a version of the current policy that has been explicitly prompted to satisfy a missing criterion. The model then learns from this teacher using a localized forward-KL loss, injecting the missing behavior directly into the main policy without relying on external guidance at rollout time.

For SC, CriPO identifies which tokens in a rollout are "responsible" for satisfying a suppressed criterion using a counterfactual self-teacher. Once identified, those tokens' advantages are flipped from negative to positive, preserving the useful patterns that would otherwise be lost to scalar aggregation. The key insight is token-level granularity: instead of discarding an entire rollout, the method salvages the specific parts that worked.

Crucially, both mechanisms use the policy itself as the teacher, so there's no train-inference mismatch—the policy only sees the kind of guidance it will encounter at test time.

Impact and Results

On medicine and science benchmarks, CriPO consistently outperforms standard rubric-based RL. Notably, it achieves stronger final performance using approximately 2× fewer optimization steps, suggesting more efficient learning of criterion-relevant behaviors. The method is straightforward to integrate into existing rubric-based RL pipelines and directly targets a genuine inefficiency in current approaches—not just exploration, but also preservation of already-discovered useful patterns.