MLToolBench: Learning Tool-Augmented Agents for Machine Learning Development

Xin Yu; Lizhu Zhang; Jiamu Bai; Yanhong Wu; Zellux Wang; Serena Li; Weiwei Li; Lingzhou Xue

cs.AISep 29, 2026

MLToolBench: Learning Tool-Augmented Agents for Machine Learning Development

Authors: Xin Yu, Lizhu Zhang, Jiamu Bai, Yanhong Wu, Zellux Wang, Serena Li, Weiwei Li, Lingzhou Xue, +2 more

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.

Figures & tables

Figure 1: Learning to find and utilize diagnostic evidence. (a) In-domain success under three setups: Gemini 3.6 Flash with and without guidance, and Qwen3.5-35B-A3B after SFT + RL, all using MLToolBench. Guided Gemini receives solution-informed hints. Signed pp labels show Qwen’s gains over unguided Gemini. (b) Failure-driven tool construction, guided SFT, and RL with Spice teach agents to acquire and use diagnostic evidence.

Figure 2: Schematic motivating cases for MLToolBench. Top: local feature diagnostics guide preprocessing checks before an architectural change. Bottom: missingness and imbalance motivate checking an apparent validation gain before submission. The contrast illustrates why evidence acquisition should precede intervention; Section 3 separately evaluates aggregate task performance.

Figure 3: Construction of MLToolBench. Failure analysis motivates diagnostic tool candidates. The agent re-runs the original task with each candidate; tools that help resolve the original failure are retained and organized by domain, while unsuccessful candidates are refined and re-tested.

Closed-Source Models
		In-domain		Out-of-domain
Backbone	Model / training	Original	+ MLToolBench	Original	+ MLToolBench
Gemini 3.6 Flash		65.20%	61.20%	52.00%	46.00%
Claude Sonnet 5		64.80%	62.40%	45.00%	43.00%
Open-Source Models
Qwen2.5-Coder-14B		45.20%	40.40%	34.00%	32.00%

Table 1: Main success rates (%; 250 ID and 100 OOD runs per condition). Fine-tuned models use MLToolBench during training. Original uses standard MLAgentBench tools without MLToolBench diagnostics; + MLToolBench adds these diagnostics at evaluation. Bold marks column maxima.

Figure 4: Feedback magnitude on fixed SFT trajectories. Mean absolute unscaled scores: TIPS potential differences and Spice

d_{i,s}

. Panels use frozen, task-unadapted scorers; raw scales differ. Editing/execution includes code edits and execution. Absolute values neither reveal score signs nor establish correct credit assignment.

Figure 5: Training and guidance ablations. (a) Mean binary outcome reward $R_{i}$ for 8B (blue) and 35B (orange): outcome-only GRPO (dashed) and SPICE (solid). Both curves exclude shaping rewards. Final checkpoint: step 100. (b) Gemini ID/OOD success (%); unguided results match Table 1 . Training outcome success is distinct from held-out success.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Partition	Instances	Runs per task	Total / purpose
Training	80	100	8,000 raw demonstrations per SFT condition
ID evaluation	25	10	250 evaluation runs per condition
OOD evaluation	10	10	100 evaluation runs per condition

Appendix

Table 2: Collection and evaluation budgets. RL collects fresh on-policy rollouts.

RL configuration	Value
Optimizer updates / reported checkpoint	100 / final step 100
Rollout batch	8 tasks $\times$ 8 rollouts = 64 trajectories
Optimizer / learning rate / scheduler	AdamW / $1\times 10^{-6}$ / cosine
Minibatch size / update epochs / accumulation	128 / 1 / 1
Advantage estimation	Task-group mean centering; no standard-deviation scaling
Clip parameter / KL coefficient	0.2 / 0 (KL disabled)

Appendix

Table 3: RL configuration and shaping-scale calibration protocol. Rollout collection and optimization batch sizes are distinct. The calibration range refers to reward magnitude, not the coefficient.

Figure 6: Full-SFT training conditions and evaluation interfaces. Each group compares SFT with and without MLToolBench demonstrations at a fixed evaluation interface. Original uses standard MLAgentBench tools without MLToolBench diagnostics; + MLToolBench adds these diagnostics at evaluation. Both panels share a scale and reproduce the tool-enabled SFT results in Table 1 .

Figure 7: RL-stage and reward comparisons with MLToolBench. Full SFT and RL + Spice values reproduce Table 1 . Outcome-only RL (

\beta=0

), RL + TIPS, and RL + Spice are shown for both backbones. Gains over SFT measure the entire RL stage.

35B training condition	ID success (%)	OOD success (%)
Outcome-only RL	48.00 (120/250)	44.00 (44/100)
SPICE, full solution summaries	69.20 (173/250)	48.00 (48/100)
SPICE, reduced solution summaries	67.20 (168/250)	41.00 (41/100)
SPICE, shuffled scores	59.20 (148/250)	45.00 (45/100)

Appendix

Table 4: Privileged-summary and score-assignment ablations. Success rates for Qwen3.5-35B-A3B with MLToolBench. Shuffling permutes scores within tasks and tool categories. Full-summary SPICE reproduces the main result.

Figure 8: Tool-use coverage by model and training condition. Mean number of rollouts, out of five per task, invoking data-management or verification tools at least once. Separators distinguish unadapted baselines, full SFT, and full SFT + RL with Spice .

	Coverage (out of 5)		Mean per rollout
Model / training	Data management	Verification	Calls	Distinct tools
8B Full SFT	2.5	3.4	5	2
8B RL + Spice	2.5	3.3	11	6
35B Full SFT	2.6	3.9	7	4
35B Outcome-only RL	0.8	1.3	9	2
35B RL + TIPS	1.4	1.6	10	3

Appendix

Table 5: MLToolBench invocation, frequency, and diversity. Coverage counts rollouts invoking each category at least once; calls count invocations, and distinct tools count unique MLToolBench tool IDs within a rollout. Editing/execution actions are excluded.

35B training objective	Observation	ID success (%)
Outcome-only RL	True	48.00
	Masked	46.40
	Paraphrased	47.60
RL + Spice	True	69.20
	Masked	54.40
	Paraphrased	68.40

Appendix

Table 6: Diagnostic-observation intervention on ID tasks. Both 35B policies retain all tools. All 250 ID rollouts are included. We randomly intervene on one diagnostic return where available; no-call trajectories retain their original outcomes in all arms.

Tool	Returned evidence / behavior
Discover Log Metrics	Metric names, sample values, and occurrence counts in logs or captured stdout.
Extract Training Metric Series	Value–step series for requested logged metrics.
Verify Hyperparam Effective	Static parameter-definition sites, conflicting values, and environment overrides.
Estimate Peak Memory	Observed peak GPU memory during a short training-script probe; OOM status.
Trace Tensor Shapes	Per-layer tensor shapes or failure traces from a CPU dummy forward pass.
Check Config Consistency	Rule-based checks of configuration relations and incompatible settings.

Appendix

Table 7: Seven shared diagnostic tools. Names match the registered action interface.

Tool	Returned evidence / behavior
Tabular (12 tools)
Detect Missing Values	Per-column missingness and suggested imputation strategies.
Identify Feature Types	Inferred column types, unique counts, and ID/high-cardinality flags.
Detect Label Imbalance	Class counts, imbalance ratios, and resampling suggestions.
Check Feature Normalization	Numeric ranges, means, standard deviations, and scaling suggestions.
Detect Outliers	Per-column outlier counts and IQR bounds.

Appendix

Table 8: Complete domain-module inventory: 60 tools across five modules. Descriptions summarize implemented checks, including heuristic suggestions.

Task

Task-specific diagnostic guidance

ogbn-arxiv

Context. Improve paper-topic node classification from an MLP, following the diagnostic example in Figure 2 . The presence of graph edges alone does not establish that replacing the MLP with GraphSAGE will improve performance. Inspect. Check node-feature ranges, per-column scales, and the active preprocessing pipeline. Compare training and validation curves, then assess whether graph connectivity and label homophily on permitted training labels motivate neighborhood aggregation. Act on evidence. If features have disproportionate scales or missing normalization, fit the transformation on permitted training data and re-run the MLP first. Compare the corrected MLP with the original baseline under the same split, evaluator, and budget. Only then test GraphSAGE if graph diagnostics motivate it; retain the same feature preprocessing and assess the additional gain from neighborhood information. Do not use hidden labels for diagnosis or treat a more complex architecture as an automatic improvement.

cifar10

Context. Improve ten-class image classification from a small CNN trained for five epochs; configured instances vary input scaling, training class balance, or cropping strength. Inspect. Check pixel ranges and channel statistics after preprocessing, inspect transformed images for retained object content, and compare training class counts with per-class validation performance. Read training and validation curves to distinguish optimization problems from overfitting. Act on evidence. If inputs retain an unintended integer scale, correct scaling and use training-derived normalization. If crops remove most object content, reduce cropping strength while preserving the expected input shape. If minority classes are underrepresented and perform poorly, test balanced sampling or class-weighted loss. If these checks are healthy, retain preprocessing and investigate optimization. Re-run under the same split and budget, comparing validation accuracy with the starter before accepting an edit.

Appendix

Table 9: Task-specific guidance examples for SFT demonstration collection. Checks and conditional actions reflect the task’s starter, data conditions, and evaluation constraints; diagnostic observations must be obtained through actual tool calls.

SFT configuration	Value
Adaptation	Full parameters
Epochs / optimizer	3 / AdamW
Learning rate / scheduler / warmup	$1\times 10^{-5}$ / cosine / 0
Effective batch size	64, using gradient accumulation
Maximum sequence length	10,240 tokens
Hardware / precision	8 $\times$ H200 / BF16

Appendix

Table 10: SFT optimization configuration.

Condition	Agent-visible prompt	Training signal
Outcome-only RL	Shared system + task + registry + rollout history	Outcome reward only
RL + TIPS	Same agent prompt as outcome-only RL	Outcome + TIPS shaping
RL + SPICE	Same agent prompt as outcome-only RL	Outcome + SPICE shaping
SFT demonstration collection	Shared system + task + registry + history; guidance with MLToolBench	Full SFT on filtered trajectories
Unguided evaluation, both interfaces	Shared system + task + condition-specific registry + history	None
Guided Gemini 3.6 Flash evaluation	Shared system + task + registry + teacher guidance	None

Appendix

Table 11: Prompt composition used in the experiments. Outcome-only, TIPS, and SPICE share the acting-policy prompt; reward computation differs externally. Privileged scoring context is not appended to the acting policy’s instructions.

Instance suffix	Source	Held-out source / target
cv-ood-01	APTOS 2019	Unseen medical-image source and ordinal target
cv-ood-02	RSNA Breast Cancer	Unseen mammography source and probabilistic-F1 metric
nlp-ood-01	AI4Code	Unseen code-notebook source and ranking target
nlp-ood-02	BabyLM	Unseen language-modeling source and next-token target
tabular-ood-01	California Housing	Unseen dataset and reconstruction target; downstream R² is secondary
tabular-ood-02	Breast Cancer Wisconsin	Unseen clinical tabular source and calibration target

Appendix

Table 12: OOD instances and the source or objective held out from training. Source links and exact construction parameters are included in the JSON manifest.

Instance suffix	Template	Configuration	Seed
cv-train-01	cv-cifar10	healthy	10000
cv-train-02	cv-cifar10	input scale	10001
cv-train-03	cv-cifar10	rare class sampling	10002
cv-train-04	cv-cifar10	overaggressive crop	10003
cv-train-05	cv-cifar100	healthy	10010
cv-train-06	cv-cifar100	normalization scale	10011

Appendix

Table 13: Complete 80/25/10 task-instance index. Condition names describe author-side configuration metadata and are not exposed in the actor prompt.

Explore similar work

Jul 6, 2026cs.CL

ToolFailBench: Diagnosing Tool-Use Failures in LLM Agents

Tool calling is central to modern language model agents, but aggregate benchmark scores often hide where tool use fails. A model that never calls a needed tool and a model that calls the tool but ignores the result can look similar under final task accuracy. We introduce ToolFailBench, a diagnostic benchmark for measuring tool-use failures across 1,000 tasks in finance, medicine, law, cybersecurity, and real estate. Tool-required tasks return values the model wouldn't guess, forcing it to trust the tool while control tasks attach the same tools but should be answered directly. We label each trace with Tool-Skip, Result-Ignore, Output-Fabrication, and Unnecessary-Tool-Use, using a rule classifier and two LLM judges aggregated by majority vote. Across 19 headline models, the best reaches 86.33% Clean Tool-Use Rate, showing that faithful tool use is not saturated. More importantly, models with similar aggregate scores fail in different ways: most stay disciplined on no-tool controls, while Llama-3.1 models show an Always-Call pattern, and at the same parameter scale Llama-3.1-70B and Qwen2.5-72B differ by 89 percentage points on control-task accuracy. Tool-use evaluation should measure not only whether agents call tools, but whether they use tool outputs correctly and avoid tools when none is needed.

Harsh Soni

UC Berkeley

Large Language Model Agents Large Language Model Tool Use

Ask PDF

Jun 11, 2026cs.CL

HyperTool: Beyond Step-Wise Tool Calls for Tool-Augmented Agents

Tool-augmented LLM agents commonly rely on step-wise atomic tool calls, where each invocation, observation, and value transfer is exposed in the main reasoning trace. This creates an \emph{execution-granularity mismatch}: locally deterministic tool workflows are unfolded into repeated model-visible decisions, consuming context and forcing the model to manage low-level dataflow in the trace. We introduce \textbf{HyperTool}, a unified executable MCP-style tool interface that changes the model-visible unit of tool execution. A model invokes HyperTool with a code block that can call existing tools through their original schemas, manipulate returned values, and pass intermediate results locally, folding deterministic tool subroutines into a single outer call. To train models to use this interface, we synthesize HyperTool-format trajectories from cross-tool compositional tasks and verify them in real MCP environments. On MCP-Universe, HyperTool improves average accuracy from 15.69% to 35.29% on Qwen3-32B and from 9.93% to 33.33% on Qwen3-8B, and surpass GPT-OSS and Kimi-k2.5 on average accuracy, showing that our HyperTool can substantially improve multi-step tool use.

Yaxin Du, Yifan Zhou, Yujie Ge +7

Shanghai Jiao Tong University · IQuest Research · Beijing University of Aeronautics and Astronautics

Tool Invocation Agentic Workflow Design

Ask PDF

Sep 16, 2026cs.AI

TuiML: Machine Learning for AI Agents

Machine-learning libraries such as Weka and scikit-learn were designed for human programmers. Language-model agents now use these same libraries by recalling APIs from memory and writing code, an approach that hides what a library offers, delays errors until runtime, and loses experimental state between turns. We present TuiML, a self-contained machine-learning library built for AI agents, with native algorithms across supervised, unsupervised, time-series, data handling, tuning, and evaluation tasks. Every component describes itself through machine-readable metadata and parameter schemas, so an agent can search the library, inspect components, compose validated workflows, and register new ones that become discoverable in turn. Every call is validated, seeded, and traced, and sessions export as runnable notebooks, making experiments reproducible by construction. One specification layer drives the Model Context Protocol (MCP), agent-framework adapters, a Python API, a CLI, and local model serving, while data and models never leave the machine. Benchmarks show TuiML remains predictively competitive with scikit-learn and Weka. While looking like a conventional library to a human user, TuiML is designed for agents first, allowing them to read, extend, and operate machine learning autonomously. TuiML is open source, with documentation at https://tuiml.ai.

Nilesh Verma, Nick Lim, Albert Bifet +1

AI Institute, University of Waikato Hamilton 3216, New Zealand · LTCI, T´el´ecom Paris, Institut Polytechnique de Paris 19 place Marguerite Perey, 91120 Palaiseau, France

Artificial Intelligence Agents Model Context Protocol

Ask PDF