CRAFT: Counterfactual Credit Assignment from Free Sibling Rollouts for Self-Distilled Agentic Reinforcement Learning

Abstract: Self-distilled agentic reinforcement learning augments trajectory-level reward with a token-level distillation loss, using as its teacher the same policy conditioned on privileged context. The prevailing recipe gates this loss by a single scalar, the teacher-student log-probability gap. This signal is doubly limited: it is retrospective, scoring only the realised rollout and never the counterfactual ones, and it is sign-blind, never signalling when a teacher-preferred action would have harmed the trajectory. We introduce CRAFT, a three-pillar credit-assignment scheme that addresses both limitations. Pillar 1, Counterfactual Token Importance, reuses the G-1 sibling rollouts that GRPO already samples and importance-weights them by the log-probability gap to form a self-normalised estimate of the group-level counterfactual change in advantage from up-weighting teacher-preferred actions at each step; this yields a signed per-token credit at near-zero extra compute. Pillar 2 is an asymmetric controller that raises the distillation weight as it lowers the reference-KL weight along an exponential moving average of gate activity, and conversely. Pillar 3 polarises the KL penalty token by token, switching between a mode-seeking and a mode-covering update according to the sign of the credit. Each pillar has an independent switch that, when disabled, renders the loss and gradient byte-identical to the baseline in IEEE-754 arithmetic, so any measured gain is attributable to algorithmic change rather than implementation drift. We prove the estimator's consistency and a variance bound, give structural and bit-exact reproducibility guarantees, and evaluate CRAFT across three agentic environments, four model scales, and five end-to-end methods, plus two tabulated prior-work baselines. Among these is Adaptive-CRINGE, a comparator sharing Pillar 2 with CRAFT, isolating the counterfactual contribution.

• The paper introduces a counterfactual credit assignment framework (CRAFT) that leverages free sibling rollouts to produce signed per-token feedback in self-distilled agentic RL.
• It proposes a three-pillar method—including CTI, phase-aware adaptive control, and direction-polarized KL penalty—to improve policy learning and overcome sign-blindness.
• Empirical results across diverse environments show that CRAFT outperforms prior baselines, enhances generalization, and yields significant gains especially on smaller models.

Counterfactual Credit Assignment in Self-Distilled Agentic RL: The CRAFT Method

Motivation and Problem Statement

Self-distilled agentic reinforcement learning (SDAR) is a dominant framework for LLM-based agents operating in complex, multi-turn environments. In SDAR, the agent is trained not only on a trajectory-level RL signal (e.g., group-normalised advantage) but also on a token-level distillation loss, with the current policy under privileged context acting as the teacher. However, standard SDAR uses a token-level distillation gate based solely on the teacher-student log-prob gap ($A_t$), resulting in two critical limitations:

1. Retrospective Credit Assignment: The gate only scores realized rollouts, neglecting what would have happened had different actions been taken at each step—despite the parallel rollouts ("siblings") produced by Grouped Relative Policy Optimization (GRPO) being readily available.
2. Sign-Blindness: The gate's monotonic function never signals when following the teacher would harm return, thus failing to penalize deleterious teacher-favored actions.

The paper "CRAFT: Counterfactual Credit Assignment from Free Sibling Rollouts for Self-Distilled Agentic Reinforcement Learning" (2606.29476) systematically addresses these issues.

The CRAFT Framework

CRAFT introduces a three-pillar credit assignment framework that systematically overcomes the retrospective and sign-blind nature of SDAR's per-token feedback.

Pillar 1: Counterfactual Token Importance (CTI)

CTI exploits the G-1 sibling rollouts already sampled by GRPO to estimate, at essentially no extra computational cost, the counterfactual change in expected trajectory advantage had a teacher-preferred action been taken at step $t$. This is computed as a self-normalized importance-weighted average over sibling advantages, with importance weights given by the exponentiated teacher-student log-prob gap. The resulting estimate provides a signed per-token credit; its consistency and $O(1/G)$ variance bound are theoretically established.

Pillar 2: Phase-Aware Adaptive Controller

To handle the disparate timescales and objectives of distillation and KL regularization, Pillar 2 uses an exponential moving average of the active gating signal (“Pgate”) to co-regulate the magnitudes of the distillation (CTI) and reference-KL losses. The coefficients adaptively move in opposite directions: as more disagreement tokens are found (higher Pgate), the distillation loss weight increases while the KL anchor decreases, relaxing the policy to absorb new signal when needed and restraining it as convergence nears.

Pillar 3: Direction-Polarised KL Penalty

Whereas prior work applies the same KL penalty regardless of the sign of the per-token signal, Pillar 3 routes the penalty branch based on the sign of the CTI credit:

• Mode-seeking KL (reverse-KL) when following the teacher is beneficial.
• Mode-covering KL (forward-KL-inspired) when diverging from the teacher is advantageous.
• Neutral KL for near-zero credit.

Each pillar's contribution can be ablated via a master switch, guaranteeing that gains are attributable to algorithmic change, not implementation artifacts.

Theoretical Properties

The CTI estimator is shown to be a consistent, self-normalized importance-sampling estimator of the group-level counterfactual advantage, with a variance bound scaling as $O(1/G)$. Degeneracy proofs guarantee that when disabled, CRAFT recovers the byte-exact behavior (including gradients) of the underlying SDAR baseline, ensuring result attribution is always unambiguous.

CRAFT's policy-gradient surrogate and KL augmentations commute cleanly with PPO clipping, and all auxiliary terms are detached in accordance with policy gradient principles.

Empirical Results

Experiments span three challenging agentic RL environments (ALFWorld, Search-QA, and WebShop), four model sizes (from 1.7B to 8B parameters, Qwen2.5 and Qwen3 families), and five major baseline and ablation methods (including GRPO, RLSD, SDAR, Adaptive-CRINGE, and CRAFT-Full).

Key findings are:

• CRAFT-Full outperforms all baselines and ablations in every environment and at every model size, with consistent gains over both the strongest prior RLSD/SDAR baselines and the harder Adaptive-CRINGE (which isolates the effect of counterfactual vs. retrospective credit).
• On smaller models, CRAFT yields marked improvements (up to +5.5% absolute on WebShop, 1.7B); the absolute gain shrinks for larger models where fewer tokens have disagreement.
• Pillar 1 (CTI) alone accounts for the majority of the gain; Pillar 2 (adaptive control) and Pillar 3 (KL routing) contribute smaller, but complementary, improvements.
• CRAFT demonstrates improved generalization: the OOD performance gap (in-distribution to out-of-distribution slices) is systematically lower than SDAR and Adaptive-CRINGE.
• Robustness tests (randomized teacher context) show CRAFT degrades gracefully rather than catastrophically.

Comparative Analysis

CRAFT's innovations distinguish it methodologically from all prior token-level credit-assignment and self-distillation frameworks:

• Unlike RLSD and SDAR, CRAFT utilizes counterfactual rather than retrospective per-token signals, introducing forward-looking, reward- and sign-aware feedback.
• The importance reweighting over GRPO siblings reclaims otherwise wasted compute, ensuring the estimator is “free” in practice.
• The CTI framework preserves compatibility with sequence-level PPO/GRPO surrogates and leverages standard batch processing efficiently.
• Adaptive KL approaches (including InstructGPT, RLHF pipelines) only adapt a global coefficient, whereas CRAFT adapts both distillation and KL in a coupled, token-wise, phase-aware manner.
• Mode-seeking/mode-covering asymmetry at the per-token level enhances local policy optimization and avoids over-anchoring typical of global KL schemes.

Limitations and Implications

There are noteworthy constraints: CRAFT requires group rollouts (G ≥ 2), introducing a reliance on parallel batch environments, and nine new hyperparameters, though reasonable defaults are provided. Benchmarks are restricted to text environments, leaving cross-modal and cross-domain generalization as future work.

Practically, CRAFT provides a straightforward drop-in replacement for SDAR/GRPO token-level RL pipelines with bit-exact reproducibility, enabling precise algorithmic attribution. Theoretically, it suggests that leveraging parallel rollouts for counterfactual reasoning is critically underexplored in sequence modeling RL, and that phase-aware, sign-sensitive losses are essential for fine-grained credit assignment.

Future work could involve extension to synthetic-sibling settings, explicit modeling/estimation of the within-group exchangeability gap, and generalization to multi-modal, real-world agentic environments.

Conclusion

CRAFT presents a theoretically grounded and empirically validated approach to per-token counterfactual credit assignment in self-distilled agentic reinforcement learning (2606.29476). Through judicious leverage of GRPO's sibling rollouts, adaptive scheduling, and direction-polarized KL, CRAFT achieves consistent improvements over prevailing baselines while maintaining strict reproducibility. This establishes a new protocol for integrating counterfactual reasoning with self-distillation, with implications for the ongoing development of robust, sample-efficient LLM agents.