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

TL;DR

This paper introduces ML-Agent, a reinforcement learning framework for autonomous ML engineering using a 7B-sized LLM, achieving competitive results with lower computational costs across multiple tasks.

Contribution

It proposes a novel learning-based agentic ML paradigm with exploration-enriched fine-tuning, step-wise RL, and a unified reward module, enabling efficient autonomous ML development.

Findings

ML-Agent trained on 9 tasks performs comparably to larger proprietary LLM-based agents.

The framework reduces computational costs while maintaining strong cross-task generalization.

The approach demonstrates effective reinforcement learning for autonomous ML engineering.

Abstract

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.