MLToolBench: Learning Tool-Augmented Agents for Machine Learning Development
cs.AI
Submitted: 2026-09-29
Updated: 2026-10-06
License: http://creativecommons.org/licenses/by/4.0/
The gist: Machine learning engineering (MLE) agents have made substantial progress, but learning through ML experimentation remains costly in time and computation.
Abstract
Machine learning engineering (MLE) agents have made substantial progress, but learning through ML experimentation remains costly in time and computation. Synthetic environments reduce these costs while introducing variations in data and experimental settings that require task-specific diagnosis. Access to diagnostic tools alone does not ensure that agents learn when to use them or how to act on their findings. We introduce ToolMLBench, a suite of executable tools for data inspection, code verification, and experiment diagnosis, together with an SFT and RL pipeline for learning their use. Diagnostic calls acquire evidence whose value depends on subsequent decisions, so final outcomes provide limited guidance on which calls to reinforce. We address this challenge with SPICE, which measures how privileged context changes the likelihood of a sampled tool action and uses this difference as a turn-level reward alongside the final outcome. We train on 80 synthetic tasks and evaluate on 25 in-domain and 10 out-of-domain tasks. Providing tool interfaces and descriptions alone yields inconsistent gains across unadapted models. With the same diagnostic interface, our training pipeline raises in-domain success from 24.8% to 52.4% for Qwen3-8B and from 35.6% to 69.2% for Qwen3.5-35B-A3B. The latter also improves from 31% to 48% out-of-domain, supporting learned diagnostic tool use on held-out sources and targets.
Sources
- AceGRPO: Adaptive Curriculum Enhanced Group Relative Policy Optimization for Autonomous Machine Learning Engineering
- AI Scientist via Synthetic Task Scaling
- When Should Service Agents Reconsider? Difficulty-Routed Control in Customer-Service Operations
- TACO: Tool-Augmented Credit Optimization for Agentic Tool Use
- DeepSeek-Coder: When the Large Language Model Meets Programming -- The Rise of Code Intelligence
- AIDE: AI-Driven Exploration in the Space of Code
- Learning the ARTS of Search for Automated Discovery
- Agent0-VL: Exploring Self-Evolving Agent for Tool-Integrated Vision-Language Reasoning
- Matryoshka Agent: Unfolding Sub-Agents for Long-Horizon Machine Learning Engineering
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- Hindsight Credit Assignment for Long-Horizon LLM Agents
- SEED: Self-Evolving On-Policy Distillation for Agentic Reinforcement Learning
- CurveShift: Is Agent Progress Scalar? Separating Level from Shape
- TRIAGE: Role-Typed Credit Assignment for Agentic Reinforcement Learning
- Self-Distilled RLVR
- Preference-Based Self-Distillation: Beyond KL Matching via Reward Regularization
- Synthetic Sandbox for Training Machine Learning Engineering Agents
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems