ScalingIntelligence/TRACE
GitHub: ScalingIntelligence/TRACE
TRACE 是一套端到端的智能体自我提升系统，通过自动识别 LLM 缺失能力、生成针对性训练环境并训练专用 LoRA 适配器来提升模型在复杂任务中的表现。
Stars: 92 | Forks: 10
# TRACE：面向能力的智能体训练
[](https://arxiv.org/abs/2604.05336)
[](https://scalingintelligence.stanford.edu/blogs/trace/)
本仓库包含 **TRACE — Turning Recurrent Agent failures into Capability-targeted training Environments**（将循环智能体失败转化为面向能力的训练环境）的官方实现。TRACE 是一个用于特定环境智能体自我提升的端到端系统。
作为智能体部署在复杂环境中的大型语言模型（LLM）通常会因为没有具备特定的底层能力而失败。现有的方法通常依赖于通用的合成数据，或者直接在目标环境中使用 RL，这迫使模型以隐式且低效的方式学习缺失的能力。TRACE 通过以下方式解决了这个问题：
1. **自动识别**智能体缺失的具体能力，
2. **合成**能够隔离出每种缺失能力的针对性训练环境，
3. 通过 GRPO 为每种能力单独**训练**一个 LoRA adapter，以及
4. 在推理时将任务**路由**到合适的 adapter。
## TRACE 的工作原理
该 pipeline 包含四个自动化步骤。每一步都由 LLM 编码智能体（Claude Code、Codex 等）根据 `prompts/` 下的 markdown 提示词来驱动。
| # | 步骤 | 提示词 | 输出 |
|---|---|---|---|
| 1 | **能力选择** | `prompts//capability_selection.md` | `selected_capabilities.json` — 缺失能力的排序列表 |
| 2 | **环境生成** | `prompts//environment_generation.md` | 每种能力对应一个合成训练环境 |
| 3 | **GRPO 训练** | `train/train_grpo.py` | 每种能力对应一个 LoRA adapter |
| 4 | **选择与适配** | `moe_gate/` | 将推理路由到正确 adapter 的门控 |
## 参考结果
TRACE 已在两个模型规模的三个 benchmark 上进行了验证：
| Benchmark | 基础模型 | 基线 | + TRACE | Δ | 备注 |
|---|---|---:|---:|---:|---|
| τ²-Bench (Airline + Retail) | Qwen3-30B-a3b-instruct | 32.9% | **48.3%** | **+15.4 pp** | 来自论文 |
| **SWE-bench Verified** | **Qwen3.6-27B** | **68.0%** | **73.2%** | **+5.2 pp** | 最佳迭代检查点；参见 `swebench/qwen3.6_self_trace/` |
Qwen3.6 SWE-bench 的参考运行位于 `swebench/qwen3.6_self_trace/` 中，包含阶段脚本、生成环境的包装器以及复现说明。这是将 TRACE 适配到新 benchmark 的推荐起点。
## 仓库结构
```
TRACE/
├── prompts/ Prompts for the LLM-driven pipeline stages.
│ ├── general/ Env-agnostic templates (use as starting point
│ │ when adapting TRACE to a new benchmark).
│ ├── tau-bench/ Filled-in version for τ²-Bench.
│ └── swebench/ Filled-in version for SWE-bench Verified.
│
├── pipeline/ Pipeline-stage scripts (benchmark-agnostic).
│ ├── aggregate_capabilities.py Aggregation step for capability selection.
│ ├── calibrate_environment.py Rollout statistics and pass@k threshold check.
│ └── _summarize.py Trajectory compaction for capability selection.
│
├── swebench/ SWE-bench reference runs (one folder per run).
│ └── qwen3.6_self_trace/ End-to-end Qwen3.6-27B run (phase1-phase6).
│
├── train/ GRPO training + gate training.
│ ├── train_grpo.py Train one capability-specific LoRA.
│ ├── train_router_sft.py SFT warm-start of the MoE gate.
│ ├── warmstart_gater.py Gate-only warm-start with frozen base + LoRAs.
│ ├── grpo_gater.py GRPO fine-tune of the MoE gate.
│ ├── collect_rollouts.py Rollout collection against vLLM.
│ └── ... model.py, ppo.py, inference.py, config.py
│
├── moe_gate/ Inference-time routing across capability LoRAs.
│ ├── README.md How the gate works (per-layer / global modes).
│ ├── gater.py LayerCapabilityGater / TrajectoryGlobalRouter.
│ ├── multi_lora.py MultiLoRAGatedLinear (weighted-mix forward).
│ ├── model_builder.py build_capability_model(...).
│ ├── hf_wrapper.py / freeze_utils.py / smoke_test.py
│
├── configs/ YAML configs for render_pipeline.py.
│ ├── capability_selection.yaml
│ ├── environment_generation.yaml
│ └── moe_gate.yaml
│
├── render_pipeline.py Renders prompt templates with config values.
├── game_registry.py Game/environment registry.
└── requirements.txt
```
## 快速开始
### 前置条件
- Python 3.11+
- 支持 CUDA 的 GPU：
- 在能力选择阶段，至少需要 **1 个 GPU** 来运行 vLLM 推理服务器。
- 进行 GRPO 训练建议使用 **4 个或更多 GPU**（在 128k 上下文下，设置 TP=4 可以让 Qwen3.6-27B 舒适地驻留在内存中）。
- [conda](https://docs.conda.io/en/latest/) 或同等环境。
### 安装
```
conda create -n trace python=3.11 -y
conda activate trace
pip install -r requirements.txt
```
## 快速上手
TRACE 的核心是将各阶段的 markdown 提示词交给 LLM 编码智能体（Claude Code、Codex 等）来驱动。对于每个阶段，您有两种选择：
1. 直接使用 `prompts//` 中**已填充**的提示词 —— 如果您的 benchmark 已有对应文件夹，这是最快的方法。
2. 通过 YAML 配置渲染 `prompts/general/` 中的**模板** —— 最适合适配到新的 benchmark。
### 路径 A — 使用特定 benchmark 的提示词
`tau-bench` 和 `swebench` 均附带了已渲染并经过审查的提示词：
```
# 能力选择（第 1 步）
cat prompts/swebench/capability_selection.md # read it first
# 然后将其交给 Claude Code / Codex / 等
# Environment 生成（第 2 步）
cat prompts/swebench/environment_generation.md
```
### 路径 B — 通过 YAML 配置渲染
使用您的模型和评估结果编辑 `configs/*.yaml` 中的任意一个，然后进行渲染：
```
python render_pipeline.py configs/capability_selection.yaml --stage capability
python render_pipeline.py configs/environment_generation.yaml --stage environment
```
这会将渲染后的提示词写入 `prompts/rendered/_.md`。然后将它们交给您的编码智能体。
## 分步指南
### 步骤 1 — 能力选择
运行 `prompts//capability_selection.md` 中描述的 **发现 + 标记 + 聚合** 流程。智能体将会：
1. 读取您的评估结果（通过/失败的轨迹）。
2. 阶段 1：提出 5-10 个候选能力。
3. 阶段 2：生成 10 个并行的标记子智能体，每个子智能体独立将每个（轨迹，能力）对标记为 `NA`、`PRESENT` 或 `LACKING`。
4. 阶段 3：运行 `pipeline/aggregate_capabilities.py` 来计算每次运行的覆盖率 `Cov(c)` 和对比差距 `Δ(c)`，应用双重阈值过滤（`Cov ≥ ρ`，`Δ ≥ δ`）以及跨运行一致性（`N` 次运行中的 `K` 次）。
**输出：** `pipeline/selected_capabilities.json` — 每个条目包含名称、描述、`mean_cov`、`mean_delta`、失败的案例示例，以及 `status: "PENDING"`。
### 步骤 2 — 环境生成
针对**每种能力**运行一次 `prompts//environment_generation.md`。每次调用都会选择处于 PENDING 状态且优先级最高的能力（平均 Δ 值最高），并执行以下操作：
1. 生成一个合成环境文件（对于工具使用的 benchmark，这是一个游戏类；对于 SWE-bench，这是一个由 pytest 评分的场景 JSON 集合）。
2. 对环境进行自测（断言目标测试在修复前失败，并在应用 oracle 后通过）。
3. 启动一个 vLLM 服务器，并在 temperature 为 1.0 时收集验证 rollouts。
4. **通过变异校准难度**：如果基础模型的成功率太高，智能体会应用表面伪装变异（重命名标识符、内联 helper、添加诱饵编辑点）并重新测试，最多进行 5 轮。
5. 将该能力标记为 `DONE` 并停止。用户需重新调用以处理下一个能力。
**输出（工具使用 benchmark）：** 位于项目根目录的 `capability__game.py`。
**输出（SWE-bench）：** `pipeline/swebench//scenarios_parsed.json`（每种能力包含 10 个校准好的场景）。
### 步骤 3 — GRPO 训练
为每种能力训练一个 LoRA adapter。
**3a. 启动 vLLM 推理服务器**：
```
conda activate trace
export CUDA_VISIBLE_DEVICES=0,1,2,3 # GPUs for inference
export VLLM_ALLOW_RUNTIME_LORA_UPDATING=True
vllm serve \
--host 0.0.0.0 --port 8000 \
--dtype bfloat16 \
--max-model-len 65536 \
--tensor-parallel-size 4 \
--gpu-memory-utilization 0.85 \
--enable-prefix-caching \
--enable-chunked-prefill \
--enable-lora \
--max-loras 2 \
--max-lora-rank 32 \
--enable-auto-tool-choice \
--tool-call-parser hermes # use 'qwen3_xml' for Qwen3.6
```
**3b. 在剩余的 GPU 上运行 GRPO 训练**：
```
conda activate trace
export CUDA_VISIBLE_DEVICES=4,5,6,7
export VLLM_BASE_URLS=http://localhost:8000
export VLLM_MODEL=
export VLLM_ALLOW_RUNTIME_LORA_UPDATING=True
torchrun --nproc_per_node= --master-port=29501 -m train \
--game capability_ \
--model
```
针对每种能力重复此操作。输出为每种能力对应的 LoRA adapter（`pipeline/adapters/.safetensors`）。
### 步骤 4 — 选择与适配 (MoE Gate)
一旦您获得了 N 个特定能力的 LoRA，就可以训练在推理时选择正确 LoRA 的路由门控。请参阅 `moe_gate/README.md` 了解架构，并在 `configs/moe_gate.yaml` 中查看超参数。
```
# 在标注的 (prompt, capability) 对上用 SFT warm-start 该 gate
python -m train.train_router_sft \
--config configs/moe_gate.yaml
# 可选：在合成的 envs 上对 gate 进行端到端的 GRPO fine-tune
python -m train.grpo_gater \
--config configs/moe_gate.yaml
```
**在推理时**，`moe_gate/` 中的 `build_capability_model(...)` 会加载基础模型 + 所有能力的 LoRA + 训练好的门控。该门控通过正确的 adapter（或软混合）为每一层路由每次前向传播。
## 将 TRACE 适配到新 benchmark
1. **从 `prompts/general/` 开始。** 将两个文件复制到 `prompts//` 中，并将与环境无关的描述替换为特定于 benchmark 的描述（轨迹格式、“通过”的含义、场景结构）。`prompts/swebench/` 文件夹是针对单轮代码编辑 benchmark 执行此操作的一个完整示例；`prompts/tau-bench/` 则展示了针对多轮工具使用 benchmark 的相同操作。
2. **精简您的轨迹。** `pipeline/_summarize.py` 提供了通用的压缩功能；如果需要，可针对您 benchmark 的轨迹 schema 进行扩展。
3. **选择阈值。** 默认值（`ρ = 0.10`、`δ = 0.20`、`K = 8 of 10`）取自论文并且在实践中效果良好。仅当您的 benchmark 具有异常结构（失败极少、轨迹极长等）时才需要调整。
4. **运行。** pipeline 的其余部分（能力选择 → 环境生成 → GRPO → 路由）是与 benchmark 无关的。
## 引用
如果您的研究中觉得这项工作有帮助，请考虑引用我们的论文：
```
@misc{kang2026tracecapabilitytargetedagentictraining,
title={TRACE: Capability-Targeted Agentic Training},
author={Hangoo Kang and Tarun Suresh and Jon Saad-Falcon and Azalia Mirhoseini},
year={2026},
eprint={2604.05336},
archivePrefix={arXiv},
primaryClass={cs.AI},
url={https://arxiv.org/abs/2604.05336},
}
```
标签：AI智能体, DLL 劫持, LoRA, Vectored Exception Handling, 人工智能, 凭据扫描, 大语言模型, 强化学习, 模型训练, 用户模式Hook绕过, 自动化代码生成, 逆向工具