🤖 AI Training AI: A Meta-RL Experiment

Project Repository: Danau5tin/ai-trains-ai — RL-training an AI agent to RL-train AI agents.

This project explores a recursive approach to machine learning: using Reinforcement Learning (RL) to optimize an agent whose sole purpose is to design RL training pipelines for other models.

Open Source Status: Everything is available, including the trained agent's weights (via a LoRA adapter on Hugging Face), the agent harness, reward logic, task families, GPU orchestration scripts, and detailed write-ups of every pilot—even the ones that failed including the failures.

---

🎯 TL;DR

I developed a pipeline where a "Trainer Agent" is given a goal (e.g., "Teach a model to do X"). The agent then generates a full prime-rl training configuration, encompassing the environment, the reward rubric, the dataset, and hyperparameters.

By utilizing Tinker, I RL-trained the Trainer Agent itself, providing rewards based on the performance of the models it produced. Over 54 training steps, the reward climbed from $\approx 0.0$ to a peak of $\approx 0.63$.

The essence: An AI operating within an RL loop, where its primary action is to initiate another RL loop to train a different AI.

---

🔁 How it Works

The architecture consists of two distinct RL loops running on separate training stacks.

The Dual-Loop Architecture

Anatomy of an Episode

One episode represents a single attempt by the Trainer Agent to create a high-quality training job:

• Task Specification: The agent receives a description (e.g., "teach a small model to resolve multi-hop persona queries"), hard constraints, the eval tool API, and a few development examples.
• Agent Iteration: The agent interacts with a sandboxed environment using tools like read_file, write_file, edit_file, and list_files. It can also call get_baseline_scores to evaluate the current untrained base models.
• Dispatch: Once validated, the job is sent to a pool of warm Runpod GPU pods. These pods execute GRPO training via prime-rl and measure the checkpoint's performance before and after training.
• Reward Calculation: The system calculates a reward based on how efficiently the job was validated and the resulting performance uplift.

Note: Each outer-loop batch triggers 40 inner training jobs across up to 16 GPU pods.

---

⚖️ Reward Design

The episode reward is a weighted sum of three components: $\text{Total Reward} = (0.35 \cdot \text{Validation}) + (0.60 \cdot \text{Job Quality}) + (0.05 \cdot \text{Train Speed})$

1. Validation (0.35):

   • $1.0$ for a first-try valid submission.
   • Decays with each additional attempt.
   • $0$ if it never validates.
   • Penalty: The outer loop applies a $-0.1$ score if no parseable submission is made. (Note: $-1.0$ in CSVs is a logging sentinel for "no post score," not a reward).

2. Job Quality (0.60): A hybrid metric focusing on the absolute score and the improvement over the best untrained baseline ($\text{best\_pre}$): $\text{Quality} = 0.25 \cdot \text{post\_score} + 0.75 \cdot (\text{post\_score} - \text{best\_pre})$

   • If the GPU job crashes, this scores $0$.

3. Train Speed (0.05): A minor tie-breaker to encourage efficiency, provided the job succeeded.

Nuance Alert: The agent's internal prompt (template/INSTRUCTIONS.md) uses a simplified version of this logic (a 75/25 uplift/absolute split and a nudge for fewer attempts) rather than the full $0.35/0.60/0.05$ breakdown. The published adapter was trained against this simplified prompt.

---

🧩 Task Families

The agent was tested on six task families requiring multi-step reasoning and tool use. These were specifically chosen because base models struggle with them.

The agent was trained on the first five families; triage was used exclusively for testing generalization.

---

📈 Results & Infrastructure

Key Findings:

• Reward Growth: The reward climbed in two distinct "rungs," eventually peaking at $0.63$.
• Generalization: The skills learned transferred successfully to the held-out triage family.
• Model Selection: The agent learned to intelligently select the superior base model for the task.

Technical Setup:

• Infrastructure: Runpod GPUs.
• Training Frameworks: Tinker (Outer) and prime-rl (Inner).
• Cost: Approximately -\1.3\text{k}$.

Example Configuration Snippet:

[prime_rl]
model = "qwen3-0.6b"
learning_rate = 1e-6
batch_size = 64
# Agent-generated reward rubric follows...

🚀 Getting Started & Future

The project is fully open-sourced. Future improvements will focus on scaling the task diversity and refining the agent's ability to debug its own training jobs in real-time.