MLGym: A New Framework and Benchmark for Advancing AI Research Agents | Summary

This article explains the key points of MLGym: A New Framework and Benchmark for Advancing AI Research Agents.

• 2025-02-20 (arXiv)
• Nathani, Deepak, Madaan, Lovish, Roberts, Nicholas, Bashlykov, Nikolay, Menon, Ajay, Moens, Vincent, Budhiraja, Amar, Magka, Despoina, Vorotilov, Vladislav, Chaurasia, Gaurav, et al.
• University of California, Santa Barbara, University College London, University of Wisconsin–Madison, University of Oxford, PyTorch Core Libraries at Meta, FAIR at Meta, GenAI at Meta
• Paper
• Github

Read this article in Korean

Summary

• MLGym addresses the lack of a common environment for evaluating large language model (LLM) agents, systems that execute commands and use feedback to improve solutions, across diverse machine learning research tasks. Using the same SWE-Agent command-execution framework and best-of-four results from independent runs on each of 13 tasks, OpenAI O1-preview achieves the highest aggregate scores among five evaluated models, while the authors report that improvements usually come from tuning existing methods.
• The framework separates agents, environments, datasets, and tasks behind a Gymnasium interface, a standard interface for reinforcement learning, where policies learn through interaction and reward. It supports executable evaluation of model weights, algorithms, and game strategies, allowing tasks with different output formats and performance measures to share an evaluation platform.
• OpenAI O1-preview records Best Attempt AUP@4 of 1.150, versus 1.142 for Claude-3.5-Sonnet and 1.140 for Gemini-1.5-Pro; Best Attempt@4 selects the best intermediate evaluation across four runs, and AUP is the area under a performance profile summarizing relative performance across tasks. The paper reports that Gemini-1.5-Pro reaches approximately 99% of O1-preview’s Best Attempt AUP at approximately 9× lower model-call cost. However, applying the Appendix A.1 token counts and prices gives about $9.18 for O1-preview and $0.36 for Gemini-1.5-Pro, a roughly 25.4× difference; the paper does not reconcile these figures. The benchmark measures improvement over supplied baselines and does not establish autonomous scientific discovery.

1 Introduction

MLGym: A New Framework and Benchmark for Advancing AI Research Agents proposes a common environment for iterative workflows involving implementation, training, evaluation, and analysis. The authors identify gaps in existing benchmarks’ domain coverage, support for training research agents, and flexibility of evaluated artifacts, meaning outputs such as trained models, algorithms, or predictions. The framework and benchmark are released at https://github.com/facebookresearch/MLGym.

The environment mediates commands and feedback between the agent and computer while supplying tools, code, data, and dependencies. This separation allows different agents to operate on the same task infrastructure.

1.1 Capability Levels for AI Research Agents

The proposed hierarchy contains Level 0: Reproduction, Level 1: Baseline Improvement, Level 2: SOTA Achievement, Level 3: Novel Scientific Contribution, Level 4: Groundbreaking Scientific Contribution, and Level 5: Long-Term Research Agenda; SOTA means state-of-the-art performance. Level 2 permits only a task description and literature predating the SOTA approach, excluding that approach’s paper and code, while higher levels require increasingly substantial and sustained scientific contributions. MLGym-Bench targets Level 1, improving supplied code that is below the state of the art.

2 Related Work

The related work covers research frameworks and benchmarks, general tool-using agents, software engineering and data science agents, and scientific research systems. These comparisons motivate evaluating measurable research outcomes while allowing agents to change methods and produce different kinds of artifacts.

2.1 AI Research Frameworks and Benchmarks

The authors compare MLGym with MLE-Bench, a machine learning engineering benchmark; SWE-Bench/Agent, a software engineering benchmark and agent; and the research benchmarks MLAgentBench, RE-Bench, and ScienceAgentBench. The comparison covers Gym interfaces, algorithmic tasks, open-ended research, flexible artifacts, and default agent execution frameworks. MLGym combines these features, although the paper acknowledges that several other research benchmarks already use script-based evaluation.

2.2 LLM Agents

General agent benchmarks evaluate interaction with tools, websites, operating systems, games, and simulated or physical environments. AssistantBench tests realistic web tasks, AgentBench covers eight interactive settings, GAIA tests real-world questions requiring tools and multimodal reasoning, and AppWorld evaluates iterative interaction with applications. The authors distinguish these settings from open-ended AI research and emphasize evaluation protocols that report both performance and cost.

2.3 Agents for Software Engineering and Data Science

Software engineering systems such as SWE-agent and OpenHands provide controlled execution environments, while AutoKaggle and AgentK v1.0 automate data science competition workflows and the Spider systems address database queries and enterprise workflows. Search-based methods explore alternative execution paths, but the cited Agentless results also show that a simple three-step procedure can outperform more elaborate agents on SWE-bench Lite. These studies concern different agent procedures as well as different underlying models.

2.4 Agents for Scientific Research

Scientific research evaluation raises questions about novelty, falsifiability, meaning whether a claim can be challenged through tests, and generalization beyond producing functioning code.

• Data Science: ML-Bench emphasizes repository tasks, Data Interpreter covers broader applications, and SUPER requires agents to formulate and execute tasks within research repositories.
• AI Research: automated machine learning and neural architecture search automate method and architecture choices within specified search spaces, whereas research agents can propose less constrained solutions. MLAgentBench, RE-Bench, and MLE-bench provide related experimental settings with different task coverage and agent execution frameworks.
• Scientific Discovery: DiscoveryWorld evaluates discovery in simulated environments, ScienceAgentBench evaluates publication-derived tasks, and ResearchAgent iteratively develops literature-based ideas. The authors argue that subjective preference judgments need support from objective outcome measures.

3 MLGym

An agent receives a task description, starter code, and execution history, then iteratively issues shell commands and observes their outputs. MLGym implements this interaction in a local Docker container, an isolated software workspace, through four modular components: Agents, Environment, Datasets, and Tasks. Researchers can replace the agent execution framework, add tools, reuse datasets, or introduce tasks independently.

3.1 Agents

The Agent class wraps a base LLM and manages model integration, history processing, and costs. It returns an action from prior observations and actions, while the environment executes that action and supplies the next observation. The default agent is adapted from SWE-Agent, enabling comparisons between base models within the same execution procedure.

3.2 Environment

The Environment initializes task dependencies, tools, code, and data in a separate Docker workspace and manages command execution. A non-root user named “agent” and file permissions restrict access to protected resources. Tools are initialized separately from the agent, and their documentation can be inserted into its prompt; a footnote acknowledges that the latest SWE-Agent release also separates tools from agents.

3.3 Datasets

Dataset configuration is independent of task configuration, supporting both locally stored files and datasets hosted on Hugging Face. A dataset can serve multiple tasks, and a task can contain multiple datasets to evaluate whether a method generalizes. Local dataset files are copied with read-only permissions; remotely hosted datasets are identified through URLs supplied in prompts or starter code.

3.4 Tasks

Task configurations specify datasets, read-only evaluation scripts, dependencies, optional starter code, training timeouts, and memory settings. Evaluation artifacts can include prediction files, trained models, or executable strategy functions, depending on the task. Varying starter code availability or introducing bugs can define different difficulty levels, although the reported benchmark supplies baselines.

3.5 Tools and ACI

The ACI, or agent-computer interface, exposes file search, navigation, editing, validation, and submission as shell-accessible tools, with textual feedback when permissions block an operation. The validate command evaluates the current artifact on the test set without ending the run, whereas submit evaluates it and terminates the trajectory, the sequence of actions and observations.

• Extended tools include literature_search and parse_pdf_url for finding and extracting papers, plus memory_write and memory_read for storing research logs and retrieving entries through similarity between numerical text representations.
• The memory examples illustrate returning to a successful configuration after earlier experiments. However, the paper states that its reported experiments use only SWE-Agent tools and validation, so the principal model comparison does not establish benefits from literature search or memory.

4 MLGym-Bench

MLGym-Bench contains 13 tasks spanning regression, image classification, image captioning, language inference, language modeling, reinforcement learning, satisfiability solving, and repeated games. Each task includes a baseline implementation and a standardized evaluator. The tasks exercise data handling, training decisions, code modification, and strategic reasoning within one environment.

4.1 Data Science

House Price Prediction asks the agent to predict housing prices from tabular features, starting from Ridge Regression, a regression model, with minimal feature engineering. The task description names root mean squared error, a prediction-error measure, and R², a regression fit score; the results tables report R². The dataset URL provided in the paper is https://www.kaggle.com/datasets/yasserh/housing-prices-dataset.

4.2 3-SAT

The 3-SAT task, a Boolean satisfiability problem, asks agents to improve variable selection in supplied DPLL solver code. The solver itself is read-only, while the variable-selection heuristic, the rule choosing the next variable, can be changed from its random-selection baseline. Performance is total wall-clock time in seconds for solving 100 generated instances.

4.3 Game Theory

Game tasks ask agents to write a best response, a strategy maximizing expected payoff against a fixed opponent, for Prisoner’s Dilemma, Battle of Sexes, and Colonel Blotto. Prisoner’s Dilemma involves choosing cooperation or defection, Battle of Sexes involves coordinating despite different preferences, and Colonel Blotto involves allocating limited resources across competing battlefields; the first two use k = 20 rounds. Agents can inspect opponent code, so these tasks evaluate understanding and exploitation of known strategies.

The inventory combines supervised and self-supervised model training, reinforcement learning, satisfiability solving, and game strategies. These tasks require different submission artifacts and performance measures.

4.4 Computer Vision

CIFAR-10 evaluates ten-class image classification from a baseline accuracy of 49.71%, while Fashion MNIST classifies clothing images from a two-layer convolutional neural network baseline, a model using spatial filters on images. Agents may modify architectures and training hyperparameters, settings governing model training. MS-COCO evaluates image captioning using BLEU, a measure comparing generated captions with reference captions, after removing all images containing humans from the training and validation sets.

4.5 Natural Language Processing

Natural Language Inference starts from pretrained BERT, a language-understanding model, and baseline fine-tuning code for MNLI, a benchmark for inference between text pairs. Language Modeling trains next-token prediction from version #8 of modded-nanogpt at https://github.com/KellerJordan/modded-nanogpt, using FineWeb text subsets containing 1.773B training tokens and 100M validation tokens. The task description names perplexity, a measure of predictive uncertainty, while the results tables label their reported values Validation Loss.

4.6 Reinforcement Learning

Reinforcement learning is represented by MetaMaze Navigation, Mountain Car Continuous, and Breakout MinAtar. These require navigating from local observations, driving a car up a hill with continuous controls, and playing a simplified arcade game. All use Gymnax, a reinforcement learning environment library, and PPO, the supplied policy-training algorithm, from https://github.com/RobertTLange/gymnax-blines.

5 Experimental Setup

The experimental design fixes the SWE-Agent-derived execution framework while changing the backbone model, the LLM generating the agent’s actions. Task-specific training timeouts constrain training, and a common interaction limit bounds the commands an agent can issue. Four independent runs per model and task support the reported best-of-four comparisons.

5.1 Agent and Models

The agent uses a ReAct-style loop, alternating reasoning and actions, with one command per step and no interactive sessions such as a Python REPL, an interactive interpreter, or vim. Evaluated models are OpenAI O1-preview, Gemini 1.5 Pro, Claude-3.5-sonnet-20241022, Llama-3.1-405b-Instruct as identified in the results, and GPT-4o. Decoding uses temperature=0.0 and top-p=0.95, settings controlling output sampling, except for O1-preview, which uses its fixed default temperature=1.0.

5.2 Environment Configuration

File viewing uses a 1000-line window with 2-line overlap, and history processing retains the five most recent action-observation interactions. Runs stop after 50 steps, at which point the last codebase state is automatically submitted. Task-specific training timeouts limit individual training commands and discourage improvements obtained solely by increasing model size.

6 Evaluation

Averaging raw accuracy, reward, loss, and runtime scores would weight tasks according to their metric scales. Following Dolan and Moré (2002), performance profiles instead summarize each method’s relative distance from the best method on each task. The area under each profile supplies an aggregate AUP score, with larger values indicating better performance under the chosen construction.

6.1 Performance Profiles and the AUP Score

Equation (1) defines ρ_m(τ) = |{t ∈ T : log₁₀ r_t,m ≤ τ}| / |T|, with r_t,m = ℓ_t,m / min{ℓ_t,m : m ∈ M} for lower-is-better metrics; T denotes tasks, M denotes methods, and ℓ_t,m is a task score. Equation (2) prints AUP_m = ∫₁^τmax ρ_m(τ)dτ, where τmax is the smallest threshold at which all methods’ profiles reach 1. The printed lower integration limit of 1 warrants clarification because the profile uses logarithmic ratios and the plotted thresholds begin at 0.

6.2 Usage in MLGym

For higher-is-better metrics, Equation (3) uses r_t,m = max{ℓ_t,m : m ∈ M} / ℓ_t,m. A method without a valid baseline-beating solution is treated as infeasible and assigned a ratio of (1 + ε) × r_t,m_baseline, with ε = 0.05.

• Best Submission@4 selects the best final artifact across four runs, including automatic submissions after termination.
• Best Attempt@4 selects the best valid intermediate evaluation across four runs. Separating the two shows whether successful experiments survive final submission; neither best-of-four measure estimates average-run reliability.

7 Results

Results combine aggregate profiles, individual task scores, model-call costs, and analyses of execution failures and actions. Agents improve supplied baselines across several domains, with differences in outcome quality, reliability, and cost.

7.1 AUP Scores and Performance Profiles

OpenAI O1-preview leads both aggregate measures: Best Attempt AUP@4 is 1.150 and Best Submission AUP@4 is 1.176. Claude-3.5-Sonnet records 1.142 and 1.135, and Gemini-1.5-Pro records 1.140 and 1.125; Llama3.1-405b-instruct and GPT-4o score lower. Attempt and submission profiles are constructed separately, so their aggregate values do not directly measure improvement between intermediate and final artifacts.

O1-preview leads Best Attempt AUP@4 at 1.150 and Best Submission AUP@4 at 1.176. Claude and Gemini are close on Best Attempt, with scores of 1.142 and 1.140; these aggregate differences are not percentage improvements in individual task performance.

7.2 Raw Performance Scores

Task-level leadership varies: Claude-3.5-Sonnet achieves Best Attempt CIFAR-10 accuracy of 0.895 versus the tabulated baseline of 0.497, Gemini-1.5-Pro achieves Breakout score 71.389 versus 48.817, and O1-preview achieves language-model validation loss 3.966 versus 4.673. Llama-3.1-405b-Instruct produces no valid language-modeling solution, and GPT-4o produces no valid Breakout solution.

• Successful intermediate results are sometimes lost: Claude’s MS-COCO BLEU falls from Best Attempt@4 of 0.298 to Best Submission@4 of 0.125, below the 0.279 baseline.
• O1-preview’s 3-SAT runtime is 13.652 s for Best Attempt@4 and 13.83 s for Best Submission@4, against a 16.158 s baseline. These compare improvements over the supplied heuristic, without testing against advanced satisfiability solvers.

Different models lead different tasks despite O1-preview’s aggregate lead. The infinity entries identify the absence of any valid language-modeling solution for Llama and any valid Breakout solution for GPT-4o.

Comparison with intermediate scores reveals failures to retain successful artifacts. Claude’s MS-COCO BLEU drops from 0.298 in Best Attempt@4 to 0.125 in Best Submission@4, illustrating why final submission quality needs separate measurement.

7.3 Computational Cost

The cost comparison plots Best Attempt AUP@4 against average model application programming interface (API) cost, the charge for model calls, using prices reported in the paper. The paper reports that Gemini-1.5-Pro achieves approximately 99% of O1-preview’s AUP at approximately 9× lower cost. Yet its Appendix A.1 token counts and prices imply estimated model-call costs of about $9.18 for O1-preview and $0.36 for Gemini-1.5-Pro, or roughly 25.4× lower; the paper does not reconcile this discrepancy. These figures exclude a complete accounting of task-training compute.

Gemini lies near Claude and O1-preview in aggregate performance while incurring lower model-call costs. The plot supports the reported cost tradeoff under the paper’s provider pricing but does not account for total task-training compute.

7.4 Agent Behavior Analysis

7.4.1 Failure Mode Analysis examines 220 trajectories from 11 tasks, five models, and four seeds, with evaluation errors accounting for the reported 75% of termination errors. A failed run produces no valid intermediate artifact before an error, while an incomplete run produces at least one; validity here does not require beating the baseline. 7.4.2 Action Analysis shows that editing dominates activity, with training and validation interleaved throughout execution.

• Figure 5 records failed/incomplete counts of 12/3 for Llama, 12/3 for Claude, 11/2 for GPT-4o, 6/4 for Gemini, and 2/9 for O1-preview. These counts conflict with the prose statement that GPT-4o has the highest failure rate.
• Figure 6 reports 3,708 Edit, 1,458 Python, 1,066 Validate, 835 View, 449 Bash, 226 Submit, and 11 Search actions. The stated 1% Search share does not match the displayed counts.
• Per-model and per-step analyses show fewer actions for GPT-4o, extensive editing by Claude and O1-preview, initial environment inspection, and submissions as early as Step 5. These distributions describe behavior but do not establish the causes of performance differences.

O1-preview has 2 wholly failed runs and 9 incomplete runs with valid intermediate artifacts. Llama and Claude each have 12 failed runs, compared with GPT-4o’s 11, conflicting with the surrounding prose’s highest-failure claim.

Editing accounts for 3,708 actions, followed by Python execution and validation. Only 11 Search actions are displayed, confirming rare use while exposing a mismatch with the prose’s stated 1% share.

8 Discussion and Limitations

The authors propose broader task coverage, interdisciplinary generalization tests, and ablation studies, experiments removing or changing components to measure their effects. Scientific novelty remains difficult to define and automate, and the authors emphasize open scientific data and code for verification. The present baseline-improvement setting does not evaluate the higher capability levels proposed in Section 1.1.

9 Ethical Considerations

The ethical discussion considers possible scientific benefits alongside risks from agents improving AI capabilities faster than evaluation and control methods advance. The authors propose MLGym-Bench as one possible input to assessments of autonomous research capabilities and release it to support transparent measurement. They encourage additional evaluations tailored to frontier-model research workflows.

10 Conclusions

The conclusion presents MLGym and MLGym-Bench as infrastructure for developing and evaluating research agents. It identifies long-context reasoning, agent architectures, training and inference methods, and richer evaluation as areas for further work. The experiments demonstrate baseline improvement under the reported settings, while broader scientific-discovery objectives remain prospective.

Appendix

• A Additional Results and Analysis, A.1 Computational Cost: Table 7 specifies 30m training-command timeouts for most tasks and 40m for MS-COCO, MNLI, and Language Modeling. Allocations range from 0 graphics processing units (GPUs) for games and 3-SAT to 2 GPUs for language modeling and reinforcement learning; reported average agent runtimes range from 30m to 4h, with MS-COCO and MNLI runtime cells left blank.
• A.1 Computational Cost also reports token usage, context lengths, and USD pricing per 1M tokens. O1-preview averages 368898 input and 60704 output tokens at prices of 15.0 and 60.0, while Gemini-1.5-Pro averages 282613 input and 1633 output tokens at 1.25 and 5.00. Llama uses an internally hosted FP-8 checkpoint, an eight-bit floating-point model representation, with Together AI pricing used for cost estimation.
• A.2 Failure Mode Analysis identifies language modeling and reinforcement learning as difficult execution settings; Fashion MNIST and Prisoner’s Dilemma produce valid intermediate and final solutions for every model and seed. A.3 Action Analysis shows the largest action totals on reinforcement learning tasks, while games have lower totals and algorithmic tasks have frequent validation.
• A.4 Model Rankings supplies per-task rankings and aggregates them with BORDA count, a method combining ranking positions. O1-preview ranks first overall under both Best Attempt@4 and Best Submission@4. The tables omit MS-COCO and MNLI, and their aggregate order differs from Table 4 despite the text’s claim of agreement with AUP.
• A.5 Memory Utilization illustrates retrieving successful Fashion MNIST configurations and continuing experimentation before submission. These examples show the intended workflow, but the appendix provides no controlled comparison with memory disabled to establish its average benefit.
• B Prompts reproduces the experimental agent’s tool documentation, response format, execution restrictions, and advice about validating and submitting solutions. The listing includes a 50-step example and a stated 1800-second action limit. Its submission.csv instruction differs from the framework’s broader support for task-specific artifacts and should be read as part of the reproduced prompt example.

Brief Thoughts

MLGym’s reusable environment makes model comparisons possible within one agent procedure and exposes failures to preserve successful intermediate results. The evidence is limited to four-run maxima, one execution framework, supplied baselines, and repeated test-set feedback, leaving independent generalization unestablished. The printed AUP definition, differing analysis task sets, and discrepancies between plotted counts and prose warrant checking against the implementation before reproducing the aggregate rankings.