MLE-Dojo: Entornos Interactivos para Potenciar Agentes de LLM en Ingeniería de Aprendizaje Automático

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

Resumen

Presentamos MLE-Dojo, un framework estilo Gym para el aprendizaje por refuerzo sistemático, evaluación y mejora de agentes autónomos basados en modelos de lenguaje grande (LLM) en flujos de trabajo iterativos de ingeniería de aprendizaje automático (MLE). A diferencia de los benchmarks existentes que dependen principalmente de conjuntos de datos estáticos o evaluaciones de un solo intento, MLE-Dojo proporciona un entorno interactivo que permite a los agentes experimentar, depurar y refinar soluciones de manera iterativa a través de bucles de retroalimentación estructurados. Construido sobre más de 200 desafíos reales de Kaggle, MLE-Dojo abarca diversas tareas abiertas de MLE cuidadosamente seleccionadas para reflejar escenarios de ingeniería realistas, como procesamiento de datos, búsqueda de arquitecturas, ajuste de hiperparámetros y depuración de código. Su entorno completamente ejecutable permite un entrenamiento integral de agentes mediante ajuste fino supervisado y aprendizaje por refuerzo, facilitando la experimentación iterativa, el muestreo realista de datos y la verificación de resultados en tiempo real. Evaluaciones exhaustivas de ocho LLM de vanguardia revelan que, aunque los modelos actuales logran mejoras iterativas significativas, aún presentan limitaciones importantes en la generación autónoma de soluciones a largo plazo y en la resolución eficiente de errores complejos. Además, la arquitectura flexible y extensible de MLE-Dojo integra sin problemas diversas fuentes de datos, herramientas y protocolos de evaluación, permitiendo de manera única el ajuste de agentes basados en modelos y promoviendo la interoperabilidad, escalabilidad y reproducibilidad. Hacemos público nuestro framework y benchmarks para fomentar la innovación impulsada por la comunidad hacia la próxima generación 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.