ai-trains-ai: 通过强化学习训练一个 AI Agent 来训练其他 AI 模型

ai-trains-ai: 通过强化学习训练一个 AI Agent 来训练其他 AI 模型

ai-trains-ai 项目实现了一种嵌套强化学习 (RL) 架构,其中“外环” RL Agent 被训练用于优化“内环”模型的训练过程。通过根据模型实际性能的提升来奖励 Agent,该系统成功学会了提高流程可靠性、选择更优的基础模型并优化超参数,并最终将这些技能迁移到训练期间从未遇到过的任务族中。

嵌套 RL 架构:外环与内环

该系统使用两个不同的 RL 训练栈来创建一个递归优化循环。

外环 (训练器/The Trainer)

  • 模型: Qwen3.6-35B-A3B (使用 LoRA adapter)。
  • 目标: 为特定任务编写完整、高质量的训练任务 (job)。
  • 技术栈: 使用 Tinker (Thinking Machines 的托管 RL API) 配合重要性采样 GRPO 进行训练。
  • 动作: Agent 编辑沙盒工作区以定义训练任务的环境、奖励准则 (reward rubric) 和超参数。

内环 (受训者/The Trainee)

  • 模型: 小型基础模型 (Qwen3-0.6B 或 Qwen3-1.7B)。
  • 目标: 解决特定任务 (例如,链式算术或多跳查询)。
  • 技术栈: 使用 prime-rl (GRPO) 在 Runpod GPU 上进行训练。
  • 反馈: 内环模型在隐藏评估集上的性能表现将作为主要奖励信号反馈给外环。

奖励设计与优化

为了引导训练器 Agent,系统使用了加权奖励总和 (0.35 验证 / 0.60 任务质量 / 0.05 训练速度)。

  • 验证 (35%): 奖励 Agent 在第一次尝试时就生成可解析、有效的提交内容。
  • 任务质量 (60%): 一个混合指标,计算方式为 0.25 * post-training score + 0.75 * uplift over the best untrained baseline。这确保了 Agent 是因为真正的性能提升而获得奖励,而不仅仅是选择一个强大的基础模型。
  • 训练速度 (5%): 一个微小的平局决胜项 (tie-breaker),用以鼓励效率。

关键结果与行为转变

在 54 个训练步长中,训练器 Agent 的奖励从 ~0.0 攀升至 ~0.63 的峰值。这一进展分为两个明显的阶段:

  1. 流程可靠性: 初期的提升是由 Agent 学会避免验证失败和 GPU 崩溃驱动的。在此阶段,任务质量保持平稳,而总奖励上升至 ~0.26。
  2. 模型改进: 一旦可靠性达到饱和,Agent 开始优化实际的训练质量。隐藏评估集的训练后得分从接近零的噪声水平上升到持续的 0.22–0.48。

泛化能力与模型选择

  • 零样本迁移: Agent 在一个“分诊”任务族 (on-call incident triage) 上进行了测试,该任务在训练期间从未出现过。性能有所提升后趋于平稳,证明了 Agent 学会了设计 RL 任务的通用技能。
  • 智能选择: 在获得基准分数访问权限后,Agent 的行为从早期阶段 77% 的情况选择较弱的 0.6B 模型,转变为在 95% 的情况选择 1.7B 模型。
  • 超参数微调: Agent 对 [prime_rl] 配置表面的使用率从 21% 增加到 ~78%,优化了采样温度、优化器选择和调度器。

基础设施与成本分析

该项目利用分布式 GPU 集群来处理大量的内环任务 (总计约 1,750 个任务)。

  • 计算集群: 最多包含 16 个热启动 Runpod GPU pod 的池。大部分训练发生在 A40 GPU (64%) 和 RTX 4090 (32%) 上。
  • 成本效率: 经确定,最具成本效益的配置是 2x RTX A5000,每个任务约 $0.13,尽管由于可用性问题,通常不得不使用 A40s。
  • 总支出: 核心训练阶段的成本约为 $1,275,由 Runpod ($810) 和 Tinker ($465) 分摊。

项目资源

  • 权重: 训练器 Agent 的 LoRA adapter (step-34 checkpoint) 可在 Hugging Face 上通过 Danau5tin/ai-trains-ai-trainer 获取。
  • 代码: 完整的 Agent 框架、奖励代码和 GPU 编排工具已在 GitHub 上开源,地址为 Danau5tin/ai-trains-ai
  • 工具使用: 该项目大量利用了 prime-rl 进行 GRPO 训练,使用 verifiers 进行环境准则定义,并使用 Tinker 进行托管 RL 基础设施。

Sources