PaperlayerView exact revision on arXiv

Source-linked AI summary

Synthetic Sandbox for Training Machine Learning Engineering Agents

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

arXiv:2604.04872v1cs.CLcs.LG

TL;DR

MLE trajectory-wise on-policy RL is bottlenecked by expensive full-pipeline execution and has often been replaced by SFT or offline proxy rewards. SandMLE generates diverse, verifiable synthetic environments with 50–200-sample datasets and combines them with dense milestone rewards for trajectory-wise RL. It cuts execution time by over 13×, improves Any Medal rates over Seed-SFT by 20.3% to 66.9%, and generalizes across unseen scaffolds with up to 32.4% better HumanRank.

  • Problem

    Full MLE-pipeline execution makes trajectory-wise on-policy RL prohibitively slow, leading prior approaches toward SFT or offline proxy rewards.

  • Method

    SandMLE uses multi-agent generation of diverse, verifiable synthetic MLE environments with 50–200-sample datasets and applies trajectory-wise RL with dense milestone rewards.

  • Results

    20.3% to 66.9% relative improvement in Any Medal rate over Seed-SFT is achieved across Qwen3 models, with up to 32.4% better HumanRank on MLE-Dojo.

  • Takeaways & Limitations

    SandMLE makes large-scale on-policy trajectory-wise RL practical for MLE while preserving structural complexity in synthetic tasks.

Abstract

from arXiv · show

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.

1 Introduction

MLE agents require long-horizon interaction with costly full-pipeline execution, making trajectory-wise on-policy RL impractical. SandMLE addresses this bottleneck with verifiable synthetic micro-scale environments and applies trajectory-wise RL, improving benchmark performance and cross-scaffold generalization.

  • MLE tasks require iterative coding, execution, feedback analysis, and strategy adjustment rather than single-turn responses.
  • Nearly 200 seconds per MLE code execution makes large-scale trajectory-wise on-policy RL impractical.Each step runs preprocessing, training, and metric evaluation on associated data.
  • SandMLE constrains datasets to 50–200 samples, reducing average execution time to under 15 seconds while preserving structural and mathematical complexity.The framework uses multi-agent collaboration to generate diverse, verifiable synthetic tasks from seed tasks.
  • SandMLE orchestrates Data Strategist, MLE Developer, MLOps Engineer, and Technical Writer roles to generate synthetic environments and task specifications.
  • Trajectory-wise RL with milestone-based rewards enables models to learn long-horizon trial-and-error strategies in MLE.
  • 20.3% to 66.9% relative improvement in Any Medal rate over Seed-SFT is achieved across Qwen3-8B, 14B, and 30B-A3B on MLE-bench-lite, with up to 32.4% HumanRank improvement on MLE-Dojo.

2 Related Work

Prior MLE-agent work improves performance through sophisticated scaffolds and test-time iteration, while trajectory-wise RL has succeeded in other long-horizon domains. SandMLE targets MLE’s execution-cost barrier and preserves strictly on-policy training using fast synthetic environments.

  • MLE benchmarks such as MLE-bench, MLE-Dojo, and MLE-Smith expand standardized evaluation volume and task diversity.
  • Existing MLE systems use multi-turn feedback, evolutionary procedures, and test-time compute scaling to iteratively refine agent trajectories.
  • Trajectory-wise RL has improved long-horizon software-engineering and web-search agents, but extending it to MLE introduces severe computational cost.
  • SFT and asynchronous step-wise RL reduce latency but remain off-policy, whereas SandMLE uses fast synthetic environments to maintain strictly on-policy training.

3 Preliminaries

The paper models MLE agents as finite-horizon sequential decision processes optimized with GRPO from verifiable rewards. Multi-turn trajectories create reward-sparsity and execution-cost challenges, which SandMLE addresses through micro-scale environments and dense milestone rewards.

  • 3.1 MLE as a Sequential Decision Process: An MLE task is represented by initial specification I, tool set T, and interactive execution environment E.
  • 3.1 MLE as a Sequential Decision Process: At each step, the policy observes the trajectory history, generates an action such as code, and receives environment observations including errors or metrics.
  • 3.1 MLE as a Sequential Decision Process: Completed trajectories are evaluated on a hidden test set to produce a scalar score S, with termination by submission or a step limit.
  • 3.1 MLE as a Sequential Decision Process: Environment latency makes on-policy optimization prohibitively expensive because every rollout step executes the MLE environment for every sampled group member.
  • 3.2 Reinforcement Learning from Verifiable Rewards: GRPO samples groups of candidate outputs, normalizes each verifiable reward relative to its group, and penalizes divergence from a reference policy.
  • 3.2 Reinforcement Learning from Verifiable Rewards: For agentic MLE, GRPO treats each output as a multi-turn trajectory, computes gradients over action tokens, and masks environment observations from the loss.
  • 3.2 Reinforcement Learning from Verifiable Rewards: SandMLE reduces execution cost by over 13× and uses dense milestone-based rewards to address the two practical challenges of trajectory-level GRPO.

4 SandMLE

SandMLE constructs verifiable synthetic MLE environments from seed tasks while avoiding the evaluation corruption and task scarcity of simple dataset downsampling. Its micro-scale environments support trajectory-wise RL through automated generation, sanity verification, and dense rewards.

  • 4.1 Synthetic Environment Generation: SandMLE generates lightweight synthetic environments from curated seed tasks instead of downsampling fixed real-world datasets.The pipeline extracts task structure, changes domains, defines hidden rules and noise, and procedurally creates small datasets with aligned evaluation environments.
  • 4.1 Synthetic Environment Generation: The Data Strategist abstracts each seed task into Task DNA, maps it to new scenarios, adds adversarial noise, and compiles a hidden rule linking features to labels.The hidden rule is represented as H : l = f(z) + ϵ, while baseline methods establish progressive milestone thresholds for dense rewards.
  • 4.1 Synthetic Environment Generation: The ML Developer and MLOps Engineer create deterministic data and scoring scripts, lock milestone thresholds, and test evaluators with dummy submissions.The generated scripts partition data, apply hidden rules, evaluate baselines, standardize metrics, and iteratively debug runtime errors.
  • 4.2 Environment Sanity Verification: A synthetic task is retained only when its milestone thresholds and dummy-sample baseline satisfy the metric-specific strict monotonic ordering.Tasks violating the ordering are discarded so the RL curriculum contains valid monotonic reward signals.
  • 4.3 Trajectory-Level GRPO: Micro-scale environments transform trajectory-level GRPO from computationally infeasible to practical for MLE agents.Agents interact through ReAct over multiple turns, receiving execution outputs, errors, and intermediate metrics while using a dense reward combining format and milestone components.

5 Experiments

Experiments evaluate SandMLE across synthetic-task diversity, execution efficiency, model scales, baselines, and training configurations. The results show that micro-scale synthetic environments support faster trajectory-wise RL and improve performance, reliability, and generalization.

  • Experimental Setup: 848 synthetic training tasks and 64 held-out tasks are generated from 60 MLE-bench seed questions spanning diverse reasoning patterns.The evaluation uses Valid Submission, Above Median, Bronze, Silver, Gold, and Any Medal success tiers.
  • Synthetic Training Data: The corpus covers multiple domains, modalities, and formulations, including healthcare, image-based tasks, tabular tasks, classification, regression, ranking, forecasting, and reconstruction.Classification accounts for 56.2% of formulations, while image-based tasks account for 48.7% of modalities.
  • Synthetic Training Data: 92.9% win rate for Claude-4.5-Sonnet and lower rates for DeepSeek-V3, Gemini-2.5-Flash, and GPT-4o-mini show that the 64 synthetic tasks separate models by capability.The reported rates are 39.9%, 35.6%, and 25.5%, respectively, using pairwise wins with ties counted as 0.5 wins.
  • Execution Efficiency: 14.31 seconds versus 196.17 seconds reduces average execution time by over 13× on synthetic versus original MLE-bench tasks.The micro-scale corpus mostly contains 120–150 samples per task, compared with approximately 4.09 million samples per original seed task.
  • Main Results: +66.9%, +24.7%, and +100.7% relative Any Medal improvements over Base are achieved by SandMLE at 8B, 14B, and 30B, respectively.Against Seed-SFT, SandMLE improves Any Medal rate relatively by 20.3% to 66.9%; the 8B model matches the 22.7% rate of DeepSeek-V3.1 and Gemini-2.5-Flash, while 14B and 30B reach 27.3% versus Claude-4.5-Sonnet at 31.8%.
  • Main Results: SFT-SandMLE improves operational reliability while preserving medal performance, raising 8B Valid Submission from 63.6% to 90.9% and achieving 95.5% Valid Submission at 14B.Pure SandMLE Valid Submission scales from 63.6% at 8B to 77.3% at 14B and 100% at 30B, while Above Median rises from 27.3% to 36.4%.

6 Analysis

SandMLE gains transfer across unseen agentic scaffolds, while dense milestone-based rewards outperform sparse feedback on long-horizon MLE tasks.

  • 6.1 Framework Generalization: SandMLE consistently exceeds Base performance across all evaluated scaffold-benchmark combinations.Evaluations include AIDE, AIRA, and MLE-Agent settings.
  • 6.1 Framework Generalization: 83.9% Valid Submission rate and 38.56 HumanRank score are achieved by Qwen3-30B-SandMLE with MLE-Agent on MLE-Dojo.This configuration outperforms both Base and SFT variants.
  • 6.1 Framework Generalization: 17.7% Valid Submission rate is reached by the 30B SFT model with MLE-Agent on MLE-Dojo, illustrating brittleness outside its training scaffold.SandMLE models adapt robustly in the same cross-scaffold evaluation.
  • 6.2 Effectiveness of Milestone-Based Rewards: 27.3% Any Medal rate with milestone rewards collapses to 13.6% under sparse rewards for Qwen3-30B.The ablation compares the proposed dense reward with a sparse formulation relying on limited feedback.
  • 6.2 Effectiveness of Milestone-Based Rewards: 36.4% Above Median rate with milestone rewards halves to 18.2% under sparse rewards for Qwen3-30B.The sparse condition also reduces valid submission rate from 100.0% to 86.4%.
  • 6.2 Effectiveness of Milestone-Based Rewards: Hierarchical feedback validates format, execution, and progressive performance tiers to stabilize policy optimization.Sparse rewards provide insufficient gradients for complex, long-horizon exploration.

7 Conclusion

SandMLE makes trajectory-wise on-policy RL practical for MLE agents by generating verifiable synthetic environments with micro-scale datasets. Combined with dense milestone rewards, it improves medal rates across model scales and generalizes across unseen agentic scaffolds.

  • 7 Conclusion: SandMLE generates diverse, verifiable synthetic MLE environments with 50–200-sample datasets for trajectory-wise on-policy RL.The framework reduces per-step execution time by over 13× while preserving structural complexity for policy optimization.
  • 7 Conclusion: 20.3% to 66.9% relative medal-rate improvement over SFT is achieved across model scales on MLE-bench-lite.The conclusion attributes this result to SandMLE combined with dense, milestone-based rewards.
  • 7 Conclusion: SandMLE generalizes robustly across unseen agentic scaffolds.The paper presents synthetic micro-scale environments as proxies for real-world MLE tasks.

A.1 Metric Definition Details

The evaluation combines submission validity, threshold-based medal rates, and relative leaderboard rankings, with separate metric definitions for MLE-bench-lite and MLE-Dojo.

  • MLE-bench-lite Metrics: Valid Submission Rate measures tasks producing a properly formatted, parseable submission.csv without runtime or schema errors.It is the basic completion-validity metric for MLE-bench-lite.
  • MLE-bench-lite Metrics: Above Median measures valid submissions scoring above the original human leaderboard’s 50th percentile for a competition.The threshold is strictly better than the median performance.
  • MLE-bench-lite Metrics: Bronze, Silver, and Gold measure submissions meeting or exceeding the corresponding Kaggle medal thresholds.These are discrete performance thresholds established for each competition.
  • MLE-bench-lite Metrics: Any Medal is the union of Bronze, Silver, and Gold successes and serves as the primary MLE-bench-lite metric.It counts tasks achieving at least Bronze-level performance.
  • MLE-Dojo Metrics: HumanRank Score measures relative agent performance by ranking submissions against historical competition participants.Public and private leaderboard position scores are computed independently and averaged.
  • Training and Evaluation Setup: GRPO training samples trajectories with group size n = 4 and uses 100 steps, learning rate 1 × 10^-6, and batch size 16.Each task uses one NVIDIA H200 GPU with a 90-second rollout execution limit.

B.1 Training Dynamics

GRPO training produces upward reward trends and stabilizing validation rewards across Qwen3-8B, 14B, and 30B. Larger models show higher and more stable rewards, with the 30B model exhibiting the strongest submission-rate stability.

  • B.1 Training Dynamics: 80 GRPO steps track valid submission rate, training reward, and validation reward for Qwen3-8B, 14B, and 30B.The models occupy the top, middle, and bottom rows respectively in Figure 8.
  • B.1 Training Dynamics: All three models show upward reward trends and stabilizing final validation rewards without severe overfitting.The validation curves indicate continuous policy improvement during training.
  • B.1 Training Dynamics: 0.67 peak training reward for the 30B model exceeds the 14B model’s approximately 0.62 peak.The 30B model’s rewards are also more stable, reflecting higher performance ceilings with increased capacity.
  • B.1 Training Dynamics: The 8B model’s valid submission rate fluctuates between 0.1 and 0.8, while the 14B model reaches 1.0 intermittently.The 30B model shows the strongest submission-rate stability.

C.2 Qualitative Analysis

SandMLE generates a compact but technically structured synthetic MLE task by mutating a real seed problem, enforcing deterministic hidden rules, and aligning evaluation with class imbalance. The example demonstrates a coherent, verifiable environment for trajectory-wise reinforcement learning.

  • Task Derivation and Mutation Strategy: SandMLE transforms a 23-class animal-image seed task into motion-blurred road-surface damage classification while preserving its semantic goal and taxonomy.The mutation shifts difficulty toward blur-invariant feature extraction and temporal reasoning in a smart-city vehicle-camera setting.
  • Data Generation and Hidden Rules: 147 synthetic RGB images at 1024×768 resolution satisfy the micro-scale constraint for rapid reinforcement-learning rollouts.The generated images are divided into training and testing sets and use Perlin-noise road textures with overlaid damage types.
  • Data Generation and Hidden Rules: Motion blur increases linearly across 1–5 temporal frames through a directional line-kernel that simulates a moving vehicle.The sequence-based rule explicitly encodes progressively stronger blur as frame number increases.
  • Data Generation and Hidden Rules: Deterministic labels use edge ratio, contrast, frequency, and texture interactions, with blur deliberately degrading minor damage into the intact class.This creates structured label noise while preserving rule-based ground truth generation.
  • Evaluation and Alignment: Macro-F1 and blur-adjusted medal thresholds address extreme class imbalance and reward performance on rare damage classes.The example specifies Gold = 0.68 and Silver = 0.55 thresholds.
  • Evaluation and Alignment: The resulting environment combines domain consistency, complex mathematical rules, and deterministic verification as a testbed for trajectory-wise reinforcement learning.The qualitative example is presented as evidence that SandMLE can produce rigorous synthetic MLE environments.

D Prompt Details

The appendix specifies ReAct execution constraints and prompt templates for SandMLE’s specialized agents. Together, the prompts enforce practical runtime and data-handling requirements while guiding scenario design, mutation, task specification, and structured outputs.

  • D.1 React framework: The ReAct framework uses system and user prompt templates to structure agent interaction.The appendix identifies these templates in Table 5.
  • D.1 React framework: Agents must use GPU acceleration when available, respect interpreter and turn limits, inspect dataset files, and produce a correctly formatted submission.The requirements also mandate common preinstalled libraries, silent progress behavior, and ignoring outdated competition timelines.
  • D.2 SandMLE Method: The SandMLE method provides prompt templates for its specialized generation agents.Tables 6–9 cover the Data Strategist, MLE Developer, MLOps Engineer, and Technical Writer.
  • D.2 SandMLE Method: The Data Strategist brainstorms five industry scenarios whose domains must justify the structural DNA’s modalities and feature relationships.Examples include contract comparison for paired categoricals and product duplicate detection for category and text features.
  • D.2 SandMLE Method: The Data Simulation Engineer detects modality, selects three corresponding difficulty mutagens, and returns only a JSON mutation configuration.The prompt gives blur, noise, and typos as modality-specific examples.
  • D.2 SandMLE Method: The Synthetic Data Architect merges abstract data DNA, target domain, and noise configuration into a concrete task specification.It renames features, sets dimensions, and embeds a hidden ground-truth function incorporating requested noise.
  • D.2 SandMLE Method: The task specification is returned as JSON with task metadata, dimensions, feature mappings, hidden-rule logic, metric, and medal-threshold logic.The schema supports different modalities and requires valid JSON output.
  • D.2 SandMLE Method: Synthetic tasks must contain 50–200 samples, preserve structural consistency in feature mappings, and explicitly use the renamed features in hidden-rule logic.These constraints tie generation speed to semantic and logical consistency.
Loading 2604.04872v1…