# Reward-rate Policy Gradient for Efficient Machine Learning Engineering Agents

arXiv: https://arxiv.org/abs/2609.36393 · HTML: https://mlsanity.com/paper/2609.36393/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reinforcement learning for LLM agents, one of 16 directions in LLM agents and tool use.
- Builds on 27 of its 45 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Placed in

- LLM agents and tool use / [Reinforcement learning for LLM agents](https://mlsanity.com/agents/reinforcement-learning-for-llm-agents/)
- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 27 of 27)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Reflexion: Language Agents with Verbal Reinforcement Learning](https://mlsanity.com/paper/2303.11366/) (2023)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Teaching Large Language Models to Self-Debug](https://mlsanity.com/paper/2304.05128/) (2023)
- [MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering](https://mlsanity.com/paper/2410.07095/) (2024)
- [Darwin Godel Machine: Open-Ended Evolution of Self-Improving Agents](https://mlsanity.com/paper/2505.22954/) (2025)
- [A Practical Guide to Multi-Objective Reinforcement Learning and Planning](https://mlsanity.com/paper/2103.09568/) (2021)
- [A Survey of Multi-Objective Sequential Decision-Making](https://mlsanity.com/paper/1402.0590/) (2014)
- [AIDE: AI-Driven Exploration in the Space of Code](https://mlsanity.com/paper/2502.13138/) (2025)
- [The Landscape of Agentic Reinforcement Learning for LLMs: A Survey](https://mlsanity.com/paper/2509.02547/) (2025)
- [Bandits with Knapsacks](https://mlsanity.com/paper/1305.2545/) (2013)
- [Self-Taught Optimizer (STOP): Recursively Self-Improving Code Generation](https://mlsanity.com/paper/2310.02304/) (2023)
- [Learning to Discover at Test Time](https://mlsanity.com/paper/2601.16175/) (2026)
- [AutoML-GPT: Automatic Machine Learning with GPT](https://mlsanity.com/paper/2305.02499/) (2023)
- [AutoKaggle: A Multi-Agent Framework for Autonomous Data Science Competitions](https://mlsanity.com/paper/2410.20424/) (2024)
- [ML-Bench: Evaluating Large Language Models and Agents for Machine Learning Tasks on Repository-Level Code](https://mlsanity.com/paper/2311.09835/) (2023)
- [Large Language Model Agent for Hyper-Parameter Optimization](https://mlsanity.com/paper/2402.01881/) (2024)
- [MLE-STAR: Machine Learning Engineering Agent via Search and Targeted Refinement](https://mlsanity.com/paper/2506.15692/) (2025)
- [MLE-Dojo: Interactive Environments for Empowering LLM Agents in Machine Learning Engineering](https://mlsanity.com/paper/2505.07782/) (2025)
- [MLRC-Bench: Can Language Agents Solve Machine Learning Research Challenges?](https://mlsanity.com/paper/2504.09702/) (2025)
- [B-STaR: Monitoring and Balancing Exploration and Exploitation in Self-Taught Reasoners](https://mlsanity.com/paper/2412.17256/) (2024)
- [Exploring LLM Agents for Cleaning Tabular Machine Learning Datasets](https://mlsanity.com/paper/2503.06664/) (2025)
- [Reinforcement Learning for Machine Learning Engineering Agents](https://mlsanity.com/paper/2509.01684/) (2025)
- [Dynamic Speculative Agent Planning](https://mlsanity.com/paper/2509.01920/) (2025)

## Abstract (arXiv)

Traditional reinforcement learning (RL) techniques focus on maximizing expected cumulative reward, where each action assumes to take a constant unit of time. However, this assumption does not hold for agentic RL tasks such as machine learning engineering (MLE) agents, where actions involve data loading, feature engineering, and model training that take variable durations. Efficiency matters in modern agentic RL where actions are costly. To address this limitation, we adapt from continuous-time RL and Semi-Markov Decision Process (SMDP) formulation and propose Reward-rate Policy Gradient (RPG), where we focus on optimizing the reward rate -- the long-term reward per unit of time. RPG estimates the reward rate from off-policy samples, then charges each action for the time it consumes at that rate. We first conduct theoretical analysis in the bandit setting to establish that RPG approximates the optimal reward rate and empirically demonstrate it outperforms baselines while avoiding enumeration over the policy space, a known issue for an existing method. We then further apply RPG on a small language model (Qwen3.5-4B) with self-improvement loops and empirically show it obtains higher rewards within a fixed time budget than vanilla RL on MLE-Bench and NanoGPT, with a 19.2% and 85.7% margin, respectively. Our method provides a practical solution for optimizing performance under wait time considerations in modern agentic RL tasks, where actions interact with external environments and cost time.