ML-Dev-Bench: Comparative Analysis of AI Agents on ML development workflows

In this report, we present ML-Dev-Bench¹¹1https://github.com/ml-dev-bench/ml-dev-bench, a benchmark aimed at testing agentic capabilities on applied Machine Learning development tasks. While existing benchmarks focus on isolated coding tasks or Kaggle-style competitions, ML-Dev-Bench tests agents’ ability to handle the full complexity of ML development workflows. The benchmark assesses performance across critical aspects including dataset handling, model training, improving existing models, debugging, and API integration with popular ML tools. We evaluate three agents - ReAct, Openhands, and AIDE - on a diverse set of 30 tasks, providing insights into their strengths and limitations in handling practical ML development challenges. We open source the benchmark for the benefit of the community at https://github.com/ml-dev-bench/ml-dev-bench.

Recent advances in Large Language Models (LLMs) have demonstrated impressive capabilities in code generation and software engineering tasks. This has led to the development of various benchmarks like HumanEval [2], MBPP [1] that evaluate coding abilities, and others like SWE-Bench [5], that test LLM-based agents on software engineering tasks. However, while these benchmarks effectively assess general programming capabilities, they don’t capture the unique challenges of Machine Learning development workflows,

Benchmarks such as ML-Bench [6], test agents’ abilities to generate code and commands to interact with popular ML repositories, while MLE-Bench [2] and MLAgentBench [4] focus on Kaggle-style tasks to evaluate the iterative and open-ended nature of ML development. However, real-world ML development extends far beyond that, including the complexity of working on top of existing codebases and models, integrating with third-party tools, debugging complex issues that span multiple components of the ML pipeline, and understanding and balancing trade-offs like model performance and cost to come up with optimal design.

ML-Dev-Bench addresses this gap by providing a comprehensive evaluation framework that tests an agent’s ability to handle real-world ML development scenarios. Our benchmark is particularly relevant,, as ML development increasingly relies on large language models and AI agents to assist developers. Understanding the capabilities and limitations of these agents in handling practical ML development tasks is crucial for their effective deployment in production environments.

ML-Dev-Bench comprises of 30 carefully designed tasks that evaluate various aspects of ML development. These tasks are structured to assess both specific technical capabilities (like handling datasets, model implementation) and broader problem-solving skills (like model training and performance improvement) that are essential in real-world ML development. The tasks span several key categories of ML development shown in Table 1:

1. 1.

   Dataset Handling focuses on evaluating the ability to work with large datasets, inspect them and apply pre-processing pipelines. An example is the noisy imagenette [3] dataset download task, where the agent needs to download the dataset, inspect its contents to identify the labels file, only load the 50% noisy labels from it and generate class summary statistics.

2. 2.

   Model Training tests an agent’s ability to work with existing models, from loading pretrained weights to implementing training loops, logging metrics and managing the training process. These tasks assess both technical skills and the ability to handle long-running tasks.

3. 3.

   Debugging presents common scenarios including shape errors, exploding gradients, incorrect implementations, and integration errors. Agents must analyze large training logs, metrics, and code across multiple files to identify and resolve issues.

4. 4.

   Model Implementation tests the ability to modify existing architectures and implement new features. An example is the ChannelViT related tasks, which follow three levels of increasing difficulty: Level 1 provides complete specifications with examples and tests; Level 2 includes specifications and tests but omits examples; Level 3 gives specifications but tests and examples are hidden

5. 5.

   API Integration assesses the ability to work with essential ML development tools, particularly for logging and experiment tracking.

6. 6.

   Model Performance tasks challenge agents to improve baseline implementations through iterative experimentation and hypothesis testing.

In this section we briefly describe the design of our evaluation framework, called Calipers, for running the benchmark. The framework consists of three components: agents, evaluation tasks, and metrics. Agents are evaluated on various Machine Learning tasks to generate metrics. We designed Calipers to allow easy addition of new evaluation tasks, agents and metrics, ensuring the benchmark can evolve alongside advances in ML development practices and tooling.

We evaluate three agents on ML-Dev-bench. The agents and their setup is described below. Each agent uses an LLM and a set of tools to execute various actions. All agents execute their code in a runtime environment which is either a local python or docker environment depending on the agent. We customized the runtime environments for all agents to pre-install common ML frameworks like scikit-learn, pytorch, transformers, lightning, wandb, etc to ensure smooth execution.

1. 1.

   ReAct: We created a simple ReAct agent [8] as a baseline which takes actions by calling tools. We used the LangGraph framework for the agent and Composio toolset which provides tools for common use cases. We customized the tools to reliably capture large command outputs, handle long running commands and ensure consistency across different tools like file and shell tools. All the tool calls were executed in a local python environment which was pre-installed with common ML frameworks as mentioned earlier and had access to the relevant api keys. No custom prompts were used, and the agent was allowed to run for a maximum of 50 steps. We tested the agent with Claude Sonnet 3.5 10-2022 and OpenAI GPT-4o.

   1. (a)

      Command line tools

      1. i.

         Shell Tool - to execute short running commands

      2. ii.

         Spawn Tool - to execute long running commands like training in the background

      3. iii.

         Sleep and execute tool - to wait and monitor long running processes

   2. (b)

      File tools like create files, list files and edit files

2. 2.

   Openhands: Openhands [7] is a popular open-source coding agent with state-of-the-art performance on SWE-Bench-Full [5]. We used Openhands agent v0.21.1 and customized the runtime build to install common ML frameworks listed above. We tested the agent with Claude Sonnet 3.5 10-2022 model which is the current best performing model with the agent on SWE Bench and Gemini 2.0 Flash. The agent was allowed to run for a maximum of 50 steps.

3. 3.

   AIDE: AIDE is an agent purpose-built for data science tasks like Kaggle competitions [2] and performs a tree search over solutions. AIDE scaffolding performs better in comparison to other agents like Openhands on MLEBench using o1, GPT-4o. Unlike other general purpose agents which output any artifact, AIDE outputs an evaluation metric and code as its final output. All other artifacts are considered intermediate outputs and saved in a custom working directory. Since not all tasks in ML-Dev-Bench require outputting a score, we access the artifacts from its custom working directory to validate the agent’s performance. Given the high costs of o1, we evaluated the agent with GPT-4o.

Performance of the agents across different task categories, Table 2 and individual tasks Table 3 reveals a consistent pattern. Performance decreases as tasks become more open-ended and complex. The success rates are highest in well-defined categories like dataset handling and basic debugging with clear instructions, but drop significantly in open-ended and long-running tasks like model performance tasks where no agent succeeded.

OpenHands-Sonnet (OH-Sonnet) and ReAct-Sonnet are the two best performing agents performing agents with 50% and 47% success rate respectively, while OH-Gemini, AIDE-4o and React-4o achieve 17% success rate.

Across the 14 common successful tasks between OH-Sonnet and ReAct-Sonnet, we did not observe any discernible trends in cost. In long running tasks, OH-Sonnet used more tokens on average, especially in debugging and training tasks, while ReAct had higher usage in model implementation tasks like ChannelViT. Even in tasks with higher token usage, the costs dont scale proportionally due to efficient prompt caching in Openhands.

We presented ML-Dev-Bench, a benchmark focused on ML development workflows consisting of 30 tasks. We evaluated 3 agents on this benchmark - ReAct (with Claude Sonnet and GPT-4o), Openhands and AIDE; Openhands with Claude Sonnet performed the best out of these. Our open-source evaluation framework provides a foundation for the community to build upon. We encourage researchers and practitioners to explore various directions: analyzing the impact of scaling compute on these agents, studying variance in success metrics across multiple runs, evaluating emerging reasoning models like DeepSeek-R1 and o1/o3, and expanding the problem categories to include areas such as label and data collection. We welcome contributions to enhance the benchmark’s scope, methodology, and agent evaluation pipeline.

We share the traces of agentic runs for different tasks in Google Drive here https://drive.google.com/drive/folders/1o1FCvx_n9XVKgvkSWL97LlsfOh_KfG-S?usp=sharing.

We also share the comparison of token metrics across agentic frameworks while keeping the base model constant, and a detailed breakdown of performance at a per-task level below.