Agentic Reinforcement Learning with Observation-Calibrated Self-Distillation

AI summary

Overview

Research area: Reinforcement learning for large language model (LLM) agents, specifically the integration of privileged on-policy self-distillation signals with trajectory-level policy optimization.

Technical level: Advanced. The paper assumes familiarity with policy-gradient methods, group-relative advantage estimation, importance-ratio clipping, and knowledge-distillation objectives.

Scope: One sentence: the paper identifies an attribution confound in agentic privileged replay scoring — where the "support" a replay view assigns to a token mixes the effect of the future observation with the effect of the replay scaffold itself — and proposes Observation-Calibrated Self-Distillation (OCSD), which subtracts a structurally matched Observation-Ablated view to isolate an observation residual and uses that residual to modulate token-level GRPO updates.

What This Paper Is About

LLM agents are typically trained with reinforcement learning that gives one reward per trajectory, so every token in a multi-step trajectory receives the same update signal and no token gets individualized guidance. On-Policy Self-Distillation (OPSD) tries to fix this by re-scoring the agent's own generated tokens under a privileged "replay" view that includes future information, then treating the score change as dense token-level supervision. The problem the authors identify is that this score change also includes shifts caused by the replay scaffold itself, so it is unclear how much of the signal actually comes from the future observation; OCSD's goal is to separate the two.

Key Contributions

1. Identifying a confounding effect in agentic privileged replay scoring. The authors show that Full-view support can reflect both the information contained in future environment observations and score changes induced by the shared replay scaffold (replay format plus future-action schemas), making attribution to the future observation ambiguous.

2. Proposing OCSD. The method constructs two structurally matched replay views — Full and Observation-Ablated — that differ only in whether the actual future observation is present, contrasts their token scores to derive an observation residual, and applies that residual to modulate token-level GRPO updates at high-uncertainty steps while preserving the trajectory-level update direction.

3. Diagnostic evidence that the calibrated residual is better aligned with environment feedback. Fine-grained analysis of Qwen3-1.7B over 150 training steps on ALFWorld shows the residual achieves higher AUROC than Full support for distinguishing valid from invalid actions, and that this advantage survives normalization, stratified matching, and different training stages.

4. Comprehensive evaluation across three benchmarks and three model scales. Experiments on ALFWorld, WebShop, and Search-QA with Qwen3-1.7B, Qwen3-4B, and Qwen3-8B show OCSD achieving the best overall performance in every task-and-scale combination reported.

Main Findings

• OCSD achieves the strongest overall results across all benchmarks. On ALFWorld, OCSD reaches overall success rates of 55.5, 82.8, and 87.2 with Qwen3-1.7B, 4B, and 8B respectively, outperforming GRPO by 8.9, 12.2, and 14.0 points. It also reaches WebShop success rates of 54.4, 73.7, and 78.1, and Search-QA average Exact Match scores of 43.1, 47.5, and 49.1.

• Direct use of privileged replay support yields inconsistent gains. Standalone OPSD performs substantially worse than GRPO across models and environments, and is particularly poor and unstable on Search-QA (for example, 0.6 average EM with Qwen3-4B versus GRPO's 45.5). Combining OPSD directly with GRPO still fails to guarantee stable gains: GRPO+OPSD underperforms the original GRPO on ALFWorld with Qwen3-1.7B and Qwen3-4B.

• The observation residual aligns better with local environment feedback than Full support. Over 150 training steps of Qwen3-1.7B on ALFWorld, using the top-20% high-NLL interaction steps within each trajectory, step-averaged AUROC for distinguishing valid from invalid actions is 0.464 for the Observation-Ablated signal, 0.654 for the Full signal, and 0.707 for the residual. The residual's improvement over Full support remains stable under within-checkpoint standardization, stratified matching for training stage and action operator, and across training stages.

• Full and Observation-Ablated views produce overlapping but distinguishable token signals. Word-level analysis using upper-5% tail hit rates (and the lower-5% tail for the negative residual) shows the upper tails of Full and Observation-Ablated support contain substantially overlapping words, while the observation residual highlights a distinct set. A word-level relative-residual-contribution measure D shows reasoning and environment entity/state words have higher D than action and format words.

• Ablations show every component of OCSD matters. On Qwen3-4B, the full method reaches 82.8 ALFWorld overall, 86.9 WebShop Score, and 73.7 WebShop Success, versus GRPO's 70.6, 81.9, and 69.5. Removing the ablated teacher (using only Full-teacher support) gives 74.8 / 85.0 / 72.9 — better than GRPO, showing future observations carry useful supervision. Random step selection gives 78.9 / 84.2 / 70.6. Removing step selection entirely gives 66.4 / 82.4 / 69.8, falling below GRPO on ALFWorld. Removing sign alignment gives 71.4 / 81.5 / 71.6, near the GRPO level.

• The residual depends on the correspondence between future observation and current action, not just on extra observation text. When the realized future observation in the Full view is replaced by a randomly swapped observation from another interaction step, AUROC for distinguishing valid from invalid steps drops from 0.803, 0.668, and 0.664 (realized) to 0.542, 0.484, and 0.552 (swapped), close to random ranking.

• OCSD performs consistently well under distribution shift. Across all three model scales it achieves the best out-of-domain Search-QA results among compared methods. In Search-QA, NQ and HotpotQA serve as in-domain training distributions while the remaining five datasets are used for out-of-domain evaluation.

• Training overhead is small. Dual-view replay scoring requires 7.52 seconds per iteration and residual modulation 0.12 seconds, for a combined overhead of 7.64 seconds — only 1.40% relative to the shared GRPO pipeline, measured on ALFWorld with Qwen3-4B using 8 A100 GPUs.

• OCSD improves more steadily than competing methods across training. In ALFWorld validation success rate curves, methods perform similarly and fluctuate early on, while OCSD improves more steadily and gradually separates from GRPO, RLSD, and SDAR as training progresses.

Methodology in Plain English

The method builds on GRPO, which samples a group of G trajectories per task and computes a group-relative advantage from their terminal rewards, then applies that same advantage to every token. OCSD keeps that trajectory-level direction but changes the per-token update magnitude, in three steps.

First, it decides where to intervene. It measures the average negative log-likelihood of the tokens generated at each interaction step under the old policy, treating high values as high uncertainty, ranks steps within each trajectory, and keeps the top ρ fraction (ρ = 0.2 in the main experiments, with at least one step per trajectory).

Second, it builds two replay views for each selected step. The Full view contains the actual future observation plus "naturalized" schemas of future actions, where the schema keeps the coarse action intent and syntax but replaces specific entities, locations, queries, and answers with natural-language descriptions of their roles. The Observation-Ablated view is identical except the future observation content is replaced by the fixed phrase "Observation: not provided," with field order, replay format, and future-action scaffold unchanged. Both views re-score the same tokens the student already generated, using the same model parameters but different input contexts.

Third, it differences them. The observation residual is the Full-minus-Ablated log-probability difference for each token; because both terms share the same student prediction, the student log-probability cancels. This residual is squashed through a tanh to a bounded signal in [-1, 1]. The sign says whether the Full teacher likes the token more or less than the ablated teacher; the magnitude says how much the two views disagree. That signal multiplies the trajectory advantage through a sign-preserving factor at selected steps only, so the modulation is always non-negative and the update direction set by the trajectory reward is never flipped. The resulting objective is the standard GRPO clipped surrogate with the calibrated token-level advantage and a KL penalty to the reference policy. No extra imitation objective is added, and no actions are generated from the replay views — they are used only to score already-generated tokens.

Why This Matters

Impact on research: The paper reframes a widely used technique — scoring student tokens under a privileged hindsight view — as a measurement problem with a confounding variable. Many recent agentic distillation methods build privileged context from future observations, and this work shows that a naive Full-versus-Student difference conflates the target information with artifacts of the replay scaffold. The structurally matched control-view design is a reusable methodological template for anyone constructing privileged supervision.

Real-world applications:

• Web and e-commerce agents. WebShop-style product search, filtering, and purchase under user-specified requirements is a direct application, where an agent must learn from each page it lands on.
• Embodied and household task agents. ALFWorld-style text-based manipulation tasks — picking, looking, cleaning, heating, cooling, and multi-object Pick2 variants — benefit from step-local supervision when a single wrong object reference derails a long task.
• Retrieval-augmented and search-based question answering. Search-QA-style external information retrieval and integration, including multi-hop questions where the agent must decide which retrieved evidence to trust.
• Any multi-turn tool-using agent where environment feedback arrives after each action and sparse terminal rewards leave individual decisions under-specified.

Industry relevance: The reported overhead is 1.40% relative to the shared GRPO pipeline, and OCSD introduces no additional imitation objective or replay-conditioned action generation. That makes it a practical drop-in modification for teams already running GRPO-style training on agentic workloads, rather than requiring a separate training paradigm. The gains are also larger at smaller model scales in relative terms and consistent across three scales, which matters for cost-sensitive deployment.

Future Directions

• Extending the matched-view design to other privileged signals. The paper's own related-work discussion notes that privileged signals can shift for reasons other than scaffolding — hints, references, teacher specialization in reasoning tasks. Whether the same contrast-based deconfounding applies to those sources is open.
• Understanding and removing the remaining scaffold sensitivity. The authors report that even the Observation-Ablated view alone achieves only 0.464 AUROC for validity discrimination, and that ablations such as random step selection still beat GRPO. Where the residual's remaining error comes from, and how much of it could be removed with better view construction, is not resolved.
• Tuning the intervention more finely. The paper fixes ρ = 0.2 and β = 0.5 for the main experiments and states that hyperparameter sensitivity analyses are provided in supplementary material (the appendix content in the available text is truncated mid-sentence). Adaptive step selection or per-step calibration strength is a natural next question.
• Scaling and generalization beyond the reported settings. All results are on Qwen3-1.7B/4B/8B with three benchmarks. Whether the gains persist at much larger scales, on model families with different pretraining, or on longer-horizon tasks with more interaction steps is not reported.

Target Audience

Researchers and engineers working on reinforcement learning for LLM agents, particularly those implementing or evaluating on-policy self-distillation, hindsight re-scoring, or token-level credit assignment on top of GRPO. The paper is also relevant to practitioners deploying multi-turn tool-using agents who need an efficient way to add dense supervision without a separate distillation objective, and to methodologists interested in experimental controls for privileged-information training signals.

Authors’ abstract

Large language model agents are commonly trained through reinforcement learning with sparse trajectory-level rewards, which offer limited guidance on how strongly individual tokens should be updated. On-Policy Self-Distillation (OPSD) addresses this by re-scoring generated tokens under a privileged replay view to obtain dense, token-level supervision. However, we identify a confounding issue: the resulting support may reflect both the privileged information contained in the replay view and score shifts induced by the replay scaffold, making it difficult to attribute the support specifically to that information. This issue is especially pronounced when future environment observations serve as privileged information, since replaying them requires reconstructing an extended scaffold that itself perturbs token scores. To resolve this confounding, we propose Observation-Calibrated Self-Distillation (OCSD), which contrasts two structurally matched replay views, Full and Observation-Ablated, differing only in whether the actual future observation is present, to derive an observation residual that discounts score changes shared by the replay scaffold. OCSD then applies this residual to modulate token-level GRPO updates at high-uncertainty steps, while preserving the trajectory-level update direction. Experiments on ALFWorld, WebShop, and Search-QA across three Qwen3 model scales show that OCSD consistently outperforms strong baselines. Diagnostic analyses further confirm that the calibrated residual aligns better with local environment feedback. Our code is publicly available at https://github.com/yiy1x/OCSD.