MLE-Dojo: Ambientes Interativos para Capacitar Agentes de LLM em Engenharia de Aprendizado de Máquina

MLE-Dojo: Interactive Environments for Empowering LLM Agents in Machine Learning Engineering

May 12, 2025

Autores: Rushi Qiang, Yuchen Zhuang, Yinghao Li, Dingu Sagar V K, Rongzhi Zhang, Changhao Li, Ian Shu-Hei Wong, Sherry Yang, Percy Liang, Chao Zhang, Bo Dai

cs.AI

Resumo

Apresentamos o MLE-Dojo, um framework no estilo Gym para o aprendizado por reforço sistemático, avaliação e aprimoramento de agentes autônomos de modelos de linguagem de grande escala (LLM) em fluxos de trabalho iterativos de engenharia de aprendizado de máquina (MLE). Diferente dos benchmarks existentes que dependem principalmente de conjuntos de dados estáticos ou avaliações de tentativa única, o MLE-Dojo fornece um ambiente interativo que permite aos agentes experimentar, depurar e refinar soluções de forma iterativa por meio de loops de feedback estruturados. Construído com base em mais de 200 desafios reais do Kaggle, o MLE-Dojo abrange diversas tarefas de MLE de natureza aberta, cuidadosamente curadas para refletir cenários realistas de engenharia, como processamento de dados, busca de arquitetura, ajuste de hiperparâmetros e depuração de código. Seu ambiente totalmente executável suporta treinamento abrangente de agentes por meio de ajuste fino supervisionado e aprendizado por reforço, facilitando experimentação iterativa, amostragem realista de dados e verificação de resultados em tempo real. Avaliações extensivas de oito LLMs de ponta revelam que, embora os modelos atuais alcancem melhorias iterativas significativas, eles ainda apresentam limitações importantes na geração autônoma de soluções de longo prazo e na resolução eficiente de erros complexos. Além disso, a arquitetura flexível e extensível do MLE-Dojo integra de forma contínua diversas fontes de dados, ferramentas e protocolos de avaliação, permitindo de maneira única o ajuste de agentes baseados em modelos e promovendo interoperabilidade, escalabilidade e reprodutibilidade. Disponibilizamos nosso framework e benchmarks em código aberto para fomentar a inovação impulsionada pela comunidade em direção à próxima geração de agentes de MLE.

English

We introduce MLE-Dojo, a Gym-style framework for systematically reinforcement learning, evaluating, and improving autonomous large language model (LLM) agents in iterative machine learning engineering (MLE) workflows. Unlike existing benchmarks that primarily rely on static datasets or single-attempt evaluations, MLE-Dojo provides an interactive environment enabling agents to iteratively experiment, debug, and refine solutions through structured feedback loops. Built upon 200+ real-world Kaggle challenges, MLE-Dojo covers diverse, open-ended MLE tasks carefully curated to reflect realistic engineering scenarios such as data processing, architecture search, hyperparameter tuning, and code debugging. Its fully executable environment supports comprehensive agent training via both supervised fine-tuning and reinforcement learning, facilitating iterative experimentation, realistic data sampling, and real-time outcome verification. Extensive evaluations of eight frontier LLMs reveal that while current models achieve meaningful iterative improvements, they still exhibit significant limitations in autonomously generating long-horizon solutions and efficiently resolving complex errors. Furthermore, MLE-Dojo's flexible and extensible architecture seamlessly integrates diverse data sources, tools, and evaluation protocols, uniquely enabling model-based agent tuning and promoting interoperability, scalability, and reproducibility. We open-source our framework and benchmarks to foster community-driven innovation towards next-generation MLE agents.