Harness-Aware Distillation for Small Language Model Agents

AI summary

Overview

Research area: Knowledge distillation for language model agents, specifically on-policy distillation of small language models that operate inside a software "harness."

Technical level: Intermediate. The paper assumes familiarity with on-policy distillation, preference learning (logistic/DPO-style losses), and multi-turn agent benchmarks, but the core argument is stated plainly.

Scope: The paper proposes Harness-Aware Distillation (HAD), a method that trains small agents to act on information their harness supplies, and evaluates it against four on-policy distillation baselines on three long-horizon agent benchmarks and two model families.

What This Paper Is About

Language model agents are deployed with a "harness" — the surrounding software that tracks environment state, lists available actions, and reports failures or stalls. When such an agent is distilled into a smaller model, the harness stays in place, so the student only needs to learn what the teacher adds beyond the harness, such as acting correctly on the harness's reports. Standard on-policy distillation (OPD) treats the harness as just another part of the input and imitates the teacher's full outputs, which the authors argue fails to transfer this ability; HAD instead supervises exactly the decisions where the harness changes what the teacher does.

Key Contributions

1. A reframing of agent distillation as distillation between two systems that share a harness. The paper argues that because the harness is fixed across teacher and student, the student should learn teacher-specific abilities the harness cannot supply, and shows empirically that plain on-policy distillation with the shared harness does not reliably transfer this.

2. Harness-Aware Distillation (HAD), combining two components. A harness-aware action preference that queries the same teacher with and without the harness information and trains the student to prefer the harness-conditioned action, scored after the student's own reasoning; and a harness-aware action validity filter that drops preference pairs whose preferred action contradicts the harness records. Response distillation is left unchanged. HAD uses no task rewards, success labels, or future information.

3. A formal characterization of what the contrast adds. The paper shows that the isolated unfiltered preference loss has optimal unrestricted scores of the form log(P+(a)/P−(a)) plus an arbitrary constant, meaning the term emphasizes actions that are likely with the harness relative to without it — something response distillation, whose risk depends only on the positive response law, cannot identify.

4. Empirical validation across benchmarks and model families. HAD outperforms on-policy distillation baselines with the same fixed harness on ALFWorld, WebShop, and ScienceWorld, and on both the Qwen3 and Gemma-4 families, with ablations isolating each design choice.

Main Findings

• Best performance in every evaluated environment. On ALFWorld (1.7B student from 8B teacher) HAD reaches 63.43% unseen and 51.43% seen success rate with the harness, versus 47.01% and 41.43% for the best baseline, and it exceeds its 8B teacher (51.49% unseen, 50.00% seen). On WebShop (0.6B from 30B-A3B) HAD scores 47.19 with the harness and 19.60% success rate, against 41.59 and 19.00 for the best baseline. On ScienceWorld HAD reaches 9.40% success rate and −4.17 score, against 6.71% and −8.45 for the best baseline.

• The harness matters more for HAD than for baselines. Removing the harness at evaluation lowers HAD from 63.43% to 40.30% on unseen ALFWorld tasks, the largest drop among all methods, yet HAD without the harness still outperforms the majority of baselines and is on par with or above the teacher without the harness.

• Adding the harness to ordinary distillation helps harness utilization but not performance. The paper reports harness utilization rising from 65.7% to 73.1% while performance moves only from 43.1% to 43.5%, whereas HAD improves both (65.7% to 81.0%; 43.1% to 57.4%).

• HAD both avoids and escapes stalls. On ALFWorld, HAD achieves 75.9% prevention (episodes without stalls) and 59.7% recovery (stalls escaped), against at most 64.6% prevention for any other model (teacher) and 46.8% recovery for the untrained student. No baseline clearly improves recovery over the untrained student, while HAD escapes 59.7% of its stalls. Every counted stall is flagged by the harness.

• Selective internalization of harness components. Hiding state tracking drops HAD from 63.43% to 37.31% on unseen ALFWorld tasks and from 51.43% to 40.71% on seen tasks, while hiding behavior monitoring drops it only to 55.97% and 45.00%. Blocking the actions the harness rejects lowers success to 55.22% and 46.43%, which the authors read as evidence that the harness helps HAD more as information to act on than as a hard constraint.

• Cross-family generalization. Distilling Gemma-4-E2B from Gemma-4-12B on ALFWorld, HAD attains 37.31% unseen and 36.43% seen success rate, versus 32.09% and 34.29% for the strongest baseline, and its harness utilization exceeds even the teacher's.

• Every component contributes, and the gain is not explained by extra teacher queries. Ablations on ALFWorld show HAD at 57.4% SR and 81.0% HU; removing the validity check gives 50.8%/78.9%, removing the negative (positive-only imitation) gives 52.4%/74.4%, and removing the preference term entirely gives 43.5%/73.1%. Giving OPD and SOPD twice the teacher budget raises the average of seen and unseen success only from 43.5% to 46.7% and 46.8% respectively, far below HAD's 57.4%.

• Design choices validated by ablation. Scoring only the action after the student's own reasoning beats full-response preferences (which let the loss drop by distinguishing teacher reasoning styles, with long reasoning diluting the short action signal), and using the teacher as a reference model does not help. The success rate peaks at ρ = 0.5.

Methodology in Plain English

The authors start from a simple observation: teacher and student run inside the same harness, and the harness is unchanged when the teacher is swapped for a smaller model. So the student does not need to relearn what the harness already provides — it needs the teacher's judgment about what to do with that information.

Their setup has the student act through the harness, collecting its own rollout trajectories rather than training on teacher trajectories. At each visited state, they query the same teacher twice: once with the harness information present, and once with that information stripped from the history (the stripped version is not a separately generated trajectory — only the harness text is removed, not its influence on what already happened). The teacher's response with the harness is then imitated token by token, exactly as in standard response distillation.

On top of that, the student gets a pairwise preference signal: the action the teacher chose with the harness should be preferred over the action it chose without it. Both are scored under the student's own reasoning trace, and only the action tokens count, with log-likelihoods averaged per token. The loss is a reference-free logistic loss with margin scale β.

Because a teacher can still be wrong, a validity checker consults only the harness records at that state — for example, an action outside the listed available actions, an action using an object the agent does not hold, or a repeated action the harness already reported as ineffective. Pairs whose preferred action fails the check are dropped from the preference term only; the teacher response still enters the distillation term. Negative actions are not checked, since lowering an invalid action's likelihood does no harm.

The two loss terms are balanced automatically per minibatch: λ is set so the preference gradient's norm is ρ = 0.5 times the distillation gradient's norm.

Experiments use ALFWorld, WebShop, and ScienceWorld, with a hand-designed, task-agnostic harness shared across environments whose components are state tracking (available actions, location and inventory, visit history, step counter) and behavior monitoring (failure, stall). Baselines are OPD, SAGE-OPD, Guided-OPD, and SOPD, all trained and evaluated with the same harness, same training loop, task pool, and number of updates. Metrics are success rate, WebShop score and ScienceWorld score, average turns in successful episodes, and harness utilization. Training uses full fine-tuning with AdamW, learning rate 10⁻⁵, effective batch size 60, 250 updates on ALFWorld, and 150 on WebShop and ScienceWorld, with evaluation at temperature 0.4 and thinking disabled.

Why This Matters

Impact on research. The paper argues a general principle: when distilling into a system that contains a fixed component, supervision should target what that component cannot provide. It also gives an alternative to reward-based or success-label-based agent training, and it reframes harness utilization as both a design target and a diagnostic metric. The analysis connecting the preference term to the log-ratio of harness-conditioned to stripped action marginals gives a concrete account of why contrastive supervision carries information imitation cannot.

Real-world applications.

• Deploying small, cheap agents with long-horizon tool use in environments where a fixed orchestration layer already tracks state and errors.
• Web navigation and shopping assistants where a wrapper supplies site structure, available links, and budget limits.
• Embodied or simulated task agents (household, scientific procedure) that must recover from repeated ineffective actions.
• Any retrieval- or tool-augmented agent, which the authors suggest may follow the same principle of focusing distillation on what the fixed component cannot supply.

Industry relevance. The headline motivation is cost: large models are expensive to serve, and on-policy distillation is a common path to cheaper agents. HAD targets the deployment configuration practitioners already use — model plus harness — and requires no task rewards, no success labels, and no future information, which matters when reward signals are hard to specify. The result that HAD exceeds its 8B teacher on ALFWorld makes the cost argument concrete.

Future Directions

• Scaling beyond the evaluated settings: students up to 2B parameters and teachers up to 30B parameters in text-based environments; the authors explicitly leave larger models and tool-calling or software agents to future work.
• Understanding how HAD interacts with harnesses that are learned or automatically optimized, rather than the hand-designed one used here.
• Extending the "focus on what the fixed component cannot provide" principle to agents that rely on retrieval or external tools.
• Improving the validity checker, which is limited by construction: it can only remove actions that contradict harness records and cannot detect actions that are valid yet fail to advance the task.

Target Audience

Researchers and engineers working on knowledge distillation, agent training, or small-model deployment. It is most useful to readers who already understand on-policy distillation and preference-based objectives and want a method for the specific case where a fixed orchestration layer surrounds both teacher and student. Practitioners building cost-sensitive agent products will find the benchmark comparisons and the harness-utilization framing directly applicable.

Authors’ abstract

Language model agents are deployed with a harness, the software around the model that manages its context, tools, and feedback. When such an agent is distilled into a smaller one, the harness stays in place, so the student mainly needs the teacher-specific abilities that the harness cannot provide, such as acting correctly on harness information. Standard distillation, however, imitates the teacher's full outputs and treats the harness as part of the input. We propose Harness-Aware Distillation (HAD), which focuses distillation on what the teacher adds beyond the harness. HAD complements on-policy distillation with two components: an action preference that contrasts the same teacher's actions with and without the harness information, scored after the student's own reasoning, and a validity check that drops preference pairs whose preferred action contradicts the harness records. We show that the contrast gives the student information that imitating the teacher alone cannot provide, and HAD needs no task rewards, success labels, or future information. Across multiple long-horizon agent benchmarks and models, HAD outperforms on-policy distillation baselines with the same fixed harness. Our analysis shows that HAD enters fewer unproductive loops and recovers from errors more often than the baselines, and suggests that it adaptively keeps learnable feedback in its weights while reading state information from the harness.