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) 아키텍처를 구현합니다. 에이전트가 학습시키는 모델의 실제 성능 향상에 따라 보상을 제공함으로써, 시스템은 프로세스 신뢰성을 개선하고, 우수한 베이스 모델을 선택하며, 하이퍼파라미터를 최적화하는 방법을 성공적으로 학습했으며, 결과적으로 학습 중에 접해보지 못한 작업 유형(task families)으로 이러한 기술을 전이할 수 있었습니다.

Nested RL Architecture: The Outer and Inner Loops

시스템은 재귀적 최적화 루프를 생성하기 위해 두 개의 별도 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: 에이전트는 샌드박스된 워크스페이스를 편집하여 학습 작업에 대한 환경, 보상 루브릭(reward rubric), 하이퍼파라미터를 정의합니다.

The Inner Loop (The Trainee)

  • Model: Small base models (Qwen3-0.6B or Qwen3-1.7B).
  • Objective: 특정 작업을 해결하는 것 (예: chained arithmetic 또는 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회의 학습 단계(training steps) 동안, 트레이너 에이전트의 보상은 ~0.0에서 최고 ~0.63까지 상승했습니다. 이 진전은 두 가지 뚜렷렷한 단계로 나타났습니다:

  1. Process Reliability: 초기 이득은 에이전트가 검증 실패 및 GPU 크래시를 피하는 법을 배우면서 발생했습니다. 이 단계 동안 job quality는 평탄하게 유지되었으나 total reward는 ~0.26까지 상승했습니다.
  2. Model Improvement: 신뢰성이 포화 상태에 도달한 후, 에이전트는 실제 학습 품질을 최적화하기 시작했습니다. hidden-eval post-training scores는 거의 0에 가까운 노이즈 수준에서 지속적인 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 작업을 처리하기 위해 분산 GPU 플릿(fleet)을 활용했습니다 (총 약 1,750개의 작업).

  • Compute Fleet: 최대 16개의 warm Runpod GPU pods. 대부분의 학습은 A40 GPU (64%)와 RTX 4090 (32%)에서 수행되었습니다.
  • Cost Efficiency: 가장 비용 효율적인 설정은 작업당 ~$0.13인 2x RTX A5000으로 확인되었으나, 가용성 문제로 인해 A40를 사용하는 경우가 많았습니다.
  • Total Expenditure: 주요 학습 과정의 비용은 약 $1,275였으며, Runpod ($810) 및 Tinker ($465)로 나뉘었습니다.

Project Resources

  • Weights: 트레이너 에이전트의 LoRA adapter (step-34 checkpoint)는 Hugging Face의 Danau5tin/ai-trains-ai-trainer에서 사용할 수 있습니다.
  • Code: 전체 에이전트 하네스(agent harness), 보상 코드, 및 GPU 오케스트레이션은 GitHub의 Danau5tin/ai-trains-ai에서 오픈 소스로 공개되어 있습니다.
  • Tools Used: 프로젝트는 GRPO 학습을 위한 prime-rl, 환경 루브릭을 위한 verifiers, 그리고 관리형 RL 인프라를 위한 Tinker를 적극적으로 활용했습니다.

Sources