Synthetic Sandbox for Training Machine Learning Engineering Agents

Published
Source
arXiv
Paper number
1269
Field
Research Agents
arXiv ID
2604.04872

Key points

  • Problem: An MLE (machine learning engineering) agent must run a full ML pipeline of data preprocessing, model training, and metric evaluation at every rollout step, and a single code execution averages nearly 200 seconds on standard problems. This latency makes trajectory-level on-policy RL far too slow, so prior work retreated to SFT or step-wise RL with offline proxy rewards.
  • Input and output: This is a training method paper with an implementation and experiments. The input is 60 seed tasks taken from MLE-bench, and the output is a set of synthetic MLE tasks whose training data hold only 50 to 200 samples (data, hidden rules, evaluation scripts, task descriptions), together with Qwen3 policy models trained with RL on these tasks.
  • Method: A Data Strategist extracts a Structural DNA (the mathematical structure of the task stripped of domain context) from a seed task, adds a new domain and noise, and defines a hidden rule linking features to labels, while an ML Developer generates the data and trains and evaluates baseline methods to fix milestone thresholds. An MLOps (machine learning operations) Engineer writes an evaluator that scores against a hidden test set, a Technical Writer writes the task description, and only tasks whose thresholds keep the required order stay in training. Training is trajectory-level GRPO with ReAct rollouts, and the dense reward splits weights across format compliance, successful execution, and passing the median, bronze, silver, and gold thresholds.
  • Implementation: Starting from 1,200 synthetic tasks built from the 60 seeds, execution checks and threshold checks leave 912, of which 848 are used for training and 64 for validation. Qwen3-8B, Qwen3-14B, and Qwen3-30B-A3B-2507 are trained with the RLLM framework, with at most 20 turns and a 90-second execution limit per step during training and one NVIDIA H200 GPU per task. The reward weights are 0.1 for format, 0.3 for execution, 0.1 for median, 0.2 for bronze, 0.2 for silver, and 0.1 for gold.
  • Results: With code written by Gemini-2.5-Flash, the average execution time falls from 196.17 seconds on the original MLE-bench tasks to 14.31 seconds on the synthetic tasks. Measured with ReAct on MLE-Bench-Lite, the Any Medal rate rises from 13.6% (Base, and Seed-SFT trained by SFT on Claude-4.5-Sonnet trajectories) to 22.7% for 8B and from 13.6% for Base to 27.3% for 30B, while Claude-4.5-Sonnet reaches 31.8%. On MLE-Dojo, the 30B SandMLE model with the MLE-Agent scaffold reaches a HumanRank score (based on rank among human competitors on the leaderboard) of 38.56, above 29.12 for Base.
  • Limits: The authors report that in a test-time experiment with a growing turn limit, performance peaks at 30 turns and then drops because the context window overflows frequently. The 8B SandMLE model has a valid submission rate of 63.6%, below 68.2% for Base, and on MLE-Bench-Lite with the AIDE scaffold the 30B model's Any Medal rate stays at 13.6%, the same as Base. Evaluation covers MLE-Bench-Lite with 22 questions and MLE-Dojo with 62 tasks.

Paper links

External research summaries. These are not HDATF publications or measured product results.

Read original (opens in a new tab)