Branch-Grafted On-Policy Distillation for Agents Solving Capture-the-Flag Challenges

On-policy distillation (OPD) trains a student policy using teacher supervision on trajectories generated by the student, thereby reducing the distribution mismatch between offline supervision and the states encountered during deployment. For multi-turn agents, we study how to make fuller use of teacher supervision at interaction histories already visited by the student. Standard OPD follows the response sampled by the student, while alternative teacher responses from the same history are not directly used for distillation. We introduce **Branch-Grafted On-Policy Distillation (BG-OPD)**, which augments student trajectories with single-turn teacher-generated response branches at selected interaction histories. BG-OPD uses teacher–student residual disagreement to allocate branch supervision and distills the resulting teacher responses without executing their proposed actions in the environment. We further show that the branch objective can be interpreted as a turn-level approximation to residual-weighted teacher supervision and characterize the corresponding approximation error. We evaluate BG-OPD on interactive cybersecurity tasks, where long-horizon hypothesis testing and tool use make alternative response paths particularly important. To support on-policy training at scale, we construct 1,370 containerized CTF environments whose solvability is verified through independently executed solutions. Across Cybench, NYU CTF Bench, and XBOW, BG-OPD consistently outperforms offline SFT and existing OPD baselines for both 4B and 9B students, improving mean task success from 13.20% to 21.01% over the strongest baseline for the 9B model and from 7.70% to 14.76% for the 4B model. Ablations further attribute the improvement to teacher-generated branches and disagreement-based branch selection.