AHEADNot every step deserves the same feedback. AHEAD gives multi-turn agents step-aware supervision: grounded environment feedback everywhere, and an LLM corrective hint exactly where the agent went wrong.
1AWS AI Labs·2Purdue University
† Work done during an internship at Amazon.
Across ALFWorld, WebShop, and Search-based QA, and three model scales (Qwen2.5-3B/7B, Qwen3-1.7B).
Trajectory-level RL broadcasts one advantage to every token, and prior self-distillation injects the same privileged information (PI) at every step. Both ignore a basic asymmetry between step types.
Task-level PI (retrieved skills, reference answers) applied identically at every step. On an error step, environment feedback like "Nothing happens" only signals failure, with no direction on what to do instead.
Environment feedback on all steps; on error steps an LLM corrective hint ("go to diningtable 1 first") adds the missing direction, a stronger signal exactly where it matters, with no explicit gating.
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 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.
See the main figure below for the full pipeline.
An LLM analyzer reads each failed trajectory and marks the steps whose actions were critical errors: a wrong object, an invalid action, an unproductive query.
Environment feedback is the base PI on every step. Error steps additionally get an LLM corrective hint, so they carry a naturally stronger signal, with no per-step coefficient or gate.
A token-level log-prob gap between the original and PI-augmented contexts becomes a bounded multiplicative reweight of the GRPO advantage that preserves the reward's sign.
The teacher–student log-probability gap measures how the PI revises each token:
$$\delta_{t,\ell} = \log \pi_{\theta_{\text{old}}}(y_{t,\ell}\mid \tilde{h}_t, y_{t,<\ell}) - \log \pi_{\theta_{\text{old}}}(y_{t,\ell}\mid h_t, y_{t,<\ell})$$It becomes a sign-preserving reweight of the episode advantage, with a mixing weight $\lambda_k$ that decays $0.5 \to 0$ over training:
$$w_{t,\ell} = \text{clip}\!\big(\exp(\text{sgn}(A^{\text{ep}})\,\text{sg}(\delta_{t,\ell})),\; 1-\varepsilon,\; 1+\varepsilon\big), \qquad \tilde{A}_{t,\ell} = A^{\text{ep}}\cdot\big[(1-\lambda_k) + \lambda_k\, w_{t,\ell}\big]$$Because the reweight is positive, $\text{sign}(\tilde{A}) = \text{sign}(A^{\text{ep}})$, so the environment reward always controls the update direction.
Headline metrics per scale; the full 18-column breakdown is in the paper. All baseline numbers follow a common backbone, environment wrapper, and rollout budget.
Table 1: Main results across benchmarks and scales
ALFWorld success (%, Val-128) · Search-QA macro-average accuracy (%) · WebShop success (%). AHEAD rows are shaded; the best value per column-block is highlighted.
| Method | ALFWorld | Search-QA | WebShop Succ. |
|---|---|---|---|
| Qwen2.5-3B-Instruct | |||
| GRPO | 75.0 | 36.4 | 63.3 |
| RLSD | 79.7 | 43.8 | 66.4 |
| Skill-SD | 73.4 | 44.1 | 64.0 |
| SDAR | 84.4 | 43.4 | 68.0 |
| AHEAD | 87.5 | 44.3 | 73.4 |
| Qwen2.5-7B-Instruct | |||
| GRPO | 81.2 | 42.0 | 72.6 |
| RLSD | 82.0 | 49.0 | 77.3 |
| Skill-SD | 85.1 | 47.8 | 76.5 |
| SDAR | 85.9 | 49.0 | 82.8 |
| AHEAD | 94.5 | 48.5 | 83.6 |
| Qwen3-1.7B-Instruct | |||
| GRPO | 46.1 | 40.8 | 38.3 |
| RLSD | 42.2 | 40.6 | 50.8 |
| Skill-SD | 52.3 | 40.8 | 53.9 |
| SDAR | 53.9 | 41.9 | 58.6 |
| AHEAD | 67.2 | 41.9 | 64.8 |
Table 2: Component ablation (ALFWorld success, %)
Removing any single component hurts; no variant matches the full method. Dropping the hints costs 7.0–14.1 points and dropping environment feedback costs 4.7–7.8, so both PI sources contribute while the corrective direction is the larger effect. "Env. feedback only" drops both the hints and the failure-only filter.
| Configuration | 1.7B | 3B | 7B |
|---|---|---|---|
| AHEAD (full) | 67.2 | 87.5 | 94.5 |
| − decay | 60.2 | 71.1 | 85.2 |
| − failure-only | 48.4 | 82.0 | 93.0 |
| − multi-step | 61.7 | 81.2 | 86.7 |
| − hints | 53.1 | 73.4 | 87.5 |
| − env. feedback | 59.4 | 81.2 | 89.8 |
| Env. feedback only | 56.2 | 72.7 | 88.3 |
Swapping the LLM analyzer (Opus 4.7, Sonnet 5, Kimi, GLM-5) leaves every variant well above GRPO; at 1.7B the four fall within 0.8 points of each other, so a cheaper analyzer suffices for smaller backbones.
@inproceedings{jin2026ahead,
title = {AHEAD: Adaptive Hindsight with Environment-Augmented
Distillation for Agentic RL},
author = {Jin, Xiaolong and Wang, Dingmin and Lingam, Vijay and Kumar, Varun},
year = {2026},
note = {Under review}
}