AHEAD: Adaptive Hindsight with Environment-Augmented Distillation for Agentic RL

Xiaolong Jin; Dingmin Wang; Vijay Lingam; Varun Kumar

cs.AIAug 25, 2026

AHEAD: Adaptive Hindsight with Environment-Augmented Distillation for Agentic RL

Authors: Xiaolong Jin, Dingmin Wang, Vijay Lingam, Varun Kumar

Organizations: AWS AI Labs · Purdue University

Abstract

Training multi-turn LLM agents with reinforcement learning typically relies on trajectory-level rewards, which assign a uniform advantage to every step and cannot identify which decisions led to success or failure. Self-distillation methods can provide finer-grained supervision by augmenting RL with privileged information. However, existing approaches usually apply the same type of privileged information to every step in an indistinguishable manner, ignoring a key asymmetry: routine steps need little additional guidance, while critical error steps require corrective direction that environment feedback alone cannot provide. We propose AHEAD, a step-aware framework that matches different supervision sources to different step types. The teacher receives environment feedback on all steps as a grounded dense signal, and additionally receives LLM-generated corrective hints on error steps to supply the direction that environment feedback lacks. The method introduces minimal changes to the standard GRPO algorithm. Across ALFWorld, WebShop, and Search-based QA, and across three model scales, AHEAD raises task success (+13.3 points on ALFWorld and +11.0 on WebShop at 7B over GRPO), reaches a given success rate in fewer training steps, and solves tasks within tighter interaction budgets than outcome-only RL and prior self-distillation baselines.

Figures & tables

Figure 1: Existing self-distillation methods apply a single source of privileged information (PI) uniformly across all steps (top): on error steps such as

t_{3}

and

t_{5}

, environment feedback (e.g., “Nothing happens”) only signals failure and provides no corrective direction. Ahead (bottom) uses environment feedback as base PI on all steps, and adaptively provides an LLM-generated corrective hint on error steps, telling the agent what it should have done.

Figure 2: Overview of Ahead . Stage 1 identifies error steps in failed trajectories via an LLM analyzer. Stage 2 constructs step-aware privileged information: environment feedback for routine steps, environment feedback combined with an LLM corrective hint for error steps. Stage 3 computes a token-level self-distillation signal $\delta_{t,\ell}$ by comparing the policy’s log-probabilities under original and PI-augmented contexts, and converts it into a bounded reweight of the GRPO advantage. Successful trajectories bypass the PI pipeline entirely and retain the vanilla GRPO advantage.

Qwen2.5-3B-Instruct
	ALFWorld							Search-based QA								WebShop
Method	Pick	Look	Clean	Heat	Cool	Pick2	Avg	NQ	Triv	Pop	Hotp	2Wk	MuS	Bam	Avg	Score	Succ.
Vanilla	44.4	11.1	6.2	15.4	28.6	12.5	21.9	24.6	48.1	31.0	26.3	25.3	7.2	59.7	31.7	6.7	0.8
Skill-Prompt*	51.7	66.7	48.4	0.0	4.3	10.0	28.9	23.7	46.2	30.6	24.4	22.1	7.5	12.5	23.9	0.2	0.8
OPSD	48.8	41.7	16.7	0.0	15.8	16.7	28.1	0.1	0.1	0.1	0.0	0.0	0.0	0.0	0.0	11.3	3.1
GRPO	91.2	62.5	96.2	61.9	65.0	47.4	75.0	39.3	60.6	41.1	37.4	34.6	15.4	26.4	36.4	79.8	63.3

Table 1: Performance Comparison on long-horizon benchmarks. We report the success rate (%) on ALFWorld, accuracy on search-based QA, and task-completion score/success rate on WebShop. An asterisk (*) denotes validation with skills. Within each backbone block, the best result in a column is in bold and the second-best is underlined. Baseline numbers are taken from Lu et al. (2026b) .

	Components					ALFWorld			WebShop
Configuration	Env. fb.	Hints	Multi	Fail-only	Decay	1.7B	3B	7B	1.7B	3B	7B
Ahead	✓	✓	✓	✓	✓	67.2	87.5	94.5	64.8	73.4	83.6
w/o decay	✓	✓	✓	✓	$\times$	60.2	71.1	85.2	54.7	66.4	83.6
w/o failure-only	✓	✓	✓	$\times$	✓	48.4	82.0	93.0	64.1	72.7	75.8
w/o multi-step	✓	✓	$\times$	✓	✓	61.7	81.2	86.7	57.0	47.7	75.0
w/o hints	✓	$\times$	–	✓	✓	53.1	73.4	87.5	63.3	68.0	74.2

Table 2: Component ablation. Success rate (%) on ALFWorld and WebShop. Ticks give the configuration of each run, so rows differing in a single tick are one-component comparisons. Env. fb. : environment feedback as base PI. Hints : the LLM-generated corrective hint

\Phi^{\mathrm{llm}}

on error steps. Multi : the analyzer may mark several error steps per trajectory rather than one. Fail-only : PI is applied to failed trajectories only (Sec. 2.3 ). Decay : the linear schedule on

\lambda_{k}

. The last row removes both the hints and the failure-only filter, leaving environment feedback applied uniformly.

Analyzer	3B	1.7B
Opus 4.7 (default)	87.5	67.2
Sonnet 5	89.1	68.0
Kimi	83.6	68.0
GLM-5	81.2	67.2
GRPO (no analyzer)	75.0	46.1

Table 3: ALFWorld success rate (%) when the error-step analyzer is swapped, holding everything else fixed. GRPO, which uses no analyzer, is repeated from Table 1 as a floor.

Figure 5: Trajectory comparison on an ALFWorld task put a hot plate in cabinet . Left: a GRPO-trained policy overlooks the visible plates, revisits the same countertops and empty cabinets, and never takes, heats, or moves an object before the 50-step limit (failed trajectory). Right: a Ahead -trained policy grounds each action in the current observation and completes the find

\rightarrow

pick

\rightarrow

heat

\rightarrow

place sequence in nine steps. Reasoning and observation excerpts are lightly condensed for space; “

\ldots

” marks omissions and retained text is verbatim.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Benchmark	#Train	#Test
ALFWorld	2,400	128 (Val-128) / 140 (Seen-140) / 134 (Unseen-134)
WebShop	2,400	128
NQ, TriviaQA, PopQA, HotpotQA,	19,200	51,713
2WikiMultiHopQA, MuSiQue, Bamboogle

Appendix

Table 4: Datasets used in our experiments: ALFWorld (embodied reasoning), WebShop (web navigation), and Search-based QA. For Search-based QA we train only on NQ and HotpotQA; the remaining five datasets are held out and never seen during training.

Table 9

Figure 8: Training dynamics on ALFWorld (7B); the success-rate panel is measured on the seen split. AHEAD attains a higher success rate and train reward throughout training, while driving average trajectory length down faster. The agent solves more tasks and does so in fewer steps.

Figure 9: Training dynamics on ALFWorld (Qwen3-1.7B-Instruct); the success-rate panel is measured on the seen split. Companion to Figure 8 .

Panel of Figure 11 — Figure 11: Token-level credit on one failed trajectory (ALFWorld). Top: signed gap $\delta_{t,\ell}$ ; red marks tokens the teacher disfavors ( $\delta\!<\!0$ ). At the error step ( $t\!\in\!\mathcal{E}_{\tau}$ ) the tokens for the wrong decision have the largest $|\delta|$ ; the routine step ( $t\!\notin\!\mathcal{E}_{\tau}$ ) stays near zero. Bottom: the weight $w_{t,\ell}=\mathrm{clip}(\exp(\mathrm{sgn}(A_{\mathrm{ep}})\delta_{t,\ell}),1\!-\!\epsilon,1\!+\!\epsilon)$ . Since $A_{\mathrm{ep}}\!<\!0$ on failed trajectories, these tokens get the largest amplification ( $w\!>\!1$ ).

Hyperparameter	Value
Training steps	150
Training batch size	16 for ALFWorld and WebShop; 128 for Search-based QA
Rollout group size $N$	8
Learning rate	$1\times 10^{-6}$
Hardware	8 $\times$ H100
GRPO clipping range $\epsilon$	0.2

Appendix

Table 7: Training configuration. The block in the middle lists the parameters introduced by Ahead (Sec. 2.4 ); everything else is inherited unchanged from the GRPO baseline, so the comparison in Table 1 isolates the effect of step-aware distillation.

Panel of Figure 12 — Figure 12: Prompt template for the ALFWorld environment.

Panel of Figure 13 — Figure 13: Prompt template for the WebShop environment.

Panel of Figure 14 — Figure 14: Prompt template for the Search-based QA environment.

Panel of Figure 15 — Figure 15: Prompt template for the LLM error-step analyzer.

Explore similar work

What and When to Distill: Selective Hindsight Distillation for Multi-Turn Agents
May 19, 2026Xiaozhe Li, Tianyi Lyu, Yang Li +6Reinforcement Learning Credit Assignment in RL
1Tongji University · 2Shanghai AI Laboratory · 5Independent +2
HINT-SD: Targeted Hindsight Self-Distillation for Long-Horizon Agents
May 18, 2026Woongyeng Yeo, Yumin Choi, Taekyung Ki +1Credit Assignment in RL LLM Agent Training
KAIST
Skill-SD: Skill-Conditioned Self-Distillation for Multi-turn LLM Agents
Apr 12, 2026Hao Wang, Guozhi Wang, Han Xiao +8RL for LLM Agents LLM Agent Skill Learning
Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences · vivo AI Lab · The Chinese University of Hong Kong +2

May 19, 2026J/K move · Enter open · S save

What and When to Distill: Selective Hindsight Distillation for Multi-Turn Agents

Xiaozhe Li, Tianyi Lyu, Yang Li, Yichuan Ma+5

1Tongji University · 2Shanghai AI Laboratory · 5Independent+2

Reinforcement Learning Credit Assignment in RL Step-Level Credit Assignment in RL RL for LLM Agents

Open paper AskPDF arXiv

Reinforcement learning can train LLM agents from sparse task rewards, but long-horizon credit assignment remains challenging: a single success-or-failure signal must be distributed across many actions. Existing methods rely on trajectory-level rewards or proxy signals, without fully leveraging per-step environmental feedback. Multi-turn agent settings are underexplored, where feedback can include error messages, page changes, observations, or reference trajectories. We systematically study five feedback sources and two insertion granularities and introduce SERL, a selective environment-reweighted learning framework. SERL uses the task reward to determine update direction, while environment feedback adjusts placement and magnitude, focusing on critical actions. On ALFWorld and WebShop, SERL achieves 90.0% and 80.1% success, outperforming strong RL and distillation baselines. Analysis shows that grounded, action-relevant feedback at meaningful points consistently outperforms indiscriminate use of longer or richer context.

Abstract

Explore similar work

What and When to Distill: Selective Hindsight Distillation for Multi-Turn Agents

HINT-SD: Targeted Hindsight Self-Distillation for Long-Horizon Agents

Skill-SD: Skill-Conditioned Self-Distillation for Multi-turn LLM Agents