ai-trains-ai: RL-Training an AI Agent to Train Other AI Models

ai-trains-ai: RL-Training an AI Agent to Train Other AI Models

ai-trains-ai プロジェクトは、"outer loop" RL エージェントが "inner loop" モデルのトレーニングプロセスを最適化するようにトレーニングされる、入れ子状の強化学習 (RL) アーキテクチャを実装しています。エージェントがトレーニングするモデルの実際のパフォーマンス向上に基づいてエージェントに報酬を与えることで、システムはプロセスの信頼性を向上させ、優れたベースモデルを選択し、ハイパーパラメータを最適化することに成功し、最終的にはトレーニング中に遭遇したことのないタスクファミリーへとこれらのスキルを転移させることができました。

Nested RL Architecture: The Outer and Inner Loops

システムは、再帰的な最適化ループを作成するために、2つの異なる RL トレーニングスタックを使用します。

The Outer Loop (The Trainer)

  • Model: Qwen3.6-35B-A3B (using a LoRA adapter).
  • Objective: 特定のタスクのための、完全で高品質なトレーニングジョブを記述すること。
  • Stack: Tinker (Thinking Machines' managed RL API) と importance-sampling GRPO を使用してトレーニング。
  • Action: エージェントはサンドボックス化されたワークスペースを編集して、トレーニングジョブの環境、報酬ルーブリック、およびハイパーパラメータを定義します。

The Inner Loop (The Trainee)

  • Model: Small base models (Qwen3-0.6B or Qwen3-1.7B).
  • Objective: 特定のタスク (e.g., chained arithmetic or multi-hop queries) を解決すること。
  • Stack: prime-rl (GRPO) を Runpod GPU 上でトレーニング。
  • Feedback: 非公開の評価セットにおける inner model のパフォーマンスが、主要な報酬信号として outer loop にフィードバックされます。

Reward Design and Optimization

トレーナーエージェントを導くために、システムは重み付き報酬合計 (0.35 validation / 0.60 job quality / 0.05 train speed) を使用します。

  • Validation (35%): エージェントが初回試行でパース可能な有効なサブミッションを生成することに対して報酬を与えます。
  • Job Quality (60%): 0.25 * post-training score + 0.75 * uplift over the best untrained baseline というハイブリッドメトリクスとして計算されます。これにより、エージェントが単に強力なベースモデルを選択するのではなく、真の改善をもたらすように報酬が与えられます。
  • Train Speed (5%): 効率性を促進するための、軽微なタイブレーカー。

Key Results and Behavioral Shifts

54 回のトレーニングステップを経て、トレーナーエージェントの報酬は ~0.0 からピークの ~0.63 まで上昇しました。この進展は、2つの異なる段階に分かれていました。

  1. Process Reliability: 初期段階の利得は、エージェントが検証失敗や GPU クラッシュを回避することを学習することによってもたらされました。このフェーズでは、job quality は平坦でしたが、total reward は ~0.26 まで上昇しました。
  2. Model Improvement: 信頼性が飽和した後、エージェントは実際のトレーニング品質を最適化し始めました。hidden-eval post-training スコアは、ほぼゼロのノイズから、持続的な 0.22–0.48 へと上昇しました。

Generalization and Model Selection

  • Zero-Shot Transfer: エージェントは、トレーニング中に遭遇したことのない "triage" タスクファミリー (on-call incident triage) でテストされました。パフォーマンスは上昇し、その後プラトーに達しました。これは、エージェントが RL ジョブを設計するための汎用的なスキルを学習したことを証明しています。
  • Intelligent Selection: エージェントは、ベースラインスコアへのアクセス権を与えられた後、初期のエピソードにおける 77% のケースでより弱い 0.6B モデルを選択していた状態から、エピソードの 95% で 1.7B モデルを選択するように移行しました。
  • Hyperparameter Tuning: エージェントは、[prime_rl] config surface の使用率を 21% から ~78% に増加させ、sampling temperature, optimizer choice, および schedulers を最適化しました。

Infrastructure and Cost Analysis

プロジェクトは、大量の inner-loop ジョブ (合計約 1,750 ジョブ) を扱うために、分散 GPU フリートに利用しました。

  • Compute Fleet: 最大 16 個の warm Runpod GPU pods。トレーニングの大部分は A40 GPU (64%) と RTX 4090 (32%) で行われました。

  • Cost Efficiency: 最もコスト効率の高いセットアップは 2x RTX A5000 がジョブあたり ~$0.13 でしたが、可用性の問題で A40 を使用せざるを得ないことがよくありました。

  • Total Expenditure: メインのトレーニングアークは、Runpod ($810) と Tinker ($465) に分かれ、合計約 $1,275 かかりました。

Project Resources

  • Weights: トレーナーエージェントの LoRA adapter (step-34 checkpoint) は、Hugging Face 上の Danau5tin/ai-trains-ai-trainer で利用可能です。
  • Code: 完全なエージェント・ハーネス、報酬コード、および GPU オーケストレーションは、GitHub の Danau5tin/ai-trains-ai でオープンソース化されています。
  • Tools Used: プロジェクトは、GRPO トレーニングのための prime-rl、環境ルーブリックのための verifiers、および管理された RL インフラストラクチャのための Tinker を多用用しました。

Sources