🤖 machine learning

Agentic Reinforcement Learning with Self-Distilled Reward Shaping

This paper introduces Agentic Reinforcement Learning with Self-Distilled Reward Shaping (ADRS), a framework that enhances long-horizon language agent performance by generating return-associated token-level credit through self-distilled, privileged skill signals that are jointly calibrated and gated by teacher confidence and realized returns.

Original authors: Ranxu Zhang, Guinan Chen, Chenshaodong, Jinghao Lin, Xiaozhou Xu, Sunzhe, Yanyong Zhang, Chao Wang

Published 2026-08-05

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Ranxu Zhang, Guinan Chen, Chenshaodong, Jinghao Lin, Xiaozhou Xu, Sunzhe, Yanyong Zhang, Chao Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to navigate a giant, messy house to find a specific item, like a warm cup of tea. You can't just tell the robot, "Good job!" only after it finally finds the tea. If you wait that long, the robot won't know which specific steps were helpful—did it open the right door? Did it pick up the right mug? Or did it just wander around for ten minutes before getting lucky? This is the core puzzle of Reinforcement Learning: how do you give credit to the tiny, intermediate decisions that lead to a big success, especially when the final reward is rare and far away?

To solve this, scientists often use a trick called Self-Distillation. Think of it like a student taking a test, and then immediately after, the same student acts as a strict teacher to grade their own answers. The "teacher" version has a special reference sheet (called "privileged information") that the "student" didn't have while taking the test. By comparing what the student did with what the teacher would have done with the reference sheet, we can figure out which moves were smart and which were lucky. However, there's a catch: just because the teacher thinks an answer is "good" doesn't mean it actually led to the right outcome. The teacher might be confident but wrong, or confident about the wrong thing.

This brings us to a new study by researchers at the University of Science and Technology of China and Alibaba Group. They tackled the problem of how to take that "teacher's confidence" and turn it into a reliable guide for the robot's next move, without ever letting the robot use the reference sheet during the actual game. They call their new method ADRS (Agentic Reinforcement Learning with Self-Distilled Reward Shaping).

The researchers found that simply copying the teacher's score doesn't work because the teacher's confidence can be misleading. Sometimes the teacher is very sure about a move that leads to failure, or unsure about a move that leads to success. To fix this, ADRS acts like a smart filter. First, it normalizes the teacher's scores so they are fair to compare across different parts of the task. Then, it checks if the teacher's confidence actually matches the final result (did the confident move lead to a win?). If the confidence and the result agree, the filter boosts that move's score; if they disagree, it dampens the signal. Finally, it injects this refined signal directly into the robot's learning process before the robot decides what to do next.

The results suggest this method is quite effective. When tested on three different challenging tasks—navigating a virtual house, shopping on a simulated website, and searching for answers online—ADRS consistently helped the AI agents perform better, especially on long, complicated tasks. For example, on the house-navigating task, the method helped a specific AI model achieve a 94.5% success rate, beating previous best attempts by a significant margin. The study also showed that the AI learned these skills using less training data than usual and could still perform well even when the tasks got slightly different or harder. Crucially, the researchers demonstrated that the AI didn't need the "reference sheet" during the actual game; it only needed it during the training phase to learn how to make better decisions on its own. This suggests that by carefully calibrating how much we trust a "teacher's" opinion based on real-world results, we can teach AI agents to be much more reliable and efficient at solving complex, multi-step problems.

Technical Summary: Agentic Reinforcement Learning with Self-Distilled Reward Shaping (ADRS)

Problem Statement

Agentic Reinforcement Learning (RL) enables Large Language Model (LLM) agents to learn through interaction with environments. However, standard agentic RL typically relies on sparse, trajectory-level terminal rewards. While these rewards indicate whether a trajectory succeeded, they fail to identify which intermediate decisions (queries or actions) deserve credit. This sparsity makes token-level temporal credit assignment difficult, even for critic-free group-relative methods like GRPO and GiGPO.

While natural-language skills (e.g., subgoals, action rules) can provide denser guidance, existing methods struggle to integrate this information effectively. Specifically, current approaches face three technical challenges when using "privileged" information (skills available only during training):

Score Calibration: Raw teacher log-probabilities vary in offset and scale across interaction steps due to context length and response composition, making them incomparable.
Reliability Estimation: High teacher confidence does not necessarily correlate with higher realized returns; a teacher may prefer linguistically canonical decisions that fail to achieve the task outcome.
Credit Integration: Existing methods often treat privileged supervision as a separate auxiliary objective, which may misalign token-level updates with the native reward-to-advantage construction of the RL backbone.

Methodology: ADRS Framework

The authors propose Agentic Reinforcement Learning with Self-Distilled Reward Shaping (ADRS), a framework that converts task-matched procedural skills into return-associated token-level rewards without requiring skills during rollout or inference. ADRS operates as a credit-construction layer between trajectory collection and actor optimization, addressing the three challenges via three coupled modules:

1. Within-Step Relative Privileged Scoring (Score Calibration)

To address the incomparability of raw scores across steps, ADRS rescores student-generated tokens using a frozen policy snapshot under a privileged context (task-matched skills). It then applies within-step centering and scale normalization:

It calculates the mean log-probability ( $b_s$ ) of the privileged scores within a specific interaction step.
It computes the signed deviation ( $q_{s,j}$ ) of each token from this mean.
It normalizes by the standard deviation ( $\sigma_q$ ) of deviations within the step.
This yields a centered, scale-controlled score ( $\hat{q}_{s,j}$ ) that represents relative token preference within a decision step, suppressing step-dependent drift while retaining the raw scores for subsequent reliability estimation.

2. Return-Associated TVA Reliability Gating (Reliability Estimation)

To ensure the privileged signal aligns with actual task success, ADRS introduces a Teacher Value Advantage (TVA) gate. This gate modulates the normalized signal based on the association between privileged confidence and realized returns within a comparison group (e.g., trajectories for the same prompt in GRPO, or anchor states in GiGPO).

The method computes a soft high-confidence weight ( $\alpha_u$ ) for each unit (trajectory or step) based on its average privileged log-probability.
It contrasts the returns of high-confidence units ( $\mu^+_g$ ) against low-confidence units ( $\mu^-_g$ ).
The difference ( $d_g$ ) serves as a scalar gate ( $m_u$ ). If higher confidence correlates with higher returns, the gate amplifies the signal; if not, it suppresses it.
Proposition 1 establishes that the sign of the TVA gate corresponds to the empirical covariance between privileged confidence and return within the group.

3. Teacher-Guided Per-Token Advantage Modulation (Credit Integration)

Instead of using a separate auxiliary loss, ADRS integrates the gated signal into the native reward-to-advantage path before advantage construction.

Teacher Reward: The gated, normalized score is converted into a token-level reward ( $r^T_{s,j} = \eta m_{s,j} \hat{q}_{s,j}$ ).
Shaped Reward: This is added to the base environment reward ( $\tilde{r}_{s,j} = r^{base}_{s,j} + r^T_{s,j}$ ).
Advantage Construction: The shaped reward is passed through the standard advantage operator (e.g., GRPO or GiGPO).
Dual Path: Because the step-centered teacher reward sums to zero (Proposition 2), the trajectory-level sum remains unchanged, preserving the environment-defined outcome direction. However, the token-level advantage is redistributed ( $A^{ADRS} = A^{traj} + \eta Z^T$ ), allowing the policy to learn which specific tokens within a successful trajectory were most critical.

Key Contributions

ADRS Framework: A self-distilled reward-shaping method that transforms training-only procedural skills into return-associated token rewards. It uniquely combines within-step score calibration and TVA gating while keeping rollouts and inference skill-free.
Native Integration: The authors formulate the gated signal as pre-advantage reward shaping, integrating privileged guidance directly into the native RL credit construction rather than as a separate auxiliary objective. They characterize its compatibility with GRPO/GiGPO and its local first-order equivalence to sampled-token auxiliary updates.
Empirical Validation: Extensive experiments across three interactive benchmarks (ALFWorld, Search-based QA, WebShop) demonstrate consistent performance improvements.

Experimental Results

The paper evaluates ADRS on Qwen2.5 (3B, 7B) and Qwen3 (1.7B) models across three benchmarks:

State-of-the-Art Performance: ADRS achieves new state-of-the-art results on all three benchmarks. For example, with Qwen2.5-3B, it reaches 94.5% success on ALFWorld (a 10.1 percentage point improvement over the strongest baseline) and 45.0% macro-average on Search-based QA.
Long-Horizon Tasks: ADRS shows particular strength in long-horizon procedural tasks (ALFWorld), sustaining performance gains over extended training (300 steps) and across different RL backbones (GRPO and GiGPO).
Data Efficiency: ADRS surpasses full-data GRPO baselines using only 60% of the available training data.
Generalization: The method transfers effectively to unseen task splits in ALFWorld, improving success rates across most categories.
Mechanism Analysis: Diagnostic analysis reveals that ADRS concentrates update contrast on task-bearing action and object tokens, separating them from generic style tokens, which explains its efficacy in long-horizon interactions.

Significance and Claims

The paper claims that ADRS successfully addresses the fundamental challenge of assigning credit to intermediate decisions in agentic RL. By converting training-only privileged knowledge into a signal that is calibrated (comparable across steps), reliable (associated with returns), and integrated (native to the RL objective), ADRS enables agents to learn more effectively from sparse feedback.

The authors emphasize that ADRS does not require privileged information during inference, making it a practical framework for deploying skill-free agents that have learned from skill-guided training. The results suggest that training-only procedural knowledge can be effectively converted into token-level credit and directly incorporated into the native reward-to-advantage construction, offering a robust path for improving multi-turn language agents.