Title: ML-Agent: Reinforcing LLM Agents for Autonomous Machine Learning Engineering
ArXiv: 2505.23723
Authors: Zexi Liu, Jingyi Chai, Xinyu Zhu, Shuo Tang, Rui Ye, Bo Zhang, Lei Bai, Siheng Chen
Sections: 29
Estimated tokens: 23.3k

## Contents
- 1 Introduction
- 2 Related Work
- 3 Problem Setup and Preliminaries
- 4 Agentic ML Training Framework
  - 4.1 Exploration-enriched fine-tuning
  - 4.2 Step-wise RL paradigm
  - 4.3 Agentic ML-specific reward
- 5 Experiments
  - 5.1 Experimental setups
  - 5.2 Main results
  - 5.3 Ablation Study and Analysis
- 6 Conclusion
- Impact Statement
- References
- Appendix A Problem formulation
- Appendix B Machine Learning Tasks and Data Collection Process
  - B.1 Data Collecting Pipeline for Exploration-Enriched Fine-tuning
  - B.2 Details of Machine Learning Tasks
  - B.3 Details of Data Collection
- Appendix C Experimental Details
  - C.1 Details of Experimental Set-up
  - C.2 Additional Ablation Study and Analysis
  - C.3 Additional ML tasks results
  - C.4 Cost analysis with amortization
  - C.5 Case study
    - C.5.1 Task Description for denoising-dirty-documents
    - C.5.2 Initial Script for denoising-dirty-documents
    - C.5.3 ML-Agent Trajectory on denoising-dirty-documents
    - C.5.4 Analysis

## Abstract

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.

## 1 Introduction

Machine Learning (ML) engineering is a critical yet labor-intensive process, requiring expert researchers to invest significant time—potentially days or even months—designing architectures, tuning parameters, and iteratively refining models through trial and error (Bergstra and Bengio, 2012).
This challenge has sparked an ambitious vision of autonomous ML: building autonomous AI systems that independently orchestrate the entire ML lifecycle, from conceptual design and code implementation to refinement.

Fortunately, the advent of LLM-based agents, equipped with capabilities of interaction (Du et al., 2023; Pang et al., 2024), coding (Hong et al., 2023; Qian et al., 2023; Hu et al., 2024) and tool-calling (Masterman et al., 2024), has propelled us significantly closer to realizing this vision (Agentic ML) (Huang et al., 2023; Chan et al., 2024).
Unlike traditional automated ML with pre-defined limited search and action spaces (Tang et al., 2024; LeDell and Poirier, 2020; Jin et al., 2023), these LLM agents, when provided with instructions in natural language, can autonomously propose effective actions, generate executable codes, and iteratively improve solutions based on environmental feedback (Huang et al., 2023; Jiang et al., 2025).
For example, AIDE (Jiang et al., 2025) and ML-Master (Liu et al., 2025a) both leverage LLM agents together with experimental environments to automate ML process.

Currently, the dominant paradigm in agentic ML relies on prompt-based design, where agents are constructed through heuristic prompt engineering. This approach offers practical advantages, as it allows rapid deployment without parameter updates or extensive retraining. However, it also exhibits notable limitations: when driven with smaller language models, such agents lack the capacity to learn from and internalize execution trajectories, causing limited generalization across diverse tasks; conversely, when implemented with large-scale proprietary models, the paradigm incurs substantial computational overhead and resource consumption, thereby restricting accessibility and undermining sustainable scalability (Belcak et al., 2025).

To address these limitations, we propose moving beyond the prompt-based paradigm toward a new research trajectory: learning-based agentic ML. In this paradigm, agents are no longer constrained to static prompt instructions but instead learn adaptively from task-solving trajectories via reinforcement learning (RL). Such a formulation empowers agents to systematically explore diverse strategies, accumulate knowledge across successive runs, and progressively refine their decision-making processes (Xiong et al., 2024). Importantly, this learning-based approach endows even relatively small language models with the capacity to achieve strong generalization, while substantially reducing computational and resource demands. As a result, it opens a more accessible, efficient, and sustainable path for advancing the frontier of autonomous machine learning.

While being straightforward, employing RL to train autonomous ML agents poses three key challenges.
(1) Limited exploration: agents often propose similar actions for the same ML task across runs, leading to narrow exploration trajectories in RL (Park et al., 2024).
(2) Slow experience collection: ML experiments can take minutes to hours, making RL data gathering inefficient and thus limiting feedback-driven training samples (Chan et al., 2024).
(3) Complex reward design: agentic ML involves various outcomes, such as task-specific metrics, out-of-memory failures, and compilation errors. This requires a unified reward function to reconcile varied feedback signals (Eschmann, 2021).

In response to these challenges, we propose a novel agentic ML training framework, the first designed to train LLM agents for autonomous ML engineering using RL.
This framework enables agents to explore diverse ML trajectories, collect rewards efficiently, and iteratively enhance their capabilities through learned experience.
(1) To improve exploration diversity, we introduce exploration-enriched fine-tuning, generating a diverse action pool from fast-executable ML tasks to finetune agents for broader RL exploration.
(2) To accelerate experience collection, we design a step-wise RL paradigm, evaluating atomic actions using expert trajectories as single-step queries, significantly boosting training efficiency.
(3) To tackle reward design, we develop an agentic ML-specific reward module that dynamically handles errors (e.g., runtime failures) and quantifies performance via normalized, task-specific metrics (e.g., accuracy gains).

By leveraging our proposed agentic ML training framework, we train ML-Agent, an agent driven by a 7B-sized Qwen2.5 LLM for autonomous ML.
During training, our ML-Agent can efficiently explore the environment, learn from experience, and achieve continuous performance improvement through iterative exploration across various ML tasks.
Surprisingly, despite its modest size and training on only 9 ML tasks, ML-Agent demonstrates strong performance and cross-task generalization, outperforming 671B-sized DeepSeek-R1 agent on 3 held-in and 10 held-out tasks across diverse data modalities and objectives. Notably, it achieves results comparable to agents using the most advanced proprietary LLMs (GPT-5) but at significantly lower computational cost.

In summary, our work makes the following significant contributions to the field:

- •
We introduce a new paradigm for autonomous ML: learning-based agentic ML, where an LLM agent learns through interactive experimentation on ML tasks via reinforcement learning.
- •
We propose a novel training framework for agentic ML, which incorporates three technical designs: exploration-enriched fine-tuning, step-wise RL, and agentic ML-specific rewards.
- •
Extensive experiments show that despite training on only 9 ML tasks, our 7B-sized ML-Agent surpasses agents driven by much larger LLMs and even matches agents driven by proprietary LLMs (e.g., GPT-5) with much lower cost.

## 2 Related Work

Autonomous Machine Learning.
Autonomous machine learning aims to automate the manual and expertise-intensive aspects of machine learning, including data preprocessing, model selection and hyperparameter tuning.
Autonomous machine learning has evolved from classical hyperparameter and pipeline search to agentic frameworks powered by large language models. Classical autonomous machine learning frameworks focus on automating model selection, hyperparameter optimization, and pipeline construction within a fixed search space (Tang et al., 2024; Olson and Moore, 2016; Feurer et al., 2022; Mohr et al., 2018; Erickson et al., 2020; Liu et al., 2020). For example, AutoGluon-Tabular (Erickson et al., 2020) ensembles multiple models and stackings to deliver state-of-the-art performance on tabular data with minimal user effort.
These classical autonomous machine learning works remain constrained by predefined search spaces and static configurations, lacking the adaptability and continuous learning capabilities.

LLM Agents in Autonomous Machine Learning. Recent advancements in LLMs have empowered them to autonomously generate and refine machine learning solutions, opening new possibilities in machine learning. Methods such as AutoML-GPT (Zhang et al., 2023b) and MLCopilot (Zhang et al., 2023a) prompt LLMs to automate the entire machine learning pipeline, where MLCopilot introduces past experience retrieval to help decision-making. AIDE (Jiang et al., 2025) and ML-Master (Liu et al., 2025a) focus on optimizing the ML engineering process through iterative search and refinement strategies. Other works like AutoKaggle (Li et al., 2024) and AutoML-Agent (Trirat et al., 2024) employ a multi-agent framework to address ML problems.
However, these approaches are fundamentally constrained by a prompt-based paradigm. While agents may leverage past experience, their underlying models are not trained on these interaction histories. Consequently, their problem-solving strategies remain static and rely on costly advanced models. This limitation motivates our shift toward a learning-based paradigm where agents adapt and improve over time.

Reinforcement Learning for LLMs.
Reinforcement learning (RL) significantly enhances the ability of LLMs, particularly in preference alignment and complex reasoning (Xu et al., 2025; Wang et al., 2024; Zheng et al., 2023).
By facilitating exploration and exploitation, RL trains LLMs to adapt and improve their policy based on feedback, thus refining their performance in dynamic environments. One line of work is preference optimization (Kaufmann et al., 2023), with methods such as Reinforcement Learning from Human Feedback (RLHF)(Ouyang et al., 2022). RL is also utilized to train LLMs for complex reasoning tasks (Guo et al., 2025a; Liu et al., 2025b).
Another line of research involves training LLM agents for specific tasks using RL (Zhang et al., 2025). For example, IPR (Xiong et al., 2024) and AgentQ (Putta et al., 2024) use DPO (Rafailov et al., 2023) to iteratively refine their policy. While StarPO (Wang et al., 2025) discusses the multi-turn reinforcement learning considering episode-wise reward. However, applying RL to train LLM agents for autonomous machine learning remains unexplored.

## 3 Problem Setup and Preliminaries

Problem Formulation. Agentic ML leverages an LLM agent to autonomously orchestrate the ML lifecycle by interacting with the experimental environment.
This environment includes editable task-related code files together with an interpreter executing code and provides explicit experimental feedback (e.g., code execution results or error messages). Given an initial ML task specification (e.g., dataset description and evaluation metric), the agent begins interacting with the environment to iteratively refine its solution. At each step, the agent takes actions (e.g., add BN layers in the model architecture) and receives feedback (e.g., code execution output or error messages) from the environment. This loop continues until a step or time limit is reached. We follow the action space from prior work (Huang et al., 2023) (The details are provided in Table [5](#A2.T5)).

Agentic ML as a MDP. We format agentic ML as a Markov Decision Process (MDP) $\mathcal{M}=\left(\mathcal{S},\mathcal{A},\mathcal{P}\right)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ the action space and $\mathcal{P}$ the state transition dynamics. Let the environment feedback at time $t$ be $f_{t}\in\mathcal{F}$, where $\mathcal{F}$ denotes the feedback space. We employ a history-based state representation $s_{t}=(s_{0},a_{0},f_{0},a_{1},f_{1},\ldots,a_{t-1},f_{t-1})$ to capture richer contextual information from past feedback, in which $s_{0}$ encodes the initial ML task specification and each pair $(a_{i},f_{i})$ represents the agent’s action and corresponding environment feedback. The agent policy $\pi_{\theta}$ generates an action $a_{t}\in\mathcal{A}$ conditioned on current state $s_{t}$, forming a trajectory of interactions $\tau=(s_{0},a_{0},s_{1},\ldots,a_{n-1},s_{n})$. Note that $\theta$ is the LLM’s parameters within the agent and $n$ is the trajectory length. The goal is to maximize the expected trajectory reward:

$$ $\small\mathcal{J}(\theta)=\mathbb{E}_{\tau\sim\pi_{\theta}}\left[R(\tau)\right],$ (1) $$

where the reward function $R(\tau)$ denotes the cumulative reward over the entire trajectory.

Challenges.
Although the formulation of agentic ML is relatively straightforward, employing RL to train LLM agents for autonomous machine learning poses several key challenges, including:
(1) Limited exploration. Agents often repeat similar actions across episodes, narrowing their exploration and limiting their ability to discover innovative ML solutions.
(2) Slow experience collection. ML experiments can take minutes to hours, slowing down the online data collection process for RL training.
(3) Complex reward design. Agentic ML produces varied outcomes (e.g., execution results or resource errors), making it challenging to design a unified reward function that effectively guides the agent.
The subsequent section presents our agentic ML training framework designed to overcome these challenges, with the overall architecture illustrated in Figure [1](#S4.F1).

## 4 Agentic ML Training Framework

Our agentic ML training framework is designed to train LLM agents for autonomous machine learning.
As shown in Figure [1](#S4.F1), it comprises three key steps for effective learning.
First, exploration-enriched fine-tuning builds a diverse action pool to enhance RL exploration. Second, a step-wise RL paradigm uses expert trajectories as single-step queries to accelerate experience collection in RL. Third, an agentic ML-specific reward module handles errors and quantifies agentic ML task-specific performance. These steps sequentially enable diverse exploration, efficient training, and unified feedback, enabling agents to iteratively improve agentic ML performance across varied ML tasks.

Figure: Figure 1: 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.
Refer to caption: 2505.23723v2/x1.png

### 4.1 Exploration-enriched fine-tuning

In agentic ML, limited exploration hinders autonomous machine learning workflows. Agents often repeat similar actions (e.g., small code edits) across episodes, leading to narrow exploration and preventing the discovery of innovative architectures or optimization strategies.

To address this, we introduce exploration-enriched fine-tuning with an automated data collection pipeline. It organizes ML optimization strategies into 3 semantic categories: data, model and learning. For each category, an LLM generates a large set of candidate ideas and an embedding-based diversity filter selects a compact and diverse pool. During trajectory generation, the system automatically samples 1–3 categories, shuffles their order, and draws one idea from each corresponding pool to form the initial action sequence (See Appendix [B.1](#A2.SS1)). An expert LLM with policy $\pi_{e}$ then executes the full workflow on fast-executable ML tasks, producing expert trajectories $\mathcal{D}=\{\tau^{(i)}\}_{i=1}^{|\mathcal{D}|}$. We fine-tune the agent policy $\pi_{\theta}$ via supervised fine-tuning (SFT):

$$ $\begin{split}\mathcal{L}_{\text{SFT}}(\theta)&=-\mathbb{E}_{\tau\sim\mathcal{D}}\left[\log P_{\pi_{\theta}}(\tau|s_{0})\right]\\ &=-\mathbb{E}_{\tau\sim\mathcal{D}}\left[\log\prod_{t=0}^{n-1}\pi_{\theta}(a_{t}|s_{t})\right]\\ &=-\mathbb{E}_{\tau\sim\mathcal{D}}\bigg[\sum_{t=0}^{n-1}\log\pi_{\theta}(a_{t}|s_{t})\bigg].\end{split}$ (2) $$

This exploration-enriched fine-tuning approach preserves action format compliance while enabling agents to learn diverse strategies, significantly broadening the exploration scope in subsequent RL.

### 4.2 Step-wise RL paradigm

Objective. Due to the time-consuming nature of AI experiments, directly applying RL methods (e.g., PPO) is impractical, as sampling a single trajectory during rollout takes hours. To address this issue, we propose a step-wise RL approach that reformulates the objective function equation [1](#S3.E1), where we sample only a single step of action during the rollout phase instead of the entire trajectory. This approach extensively reduces the computational cost of the rollout phase and makes the overall training process more efficient.
Specifically, we expand equation [1](#S3.E1) into steps according to the state distribution $d^{\pi_{\theta}}(s)$:

$$ $\small\mathcal{J}(\theta)=\sum_{t=0}^{n-1}\sum_{s_{t}\in\mathcal{S}}d^{\pi_{\theta}}(s_{t})\left[\sum_{a_{t}\in\mathcal{A}}\pi_{\theta}\left(a_{t}\big|s_{t}\right)R(s_{t},a_{t})\right],$ (3) $$

where $R(s_{t},a_{t})$ is the step-wise reward at time $t$, and $d^{\pi_{\theta}}(\cdot)$ is the state distribution at time $t$ under policy $\pi_{\theta}$. This distribution can be calculated recursively based on the policy $\pi_{\theta}$ and the state transition dynamics $\mathcal{P}$; see Appendix [A](#A1) for details. The time-consuming components in equation [3](#S4.E3) include: 1) $d^{\pi_{\theta}}(s_{t})$, which involves multiple state transition dynamics from $s_{t}$ to $s_{t+1}$, and 2) $R(s_{t},a_{t})$, where the reward is determined based on feedback from the environment (e.g., code execution platform). Since $d^{\pi_{\theta}}(\cdot)$ relies on $\pi_{\theta}$, the trajectory sampling process operates repeatedly in standard RL training, making the computational cost even higher. However, using $d^{\pi_{\theta}}$ to sample state distribution is not necessary for two reasons: 1) $\pi_{\theta}$ poorly aligns with the environment format during the early stage of RL training, hindering effective state exploration; 2) Once $\pi_{\theta}$ can interact properly with the environment, the set of states it could explore tends to vary only slightly as $\pi_{\theta}$ updates. Hence, we sample the states from a states pool according to a fixed expert distribution $d^{\pi_{e}}(s_{t})$, which forms the step-wise objective function

$$ $\begin{split}\mathcal{J}_{\text{step}}(\theta)&=\sum_{s_{t}\in\mathcal{S}}d^{\pi_{e}}(s_{t})\left[\sum_{a_{t}\in\mathcal{A}}\pi_{\theta}\left(a_{t}\big|s_{t}\right)R(s_{t},a_{t})\right]\\ &=\mathbb{E}_{s_{t}\sim d^{\pi_{e}},a_{t}\sim\pi_{\theta}(\cdot|s_{t})}\left[R(s_{t},a_{t})\right].\end{split}$ (4) $$

This objective function $\mathcal{J}_{step}(\theta)$ reformulates multi-step trajectory RL into step-wise training. This reformulation offers two advantages: 1) The state sampling process is decoupled from the RL of the model. This allows us to directly sample states from a pre-collected set and avoids expensive online sampling during training, significantly reducing the overall training time. 2) The state sampling process is performed before RL training, rather than during the rollout phase. This enables us to perform extensive sampling from the expert distribution, making training more scalable.

Training approach. Based on the step-wise RL formulation, our goal is to maximize the expected reward $R(s_{t},a_{t})$ shown in equation [4](#S4.E4) according to the state distribution $d^{\pi_{e}}$ and $\pi_{\theta}$. This aligns with the approach used in RLVR methods (Guo et al., 2025b), where the policy represents a token generation process and $R(s_{t},a_{t})$ is the outcome reward of $\pi_{\theta}(a_{t}|s_{t})$. Hence, any RL training approach can be applied to this objective $\mathcal{J}_{\text{step}}(\theta)$. For our implementation, we choose PPO (Schulman et al., 2017) as the training algorithm because of its widespread use and proven effectiveness. Specifically, suppose we expand the token generating process of $\pi_{\theta}(a_{t}|s_{t})$, our PPO loss function can be defined as follows:

$$ $\displaystyle\mathcal{J}_{\text{step}}^{\text{PPO}}(\theta)=\mathbb{E}_{s\sim d^{e},o_{\leq i}\sim\pi_{\theta_{\text{old}}}(\cdot\mid s)}\Bigg[\min\Bigg(\frac{\pi_{\theta}(o_{i}\mid s,o_{<i})}{\pi_{\theta_{\text{old}}}(o_{i}\mid s,o_{<i})}\hat{A}_{i},$ (5) $\displaystyle\text{clip}\Bigg(\frac{\pi_{\theta}(o_{i}\mid s,o_{<i})}{\pi_{\theta_{\text{old}}}(o_{i}\mid s,o_{<i})},1-\varepsilon,1+\varepsilon\Bigg)\hat{A}_{i}\Bigg)\Bigg],$ $$

where $o_{i}$ is the $i$th token of $a_{t}$ and $\hat{A}_{i}$ is an estimator of the advantage at the token generation step $i$.

### 4.3 Agentic ML-specific reward

Having enabled efficient RL for agentic ML via the step-wise RL paradigm, the next crucial step is to convert the varied feedback into a unified, meaningful reward. While numerical metrics like validation accuracy or loss naturally serve as RL rewards, non-numerical feedback, such as compilation errors or out-of-memory failures, must be carefully incorporated to ensure the reward is coherent.

To address this, we propose an agentic ML-specific reward module that dynamically processes these diverse signals while quantifying performance improvements through scaled task-specific metrics. The key idea is to translate every execution outcome into a unified scalar value.
Define $\mathcal{A}_{\text{valid}}$
as valid actions, $\mathcal{A}_{\text{edit}}\subset\mathcal{A}_{\text{valid}}$
as editing actions for ML code, $\mathcal{F}_{\text{error}}$ as error feedback (e.g., compilation failures), $\mathcal{F}_{\text{corner}}$ as corner cases (e.g., resource exhaustion), and $\mathcal{F}_{\text{success}}$
as successful executions. Let $m_{t}$ be the task-specific metric(^1^11We follow the official Kaggle evaluation protocol which defines a scalar metric for each ML tasks.) at state $s_{t}$ (e.g., loss or accuracy), with $m_{\mathrm{init}}$ and $m_{\mathrm{best}}$ as the baseline and best human-achievable scores. The reward $R(s_{t},a_{t})$ is:

$$ $\small{R(s_{t},a_{t})=\begin{cases}-1&,\text{if }a_{t}\notin\mathcal{A}_{\text{valid}}\text{ or }f_{t}\in\mathcal{F}_{\text{error}}\\ 0&,\text{if }a_{t}\in\mathcal{A}_{\text{valid}}/\mathcal{A}_{\text{edit}}\text{ or }f_{t}\in\mathcal{F}_{\text{corner}}\\ \frac{m_{t+1}-m_{t}}{m_{\mathrm{best}}-m_{\mathrm{init}}}&,\text{if }a_{t}\in\mathcal{A}_{\text{edit}}\text{ and }f_{t}\in\mathcal{F}_{\text{success}}.\end{cases}}$ (6) $$

This reward module handles all possible agentic ML scenarios:
(1) Invalid actions or errors receive -1 to penalize faulty outputs;
(2) Valid non-editing actions or corner cases receive 0 as a neutral acknowledgment of legitimacy while recognizing external constraints;
(3) Success edits yield a scaled metric improvement for task-driven refinement.
By unifying penalties for errors, neutrality for non-editing actions, and task-driven rewards for edits, the module provides consistent, informative feedback for iterative refinement and continuous improvement across diverse ML tasks.

## 5 Experiments

### 5.1 Experimental setups

Training. For training data collection, we adopt a GPT-4o-mini-driven (OpenAI, 2024) agent scaffolded by MLAB (Huang et al., 2023) to interact with the MLAgentBench environment and generate expert trajectories.
Collecting step-wise expert trajectories for autonomous ML agents is computationally expensive, as each trajectory requires executing a full ML pipeline, including data preprocessing, model training, and evaluation.
Under a fixed compute budget, we therefore prioritize fast-executable tasks while maintaining broad coverage of heterogeneous ML settings. Specifically, we collect 10k expert trajectories across 9 ML tasks, comprising 4 tasks from MLAgentBench and 5 from MLE-Bench (Chan et al., 2024).
These tasks jointly span multiple data modalities (image, tabular, text, and graph),
task types (classification and regression),
and evaluation metrics (e.g., accuracy, AUC, log-loss, and MAE),
aiming to balance computational feasibility with exposure to diverse and realistic ML pipelines rather than optimizing for any single benchmark.
Each trajectory is limited to 15 steps and 30 minutes of runtime.
Additional data collection details are provided in Appendix [B](#A2).
For exploration-enhanced fine-tuning, we train Qwen2.5-7B (Yang et al., 2024) using these 10k expert trajectories via supervised fine-tuning (SFT). For step-wise RL, we select 10k states sampled from expert trajectories to further train the SFT model using Proximal Policy Optimization (PPO). All training is conducted on 8 A100 GPUs. The fine-tuning stage runs for 2 epochs with a learning rate of $2e-5$, while the RL stage runs for 1 epoch with an actor learning rate of $1e-6$ and a critic learning rate of $1e-5$. See additional training details in Appendix [C.1](#A3.SS1).
Importantly, our training procedure does not rely on any special property of these selected tasks and can be directly scaled to more tasks given additional computational budget.
We study the effect of the number of training tasks in Section [5.3](#S5.SS3).

**Table 1: Comparing 7B ML-Agent with baselines across different agent frameworks driven by proprietary/open-source LLMs on 3 held-in tasks (included in training) and 10 held-out tasks (unseen during training) from MLE-bench. For each task, we report average performance gain (%) over 8 trajectories. The best open-source LLM-driven agent performance is marked in bold.**
| Method | Model | #Params | cifar10 | house | feedback | denoising | leaf | statoil | whale |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Prompt-Based Method |  |  |  |  |  |  |  |  |  |
| MLAB | Qwen2.5-7B-Instruct | 7B | $1.37\pm 7.98$ | $0.23\pm 0.43$ | $1.39\pm 3.42$ | $2.10\pm 40.99$ | $2.52\pm 14.85$ | $-6.32\pm 8.68$ | $12.25\pm 31.63$ |
| Qwen3-235B | 235B | $57.61\pm 7.35$ | $3.01\pm 1.53$ | $6.70\pm 3.98$ | $62.60\pm 11.12$ | $-2.12\pm 1.97$ | $-16.36\pm 16.38$ | $26.68\pm 45.31$ |  |
| DeepSeek-R1 | 671B | $28.96\pm 23.63$ | $3.45\pm 2.16$ | $5.53\pm 5.10$ | $8.83\pm 3.36$ | $4.85\pm 12.01$ | $0.04\pm 4.08$ | $33.44\pm 45.96$ |  |
| GPT-5 | N/A | $61.46\pm 36.20$ | $12.15\pm 4.10$ | $12.74\pm 4.08$ | $66.00\pm 25.14$ | $-45.63\pm 116.99$ | $-6.43\pm 5.71$ | $89.59\pm 31.34$ |  |
| Gemini-2.5-Pro | N/A | $16.78\pm 29.14$ | $1.16\pm 2.06$ | $0.10\pm 0.00$ | $37.85\pm 31.02$ | $-4.38\pm 7.99$ | $-4.26\pm 5.48$ | $22.38\pm 37.13$ |  |
| AIDE | Qwen2.5-7B-Instruct | 7B | $11.36\pm 16.97$ | $2.42\pm 2.49$ | $7.52\pm 5.66$ | $7.33\pm 12.94$ | $-4.75\pm 2.95$ | $-4.33\pm 10.71$ | $0.52\pm 0.43$ |
| Qwen3-235B | 235B | $-0.10\pm 2.71$ | $2.04\pm 2.12$ | $11.10\pm 6.76$ | $41.65\pm 16.38$ | $4.75\pm 16.08$ | $-2.89\pm 4.31$ | $8.26\pm 20.27$ |  |
| DeepSeek-R1 | 671B | $72.55\pm 1.99$ | $5.35\pm 2.43$ | $13.07\pm 2.98$ | $33.23\pm 22.78$ | $-10.25\pm 18.29$ | $-4.54\pm 4.99$ | $30.77\pm 36.23$ |  |
| GPT-5 | N/A | $76.53\pm 8.08$ | $22.15\pm 4.72$ | $8.77\pm 9.44$ | $77.38\pm 10.77$ | $31.50\pm 23.20$ | $-9.18\pm 10.99$ | $26.42\pm 35.38$ |  |
| Gemini-2.5-Pro | N/A | $53.59\pm 31.32$ | $11.13\pm 6.56$ | $9.44\pm 3.27$ | $62.72\pm 22.94$ | $-84.25\pm 223.17$ | $-6.08\pm 8.29$ | $56.45\pm 42.00$ |  |
| ML-Master | Qwen2.5-7B-Instruct | 7B | $1.03\pm 1.49$ | $0.00\pm 0.00$ | $0.10\pm 0.00$ | $2.44\pm 6.23$ | $-1.38\pm 0.00$ | $-3.99\pm 7.06$ | $1.12\pm 0.49$ |
| DeepSeek-R1 | 671B | $73.43\pm 1.71$ | $18.25\pm 5.92$ | $12.07\pm 0.04$ | $14.56\pm 7.21$ | $-14.75\pm 35.32$ | $-2.78\pm 4.72$ | $33.39\pm 34.94$ |  |
| GPT-5 | N/A | $71.64\pm 6.34$ | $22.3\pm 4.37$ | $10.54\pm 12.27$ | $10.96\pm 27.28$ | $23.88\pm 20.12$ | $-2.48\pm 5.23$ | $67.07\pm 19.30$ |  |
| Learning-Based Method |  |  |  |  |  |  |  |  |  |
| ML-Agent(Ours) | 7B | $33.80\pm 11.27$ | $6.77\pm 3.47$ | $13.47\pm 2.36$ | $52.38\pm 10.07$ | $13.87\pm 20.39$ | $1.41\pm 8.16$ | $72.89\pm 24.43$ |  |

| Method | Model | #Params | learning | detecting | spooky | jigsaw | us | tabular | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Prompt-Based Method |  |  |  |  |  |  |  |  |  |
| MLAB | Qwen2.5-7B-Instruct | 7B | $1.23\pm 1.11$ | $0.51\pm 0.43$ | $-0.46\pm 3.94$ | $-0.06\pm 0.15$ | $3.75\pm 2.82$ | $0.04\pm 0.06$ | $1.43$ |
| Qwen3-235B | 235B | $0.30\pm 2.53$ | $1.02\pm 1.53$ | $0.80\pm 1.55$ | $0.01\pm 0.03$ | $1.96\pm 2.19$ | $-0.07\pm 0.26$ | $10.93$ |  |
| DeepSeek-R1 | 671B | $0.05\pm 0.89$ | $0.25\pm 0.25$ | $0.89\pm 1.81$ | $0.00\pm 0.00$ | $2.67\pm 2.56$ | $-0.13\pm 0.59$ | $6.83$ |  |
| GPT-5 | N/A | $4.36\pm 2.40$ | $11.20\pm 8.41$ | $6.79\pm 7.56$ | $0.00\pm 0.00$ | $23.38\pm 4.64$ | $0.23\pm 0.02$ | ${18.14}$ |  |
| Gemini-2.5-Pro | N/A | $0.00\pm 0.00$ | $0.13\pm 0.00$ | $0.04\pm 0.02$ | $0.00\pm 0.00$ | $0.13\pm 0.36$ | $0.00\pm 0.00$ | $5.38$ |  |
| AIDE | Qwen2.5-7B-Instruct | 7B | $-9.78\pm 30.17$ | $-0.38\pm 1.45$ | $0.07\pm 0.00$ | $0.01\pm 0.01$ | $0.00\pm 0.00$ | $0.08\pm 0.09$ | $0.77$ |
| Qwen3-235B | 235B | $2.37\pm 3.08$ | $0.43\pm 0.70$ | $0.96\pm 2.36$ | $-12.15\pm 21.06$ | $0.51\pm 3.74$ | $0.00\pm 0.00$ | $4.38$ |  |
| DeepSeek-R1 | 671B | $1.38\pm 0.63$ | $0.31\pm 0.92$ | $0.36\pm 0.78$ | $0.01\pm 0.04$ | $5.78\pm 3.05$ | $0.14\pm 0.07$ | $11.40$ |  |
| GPT-5 | N/A | $4.51\pm 2.96$ | $0.13\pm 0.00$ | $4.25\pm 6.74$ | $0.14\pm 0.21$ | $29.69\pm 12.44$ | $0.11\pm 0.10$ | ${20.95}$ |  |
| Gemini-2.5-pro | N/A | $7.35\pm 0.60$ | $0.74\pm 1.76$ | $4.34\pm 7.51$ | $0.04\pm 0.04$ | $31.92\pm 28.18$ | $0.13\pm 0.14$ | $11.35$ |  |
| ML-Master | Qwen2.5-7B-Instruct | 7B | $1.79\pm 0.00$ | $0.26\pm 0.29$ | $-0.04\pm 0.00$ | $-0.02\pm 0.04$ | $-0.02\pm 0.00$ | $0.00\pm 0.00$ | $0.10$ |
| DeepSeek-R1 | 671B | $3.03\pm 3.29$ | $0.00\pm 0.00$ | $4.01\pm 4.70$ | $-0.04\pm 0.10$ | $29.27\pm 22.65$ | $0.22\pm 0.10$ | $13.13$ |  |
| GPT-5 | N/A | $6.38\pm 1.17$ | $0.79\pm 2.09$ | $10.41\pm 7.10$ | $0.35\pm 0.31$ | $26.49\pm 3.20$ | $0.25\pm 0.10$ | ${19.12}$ |  |
| Learning-Based Method |  |  |  |  |  |  |  |  |  |
| ML-Agent(Ours) | 7B | $1.91\pm 1.24$ | $1.74\pm 0.97$ | $1.76\pm 5.39$ | $0.01\pm 0.04$ | $12.96\pm 8.16$ | $0.20\pm 0.02$ | $\mathbf{16.40}$ |  |

Testing.
To verify the generalization ability of ML-Agent across ML tasks, we evaluate on 10 held-out tasks from MLE-Bench that are never seen during training. Due to the high computational cost of running full autonomous ML pipelines, we do not exhaustively evaluate on the entire benchmark. Instead, we select a representative subset of unused tasks that cover diverse ML pipelines and evaluation metrics, including settings that are generally more challenging than the training tasks.
These held-out tasks span heterogeneous input modalities and task formulations, such as image, text, tabular, time series, audio, and multi-modal settings, as well as different evaluation criteria (e.g., RMSE, log-loss, MAP@K, and quadratic weighted kappa). Details of these tasks are provided in Appendix [B.2](#A2.SS2). During testing, the MLAgentBench environment settings remain consistent with those used in training. To comprehensively assess the LLM agent’s ability in autonomous ML, we propose Performance gain $\Delta_{r}$, the relative improvement over the initial script, defined as $\Delta_{r}=\beta\,\frac{m_{\textit{avg@8}}-m_{\text{init}}}{m_{\text{init}}}$
where $m_{\textit{avg@8}}$ is the mean score over 8 trajectories, $m_{\text{init}}$ is the initial script’s score, and $\beta\in\{-1,1\}$ adjusts for metrics (e.g. MAE, RMSE) to ensure positive $\Delta_{r}$ indicates improvement.

Baselines.
To provide a comprehensive comparison, we evaluate ML-Agent against 3 prompted-based agentic ML methods: MLAB (Huang et al., 2023), AIDE (Jiang et al., 2025), and ML-Master (Liu et al., 2025a). All agents are tested using a diverse set of backbone LLMs, spanning small-scale open-source models (e.g., Qwen2.5-7B-Instruct (Yang et al., 2024)), medium-scale models (e.g., Qwen3-235B (Yang et al., 2025)), large-scale open-source models (e.g., DeepSeek-R1 (Guo et al., 2025b)), and state-of-the-art proprietary LLMs (Gemini-2.5-Pro (Comanici et al., 2025) and GPT-5 (OpenAI, 2025)). We keep the same time limit and number of ML code modifications for a fair comparison between agents with different scaffolds.

### 5.2 Main results

We conduct extensive experiments to evaluate the performance of ML-Agent, a learning-based LLM agent trained through our proposed framework for autonomous ML. Our results demonstrate that ML-Agent achieves strong and consistent performance across both held-in and held-out tasks, and exhibits continuous performance improvements during RL training.

ML-Agent achieves superior performance across both held-in and held-out tasks. We compare ML-Agent with 5 powerful LLM-based agents in 3 scaffolds across 3 held-in and 10 held-out tasks. As shown in Table [1](#S5.T1), ML-Agent significantly outperforms other large open-source models, such as the 671B DeepSeek-R1. For closed-source GPT-5, our agent remains remarkably competitive.
Notably, despite being trained on only 9 tasks, ML-Agent delivers top-tier results across all 10 held-out tasks, demonstrating strong generalization and effective learning from limited experience.

ML-Agent efficiently achieves good performance with much lower cost. As illustrated in Figure [2](#S5.F2), we plot the average performance gain against the average cost per trajectory for various agents. Our proposed ML-Agent (the star) is a clear outlier, positioned in the optimal top-left corner. It achieves highly competitive performance gain of over 15% while maintaining an exceptionally low cost of less than 0.01$\mathdollar$ per trajectory. In contrast, baseline agents like MLAB using powerful models such as GPT-5 incur costs that are more than 20 times higher for similar or even lower performance. This result highlights the significant efficiency of learning-based paradigm, proving it can produce a state-of-the-art agent without relying on expensive, large-scale models. See detailed cost in Appendix [C.4](#A3.SS4).

Figure: Figure 2: Comparison of average performance gain (%) vs. cost ($\mathdollar$) across different models and scaffolds on 10 held-out tasks. Our ML-Agent significantly outperforms other baselines with a competitive gain at a lower cost.
Refer to caption: 2505.23723v2/x2.png

ML-Agent achieves continuous performance improvements. Figure [3](#S5.F3) shows that ML-Agent demonstrates consistent performance improvement across both held-in and held-out tasks as training progresses. This highlights the effectiveness of our step-wise RL paradigm and exploration-enriched fine-tuning in enabling continuous learning from ML environmental feedback, ultimately allowing ML-Agent to outperform all baseline methods.

### 5.3 Ablation Study and Analysis

Ablations on our proposed training framework.
We perform a comprehensive ablation study to assess the contribution of each module in our proposed framework, as summarized in Table [2](#S5.T2). The results indicate that all three module designs are indispensable for optimal performance. For example, removing exploration-enriched fine-tuning leads to a collapse in performance, with the average gain dropping to -0.66% in held-in tasks and -6.20% in held-out tasks, highlighting its critical role in providing a diverse action initialization. Furthermore, agents trained without step-wise RL (episode-wise RL) exhibit poor generalization, achieving only 2.86% average gains in held-out tasks. In contrast, our full framework achieves the highest performance gain at both held-in and held-out tasks, confirming that these components are essential for successful autonomous ML engineering.

Figure: Figure 3: ML-Agent achieves continuous performance improvements; Step-wise RL (evaluated every 5 steps) is more efficient than episode-wise RL (standard PPO, evaluated every 1 step) on both held-in and held-out tasks.
Refer to caption: 2505.23723v2/x3.png

**Table 2: Ablation study of ML-Agent on held-in and held-out ML tasks. We report the average performance gain (%) for each task.**
| Method | Held-in | Held-out |
| --- | --- | --- |
| Tasks | Tasks |  |
| ML-Agent | 18.01 | 15.91 |
| w/o Exploration-enriched fine-tuning | -0.66 | -6.20 |
| w/o Step-wise RL | 10.71 | 2.86 |
| w/o Agentic ML-specific reward | 9.77 | -1.65 |

Effectiveness of step-wise RL training.
To improve training efficiency and scalability, we propose a step-wise RL approach that samples single states from expert trajectories and evaluates atomic actions. To validate this, we implement an alternative episode-wise RL approach using standard PPO, where the policy rolls out the entire trajectory from the task description during data collecting phase in RL. Both methods are initialized from the same ML-Agent-SFT model and trained for 39 steps. We measure GPU time every 5 steps for step-wise RL and 1 step for episode-wise RL. As shown in Figure [3](#S5.F3), step-wise RL adapts more quickly and achieves faster performance gains on both held-in and held-out tasks, while the performance of episode-wise RL improves slowly and incurs much higher time cost. These results demonstrate that step-wise RL not only improves training efficiency by avoiding expensive online rollouts, but also leads to improved performance through targeted single-step updates.

Effectiveness of exploration-enriched fine-tuning.
To validate the efficacy of exploration-enriched fine-tuning in enhancing subsequent RL training, we replace our exploration-enriched fine-tuned model (ML-Agent-SFT) with Qwen2.5-7B (Qwen-7B-Base), Qwen2.5-7B-Instruct (Qwen-7B-Instruct), and DeepSeek-R1-Distill-Qwen-7B (Guo et al., 2025b)(Qwen-7B-Distill) as base models for the RL training. We evaluate the average performance gain of the resulting RL-trained agents on held-in and held-out tasks (Figure [4](#S5.F4)). The agent trained from Qwen-7B-Distill fails to generate valid actions due to distillation-induced format issues, resulting in ineffective learning. The agent trained from Qwen-7B-Base shows overall performance degradation from limited instruction-following capabilities. The agent trained from Qwen-7B-Instruct achieves +13% gains on held-in tasks but -12% on held-out tasks, indicating poor generalization. In contrast, the agent trained from our ML-Agent-SFT achieves +18% and +16% improvement on held-in and held-out tasks, respectively, with greater action diversity during autonomous ML experimentation (Figure [7](#A3.F7)).
These results confirm that exploration-enriched fine-tuning promotes format-compliant, diverse actions, enhancing exploration and generalization in step-wise RL.

Figure: (a) Held-in tasks
Refer to caption: 2505.23723v2/x4.png

**Table 3: Ablation study on the designs of ML-specific reward module: (1) normalized performance reward ($R_{\text{perf.}}$), (2) format reward ($R_{\text{format}}$), and (3) corner cases reward ($R_{\text{corner}}$). Results indicates the necessity of each reward design. We report the average performance gain (%) for each task.**
| Task | cifar10 | house | feedback | leaf | detecting | us | tabular | whale |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $R_{\text{perf.}}$ | $R_{\text{format}}$ | $R_{\text{corner}}$ |  |  |  |  |  |  |  |  |
| ✗ | ✓ | ✓ | 17.58 | 3.94 | 7.79 | 4.75 | 0.26 | 6.40 | -24.96 | 23.24 |
| ✓ | ✗ | ✓ | 10.98 | 6.17 | 8.34 | -30.25 | 0.03 | 6.27 | -12.54 | 2.84 |
| ✓ | ✓ | ✗ | 13.56 | 6.64 | 7.67 | 8.50 | 0.58 | 8.67 | -0.48 | 28.06 |
| ✓ | ✓ | ✓ | 33.80 | 6.77 | 13.47 | 13.87 | 1.74 | 12.96 | 0.20 | 72.89 |

Effectiveness of agentic-ML specific reward module.
Ablation studies in Table [3](#S5.T3) show each reward component is essential:
(1) Performance($R_{\text{perf.}}$): Replacing the scaled performance difference with binary reward leads to noticeable performance drops. This confirms that fine-grained reward signals are more informative for learning meaningful improvements.
(2) Format($R_{\text{format}}$): Removing format constraints causes the largest degradation (e.g., -11.75% on cifar-10), emphasizing the necessity of syntactic and semantic correctness of agent’s output format.
(3) Corner cases($R_{\text{corner}}$): Disabling the neutral reward for corner cases has minimal impact due to their rarity, but improves training stability by preventing over-penalization of non-fatal issues.
In summary, each component of the reward module plays a distinct and complementary role: $R_{\text{perf.}}$ drives performance improvement, $R_{\text{format.}}$ ensures actions validity, and $R_{\text{corner}}$ maintains robustness under real-world limitations. Together, they form a coherent and comprehensive reward structure during RL training for agentic ML.

Figure: Figure 5: Effects of training task number on RL performance. Expanding the number of ML tasks during RL not only refines the agent’s ability on familar tasks but also significantly improves the agent’s ability to generalize across unseen tasks.
Refer to caption: 2505.23723v2/x6.png

Effects of task numbers in RL.
We investigate the impact of using different numbers of ML tasks (0, 3, 6, 9) during step-wise RL training, where the "0 task" condition corresponds to ML-Agent-SFT. We evaluate performance in terms of average performance gain on held-in and held-out tasks. As shown in Figure [5](#S5.F5), performance on both task types improves monotonically as the number of ML tasks increases during RL training. Specifically, training with 3, 6, and 9 ML tasks using step-wise RL lifts the average performance gain on held-out tasks from nearly 0% to approximately 3%, 6%, and 16%, respectively. These results indicate that expanding the diversity of ML tasks during RL not only refines the agent’s ability on familar tasks but also significantly improves the agent’s ability to generalize across unseen tasks.

Case study.
To provide an intuitive understanding, we present several examples in the Appendix [C.5](#A3.SS5), demonstrating task specifications, initial code implementations, baseline and our model’s execution trajectories. These demonstrate that our methodology: (1) generates diverse action sequences through comprehensive reasoning, (2) automatically initiates backtracking mechanisms when performance metrics remain unimproved by proposed modifications, and (3) maintains operational effectiveness on novel tasks through generalized reasoning capabilities.

## 6 Conclusion

In this paper, we propose a training framework for an LLM-based agent on autonomous machine learning tasks. Unlike heuristic prompt-based methods, our method enables agents to learn from task-solving experiences, iteratively refine strategies, and generalize across tasks. The framework involved exploration-enriched fine-tuning, efficient step-wise RL training, and agentic ML-specific reward module. Extensive experiments demonstrate that ML-Agent, powered by a 7B-parameter LLM, surpasses agents using 671B models and delivers performance comparable to proprietary models like GPT-5 at much lower computational cost. This suggests that equipping smaller, open-weight LLM-driven agents with the capacity to learn from experience offers a scalable and effective alternative to proprietary models for autonomous Machine Learning.

## Impact Statement

This paper presents work whose goal is to advance the field of machine
learning. There are many potential societal consequences of our work, none
which we feel must be specifically highlighted here.

## References

- P. Belcak, G. Heinrich, S. Diao, Y. Fu, X. Dong, S. Muralidharan, Y. C. Lin, and P. Molchanov (2025)
Small language models are the future of agentic ai.
External Links: 2506.02153,
[Link](https://arxiv.org/abs/2506.02153)
Cited by: [§1](#S1.p3.1).
- J. Bergstra and Y. Bengio (2012)
Random search for hyper-parameter optimization.
The journal of machine learning research 13 (1), pp. 281–305.
Cited by: [§1](#S1.p1.1).
- J. S. Chan, N. Chowdhury, O. Jaffe, J. Aung, D. Sherburn, E. Mays, G. Starace, K. Liu, L. Maksin, T. Patwardhan, et al. (2024)
Mle-bench: evaluating machine learning agents on machine learning engineering.
arXiv preprint arXiv:2410.07095.
Cited by: [§B.2](#A2.SS2.p1.1),
[Table 4](#A2.T4),
[Table 4](#A2.T4.3.2),
[§1](#S1.p2.1),
[§1](#S1.p5.1),
[§5.1](#S5.SS1.p1.3).
- G. Comanici, E. Bieber, M. Schaekermann, I. Pasupat, N. Sachdeva, I. Dhillon, M. Blistein, O. Ram, D. Zhang, E. Rosen, et al. (2025)
Gemini 2.5: pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities.
arXiv preprint arXiv:2507.06261.
Cited by: [§5.1](#S5.SS1.p3.1).
- Y. Du, S. Li, A. Torralba, J. B. Tenenbaum, and I. Mordatch (2023)
Improving factuality and reasoning in language models through multiagent debate.
In Forty-first International Conference on Machine Learning,
Cited by: [§1](#S1.p2.1).
- N. Erickson, J. Mueller, A. Shirkov, H. Zhang, P. Larroy, M. Li, and A. Smola (2020)
Autogluon-tabular: robust and accurate automl for structured data.
arXiv preprint arXiv:2003.06505.
Cited by: [§2](#S2.p1.1).
- J. Eschmann (2021)
Reward function design in reinforcement learning.
Reinforcement learning algorithms: Analysis and Applications, pp. 25–33.
Cited by: [§1](#S1.p5.1).
- M. Feurer, K. Eggensperger, S. Falkner, M. Lindauer, and F. Hutter (2022)
Auto-sklearn 2.0: hands-free automl via meta-learning. arxiv 2020.
arXiv preprint arXiv:2007.04074.
Cited by: [§2](#S2.p1.1).
- D. Guo, D. Yang, H. Zhang, J. Song, R. Zhang, R. Xu, Q. Zhu, S. Ma, P. Wang, X. Bi, et al. (2025a)
Deepseek-r1: incentivizing reasoning capability in llms via reinforcement learning.
arXiv preprint arXiv:2501.12948.
Cited by: [§2](#S2.p3.1).
- D. Guo, D. Yang, H. Zhang, J. Song, R. Zhang, R. Xu, Q. Zhu, S. Ma, P. Wang, X. Bi, et al. (2025b)
Deepseek-r1: incentivizing reasoning capability in llms via reinforcement learning.
arXiv preprint arXiv:2501.12948.
Cited by: [§4.2](#S4.SS2.p6.7),
[§5.1](#S5.SS1.p3.1),
[§5.3](#S5.SS3.p3.1).
- S. Hong, X. Zheng, J. Chen, Y. Cheng, J. Wang, C. Zhang, Z. Wang, S. K. S. Yau, Z. Lin, L. Zhou, et al. (2023)
Metagpt: meta programming for multi-agent collaborative framework.
arXiv preprint arXiv:2308.00352 3 (4), pp. 6.
Cited by: [§1](#S1.p2.1).
- Y. Hu, Y. Cai, Y. Du, X. Zhu, X. Liu, Z. Yu, Y. Hou, S. Tang, and S. Chen (2024)
Self-evolving multi-agent collaboration networks for software development.
arXiv preprint arXiv:2410.16946.
Cited by: [§1](#S1.p2.1).
- Q. Huang, J. Vora, P. Liang, and J. Leskovec (2023)
Mlagentbench: evaluating language agents on machine learning experimentation.
arXiv preprint arXiv:2310.03302.
Cited by: [§B.3](#A2.SS3.p1.1),
[Table 4](#A2.T4),
[Table 4](#A2.T4.3.2),
[Table 5](#A2.T5),
[Table 5](#A2.T5.3.2),
[§1](#S1.p2.1),
[§3](#S3.p1.1),
[§5.1](#S5.SS1.p1.3),
[§5.1](#S5.SS1.p3.1).
- Z. Jiang, D. Schmidt, D. Srikanth, D. Xu, I. Kaplan, D. Jacenko, and Y. Wu (2025)
AIDE: ai-driven exploration in the space of code.
External Links: 2502.13138,
[Link](https://arxiv.org/abs/2502.13138)
Cited by: [§1](#S1.p2.1),
[§2](#S2.p2.1),
[§5.1](#S5.SS1.p3.1).
- H. Jin, F. Chollet, Q. Song, and X. Hu (2023)
Autokeras: an automl library for deep learning.
Journal of machine Learning research 24 (6), pp. 1–6.
Cited by: [§1](#S1.p2.1).
- T. Kaufmann, P. Weng, V. Bengs, and E. Hüllermeier (2023)
A survey of reinforcement learning from human feedback.
arXiv preprint arXiv:2312.14925 10.
Cited by: [§2](#S2.p3.1).
- E. LeDell and S. Poirier (2020)
H2o automl: scalable automatic machine learning.
In Proceedings of the AutoML Workshop at ICML,
Vol. 2020, pp. 24.
Cited by: [§1](#S1.p2.1).
- Z. Li, Q. Zang, D. Ma, J. Guo, T. Zheng, M. Liu, X. Niu, Y. Wang, J. Yang, J. Liu, et al. (2024)
Autokaggle: a multi-agent framework for autonomous data science competitions.
arXiv preprint arXiv:2410.20424.
Cited by: [§2](#S2.p2.1).
- S. Liu, P. Ram, D. Vijaykeerthy, D. Bouneffouf, G. Bramble, H. Samulowitz, D. Wang, A. Conn, and A. Gray (2020)
An admm based framework for automl pipeline configuration.
In Proceedings of the AAAI Conference on Artificial Intelligence,
Vol. 34, pp. 4892–4899.
Cited by: [§2](#S2.p1.1).
- Z. Liu, Y. Cai, X. Zhu, Y. Zheng, R. Chen, Y. Wen, Y. Wang, W. E, and S. Chen (2025a)
ML-master: towards ai-for-ai via integration of exploration and reasoning.
arXiv preprint arXiv:2506.16499.
Cited by: [§1](#S1.p2.1),
[§2](#S2.p2.1),
[§5.1](#S5.SS1.p3.1).
- Z. Liu, C. Chen, W. Li, P. Qi, T. Pang, C. Du, W. S. Lee, and M. Lin (2025b)
Understanding r1-zero-like training: a critical perspective.
arXiv preprint arXiv:2503.20783.
Cited by: [§2](#S2.p3.1).
- A. L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts (2011)
Learning word vectors for sentiment analysis.
In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, D. Lin, Y. Matsumoto, and R. Mihalcea (Eds.),
Portland, Oregon, USA, pp. 142–150.
External Links: [Link](https://aclanthology.org/P11-1015/)
Cited by: [Table 4](#A2.T4.4.11.1).
- T. Masterman, S. Besen, M. Sawtell, and A. Chao (2024)
The landscape of emerging ai agent architectures for reasoning, planning, and tool calling: a survey.
arXiv preprint arXiv:2404.11584.
Cited by: [§1](#S1.p2.1).
- F. Mohr, M. Wever, and E. Hüllermeier (2018)
ML-plan: automated machine learning via hierarchical planning.
Machine Learning 107, pp. 1495–1515.
Cited by: [§2](#S2.p1.1).
- R. S. Olson and J. H. Moore (2016)
TPOT: a tree-based pipeline optimization tool for automating machine learning.
In Workshop on automatic machine learning,
pp. 66–74.
Cited by: [§2](#S2.p1.1).
- OpenAI (2024)
GPT-4o mini: advancing cost-efficient intelligence.
Note: [https://openai.com/index/gpt-4o-mini-advancing-cost-efficient-intelligence/](https://openai.com/index/gpt-4o-mini-advancing-cost-efficient-intelligence/)Accessed: 2025-01-23
Cited by: [§B.3](#A2.SS3.p1.1),
[§5.1](#S5.SS1.p1.3).
- OpenAI (2025)
Introducing gpt 5.
Note: [https://openai.com/index/introducing-gpt-5/](https://openai.com/index/introducing-gpt-5/)Accessed: 2025-09-25
Cited by: [§5.1](#S5.SS1.p3.1).
- L. Ouyang, J. Wu, X. Jiang, D. Almeida, C. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, et al. (2022)
Training language models to follow instructions with human feedback.
Advances in neural information processing systems 35, pp. 27730–27744.
Cited by: [§2](#S2.p3.1).
- X. Pang, S. Tang, R. Ye, Y. Xiong, B. Zhang, Y. Wang, and S. Chen (2024)
Self-alignment of large language models via monopolylogue-based social scene simulation.
arXiv preprint arXiv:2402.05699.
Cited by: [§1](#S1.p2.1).
- P. S. Park, P. Schoenegger, and C. Zhu (2024)
Diminished diversity-of-thought in a standard large language model.
Behavior Research Methods 56 (6), pp. 5754–5770.
Cited by: [§1](#S1.p5.1).
- P. Putta, E. Mills, N. Garg, S. Motwani, C. Finn, D. Garg, and R. Rafailov (2024)
Agent q: advanced reasoning and learning for autonomous ai agents.
arXiv preprint arXiv:2408.07199.
Cited by: [§2](#S2.p3.1).
- C. Qian, W. Liu, H. Liu, N. Chen, Y. Dang, J. Li, C. Yang, W. Chen, Y. Su, X. Cong, et al. (2023)
Chatdev: communicative agents for software development.
arXiv preprint arXiv:2307.07924.
Cited by: [§1](#S1.p2.1).
- R. Rafailov, A. Sharma, E. Mitchell, C. D. Manning, S. Ermon, and C. Finn (2023)
Direct preference optimization: your language model is secretly a reward model.
Advances in Neural Information Processing Systems 36, pp. 53728–53741.
Cited by: [§2](#S2.p3.1).
- J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017)
Proximal policy optimization algorithms.
arXiv preprint arXiv:1707.06347.
Cited by: [§4.2](#S4.SS2.p6.7).
- G. Sheng, C. Zhang, Z. Ye, X. Wu, W. Zhang, R. Zhang, Y. Peng, H. Lin, and C. Wu (2024)
HybridFlow: a flexible and efficient rlhf framework.
arXiv preprint arXiv: 2409.19256.
Cited by: [§C.1](#A3.SS1.p1.4).
- Z. Tang, H. Fang, S. Zhou, T. Yang, Z. Zhong, T. Hu, K. Kirchhoff, and G. Karypis (2024)
Autogluon-multimodal (automm): supercharging multimodal automl with foundation models.
arXiv preprint arXiv:2404.16233.
Cited by: [§1](#S1.p2.1),
[§2](#S2.p1.1).
- P. Trirat, W. Jeong, and S. J. Hwang (2024)
Automl-agent: a multi-agent llm framework for full-pipeline automl.
arXiv preprint arXiv:2410.02958.
Cited by: [§2](#S2.p2.1).
- S. Wang, S. Zhang, J. Zhang, R. Hu, X. Li, T. Zhang, J. Li, F. Wu, G. Wang, and E. Hovy (2024)
Reinforcement learning enhanced llms: a survey.
arXiv preprint arXiv:2412.10400.
Cited by: [§2](#S2.p3.1).
- Z. Wang, K. Wang, Q. Wang, P. Zhang, L. Li, Z. Yang, K. Yu, M. N. Nguyen, L. Liu, E. Gottlieb, et al. (2025)
Ragen: understanding self-evolution in llm agents via multi-turn reinforcement learning.
arXiv preprint arXiv:2504.20073.
Cited by: [§2](#S2.p3.1).
- W. Xiong, Y. Song, X. Zhao, W. Wu, X. Wang, K. Wang, C. Li, W. Peng, and S. Li (2024)
Watch every step! llm agent learning via iterative step-level process refinement.
In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing,
pp. 1556–1572.
Cited by: [§1](#S1.p4.1),
[§2](#S2.p3.1).
- F. Xu, Q. Hao, Z. Zong, J. Wang, Y. Zhang, J. Wang, X. Lan, J. Gong, T. Ouyang, F. Meng, et al. (2025)
Towards large reasoning models: a survey of reinforced reasoning with large language models.
arXiv preprint arXiv:2501.09686.
Cited by: [§2](#S2.p3.1).
- A. Yang, A. Li, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Gao, C. Huang, C. Lv, C. Zheng, D. Liu, F. Zhou, F. Huang, F. Hu, H. Ge, H. Wei, H. Lin, J. Tang, J. Yang, J. Tu, J. Zhang, J. Yang, J. Yang, J. Zhou, J. Zhou, J. Lin, K. Dang, K. Bao, K. Yang, L. Yu, L. Deng, M. Li, M. Xue, M. Li, P. Zhang, P. Wang, Q. Zhu, R. Men, R. Gao, S. Liu, S. Luo, T. Li, T. Tang, W. Yin, X. Ren, X. Wang, X. Zhang, X. Ren, Y. Fan, Y. Su, Y. Zhang, Y. Zhang, Y. Wan, Y. Liu, Z. Wang, Z. Cui, Z. Zhang, Z. Zhou, and Z. Qiu (2025)
Qwen3 technical report.
External Links: 2505.09388,
[Link](https://arxiv.org/abs/2505.09388)
Cited by: [§5.1](#S5.SS1.p3.1).
- A. Yang, B. Yang, B. Zhang, B. Hui, B. Zheng, B. Yu, C. Li, D. Liu, F. Huang, H. Wei, et al. (2024)
Qwen2. 5 technical report.
arXiv preprint arXiv:2412.15115.
Cited by: [§B.3](#A2.SS3.p1.1),
[§5.1](#S5.SS1.p1.3),
[§5.1](#S5.SS1.p3.1).
- G. Zhang, H. Geng, X. Yu, Z. Yin, Z. Zhang, Z. Tan, H. Zhou, Z. Li, X. Xue, Y. Li, et al. (2025)
The landscape of agentic reinforcement learning for llms: a survey.
arXiv preprint arXiv:2509.02547.
Cited by: [§2](#S2.p3.1).
- L. Zhang, Y. Zhang, K. Ren, D. Li, and Y. Yang (2023a)
Mlcopilot: unleashing the power of large language models in solving machine learning tasks.
arXiv preprint arXiv:2304.14979.
Cited by: [§2](#S2.p2.1).
- S. Zhang, C. Gong, L. Wu, X. Liu, and M. Zhou (2023b)
Automl-gpt: automatic machine learning with gpt.
arXiv preprint arXiv:2305.02499.
Cited by: [§2](#S2.p2.1).
- R. Zheng, S. Dou, S. Gao, Y. Hua, W. Shen, B. Wang, Y. Liu, S. Jin, Q. Liu, Y. Zhou, et al. (2023)
Secrets of rlhf in large language models part i: ppo.
arXiv preprint arXiv:2307.04964.
Cited by: [§2](#S2.p3.1).
- Y. Zheng, R. Zhang, J. Zhang, Y. Ye, Z. Luo, Z. Feng, and Y. Ma (2024)
LlamaFactory: unified efficient fine-tuning of 100+ language models.
In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 3: System Demonstrations),
Bangkok, Thailand.
External Links: [Link](http://arxiv.org/abs/2403.13372)
Cited by: [§C.1](#A3.SS1.p1.4).

## Appendix A Problem formulation

Reformulation from equation [1](#S3.E1) to equation [3](#S4.E3). Starting from equation [1](#S3.E1), suppose the initial state distribution is $d^{\pi_{\theta}}(s_{0})$, the state transition probability is $p_{\pi}(s_{t+1}|s_{t},a_{t})$, then we have

$$ $\mathcal{P}_{\pi_{\theta}}(\tau)=d^{\pi_{\theta}}(s_{0})\prod_{t=0}^{n-1}p(s_{t+1}|s_{t},a_{t})\pi_{\theta}(a_{t}|s_{t}).$ (7) $$

Hence the reformulation is:

$$ $\displaystyle\mathcal{J}(\theta)$ $\displaystyle=\mathbb{E}_{\pi_{\theta}}[R(\tau)]$ (8) $\displaystyle=\sum_{\tau}\mathcal{P}_{\pi_{\theta}}(\tau)R(\tau)$ $\displaystyle=\sum_{s_{0},a_{0},\ldots,s_{n}}\left(d^{\pi_{\theta}}(s_{0})\prod_{t=0}^{n-1}p(s_{t+1}|s_{t},a_{t})\pi_{\theta}(a_{t}|s_{t})\right)\left(\sum_{t=0}^{n}R(s_{t},a_{t})\right)$ $\displaystyle=\sum_{t=0}^{n-1}\sum_{s_{0},a_{0},\ldots,s_{n}}\left(d^{\pi_{\theta}}(s_{0})\prod_{k=0}^{n-1}p(s_{k+1}|s_{k},a_{k})\pi_{\theta}(a_{k}|s_{k})\right)R(s_{t},a_{t})$ $\displaystyle=\sum_{t=0}^{n-1}\sum_{s_{t},a_{t}}\left(\sum_{s_{0},a_{0},\ldots s_{t-1},a_{t-1}}d^{\pi_{\theta}}(s_{0})\prod_{k=0}^{t-1}p(s_{k+1}|s_{k},a_{k})\pi_{\theta}(a_{k}|s_{k})\right)\pi_{\theta}(a_{t}|s_{t})R(s_{t},a_{t})$ $$

However, we can define the state distribution $d^{\pi_{\theta}}(s_{t})$ as the probability agent visits state $s_{t}$ at time $t$. Then according to this definition, this probability can be written as:

$$ $d^{\pi_{\theta}}(s_{t})=\sum_{s_{0},a_{0},\ldots s_{t-1},a_{t-1}}d^{\pi_{\theta}}(s_{0})\prod_{k=0}^{t-1}p(s_{k+1}|s_{k},a_{k})\pi_{\theta}(a_{k}|s_{k}).$ (9) $$

Then we have

$$ $\displaystyle\mathcal{J}(\theta)$ $\displaystyle=\sum_{t=0}^{n-1}\sum_{s_{t},a_{t}}d^{\pi_{\theta}}(s_{t})\pi_{\theta}(a_{t}|s_{t})R(a_{t},s_{t})$ (10) $\displaystyle=\sum_{s_{t}\in\mathcal{S}}d^{\pi_{\theta}}(s_{t})\sum_{a_{t}\in\mathcal{A}}\pi_{\theta}(a_{t}|s_{t})R(a_{t},s_{t})$ $\displaystyle=\mathcal{J}_{\text{step}}(\theta)$ $$

## Appendix B Machine Learning Tasks and Data Collection Process

### B.1 Data Collecting Pipeline for Exploration-Enriched Fine-tuning

We construct diverse action pools along three semantic axes—Data, Model, and Learning—to support structured exploration. For each axis, we prompt a frozen LLM (GPT-4o-mini) to generate a large set of candidate actions (e.g., “Add MixUp augmentation”, “Switch to AdamW optimizer”). To promote diversity, we embed all candidates using a sentence transformer and apply farthest-point sampling (FPS) to select a compact, representative subset. The resulting pools $\mathcal{P}_{\text{Data}}$, $\mathcal{P}_{\text{Model}}$, and $\mathcal{P}_{\text{Learning}}$ are fixed during training.

During data collection, we form exploration-enriched prompts by randomly selecting 1–3 axes, shuffling their order, and drawing one action from each corresponding pool. These actions are concatenated into an initial instruction for the expert agent, which then interacts with a fast-executable ML environment (e.g., small-scale tabular or vision tasks) to produce a full trajectory.
The complete pipeline is summarized in Algorithm [1](#alg1).

Figure: Algorithm 1 Exploration-Enriched Trajectory Generation

### B.2 Details of Machine Learning Tasks

The machine learning tasks utilized in our paper are all from MLAgentBench or MLE-bench.
Table [4](#A2.T4) shows all 9 training tasks and 10 testing tasks.
The 9 training tasks contain 4 tasks from MLAgentBench and 5 from MLE-bench (Chan et al., 2024); while the 10 testing tasks are all from MLE-bench.

The selection strategy of training tasks aims to enhance data collection efficiency. Specifically, we select relatively simpler machine learning tasks (e.g. tasks labeled with low complexity in MLE-bench) for training. These training tasks typically involve smaller datasets, which enable faster iterations. For testing, we select relatively more complex tasks to evaluate the generalization capability.
In addition, the training tasks and test tasks span three machine learning data types (image, text and tabular) and two general task categories (regression and classification).

**Table 4: All training and testing tasks used in our experiments. MLA and MLE stand for MLAgentBench (Huang et al., 2023) and MLE-bench (Chan et al., 2024) respectively.**
| Task Name | Data Type | Task Type | Metric | Source |
| --- | --- | --- | --- | --- |
| Training |  |  |  |  |
| [cifar-10](https://www.kaggle.com/competitions/cifar-10) | Image | Classification | Acc. (%) ↑ | MLA |
| [aerial-cactus-identification](https://www.kaggle.com/competitions/aerial-cactus-identification) | Image | Classification | AUC ↑ | MLE |
| [dogs-vs-cats-redux-kernels-edition](https://www.kaggle.com/competitions/dogs-vs-cats-redux-kernels-edition) | Image | Classification | Logloss ↓ | MLE |
| [plant-pathology-2020-fgvc7](https://www.kaggle.com/competitions/plant-pathology-2020-fgvc7) | Image | Classification | AUC ↑ | MLE |
| [home-data-for-ml-course](https://www.kaggle.com/competitions/home-data-for-ml-course) | Tabular | Regression | MAE ↓ | MLA |
| [spaceship-titanic](https://www.kaggle.com/competitions/spaceship-titanic) | Tabular | Regression | Acc. (%) ↑ | MLA |
| [nomad2018-predict-transparent-conductors](https://www.kaggle.com/competitions/nomad2018-predict-transparent-conductors) | Tabular | Regression | RMSLE ↓ | MLE |
| [feedback-prize-english-language-learning](https://www.kaggle.com/competitions/feedback-prize-english-language-learning) | Text | Classification | MCRMSE ↓ | MLA |
| ogbn-arxiv (Maas et al., 2011) | Graph | Classification | Acc. (%) ↑ | MLA |
| Testing |  |  |  |  |
| [denoising-dirty-documents](https://www.kaggle.com/competitions/denoising-dirty-documents) | Image | Generation | RMSE↓ | MLE |
| [leaf-classification](https://www.kaggle.com/competitions/leaf-classification) | Image | Classification | Logloss ↓ | MLE |
| [statoil-iceberg-classifier-challenge](https://www.kaggle.com/competitions/statoil-iceberg-classifier-challenge) | Image | Classification | Logloss ↓ | MLE |
| [whale-categorization-playground](https://www.kaggle.com/competitions/whale-categorization-playground) | Image | Classification | MAP@5 ↑ | MLE |
| [learning-agency-lab-automated-essay-scoring-2](https://www.kaggle.com/competitions/learning-agency-lab-automated-essay-scoring-2) | Text | Regression | QWK ↑ | MLE |
| [detecting-insults-in-social-commentary](https://www.kaggle.com/competitions/detecting-insults-in-social-commentary) | Text | Classification | Acc. (%) ↑ | MLE |
| [spooky-author-identification](https://www.kaggle.com/competitions/spooky-author-identification) | Text | Classification | Logloss ↓ | MLE |
| [jigsaw-toxic-comment-classification-challenge](https://www.kaggle.com/competitions/jigsaw-toxic-comment-classification-challenge) | Text | Classification | AUC ↑ | MLE |
| [us-patent-phrase-to-phrase-matching](https://www.kaggle.com/competitions/us-patent-phrase-to-phrase-matching) | Tabular | Regression | PCC ↑ | MLE |
| [tabular-playground-series-dec-2021](https://www.kaggle.com/competitions/tabular-playground-series-dec-2021) | Tabular | Regression | Acc. (%) ↑ | MLE |

**Table 5: Actions in MLAgentBench (Huang et al., 2023), where each action has a name, input and output. Most of the actions are primitive actions that include file system operations and python script execution. The last two are compound actions that is composed of multiple primitive actions and LM calls.**
| Action Name | Input | Observation |
| --- | --- | --- |
| List Files | directory (e.g. .) | list of files in the directory |
| Copy File | Source (e.g. train.py), destination (e.g. train_copy.py) | A success or error message |
| Inspect Script Lines | file name, start line number, end line number | the file content between start and end line numbers |
| Execute Script | file name (e.g. train.py) | Any output from the execution |
| Final Answer | None | None |
| Understand File | file name, a query (e.g. the model architecture) | retrieved content from the file relevant to the query |
| Edit Script | file name, edit instruction (e.g. change epoch to 20), save file name | The diff of the edited file based on the instruction |

Specifically, Each task consists of the following components: (1) training, validation, and test data; (2) an initial bug-free script, "train.py", generated by GPT-4o-mini; (3) an evaluation script, "eval.py", which is used to calculate the test score from the submitted results; (4) a problem description file, "research_problem.txt"; and (5) a "prepare.py" script to download the data if necessary.
An example file structure and related problem descriptions are shown in Figure [6](#A2.F6). To ensure clarity regarding the task details and training objectives, we have refined some initial prompts from MLAgentBench by incorporating specific targets, such as "try your best to increase the test accuracy to 99.99%" (see in the right box in Figure [6](#A2.F6)). The format for the initial prompt, including the tool and format prompts, follows actions defined by MLAgentBench (see Table [6](#A2.T6)).

Figure: Figure 6: Task file structure and prompt about the machine learning problem of cifar-10 task, for instance.
Refer to caption: 2505.23723v2/figs/task_files.png

Table: Table 6: Initial prompt template for agents on autonomous machine learning.

Figure: Tools prompt ({tools_prompt}) in initial prompt.

Table: Table 7: Response format requirement ({format_prompt}) in the initial prompt.

### B.3 Details of Data Collection

In this paper, we use the MLAgentBench (Huang et al., 2023) environment to collect training trajectories across 9 machine learning tasks. The environment needs an LLM-based agent to take actions and send feedback to the agent. This will iterate for certain steps.
We employ GPT-4o-mini (OpenAI, 2024) as the LLM-based agent to generate thinking and action following Table [B.2](#LSTx1). This agent interacts with the environment, while Qwen2.5-Coder-32B-Instruct (Yang et al., 2024) powers the coder agent, which is responsible for writing code and understanding files within the environment.

Each trajectory comprises a multi-turn conversation between the agent and the environment. For each trajectory, we set the maximum number of steps as 15 and the time limit as 30 minutes to control the length and duration of interactions.
Finally, we generated 10k trajectories on 9 tasks. These trajectories are utilized both in SFT training and PPO training.

Since each task in the MLAgentBench environment requires an initial script, tasks sourced from MLE-bench do not have a natural initial script. To address this, we generate simple, bug-free initial scripts for those tasks using GPT-4o-mini to meet the environment’s requirements.

To diversify the trajectories we collect for SFT training, we curate an initial idea pool of at least 100 diverse ideas which may potentially improve the performance of our initial script. We calculate the embedding distance of each idea in initial idea pool and filter out the top 10 initial ideas whose average embedding distance is farthest to others. These ideas form a defined idea pool, which guides the first step of each trajectory. For the generation of each trajectory, we randomly select 1 to 3 idea combinations from this idea pool and prioritize their implementation in the initial step by including the relevant instructions in the file research_problem.txt (see Figure [6](#A2.F6)). Table [8](#A2.T8) shows the prompt we use and Table [9](#A2.T9) shows an example of defined idea pool for the first step.

Table: Table 8: The prompt we use to generate the data-preprocessing idea pool.

Table: Table 9: An example of the first step action space(after filtering) when collecting training trajectories.

## Appendix C Experimental Details

### C.1 Details of Experimental Set-up

Training details. We implement our supervised fine-tuning (SFT) and proximal policy optimization (PPO) training using 8 A100s. For the SFT, the code base is LLama-Factory (Zheng et al., 2024), where we fully fine-tune the qwen2.5-7b model for 2 epochs with batch size 64 and learning rate $2e-5$.
For the PPO, the code base is VeRL (Sheng et al., 2024). The PPO training setup involves the following hyperparameters and configurations: the training batch size is set to 256, and the number of epochs is 1. Additionally, the learning rate of actor and critic is set as $1e-6$ and $1e-5$, respectively, and the coefficient of KL is $0.001$.

Baseline details. We show the specific versions of baselines in Table [10](#A3.T10).

**Table 10: Model Version and Identifier Mapping**
| Model Name | Version |
| --- | --- |
| GPT-4o-mini | GPT-4o-mini-2024-07-18 |
| GPT-4o | GPT-4o-2024-08-06 |
| Qwen-7B-Base | Qwen2.5-7B |
| Qwen-7B-Instruct | Qwen2.5-7B-Instruct |
| Qwen-32B-Instruct | Qwen2.5-32B-Instruct |
| GPT-5 | GPT-5-2025-08-07 |

### C.2 Additional Ablation Study and Analysis

Figure: Figure 7: Unique noun counts per verb across 100 randomly sampled edit instructions, comparing the Qwen2.5-7B-Instruct model (blue) with the ML-Agent-SFT model (red).
Refer to caption: 2505.23723v2/x7.png

Diversity.
Figure [7](#A3.F7) compares the number of unique nouns associated with each editing verb in two models: Qwen2.5-7B-Instruct and ours-sft (ML-Agent-SFT). To generate these counts, we randomly sampled 100 edit_instruction sentences from the recorded expert trajectories. Then, we utilize an open-source NLP toolkit SpaCy to obtain the verb and noun for each edit_instruction sentence. Results show that after supervised fine-tuning with expert’s trajectories, the model can output a broader variety of actions, evidenced by the higher counts of unique nouns per verb.

Figure: Figure 8: Different number of training samples in RL, starting from our sft model.
Refer to caption: 2505.23723v2/x8.png

Effects of training data size in RL.

Here, we investigate how varying the number of training data samples (0k, 3k, 10k, 30k) affects the performance in RL. The 0k scenario represents ML-Agent-SFT model. For every model, we run 128 trajectories for each task and evaluate the best@K, where K ranges over [4, 8, 16, 32, 64, 128], as shown in Figure [8](#A3.F8). In particular, for the two held-out tasks (second row), training with ppo (3k, 6k, and 9k) improves performance faster than 0k as the number of sampled trajectories increases.

Is thought helpful?
In Table [11](#A3.T11), we compare the performance of models with and without the requirement of thought before taking action across 13 tasks. The thought process includes several key components such as "Reflection," "Research Plan and Status," "Fact Check," "Thought," "Action," and "Action Input,". In contrast, the model without thought only requires "Action" and "Action Input." Note that the supervised fine-tuning data is also consistent with the key requirement.
The models with thought generally exhibit higher improved performance on both held-in and held-out tasks. For instance, on the held-in cifar-10, the model with thought reaches 33.80% performance gain compared to 13.03% when thought is not required.
This trend continues across the held-out tasks, where the model with thought shows higher accuracy and lower loss, demonstrating the importance of incorporating a thoughtful reflection and planning phase for Autonomous machine learning via RL.

**Table 11: Performance comparison of reinforcement learning models with and without the requirement of thought prior to action. Average performance gains (%) are shown for both held-in and held-out tasks, highlighting improvements in various tasks when thought is incorporated.**
| Thought? | Held-in tasks | Held-out tasks |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| cifar-10 | house. | feedback | denoising. | leaf. | statoil. | learning. | detecting. | spooky. | jigsaw. | us. | tabular. |  |
| ✗ | 13.03 | 5.68 | 9.88 | 28.66 | 2.50 | -0.03 | 1.27 | 0.64 | -3.40 | 0.00 | 7.15 | -0.02 |
| ✓ | 33.80 | 6.77 | 13.47 | 52.38 | 13.87 | 1.41 | 1.91 | 1.74 | 1.76 | 0.01 | 12.96 | 0.20 |

### C.3 Additional ML tasks results

To further demonstrate the generalization of our ML-Agent, we expand our evaluation with additional held-out tasks from MLE-Bench and the Kaggle platform, covering medical image, audio, time series, and multi-modal settings. The results in Table [12](#A3.T12) show that ML-Agent maintains strong cross-task generalization compared with both same-scale open-source baselines (e.g., Qwen2.5-7B-Instruct) and much larger proprietary models.

**Table 12: Cross-task generalization of ML-Agent vs. baselines on additional held-out ML tasks over 8 trajectories.**
| Task Name | Task Type | Metric | Qwen2.5-7B | DeepSeek-R1 | Gemini-2.5-Pro | GPT-5 | ML-Agent |
| --- | --- | --- | --- | --- | --- | --- | --- |
| APTOS 2019 Blindness Detection | medical image classification | QWK | -1.1040 | -0.6148 | -1.5141 | -3.3470 | -0.0455 |
| H&M Personalized Fashion Recommendations | multi-modal regression | MAP@12 | -4.2553 | -1.0638 | 2.1277 | 20.2128 | 8.5106 |
| Optiver – Trading at the Close | time series regression | MAE | 0.0229 | 0.0382 | 0.0573 | 0.1412 | 0.1145 |
| ICML 2013 Whale Challenge: Right Whale Redux | audio classification | AUC | 0.1052 | 0.1804 | 5.8770 | 0.5712 | 0.7215 |
| Text Normalization Challenge: English | text normalization | Acc. | -2.4940 | 0.0839 | 0.1345 | 0.1243 | 0.0839 |

### C.4 Cost analysis with amortization

Here we provide the pre-inference computation cost incurred by (i) expert trajectory collection with GPT-4o-mini, (ii) supervised fine-tuning (SFT), and (iii) step-wise RL training. Table [13](#A3.T13) summarizes the total pre-inference cost. Specifically, for GPT-4o-mini, we utilize offical api price to calculate the total api cost during expert trajectories collection; for supervised fine-tuning, we estimate the computational cost for training a 7B parameters model with 10k training samples; For RL training, we estimate step-wise rollout costs.

Given this one-time cost, we estimate the number of trajectories required to amortize it when compared against strong prompt-based baselines. Specifically, we compute the break-even number of trajectories as
$N_{\text{amort}}=\frac{C_{\text{pre}}}{\Delta c},$
where $C_{\text{pre}}$ is the total pre-inference cost and $\Delta c$ is the per-trajectory inference cost gap between a baseline and ML-Agent. Table [14](#A3.T14) reports the resulting $N_{\text{amort}}$ across baselines and foundation models. Overall, owing to its low per-trajectory inference cost, ML-Agent typically needs only a few thousand trajectories to amortize the pre-inference computation cost.

**Table 13: Computation costs ($) incurred before inference, including expert trajectory collection (GPT-4o-mini), SFT, and RL training.**
|  | Expert trajectory collection | SFT | RL | Total |
| --- | --- | --- | --- | --- |
| Cost ($) | 214.47 | 79.20 | 29.14 | 322.81 |

**Table 14: Number of trajectories required to amortize the pre-inference cost (Table [13](#A3.T13)) compared to strong prompt-based methods under different foundation models.**
| Method | Model | #Trajectories to amortize cost |
| --- | --- | --- |
| MLAB | Deepseek-R1 | 4262 |
| MLAB | Gemini-2.5-Pro | 1518 |
| MLAB | GPT-5 | 1633 |
| AIDE | Deepseek-R1 | 3488 |
| AIDE | Gemini-2.5-Pro | 2106 |
| AIDE | GPT-5 | 1901 |
| ML-Master | Deepseek-R1 | 3284 |
| ML-Master | GPT-5 | 3945 |

### C.5 Case study

In this section, we will present more detailed case study on trajectories generated by ML-Agent on some test tasks such as [denoising-dirty-documents](https://www.kaggle.com/competitions/denoising-dirty-documents). In Appendix [C.5.1](#A3.SS5.SSS1), we show the task description for denoising-dirty-documents. In Appendix [C.5.2](#A3.SS5.SSS2), we show the initial script for denoising-dirty-documents. In Appendix [C.5.3](#A3.SS5.SSS3), we show partial trajectory generated by ML-Agent. We give an analysis in Appendix [C.5.4](#A3.SS5.SSS4).

#### C.5.1 Task Description for denoising-dirty-documents

Figure: Task description for denoising-dirty-documents

#### C.5.2 Initial Script for denoising-dirty-documents

Figure: Initial script for denoising-dirty-documents

#### C.5.3 ML-Agent Trajectory on denoising-dirty-documents

Figure: Partial trajectory of ML-Agent on denoising-dirty-documents

#### C.5.4 Analysis

As shown in  [C.5.3](#A3.SS5.SSS3), we see that ML-Agent automatically initiates backtracking mechanisms when performance metrics remain unimproved by proposed modifications. It tries to use a smaller learning rate in step 7 but finds the RMSE increases after executing the code in step 8. Thus, it tries to revert the learning rate back out its initial value in step 9. We also show some step examples on different tasks in this section, which shows that ML-Agent can generate diverse action sequences through comprehensive reasoning.
As a comparison, other models (such as GPT-4o) sometimes fail to understand the requirement and submit the answer after only one modification(sometimes even no modification).

Figure: Some step examples on different tasks

Figure: Some step examples on different tasks

Figure: Some step examples on different tasks

Figure: Some step examples on different tasks

Figure: gpt-4o trajectory on detecting-insults-in-social-commentary