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.
Figures & tables
Figure 1: Resampling from one prefix on a mathematics problem and a logged ALFWorld trajectory. Scores mark the fork; q is the student probability of each continuation’s first token.
Figure 2: Overview of BranchOPD. Fork points are selected on the main trajectory, suffixes are resampled from the student, and the teacher scores every generated token; shared prefixes are counted once. The bottom examples show concrete PBR constraints used in Section 5.5 .
Figure 3: ALFWorld generation cost and success with PBR, for three (left) and seven (right) trajectory slots. OPD, Independent, and Reweighting denote Vanilla OPD, Multi-Rollout OPD, and Entropy-Weighted OPD.
Figure 4: Training dynamics with PBR: mathematics response length and truncation rate (fraction of responses stopped at the length limit), and ALFWorld turn-level protocol violation (fraction of turns with a violation). Repetition penalties delay budget exhaustion but leave generation near the cap; PBR reduces both truncation and violations.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Trajectory-level protocol violation (PVP): the fraction of the 192 main and branch trajectories that contain a violation. All methods shown use PBR; BranchOPD-F denotes forced-token expansion (Table 8 ).
Figure 6: Probability mass available for alternative decisions. (a) “Mass kept by the main path” is the probability of the sampled main token, πθold(a0∣s) . (b) “Mass discarded” and MK denote the alternative-token mass 1−πθold(a0∣s) , i.e., the mass left unrealized by the main draw. (c) All supervised positions compared with selected forks. ALFWorld and WebShop use training dumps; mathematics and Search-QA use reconstructed summaries. In the plot, q denotes πθold and turn-open-RS denotes Turn-initial branching.
Figure 7: Empirical CDF of alternative mass over supervised positions (ALFWorld step 200, WebShop step 210), accumulated over 20 equal-width bins. BranchOPD-F denotes forced-token expansion (Table 8 ), and “uniform” denotes uniform weights on its branches.
Figure 8: Behavior coverage of resampled branches over training. “New state” in the plot denotes new-action coverage (NSR). High disagreement and NSR mean that branches reach behavior not present in the main trajectory.
Figure 9: Rescue and harm rates during training, each conditioned on the outcome of the main trajectory. “Mean net” is R↑−R↓ averaged over training steps. “Spoil” means that all branches fail although the main trajectory succeeds; the main trajectory itself stays in the training batch.
Figure 10: Fork location and distance to the first changed action. (a) Fork position relative to the full trajectory and (b) relative to the selected turn, where 0 and 1 mark the start and end. (c) Number of generated student tokens from the fork to the first changed action.