Small agents learn tasks faster using explicit task progress
Learning to Act with Task Progress: Distilling Small Agents from Compact Teacher Supervision
Computation and Language
Summary
Teaching small AI agents to perform tasks is hard because they often rely on big models that are slow to use every step. The authors show that by giving small agents short labels about what stage of the task they are in, instead of just actions, the agents learn faster and better. Their method, called Task-Progress Distillation, helps students catch on quickly, especially when there aren’t many examples to learn from. This approach also seems to help agents decide when to switch between important parts of a task.
What this means in practice
- •For robotic system developers: Train lightweight robot controllers to complete recurring tasks efficiently without relying on heavy computation at runtime.
- •For virtual assistant engineers: Improve virtual assistants’ task performance by incorporating explicit task stage labels to help them manage multi-step interactions.
Tested on one dataset.
Authors
Wenxi Gan
Abstract
Learning from large-model demonstrations offers a way to train small agents that can complete recurring tasks without calling a large model at every step. A central design choice is what to retain from teacher trajectories that contain reasoning, actions, and information about task progress. We introduce Task-Progress Distillation (TPD), an offline approach that pairs each demonstrated action with a short label describing the current task stage. The student learns these compact targets and selects actions by jointly scoring admissible stage--action pairs, which a deterministic harness executes in the environment. On ALFWorld, a 1.7B student trained with 404 demonstrations achieves 72.4\% mean unseen task success with either TPD or action-only supervision, compared with 48.3\% for a reasoning-trained student using constrained action selection. Explicit stages provide an additional benefit at 200 demonstrations, improving success from 48.0\% to 67.7\% over action-only supervision. With more demonstrations, the action-only student closes the gap, and both approaches reach 76.9\% at 808 demonstrations. Shared-history analyses link part of TPD's local advantage to better decisions when moving between subgoals, particularly from object acquisition to processing. These results show that compact supervision can train effective small task agents, while explicit task progress provides additional guidance at an intermediate demonstration budget.