PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Deepak Nathani, Lovish Madaan, Nicholas Roberts, Nikolay Bashlykov, Ajay Menon, Vincent Moens, Amar Budhiraja, Despoina Magka, Vladislav Vorotilov, Gaurav Chaurasia, Dieuwke Hupkes, Ricardo Silveira Cabral, Tatiana Shavrina, Jakob Foerster, Yoram Bachrach, William Yang Wang, Roberta Raileanu

arXiv:2502.14499v1cs.CLcs.AIcs.LG

TL;DR

AI research agents need benchmarks and environments for diverse, open-ended ML research workflows. The paper introduces MLGym and MLGym-Bench to provide that infrastructure and evaluates frontier models, finding broad quantitative-task competence alongside capability gaps and limited scientific novelty.

  • Problem

    The field lacks comprehensive frameworks and benchmarks for evaluating agents conducting open-ended AI research across diverse domains.

  • Method

    MLGym provides a unified Gym environment, while MLGym-Bench supplies diverse open-ended tasks and flexible artifact-specific evaluation for developing and assessing LLM agents.

  • Results

    Modern LLM agents can tackle diverse quantitative experiments, but the findings also reveal notable capability gaps.

  • Takeaways & Limitations

    Further progress requires stronger long-context reasoning, agent architectures, training and inference algorithms, and richer evaluation methodologies.

  • Takeaways & Limitations

    Scientific novelty is not yet clearly automatable or formally definable for agents, and progress depends partly on open data, methods, metrics, and definitions.

Abstract

from arXiv · show

We introduce Meta MLGym and MLGym-Bench, a new framework and benchmark for evaluating and developing LLM agents on AI research tasks. This is the first Gym environment for machine learning (ML) tasks, enabling research on reinforcement learning (RL) algorithms for training such agents. MLGym-bench consists of 13 diverse and open-ended AI research tasks from diverse domains such as computer vision, natural language processing, reinforcement learning, and game theory. Solving these tasks requires real-world AI research skills such as generating new ideas and hypotheses, creating and processing data, implementing ML methods, training models, running experiments, analyzing the results, and iterating through this process to improve on a given task. We evaluate a number of frontier large language models (LLMs) on our benchmarks such as Claude-3.5-Sonnet, Llama-3.1 405B, GPT-4o, o1-preview, and Gemini-1.5 Pro. Our MLGym framework makes it easy to add new tasks, integrate and evaluate models or agents, generate synthetic data at scale, as well as develop new learning algorithms for training agents on AI research tasks. We find that current frontier models can improve on the given baselines, usually by finding better hyperparameters, but do not generate novel hypotheses, algorithms, architectures, or substantial improvements. We open-source our framework and benchmark to facilitate future research in advancing the AI research capabilities of LLM agents.

1 Introduction

AI research agents could automate broad scientific workflows, but progress is constrained by the lack of benchmarks for open-ended research. MLGym and MLGym-Bench provide a unified environment and diverse task suite focused on baseline improvement.

  • Motivation: AI Research Agents are envisioned to search literature, generate hypotheses, design experiments, implement methods, analyze results, and communicate findings.ML research provides an empirical testbed because it emphasizes systematic experimentation and validation.
  • Research gap: Existing evaluation tools lack comprehensive frameworks and benchmarks for assessing open-ended AI research across diverse domains.Existing benchmarks either exclude open-ended research tasks or cover only narrow research areas, while frameworks do not support training-algorithm research.
  • Framework: MLGym is the first Gym environment for AI Research Agents, unifying diverse open-ended tasks while enabling research on RL, curriculum learning, and open-ended learning.The framework is designed for developing and evaluating LLM agents on AI research tasks.
  • Benchmark: MLGym-Bench contains 13 open-ended research tasks spanning computer vision, natural language processing, reinforcement learning, and game theory.The benchmark is intended to evaluate realistic, multifaceted research workflows and supports varied artifacts such as model weights, algorithms, and strategy code.
  • Capability scope: MLGym-Bench focuses on Level 1, where agents improve a non-state-of-the-art baseline rather than achieve novel scientific contributions.The capability framework ranges from reproducing existing work through long-term, paradigm-shifting research agendas.

2 Related Work

Prior agent benchmarks generally emphasize constrained software, data-science, or narrowly scoped ML tasks. MLGym extends this landscape with algorithmic, open-ended research tasks and flexible evaluation artifacts, while related work exposes both capability gaps and evaluation limitations.

  • Distinctive framework features: MLGym is presented as the first AI Research Agent framework with a Gym interface, supporting integration and training with RL algorithms.Its benchmark also includes algorithmic research tasks across domains such as reinforcement learning, game theory, and SAT.
  • Benchmark scope: MLGym-Bench covers open-ended tasks across supervised learning, language modeling, reinforcement learning, game theory, and SAT, unlike narrower repository or Kaggle benchmarks.These tasks target research objectives rather than only code correctness or already-solved competition problems.
  • Evaluation: MLGym supports flexible evaluation artifacts, including model checkpoints, RL algorithms, and strategy code, rather than requiring a single CSV submission format.A Python evaluator can examine the quality of the current solution for each task.
  • Related agent research: Related agent work spans tool use, web navigation, operating systems, games, software engineering, data science, search, and text-to-SQL workflows.These settings provide increasingly complex automation testbeds while generally retaining controlled interfaces and measurable outcomes.
  • Capability evidence: 0% success was reported for current state-of-the-art LLMs on the most difficult MLAgentBench tasks.Other evaluations show agents can outperform humans under short budgets, while humans achieve higher returns with substantially more time.
  • Evaluation limitations: Human-preference-only evaluation is insufficient for assessing end-to-end experimental solution generation without objective metrics.The cited evaluation combines subjective annotations and automatic human-preference evaluations but lacks supporting objective measures.

3 MLGym

MLGym provides a modular Gym-based environment for developing and evaluating LLM agents on open-ended ML research tasks. Its shell-based interaction, extensible task abstractions, reproducibility safeguards, external tools, and persistent memory support long-horizon experimentation.

  • Agent and environment: MLGym agents interact with a shell environment through commands to generate ideas, process data, implement methods, train and evaluate models, analyze results, and plan subsequent experiments.Execution feedback lets agents iteratively self-refine solutions in context.
  • Framework design: MLGym separates agents from environments and exposes modular Agents, Environment, Datasets, and Tasks components for integrating models, harnesses, tools, datasets, and new benchmark tasks.The separation also supports comparing base models under the same agentic harness.
  • Task abstraction: Task configurations can specify datasets, read-only evaluation scripts, task-specific environments, starter code, timeouts, and memory settings, enabling varied difficulty and curriculum construction.Baseline code can define easier tasks, while missing or buggy starter code creates harder versions.
  • Evaluation: Evaluation uses task-specific artifacts and read-only scripts rather than a universal CSV format, accommodating reinforcement-learning models and game-theoretic strategy code while protecting evaluation logic.The framework explicitly supports diverse open-ended task protocols.
  • Tools and reproducibility: MLGym extends tool support with permission management, literature search, PDF parsing, and experimental memory, while read-only datasets and evaluation scripts improve reproducibility and cheating prevention.Tools return textual feedback when agents attempt disallowed file operations.
  • Memory Module: The Memory Module stores findings and successful training configurations so agents can retrieve strong configurations and continue iterating during long-horizon tasks.Authors report that memory-equipped agents sustain progress over extended trial sequences and reduce constrained-context limitations.

4 MLGym-Bench

MLGym-Bench is a diverse testbed of open-ended ML research tasks spanning multiple domains and research abilities. Tasks use standardized baselines and evaluation procedures, while game-theoretic settings emphasize best responses to provided opponents.

  • Benchmark scope: MLGym-Bench evaluates data handling, model architecture, and strategic decision-making across data science, game theory, computer vision, natural language processing, and reinforcement learning.The benchmark is intended to test generalization across varied real-world scenarios.
  • Benchmark scope: The benchmark provides standardized evaluation scripts and baseline implementations for comparing performance across its selected tasks.Task selection is intended to represent real-world challenges.
  • Data Science: House Price Prediction uses Kaggle housing data and evaluates Ridge Regression with RMSE and R2 after minimal feature engineering.The task focuses on predicting prices from housing features.
  • Game Theory and algorithmic tasks: The 3-SAT task asks agents to optimize a variable-selection heuristic around read-only DPLL code, measuring total wall-clock time on 100 generated instances against random selection.The agent can inspect but not modify the DPLL implementation.
  • Game Theory: Repeated-game tasks ask agents to produce strategies that best respond to provided opponent bots, using prior action history to select actions over k rounds.The suite includes iterated Prisoner’s Dilemma, Battle of the Sexes, and Colonel Blotto.
  • Game Theory: Battle of the Sexes uses 20-round repeated play with a random-choice opponent informed by the previous round.The game captures coordination between players with different venue preferences.
  • Game Theory: All game-theoretic tasks expose the opponent strategy, so they measure code understanding and exploitation of that strategy rather than performance against undisclosed opponents.Future tasks are planned without provided strategies and with round-robin opponents.
  • Computer Vision and NLP: The vision and language tasks include CIFAR-10 image classification, Fashion MNIST classification, and FineWeb language modeling using supplied baselines and datasets.The language-modeling task starts from NanoGPT and uses 1.773B training tokens and 100M validation tokens.

5 Experimental Setup

The experiments evaluate a SWE-Agent-based model using five frontier LLMs within a constrained, command-driven MLGym environment. Runs use fixed interaction limits, rolling context management, specialized tools, and task-specific training timeouts.

  • Agent configuration: The experimental agent follows a ReAct-style thought-action loop, receives task and dataset descriptions plus ACI documentation, and issues one non-interactive command per step.Interactive sessions such as Python REPL and vim are disallowed.
  • Models: Experiments use OpenAI O1-preview, Gemini 1.5 Pro, Claude-3.5-sonnet, Llama-3-405b-instruct, and GPT-4o.Models use temperature 0.0 and top-p 0.95 except O1-preview, whose defaults are retained.
  • Environment configuration: The environment uses a 1000-line window with 2-line overlap and retains the five most recent action-observation interactions.These settings support navigation of large files while keeping recent context manageable.
  • Environment configuration: Specialized commands support file navigation, editing with linting, file and directory search, and validation or submission.The command interface extends standard bash operations.
  • Run limits: Each agent run is limited to 50 environment interactions, after which the final codebase state is autosubmitted, and training commands receive task-specific timeouts.These limits control runtime and discourage simply increasing model parameter counts.

6 Evaluation

MLGym evaluates methods across diverse tasks using performance profiles and AUP scores rather than relying only on aggregate averages. The evaluation adapts these measures to MLGym-specific metric directions, infeasible methods, repeated runs, and final-versus-best attempts.

  • Aggregate averages can weight metrics undesirably and disproportionately penalize certain methods.
  • AUP score: AUP converts each method’s performance-profile curve into a scalar score by integrating up to τmax, with higher values indicating better performance.τmax is the minimum threshold at which every method reaches a profile value of 1.
  • Performance profiles: Performance profiles measure the proportion of tasks where a method is within threshold τ of the best method.The performance ratio compares each method’s task score with the best score on that task.
  • MLGym adaptations: MLGym defines a method as an agent scaffolding paired with a backbone model and adapts performance profiles and AUP scores for task-specific edge cases.Because the experiments use one scaffolding, SWE-Agent, comparisons focus on backbone models.
  • MLGym adaptations: Higher-is-better metrics invert the performance-ratio calculation and use the maximum task score instead of the minimum.This handling applies to metrics such as accuracy and R2.
  • MLGym adaptations: Methods are marked infeasible when they produce no valid solution or fail to outperform the baseline, and their scores receive an ε-adjusted baseline penalty.The infeasible score is set to (1 + ε) × r_t,m_baseline.
  • Run aggregation: Metrics are reported over four independent runs, with separate profiles for best final submissions and best attempts observed during validation.Final submissions use the best final submission across runs, whereas best attempts use the best validated attempt.
  • AUP variants: AUPbs measures consistently submitting the best attempt, while AUPba reflects exploration capability and indicates the model’s performance ceiling.AUPbs also requires preserving the best attempt and recovering from later mistakes.

7 Results

Across diverse tasks, OpenAI O1-preview leads aggregate performance, while Gemini-1.5-Pro offers the strongest performance–cost balance. Model behavior also reveals substantial differences in completion, error handling, and experimental trajectories.

  • 7.1 AUP Scores and Performance Profiles: OpenAI O1-preview achieves the best aggregate performance for both Best Attempt and Best Submission, with Gemini 1.5 Pro and Claude-3.5-Sonnet close behind.
  • 7.1 AUP Scores and Performance Profiles: OpenAI O1-preview is not dominant on every task but remains consistently among the top-performing models, yielding the highest AUP scores and performance profiles.
  • 7.3 Computational Cost: Gemini-1.5-Pro is the cheapest model to run, approximately 9× cheaper than OpenAI O1, while achieving 99% of O1’s AUP and the best overall performance–cost balance.
  • 7.4.1 Failure Mode Analysis: Evaluation Error accounts for 75% of termination errors, usually because submission artifacts are missing or incorrectly formatted; Gemini-1.5-Pro submits no invalid solutions.
  • 7.4.1 Failure Mode Analysis: GPT-4o has the highest failure rate, whereas Gemini-1.5-Pro and OpenAI O1-preview have the best completion rates; O1 nevertheless produces at least one valid solution for every task.
  • 7.4.2 Action Analysis: Agents primarily edit and view files while repeatedly running experiments and validation, but some submit as early as Step 5 and likely miss better solutions.

8 Discussion and Limitations

MLGym-Bench shows that modern LLM agents can handle diverse quantitative experiments, while exposing gaps in generalization, scientific novelty, and evaluation beyond ML tasks.

  • Modern LLM agents successfully tackle diverse quantitative experiments, demonstrating advanced skills and domain adaptability.
  • Scaling beyond ML tasks: Scaling beyond ML tasks is needed to assess robustness and generalizability on larger datasets, more complex tasks, and domains outside AI.
  • Interdisciplinary Ablations and Generalization: Interdisciplinary evaluation can test method applicability across domains and use ablations to identify factors contributing to cross-domain generalization.
  • Addressing Scientific Novelty: Scientific novelty remains difficult to automate or formally define, with progress varying by discipline and depending on available data, methods, metrics, and definitions.
  • Data Openness Imperative: Removing public resources after their inclusion in LLM training can make factual sources and attribution impossible to identify, hindering reproducibility and interpretation of results.

9 Ethical Considerations

MLGym-Bench is framed as a tool for understanding the autonomy and risks of AI research agents. Greater agent capability could accelerate science, but may also outpace human oversight and increase misuse risks.

  • MLGym-Bench can help measure model autonomy and support transparency regarding acceleration risks in frontier AI labs.
  • Autonomous AI research could accelerate progress in healthcare, climate science, safety research, and novel product development.
  • Agents that improve their own training code could advance cutting-edge models faster than human researchers can understand their consequences.
  • The authors acknowledge MLGym-Bench’s limitations and encourage evaluations tailored to researchers training frontier models.

10 Conclusions

The paper presents MLGym and MLGym-Bench as initial steps toward robust, flexible, and transparent LLM agents for AI research. It identifies improved reasoning, agent design, algorithms, and evaluation as necessary for broader scientific impact.

  • MLGym and MLGym-Bench are initial steps toward robust, flexible, and transparent LLM agents for AI research.
  • Long-context reasoning, agent architectures, training and inference algorithms, and richer evaluation methodologies are identified as essential future improvements.
  • The intended direction is AI-driven scientific research that maintains verifiability, reproducibility, and integrity.

A.1 Computational Cost

The appendix reports computational resources, runtime measures, and model token-pricing details for MLGym-Bench experiments.

  • Training Timeout limits Python-command duration to prevent continuous scaling of model parameters.
  • Average agent runtime and Baseline runtime report wall-clock time for agent runs and provided baseline code, respectively.
  • Table 7 lists the computational resources required for each MLGym-Bench task.
  • Table 8 reports model pricing, token usage, and context-length details, with pricing measured in USD per 1M tokens.

A.2 Failure Mode Analysis

Failure rates vary substantially across tasks: Language Modeling and the reinforcement-learning tasks are most challenging, while Fashion MNIST and Prisoner’s Dilemma are most reliable. These patterns align with performance scores and suggest that complex architectural or algorithmic demands drive difficulty.

  • Task difficulty: Language Modeling and all reinforcement-learning tasks have the highest failure rates.The reinforcement-learning tasks are Meta Maze, Mountain Car Continuous, and Breakout.
  • Task difficulty: Fashion MNIST and Prisoner’s Dilemma have the lowest failure rates, with valid intermediate solutions and submissions for all seeds.
  • Sources of difficulty: Tasks requiring complex architectural decisions or algorithms are more difficult than traditional supervised-learning tasks.Language Modeling exemplifies architectural complexity, while Breakout, Meta Maze, and Mountain Car Continuous require complex algorithms.
  • Model behavior: More advanced models demonstrate better error handling and completion rates overall.

A.3 Action Analysis

Action patterns differ across task families: reinforcement-learning tasks require the most actions, while game-theoretic tasks require the fewest and algorithmic tasks show many validation actions. Rankings based on Best Attempt@4 and Best Submission@4 use BORDA aggregation, while memory retrieval supports restarting from strong configurations and submitting the best model.

  • Action distribution: Reinforcement-learning tasks have the highest action counts, whereas game-theoretic tasks have the lowest.Figure 10 groups actions according to the categories defined in Table 2 and Section 7.4.2.
  • Action distribution: 3-SAT and game-theoretic tasks have many validation actions, indicating a quick experimental cycle.The game-theoretic tasks are Blotto, Prisoner’s Dilemma, and Battle of the Sexes.
  • Model rankings: BORDA aggregate rankings based on Best Attempt@4 and Best Submission@4 align with AUP results but do not convey relative performance differences.Tables 9 and 10 report the two ranking schemes and their aggregate ranks.
  • Memory-assisted exploration: The memory module stores and retrieves experimental results so agents can restart exploration from the best training configuration and submit the best model.Figures 11 and 12 illustrate this retrieval-and-restart process.
  • Agent interface: The agent operates through a command-line interface with file editing, validation, and submission commands.The interface requires one discussion and one command field, and editing commands require proper indentation.
Loading 2502.14499v1…