GraphOPD: Graph-Augmented On-Policy Distillation for LLM Agents

2026-10-06

AI Summary

This study addresses the issue wherein traditional distillation misidentifies critical steps due to trajectory drift in LLM agents operating under sparse rewards. To this end, we propose a graph-structure-enhanced online distillation method that constructs dependency graphs from environment states and integrates random walk scoring with divergence signals to precisely identify high-value steps. This work represents the first introduction of structural graph modeling into agent online distillation for causal credit assignment. Experimental results demonstrate that the proposed approach outperforms the strongest baseline by 5.8 percentage points on benchmarks such as ALFWorld, validating both the causal efficacy of structure-based credit assignment and its cross-domain transferability.

Abstract

On-policy distillation post-trains large language model agents by supplying dense, step-level guidance from a teacher policy when the reinforcement-learning reward is sparse and arrives only once per trajectory. Existing instantiations allocate this guidance by the size of the teacher-student divergence at each step, on the single-turn intuition that a large disagreement marks a mistake worth correcting. Once decisions chain over many turns, that rule misfires, since an early drift enters every later context both policies condition on, leaving the teacher consistent with the drifted trajectory instead of flagging its cause, while interchangeable steps register large but outcome-irrelevant divergences. We demonstrate this on an agentic benchmark, where distilling the highest-divergence steps brings no consistent benefit over random selection. To this end, we introduce GraphOPD, the first method to bring graph-based structural augmentation into on-policy distillation for agent capabilities. It reads which steps enabled which later ones from the environment's own record of state changes, immune to the drift that corrupts the teacher-student gap, organizes them into a dependency graph, scores each step by a random-walk stationary distribution over it, and fuses that structural credit with the divergence signal into a trajectory-relative mask concentrating supervision on each rollout's highest-aptitude steps. Across three model scales and eleven baselines on ALFWorld, WebShop, and SearchQA, GraphOPD shows competitive performance throughout, improving over the strongest baseline by up to +5.8 pp. An executed-replay audit further shows that this structural credit score tracks true causal impact far above chance, that both fused signals are independently necessary, and that the same signal transfers to out-of-domain tool-integrated reasoning.

Problem

• On-policy distillation
• LLM agents
• Step-level guidance
• Trajectory drift
• Teacher-student divergence

Innovation

• On-Policy Distillation
• Graph-Augmented
• Dependency Graph
• Random Walk
• Structural Credit