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

> CoRT uses counterfactual replay to redistribute advantage by token-level rubric dependence, improving GRPO by 4.4 points without extra training.

- **Source:** [arXiv](https://arxiv.org/abs/2607.25659)
- **Published:** 2026-07-31
- **Permalink:** https://picx.dev/p/15gUZ4
- **Whiteboard:** https://picx.dev/p/15gUZ4/image

## Summary

## Summary (Overview)
- **CoRT** (Counterfactual Replay for Token-Level Rubric-Guided Policy Optimization) addresses the token-level credit assignment problem in rubric-based RL for LLMs, where GRPO uniformly broadcasts response-level advantages to all tokens.
- The method uses a **counterfactual replay** pass: the same sampled response is rescored under a criteria-free prompt, and the log-likelihood contrast between the two contexts serves as a proxy for rubric dependence.
- This contrast is mapped to bounded, response-normalized **token weights** that redistribute the signed GRPO advantage toward tokens more dependent on the rubric criteria, without training an auxiliary token scorer or changing the response-level reward.
- Experiments across Qwen3-4B, Qwen2.5-7B, and Qwen3-14B on instruction-following benchmarks (IFBench, IFEval, MultiDimIF, AdvancedIF) show that CoRT improves over response-level GRPO in the vast majority of comparisons, with an average gain of **4.4 percentage points**, and remains competitive with learned token-level credit baselines (RTT) while avoiding a separate relevance-learning stage.
- CoRT is also compatible with different policy optimization objectives (DAPO, GSPO) and training diagnostics confirm that response normalization and scheduled activation (SmoothStep ramp) stabilize training dynamics.

## Introduction and Theoretical Foundation
**Background**: Reinforcement learning (RL) is central to LLM post-training, enabling models to follow instructions with complex constraints. Rubric-based RL decomposes response behavior into explicit criteria (e.g., correctness, formatting, safety), providing structured supervision beyond scalar rewards. However, in GRPO-style pipelines, these structured judgments are aggregated into a response-level advantage and applied uniformly to all generated tokens, leaving potential attribution information underused.

**Motivation**: The paper identifies a mismatch: rubric criteria are evaluated at the criterion level, but the policy update remains uniform at the token level. Previous work like Rubrics-to-Tokens (RTT) learns a separate token-level relevance discriminator, but this requires an additional data-generation pipeline and training stage. The central question: *Can token-level credit be derived directly from policy-internal signals, without training a separate relevance estimator?*

**Key insight**: Fixed-response replay under a counterfactual prompt (removing rubric criteria) reveals token-level dependence: tokens whose likelihood depends on the rubric criteria show larger log-probability changes than generic content tokens (see Figure 1 in the paper). This contrast is a natural, policy-internal signal for credit assignment.

**Theoretical basis**: The paper builds on group-relative policy optimization (GRPO), where for each prompt group, the old policy samples \(G\) responses and the reward is normalized into a group-relative advantage:
\[
A_i = \frac{r_i - \mu}{\sigma + \epsilon_r},
\]
where \(\mu\) and \(\sigma\) are the mean and standard deviation of rewards in the group. The GRPO loss is:
\[
\mathcal{L}_{\text{GRPO}} = -\mathbb{E}_{i,t}\left[\min(\rho_{i,t} A_i, \bar{\rho}_{i,t} A_i)\right],
\]
with \(\rho_{i,t} = \frac{\pi_\theta(y_{i,t}|x^+, y_{i,<t})}{\pi_{\theta_{\text{old}}}(y_{i,t}|x^+, y_{i,<t})}\) and \(\bar{\rho}_{i,t} = \text{clip}(\rho_{i,t}, 1-\epsilon_{\text{clip}}, 1+\epsilon_{\text{clip}})\). This objective is agnostic to which tokens express the rubric criteria.

## Methodology
**Core idea**: CoRT redistributes the same response-level advantage \(A_i\) across tokens using token weights \(w_{i,t}\) that satisfy \(\frac{1}{T_i}\sum_{t=1}^{T_i} w_{i,t} = 1\), so the average shaped advantage remains \(A_i\). The weights are derived from counterfactual replay without training an auxiliary model.

**Algorithm steps** (per prompt group):
1. **Standard GRPO**: Sample responses \(y_i\) from prompt \(x^+\) (with rubric criteria), compute rewards \(r_i\), and group-relative advantages \(A_i\).
2. **Counterfactual scoring**: For each response \(y_i\), compute log-probabilities under the full prompt \(\ell_{i,t}^+ = \log \pi_{\bar{\theta}}(y_{i,t} | x^+, y_{i,<t})\) (already available from rollout) and under the criteria-free prompt \(\ell_{i,t}^- = \log \pi_{\bar{\theta}}(y_{i,t} | x^-, y_{i,<t})\) (by replaying fixed tokens).
3. **Contrast computation**: \(\Delta_{i,t} = \ell_{i,t}^+ - \ell_{i,t}^-\).
4. **Replay-margin token credit**: Map contrasts to bounded scores:
   \[
   s_{i,t} = \text{sigmoid}(\tau (\Delta_{i,t} - b)) - \frac{1}{2}, \quad \tilde{w}_{i,t} = 1 + \eta \lambda_k s_{i,t},
   \]
   where \(b\) centers the contrast, \(\tau\) controls sharpness, \(\lambda_k \in [0,1]\) is a scheduled coefficient, and \(\eta\) controls weighting strength. Provisional weights lie in \((1 - \frac{\eta\lambda_k}{2}, 1 + \frac{\eta\lambda_k}{2})\).
5. **Scheduled response normalization**: Use a cubic Hermite SmoothStep schedule:
   \[
   u_k = \text{clip}\left(\frac{k - k_0}{K}, 0, 1\right), \quad \lambda_k = 3u_k^2 - 2u_k^3,
   \]
   where \(k_0\) is warmup offset and \(K\) is ramp length. Then normalize within each response:
   \[
   \bar{w}_i = \frac{1}{T_i}\sum_{t=1}^{T_i} \tilde{w}_{i,t}, \quad w_{i,t} = \frac{\tilde{w}_{i,t}}{\bar{w}_i}.
   \]
   This ensures \(\frac{1}{T_i}\sum_t w_{i,t} = 1\) and the average shaped advantage remains \(A_i\).
6. **Policy objective**: The token-shaped advantage is \(\hat{A}_{i,t} = \text{sg}(w_{i,t}) A_i\), and the loss is:
   \[
   \mathcal{L}_{\text{CoRT}} = -\mathbb{E}_{i,t}\left[\min(\rho_{i,t} \hat{A}_{i,t}, \bar{\rho}_{i,t} \hat{A}_{i,t})\right].
   \]
   Gradients are stopped through the weight \(w_{i,t}\).

**Computational cost**: CoRT adds one additional forward scoring pass per sampled response under the criteria-free prompt, without extra generation, verifier calls, or token-level labels. It is a lightweight modification to GRPO.

## Empirical Validation / Results
**Experimental setup**:
- **Training data**: HIR-16k, providing prompts with and without instruction lists.
- **Reward modes**: Constraint Satisfaction Rate (CSR) – fraction of satisfied criteria; All-or-Nothing (AON) – reward 1 only if all criteria satisfied.
- **Models**: Qwen3-4B-Instruct, Qwen2.5-7B-Instruct, Qwen3-14B-Instruct.
- **Evaluation benchmarks**: IFBench, IFEval, MultiDimIF, AdvancedIF (with o3-mini judge). All methods evaluated with 5 sampled responses per prompt, temperature 0.7.
- **Comparison baselines**: Instruct (base), SFT, DPO, GRPO, RTT (learned token relevance). CoRT uses the same rollout and reward as GRPO.

**Main results** (Table 1, selected):
| Method | IFEval Prompt | IFEval Inst. | IFBench Prompt | IFBench Inst. | MultiDimIF Acc. | AdvancedIF Overall | AdvancedIF Rubric |
|--------|---------------|--------------|----------------|---------------|-----------------|-------------------|-------------------|
| **Qwen3-4B CSR** | | | | | | | |
| GRPO | 84.29 | 89.21 | 32.31 | 33.43 | 74.38 | 46.02 | 79.81 |
| +CoRT | **86.06** | **90.48** | **34.49** | **36.24** | **80.48** | **49.18** | **81.14** |
| Delta | +1.77 | +1.27 | +2.18 | +2.81 | +6.10 | +3.16 | +1.33 |
| **Qwen2.5-7B CSR** | | | | | | | |
| GRPO | 78.56 | 84.41 | 28.23 | 29.85 | 66.67 | 30.91 | 72.00 |
| +CoRT | **81.92** | **87.51** | **34.63** | **36.30** | **78.52** | **34.18** | 73.28 |
| Delta | +3.36 | +3.10 | +6.40 | +6.45 | +11.85 | +3.27 | +1.28 |

- CoRT improves over GRPO in all metrics for Qwen3-4B (CSR) and Qwen2.5-7B (CSR), and in most metrics for AON settings.
- CoRT is competitive with RTT (often stronger), while avoiding a separate relevance-learning stage.
- On Qwen3-14B (Table 2), CoRT improves all metrics under CSR and most under AON.
- Compatibility with DAPO and GSPO (Table 3): CoRT improves all metrics over DAPO and four of five over GSPO.

**Training diagnostics** (Figure 3, Figure 4):
- Ablation study on Qwen2.5-7B CSR: full CoRT (with response normalization and SmoothStep ramp) maintains mean token weight close to 1, low length clipping, stable gradient norm and entropy.
- Removing response normalization causes mean token weight drift and late training spikes.
- Removing the ramp leads to abrupt weighting and instability.
- The full recipe achieves the strongest validation average and training reward around the 500-step horizon.

## Theoretical and Practical Implications
**Theoretical implications**:
- The paper demonstrates that **policy-internal counterfactual likelihood contrasts** provide a sufficient training signal for within-response credit allocation, without needing an external relevance model. This connects to causal inference (Pearl, 2009) and attribution methods, but uses contrasts for training rather than explanation.
- The method preserves the original reward structure and group-relative normalization, only modifying the distribution of the signed advantage across tokens. This shows that structural criteria can guide not only scoring but also credit assignment.

**Practical implications**:
- CoRT is a **lightweight, plug-and-play addition** to GRPO-style pipelines, requiring only a single additional forward pass per response. It avoids the complexity of training a separate token-level relevance model (as in RTT).
- The method is **compatible with multiple policy optimization algorithms** (GRPO, DAPO, GSPO), making it broadly applicable.
- The scheduled activation and response normalization are crucial for stability, providing practical guidelines for implementing counterfactual-based credit assignment.

## Conclusion
**Main takeaways**:
- CoRT addresses the token-level credit assignment mismatch in rubric-based RL by using counterfactual replay to derive token weights from the policy's own likelihood contrasts.
- It improves over response-level GRPO consistently across models, reward granularities, and benchmarks, with an average gain of 4.4 percentage points.
- It remains competitive with learned token-level methods while being simpler and more lightweight.
- The method is stable in practice due to integration controls (response normalization and SmoothStep ramp).

**Future directions**:
- Finer-grained interventions (e.g., criterion-specific allocation).
- Longer-horizon credit assignment across turns, actions, and tool calls.
- Extending to settings where criteria have little effect on token likelihoods or rewards are noisy.

---

_Markdown view of https://picx.dev/p/15gUZ4, served by PicX — AI-generated visual whiteboard summaries of research papers._