💬 NLP

Synthetic Sandbox for Training Machine Learning Engineering Agents

The paper introduces SandMLE, a multi-agent framework that generates micro-scale synthetic MLE environments to overcome the prohibitive cost of full-pipeline verification, thereby enabling efficient on-policy reinforcement learning that significantly outperforms supervised fine-tuning baselines and generalizes across unseen scaffolds.

Original authors: Yuhang Zhou, Lizhu Zhang, Yifan Wu, Jiayi Liu, Xiangjun Fan, Zhuokai Zhao, Hong Yan

Published 2026-04-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuhang Zhou, Lizhu Zhang, Yifan Wu, Jiayi Liu, Xiangjun Fan, Zhuokai Zhao, Hong Yan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Teaching AI to Cook by Burning the Kitchen Down

Imagine you want to teach a robot chef (an AI agent) how to cook complex meals (Machine Learning tasks).

In the past, for simple tasks like writing a short email or fixing a typo in code, we could test the robot instantly. It writes a sentence, we check it, and say "Good job!" or "Try again." This is fast, like checking a math homework problem.

But Machine Learning Engineering (MLE) is different. It's like asking the robot to grow a new species of tomato from scratch.

The robot has to plan the recipe.
It has to plant the seeds (data).
It has to water them and wait for them to grow (training the model).
Finally, it tastes the tomato to see if it's good (evaluation).

The Bottleneck: In the real world, growing that tomato takes hours or days. If you want to teach the robot using Reinforcement Learning (RL)—a method where the robot learns by trial and error—you need to let it try thousands of times.

The Reality: If one "try" takes 3 hours, you can only do 8 tries a day. That's too slow to learn anything useful.
The Old Solution: Scientists stopped using trial-and-error. Instead, they just showed the robot a few videos of human chefs cooking (Supervised Fine-Tuning). The robot memorized the moves but didn't truly learn how to adapt when things went wrong.

The Solution: SandMLE (The "Toy Kitchen" Sandbox)

The authors realized the problem wasn't the robot; it was the size of the garden. Real datasets have millions of data points (like millions of tomato seeds).

SandMLE is a new framework that builds a Micro-Scale Sandbox.

The Analogy: The "Toy Car" vs. The "Real Truck"

Imagine you want to teach a child how to drive a semi-truck.

The Old Way: You put them in a real semi-truck on a highway. It takes 20 minutes just to get the engine started. You can only practice once a day.
The SandMLE Way: You build a perfectly realistic toy version of the truck on a small table.
- The toy truck has the same steering wheel, the same gears, and the same physics.
- But instead of a 10-ton load, it carries a tiny 1-ounce weight.
- Instead of driving 100 miles, it drives 10 feet.
- The Result: The child can practice driving, crashing, and correcting their steering thousands of times in an hour. They learn the principles of driving perfectly. When they finally get in the real truck, they already know how to steer.

How SandMLE Works (The "Factory")

The paper describes a team of AI agents working together to build these "Toy Kitchens" automatically. Think of them as a construction crew:

The Data Strategist (The Architect): Looks at a real, complex problem (like predicting stock prices) and strips away the boring details. It says, "Okay, the structure is: Input A + Input B = Result C. Let's make a tiny version of this."
The MLE Developer (The Builder): Writes code to generate a tiny dataset (only 50 to 200 items instead of millions). It creates a "hidden rule" (e.g., "If the sky is blue, the tomato is red") that the robot must discover.
The MLOps Engineer (The Referee): Builds a quick scoring system. Because the dataset is tiny, the referee can grade the robot's work in 15 seconds instead of 200 seconds.
The Technical Writer (The Storyteller): Writes the instructions so the robot thinks it's solving a real, serious problem, even though the data is tiny.

The Results: Why It Matters

Because the "Toy Kitchen" is so fast, the researchers could finally use On-Policy Reinforcement Learning (the "Trial and Error" method) for Machine Learning tasks for the first time.

Speed: They made the process 13 times faster.
Learning: The robots trained in these sandboxes didn't just memorize; they learned how to think and adapt.
Generalization: When they took these trained robots and put them in real environments (with real, huge datasets) and even gave them different "tools" (different software frameworks), the robots still performed incredibly well.
- Analogy: It's like the child who learned on the toy car didn't just learn to turn the wheel; they learned the concept of driving. So, when they got a motorcycle, a boat, or a real truck, they could drive them all.

The Secret Sauce: "Milestone Rewards"

One of the biggest challenges in teaching AI is that the reward is often too far away.

Bad Reward: "You failed. The tomato is burnt." (This happens after 3 hours).
SandMLE Reward: "Good job! You planted the seeds. Good job! You watered them. Good job! The sprout is green."

The paper introduces a Dense Reward System. Instead of waiting for the final grade, the AI gets small "points" for every correct step (formatting the code, running the script, getting a valid result). This keeps the AI motivated and guides it step-by-step toward the solution, preventing it from getting lost in the dark.

Summary

SandMLE is a breakthrough because it realized that to teach AI complex engineering skills, you don't need a massive, slow, real-world dataset. You need a fast, tiny, synthetic sandbox that preserves the logic of the problem but shrinks the size of the data.

By shrinking the data, they unlocked the power of trial-and-error learning, allowing AI agents to become true Machine Learning Engineers rather than just mimics.

1. Problem Statement

The paper addresses a critical bottleneck in training Large Language Model (LLM) agents for Machine Learning Engineering (MLE) tasks using Reinforcement Learning (RL).

The Latency Bottleneck: While RL with Verifiable Rewards (RLVR) has successfully improved LLM reasoning in Software Engineering (SWE) and web search, applying it to MLE is prohibitively slow. In SWE, verification involves fast unit tests (seconds). In MLE, verification requires running full ML pipelines (data preprocessing, model training, inference, and metric evaluation) on large datasets.
The Cost: A single code execution step in standard MLE benchmarks (like MLE-bench) averages ~200 seconds due to large dataset sizes (millions of samples).
The Consequence: This latency makes on-policy trajectory-wise RL (e.g., Group Relative Policy Optimization or GRPO) computationally infeasible. Existing approaches resort to Supervised Fine-Tuning (SFT) or offline proxy rewards, which sacrifice the exploration and generalization benefits of on-policy learning.

2. Methodology: SandMLE

The authors propose SandMLE, a framework that enables scalable on-policy RL for MLE by generating synthetic, micro-scale environments. The core insight is that execution latency is driven primarily by dataset size, not the complexity of the logic itself.

A. Synthetic Environment Generation Pipeline

SandMLE uses a multi-agent system to procedurally generate diverse MLE tasks from a small set of seed tasks (60 seeds from MLE-bench). The pipeline consists of four specialized agents:

Data Strategist: Extracts the "Structural DNA" (mathematical schema, modality, label distribution) from a seed task, mutates the domain (e.g., changing "animal classification" to "road damage detection"), and injects adversarial noise to increase difficulty. Crucially, it constrains the dataset size to 50–200 samples.
MLE Developer: Procedurally generates the raw data assets (images, tabular data, etc.) and implements the "hidden rules" (ground truth logic) to ensure the task is learnable. It also trains baseline models to establish performance thresholds.
MLOps Engineer: Constructs a deterministic evaluation sandbox. It hard-codes evaluation metrics and establishes progressive milestone thresholds (e.g., Bronze, Silver, Gold scores) based on baseline performance.
Technical Writer: Synthesizes the final task description, ensuring the prompt aligns perfectly with the generated data and evaluation code.

B. Verification and Sanity Checks

Before RL training, the system performs automated sanity checks:

Logical Consistency: Ensures milestone thresholds are strictly monotonic (e.g., Gold > Silver > Bronze) and that a dummy submission yields a score below the highest threshold.
Execution Verification: Runs the generation and evaluation scripts to ensure they execute without errors.

C. Training Paradigm: Trajectory-wise GRPO

SandMLE applies Group Relative Policy Optimization (GRPO) to the synthetic environments.

Dense Reward Formulation: To combat reward sparsity in long-horizon tasks, the authors design a dense reward function:
$r = w_{format} \cdot r_{format} + w_{execute} \cdot I_{execute} + \sum w_{si} \cdot I_{si}$
This includes rewards for:
- Format: Correct use of reasoning tags (<thought>).
- Execution: Successfully generating a valid submission file.
- Milestones: Achieving specific performance tiers (Median, Bronze, Silver, Gold) defined by the synthetic thresholds.
Selective Masking: During backpropagation, the loss is masked for environment observations, updating only the agent's generated action tokens.

3. Key Contributions

Scalable Synthetic MLE Sandbox: A novel multi-agent framework that generates diverse, verifiable MLE tasks with micro-scale datasets (50–200 samples), reducing execution time by 13× (from ~200s to ~14s per step).
First On-Policy RL for MLE: Successfully demonstrates that trajectory-wise RL is feasible in the MLE domain, enabling thousands of rollout updates where previously only a handful were possible.
Framework-Agnostic Generalization: Proves that models trained via SandMLE generalize robustly across unseen agent scaffolds (e.g., ReAct, AIDE, AIRA) and benchmarks (MLE-bench-lite, MLE-Dojo), unlike SFT baselines which often fail outside their specific training scaffold.

4. Experimental Results

The authors evaluated the approach on Qwen3 models (8B, 14B, and 30B-A3B) against baselines (Base, Seed-SFT, and closed-source models like Claude-4.5-Sonnet).

Performance Gains:
- SandMLE achieved 20.3% to 66.9% relative improvement in "Any Medal" rate over SFT baselines across model sizes.
- The 30B SandMLE model achieved a 100% Valid Submission rate and a 27.3% Any Medal rate, closing the gap with top closed-source models (Claude-4.5-Sonnet at 31.8%).
Generalization:
- Models trained with SandMLE maintained high performance when tested on different agent frameworks (AIDE, AIRA) and the MLE-Dojo benchmark.
- In contrast, SFT models suffered significant performance drops (e.g., Valid Submission dropping from 83.9% to 17.7% on MLE-Dojo when switching scaffolds).
Ablation Studies:
- Reward Design: Removing the dense, milestone-based rewards and using only sparse rewards (format + final gold score) caused performance to collapse, confirming the necessity of intermediate feedback for long-horizon tasks.
- Scaling: Larger models (30B) learned to produce valid outputs purely through RL, whereas smaller models (8B) benefited from SFT initialization for stability.

5. Significance and Conclusion

SandMLE fundamentally shifts the paradigm for training MLE agents. By identifying dataset size as the primary bottleneck and solving it through synthetic micro-scale environments, the authors make on-policy RL practical for complex engineering tasks.

Efficiency: Reduces training time by over 13×, enabling large-scale exploration.
Quality: The synthetic tasks preserve the structural complexity and mathematical rules of real-world problems while discarding the computational overhead of massive data.
Future Impact: This approach suggests that synthetic, verifiable sandboxes can serve as effective proxies for real-world training, allowing agents to learn complex trial-and-error strategies through direct environment interaction rather than mere imitation of expert behavior. This opens a scalable path toward autonomous AI researchers and engineers.