arxiv: 2604.05336 · v1 · submitted 2026-04-07 · 💻 cs.AI

Recognition: no theorem link

TRACE: Capability-Targeted Agentic Training

Hangoo Kang , Tarun Suresh , Jon Saad-Falcon , Azalia Mirhoseini

Authors on Pith no claims yet

Pith reviewed 2026-05-10 20:07 UTC · model grok-4.3

classification 💻 cs.AI

keywords agentic trainingcapability identificationsynthetic environmentsLoRA adapterstrajectory contrastreinforcement learningLLM agentsself-improvement

0 comments

The pith

TRACE improves LLM agents by identifying missing capabilities from failures and training targeted synthetic environments.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper presents TRACE as a method for LLM agents to improve on specific tasks by first spotting the exact capabilities they lack. It does so through automatic comparison of successful and failed attempts in an environment, then builds custom training setups that reward exercising those capabilities. Small LoRA adapters are trained on each setup with reinforcement learning and selected at inference time. This yields gains on customer service and tool-use benchmarks while using rollouts more efficiently than direct training or untargeted methods. A reader would care because it offers a concrete way to move beyond broad training and focus effort on the precise gaps that cause repeated failures.

Core claim

TRACE contrasts successful and failed trajectories to automatically identify lacking capabilities, synthesizes a targeted training environment for each that rewards whether the capability was exercised, and trains a LoRA adapter via RL on each synthetic environment, routing to the relevant adapter at inference. This produces +14.1 points on τ²-bench and +7 perfect scores on ToolSandbox while outperforming the strongest baseline.

What carries the argument

The pipeline that contrasts trajectories to generate capability-targeted synthetic environments and trains routed LoRA adapters on them.

If this is right

The approach generalizes across customer-service and tool-use environments.
It delivers larger gains than GRPO or GEPA given the same number of rollouts.
Separate LoRA adapters per capability allow modular routing at inference.
Performance scales more efficiently than training directly on the target distribution.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

The method could reduce reliance on hand-crafted or generic synthetic data when adapting agents to new domains.
Modular adapters might support combining capabilities across tasks without retraining a single large model.
If capability isolation proves reliable, similar contrast-based synthesis could apply to other sequential decision settings.
The routing step suggests a path toward agents that dynamically activate skill modules rather than relying on one monolithic policy.

Load-bearing premise

Automatically contrasting successful and failed trajectories correctly isolates the precise capabilities whose absence caused the failures, and training on the resulting synthetic environments transfers back to the original task distribution.

What would settle it

A controlled test in which TRACE-generated environments and adapters produce no measurable gain on tasks that require the identified capabilities, or where expert review shows the automatically flagged capabilities do not match the actual causes of failure.

Figures

Figures reproduced from arXiv: 2604.05336 by Azalia Mirhoseini, Hangoo Kang, Jon Saad-Falcon, Tarun Suresh.

**Figure 2.** Figure 2: Stability and coverage of identified capabilities. Repeated contrastive analysis [PITH_FULL_IMAGE:figures/full_fig_p009_2.png] view at source ↗

**Figure 3.** Figure 3: Scaling with the number of capabilities with TRACE and GEPA on [PITH_FULL_IMAGE:figures/full_fig_p010_3.png] view at source ↗

**Figure 4.** Figure 4: Pass rate on τ 2 -Bench and ToolSandBox with scaling the number of rollouts. exhibits instability, dropping to 35.4% at 3,840 rollouts, and ultimately stalls at 37.8%. A consistent trend appears on ToolSandBox (Figure 4b), where TRACE maintains a steady climb to a 0.552 mean similarity score, outperforming both GEPA (0.520) and GRPO (0.519). 5 Conclusion In this work, we propose TRACE, an end-to-end system… view at source ↗

read the original abstract

Large Language Models (LLMs) deployed in agentic environments must exercise multiple capabilities across different task instances, where a capability is performing one or more actions in a trajectory that are necessary for successfully solving a subset of tasks in the environment. Many existing approaches either rely on synthetic training data that is not targeted to the model's actual capability deficits in the target environment or train directly on the target environment, where the model needs to implicitly learn the capabilities across tasks. We introduce TRACE (Turning Recurrent Agent failures into Capability-targeted training Environments), an end-to-end system for environment-specific agent self-improvement. TRACE contrasts successful and failed trajectories to automatically identify lacking capabilities, synthesizes a targeted training environment for each that rewards whether the capability was exercised, and trains a LoRA adapter via RL on each synthetic environment, routing to the relevant adapter at inference. Empirically, TRACE generalizes across different environments, improving over the base agent by +14.1 points on $\tau^2$-bench (customer service) and +7 perfect scores on ToolSandbox (tool use), outperforming the strongest baseline by +7.4 points and +4 perfect scores, respectively. Given the same number of rollouts, TRACE scales more efficiently than baselines, outperforming GRPO and GEPA by +9.2 and +7.4 points on $\tau^2$-bench.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

TRACE turns trajectory contrasts into per-capability synthetic RL environments and routes LoRA adapters, which is a clean pipeline that delivers reported gains on two agent benchmarks, but the causal link from contrast to targeted capability remains unverified.

read the letter

The main point is that TRACE automates identification of missing capabilities by contrasting successful and failed trajectories, builds a separate synthetic environment for each that rewards exercising that capability, trains a LoRA adapter on it with RL, and routes the right adapter at inference. This produces +14.1 points on τ²-bench and +7 perfect scores on ToolSandbox while using the same number of rollouts more efficiently than GRPO or GEPA baselines.

Referee Report

2 major / 2 minor

Summary. The paper introduces TRACE, an end-to-end system for LLM agent self-improvement in specific environments. It automatically identifies capability deficits by contrasting successful and failed trajectories, synthesizes per-capability training environments that reward exercise of the identified capability, trains a LoRA adapter via RL on each such environment, and routes the relevant adapter at inference. The central empirical claim is that this yields generalization across environments, with gains of +14.1 points on τ²-bench (customer service) and +7 perfect scores on ToolSandbox (tool use) over the base agent, outperforming the strongest baseline by +7.4 points and +4 perfect scores respectively, while also scaling more efficiently than GRPO and GEPA under equal rollouts.

Significance. If the results and underlying assumptions hold after proper validation, TRACE would represent a meaningful advance in targeted agentic training. It moves beyond generic synthetic data or implicit learning on target tasks by attempting to isolate and directly train specific capabilities, with reported efficiency gains in rollout usage. The end-to-end automation and routing mechanism could have practical impact on deploying agents in varied environments if the contrast-based identification reliably surfaces causally necessary capabilities.

major comments (2)

[Method (capability identification and environment synthesis)] The core methodological claim—that contrasting successful and failed trajectories isolates the precise capabilities whose absence caused failures, enabling effective transfer from synthetic environments—lacks any supporting ablation, metric, or formal validation in the method description. This assumption is load-bearing for attributing the reported +14.1 point gain on τ²-bench and +7 perfect scores on ToolSandbox to the targeted training rather than generic RL or data effects.
[Experiments and Results] The experimental evaluation reports benchmark improvements and scaling advantages but supplies no information on controls (e.g., equivalent total compute or rollouts including synthetic environment generation), statistical significance, data exclusion criteria, or how synthetic environments were validated for transfer to the original task distribution. Without these, the central performance claims cannot be assessed.

minor comments (2)

[Abstract and Experiments] The notation τ²-bench is used without an explicit definition or citation in the abstract and results; a brief description of the benchmark and its scoring would improve clarity for readers unfamiliar with it.
[Overall presentation] The paper would benefit from a small table or diagram summarizing the end-to-end pipeline (trajectory contrast → capability extraction → synthetic env → LoRA training → routing) to make the flow easier to follow.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive and detailed review. The comments correctly identify areas where the manuscript would benefit from greater methodological validation and experimental transparency. We address each major comment below and have prepared revisions to incorporate additional ablations, controls, and clarifications.

read point-by-point responses

Referee: [Method (capability identification and environment synthesis)] The core methodological claim—that contrasting successful and failed trajectories isolates the precise capabilities whose absence caused failures, enabling effective transfer from synthetic environments—lacks any supporting ablation, metric, or formal validation in the method description. This assumption is load-bearing for attributing the reported +14.1 point gain on τ²-bench and +7 perfect scores on ToolSandbox to the targeted training rather than generic RL or data effects.

Authors: We agree that the manuscript does not contain a dedicated ablation study or quantitative metric (e.g., capability coverage or causal necessity score) that formally validates the contrast-based identification step. The current method description in Section 3 relies on the operational definition of capabilities as actions necessary for task subsets, with qualitative trajectory examples provided to illustrate the process. To address this, the revised manuscript will include an ablation that replaces the contrast-derived capabilities with randomly selected ones and with capabilities derived from direct failure analysis without contrast. We will also add a transfer metric measuring overlap between capabilities exercised in synthetic environments and those required on held-out original tasks. These additions will allow readers to assess whether the gains are attributable to targeted training rather than generic RL effects. revision: yes
Referee: [Experiments and Results] The experimental evaluation reports benchmark improvements and scaling advantages but supplies no information on controls (e.g., equivalent total compute or rollouts including synthetic environment generation), statistical significance, data exclusion criteria, or how synthetic environments were validated for transfer to the original task distribution. Without these, the central performance claims cannot be assessed.

Authors: We acknowledge that the original submission omitted several critical experimental details. In the revised version we have added: (1) a compute and rollout budget table showing that synthetic environment generation rollouts are included in the total and that all methods (TRACE, GRPO, GEPA) are compared under identical total rollout counts; (2) statistical significance via mean and standard error across five independent runs with different seeds; (3) explicit data exclusion criteria (trajectories with invalid tool calls or parsing failures are discarded, representing <3% of data); and (4) a transfer validation subsection that reports performance on a held-out subset of the original task distribution after training only on the synthesized environments. These changes directly support the reported gains of +14.1 points on τ²-bench and +7 perfect scores on ToolSandbox. revision: yes

Circularity Check

0 steps flagged

No circularity: empirical method with benchmark results, no derivations or fitted quantities

full rationale

The paper presents TRACE as an empirical pipeline that contrasts trajectories to identify capabilities, synthesizes environments, and trains LoRA adapters, with all claims supported by reported benchmark scores (+14.1 on τ²-bench, +7 perfect scores on ToolSandbox). No equations, parameter-fitting steps, self-citations, or uniqueness theorems appear in the text. The reported gains are external performance metrics, not quantities derived by construction from the method itself. The central assumption about capability isolation and transfer is an unverified empirical claim rather than a self-referential definition or reduction.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract supplies no equations, derivations, or parameter lists; therefore the ledger is empty.

pith-pipeline@v0.9.0 · 5547 in / 1140 out tokens · 86663 ms · 2026-05-10T20:07:43.925347+00:00 · methodology

discussion (0)

Reference graph

Works this paper leans on

9 extracted references · 4 canonical work pages · 2 internal anchors

[2]

SWE-bench: Can Language Models Resolve Real-World GitHub Issues?

URLhttps://arxiv.org/abs/2310.06770. 11 Preprint. Under review. Xiaoqian Liu, Ke Wang, Yuchuan Wu, Fei Huang, Yongbin Li, Junge Zhang, and Jianbin Jiao. Agentic reinforcement learning with implicit step rewards, 2025. URL https: //arxiv.org/abs/2509.19199. Jiarui Lu, Thomas Holleis, Yizhe Zhang, Bernhard Aumayer, Feng Nan, Felix Bai, Shuang Ma, Shen Ma, M...

work page internal anchor Pith review Pith/arXiv arXiv 2025
[3]

https://thinkingmachines.ai/blog/ on-policy-distillation/

URLhttps://arxiv.org/abs/2408.04682. Kevin Lu and Thinking Machines Lab. On-policy distillation.Thinking Machines Lab: Connectionism, 2025. doi: 10.64434/tml.20251026. https://thinkingmachines.ai/blog/on- policy-distillation. Tongxu Luo, Jiahe Lei, Fangyu Lei, Weihao Liu, Shizhu He, Jun Zhao, and Kang Liu. Moelora: Contrastive learning guided mixture of e...

work page doi:10.64434/tml.20251026 2025
[4]

Terminal-Bench: Benchmarking Agents on Hard, Realistic Tasks in Command Line Interfaces

URLhttps://arxiv.org/abs/2601.11868. Aniello Panariello, Daniel Marczak, Simone Magistri, Angelo Porrello, Bartlomiej Twar- dowski, Andrew D. Bagdanov, Simone Calderara, and Joost van de Weijer. Accurate and efficient low-rank model merging in core space. InAdvances in Neural Information Processing Systems, 2025a. Aniello Panariello, Daniel Marczak, Simon...

work page internal anchor Pith review arXiv 2017
[5]

Agentrl: Scaling agentic reinforcement learning with a multi-turn, multi-task framework,

URLhttps://arxiv.org/abs/2510.04206. Shuyan Zhou, Frank F. Xu, Hao Zhu, Xuhui Zhou, Robert Lo, Abishek Sridhar, Xianyi Cheng, Tianyue Ou, Yonatan Bisk, Daniel Fried, Uri Alon, and Graham Neubig. Webarena: A realistic web environment for building autonomous agents, 2024. URL https://arxiv. org/abs/2307.13854. 13 Preprint. Under review. A GRPO Details Given...

work page arXiv 2024
[6]

If all rollouts in a group get the same reward, the gradient is zero and the iteration is wasted

Reward Design for GRPO GRPO learns from WITHIN-GROUP reward variance. If all rollouts in a group get the same reward, the gradient is zero and the iteration is wasted. Your reward function MUST create variance: • DO: Use continuous/multi-level rewards (0.0, 0.2, 0.5, 0.8, 1.0) so that partial success is distinguishable from total failure • DO: Design scen...
[7]

Under review

Tool and Format Fidelity If the environment targets a specific benchmark (e.g., tau2-bench, ToolSandbox): 15 Preprint. Under review. • Use the exact same tool schemas (function names, parameter names, return types) • Error messages and response formats should be consistent within the envi- ronment • The system prompt should set up the skill context clearl...
[8]

Procedural Generation All scenarios MUST be procedurally generated from a seed: •reset(seed)must produce a deterministic scenario for that seed • Different seeds must produce meaningfully different scenarios (not just pa- rameter swaps) • The scenario space should be large enough that memorization is impossible • Include randomized: database contents, use...
[9]

Skill Isolation Each environment should primarily stress ONE skill. If you identify multiple skills: • Create separate scenario generators per skill • Assign probability weights (e.g., 40% skill_a, 35% skill_b, 25% skill_c) • Reward functions can differ per skill, but the interface is the same
[10]

Mean: {np.mean(rewards):.3f}

Multi-Skill Environments When combining multiple skills in one environment (like toolsandbox_multiturn): • Sample skill per scenario at reset time based on probability weights • Each skill gets its own scenario generators and reward logic • Track which skill each scenario targets (for analysis) • Aim for 1-2 skills max per environment — beyond that, creat...

2026