PaperlayerView exact revision on arXiv

Source-linked AI summary

MLE-Dojo: Interactive Environments for Empowering LLM Agents in Machine Learning Engineering

Rushi Qiang, Yuchen Zhuang, Yinghao Li, Dingu Sagar V K, Rongzhi Zhang, Changhao Li, Ian Shu-Hei Wong, Sherry Yang, Percy Liang, Chao Zhang, Bo Dai

arXiv:2505.07782v1cs.LG

TL;DR

Real-world MLE involves iterative workflows, but existing benchmarks often lack comprehensive task coverage and interactive, executable training environments. MLE-Dojo combines a broad Kaggle-based benchmark with interactive agent training and evaluation; its evaluations establish baselines and identify limitations in current models and architectures.

  • Problem

    Many MLE benchmarks focus on isolated or non-interactive tasks, while broader ones lack environments for iterative experimentation and agent training.

    Source: §1 · p. 2

  • Method

    MLE-Dojo provides an interactive, executable Gym-style environment built on 200+ Kaggle competitions, with training trajectories for supervised fine-tuning and reinforcement learning.

    Source: §1 · p. 2

  • Results

    Evaluations establish foundational baselines and highlight capabilities and critical limitations of current state-of-the-art models and agent architectures.

    Source: §6 · p. 13

  • Takeaways & Limitations

    The released framework, benchmarks, and leaderboard support transparent comparison, reproducible research, and community-driven progress.

    Source: §6 · p. 13

    Comprehensive training or evaluation across hundreds of tasks requires substantial computational and storage resources.

    Source: §A.1 · p. 16

Abstract

from arXiv · show

We introduce MLE-Dojo, a Gym-style framework for systematically reinforcement learning, evaluating, and improving autonomous large language model (LLM) agents in iterative machine learning engineering (MLE) workflows. Unlike existing benchmarks that primarily rely on static datasets or single-attempt evaluations, MLE-Dojo provides an interactive environment enabling agents to iteratively experiment, debug, and refine solutions through structured feedback loops. Built upon 200+ real-world Kaggle challenges, MLE-Dojo covers diverse, open-ended MLE tasks carefully curated to reflect realistic engineering scenarios such as data processing, architecture search, hyperparameter tuning, and code debugging. Its fully executable environment supports comprehensive agent training via both supervised fine-tuning and reinforcement learning, facilitating iterative experimentation, realistic data sampling, and real-time outcome verification. Extensive evaluations of eight frontier LLMs reveal that while current models achieve meaningful iterative improvements, they still exhibit significant limitations in autonomously generating long-horizon solutions and efficiently resolving complex errors. Furthermore, MLE-Dojo's flexible and extensible architecture seamlessly integrates diverse data sources, tools, and evaluation protocols, uniquely enabling model-based agent tuning and promoting interoperability, scalability, and reproducibility. We open-source our framework and benchmarks to foster community-driven innovation towards next-generation MLE agents.

1 Introduction

MLE-Dojo addresses the lack of comprehensive benchmarks and interactive, executable training environments for the iterative work of real-world MLE. It combines a large Kaggle-based task collection with agent training and evaluation features.

1 Introduction

  • The paper evaluates eight frontier LLMs across 50 evaluation tasks, with results made publicly accessible through a maintained leaderboard.The leaderboard is intended to support transparent comparison, reproducibility, and collaborative research.Source: Figure 1 · p. 1, §1 · p. 2
  • Existing benchmarks often isolate tasks or use narrow, non-interactive settings, leaving iterative experimentation, debugging, feedback, and training insufficiently supported.Broader benchmarks offer more diverse tasks but still lack interactive environments for agent fine-tuning or reinforcement learning.Source: §1 · p. 2
  • MLE-Dojo is an interactive, executable Gym-style environment built on more than 200 Kaggle competitions across machine-learning domains.Its 150-task training dataset supports trajectory sampling for supervised fine-tuning and reinforcement learning.Source: §1 · p. 2
  • Pre-installed dependencies, standardized evaluation scripts, and real-time outcome verification support iterative experimentation and debugging.The framework also supports model-agnostic agent tuning and integration of diverse datasets and tools.Source: §1 · p. 2

2 Related Works

Related work spans coding agents, MLE agents, MLE benchmarks, and interactive agent environments. MLE-Dojo is positioned as combining broad MLE task coverage with interactive training and evaluation.

X-Agent (Coding)

  • MLE agents have expanded from automating individual pipeline steps to end-to-end workflows, including iterative execution, debugging, and search-based optimization.Examples include CAAFE [15], DS-Agent [13], AutoKaggle [24], DataInterpreter [16], SELA [4], and AIDE [19].Source: §2 · p. 3

X-Bench

  • Earlier benchmarks targeted isolated actions, while broader Kaggle-based benchmarks expanded task coverage but still lacked robust interactive training environments.MLE-Bench [2] and DSBench cover 75 and 74 competitions, respectively.Source: §2 · p. 3

X-Gym (Dojo)

  • MLE-Dojo’s task diversity spans time series, computer vision, tabular data, and natural language processing.These domains are presented as representative examples of the benchmark’s task diversity.Source: Figure 2 · p. 4
  • ML-Gym integrates diverse ML research tasks into an interactive gym supporting reinforcement learning and iterative experimentation; MLE-Dojo extends task quantity and complexity with 200+ competitions and a dedicated 150-task training set.The training set enables trajectory sampling for supervised fine-tuning and reinforcement learning.Source: §2 · p. 4

3 Benchmark Tasks and Dataset Construction

MLE-Dojo assembles diverse practical ML tasks from Kaggle and other benchmark sources, then standardizes them for agent use. Its construction balances task coverage, validation, and training needs.

3.1 Task Description

  • The benchmark contains more than 200 tasks across ML domains and 15 task types, mostly practical challenges sourced from real-world Kaggle competitions.Examples include image classification and sentiment analysis.Source: §3.1 · p. 4

3.2 Dataset Construction

  • Tasks were selected to provide diverse training and evaluation resources while remaining neither beginner-level nor too difficult and being completed by many human competitors.The selection aimed for tasks that were representative, relatively easy to validate, and still tested expected agent capabilities.Source: §3.2 · pp. 4–5
  • After excluding seven unavailable, oversized, or package-dependent MLE-Bench tasks and removing duplicates, the collection combines sources into more than 200 unique tasks.Sources include 68 MLE-Bench competitions [2], 74 DSBench competitions, and 75 additionally prepared Kaggle competitions.Source: §3.2 · p. 5
  • Each standardized task includes a competition description, reorganized train and test data, a local evaluator, and a competition leaderboard snapshot.The interface and documentation also support adding tasks or datasets in the same structured format.Source: §3.2 · p. 5
  • MLE-Dojo divides tasks into 150 training and 50 evaluation examples, prioritizing prior-benchmark tasks in evaluation while adding competitions for diversity.The larger training set supports interactive experiences for supervised fine-tuning and reinforcement learning.Source: §3.2 · p. 5

4 MLE-Dojo

MLE-Dojo connects agents to task-specific MLE environments through a structured interaction loop. Its modular, extensible design supports varied tasks, rich feedback, executable actions, and a normalized leaderboard-based reward.

4.1 Overview

  • Each interaction is modeled as a POMDP in which the agent receives observations and rewards, then chooses an action based on prior interaction history.The task environment supplies the observation and reward; the next action depends on the agent and environment design.Source: §4.1 · p. 5

4.2 Modular and User-Friendly Interface

  • Decoupled modules can be integrated or extended through a unified registration mechanism, while users can also use the pre-designed environment through a single env.step(action_type, action_args) call.The modular design supports custom modules, new functionality, reproducibility, and rapid prototyping.Source: §4.2 · p. 6

4.3 Extensible Task Space P

  • Tasks run in separate Docker containers with configurable sandbox limits, while a standardized competition format supports plug-and-play integration of new tasks.The format includes competition information, structured datasets, and public and private leaderboards.Source: §4.3 · p. 7

4.4 Observation Space O

  • The observation space combines task information, evaluation scores and feedback, code execution outcomes, errors, and interaction history.History is available as both agent-side conversation history and environment-side records, accommodating different agent architectures.Source: §4.4 · p. 7

4.5 Expandable Action Space A

  • Agents request task information, write and validate code, execute submissions for evaluation, retrieve history, and reset; researchers can register additional actions.Validation supports syntax and runtime checks, whereas execution performs submission verification and evaluation.Source: §4.5 · pp. 7–8

4.6 Reward Space and Environmental Feedback

  • HumanRank rewards use relative leaderboard position to provide fine-grained feedback across competitions with different metrics and score ranges.The score is normalized to [0, 1] and is intended to reward incremental improvements rather than coarse-grained medals.Source: §4.6 · pp. 8–9

5 LLMs as MLE Agents in MLE-Dojo

The evaluation compares eight LLM agents across four ML tasks using complementary performance and robustness metrics. Results show different performance trajectories, action strategies, error patterns, and solution lengths across models.

5.1 Experiment Setups

  • Eight backbone LLMs are evaluated, with non-reasoning models run at temperature 0.0 and top-p 1.0 for reproducibility.The best performance of two runs per task per model is used.Source: §5.1 · p. 9
  • For a comprehensive evaluation, the study considers Area Under the Performance Profile (AUP), HumanRank Score (H-Rank, %), and Elo ranking together as metrics.Agents have 15 steps and full access to their interaction histories; sessions last at most 12 hours, GPU memory is limited to 32 GB, and submission attempts are unrestricted.Source: §5.1 · p. 9

5.2 Main Results

  • Reasoning and coding models including o3-mini, DeepSeek-r1, and Gemini-2.5-Pro consistently rank highly across the evaluated metrics and four ML tasks.Figure 8 presents performance profiles and corresponding AUP curves for robustness across the tasks.Source: §5.2 · p. 10, Figure 8 · p. 10
  • Computer-vision tasks are hardest: none of their average HumanRank scores exceed 60, and more than half are below 30.MLE-Lite task averages mostly exceed 30, while other domains have more evenly distributed difficulty.Source: §5.2 · p. 10, Figure 7 · p. 10

5.3 Cost Analysis

  • Reasoning models typically cost more because of premium pricing and longer outputs, while DeepSeek-r1 achieves competitive performance with significantly fewer tokens.Computer-vision and deep-neural-network training tasks generate longer code than classical CPU-based tabular tasks.Source: §5.3 · pp. 10–11, Figure 9 · p. 11

5.4 Step-wise Performance Dynamics

  • o3-mini reaches high performance within the first five steps and remains stable, while DeepSeek-r1 and Gemini-2.5-Pro improve gradually into later steps.Non-reasoning models may outperform reasoning models early or midway, but generally improve less and finish with lower scores.Source: §5.4 · p. 11, Figure 10 · p. 11

5.5 Action and Error Analysis

  • o3-mini executes code in over 90% of instances, whereas gpt-4o and gpt-4o-mini execute about 20% of the time and validate more.Stronger models more often encounter execution failures, while conservative validators reduce those failures but still face submission creation and validity issues.Source: §5.5 · pp. 11–12, Figure 11 · p. 12

5.6 History Length and Solution Length

  • Chat-history length closely aligns with best-solution length, and both positively correlate with overall model performance, though longer solutions do not guarantee better outcomes.Reasoning models and stronger-performing models typically generate longer solutions.Source: §5.6 · p. 12, Figure 14 · p. 12

6 Conclusion

The conclusion presents MLE-Dojo as an interactive framework for training, evaluating, and benchmarking autonomous MLE agents.

6 Conclusion

  • MLE-Dojo uses real-world Kaggle competitions to support systematic experimentation, outcome verification, and iterative feedback in LLM-driven MLE workflows.The authors report that evaluations established baselines and highlighted capabilities and limitations of current models and agent architectures.Source: §6 · p. 13
  • The framework, benchmarks, and leaderboard are publicly released to encourage transparent comparison, reproducible research, and community-driven progress.Future work includes adding domain-specific deep research and multi-agent collaboration.Source: §6 · p. 13

A.1 Limitations

The appendix notes resource and licensing responsibilities for MLE-Dojo, alongside considerations about societal impact, fairness, and responsible use.

A.1 Limitations

  • Comprehensively training or evaluating across hundreds of competition tasks requires substantial computational and storage resources.Running full MLE pipelines may require API credits, CPU/GPU resources, and large-scale disk storage.Source: §A.1 · p. 16
  • Competition datasets have differing licenses and privacy policies, and users are responsible for reviewing and complying with each dataset's terms.The framework accesses datasets through the official Kaggle API and does not redistribute proprietary data directly.Source: §A.1 · p. 16

A.2 Broader Impacts

  • The framework may lower engineering barriers for practitioners and support domain experts applying machine learning, while its open-source design promotes reproducibility and transparency.The stated potential beneficiaries include healthcare, education, and scientific research.Source: §A.2 · p. 16

A.3 Ethical Statements

  • MLE-Dojo accesses publicly available competition data through Kaggle and states that it retains no personal or sensitive information.Users must follow each competition's data-handling terms; selected tasks are reviewed where applicable for personally identifiable or ethically sensitive data.Source: §A.3 · p. 16
  • Tasks are drawn from open competitions and may not represent all application domains or demographic contexts.The authors advise critical evaluation of agent outputs in sensitive settings because agents may inherit and amplify biases.Source: §A.3 · p. 16
  • The authors commit to maintaining MLE-Dojo as an open-source project under a permissive license.They welcome contributions such as fairness metrics, safety guardrails, and responsible-use documentation.Source: §A.3 · p. 16

B Disclaimer

  • The dataset is provided without warranties, and its providers disclaim liability for damages arising from its use.Users are encouraged to verify accuracy, completeness, and currency and bear responsibility for conclusions derived from the data.Source: B Disclaimer · p. 17

C Unified Data Structure

MLE-Dojo supports integrating personalized competitions from diverse sources, using a unified format and a modular directory structure.

C Unified Data Structure

  • New tasks from diverse sources can be standardized into a unified competition format and integrated for evaluation, benchmarking, and training.The authors describe this integration as requiring minimal overhead.Source: C Unified Data Structure · p. 17
  • Each competition uses data/, utils/, and an optional info/ directory to standardize task data, preparation and evaluation scripts, and metadata.The data directory separates private and public phases; utils/ contains task-specific preparation and metric scripts.Source: C Unified Data Structure · p. 17

D Data Details

This section details MLE-Dojo’s curated Kaggle competitions, their preparation and evaluation workflows, and the datasets and evaluation competitions across four task categories. It also presents a HumanRank-based difficulty ranking as a reference, noting potential bias from dataset splits.

D Data Details

  • The first release provides 200 competitions selected from around 600 reviewed tasks for suitability and a balance of challenge complexity with computational feasibility.The review considered clarity of descriptions, data availability, and relevance of evaluation metrics, while excluding excessively large or unwieldy datasets.Source: D Data Details · pp. 17–18
  • Competition-specific scripts structure datasets, generate representative submissions, split training data into train and test subsets, and enable local testing with tailored metrics derived from a common base class.The split addresses the lack of publicly available test labels in most Kaggle competitions.Source: D Data Details · p. 17
  • Each competition uses a standardized directory structure with data, descriptions, sample submissions, and preparation and metric utilities.Optional files include leaderboard data, and the structure can include raw descriptions and data-structure documentation.Source: D Data Details · p. 18, Figure 15
  • Table 3 inventories the released competitions in ascending data-size order and tags tasks by metric, category, modality, or domain to support flexible selection.The paper describes progressive releases of processed datasets to support ongoing benchmarking.Source: D Data Details · p. 18, Table 3 · pp. 18–22
  • The catalog pairs dataset sizes with domain or modality labels and evaluation metrics across geospatial research, audio tagging, and physics/NLP challenges, including smartphone-decimeter-2022 (22.9 GB), Freesound (26.15 GB), and kaggle-llm-science-exam (364.21 kB).Source: D Data Details · pp. 23, 25
  • The evaluation benchmark uses 50 competitions spanning MLE-Lite, Tabular, NLP, and CV tasks.Table 4 lists the evaluation competitions.Source: D Data Details · p. 24, Table 4 · p. 24
  • Competition difficulty is ranked by average HumanRank across models, from easiest to hardest, but the authors call it a reference rather than a definitive measure.They note that dataset splits may bias the ranking; examples include AI4Code at 0.025269 and santander-customer-transaction-prediction at 0.000000.Source: D Data Details · pp. 26–27, Table 5 · p. 26

E.1 Evaluation Metrics Details

E.1 defines performance-profile and AUP metrics for comparing methods across benchmark tasks, then describes adjustments for higher-is-better metrics and invalid scores. It also outlines HumanRank and Elo-based comparisons.

E.1 Evaluation Metrics Details

  • AUP aggregates task-level performance profiles into a single measure of each method’s overall effectiveness across the benchmark.The profile tracks the proportion of tasks where a method performs within a specified log-distance of the best method; AUP integrates that curve.Source: §E.1 · p. 27
  • For higher-is-better metrics such as Accuracy or R2, the performance-ratio calculation is inverted; invalid scores are penalized relative to the lowest valid score, with ratios capped at 100.The invalid-score ratio is set to (1 + ε) times the lowest valid-score ratio, with ε = 1.0.Source: §E.1 · p. 27
  • HumanRank measures a submission’s relative leaderboard position, averaging its independently calculated public- and private-leaderboard scores.The position score is s = 1 − p/N, where p is the submission’s rank among N submissions.Source: §E.1 · p. 28
  • Elo ratings compare models through competition-level win, loss, or tie outcomes, using logistic regression and bootstrap resampling to estimate ratings and confidence intervals.The Bradley-Terry model assumes evaluated LLMs are static; reported ratings use the median across 100 bootstrap rounds.Source: §E.1 · p. 28

F.1 Additional Analysis of Evaluation Metrics

The section explains how complementary metrics characterize LLM performance, then describes the prompts and examples used to examine agents’ iterative work and solution diversity.

F.1 Additional Analysis of Evaluation Metrics

  • HumanRank scores normalize task-specific metrics against human performance, but equal task weighting can hide task-level strengths and weaknesses.Scores range from 0 to 1 and support comparisons across tasks with heterogeneous raw metrics.Source: §F.1 · p. 28
  • Elo rankings can diverge from absolute averages because pairwise win-loss records do not measure the magnitude of performance gaps.The metric summarizes competitive ordering across task-level pairwise comparisons.Source: §F.1 · pp. 28–29
  • Performance Profiles and AUP scores summarize performance ratios across thresholds, capturing robustness while raw-score scale differences complicate cross-task comparison.The paper describes these metrics as reflecting adaptability across performance thresholds.Source: §F.1 · p. 29

G Prompt Details

  • MLE-Agent prompts combine initial task instructions with feedback-driven Error and Reflection prompts and a Parse Error prompt for invalid formats.The initial instruction covers the setting, actions, and expected output; feedback prompts guide subsequent actions.Source: G Prompt Details · p. 29

H Code Analysis

  • The analysis examines how models’ solutions vary across competitions and how different models approach the same competition.The paper introduces these analyses as examples of model behavior during interaction with competitions.Source: H Code Analysis · p. 31

H.1 Solution diversity of models

  • On the ciphertext challenge, Gemini-2.5-Pro used CPU-based sklearn and Logistic Regression, whereas DeepSeek-r1 used GPU-based PyTorch and a deep-learning approach.The examples illustrate differences in packages, devices, and modeling strategies for the same task.Source: §H.1 · pp. 31–32

H.2 Solution diversity of competitions

  • Gemini-2.5-Pro used markedly different approaches for the ciphertext and aerial-cactus competitions.The paper uses Code 1 and Code 3 to illustrate task-specific solutions by the same model.Source: §H.1 · p. 32, §H.2 · pp. 33–34

H.3 Statistics of Method Diversity

The section reports diverse, task-dependent model choices and describes MLE-Dojo’s lightweight agent scaffold and its integration with AIDE.

H.3 Statistics of Method Diversity

  • Model choices vary across LLMs on identical tasks and correlate with task type: Tree and Sim-LN are common for tabular data, while neural-network types are favored for image and NLP tasks.The analyzed categories include classic, advanced, and self-defined neural networks.Source: §H.3 · p. 34
  • Some LLMs more often propose custom architectures, default to simpler models, or yield N/A, indicating variability in generating relevant, complete modeling code.The distribution also reflects differing strategies and potential biases across LLMs.Source: §H.3 · p. 35

I.1 MLE Agent

  • The MLE Agent is a lightweight benchmarking scaffold that gives initial task guidance, then iterates through structured actions and feedback with a fixed-length history.Its instructions cover task description, output format, objectives, available actions, and resource constraints without additional hints.Source: §I.1 · p. 35

I.2 AIDE

  • MLE-Dojo integrates AIDE’s generator–evaluator–selector loop, requiring a correctly formatted submission.csv and selecting candidates by raw score and HumanRank rather than model validation.The integration adapts AIDE’s prompts and feedback while retaining its core logic.Source: §I.2 · pp. 35–36

Related papers

Loading 2505.07782v1…