MLEBench

Description

MLEBench is an environment for evaluating machine learning engineering capabilities. Based on OpenAI's MLE-bench, agents are given Kaggle competitions and must develop, train, and submit ML models. Tasks cover diverse domains including classification, regression, computer vision, NLP, and time series across real Kaggle competition datasets.

Capabilities

Machine learning model development and training
Data exploration and feature engineering
Working with Kaggle competition formats
Iterative model improvement and submission validation

Compute Requirements

Competitions are assigned sandbox compute based on their requirements:

GPU tasks (70 competitions): nvidia-l4 (NVIDIA L4 GPU)
CPU tasks (7 competitions): 16:64 (16 CPUs, 64 GB RAM)

CPU-only competitions: new-york-city-taxi-fare-prediction, tabular-playground-series-dec-2021, tabular-playground-series-may-2022, ventilator-pressure-prediction, playground-series-s3e18, spaceship-titanic, nomad2018-predict-transparent-conductors.

Sandbox backends

Each task runs in one of two sandbox backends, chosen by the backend task-spec field:

modal (default for GPU competitions): a Modal sandbox with gpu (default L4 for GPU competitions), cpu (default 16) and memory_gb (default 64). The sandbox is terminated after sandbox_timeout_hours (default 24, the upstream MLE-bench limit and Modal's max), or after sandbox_idle_timeout_minutes (default 60, null to disable) with no tool call running, which also reaps sandboxes orphaned by a dead env server. Competition data is mounted read-only from the mlebench-data Modal Volume, a copy of gs://gr-datasets/mlebench/ made with upload_modal_volume.py. Needs the secrets modal_token_id and modal_token_secret.
openreward (default for CPU competitions): an OpenReward sandbox of machine_size (default 16:64, or nvidia-l4 for GPU tasks). Needs the secret api_key.

smoke_modal.py runs one competition end to end on Modal.

License

MIT

Tasks

The original MLE-bench contains 82 Kaggle competitions. We currently use a 77-competition subset, plus two splits drawn from it:

Split	Type	Tasks	Description
`or_initial`	`test`	77	70 GPU + 7 CPU competitions
`low`	`test`	22	Upstream MLE-bench's lite split
`cpu_only`	`test`	7	The CPU-only competitions, run in OpenReward sandboxes

Reward Structure

MLEBench is a multi-turn environment. The agent develops an ML model in a sandbox, produces a submission CSV at /home/agent/submission/submission.csv, and calls answer to submit. Reward is the leaderboard-normalized score (score - median) / (best - median), clipped to [0, 1]: 0 at or below the competition's leaderboard median, 1 at the top. Invalid submissions and grading errors give 0. The full grading report, including medal thresholds and any_medal / above_median, is returned as metadata; report medal rates from it, as upstream MLE-bench does.

Data

Competition datasets are mounted read-only from cloud storage. Each competition includes:

train.csv - training data with labels
test.csv - test data for predictions
sample_submission.csv - required submission format
description.md - competition objectives and evaluation metric

Tools

MLEBench provides 6 tools for agents:

bash - Execute shell commands in the sandbox
view - View file contents with optional line range
str_replace - Replace strings in files
insert - Insert content at a specified line
create - Create new files
answer - Submit the final submission CSV for grading

Time Horizon

MLEBench is a multi-turn environment. Agents read competition descriptions, explore data, develop models, train and iterate, validate submissions, and submit for final scoring.

Environment Difficulty

Competitions span a wide range of difficulty, from tabular classification (e.g., spaceship-titanic) to complex computer vision and NLP tasks (e.g., rsna-breast-cancer-detection, lmsys-chatbot-arena).

Other Environment Requirements

There are no further environment requirements; MLEBench works out of the box with the OpenReward endpoint without any external API keys.

Safety

Agents in MLEBench develop ML models in a sandboxed environment. The environment does not present direct safety risks.

To keep grading honest, the agent's tools run as an unprivileged agent user. The answers, grader and validation server live under root-only /private, the uv venv is read-only (pip install goes to /home/agent/.pydeps), and answer reads the submission as the agent, so it cannot be pointed at files the agent couldn't read itself.

Citations

@article{chan2024mlebench,
  title={MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering},
  author={Chan, Jun Shern and Chowdhury, Neil and Jaffe, Oliver and Aung, James and Sherburn, Dane and Mays, Evan and Starace, Giulio and Liu, Kevin and Maksin, Lukasz and Patil, Tejal and Rein, David and Beutel, Alex},
  journal={arXiv preprint arXiv:2410.07095},
  year={2024}
}