ai-trains-ai: 使用 RL 訓練一個 AI Agent 來訓練其他 AI 模型

ai-trains-ai: 使用 RL 訓練一個 AI Agent 來訓練其他 AI 模型

ai-trains-ai 專案實現了一種嵌套式強化學習 (RL) 架構，其中「外部迴圈」RL agent 被訓練用來優化「內部迴圈」模型的訓練過程。透過根據 agent 訓練的模型之實際性能提升來給予獎勵，該系統成功學會了提高流程可靠性、選擇更優質的基礎模型，並優化超參數，最終將這些技能轉移到訓練期間從未遇到過的任務家族。

嵌套式 RL 架構：外部與內部迴圈

該系統使用兩個不同的 RL 訓練堆疊來建立遞迴優化迴圈。

外部迴圈 (訓練器)

  • 模型： Qwen3.6-35B-A3B (使用 LoRA adapter)。
  • 目標： 為特定任務編寫完整且高品質的訓練任務 (job)。
  • 堆疊： 使用 Tinker (Thinking Machines 的託管式 RL API) 搭配 importance-sampling GRPO 進行訓練。
  • 動作： Agent 會編輯一個沙盒工作區，以定義訓練任務的環境、獎勵準則 (reward rubric) 和超參數。

內部迴圈 (受訓者)

  • 模型： 小型基礎模型 (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 開始優化實際的訓練品質。隱藏評估集的 post-training 分數從接近零的雜訊水平上升到持續的 0.22–0.48。

泛化能力與模型選擇

  • 零樣本轉移 (Zero-Shot Transfer)： Agent 在一個在訓練期間從未見過的「分流 (triage)」任務家族 (值班事故分流) 上進行了測試。性能表現先上升後趨於平穩，證明了 agent 學會了設計 RL 任務的通用技能。
  • 智慧選擇： 在獲得基礎模型分數的權限後，agent 從早期階段在 77% 的次數中選擇較弱的 0.6B 模型，轉變為在 95% 的次數中選擇 1.7B 模型。
  • 超參數調優： Agent 增加了對 [prime_rl] 配置介面的使用率，從 21% 增加到 ~78%，優化了採樣溫度 (sampling temperature)、優化器選擇以及調度器 (schedulers)。

基礎設施與成本分析

該專案利用分散式 GPU 集群來處理大量的內部迴圈任務 (總計約 1,750 個任務)。

  • 運算集群： 一個包含多達 16 個熱啟動 Runpod GPU pods 的池。大多數訓練發生在 A40 GPUs (64%) 和 RTX 4090s (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 取得。
  • Code： 完整的 agent harness、獎勵代碼與 GPU 編排工具皆已在 GitHub 上以 Danau5tin/ai-trains-ai 開源。
  • 工具使用： 該專案大量利用了 prime-rl 進行 GRPO 訓練、verifiers 用於環境準則，以及 Tinker 用於託管式 RL 基礎設施。

Sources