Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning to Ideate for Machine Learning Engineering Agents

Yunxiang Zhang, Kang Zhou, Zhichao Xu, Kiran Ramnath, Yun Zhou, Sangmin Woo, Haibo Ding, Lin Lee Cheong

TL;DR

MLE-Ideator presents a dual-agent architecture that decouples ideation from implementation to enable iterative optimization in machine learning engineering tasks. The Ideator is queried on-demand via a <seek_help> action, and is further enhanced through reinforcement learning with execution-based rewards using a GRPO-based objective. Empirical results on MLE-Bench show that prompting the Ideator improves over implementation-only baselines, while an RL-trained Ideator (e.g., Qwen3-8B) achieves substantial gains and can surpass strong baselines like Claude Sonnet 3.5. The work demonstrates a scalable path toward training strategic AI systems for scientific discovery, with clear trade-offs in inference cost and heavy RL compute requirements.

Abstract

Existing machine learning engineering (MLE) agents struggle to iteratively optimize their implemented algorithms for effectiveness. To address this, we introduce MLE-Ideator, a dual-agent framework that separates ideation from implementation. In our system, an implementation agent can request strategic help from a dedicated Ideator. We show this approach is effective in two ways. First, in a training-free setup, our framework significantly outperforms implementation-only agent baselines on MLE-Bench. Second, we demonstrate that the Ideator can be trained with reinforcement learning (RL) to generate more effective ideas. With only 1K training samples from 10 MLE tasks, our RL-trained Qwen3-8B Ideator achieves an 11.5% relative improvement compared to its untrained counterpart and surpasses Claude Sonnet 3.5. These results highlights a promising path toward training strategic AI systems for scientific discovery.

Learning to Ideate for Machine Learning Engineering Agents

TL;DR

MLE-Ideator presents a dual-agent architecture that decouples ideation from implementation to enable iterative optimization in machine learning engineering tasks. The Ideator is queried on-demand via a <seek_help> action, and is further enhanced through reinforcement learning with execution-based rewards using a GRPO-based objective. Empirical results on MLE-Bench show that prompting the Ideator improves over implementation-only baselines, while an RL-trained Ideator (e.g., Qwen3-8B) achieves substantial gains and can surpass strong baselines like Claude Sonnet 3.5. The work demonstrates a scalable path toward training strategic AI systems for scientific discovery, with clear trade-offs in inference cost and heavy RL compute requirements.

Abstract

Existing machine learning engineering (MLE) agents struggle to iteratively optimize their implemented algorithms for effectiveness. To address this, we introduce MLE-Ideator, a dual-agent framework that separates ideation from implementation. In our system, an implementation agent can request strategic help from a dedicated Ideator. We show this approach is effective in two ways. First, in a training-free setup, our framework significantly outperforms implementation-only agent baselines on MLE-Bench. Second, we demonstrate that the Ideator can be trained with reinforcement learning (RL) to generate more effective ideas. With only 1K training samples from 10 MLE tasks, our RL-trained Qwen3-8B Ideator achieves an 11.5% relative improvement compared to its untrained counterpart and surpasses Claude Sonnet 3.5. These results highlights a promising path toward training strategic AI systems for scientific discovery.
Paper Structure (35 sections, 4 equations, 4 figures, 3 tables)

This paper contains 35 sections, 4 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Single-Agent Frameworks for MLE.
  4. Multi-Agent Frameworks for MLE.
  5. RL-Trained MLE Agents.
  6. Methods
  7. MLE-Ideator
  8. MLE-Ideator-RL
  9. State and Action Space.
  10. Reward Function.
  11. RL Training Pipeline.
  12. Experiments
  13. Experimental Setup
  14. Dataset and Baselines.
  15. Evaluation Metric.
  16. ...and 20 more sections

Figures (4)

  • Figure 1: Overview of MLE-Ideator. The Implementer uses a <seek_help> action to solicit strategic guidance from a dedicated Ideator. This separation improves performance over single-agent baselines and enables the Ideator to be optimized via reinforcement learning with execution-based rewards.
  • Figure 2: For each Ideator model, we show the agent's best achieved normalized score so far in a trajectory averaged over all tasks w.r.t the number of steps in trajectory. We also plot the frequency of <seek_help> actions at each steps, aggregated over all Ideators.
  • Figure 3: Distribution of idea types, aggregated across Qwen3-8B, Qwen3-8B-RL and Claude Sonnet 3.5 as Ideators, paired with Claude Sonnet 3.5 as the implementer.
  • Figure 4: Proportion of effective ideas in each idea type, aggregated across Qwen3-8B, Qwen3-8B-RL and Claude Sonnet 3.5 as Ideators, paired with Claude Sonnet 3.5 as the Implementer. We define an effective idea as the performance of the refined ML solution according to the idea is better than before.