CoRT: Counterfactual Replay for Token-Level Rubric-Guided Policy Optimization

2026-07-28

AI Summary

Existing GRPO methods based on rubric scores rely solely on response-level rewards, making fine-grained credit assignment across individual tokens challenging—particularly when scoring involves multiple segments or semantic choices. To address this, this work proposes CoRT, the first approach to incorporate a counterfactual replay mechanism for token-level credit assignment. CoRT constructs a context-aware proxy signal by comparing token-level log-likelihood differences between responses generated with and without the scoring prompt, and uses this signal to redistribute advantages in GRPO. Notably, CoRT requires no additional reward models or auxiliary training stages, preserving GRPO’s simplicity while enabling fine-grained policy optimization. Experiments show that CoRT achieves an average improvement of 4.4 percentage points across various instruction-tuned models and reward settings, significantly outperforming response-level GRPO and matching the performance of token-level baselines that require extra learning components.

Abstract

Rubric-based reinforcement learning enriches language model training by evaluating model outputs against explicit criteria. Yet in GRPO-style pipelines, these structured judgments are reduced to a scalar response-level reward and converted into a response-level advantage, which is broadcast uniformly to all generated tokens. This leaves no explicit mechanism for allocating credit within a response, even when different criteria are grounded in different spans, formatting decisions, or semantic choices. We propose CoRT, a token-level credit weighting method for rubric-conditioned GRPO. Instead of training an auxiliary token scoring model, CoRT uses counterfactual replay to rescore the same sampled response under the original rubric-conditioned prompt and a matched criteria-free prompt. The resulting tokenwise log-likelihood contrasts serve as a proxy for dependence on the rubric context. CoRT maps these contrasts to bounded, response-normalized weights and uses them to redistribute the signed GRPO advantage across tokens, without introducing an auxiliary scorer or changing the response-level reward. Experiments across instruction-tuned models and reward granularities show that CoRT improves over matched response-level GRPO in the vast majority of comparisons, with an average gain of 4.4 percentage points. The method remains competitive with learned token-level credit baselines while avoiding a separate relevance-learning stage. These results suggest that policy-internal counterfactual likelihood contrasts provide an effective training signal for within-response credit allocation while retaining the simplicity and stability of GRPO.

Problem

• credit assignment
• token-level reward
• rubric-based reinforcement learning
• GRPO
• within-response evaluation

Innovation

• Counterfactual Replay
• Token-level Credit Allocation
• Rubric-guided RL
• GRPO
• Likelihood Contrast