Enhancing Rubric-based RL via Self-Distillation

ListenThe Signal You're Throwing Away: Suppressed Criteria in Rubric RL10:25

Show notes

A look at CriPO — Criterion-Distilled Policy Optimization — which names and quantifies a hidden failure mode in rubric-based reinforcement learning. The authors show that over 57% of training samples suffer from "Suppressed Criteria," where good behaviors get erased by scalar reward aggregation, and fix it with on-policy self-distillation that reaches comparable performance in roughly half the optimization steps. Maya and Ethan unpack the mechanism, the ablations, and what token-level advantage shaping might mean for visual generative modeling. **Sources** - Enhancing Rubric-based RL via Self-Distillation

Subscribe to the channel: https://podcast.lumalabs.link/27a25f51ea81d41a8253f46c6133a35e/feed.xml

TL;DR

CriPO (Criterion-Distilled Policy Optimization) targets two failure modes of rubric-based GRPO for open-ended tasks: Unexplored Criteria (UC — no rollout satisfies a criterion, so no gradient signal exists) and Suppressed Criteria (SC — the criterion is satisfied by some rollout, but scalar reward aggregation assigns that rollout non-positive advantage and its useful pattern gets pushed down anyway). The paper reports SC affects >57% of samples with an average 1.8 SC per sample during training — a previously-uncharacterized bottleneck. CriPO addresses each with an on-policy self-distillation mechanism (criterion-injection self-teacher for UC via localized forward-KL; counterfactual self-teacher for SC via token-level advantage flipping) and reports outperforming rubric-based RL baselines on medicine and science benchmarks at ~2× fewer optimization steps.

Key claims

The failure mode split under rubric-based RL is two-way, not one-way: UC (unexplored criteria) is the well-known exploration limitation, and SC (suppressed criteria) is a distinct, quantitatively larger phenomenon where scalar reward aggregation actively erases useful behaviors that are being produced by some rollouts [§Abstract, §Analysis].
SC prevalence is substantial: >57% of samples exhibit SC throughout training, with an average of 1.8 SC per sample [§Abstract].
Prior exploration-focused rubric-based RL methods that inject rubric text as external guidance during rollout introduce a train-inference mismatch — the policy is optimized on guided rollouts but deployed without guidance, causing error accumulation through autoregressive decoding [§Abstract].
CriPO’s UC fix: build a criterion-injection self-teacher (same model, rubric injected as external guidance) and compute a localized forward-KL loss on only the smallest set of student tokens contributing to X% of the total KL — concentrating the distillation signal on tokens where the criterion is meaningfully expressed [§Method].
CriPO’s SC fix: build a counterfactual self-teacher by asking the teacher to revise the student’s own answer under the criterion, then compare teacher and student logits token-by-token. Tokens where the teacher confidently prefers a different token get their per-token GRPO advantage flipped to a fixed positive value (0.1 in the paper), rescuing the useful pattern from suppression [§Method].
Both losses (localized forward-KL for UC, per-token modified GRPO for SC) are combined with a hyperparameter controlling OPSD-loss influence; the composite trains on-policy with no rubric injection at inference — avoiding the train-inference mismatch that OTC-style methods introduce [§Method].
Empirically CriPO beats OPSD baselines by a large margin and matches or exceeds GRPO at ~2× fewer optimization steps on medicine and science benchmarks [§Abstract, §Experiments].

Method

CriPO extends GRPO with two on-policy self-distillation components, each targeting a specific rubric-learning failure mode. For Unexplored Criteria, the model plays two roles: as student it generates the rollouts as normal, and as criterion-injected self-teacher it re-generates given the rubric as extra context. A per-token forward-KL is computed between the two, then filtered to keep only the smallest set of student tokens whose combined KL mass covers a fraction X of the total. The OPSD-style loss is applied on this filtered subset — the intuition is that these are exactly the tokens where the missing criterion behavior would surface, so distillation there injects it into the student without polluting the rest of the response.

For Suppressed Criteria, the model instead plays counterfactual self-teacher: given the student’s own rollout and the criterion, it produces a criterion-conforming revision. Comparing teacher and student logits token-by-token identifies tokens where the teacher confidently prefers a substitution — the “criterion-relevant tokens” that would otherwise be pushed down by a rollout with non-positive aggregate advantage. Those tokens have their advantage flipped from the group-relative value to a fixed positive hyperparameter (0.1), and the rest of the rollout keeps its GRPO advantages. The two losses are summed with a mixing coefficient. Notably, the entire pipeline is on-policy self-distillation — no external teacher model, no rubric at inference time.

Results

Headline result: CriPO matches or exceeds GRPO baseline final performance at approximately 2× fewer optimization steps, and beats OPSD by a large margin (paper phrase: “eviscerates OPSD in performance”). Reported on medicine and science benchmarks against rubric-based RL baselines. Ablations confirm both the UC branch (localized forward-KL over top-KL tokens) and the SC branch (counterfactual-teacher advantage flipping) are individually necessary, and the analysis section supports SC prevalence with the >57% / 1.8-per-sample figure that motivates the whole design.

Why it’s interesting

CriPO is the first filed paper on this wiki that names and quantifies a failure mode of rubric-based GRPO that is not exploration — Suppressed Criteria at 57% sample prevalence is a big number, and it reframes the current Reasoning RL recipe’s reliance on scalar-aggregated group-relative advantages as leaving substantial signal on the floor. It also sits at an interesting three-way intersection: it uses on-policy self-distillation (On-Policy Distillation) as the injection primitive but with the current model itself as the teacher family, complementing Learning to Replicate Expert Judgment in Financial Tasks‘s “promoted teacher” schedule with a “criterion-conditioned teacher” schedule. And it directly addresses the shaped-reward story on the Reasoning RL concept page — where prior recipes shape by task structure (hierarchical distance in vlm-gym, correctness+format in Video-Thinker, PARL’s parallelism-vs-success), CriPO shapes by which tokens in a rollout carry which criteria, i.e. a token-level rubric decomposition of the advantage signal.