Building Multi-Task Agentic LLMs via Two-Phase Distillation
This work provides a practical method for consolidating multiple RL-trained agents into one multi-task model, which is important for scaling LLMs to handle diverse tasks without performance degradation.
The paper addresses the challenge of building multi-task agentic LLMs by distilling separate RL experts into a single model. The proposed two-phase distillation (off-policy then on-policy) matches single-task expert performance across conversational agents and text-based games, while either phase alone fails.
A key step toward artificial general intelligence is to train models that can perform multiple tasks. In this paper, we study how to build such models by first training separate RL experts for individual tasks and then consolidating them via distillation, as an alternative to directly training a single model on mixed tasks. We show that off-policy distillation degrades in multi-task settings due to the mode-covering nature of forward KL: aggregating data from multiple tasks introduces a large number of behavioral modes that can exceed the student's capacity, forcing it to average across behaviors and leading to degraded performance. In contrast, on-policy distillation is mode-seeking but requires strong initialization. Inspired by these observations, we propose a two-phase approach: off-policy distillation followed by on-policy refinement. Evaluation across conversational agents and text-based games confirms that this two-phase approach matches single-task RL expert performance for each individual task, whereas off-policy or on-policy distillation alone fails to match this performance.