Learning to Act with Task Progress: Distilling Small Agents from Compact Teacher Supervision

Wenxi Gan

cs.CLOct 7, 2026

Learning to Act with Task Progress: Distilling Small Agents from Compact Teacher Supervision

Authors: Wenxi Gan

Organizations: City University of Hong Kong Shenzhen Loop Area Institute

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.

Figures & tables

Figure 1 : Offline training and closed-loop execution. (a) A frozen annotator reads each teacher-history prefix and pairs the current task stage with the demonstrated action for student SFT. (b) The student scores $3|\mathcal{A}(h_{t})|$ stage–action candidates, and the harness executes the action from the highest-scoring pair. The next input includes the action and environment feedback; the predicted stage is logged only for analysis.

Arm	Supervised target	Execution interface
A	Teacher reasoning and action	Free generation followed by action parsing
$A_{\mathrm{constr}}$	Same checkpoint as A	Generate reasoning, then score $N$ admissible actions
B / TPD	Prefix-aligned stage and action	Joint scoring over $3N$ stage–action pairs
C	Action only	Score $N$ action candidates
$B_{\mathrm{const}}$	Constant phase field and action	Score $N$ fixed-stage action candidates
$B_{\mathrm{shuffle}}$	Shuffled stage labels and unchanged actions	Same $3N$ candidate structure as B

Table 1: Training targets and execution interfaces.

N

denotes the current admissible action count. A and

A_{\mathrm{constr}}

share weights at the corresponding budget and seed; other trained arms are initialized independently from the same Qwen3-1.7B checkpoint.

	Number of training demonstrations
Configuration	98	200	404	808
A: reasoning–action	$0.0\pm 0.0$	$0.0\pm 0.0$	$5.0\pm 1.7$	$16.4\pm 1.5$
$A_{\mathrm{constr}}$	$1.5\pm 0.8$	$30.8\pm 6.5$	$48.3\pm 4.8$	$48.8\pm 5.6$
C: action only	$42.5\pm 2.0$	$48.0\pm 12.1$	$72.4\pm 1.5$	$76.9\pm 3.0$
$B_{\mathrm{const}}$	$31.6\pm 4.4$	$51.5\pm 11.7$	$\mathbf{72.6}\pm 1.6$	$\mathbf{77.1}\pm 6.8$
B: stage–action ( TPD )	$\mathbf{48.0}\pm 2.8$	$\mathbf{67.7}\pm 2.4$	$72.4\pm 4.5$	$76.9\pm 3.3$

Table 2: Unseen success (%), mean

\pm

sample standard deviation across three training seeds, each evaluated on 134 tasks. Blue and orange identify C and TPD; bold marks the highest observed mean in each column, including ties.

Panel of Figure 2 — Figure 2 : Compact students perform well, with the largest stage benefit at 200 demonstrations. (a) Constrained selection improves the reasoning-trained student; the action-only configuration achieves higher success still. (b) TPD leads C most clearly at 200 demonstrations; C catches up in mean unseen success at 404 and 808. Error bars show sample standard deviations across three training seeds, and faint points show individual B/C runs.

Arm	Success (%)	Steps / episode	Output tokens / episode
A	5.0	47.9	14,635
$A_{\mathrm{constr}}$	48.3	33.6	9,995
B / TPD	72.4	22.0	299
C	72.4	22.2	213

Table 3: Execution at 404 demonstrations on unseen tasks. Output totals are divided by 402 episodes; mean steps are averaged from the reported per-seed means. Token counts exclude the additional work of scoring candidates.

Demonstrations	Both succeed	B only	C only	Both fail
200	183	89	10	120
808	284	25	25	68

Table 4: Paired B/C unseen outcomes over 134 tasks evaluated with three training seeds (402 seed–task records per budget).

Arm	Seed 42	Seed 1234	Seed 2026	Mean
B / TPD	55.2	53.4	56.9	55.2
C	35.3	37.9	50.9	41.4
$B_{\mathrm{const}}$	31.9	49.1	50.0	43.7
$B_{\mathrm{shuffle}}$	31.0	42.2	39.7	37.6

Table 5: Stage-alignment controls at 200 demonstrations: success (%) on the 116 seen diagnostic tasks. Results are shown for three training seeds and their mean; bold marks the highest value in each column.

Task	Paired records	B processes successfully	C processes successfully
Heat	52	49/52 (94.2%)	28/52 (53.8%)
Clean	60	59/60 (98.3%)	51/60 (85.0%)
Cool	49	44/49 (89.8%)	45/49 (91.8%)

Table 6: Successful processing after correct acquisition, restricted to pairs in which both agents acquire the target (200 demonstrations, unseen). Task and training seed are paired; the acquisition-time environment states are not identical.

Criterion or coverage	Result
Both original policies choose progress	81/114
Only original B chooses progress	15/114
Only original C chooses progress	1/114
Neither original policy chooses progress	17/114
B-only cases with useful matching stage and sensitivity to a mismatch	11/15
Also require original B to select reference stage and matching action	11/15

Table 7: Local progress and stage sensitivity at 200 demonstrations: 38 seen prefixes evaluated with three training seeds. Overlap fractions use the 15 B-only progress cases as their denominator. Coverage and scoring checks show how many tasks support the overlap and whether it persists under both scoring rules.

Panel of Figure 3 — Figure 3 : Moving from object acquisition to processing. (a) On the same history after acquiring a cup, B chooses the microwave at all nine checkpoints. C chooses it at one of three 200-demonstration checkpoints and all six larger-budget checkpoints. (b) Across 38 rule-evaluable seen prefixes from 30 tasks, fixing the matching stage more often yields a progress action than fixing a mismatched stage. Curves show local progress rates averaged over three seeds at each budget. Table 7 relates these tests to the original B/C choices.

Panel of Figure 4 — Figure 4 : Reported success and demonstration counts for offline-supervised agents. External points show LM-BUTLER, ETO’s SFT stage, and the 800- and 1,600-demonstration SFT baselines reported in HSL Micheli and Fleuret (2021) ; Song et al. (2024) ; Li et al. (2026) . Our B/C results are three-seed means; their markers overlap at 404 and 808 demonstrations. Each work uses its own student model, demonstration source, and evaluation protocol. Appendix C.1 lists the plotted values and their sources.

Explore similar work

Agent Memory Distillation: Empowering Small LLM Agents with Hierarchical Teacher Memory
Aug 7, 2026Taeil Kim, Kangsan Kim, Sung Ju HwangLarge Language Model Agents Frontier Large Language Model Agents
KAIST
On-Policy Distillation with Curriculum Turn-level Guidance for Multi-turn Agents
Jun 14, 2026Gengsheng Li, Mao Zheng, Mingyang Song +8Teacher Curriculum
Foundation Model Research Center, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Large Language Model Department, Tencent +4
Look Ahead Before You Distill: Future Trajectory Validation of Teacher Guidance for Agentic On-Policy Distillation
Aug 3, 2026Chishui Chen, Yaoyou Fan, Te Sun +11Efficient On-Policy Distillation Teacher
Meituan LongCat Interaction · Peking University · Shanghai Jiao Tong University +6

Aug 7, 2026, cs.AIJ/K move · Enter open · S save

Agent Memory Distillation: Empowering Small LLM Agents with Hierarchical Teacher Memory

Taeil Kim, Kangsan Kim, Sung Ju Hwang

KAIST

Large Language Model Agents Frontier Large Language Model Agents Peak Memory Teacher Agentic Workflow Design

Open paper AskPDF arXiv

Memory systems have shown promise for improving agent performance, but their potential remains largely unexplored for small language models, which struggle to generate sufficient successful trajectories on their own. We propose Agent Memory Distillation (AMD), a training-free framework that transfers structured knowledge from a large teacher agent to a small student agent through hierarchical memory. AMD constructs three complementary memory types from successful teacher trajectories: Workflow memory encodes task-level strategies, Subtask memory provides concrete behavioral examples at an intermediate granularity, and Function memory captures per-function calling conventions and common pitfalls. Workflow and Subtask memories are injected proactively at the start of each task, while Function memory is retrieved reactively upon tool-calling errors. We evaluate AMD on three tool-use benchmarks using four student models (4B-8B parameters) with GPT-5-mini as the teacher, achieving average accuracy gains of 27.2%p, 11.2%p, and 3.4%p on AppWorld, BFCL V3, and ToolSandbox, while consistently outperforming existing memory-based baselines. Further analysis shows that Subtask memory contributes the largest gains, teacher effectiveness depends on both teacher capability and student compatibility, and 4B-sized students benefit most from AMD.

Abstract

Explore similar work

Agent Memory Distillation: Empowering Small LLM Agents with Hierarchical Teacher Memory

On-Policy Distillation with Curriculum Turn-level Guidance for Multi-turn Agents

Look Ahead Before You Distill: Future Trajectory Validation of Teacher Guidance for Agentic On-Policy Distillation