Agentic Reinforcement Learning with Self-Distilled Reward Shaping

2026-08-04

AI Summary

This work addresses the challenge of credit assignment in large language model agents under sparse trajectory-level rewards in reinforcement learning. To this end, the authors propose ADRS, a self-distillation reward shaping framework that re-evaluates tokens in skill-agnostic trajectories using frozen policy snapshots and generates dense supervision signals aligned with task-specific procedural skills. ADRS seamlessly integrates token-level credit assignment with native advantage estimation through teacher-score centralization normalization and a return-correlated confidence gating mechanism. Experimental results across three interactive benchmarks demonstrate that ADRS significantly enhances performance on long-horizon tasks and exhibits robustness across diverse reinforcement learning backbones, low-data regimes, unseen tasks, and extended training settings.

Abstract

Agentic reinforcement learning enables LLM agents to learn through interaction, but sparse trajectory-level rewards reveal success without identifying which intermediate decisions deserve credit. Training-only privileged skills can provide denser supervision by allowing the same frozen policy snapshot to rescore fixed tokens from skill-free trajectories while conditioned on task-matched procedural skills. Existing methods, however, do not jointly calibrate teacher scores across interaction steps, relate teacher confidence to realized returns, and integrate the resulting signal into native reward-to-advantage construction. We introduce Agentic Reinforcement Learning with Self-Distilled Reward Shaping (ADRS), a framework for constructing return-associated token-level credit for multi-turn language agents. ADRS centers and normalizes privileged token scores within each step, modulates them with a return-associated Teacher Value Advantage (TVA) gate based on within-group confidence--return association, and incorporates the gated token signal into native RL credit construction. Together, these components determine what the teacher prefers, when that preference is return-relevant, and how it enters the native reinforcement-learning credit path, while keeping rollouts and inference skill-free. Finally, experiments across three interactive benchmarks show that ADRS consistently improves performance on long-horizon tasks, with gains persisting across RL backbones, reduced-data settings, unseen tasks, and extended training. For anonymous review, our code is available at the following the link: https://github.com/gitrxh/ADRS-arxiv

Problem

• Agentic Reinforcement Learning
• Reward Shaping
• Credit Assignment
• Sparse Rewards
• Language Agents

Innovation

• Self-Distilled Reward Shaping
• Token-Level Credit Assignment
• Teacher Value Advantage
• Agentic Reinforcement Learning
• Privileged Skill Distillation