TRACE: Capability-Targeted Agentic Training

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

Keywords: Agents, Reinforcement Learning

Abstract: Large Language Models (LLMs) deployed in agentic environments must exercise 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, composing them at inference either by routing to a single relevant adapter or by training a mixture-of-experts (MoE) model over the capability adapters. Empirically, TRACE generalizes across different environments, improving over the base agent by +14.1 points on τ 2-Bench (customer service) and +7 perfect scores on Tool-SandBox (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 τ 2-Bench.

Track: Regular Paper (9 pages)

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 323