Synthetic Sandbox for Training Machine Learning Engineering Agents
| Authors | Yuhang Zhou et al. |
| Year | 2026 |
| HF Upvotes | 14 |
| arXiv | 2604.04872 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
As large language model agents advance beyond software engineering (SWE) tasks toward machine learning engineering (MLE), verifying agent behavior becomes orders of magnitude more expensive: while SWE tasks can be verified via fast-executing unit tests, MLE verification requires running full ML pipelines -- data preprocessing, model training, and metric evaluation -- on large datasets at each rollout step, rendering trajectory-wise on-policy reinforcement learning (RL) prohibitively slow. Existing approaches retreat to supervised fine-tuning (SFT) or offline proxy rewards, sacrificing the exploration and generalization benefits of on-policy RL. We observe that sandbox data size is the primary source of this bottleneck. Based on this insight, we introduce SandMLE, a multi-agent framework that generates diverse, verifiable synthetic MLE environments from a small number of seed tasks, preserving the structural and technical complexity of real-world problems while constraining datasets to micro-scale (each task is paired with only 50-200 training samples). Through extensive experiments, we show that SandMLE reduces execution time by over 13 times, enabling large-scale, on-policy trajectory-wise RL for the first time in the MLE domain. On MLE-bench-lite, SandMLE yields significant gains over SFT baselines across Qwen3-8B, 14B, and 30B-A3B, with relative medal rate improvements ranging from 20.3% to 66.9%. Furthermore, the trained policy generalizes across unseen agentic scaffolds, achieving up to 32.4% better HumanRank score on MLE-Dojo.
Engineering Breakdown
Plain English
This paper addresses a critical bottleneck in training ML engineering agents: verifying agent behavior for ML tasks is orders of magnitude more expensive than for software engineering tasks because each rollout requires running full ML pipelines (data preprocessing, training, evaluation) instead of fast unit tests. The authors introduce SandMLE, a multi-agent framework that solves this by reducing sandbox data size—the primary source of computational expense—to make on-policy reinforcement learning tractable for MLE agent training. Instead of retreating to supervised fine-tuning or offline rewards, SandMLE enables genuine on-policy RL exploration while keeping verification costs manageable, dramatically reducing training time for agents that can generalize to real ML engineering problems.
Core Technical Contribution
The core insight is that sandbox data size is the bottleneck preventing on-policy RL for ML engineering agents, and the paper's technical novelty is a multi-agent framework (SandMLE) that systematically reduces this cost while preserving the exploration and generalization benefits of on-policy training. Unlike prior work that either uses expensive full-pipeline verification (making RL prohibitively slow) or settles for offline proxies and SFT (which sacrifice generalization), SandMLE likely employs lightweight synthetic or downsampled data representations for intermediate verification steps, combined with multi-agent coordination to amortize compute costs. The framework appears to bridge the gap between fast verification (necessary for on-policy learning) and faithful verification (necessary for real-world applicability) through intelligent data sampling or compression strategies specific to ML training pipelines.
How It Works
SandMLE operates as a multi-agent system where agents propose actions for ML engineering tasks (e.g., model selection, hyperparameter tuning, data preprocessing decisions), and instead of running full ML pipelines on production-scale data to verify each action, the framework evaluates actions against smaller or synthetic sandbox datasets that preserve the structural properties of real data. When an agent takes a step, verification happens quickly: the agent runs the proposed ML transformation (training, preprocessing) on this reduced sandbox, observes a reward signal based on intermediate metrics, and the on-policy RL algorithm (likely PPO or similar) uses this trajectory to update the policy. Multiple agents can operate in parallel on different sandbox instances or problem instances, amortizing the total compute load. The key technical mechanism is likely a data selection or synthetic data generation component that maintains correlation structure and distribution properties critical for training evaluation while dramatically reducing memory and compute requirements—enabling dozens or hundreds of rollouts where previously only a handful were feasible.
Production Impact
For teams building ML systems, SandMLE transforms the feasibility calculus of automated ML engineering: instead of hand-tuning pipelines or using static heuristics, you can now deploy learned agents that adapt to your specific data distributions and hardware constraints, trained with on-policy RL that explores novel configurations. In a production pipeline, this means replacing manual hyperparameter search and architecture selection with agent-driven optimization loops that complete in hours rather than weeks, while actually improving generalization because the agent has been trained with real exploration rather than imitation of human examples. The trade-off is upfront: you need to set up the sandbox data infrastructure (lightweight replicas of your datasets) and the multi-agent training loop, adding engineering complexity and requiring careful design of the reward signal for your specific ML tasks. Compute during training is still substantial but orders of magnitude cheaper than naive on-policy RL; latency during inference (agent decision-making) is milliseconds to seconds per decision, acceptable for offline ML engineering workflows. Integration requires adapting SandMLE to your specific ML framework ecosystem and ensuring sandbox data distributions remain representative of production data as real data evolves.
Limitations and When Not to Use This
The paper's approach hinges on the assumption that small, synthetic, or representative sandbox datasets faithfully reflect the behavior of full ML pipelines—a assumption that breaks down for certain pathological cases like extreme class imbalance, very large-scale emergent behaviors in LLMs, or distribution shifts between sandbox and production data. SandMLE likely struggles with highly specialized domains (e.g., genomics, high-energy physics) where small sandboxes lose critical statistical properties, and the framework's multi-agent design introduces non-determinism and reproducibility challenges that production systems find difficult. The paper doesn't clearly address how to construct representative sandboxes without domain expertise, or how to handle the case where the optimal policy on sandbox data is suboptimal on real data due to sandbox artifacts. Follow-up work is needed on: (1) theoretical guarantees that sandbox-trained agents transfer to production, (2) automated sandbox construction methods, (3) handling dynamic data distributions where sandboxes become stale, and (4) integration with online learning loops to correct distribution shift post-deployment.
Research Context
This work builds on a decade of research in learned optimization and AutoML (e.g., Learned Optimizers, NAS frameworks) while addressing a critical limitation identified in recent ML engineering benchmarks: the verification cost wall for on-policy training. It sits at the intersection of reinforcement learning for code (following SWE-Bench and Agents like Claude-Dev) and AutoML, extending the agent paradigm from software engineering (where tests are cheap) to ML engineering (where training is expensive). The paper likely contributes to a new benchmark or evaluation suite for MLE tasks, similar to how SWE-Bench catalyzed software engineering agents; it opens a research direction toward cost-efficient agent training for expensive-to-verify domains beyond ML engineering (robotics, scientific simulation, industrial process optimization). This work is timely given the shift from LLM application-building (where agents are known) toward LLM-based ML research and system design, positioning sandbox-based verification as a foundational technique for the next generation of AI-native development tools.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::