Language models improve scoring with rubric-guided training and reinforcement learning
OPD Before RL: Warm-Starting Rubric-Based RL with On-Policy Distillation
Machine LearningArtificial IntelligenceComputation and Language
Summary
Some tasks for language models are hard to check because answers can be open-ended and scored by rubrics rather than exact right or wrong outcomes. The authors introduce a two-step way to train these models: first, the model learns to mimic another that has help from the rubric for detailed token-by-token guidance. Then, it uses reinforcement learning to better optimize the rubric scores given after the full answer. This approach helps the model make better responses without just gaming the rubric rules.
What this means in practice
- •For health data teams: Improve automated evaluation of open-ended medical text by training language models with rubric-guided distillation before reinforcement learning to better score complex responses.
- •For science qa platform developers: Enhance accuracy of science question answering systems by pretraining models with rubric-informed token-level guidance followed by reinforcement learning to optimize rubric-based rewards.
Authors
Xinpeng Wang, Wei Shi, Yu-Chia Chen, Maria Zontak, Yun He, Richard Yuanzhe Pang
Abstract
Many useful language-model tasks cannot be evaluated by exact outcome verification. Rubric-based reinforcement learning (RL) addresses this issue by scoring open-ended responses against explicit criteria. However, because the reward is assigned after the complete response, the training signal does not directly identify which individual decisions contributed to the final score. We propose a two-stage training framework that uses rubrics first as privileged teacher context for dense token-level supervision, then as rewards for further RL. In the first stage, rubric-privileged on-policy distillation (RP-OPD), a student without access to the rubric matches a rubric-aware teacher's next-token distributions at student-generated prefixes. In the second stage, RL directly optimizes the rubric reward and improves beyond the observed distillation plateau. We evaluate the framework on health and science tasks using open-weight models. Across HealthBench, ResearchQA, and RubricHub Science, we compare post-training methods and vary the amount of SFT or RP-OPD training before RL, finding that our two-stage framework achieves the highest scores among the methods evaluated. RP-OPD + RL shows limited signs of reward hacking on RubricHub Science, whereas the SFT + RL baseline increasingly receives high rewards for claims of rubric compliance without providing the required content. These findings support using rubrics to guide on-policy distillation before applying rubric-based RL.