Reinforcement Learning for Machine Learning Engineering Agents

2025-09-01

AI Summary

Existing ML engineering agents rely solely on large language model (LLM) prompting and lack experience-driven, continual optimization capabilities. This paper introduces the first reinforcement learning (RL) agent framework specifically designed for ML engineering tasks, breaking away from conventional prompting-based paradigms. Our method features: (1) a duration-aware gradient update mechanism to mitigate delayed reward propagation in long-action sequences; (2) static LLM-based environment instrumentation, enabling automatic logging injection to generate fine-grained, partially rewardable execution feedback; and (3) a distributed asynchronous RL training architecture for scalable and efficient policy optimization. Evaluated on 12 Kaggle tasks from MLEBench, our RL-trained Qwen2.5-3B agent achieves an average 22% performance gain over the Claude-3.5-Sonnet prompting baseline—demonstrating that lightweight models can attain superior engineering intelligence through experience-driven evolution.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: Prompt Engineering / Prompting
• Multiagent Systems: Multiagent Learning

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

Abstract

Existing agents for solving tasks such as ML engineering rely on prompting powerful language models. As a result, these agents do not improve with more experience. In this paper, we show that agents backed by weaker models that improve via reinforcement learning (RL) can outperform agents backed by much larger, but static models. We identify two major challenges with RL in this setting. First, actions can take a variable amount of time (e.g., executing code for different solutions), which leads to asynchronous policy gradient updates that favor faster but suboptimal solutions. To tackle variable-duration actions, we propose duration- aware gradient updates in a distributed asynchronous RL framework to amplify high-cost but high-reward actions. Second, using only test split performance as a reward provides limited feedback. A program that is nearly correct is treated the same as one that fails entirely. To address this, we propose environment instrumentation to offer partial credit, distinguishing almost-correct programs from those that fail early (e.g., during data loading). Environment instrumentation uses a separate static language model to insert print statement to an existing program to log the agent's experimental progress, from which partial credit can be extracted as reward signals for learning. Our experimental results on MLEBench suggest that performing gradient updates on a much smaller model (Qwen2.5-3B) trained with RL outperforms prompting a much larger model (Claude-3.5-Sonnet) with agent scaffolds, by an average of 22% across 12 Kaggle tasks.

Problem

• Improving ML engineering agents with reinforcement learning instead of static prompting
• Addressing variable-duration actions in RL through duration-aware gradient updates
• Providing partial credit rewards through environment instrumentation for better feedback

Innovation

• Duration-aware gradient updates for variable actions
• Environment instrumentation for partial credit rewards
• Distributed asynchronous RL framework for smaller models