TRACE: Capability-Targeted Agentic Training

Paper: arXiv:2604.05336 Code: ScalingIntelligence/TRACE Code reference: main @ 97c03dd8 (2026-04-13) Batch: batch_index=119, category_rank=9, Agent / Reward & Training Methodology

1. Motivation (研究动机)

现有 agentic training 的核心瓶颈不是“没有更多 RL 数据”，而是目标环境的 reward 往往只告诉模型任务成败，却不告诉模型缺少哪一种可迁移能力。论文把 capability 定义为：在某一类任务实例中，成功轨迹必须执行的一组动作或判断；例如客服环境中“检索正确 customer record”是取消航班、改座位、验证预订等任务共享的前置能力。

直接在目标环境上做 RL / SFT 会把 credit assignment 绑定到具体任务轨迹：最终成功、失败、演示 trace 都只监督了一串任务特定动作，模型需要隐式地从稀疏反馈中发现“到底缺 structured data reasoning、precondition verification，还是 tool-calling precision”。这会导致样本效率低，也会让 broad synthetic data / general agent data 方法训练到与当前失败无关的能力。

TRACE 要解决的具体问题是：给定 base agent 在目标环境中的成功与失败轨迹，自动找出“失败中显著缺、成功中不缺”的高影响 capability；为每个 capability 合成可验证、可程序化打分的 targeted environment；再分别训练 LoRA adapter，并在推理时选择最相关的 adapter。这个问题值得研究，因为 agent deployment 的长期改进不应只依赖人工分析 failure cases 或盲目扩大 synthetic environment，而应能把真实失败转化成可复用、可验证、可扩展的训练信号。

2. Idea (核心思想)

核心 insight：agent 的失败不是均匀随机的，少数高覆盖 capability deficits 会解释大量 failed trajectories；如果先用 successful-vs-failed contrastive analysis 把这些 deficits 显式命名，再为每个 deficit 构造 isolated verifiable environment，那么 RL 的 reward 会从“任务级稀疏信号”变成“能力级归因信号”。

TRACE 的关键创新不是提出新的 RL optimizer，而是把训练目标重构为四步闭环：contrastive capability identification → capability-targeted environment generation → per-capability GRPO LoRA training → training-free adapter routing。与 GRPO-on-target 相比，TRACE 不在原环境里让模型同时摸索所有能力；与 AWM / ADP 这类 general-purpose synthetic/data scaling 相比，TRACE 的环境由当前模型的真实失败驱动；与 GEPA prompt optimization 相比，TRACE 不是把 capability 描述塞进 prompt，而是把 capability 训练成独立 adapter。

3. Method (方法)

3.1 Overall framework (整体框架)

Figure 1 解读：左侧是目标环境中的 base-agent rollouts，TRACE 先收集成功与失败轨迹；中间用 LLM analysis agent 归纳 capability dictionary 并对每条轨迹标注 NA / PRESENT / LACKING；随后 LLM generation agent 为每个 retained capability 生成 synthetic verifiable environment；右侧对每个环境训练一个 LoRA adapter，推理时由 base model 根据 task prompt 和 capability descriptions 选择一个 adapter。

形式化地，目标 agentic environment 写作： $\mathcal{E}=({\mathcal{X}}_{\mathcal{E}},P_{\mathcal{E}},R_{\mathcal{E}},y_{\mathcal{E}}),J({\pi }_{\theta };\mathcal{E})={\mathbb{E}}_{x\sim {\mathcal{X}}_{\mathcal{E}},\tau \sim p_{{\pi }_{\theta }}(\cdot \mid x,\mathcal{E})}[R_{\mathcal{E}}(x,\tau )].$ 一个 capability $c$ 是成功解某个任务子集 $X_c\subseteq X_{\mathcal{E}}$ 所必要的动作/判断集合。若 ${\nu }_c(x,\tau )=1$ 表示轨迹 $\tau$ 对任务 $x$ 真的 exercise 了 capability $c$，论文要求成功蕴含必要能力被执行： $y_{\mathcal{E}}(x,\tau )=1\Rightarrow {\nu }_c(x,\tau )=1.$ TRACE 的能力分解直觉是：原环境 reward $R_{\mathcal{E}}$ 同时混合许多能力，失败归因弱；而 capability-targeted synthetic environment ${\mathcal{E}}_s^c=(G_c,P_c,R_c,y_c)$ 强制每个生成任务 $x_c=G_c(z)$ 都需要 capability $c$，同时保留目标环境的 tool schemas、state representation 和 policy constraints，并让 reward/success 可由程序自动计算。

3.2 Contrastive capability identification (对比式能力识别)

给定 base policy 在目标环境收集的数据： $\mathcal{D}=\{(x_i,{\tau }_i,r_i,y_i){\}}_{i=1}^N,$ TRACE 先按 success label 切分为 ${\mathcal{D}}^{+}$ 与 ${\mathcal{D}}^{-}$。为了避免每次 LLM 分析生成不一致的能力名称，流程分两阶段：

• Discovery phase：analysis agent 检查 tool calls、tool results、final responses，归纳候选 capability dictionary $\mathcal{C}$，每个 capability 有固定名称与自然语言描述。
• Labeling phase：用固定 dictionary 重新标注每条 trajectory 与每个 capability 的关系：NA 表示该任务不需要该能力；PRESENT 表示需要且轨迹已执行；LACKING 表示需要但轨迹未执行。

对每个 capability $c$，论文估计 successful / failed trajectories 中的 lacking rate： ${\widehat{\mathrm{ER}}}^{+}(c)=\frac{{\sum }_{i}1[{\ell }_i^c=\text{textsc}Lacking,y_i=1]}{{\sum }_{i}1[{\ell }_i^c\ne \text{textsc}NA,y_i=1]},{\widehat{\mathrm{ER}}}^{-}(c)=\frac{{\sum }_{i}1[{\ell }_i^c=\text{textsc}Lacking,y_i=0]}{{\sum }_{i}1[{\ell }_i^c\ne \text{textsc}NA,y_i=0]}.$ contrastive gap 定义为： $\hat{\Delta }(c)={\widehat{\mathrm{ER}}}^{-}(c)-{\widehat{\mathrm{ER}}}^{+}(c).$ coverage 定义为： $\widehat{\mathrm{Cov}}(c)=\frac{1}{|{\mathcal{D}}^{-}|}{\sum }_{i=1}^{N}1[{\ell }_c(x_i,{\tau }_i)=\text{textsc}Lacking\wedge y_i=0].$ 最终保留： ${\mathcal{C}}^{\ast }=\{c\in \mathcal{C}\mid \hat{\Delta }(c)\ge \delta ,\widehat{\mathrm{Cov}}(c)\ge \rho \}.$ 论文实验设定 $\rho =0.10$、$\delta =0.20$；released code 的 configs/capability_selection.yaml 还设定 n_runs=10、n_candidates=10、k_consistency=8，并由 pipeline/aggregate_capabilities.py 实现“10 次运行中至少 8 次通过阈值”的 consistency filter。

def select_capabilities_from_runs(runs, rho=0.10, delta=0.20, k=8):
    # mirrors pipeline/aggregate_capabilities.py
    summaries = []
    for cap in all_capability_names(runs):
        pass_count, covs, deltas = 0, [], []
        for run in runs:
            labels = run["labels"].get(cap, {})
            total_failed = run["totals"]["total_failed"]
            lf = len(labels.get("lacking_failed", []))
            pf = len(labels.get("present_failed", []))
            lp = len(labels.get("lacking_passed", []))
            pp = len(labels.get("present_passed", []))
            cov = lf / max(total_failed, 1)
            er_minus = lf / max(lf + pf, 1)
            er_plus = lp / max(lp + pp, 1)
            gap = er_minus - er_plus
            covs.append(cov); deltas.append(gap)
            pass_count += int(cov >= rho and gap >= delta)
        if pass_count >= k:
            summaries.append({
                "skill": cap,
                "mean_cov": float(torch.tensor(covs).mean()),
                "mean_delta": float(torch.tensor(deltas).mean()),
                "n_runs_passed": pass_count,
                "status": "PENDING",
            })
    return sorted(summaries, key=lambda x: -x["mean_delta"])

3.3 Capability-targeted environment synthesis (能力定向环境合成)

对每个 retained capability，generation agent 接收 capability description 与对应 failure patterns，生成完整 synthetic environment：seeded task generator $G_c$、transition logic、evaluation criteria。每个 seed $z$ 确定性生成一个 task instance $x=G_c(z)$，包含 synthetic user profiles、database records、task-specific parameters；这样可以同时避免 memorization，并让 rollout group 能共享同一初始状态。

对 ${\tau }^2$-Bench 的 structured data reasoning，附录中的 synthetic game 会随机生成 Airline / Retail 数据库、用户状态和 goal constraints；agent 多轮调用工具查询/修改数据库，与模拟用户对话；reward 通过 final database hash 与 gold action replay 比对，并检查最终通信内容是否包含应报告信息。对 ToolSandbox 的 error recovery，synthetic game 会构造 PermissionError 触发链，例如 low-battery setting 阻断目标服务；reward 由 action score 与 communication score 加权：$0.6\cdot \text{action}+0.4\cdot \text{communication}$。

class CapabilityGame:
    # aligned with game_registry.GameEnv and appendix synthetic games
    def reset(self, seed: int):
        self.scenario = self.generator(seed)
        self.db = copy.deepcopy(self.scenario.initial_db)
        self.done = False
        return self.scenario.initial_user_message
 
    def step(self, agent_action: str):
        if is_tool_call(agent_action):
            result = execute_tool(agent_action, self.db)
            return tool_result_to_observation(result)
        user_text = self.user_simulator.reply(agent_action)
        if should_stop(user_text):
            self.rewards[0] = self.evaluate(self.db, agent_action)
            self.done = True
        return user_text
 
    def evaluate(self, final_db, final_response):
        state_ok = hash(final_db) == hash(replay_gold_action(self.scenario))
        comm_ok = all(v in final_response for v in self.scenario.required_outputs)
        return 1.0 if state_ok and comm_ok else 0.3 if state_ok else 0.0

3.4 GRPO LoRA training (每个能力一个 adapter)

TRACE 对每个 $c\in {\mathcal{C}}^{\ast }$ 训练一个 LoRA adapter ${\Delta }_c$，base policy ${\pi }_{\theta }$ 冻结。每次 iteration 中，adapted policy ${\pi }_{\theta +{\Delta }_c}$ 在 ${\mathcal{E}}_s^c$ 中生成 $G$ 个 rollout groups；每个 group $g$ 使用同一 seed $z_g$ 采样 $K$ 条 trajectories，因此差异主要来自 stochastic decoding。论文中的 group-relative advantage 是： ${\hat{A}}_{g,k}=\frac{r_{g,k}-{\bar{r}}_g}{{\sigma }_g+ϵ},{\bar{r}}_g=\frac{1}{K}{\sum }_{j=1}^K{r}_{g,j},{\sigma }_g=\sqrt{\frac{1}{K}{\sum }_{j=1}^K(r_{g,j}-{\bar{r}}_g{)}^2}.$ adapter 用 clipped GRPO surrogate 更新： ${\mathcal{L}}_{\mathrm{GRPO}}=-\frac{1}{|\mathcal{B}|}{\sum }_{(g,k)\in \mathcal{B}}\frac{1}{T_{g,k}}{\sum }_{t=1}^{T_{g,k}}\min \left({\rho }_t{\hat{A}}_{g,k},\mathrm{clip}({\rho }_t,1-ϵ,1+ϵ){\hat{A}}_{g,k}\right),$ 其中 ${\rho }_t=\frac{{\pi }_{\theta +{\Delta }_c}(a_t\mid x,a_{<t})}{{\pi }_{\theta +{\Delta }_c}^{\mathrm{old}}(a_t\mid x,a_{<t})}.$ released code 的 train/train_grpo.py 采用 vLLM rollout worker + PyTorch distributed training：默认 group_size=16、groups_per_batch=4、lr=1e-5、mini_batch_size=4、clip_eps=0.28、filter_constant_groups=True、dynamic_sampling_max_batches=3；README.md 的 launch recipe 使用 1 张 GPU 跑 vLLM、其余 GPU 跑 torchrun -m train，并启用 runtime LoRA reload。

def grpo_train_iteration(model, old_policy, backend, game, tokenizer, args):
    # based on train/train_grpo.py
    samples, metrics = collect_grpo_rollouts(
        backend=backend,
        tokenizer=tokenizer,
        game_spec=game,
        groups_per_batch=args.groups_per_batch,
        group_size=args.group_size,
        temperature=args.temperature,
        base_seed=args.seed,
    )
    adv = compute_group_advantages(samples)          # code centers rewards per group
    if args.filter_constant_groups:
        samples, adv, _ = filter_constant_reward_groups(samples, adv)
 
    seqs, prompt_lens, action_lens = build_prompt_action_tensors(samples, tokenizer)
    old_logp = old_policy.logprob_action_tokens(seqs, prompt_lens, action_lens)
 
    for batch in make_sorted_batches(lengths(seqs), args.mini_batch_size):
        logits = model(input_ids=batch.ids, attention_mask=batch.mask).logits
        new_logp = logprob_action_tokens(logits, batch.ids, batch.prompt_lens, batch.action_lens)
        ratio = torch.exp(new_logp - old_logp[batch.idx])
        surr1 = ratio * adv[batch.idx]
        surr2 = torch.clamp(ratio, 1 - args.clip_eps, 1 + args.clip_eps) * adv[batch.idx]
        loss = -torch.mean(torch.min(surr1, surr2))
        loss.backward()
    return metrics

论文公式与 released code 实现差异：论文 Appendix 的 GRPO advantage 明确除以 group standard deviation ${\sigma }_g+ϵ$；train/train_grpo.py::compute_group_advantages 在当前 commit 中只做 reward - group_mean，没有除以标准差。论文实验还写“up to 40 iterations per capability”，但 public train_grpo.py 主循环是 range(start_iter, 10_000)，没有暴露 40-iteration stop 参数；因此复现实验需要外部 launch/停止策略或 checkpoint 选择。另一个代码内差异是 train/config.py 写 LORA_ALPHA=8，而 train/train_grpo.py CLI 默认 --lora-alpha=16；实际运行应以 launch args / trainer parser 为准。

3.5 Composing acquired capabilities (能力组合与路由)

训练完成后得到 $\{{\Delta }_c{\}}_{c\in {\mathcal{C}}^{\ast }}$。TRACE 不把所有 capabilities merge 成一个模型，而是使用 training-free routing：给 base model 一个 routing prompt $M_x$，其中包含当前 task prompt $o_1$、候选 capability descriptions、每个 capability 的一个成功 synthetic trajectory 示例，以及一个 base label token。base model 比较下一 token logits；若 base 最高则使用原模型，否则激活对应 LoRA adapter。 $c^{\ast }=\arg {\max }_{c\in {\mathcal{C}}^{\ast }\cup \{\mathrm{base}\}}{\mathrm{logit}}_{{\pi }_{\theta }}(t_c\mid M_x),W^{\prime }=W+B_{c^{\ast }}{A}_{c^{\ast }}.$

def route_adapter_for_task(base_model, task_prompt, capabilities, adapters):
    # paper routing; public repo describes this in README/paper but does not ship a full router class
    labels = make_single_token_labels(capabilities + ["base"])
    routing_prompt = build_routing_prompt(task_prompt, capabilities, labels)
    next_token_logits = base_model.next_token_logits(routing_prompt)
    choice = max(labels, key=lambda tok: next_token_logits[token_id(tok)])
    if choice == labels["base"]:
        return base_model
    adapter = adapters[capability_for_label(choice)]
    return activate_lora(base_model, adapter)  # W' = W + B @ A for adapted layers

论文公式与 released code 实现差异：paper/README 描述了 Select & Adapt 路由，但当前 public repo 主要提供 pipeline rendering、capability aggregation、environment generation instructions、GRPO training 和 vLLM LoRA reload；没有找到一个完整的 inference-time router implementation。因此本节路由伪代码以论文与 README 为依据，code mapping 只能锚到 README / Appendix routing prompt，而非具体 router source file。

Code reference: main @ 97c03dd8 (2026-04-13) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
capability selection config	configs/capability_selection.yaml	n_runs=10, rho=0.10, delta=0.20, k_consistency=8
contrastive aggregation	pipeline/aggregate_capabilities.py	compute_metrics, main
environment generation config	configs/environment_generation.yaml	group_size=16, num_seeds=100, hint_ratio=0.25-0.5
pipeline prompt rendering	render_pipeline.py	render_capability_selection, render_environment_generation
synthetic env interface	game_registry.py	GameEnv, GameSpec, register_game
rollout collection	train/train_grpo.py	collect_grpo_rollouts, UserLLMClient
group advantage / filtering	train/train_grpo.py	compute_group_advantages, filter_constant_reward_groups
LoRA GRPO trainer	train/train_grpo.py	parse_grpo_args_optimized, main training loop
vLLM LoRA sync	train/inference.py	sync_policy, vLLM backend helpers
inference routing	paper / README / Appendix routing prompt	no standalone router class found in released code

4. Experimental Setup (实验设置)

4.1 Datasets and benchmarks (数据集与规模)

• ${\tau }^2$-Bench：customer-service workflow benchmark；论文使用 50 个 Airline tasks 与 114 个 Retail tasks，总计 164 tasks，评估 policy-sensitive workflows。
• ToolSandbox：129 个 base scenarios，评估 broader stateful tool-use capabilities。
• Evaluation protocol：除非特别说明，所有模型 greedy decoding，temperature $0.0$，最大 context length 32,000 tokens，每个 episode 最多 50 interaction steps；user simulator 与 evaluated agent 使用同一 base model。

4.2 Baselines (对比方法)

主要 baselines 包括 Base Model、GRPO on Target、ADP、AWM、GEPA、Single Capability GRPO。能力合并 ablation 还比较 CORE-TSV、On Policy Distillation、SFT Synthetic、Multi Capability GRPO。公平性设置：training baselines 使用与所有 capabilities 分别训练总和相同的 training iterations；GEPA 被限制为每个 capability 使用相同 total rollout budget。

4.3 Metrics (评价指标)

${\tau }^2$-Bench 报告 Airline、Retail 与 Overall pass rate；Overall 定义为： ${\mathrm{PassRate}}_{\mathrm{overall}}=\frac{S_{\mathrm{Airline}}+S_{\mathrm{Retail}}}{N_{\mathrm{Airline}}+N_{\mathrm{Retail}}}.$ ToolSandbox 报告 Perfect 与 Mean Sim.： $\mathrm{MeanSimilarity}=\frac{1}{N_{\mathrm{TS}}}{\sum }_{i=1}^{N_{\mathrm{TS}}}{\mathrm{Sim}}_i,\mathrm{PerfectRate}=\frac{1}{N_{\mathrm{TS}}}{\sum }_{i=1}^{N_{\mathrm{TS}}}1[{\mathrm{Sim}}_i=1].$

4.4 Training config (训练配置)

论文使用 Qwen3-30B-A3B-Instruct-2507 同时作为 base agent 和 simulator，使用 LoRA + GRPO，optimizer 为 AdamW，learning rate $10^{-5}$，每个 capability 最多 40 iterations，rollout collection temperature $1.0$，gradient checkpointing + distributed data parallelism，硬件为 4—8 张 A100-80GB GPU。

released config/code 可复核的关键值：configs/capability_selection.yaml 设 model=Qwen/Qwen3-30B-A3B-Instruct-2507、n_runs=10、n_candidates=10、rho=0.10、delta=0.20、k_consistency=8；configs/environment_generation.yaml 设 group_size=16、num_seeds=100、hint_ratio=0.25-0.5、temperature=0.7、vLLM max-model-len=32000；train/train_grpo.py CLI 默认 group_size=16、groups_per_batch=4、lora_rank=16、lora_alpha=16、lr=1e-5、mini_batch_size=4、clip_eps=0.28。

5. Experimental Results (实验结果)

5.1 Main results (主结果)

Approach	Airline (%)	Retail (%)	Overall (%)
Base Model	24.0	36.8	32.9
GRPO on Target	32.0	40.4	37.8
ADP	28.0	34.2	32.3
AWM	32.0	41.2	38.4
GEPA	38.0	40.4	39.6
Single Capability GRPO (Ours)	34.0	43.0	40.3
TRACE (Ours)	44.0	48.2	47.0

TRACE 在 ${\tau }^2$-Bench 上达到 47.0 overall，比 base model 的 32.9 高 +14.1 points，比最强外部 baseline GEPA 的 39.6 高 +7.4 points。Single Capability GRPO 已达到 40.3，说明 targeted environment 本身强于 general-purpose AWM/ADP；完整 TRACE 进一步通过多 capability routing 获益。

Model	Perfect	Mean Sim.
Base (Qwen3-30B-A3B)	19/129	0.411
ADP	19/129	0.422
GRPO on Target	22/129	0.519
AWM	20/129	0.504
GEPA	22/129	0.520
Single Capability GRPO (Ours)	22/129	0.514
TRACE (Ours)	26/129	0.552

ToolSandbox 上 TRACE 达到 26/129 perfect 与 0.552 mean similarity，比 base 多 7 个 perfect、mean similarity +0.141；比 GEPA 多 4 个 perfect、mean similarity +0.032。

5.2 Capability analysis (能力识别质量)

Figure 2 解读：左图展示 10 次独立 contrastive analysis 中，analysis agent 稳定恢复同一组 top capability deficits；右图展示这些 deficits 对 failed tasks 的覆盖高度集中。论文附录说明，${\tau }^2$-Bench 上四个主要 gaps 是 Structured Data Reasoning、Tool Calling Precision、Multi-Step Task Completion、Precondition Verification；ToolSandbox 上两个 gaps 是 Permission Error Recovery 与 Datetime Reasoning。coverage 不互斥，因为同一失败轨迹可能同时缺多个 capability。

5.3 Routing vs. capability consolidation (路由优于合并)

Approach	Airline (%)	Retail (%)	Overall (%)
Base Model	24.0	36.8	32.9
Single Capability GRPO	34.0	43.0	40.3
CORE-TSV	36.0	41.2	39.6
On Policy Distillation	28.0	42.1	37.8
SFT Synthetic	36.0	38.6	37.8
Multi Capability GRPO	38.0	42.1	40.9
TRACE	44.0	48.2	47.0

合并类方法没有达到 routing 的效果：Multi Capability GRPO 为 40.9，CORE-TSV 为 39.6，均低于 TRACE 47.0。这支持论文的设计选择：不要把多个 capability 压进单一 adapter，而是在任务级动态选择单个 relevant adapter，避免能力干扰。

5.4 Scaling behavior (扩展性)

Figure 3 解读：capability 数量增加时，TRACE 的性能持续优于 GEPA-style prompt capability injection。论文正文特别强调，在 2—4 个 LoRA adapters 时，TRACE 可带来约 9.2 points 的 ${\tau }^2$-Bench 提升；这说明“能力变成 adapter”比“能力描述变成 prompt”更稳定。

Figure 4 解读：两张子图分别对应 ${\tau }^2$-Bench 与 ToolSandbox 的 rollout-budget scaling。结论是，在相同 rollout budget 下，TRACE 比 GRPO-on-target 和 GEPA 更高效；摘要给出的 ${\tau }^2$-Bench 数字是：相同 rollouts 下 TRACE 分别超过 GRPO 与 GEPA +9.2 和 +7.4 points。

5.5 Caveats and limitations (限制与注意点)

论文没有单列 Limitations section；从方法与 released code 可以确认的 caveats 是：第一，TRACE 依赖目标环境可收集 success/failure trajectories，且最好能把成功标准转化为 verifiable reward；第二，capability discovery 和 environment generation 依赖强 LLM analysis/generation agent，能力标签错误会直接影响后续训练；第三，当前实验集中在 ${\tau }^2$-Bench 与 ToolSandbox，并以 Qwen3-30B-A3B-Instruct-2507 为主，跨更多 agent domains / model families 的泛化还需要验证；第四，public repo 对训练与 pipeline 支持较完整，但 inference-time routing 没有独立实现文件，复现完整 Select & Adapt 需要按论文 prompt 自行实现或等待代码补全。