MLGym Benchmark for LLM Research

• MLGym Benchmark is a programmable, Gym API-based environment designed to evaluate LLM agents on iterative machine learning research tasks.
• It integrates modular components including agents, secure Docker execution, and a synthetic task generation pipeline to support diverse research workflows.
• The benchmark employs standardized evaluation protocols with AUP metrics to systematically measure performance across domains like NLP, computer vision, and reinforcement learning.

MLGym is a programmable benchmark and execution environment designed for evaluating and developing LLM agents on end-to-end machine learning research tasks. It distinguishes itself as the first Gym API-compatible environment for machine learning tasks that emphasize iterative protocols—such as hypothesis generation, code editing, experiment execution, and empirical validation—operating within full software engineering workflows rather than single-turn question-answering. MLGym comprises a modular framework, MLGym-Bench (a suite of 13 open-ended research tasks), and a standardized evaluation protocol that enables systematic measurement of AI agent research capabilities across diverse domains including computer vision, natural language processing, reinforcement learning, data science, game theory, and algorithmic reasoning (Nathani et al., 20 Feb 2025, Cai et al., 17 Mar 2026).

1. Architectural Design and Environment Interface

MLGym is constructed around four modular components: Agents, Environment, Datasets, and Tasks. The environment adheres to and extends the gym.Env API by representing each agent–environment interaction as an $$(\text{action}, \text{observation}, \text{reward}, \text{done})$$ tuple, modified for the demands of multi-modal, code-centric research tasks.

• Agent Interface: An Agent encapsulates any base LLM, supporting a ReAct-style loop—at each step, the model consumes the entire transcript of prior observations (outputs, tool feedback, validation metrics) and actions, emitting exactly one shell or tool command. The agent code is entirely decoupled from the environment execution logic, supporting plug-and-play model interchangeability. All interactions occur through a single “SWE-Agent” harness.
• Environment Wrapper: All tasks are executed within Docker containers configured for secure, reproducible research workflows. The non-root “agent” user is provisioned with standard bash tools and a custom suite (open, edit, search_file, validate, submit, literature_search, memory_read/write), along with task-specific conda environments and immutable starter code or datasets.
• Action and Observation Space: Actions consist of either bash commands or higher-level ACI tools. Observations aggregate verbatim stdout/stderr and tool feedback, supporting multi-line textual and numeric responses. Discrete, tokenized action spaces and rich multi-modal observation spaces are integral to supporting LLM planning.
• Episode Structure: Each episode consists of up to 50 iterated steps (“rounds”); agents alternate between “thought” (rationale, planning) and “action” (tool invocation or shell command execution). An episode concludes either when the agent issues a “submit” command or the round limit is reached (Nathani et al., 20 Feb 2025, Cai et al., 17 Mar 2026).

2. Benchmark Task Suite

MLGym-Bench contains 13 open-ended research problems, each requiring an agent to iteratively employ the scientific method (hypothesis generation, implementation, and empirical analysis) to surpass a given baseline. The domain diversity yields a comprehensive assessment of agent abilities:

Domain	Task Example	Required Skills
Data Science	House Price Prediction	Feature engineering, model tuning
Algorithmic Reasoning	3-SAT Heuristic Optimization	Heuristic design, code synthesis
Game Theory	Repeated Iterated Prisoner’s Dilemma	Strategy synthesis, payoff maximization
Computer Vision	CIFAR-10 Classification, Image Captioning	Model design, optimization, evaluation
NLP	MNLI, Language Modeling (NanoGPT)	Fine-tuning, perplexity minimization
Reinforcement Learning	MetaMaze, MountainCar, Breakout	Policy editing, cumulative reward maxim.

Each task is defined by its own YAML configuration (id, description, entrypoint, timeouts, metrics), links to public datasets (frequently HuggingFace), starter code, and an evaluation script that computes a scalar score. For algorithmic and RL tasks, on-demand synthetic instance generation is supported via problem-specific generators (SAT sampling, Gymnax procedural maps). Task interaction generally proceeds through edit–run–validate cycles, with the agent able to submit intermediate or final solutions at any round (Nathani et al., 20 Feb 2025, Cai et al., 17 Mar 2026).

3. Evaluation Protocol and Metrics

Performance measurement in MLGym departs from naïve metric averaging and leverages performance profiles and the area under profile (AUP) score, inspired by Dolan & Moré (2002) and the AutoML Decathlon.

• Performance Normalization: For task $t$ and method $m$, let $\ell_{t,m}$ denote the relevant metric (e.g., accuracy, BLEU, average reward). Depending on the metric's preferred direction, the performance ratio is computed as either $\frac{\max_{m'}\{\ell_{t,m'}\}}{\ell_{t,m}}$ or $\frac{\ell_{t,m}}{\min_{m'}\{\ell_{t,m'}\}}$.
• Profile Curve: The function $\rho_m(\tau)=\frac{1}{|T|}|\{t: r_{t,m}\le\tau\}|$ yields the fraction of tasks where performance is within a factor $\tau$ of the best.
• AUP Calculation: The area under $\rho_m$ over a range of $\tau$ yields the AUP for a model.
• Evaluation Modes: Two principal metrics are reported—Best Attempt@4 (peak validate score in any step, across 4 seeds) and Best Submission@4 (final submit or best intermediate score).

Empirical results (mean AUP@4 over 13 tasks) for major LLMs are as follows (Nathani et al., 20 Feb 2025):

Model	AUP⁽ba⁾@4	AUP⁽bs⁾@4
Llama-3.1 405B-instr	1.015	1.039
Claude-3.5-Sonnet	1.142	1.135
Gemini-1.5-Pro	1.140	1.125
GPT-4o	1.000	1.029
OpenAI O1-Preview	1.150	1.176

Task-level improvements are generally attributable to hyperparameter search, with all models exceeding baseline but not producing novel research outcomes.

4. Synthetic Task Generation and Agent Training

MLGym’s extensible abstraction for datasets and tasks enables the automated synthesis of new ML research environments. The synthetic environment generation pipeline, as introduced in "AI Scientist via Synthetic Task Scaling" (Cai et al., 17 Mar 2026), consists of:

1. Topic Sampling & Dataset Validation: GPT-5 samples ML topics, proposes task descriptions, metrics, and datasets. HuggingFace datasets are validated programmatically.
2. Configuration & Starter Code Generation: Task JSONs inform YAML configurations and baseline code, again synthesized via GPT-5.
3. Self-Debugging Loop: Iterative runs detect errors. If a trial fails, error logs are fed back for automatic code correction or regeneration, ensuring verified tasks.
4. Trajectory Sampling: For each validated synthetic task, teacher agent (GPT-5) rollouts are recorded as “rationale + action” trajectories.
5. Student Model Training: Teacher trajectories (~34k) are used to fine-tune student models (Qwen3-4B, Qwen3-8B) using next-token supervised loss.

Models fine-tuned on these synthetic trajectories achieve relative aggregate AUP gains of 9% (Qwen3-4B) and 12% (Qwen3-8B) over baseline, with statistically significant improvements on iterative, debugging-intensive tasks (Cai et al., 17 Mar 2026).

5. Extensibility and Task Addition Mechanism

MLGym is designed for rapid extensibility with minimal configuration. To add a new task, a user provides an evaluation script (specifying output metric via JSON), starter code, dataset links or HuggingFace dataset, and conda dependency specifications. Task registration into the environment is performed via a succinct Python call:

from mlgym.registry import register_task register_task( id="NewTask-v0", entry_point="mlgym.envs.newtask:NewTaskEnv", config={ "starter_code": "newtask/baseline.py", "dataset": "path/to/data", "eval_script": "eval_newtask.py", "timeout_minutes": 30, }, )

Upon instantiation, MLGym handles workspace setup, dependency installation, and exposure of validate/submit commands. Agents interact with new tasks using standard Gym semantics, facilitating reproducibility and benchmark growth without bespoke environment design (Nathani et al., 20 Feb 2025).

6. Limitations, Observed Gaps, and Future Directions

Current frontier LLM agents consistently demonstrate Level 1 research capability (baseline recovery and improvement, mostly via hyperparameter search), but none reach Level 2 (independent derivation of SOTA without prior code exposure) or introduce genuinely novel algorithms or architectures (Nathani et al., 20 Feb 2025). Key limitations are:

• Long-Horizon Tasks: Challenges persist in language modeling and RL tasks due to lengthy contexts and credit assignment.
• Game Theory Tasks: Difficulties arise in strategy synthesis under non-differentiable payoffs and extended planning requirements.
• Scientific Novelty: No clear, automatable metric for hypothesis generation or research novelty is presently implemented.

Reported limitations in (Cai et al., 17 Mar 2026) underscore possible benchmark format alignment effects and the lack of direct generalization evaluation to other ML agent benchmarks. Pipeline ablations (e.g., dataset grounding, self-debug, selective filtering) are untested.

Future research directions include hierarchical memory for multi-step campaigns, sub-agent architectures for specialized research subtasks, interdisciplinary benchmarks, and formalization of scientific contribution metrics. Cross-benchmark transferability and integration of literature-grounded hypothesis generation remain open avenues (Nathani et al., 20 Feb 2025, Cai et al., 17 Mar 2026).