ML-Agent: Reinforcing LLM Agents for Autonomous Machine Learning Engineering

The emergence of large language model (LLM)-based agents has significantly advanced the development of autonomous machine learning (ML) engineering. However, the dominant prompt-based paradigm exhibits limitations: smaller models lack the capacity to learn from execution trajectories for generalization, while large proprietary models incur high computational overhead, restricting accessibility and scalability. Focusing on this, for the first time, we explore the paradigm of learning-based agentic ML, where an LLM agent learns through interactive experimentation on ML tasks using online reinforcement learning (RL). To realize this, we propose a novel agentic ML training framework with three key components: (1) exploration-enriched fine-tuning, which enables LLM agents to generate diverse actions for enhanced RL exploration; (2) step-wise RL, which enables training on a single action step, accelerating experience collection and improving training efficiency; (3) an agentic ML-specific reward module, which unifies varied ML feedback signals into consistent rewards for RL optimization. Leveraging this framework, we train ML-Agent, driven by a 7B-sized Qwen-2.5 LLM for autonomous ML. Despite training on only 9 ML tasks, our 7B-sized ML-Agent achieves comparable performance to agents using much larger proprietary LLMs (e.g., GPT-5) but at significantly lower computational cost, demonstrating strong performance and cross-task generalization.

Building a good machine learning system often takes expert engineers many rounds of trial and error: writing code, running experiments, reading failures, and improving the design. Recent AI assistants can help with this process, but many rely on carefully written instructions rather than truly learning from their own attempts. Smaller assistants often struggle to improve across different tasks, while the strongest commercial systems can be expensive to run. In this work, we teach an AI assistant, ML-Agent, to become better at machine learning engineering through practice. The assistant tries different changes to data processing, model design, and training, receives feedback from each experiment, and learns which choices are useful, invalid, or wasteful. We also make this practice efficient by training it on individual decision points instead of requiring full, time-consuming experiment runs every time. After practicing on only nine tasks, ML-Agent performs well on many new tasks. It outperforms much larger open models and approaches the performance of leading commercial systems, while using far less computing cost. This suggests a more accessible path toward AI tools that can help researchers and engineers build better machine learning solutions.