TrajPD: Constructing Trajectory-Aware Targets for On-Policy Distillation in Tool-Integrated Reasoning

Transferring tool-integrated reasoning (TIR) capabilities to small language models (SLMs) benefits from supervision adapted to their multi-turn interactions. On-policy distillation (OPD) provides dense teacher guidance on trajectories generated by the student. In multi-turn TIR, student responses and tool feedback shape subsequent prediction contexts, and the tokens within each complete response jointly form an action or answer. This motivates constructing targets that remain close to the student’s current predictions while retaining a specified teacher contribution across each response. We propose Trajectory-Aware Policy Distillation (TrajPD), which constructs targets jointly within each complete response. TrajPD selects targets that minimize the KL divergence from the targets to the student’s predictions, averaged over positions within the response, while satisfying a response-level constraint on the retained teacher contribution. At the sampled contexts and the start of an update, the weighted distillation gradient preserves the prescribed component along the direct teacher gradient in logit space, aggregated across the response. The targets are obtained by adjusting the student’s log probabilities along directions derived from differences between teacher and student probabilities, using a single scale for each response. A weight determined by the average divergence between teacher and student predictions within the response determines this scale and also weights the distillation loss. Evaluation on four benchmarks spanning mathematical reasoning, scientific reasoning, and code generation shows improved aggregate performance at both student scales, with gains across multiple tasks and evaluation metrics.