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

TL;DR

The paper introduces a learning-based paradigm for autonomous ML engineering, where an LLM agent learns through online reinforcement learning to iteratively design and optimize ML solutions. It proposes a three-part training framework—exploration-enriched fine-tuning, step-wise RL, and an agentic ML-specific reward module—and demonstrates this approach with ML-Agent, a 7B-sized Qwen-2.5 model trained on 9 ML tasks. ML-Agent outperforms a 671B-baseline DeepSeek-R1 and shows robust cross-task generalization to held-out tasks, while achieving continuous performance gains through RL. The work highlights how experience-driven learning and unified reward signals can reduce reliance on manual prompt engineering and accelerate autonomous ML engineering in real-world settings.

Abstract

The emergence of large language model (LLM)-based agents has significantly advanced the development of autonomous machine learning (ML) engineering. However, most existing approaches rely heavily on manual prompt engineering, failing to adapt and optimize based on diverse experimental experiences. 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. Remarkably, despite being trained on merely 9 ML tasks, our 7B-sized ML-Agent outperforms the 671B-sized DeepSeek-R1 agent. Furthermore, it achieves continuous performance improvements and demonstrates exceptional cross-task generalization capabilities.

Figures (9)

• Figure 1: Overview of LLM agent interaction with environment for autonomous machine learning. The environment contains editable code files and an interpreter to execute code. In each experience, the agent takes actions (e.g., add BN layers), and the environment returns feedback (e.g., execution results). Our proposed ML-Agent learn from previous experiences and improves itself by training.
• Figure 2: Overview of the agentic ML training framework, introducing (1) exploration-enriched fine-tuning for diverse action pool creation, (2) a step-wise RL paradigm for efficient experience collection using expert trajectories, and (3) an agentic ML-specific reward module for various ML feedback handling and task-specific performance evaluation.
• Figure 3: Comparison results with AIDE framework on 13 tasks. We plot the normalized relative gain for each task. Specifically, we normalize the relative gain of each task via dividing by our ML-Agent value.
• Figure 4: Step-wise RL is more efficient than episode-wise RL on both held-in and held-out tasks. Both RL training methods begin with ML-Agent-SFT. The gap of red points is 5 steps while the gap of blue triangles is 1 step.
• Figure 5: Exploration-enriched fine-tuning is crucial for RL training. "N/A" means the training based on the model fails to generate valid results.