PaperlayerView exact revision on arXiv

Source-linked AI summary

TRACE: Capability-Targeted Agentic Training

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

arXiv:2604.05336v2cs.AI

TL;DR

Agent failures often reflect missing capabilities that recur across tasks, while target-environment optimization provides limited failure attribution. TRACE identifies these capabilities from successful and failed trajectories, trains capability-specific adapters in targeted environments, and combines them with a mixture-of-experts model, improving performance across customer-service and software-engineering benchmarks.

  • Problem

    Agent failures can reflect missing capabilities required across many tasks, while optimizing directly in the target environment provides limited failure attribution.

  • Method

    TRACE contrasts successful and failed trajectories, synthesizes capability-targeted environments, trains LoRA adapters via reinforcement learning, and composes them with a Mixture-of-Experts model.

  • Results

    TRACE improves over the base agent by +15.3 points on τ2-Bench and +15 points Pass@1 on SWE-bench Verified, while outperforming strongest external baselines by +8.6 points and +8.4 points, respectively.

  • Takeaways & Limitations

    TRACE can be effectively applied across different agentic environments and scales more efficiently than the evaluated baselines at the same number of rollouts.

  • Takeaways & Limitations

    TRACE assumes each generated capability-training task requires exercising that capability and that the synthetic environment preserves relevant target-environment aspects, including tools, state representation, and policy constraints.

Abstract

from arXiv · show

Models often fail to complete agentic tasks because they lack core capabilities required by the target environment. However, mainstream approaches for addressing these failures typically either fine-tune directly on target environments or generate synthetic data that is not targeted to the model's actual capability deficits, resulting in low sample efficiency and limited generalization. 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 missing capabilities, synthesizes a targeted training environment for each capability that rewards whether the capability is exercised, trains a LoRA adapter via reinforcement learning on each synthetic environment, and then trains a mixture-of-experts model over the capability adapters. TRACE can be effectively applied across different environments, improving over the base agent by +15.3 points on $τ^2$-Bench, a customer-service agent benchmark, and by +15.0 points Pass@1 on SWE-Bench Verified, a software-engineering benchmark. TRACE outperforms the strongest external baselines, GEPA and SWE-RL, by +8.6 points and +8.4 points, respectively. In addition, TRACE is more sample-efficient than strong fine-tuning baselines: using fewer than one-fourth the number of rollouts, TRACE outperforms the best-performing baselines, GRPO and GEPA, and achieves higher final accuracy by +10.4 and +8.6 points on $τ^2$-Bench.

1 Introduction

TRACE addresses recurrent agent failures as missing capabilities rather than isolated task errors, then targets those deficits with synthesized environments and capability-specific training. The system is designed to identify deficits, isolate them in training, and compose the resulting capabilities for target tasks.

  • Motivation: Agent failures often reflect missing capabilities needed across multiple task instances, not isolated errors.A capability is one or more trajectory actions necessary for successfully solving some subset of tasks.
  • Motivation: Direct target-environment training provides task-level feedback without explicitly identifying missing capabilities, limiting sample efficiency and generalization.The model must infer shared capabilities from sparse, heterogeneous outcomes and trajectories.
  • Motivation: Untargeted synthetic data can train behaviors the agent already performs or capabilities irrelevant to target failures, reducing environment-specific improvement.Scaling data quantity and diversity does not ensure alignment with the model’s actual capability deficits.
  • TRACE: TRACE contrasts successful and failed trajectories to identify missing capabilities, synthesizes environments that isolate and reward each capability, and trains and composes capability adapters.The system uses reinforcement learning for LoRA adapters and a Mixture-of-Experts model to combine them.
  • TRACE: TRACE is evaluated as an end-to-end system for environment-specific agent self-improvement across different agentic environments.Its stated objectives are identifying high-impact deficits, synthesizing capability-targeted environments, and learning and composing multiple capabilities.

2 Related Work

Prior work studies agents in complex interactive environments, scales reinforcement learning and synthetic data, and composes task-specific adaptations with LoRA mixtures. These approaches provide the surrounding benchmarks, training strategies, and model-composition techniques for agent capability acquisition.

  • LLM Agents and Interactive Environments: Agent benchmarks evaluate autonomous models on complex, multi-turn environments requiring tool interfaces and interaction protocols.Examples include τ2-Bench, SWE-bench, WebArena, WorkArena, Terminal-Bench, and TheAgentCompany.
  • Agentic Reinforcement Learning and Synthetic Data: Recent agent-training methods use device-control data, implicit step rewards, verifiable multi-turn reinforcement learning, and synthesized environments.Synthetic-environment approaches include AWM, EnvScaler, and ScaleEnv, while ADP uses unified public trajectories.
  • LoRA Merging and Routing: Model-merging research composes task-specific adaptations, while LoRA-MoE methods use mixtures of experts with dynamic or soft routing.This line of work avoids collapsing all adaptations into a single checkpoint.

3 Method

TRACE formulates agent improvement around identifying missing capabilities in a target environment, generating targeted synthetic environments, training capability adapters, and composing them with token-level routing.

  • Problem formulation: An agentic environment defines task instances, interaction dynamics, trajectory-level rewards, and binary success labels.
  • Problem formulation: A capability is an action or set of actions necessary for successfully solving some subset of task instances.Examples include retrieving the correct customer record or locating the relevant code function.
  • Synthetic environments: For each retained deficit, TRACE generates a synthetic environment where exercising the capability is necessary, relevant target-environment dynamics are preserved, and rewards are automatically computable.This decomposition makes the capability’s training signal denser and more attributable than the target environment’s general reward.
  • Capability identification: TRACE identifies missing capabilities by comparing how often they are lacking in failed versus successful trajectories and retaining those with substantial failure coverage.The contrastive criterion filters out uniformly unsuccessful or weakly discriminative capabilities.
  • Capability acquisition: TRACE trains each capability adapter with grouped-rollout reinforcement learning, normalizing trajectory rewards and discarding groups with identical rewards.The resulting updates use the standard clipped surrogate objective.
  • Capability composition: TRACE composes capability-specific LoRA adapters into a token-routed mixture-of-experts model whose gates select the capability used at each token.The base policy and adapters remain frozen while the gates are trained with cross-entropy and load balancing.

4 Experiments

TRACE is evaluated across customer-service and software-engineering benchmarks, with comparisons spanning agent backbones, training strategies, capability targeting, and rollout scaling. It consistently improves over baselines, while capability selection is stable and targeted training scales more reliably than prompt optimization or direct target-environment RL.

  • Experimental Setup: TRACE is evaluated on τ2-Bench and SWE-bench Verified, using pass rate and Pass@1 respectively, with two agent backbones.τ2-Bench covers AIRLINE and RETAIL customer-service domains, while SWE-bench Verified contains 500 software-engineering instances.
  • Effectiveness of Method: TRACE improves over the base model by +15.3 points on τ2-Bench and +15 points Pass@1 on SWE-bench Verified.The method consistently outperforms baselines across both evaluated benchmarks.
  • Effectiveness of Method: 59.1% overall pass rate is achieved by TRACE with the Qwen3.6-27B backbone on τ2-Bench, exceeding the base model by +9.1 points and the strongest baseline by +5.4 points.The same stronger-backbone variant reaches 73.2% Pass@1 on SWE-bench Verified and outperforms much larger frontier models on the public leaderboard.
  • Targeted vs. General-Purpose Training Environments: A single targeted capability adapter reaches 40.3% on τ2-Bench and 36.6% Pass@1 on SWE-bench Verified, surpassing general-purpose AWM and ADP baselines.The reported comparison supports targeted synthetic environments over broad general-purpose training data.
  • Capability Training vs. Prompt Optimization: Explicit capability training outperforms GEPA prompt optimization on both τ2-Bench overall pass rate and SWE-bench Verified Pass@1.GEPA improves over the base model but underperforms both single-capability training and TRACE with all capabilities.
  • Capability Analysis: Structured data reasoning, multi-step task completion, and precondition verification are identified in all 10 analysis runs, while tool-calling precision appears in 8 of 10.Failure coverage is concentrated in these capabilities, with structured data reasoning accounting for the largest share.
  • Method Scaling: TRACE’s τ2-Bench pass rate rises from 32.9% to 48.2% at 5,120 rollouts, while GEPA plateaus at 39.6% and GRPO ultimately stalls at 37.8%.TRACE improves monotonically with rollouts; GRPO is unstable and drops to 35.4% at 3,840 rollouts.

5 Conclusion

TRACE identifies missing capabilities from successful and failed trajectories, trains capability-specific adapters in targeted synthetic environments, and composes them with a Mixture-of-Experts model. Its training uses on-policy GRPO with group-relative reward normalization and clipped updates.

  • TRACE trains a separate low-rank adapter for each identified capability while keeping the base model frozen.
  • GRPO normalizes rewards within groups, discarding groups with identical rewards because they provide no learning signal.
  • Clipped GRPO updates limit excessively large policy changes when estimated advantages are high.
  • On-policy rollouts keep the importance ratio near 1 at update starts, improving optimization stability.
  • TRACE reaches a 0.560 mean similarity on ToolSandBox, above GEPA at 0.520 and GRPO at 0.519.

B ToolSandbox Evaluation

TRACE is evaluated on ToolSandBox, a stateful conversational tool-use benchmark, using mean similarity and perfect rate under matched budgets and greedy decoding. It improves over the base model and external baselines, including with a single capability adapter.

  • TRACE evaluates permission error recovery and datetime reasoning on 129 ToolSandBox scenarios.
  • 0.560 mean similarity and 26 perfect scores improve over the base model by +0.149 mean similarity and 7 perfect scores.
  • 0.040 mean similarity and 4 perfect scores separate TRACE from the strongest external baseline.
  • 0.514 mean similarity from one synthesized capability environment surpasses AWM at 0.504 and ADP at 0.422.

C Capability Consolidation

TRACE consolidates capability-specific training through a Mixture-of-Experts model, while its synthetic environments emphasize targeted, procedurally generated skill practice and reward variance for GRPO learning.

  • Capability Consolidation: TRACE MoE consolidates independently trained capability adapters into one model, outperforming alternative consolidation approaches.
  • Capability Consolidation: 40.9% from Multi-Capability GRPO only marginally exceeds the best single adapter at 40.3%, whereas MoE significantly outperforms it.
  • Parameterization: The backbone remains frozen, while each capability-specific LoRA adapter contributes approximately 1.6 billion trainable parameters.
  • Parameterization: At inference, the MoE gate selects top-k capabilities, and k=1 keeps only one adapter active per token.
  • Environment Design: GRPO requires within-group reward variance because identical rewards produce zero gradient and waste an iteration.
  • Environment Design: Continuous or multi-level rewards distinguish partial success from total failure.
  • Environment Design: Procedural generation varies database contents, user requests, constraints, and distractors across deterministic seeds.
  • Skill Isolation: Targeted environments should primarily stress one skill, with separate generators and reward logic when multiple skills are included.

F Metrics Detail

The paper defines benchmark metrics for τ2-Bench, SWE-bench Verified, and ToolSandBox, including task-level pass measures, patch success, trajectory similarity, and perfect-score frequency.

  • τ2-Bench overall pass rate divides solved Airline and Retail tasks by their corresponding benchmark counts.
  • SWE-bench Verified Pass@1 counts patches that apply cleanly and pass all held-out tests across 500 instances.
  • ToolSandBox mean similarity averages final trajectory similarity scores, while perfect rate counts scenarios scoring exactly 1.0.
  • Stochastic-decoding evaluations report mean values with 95% confidence intervals across three seeds.
  • Cross-environment transfer on ToolSandBox compares agents trained on τ2-Bench without retraining against in-domain ToolSandBox training.

G Results with Stochastic Decoding

Under stochastic decoding, TRACE consistently outperforms all baselines on both benchmarks, with non-overlapping confidence intervals. Its targeted training also transfers better than GRPO to ToolSandBox without retraining.

  • TRACE consistently outperforms all baselines on both benchmarks under temperature 0.7 decoding with three seeds.Results are reported as means with 95% confidence intervals.
  • +1 perfect score and +0.052 mean similarity distinguish TRACE from GRPO on ToolSandBox transfer.The model is trained on τ2-Bench and evaluated on ToolSandBox without retraining.
  • TRACE-trained agents transfer more effectively than GRPO-trained agents to a second tool-calling benchmark.The reported transfer result is presented as evidence that targeted training carries capabilities across environments.

I Extended Capability Analysis

TRACE’s analysis consistently selects a small set of high-frequency capabilities while discarding many competing categories. Failure coverage can overlap across capabilities, but contrastive analysis separates these overlapping deficits for targeted environment synthesis.

  • The analysis agent consistently selects the top five capabilities across ten independent runs.Competing categories such as conditional reasoning, numerical reasoning, early termination, and information communication appear much less frequently.
  • The long-tail distribution indicates that the target benchmark’s failure modes are concentrated rather than uniform.This pattern is reported across the ten independent runs.
  • A single failed trajectory can involve multiple missing capabilities, so coverage counts are not mutually exclusive.The analysis disentangles overlapping failures and enables isolated micro-environments for distinct deficits.

J.1 Training Capabilities and Trajectory Examples

Contrastive analysis identifies four primary capability gaps on τ2-Bench and two on ToolSandbox. The examples show distinct failure modes involving structured data, tool arguments, multi-step completion, precondition checks, permission recovery, and datetime reasoning.

  • τ2-Bench: Four primary capability gaps are identified on τ2-Bench, each targeting a distinct failure mode.They are structured data reasoning, tool calling precision, multi-step task completion, and precondition verification.
  • τ2-Bench: Structured data reasoning failures involve misreading nested tool records and computing incorrect booking totals.The example concerns selecting the wrong economy fare from nested cabin-class price arrays.
  • τ2-Bench: Tool calling precision failures involve selecting the correct tool but passing an incorrect argument.One example passes a gift card instead of the original credit card for a refund.
  • τ2-Bench: Multi-step task completion failures occur when the agent stops after completing only the first operation in a compound request.The example cancels one reservation but never cancels the second or modifies the third.
  • τ2-Bench: Precondition verification failures occur when the agent performs state-changing actions without checking policy eligibility.The cancellation example requires independently checking eligibility because the API does not enforce the policy.
  • ToolSandbox: Two primary capability gaps are identified on ToolSandbox: permission error recovery and datetime reasoning.The examples contrast halting after a permission error or mentally decoding timestamps with diagnosing blockers and using timestamp-conversion tools.

K.1 Structured Data Reasoning (τ2-Bench)

The structured-data and tool-use environments generate scenarios from seeded domain databases, execute agent actions, and evaluate both state changes and communication. Their rewards distinguish successful action and reporting, while ToolSandbox includes tool-then-communicate, error-recovery, and combined skills.

  • Scenario generation: The structured-data game generates seeded airline or retail scenarios with domain-specific databases and empirically weighted task types.Airline scenarios use flight and reservation data, while retail scenarios use product and order data.
  • Environment loop: The multi-turn loop executes tool calls against the environment database and returns tool outputs to the agent.Responding to the user can terminate the game when the stop marker appears.
  • Reward evaluation: Structured-data rewards are 1.0 for correct database state and communication, 0.3 for correct state with incorrect reporting, and 0.0 otherwise.Information-retrieval tasks receive 1.0 when communication passes and 0.0 otherwise.
Loading 2604.05336v2…