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 "outer loop" RL agent is trained to optimize the training process of "inner loop" models. By rewarding the trainer agent based on the performance uplift of the models it produces, the system successfully learned to improve model quality, select superior base models, and generalize its training skills to entirely new task families.

Nested RL Architecture: Outer and Inner Loops

The system operates using two distinct RL stacks to separate the trainer's learning process from the training jobs it generates.

The Outer Loop (The Trainer)

• Model: Qwen3.6-35B-A3B with a LoRA adapter.
• Objective: Learn to write high-quality training jobs (environments, rewards, datasets, and hyperparameters).
• Training Stack: Tinker using importance-sampling GRPO.
• Reward Signal: Derived from the inner loop's success and the resulting model's performance on a hidden evaluation set.

The Inner Loop (The Target Model)

• Model: Small base models (Qwen3-0.6B or Qwen3-1.7B).
• Objective: Solve specific tasks defined by the trainer agent.
• Training Stack: prime-rl using GRPO, executed on Runpod GPU pods.

The Training Episode Workflow

Each episode represents a single attempt by the trainer agent to produce a valid training job for a given task specification. The process follows these steps:

1. Task Specification: The agent receives a description of the goal (e.g., "teach a model to resolve multi-hop persona queries"), constraints, and evaluation interfaces.
2. Agent Execution: The agent uses a sandboxed workspace with tools to read, write, and edit files, and can 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 GPUs. The prime-rl framework fine-tunes the small model and scores it pre- and post-training on a hidden evaluation set.
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 absolute post-training score (25%) and the uplift over the best untrained baseline (75%).

Key Results and Performance Gains

Over 54 training steps, the trainer agent's reward climbed from approximately 0.0 to a peak of 0.63. This progress occurred in two distinct stages:

Stage 1: Process Reliability

Initial gains were driven by the agent learning to avoid validation failures and GPU crashes. During this phase, the agent focused on producing "working" jobs rather than "better" models, with rewards rising to ~0.26 while job quality remained flat.

Stage 2: Model Optimization

Once reliability saturated, the agent began optimizing for model performance. The hidden-eval post-training scores rose from near-zero noise to a sustained range of 0.22–0.48, indicating the agent had learned how to actually improve the target models.

Generalization and Model Selection

• Held-out Task Transfer: The agent was tested on a "triage" task family it had never seen during training. Performance rose and then plateaued, proving the trainer's skills transferred to new domains.
• Base Model Selection: The agent learned to prefer the stronger 1.7B model over the 0.6B model. The share of episodes selecting the 1.7B model increased from 42% to 95%.
• Hyperparameter Tuning: The agent began utilizing the [prime_rl] config surface more effectively, increasing its usage from 21% to ~78% of episodes to tune temperature, optimizers, and schedulers.

Infrastructure and Cost Analysis

The project utilized a distributed infrastructure to manage the high compute demands of the inner loop.

• GPU Fleet: A warm pool of up to 16 Runpod pods was used. Benchmarking revealed that 2x RTX A5000s were the most cost-effective at ~$0.13 per job, though A40s were used for 64% of the headline arc due to availability.
• Compute Orchestration: The outer loop was managed via Tinker, while the overall orchestration ran on a Nebius CPU box.
• Total Cost: The headline training arc cost approximately $1,275 (~£950), split between Runpod ($810) and Tinker ($465).

Technical Specifications and Open Source Availability

The project is fully open-sourced, including the agent harness, reward code, and GPU orchestration scripts. The trained trainer agent is available as a LoRA adapter on Hugging Face at Danau5tin/ai-trains-ai-trainer (derived from Qwen3.6-35B-A3B, released under Apache-2.0).

Task Families Used for Training

Family	Description	Untrained Baseline (best_pre)
calc_chain	Chained arithmetic steps	0.742
multi_hop	Dependent persona-world lookups	0.654
string_pipeline	Composed string transformations	0.545
ledger	Stateful account bookkeeping	0.242
dispatch	Conditional routing decisions	0.323
triage	Incident triage (Held-out)	0.352