ai-trains-ai: RL-training an AI agent to train other AI agents

The ai-trains-ai project demonstrates a nested Reinforcement Learning (RL) architecture where an AI agent is trained to optimize the training process of other AI models. By rewarding the agent based on the performance of the models it produces, the system creates an automated loop of model improvement.

The Nested RL Architecture

This system employs two distinct RL loops with separate training stacks to automate the training of small base models:

The Outer Loop (The Trainer Agent)

  • Model: Qwen3.6-35B-A3B (using a LoRA adapter).
  • Training Method: Importance-sampling GRPO via the Tinker API.
  • Objective: The agent is trained to write complete training jobs—including the environment, reward functions, datasets, and hyperparameters—and submit them for execution.

The Inner Loop (The Target Model)

  • Model: Small base models (Qwen3-0.6B or Qwen3-1.7B).
  • Training Method: GRPO via the prime-rl framework on Runpod GPUs.
  • Objective: The inner model is trained based on the the job specifications written by the outer loop agent.

Episode Workflow and Reward Design

Each episode consists of a single attempt by the trainer agent to produce a high-quality training job for a specific task. The process follows these steps:

  1. Task Specification: The agent receives a description of the task (e.g., "teach a model to resolve multi-hop persona queries"), constraints, and developer examples.
  2. Agent Execution: The agent uses a sandboxed workspace to edit files and check baseline scores of untrained models.
  3. Job Submission: The agent submits the job, which undergoes a validation probe. If it fails, the agent is allowed a capped number of retries.
  4. Dispatch and Training: Validated jobs are queued and executed on a warm pool of Runpod GPU pods. The prime-rl framework fine-tunes the small model.
  5. Reward Calculation: The final reward is a weighted sum of validation efficiency (35%), job quality (60%), and training speed (5%). Job quality is further split between the absolute post-training score (25%) and the uplift over the best untrained baseline (75%).

Performance Results and Generalization

Over 54 training steps, the trainer agent's reward climbed from approximately 0.0 to a peak of 0.63. The learning process occurred in two distinct "rungs":

  • Rung 1: Process Reliability. The agent first learned to avoid validation failures and GPU crashes, increasing total reward to ~0.26 while job quality remained flat.
  • Rung 2: Model Improvement. The agent transitioned from producing working jobs to producing better models, with post-training scores rising from near-zero noise to a sustained 0.22–0.48 range.

Generalization to Unseen Tasks

To test transfer learning, the agent was tested on a "triage" task family (on-call incident triage) that it never encountered during training. The mean reward on this holdout set rose from 0.399 (base model) to a peak of 0.545 at step 34, proving the agent learned a general skill for training models rather than just memorizing specific tasks.

Model and Hyperparameter Selection

The agent also learned to optimize its choices. Initially, 77% of episodes chose the weaker 0.6B model. After receiving access to baseline scores and uplift grading, the agent shifted its preference to the 1.7B model in 95% of episodes. It also began utilizing the [prime_rl] configuration surface to tune sampling temperature, optimizers, and schedulers.

Infrastructure and Costs

The system utilizes a distributed infrastructure to handle the high compute requirements:

  • Compute: A warm pool of up to 16 Runpod GPU pods (primarily A40s and RTX 4090s) handles the inner loop training.
  • Orchestration: The entire process is managed on a CPU box via Nebius.
  • Cost Efficiency: The most cost-effective GPU for short tasks was found to be the RTX A5000 (~$0.13 per job), though availability often forced the use of A40s.

Total Project Cost: The headline training arc cost approximately $1,275 (roughly $810 for Runpod and $465 for Tinker).

Technical Insights and Community Discussion

While the project demonstrates successful automation of the RL training pipeline, community members raised questions regarding the potential for reward-hacking. Specifically, there is a risk that an agent could write training data that leaks the hidden evaluation distribution rather than teaching a general skill. The project provides open-sourced weights for the trainer agent (a LoRA adapter for Qwen3.6-35B-A3B) and a full harness to allow for further verification and reproduction.

Sources

Related

  • Project
  • Dispatch
  • Dispatch
  • Project
  • Project