PaperlayerView exact revision on arXiv

Source-linked AI summary

Agentic Reinforcement Learning with Observation-Calibrated Self-Distillation

Yi Yang, Cong Qin, Xiaodan Liu, Chishui Chen, Qing Dong, Yan Zhang, Cao Liu, Zhao Yang, Lu Pan, Jiaye Lin, Yi Feng

arXiv:2608.04788v1cs.LGcs.AIcs.CL

TL;DR

Sparse trajectory rewards provide limited token-level guidance, and privileged replay can mix future-observation information with score shifts from the replay scaffold. OCSD contrasts matched Full and Observation-Ablated views to derive a calibrated residual for GRPO updates, achieving consistent gains across benchmarks and model scales while aligning better with local environment feedback.

  • Problem

    Privileged replay support can reflect both future-observation information and score changes induced by the shared replay scaffold, making attribution difficult.

  • Method

    OCSD contrasts structurally matched Full and Observation-Ablated replay views to derive an observation residual, then applies it at high-uncertainty steps to modulate token-level GRPO updates while preserving trajectory-level direction.

  • Results

    OCSD consistently outperforms strong baselines across ALFWorld, WebShop, and Search-QA using Qwen3 models at three scales.

  • Takeaways & Limitations

    The calibrated residual provides token-level supervision that better aligns with local environment feedback than Full-view support.

Abstract

from arXiv · show

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.

1 Introduction

Agentic RL methods often give every token the same trajectory-level guidance, while privileged replay can conflate future-observation information with score shifts from the replay scaffold. OCSD addresses this by contrasting matched replay views to obtain calibrated token-level supervision.

  • Motivation: GRPO-style methods apply shared trajectory-level rewards to all tokens, providing limited guidance about how strongly individual tokens should be updated.OPSD complements sparse rewards by re-evaluating on-policy tokens under privileged replay views.
  • Motivation: Future-observation replay can conflate information from the actual observation with score changes caused by the added replay format and future-action scaffold.Observation-Ablated views retain similar support for some tokens even without the true observation, exposing the attribution problem.
  • OCSD: OCSD contrasts structurally matched Full and Observation-Ablated views, deriving an observation residual that discounts score changes shared by the replay scaffold.The residual is applied at high-uncertainty steps to modulate token-level GRPO updates while leaving the trajectory-level update direction unchanged.
  • Results: Across ALFWorld, WebShop, and Search-QA with three Qwen3 model scales, OCSD achieves the best overall performance and consistently outperforms strong baselines.Diagnostic analyses also report closer alignment between the calibrated residual and local environment feedback than Full-view support.

2 Revisiting Privileged Replay Scoring

The section frames GRPO and privileged replay scoring, then diagnoses how replay scaffolds confound Full-view support. Matched-view comparisons show that the residual isolates observation-sensitive signal more effectively and aligns better with local feedback.

  • 2.1 Preliminaries: GRPO assigns the same group-relative trajectory advantage to every token, without distinguishing tokens’ individual contributions to the terminal reward.OPSD instead re-scores on-policy outputs under replay evidence to provide token-level support.
  • 2.1 Preliminaries: Full replay includes actual future observations and naturalized future-action schemas, whereas Observation-Ablated replay preserves the structure while replacing the future observation with a fixed phrase.Both views use the same model parameters with different input contexts.
  • 2.2 Fine-Grained Diagnosis of Replay Support: Full and Observation-Ablated views assign nearly identical support to generic tokens but differ more for tokens tied to the actual future observation, such as “take,” “mug,” and “shelf.”This comparison motivates treating the view difference as observation-sensitive rather than using Full support directly.
  • 2.2 Fine-Grained Diagnosis of Replay Support: Reasoning and environment entity/state words have higher relative residual contribution than action and format words.The analysis groups word occurrences into five mutually exclusive functional categories.
  • 2.2 Fine-Grained Diagnosis of Replay Support: The residual achieves AUROC 0.707 versus 0.654 for Full and 0.464 for Observation-Ablated support when distinguishing valid from invalid actions.Its improvement over Full remains stable under within-checkpoint standardization, stratified matching, and across training stages.
  • 2.2 Fine-Grained Diagnosis of Replay Support: Full support contains preferences that persist without the actual future observation, whereas the view difference aligns more closely with local environment feedback.This supports using the residual as a more targeted signal for observation-related supervision.

3 Methodology

OCSD derives a bounded observation-calibrated token signal by contrasting Full and Observation-Ablated replay views, then uses it to adjust GRPO updates at uncertain interaction steps while preserving trajectory-level direction.

  • 3.1 Observation-Calibrated Token Signal: OCSD contrasts Full and Observation-Ablated replay views to derive an observation residual that captures their incremental token-support difference.Both support terms use the same student prediction, so the student log-probability cancels in the difference.
  • 3.1 Observation-Calibrated Token Signal: The bounded signal q_i,k,t records whether the Full Teacher supports a token more or less than the Observation-Ablated Teacher, with magnitude measuring disagreement.
  • 3.2 NLL-Guided Step Selection: OCSD prioritizes interaction steps with high average token negative log-likelihood, selecting the top ρ fraction within each trajectory.The ceiling operation ensures that at least one step is selected from every trajectory.
  • 3.3 Sign-Preserving Integration with GRPO: OCSD adjusts update strength within selected steps while preserving the trajectory-level GRPO direction because its nonnegative modulation factor preserves the advantage sign.Tokens in unselected steps retain the original GRPO advantage, and β controls calibration strength.
  • 3.3 Sign-Preserving Integration with GRPO: The calibrated token-level advantage replaces the trajectory-level advantage in the GRPO surrogate for generated tokens.The objective uses the token importance ratio together with reference-policy KL regularization and clipping.

4 Experiments

Experiments across ALFWorld, WebShop, and SearchQA evaluate OCSD against reward-only, self-distillation, and RL–OPSD integration baselines at three Qwen3 scales. OCSD achieves the strongest overall performance, with gains supported by task-level, distribution-shift, ablation, training-dynamics, diagnostic, and efficiency analyses.

  • 4.2 Main Results: OCSD achieves the strongest overall results across ALFWorld, WebShop, and SearchQA at all three Qwen3 scales.On ALFWorld, its success rates are 55.5, 82.8, and 87.2 for Qwen3-1.7B, 4B, and 8B, exceeding GRPO by 8.9, 12.2, and 14.0 points; WebShop rates are 54.4, 73.7, and 78.1, while SearchQA average EM scores are 43.1, 47.5, and 49.1.
  • 4.2 Main Results: Direct OPSD support is inconsistent: standalone OPSD often underperforms GRPO, and GRPO+OPSD loses to GRPO on ALFWorld at 1.7B and 4B.This motivates calibrating replay support rather than applying it directly.
  • 4.2 Main Results: OCSD performs best or second-best on most ALFWorld task categories, including both basic Pick and multi-step Pick2 tasks.The improvement is therefore broad across task difficulty categories rather than confined to one task.
  • 4.2 Main Results: OCSD achieves the best out-of-distribution results at all three model scales, extending gains beyond the training distributions.The reported scope includes unseen retrieval and question-answering datasets.
  • 4.3 Ablation Analysis: All ablated variants underperform complete OCSD, while removing step selection or sign alignment reduces performance toward GRPO.The ablations support using matched-view calibration selectively at high-uncertainty steps and jointly with trajectory-level advantage direction.

5 Related Work

Related work develops privileged replay scoring, agentic self-distillation, and reliability calibration for reinforcement learning. OCSD differs by focusing on disentangling future-observation information from score changes caused by the shared replay scaffold.

  • On-Policy Self-Distillation for Agentic Reinforcement Learning: OPSD provides token-level supervision from privileged replay, while RLSD modulates reinforcement-learning updates and SDAR uses a gated auxiliary objective.These methods extend self-distillation to multi-turn agentic reinforcement learning but differ in how the signal enters optimization.
  • On-Policy Self-Distillation for Agentic Reinforcement Learning: Other agentic methods construct privileged supervision through hindsight re-scoring, verified answers, subsequent observations, grounded feedback, skills, or self-play question-construction paths.The literature therefore varies both the source of privileged context and the mechanism for using it.
  • Calibration and Reliability of Privileged Signals: OCSD focuses specifically on separating future observations from score changes shared by the replay scaffold.This distinguishes its objective from prior methods that primarily construct or apply privileged supervision without this attribution calibration.
  • Calibration and Reliability of Privileged Signals: Calibration studies address unrelated changes in privileged signals, including hint-induced style shifts, reference-specific shortcuts, unreachable hindsight guidance, positional reliability, and dense-distillation artifacts.These works establish reliability concerns for privileged supervision while examining different confounds from OCSD's replay-scaffold attribution issue.

6 Conclusion

The paper identifies an attribution confound in privileged replay scoring and proposes OCSD to calibrate token-level GRPO updates with a residual from matched replay views. Across three benchmarks and three Qwen3 scales, OCSD shows consistent gains, and diagnostics indicate better alignment with local environment feedback.

  • 6 Conclusion: OCSD addresses replay-scoring confounding by contrasting Full and Observation-Ablated views and using their residual to calibrate token-level GRPO updates.The approach preserves the trajectory-level update direction while separating future-observation effects from shared replay-scaffold changes.
  • 6 Conclusion: Experiments across three agentic benchmarks and three Qwen3 scales demonstrate consistent performance gains, while diagnostics show better alignment with local environment feedback.The conclusion covers both the empirical performance result and the diagnostic evidence supporting the calibrated residual.

A.1 End-to-End Training Procedure

OCSD inserts dual-view replay scoring between rollout and optimization, selecting high-NLL interaction steps and rescoring already generated tokens under matched Full and Observation-Ablated contexts. Their difference becomes a bounded calibration signal that changes token update magnitude without changing trajectory-level direction.

  • A.1 End-to-End Training Procedure: OCSD samples on-policy trajectories, computes group-relative advantages, and selects the top-ρ interaction steps by average token NLL.The selected steps receive the additional replay-based calibration procedure within the standard GRPO pipeline.
  • A.1 End-to-End Training Procedure: For each selected step, OCSD builds structurally matched Full and Observation-Ablated replay views with identical format and scaffold but different future-observation content.The main replay evidence uses a two-observation horizon and a naturalized intervening-action schema; the ablated view replaces observation fields with “Observation: not provided.”
  • A.1 End-to-End Training Procedure: The Student, Full, and Observation-Ablated views score the same student-generated tokens under shared old-policy parameters and the same generated prefix, without generating alternative responses or trajectories.OCSD therefore uses replay views only for training-time scoring of tokens already produced during rollout.
  • A.1 End-to-End Training Procedure: The observation residual contrasts Full and Observation-Ablated token scores, canceling their shared Student-view term and discounting replay-scaffold effects.The residual is mapped to the bounded calibration signal q_i,k,t for subsequent policy optimization.
  • A.1 End-to-End Training Procedure: The bounded signal modulates the magnitude of the original trajectory-level advantage for every token in selected steps, while the advantage sign continues to determine update direction.It is neither an independent reward nor an imitation target, and trajectory-boundary cases use only the available future observation without adding a future-action scaffold.

B Experimental Setup Details

Experiments use a unified post-training and evaluation protocol across three Qwen3 scales and three benchmarks, with OCSD fixed at ρ = 0.2 and β = 0.5. Diagnostic analyses focus on selected high-NLL steps, token-to-word aggregation, functional-role categories, and environment-feedback alignment.

  • Training Protocol: Qwen3-1.7B, Qwen3-4B, and Qwen3-8B are evaluated under a unified protocol with shared optimization and rollout hyperparameters.Published baselines retain their original method-specific configurations.
  • Training Protocol: ALFWorld and WebShop batches contain 16 tasks, Search-QA batches contain 128 tasks, each task uses G = 8 rollouts, and training lasts 150 optimization steps.
  • OCSD Settings: OCSD uses ρ = 0.2 and β = 0.5 across all benchmarks and model scales without benchmark- or size-specific tuning.
  • Evaluation Protocol: Evaluation reports ALFWorld success rates, WebShop task score and success rate, and Search-QA Exact Match across in-domain and out-of-domain datasets.ALFWorld and WebShop use one trajectory per task; Search-QA evaluates seven datasets with at most four interaction steps per question.
  • Diagnostic Analysis: Diagnostic analyses select the top-20% high-NLL steps, aggregate subword signals into word occurrences, and assign words to deterministic functional-role categories.Signals are centered within checkpoints, and word-cloud sizing combines tail preference with log-scaled occurrence frequency.
  • Diagnostic Analysis: Residual-versus-Full feedback alignment is evaluated with signed step-level signals and AUROC, using trajectory-cluster bootstrap confidence intervals.Invalid steps include rejected actions or explicit error/no-effect feedback; valid steps require accepted actions and environment progress.

C.4 Robustness Settings

Robustness analysis evaluates Full and residual scores under raw, standardized, and matched aggregation settings. Paired trajectory-cluster bootstrap intervals preserve within-trajectory comparisons across training stages and action operators.

  • Robustness Settings: Robustness analysis reports the AUROC improvement of the residual over Full support.
  • Robustness Settings: Raw scoring averages token signals within each interaction step and computes AUROC over pooled checkpoints, while Standard applies within-checkpoint z-score normalization before pooling.
  • Robustness Settings: Matched aggregation balances classes and checkpoints by stratifying samples jointly by training stage and action operator before downsampling.Training stages are early, middle, and late, with unmatched operators assigned to other.
  • Robustness Settings: All settings use 1,000 paired trajectory-cluster bootstrap samples, evaluating Full and residual scores on identical resampled trajectories.Confidence intervals are computed directly from paired AUROC differences.

D Observation-Swapping Sensitivity Analysis

Observation swapping tests whether residual sensitivity depends on matching an action with its realized immediate post-action observation. Across all three model scales, replacing that observation with a donor lowers AUROC substantially.

  • Observation Swapping: The swapping analysis uses checkpoints from steps 10, 50, 100, and 150 and samples high-NLL steps stratified into progress, no-progress, and invalid feedback categories.The evaluation set contains 560 steps across Qwen3-1.7B, 4B, and 8B.
  • Observation Swapping: Each donor replacement changes only the immediate observation in the Full view, while the action scaffold, later observation, and shared Observation-Ablated support remain fixed.The paired residual difference therefore isolates the immediate post-action observation supplied to the Full teacher.
  • Observation Swapping: Observation sensitivity averages residual magnitudes across five donor observations and evaluates valid-versus-invalid discrimination with AUROC, excluding no-progress steps.This magnitude-based metric differs from Figure 2(c), which uses signed residuals to test directional alignment with local feedback.
  • Observation Swapping: AUROC falls from 0.803, 0.668, and 0.664 to 0.542, 0.484, and 0.552 for Qwen3-1.7B, 4B, and 8B after observation swapping.Realized immediate post-action observations consistently produce higher AUROC than donor observations from different trajectories.

E Hyperparameter Sensitivity

On ALFWorld with Qwen3-4B, OCSD is most effective with moderate step selection and advantage modulation. The default setting ρ = 0.2 and β = 0.5 reaches the reported peak overall success rate of 82.8.

  • Step Selection Ratio: 82.8 overall success is reached at ρ = 0.2, rising from 70.6 at ρ = 0.1 before performance decreases as more steps are selected.The analysis varies one hyperparameter while fixing the other at its default value.
  • Advantage Modulation: 82.8 overall success is reached at β = 0.5, compared with 78.9 at β = 0.25 and 77.3 at β = 1.0.β = 0 recovers the GRPO baseline, while moderate modulation performs better than overly strong modulation.

F.1 Full Search-QA Results

The section reports detailed Search-QA results across seven datasets and three Qwen3 scales, alongside OCSD ablations and task-specific prompt templates. Search-QA coverage distinguishes two in-domain training datasets from five out-of-domain evaluation datasets and includes an overall micro-averaged score.

  • Prompt Design: All methods use identical interaction prompts for rollout generation and evaluation, while OCSD adds Full and Observation-Ablated variants only during post-rollout replay scoring.The interaction prompts include task instructions, history, current observations, and benchmark-specific action interfaces.
  • Prompt Design: The Full and Observation-Ablated replay prompts share field order and future-action scaffolding, differing only in whether actual future observations are included.In the ablated version, each future-observation field is replaced with “Observation: not provided.”
  • Search-QA Results: Table 6 covers two in-domain and five out-of-domain Search-QA datasets across all three Qwen3 model scales, plus an overall micro-averaged score.Results are reported as mean ± standard deviation over three random seeds for most experiments, with best and second-best methods visually marked.
  • Ablation Results: Table 7 evaluates four OCSD ablations with Qwen3-4B on ALFWorld and WebShop under the same training configuration.The variants remove the Observation-Ablated Teacher, randomize high-NLL step selection, calibrate all interaction steps, or remove sign alignment.
Loading 2608.04788v1…