TRACE: Capability-Targeted Agentic Training

Chapter 1: The Key Insight

Here is the core observation that makes TRACE work: not all failures are created equal. If you look at 100 failure trajectories, you will find that the same small set of capability deficits explains the vast majority of them.

The authors ran TRACE's capability identification step 10 independent times on the same dataset (Figure 2a in the paper). Each run uses a different random seed, different trajectory samples, different LLM prompting. And yet, across all 10 runs, the same 4 capabilities emerged every time:

1. Structured data reasoning — correctly extracting, comparing, and manipulating data from database lookups and API responses
2. Multi-step task completion — executing a sequence of dependent actions without dropping steps or losing context
3. Precondition verification — checking that required conditions are met before performing an action (e.g., verifying an order exists before refunding it)
4. Tool-calling precision — using the correct tool with the correct arguments in the correct format

This is not a coincidence. These capabilities are real bottlenecks in the agent's performance. The model already knows how to do many things well — it can hold conversations, follow instructions, generate coherent text. But these 4 specific capabilities are where it falls short on this particular task distribution.

Let's make this concrete. Suppose you analyze 50 success trajectories and 50 failure trajectories for the "precondition verification" capability:

• In the 50 successes, 8 had precondition errors → ER⁺ = 8/50 = 0.16
• In the 50 failures, 35 had precondition errors → ER⁻ = 35/50 = 0.70
• Contrastive gap: Δ̂ = 0.70 − 0.16 = 0.54

A gap of 0.54 is enormous. It means precondition verification errors are 4.4x more common in failures than in successes. This capability is a major bottleneck.

Now compare with "greeting the customer politely":

• ER⁺ = 0.02, ER⁻ = 0.03
• Δ̂ = 0.01 — essentially zero

The model already greets customers fine. Training on this capability would be wasted effort.

Chapter 2: Formal Setup

The Agentic Environment

TRACE works with any agentic environment defined as a tuple (X, P, R, y) where:

• X = the task space. Each x ∈ X is a specific task instance (e.g., "Customer wants to return order #4521")
• P = the protocol. The rules of engagement — which tools are available, what actions the agent can take, the conversation format, system instructions
• R = the reward function. Maps a trajectory τ to a scalar R(τ) ∈ {0, 1}. Binary: did the agent solve the task?
• y = the target variable (ground truth answer or state the agent must achieve)

A trajectory τ is the full sequence of agent actions and environment responses: τ = (a₁, o₁, a₂, o₂, ..., a_T, o_T) where a_t is the agent's action at step t (a text response, a tool call) and o_t is the environment's observation (tool output, user reply).

What Is a Capability?

A capability c is a natural-language description of a specific skill the agent needs. Examples:

Critically, capabilities are NOT predefined by the user. TRACE discovers them automatically from the trajectories. The user provides only the environment and a dataset of tasks.

The Capability-Targeted Synthetic Environment

For each identified capability deficit c, TRACE synthesizes a new environment E_c = (X_c, P_c, R_c, y_c) with three key properties:

1. Capability isolation: E_c primarily tests capability c. Success requires c and little else. This means RL training on E_c produces a clean gradient signal for learning c.
2. Interface preservation: P_c uses the same tool schemas, action space, and conversation format as the target environment P. The agent interacts with E_c the same way it interacts with the real environment. This ensures the learned capability transfers.
3. Verifiability: R_c can be computed automatically. No human evaluation needed. This is essential for RL training which needs thousands of reward evaluations.

The Decomposition Advantage

Why decompose at all? Consider the alternative: train one adapter on the full target environment. The RL loss mixes gradients from all capability deficits. A batch of 64 trajectories might contain 20 failures due to structured data reasoning, 15 due to precondition verification, 10 due to tool-calling precision, and 19 due to other causes. The policy gradient averages over all of these, producing a noisy signal that partially addresses all capabilities and fully addresses none.

In the mixed-capability setting, Â(τ) conflates advantages from different capabilities. Two trajectories that fail for different reasons get the same negative advantage. The gradient points in a direction that compromises between fixing structured data reasoning and fixing precondition verification — and that compromise may not be optimal for either.

With TRACE's decomposition, each adapter θ_i trains on E_ci where all failures are about capability c_i. The advantage function Â(τ) is now pure signal: "did you do c_i correctly?" The gradient is clean and converges faster.

Chapter 3: Contrastive Capability Identification

This is the first step of TRACE and arguably the most elegant. The goal: given a set of trajectories, automatically discover which capabilities the agent lacks. No human annotation. No predefined capability taxonomy. The capabilities emerge from the data.

Step 1: Split Trajectories

Run the base model on a dataset D of tasks. Each task produces a trajectory τ with reward R(τ) ∈ {0,1}. Split into:

• D⁺ = successful trajectories (R = 1)
• D⁻ = failed trajectories (R = 0)

In the paper's τ²-Bench experiments, the base model (Llama-3.1-8B-Instruct) achieves about 24% success rate on 680 tasks. That gives roughly D⁺ ≈ 163 successes and D⁻ ≈ 517 failures.

Step 2: Discovery Phase

Sample a batch of ~50 failure trajectories from D⁻. Feed them to an LLM (GPT-4.1 in the paper) with a prompt that says:

The LLM produces a capability dictionary — a list of (name, description) pairs. In the paper, this typically yields 5–8 candidate capabilities. The discovery is run multiple times with different trajectory samples, and all discovered capabilities are pooled.

Step 3: Labeling Phase

Now comes the scoring. For every trajectory τ in D⁺ ∪ D⁻ and every candidate capability c, an LLM labels the trajectory as one of:

• NA — this capability is not relevant to this trajectory (e.g., a task that never requires precondition checking)
• PRESENT — the capability is relevant and the agent demonstrated it correctly
• LACKING — the capability is relevant and the agent failed to demonstrate it

This is the most computationally expensive step. If you have 680 trajectories and 7 candidate capabilities, that is 4,760 labeling calls. The paper uses GPT-4.1 for this, which costs money but gives reliable labels.

Step 4: Compute Error Rates

For each capability c, compute two error rates. Let N⁺(c) = number of successes where c is relevant (not NA), and E⁺(c) = number of those where c is LACKING:

Think of ER⁺(c) as the "background noise" — how often this capability fails even in trajectories that ultimately succeed. Maybe the agent makes a precondition error 16% of the time but recovers from it (corrects itself on the next turn). ER⁻(c) is the "failure signal" — how often this capability fails in trajectories that ultimately fail.

Step 5: Contrastive Gap and Coverage

The contrastive gap is the key statistic:

A high Δ̂(c) means: "when this capability fails, the whole trajectory fails." A low Δ̂(c) means: "this capability fails at similar rates regardless of outcome — it is not the bottleneck."

The coverage measures what fraction of failures involve this capability:

Coverage prevents TRACE from selecting capabilities that have a high contrastive gap but only affect 3% of failures. A capability that perfectly separates 2 failures from successes is not useful if you have 500 other failures to fix.

Step 6: Threshold and Select

TRACE retains capability c if:

Let's walk through a concrete worked example with real-ish numbers from τ²-Bench:

Chapter 4: The TRACE Pipeline

Now we assemble all four steps into the full TRACE pipeline. This is the showcase visualization — step through the entire system from raw trajectories to inference-time routing.

Let's trace through each stage with actual numbers from the τ²-Bench experiments.

Stage 1: Identify

The base model (Llama-3.1-8B-Instruct) achieves 24.3% on τ²-Bench. TRACE rolls it out on 680 tasks: 165 successes, 515 failures. After contrastive analysis, 4 capabilities emerge with gaps ranging from 0.38 to 0.54.

Stage 2: Synthesize

For each of the 4 capabilities, TRACE prompts an LLM to generate a synthetic environment. The synthesis prompt includes: the capability description, the target environment's tool schemas, and example trajectories. The LLM produces Python code that implements a seeded task generator, transition logic, and an evaluation function. Each synthetic env produces unlimited training tasks on demand.

Stage 3: Train

Each capability gets its own LoRA adapter (rank 32, ~5.3% of parameters). Training uses GRPO with G=4 groups, K=4 rollouts per group per seed, AdamW optimizer (lr=1e−5), for 40 iterations on 4–8 A100 GPUs. Total: 640 rollouts per adapter (40 iterations × 4 groups × 4 rollouts).

Stage 4: Route

At inference, a new task comes in. The base model is shown the task description, the 4 capability descriptions, and a few example trajectories for each capability. It is asked: "Which capability is most relevant to this task? Choose one, or 'none' if the base model suffices." TRACE takes the argmax of the logit scores for each capability label token. The task is routed to the corresponding adapter (or the base model if "none" wins).

End-to-End Numbers

After training 4 LoRA adapters (one per capability), TRACE achieves 38.4% on τ²-Bench, up from 24.3% — a +14.1 point improvement. For comparison:

• GRPO on target environment: 30.4% (+6.1)
• AWM (Agentic Workflow Memory): 28.1% (+3.8)
• ADP (Agentic Data Pipeline): 27.9% (+3.6)
• GEPA (prompt evolution): 26.8% (+2.5)

TRACE more than doubles the improvement of any baseline. And it does this with only 4 adapters, each updating 5.3% of parameters.

Chapter 5: Environment Synthesis

This is the second stage of TRACE and the most creative. Given a capability description (e.g., "structured data reasoning") and the target environment's interface, TRACE must generate a complete, verifiable training environment that isolates that capability.

What the Synthesis LLM Receives

The synthesis prompt includes:

1. The capability description — what skill to test, from the contrastive identification step
2. The target environment's tool schemas — the exact API signatures (function names, parameter types, return types) so the synthetic env uses the same interface
3. Example failure trajectories — 3–5 trajectories from D⁻ where this specific capability was labeled LACKING, so the LLM can see what failure looks like in context
4. Example success trajectories — 2–3 from D⁺ where the capability was PRESENT, showing what correct behavior looks like

What the Synthesis LLM Produces

The LLM generates Python code implementing three components:

Concrete Example: Structured Data Reasoning

Suppose the identified capability is "structured data reasoning." The synthesis LLM might produce:

python
def generate_task(seed: int) -> Task:
    rng = random.Random(seed)
    # Create a customer with 3-5 orders
    n_orders = rng.randint(3, 5)
    orders = []
    for i in range(n_orders):
        orders.append({
            "order_id": f"ORD-{seed}-{i}",
            "product": rng.choice(PRODUCTS),
            "price": round(rng.uniform(10, 500), 2),
            "status": rng.choice(["delivered", "shipped", "returned"]),
            "date": random_date(rng),
        })
    # Task: "What is the total spent on delivered orders?"
    target = sum(o["price"] for o in orders if o["status"] == "delivered")
    return Task(orders=orders, question=..., answer=target)

The key is that this environment isolates structured data reasoning. The task is simple (look up orders, filter by status, sum prices) but requires correctly parsing JSON, filtering on a field, and computing an aggregate. There is no precondition verification, no complex multi-step planning. Just the one capability being trained.

Interface Preservation Is Critical

If the target environment exposes tools like get_order_details(order_id) and process_refund(order_id, amount), the synthetic environment must expose the SAME tools with the SAME signatures. The agent should not be able to distinguish whether it is talking to the real environment or the synthetic one.

Why? Because LoRA adapters modify the model's behavior within the same input distribution. If the synthetic env uses different tool names or parameter formats, the adapter learns to use those different names — and that skill does not transfer to the real environment. Interface preservation is what makes transfer possible.

Chapter 6: RL Training with GRPO

Once we have synthetic environments, we need to train adapters on them. TRACE uses GRPO (Group Relative Policy Optimization) — a variant of policy gradient methods designed for language model fine-tuning. Let's derive how GRPO works and why it is used here.

The REINFORCE Foundation

Standard REINFORCE computes the policy gradient as:

where R(τ) is the trajectory reward and π_θ(τ) = ∏_t π_θ(a_t|s_t) is the probability of the trajectory under policy θ. The problem: R(τ) has high variance. A reward of 1 on one trajectory does not tell you whether that trajectory was lucky or genuinely good.

GRPO's Solution: Within-Group Normalization

GRPO reduces variance by computing relative advantages within a group of rollouts from the same seed. For each seed (task instance), GRPO runs K rollouts and normalizes rewards within the group:

where μ_group and σ_group are the mean and standard deviation of rewards within the group. This makes the advantage relative: "how much better was this rollout than the average rollout on the same task?"

Let's work through a concrete example. Suppose we have one task with K=4 rollouts:

The Full GRPO Objective

GRPO uses a clipped objective similar to PPO, but without a separate value function:

where r_i(θ) = π_θ(τ_i) / π_old(τ_i) is the importance ratio, ε is the clip range (typically 0.2), and the KL term prevents the model from drifting too far from the reference policy (the base model).

TRACE's Training Configuration

The paper uses these specific hyperparameters:

The total training budget per adapter is 640 rollouts. With 4 adapters, that is 2,560 rollouts total. Compare with GRPO-on-target which uses the same 2,560 rollouts but trains a single adapter on the mixed-capability target environment — and achieves worse results.

Why LoRA Per Capability Beats Multi-Capability Training

Table 3 in the paper tests "merged" training (one adapter trained on all 4 synthetic environments mixed together). The merged adapter achieves only +8.2 points vs TRACE's +14.1 points. Why?

Because merged training suffers from the same gradient conflation as GRPO-on-target. Within a single training batch, some rollouts are about structured data reasoning, others about precondition verification. The gradient is an average over these different skills. The model learns a compromise that partially addresses all capabilities but excels at none.

Chapter 7: Routing & Composition

After training K adapters (one per identified capability), TRACE has K specialized models plus the base model. At inference time, a new task arrives and TRACE must decide: which adapter (if any) should handle this task?

Why Routing Beats Merging

The obvious alternative is to merge all adapters into one model. LoRA adapter merging is well-studied: you can add the adapter weight matrices. But Table 3 in the paper shows this performs worse than routing:

The Routing Mechanism

TRACE uses a surprisingly simple routing approach — it reuses the base model itself as the router. Here is how:

1. Construct a prompt that includes:
   • The incoming task description
   • Each capability's name and description
   • 2–3 example trajectories per capability (showing what that capability looks like in practice)
   • A "none" option for tasks the base model can already handle
2. Feed this prompt to the base model (not any adapter)
3. Look at the logit scores for the label tokens ("A", "B", "C", "D", "none")
4. Take the argmax — route to the adapter with the highest logit

The "None" Option

Not every task requires a specialized adapter. Some tasks are straightforward enough that the base model handles them correctly. TRACE includes a "none" option in the routing prompt. If the base model assigns the highest logit to "none," the base model handles the task directly.

This is important for two reasons:

• Avoiding regression: An adapter trained for structured data reasoning might perform worse on tasks that require creative conversation (because the LoRA weights slightly modify the model's conversational ability). Routing to "none" prevents this.
• Efficiency: Why load and run an adapter if the base model can already handle the task?

Routing Accuracy

The paper does not report routing accuracy directly, but the +14.1 point improvement implies the router is quite accurate. If routing were random, you would expect much smaller gains (each adapter only helps on its specific capability subset).

An informal analysis: if TRACE routes 60% of tasks correctly to the best adapter, and each adapter improves its capability by ~35% on relevant tasks, the expected gain is 0.60 × 0.35 × fraction_of_tasks_per_cap — which, summed over 4 capabilities covering ~80% of failures, gives roughly the observed +14 point gain.

Chapter 8: Results & Scaling

Main Results: τ²-Bench

τ²-Bench is a customer service benchmark with 680 tasks across 5 domains (retail, airline, hotel, insurance, telecom). Each task requires multi-turn interaction with tool calls. The metric is task success rate (binary: did the agent achieve the correct outcome?).

TRACE more than doubles the gain of the next best method (GRPO at +6.1). That gap is remarkable. Let's understand why.

Main Results: ToolSandbox

ToolSandbox tests tool-calling ability across 8 diverse scenarios (calendar management, file operations, email, etc.). The metric counts "perfect scores" — tasks where the agent gets everything right.

Scaling Behavior

Figures 3 and 4 in the paper show how TRACE and baselines scale with the number of training rollouts. This is one of the most important results:

• GRPO on target shows diminishing returns. Going from 1,280 to 2,560 rollouts adds only ~1.5 points. The mixed-capability gradients hit a ceiling.
• TRACE shows increasing returns in the same range. Going from 1,280 to 2,560 rollouts adds ~4 points. Each additional rollout on a focused capability env produces more useful learning signal.

The scaling curves cross around 640 rollouts — below that, GRPO-on-target is competitive (since TRACE spends some of its budget on capability identification). Above 640, TRACE pulls away rapidly.

Stability Analysis (Figure 2a)

One concern with LLM-based capability identification: is it reproducible? Or does each run produce random capabilities? The paper runs the identification step 10 times independently. Result: the same 4 core capabilities are identified in every single run, though their exact names and descriptions vary slightly. The contrastive gaps are consistent to within ±0.05.

This stability is a strong signal. It means the capability deficits are real, objective properties of the model-environment pair, not artifacts of the LLM's wording choices.

Ablation: What Happens Without Each Stage?

The paper ablates each component:

• Without contrastive identification (use manually specified capabilities): −4.2 points. Human-specified capabilities miss subtle bottlenecks.
• Without synthetic envs (train adapters on target env directly, one per manually identified capability): −6.8 points. Target env is too complex for focused learning.
• Without routing (merge adapters instead): −5.7 points. Adapter interference degrades performance.
• Without per-capability adapters (one merged adapter on all synthetic envs): −5.9 points. Gradient conflation returns.

Every component matters. The full pipeline is greater than the sum of its parts.

Chapter 9: Connections

Related Work

Limitations

• LLM-dependent capability identification: The quality of discovered capabilities depends on the LLM used for analysis (GPT-4.1). A weaker model might miss subtle capabilities or hallucinate non-existent ones.
• Synthetic environment fidelity: Synthesized environments may not perfectly capture the complexity of the target environment. The paper shows transfer works well, but the gap between synthetic and real performance is not zero.
• Binary rewards only: TRACE currently works with binary success/failure. Extending to continuous or multi-dimensional rewards is future work.
• Routing overhead: Each inference requires a routing step (one additional forward pass). For latency-critical applications, this adds delay.
• Tested on 8B models only: The paper uses Llama-3.1-8B-Instruct. Whether the same approach helps 70B+ models (which may already have these capabilities) is unknown.

Cheat Sheet

The Big Picture

TRACE represents a shift from "train harder" to "train smarter." Instead of throwing more compute at the same optimization problem, it decomposes the problem into subproblems that are individually easier. This is the same insight that powers mixture-of-experts models, modular neural networks, and curriculum learning — but applied to the meta-level problem of figuring out what to train on.

The most surprising finding is how few capabilities matter. Across 680 diverse customer service tasks, just 4 capabilities explain the vast majority of failures. This suggests that LLMs are closer to being great agents than raw benchmark numbers imply — they are not failing broadly, they are failing specifically, and those specific deficits can be surgically corrected.