Graphaugmented distillation improves large language model agent learning
GraphOPD: Graph-Augmented On-Policy Distillation for LLM Agents
Machine LearningArtificial Intelligence
Summary
Training AI agents to learn from sparse feedback is tricky because they only get rewards at the end of tasks, not at every step. Existing methods try to fix mistakes by focusing on steps where the learner and teacher AI disagree most, but this can miss the real causes of errors that happen earlier. The authors create GraphOPD, which uses a map of how each step depends on previous ones to better identify which steps deserve correction. This method works better than past approaches on several tasks and model sizes by combining disagreement with step importance identified in this map.
What this means in practice
- •For ai developers: Improve training efficiency of AI agents learning from sparse rewards by focusing corrections on structurally important steps.
- •For automated customer support teams: Enhance chatbot learning by systematically correcting key misunderstandings that affect multi-turn conversations.
Authors
Bohan Lin, Liyi Chen, Zhuoning Guo, Muyang Li, Qimeng Wang, Yan Gao, Yao Hu, Yudong Zhang
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.