Branching method improves teacher guidance for student agents
Branching On-Policy Distillation
Artificial Intelligence
Summary
Teaching AI agents is hard because they learn from only one possible action at a time, limiting the examples they see. The authors propose BranchOPD, which lets agents explore multiple possible next steps at key moments, guided by where the agent is uncertain or disagrees with the teacher. This approach introduces efficient reusing of shared early actions to save time and gives more detailed guidance at each step. Experiments show BranchOPD makes AI agents better at math problems and interactive tasks compared to previous methods.
What this means in practice
- •For ai system developers: Improve training of AI agents that solve interactive tasks by using branching to cover more decision scenarios efficiently.
- •For automated tutoring platforms: Create more accurate AI tutors in domains like math by enhancing trajectory supervision with branching and targeted corrections.$Commercial implications: Enables development of smarter automated tutoring systems that provide better step-by-step guidance by improving AI answer accuracy.
Authors
Dongyang Chen, Shudong Liu, Shubo Li, Duoxun Tang, Jinwei Liang, Zehuan Wu, Lewei Lu
Abstract
On-policy distillation (OPD) provides teacher supervision on student-generated trajectories, but a single trajectory realizes only one continuation at each decision, limiting supervision coverage under a finite sampling budget. We propose BranchOPD (Branching On-Policy Distillation), which selects informative decision points using student uncertainty and teacher-student disagreement, resamples student-policy continuations from shared prefixes, and applies token-level teacher supervision to both main trajectories and branches. BranchOPD operates at token-level decisions in mathematics and within-turn decisions in agent tasks. Shared-prefix reuse reduces student-generated token costs by approximately 19% on ALFWorld and 31% on WebShop relative to seven independent trajectories. We also introduce protocol-boundary repair (PBR), which uses targeted supervision under protocol constraints to mitigate repetition and wasted generation. BranchOPD outperforms Vanilla OPD and Multi-Rollout OPD (n=7) across all evaluated benchmarks, with gains over Vanilla OPD of 18.55 percentage points in mean mathematics accuracy over 16 sampled answers (avg@16) and 13.8 points in mean ALFWorld ID/OOD success. Mechanistic analyses and selection ablations show that branching at uncertain decisions expands subsequent behavior coverage and outperforms random selection and turn-initial branching.