ML-Tool-Bench: Tool-Augmented Planning for ML Tasks

TL;DR

The paper introduces ML-Tool-Bench, a tool-augmented planning benchmark for end-to-end tabular ML tasks using a curated 61-tool toolkit and 15 Kaggle challenges. It proposes a scratchpad-based in-memory named-object store to manage large artifacts and enable reuse across multi-step plans. Through systematic comparisons, Hierarchical MCTS and MCTS-Shaped outperform ReAct and LATS on GPT-4o and GPT-4.1-mini, demonstrating the value of subtask decomposition and deterministic rewards for long-horizon planning. The work establishes strong baselines for tool-augmented ML planning and highlights memory organization and structured rewards as key factors for scalability and reliability.

Abstract

The development of autonomous machine learning (ML) agents capable of end-to-end data science workflows represents a significant frontier in artificial intelligence. These agents must orchestrate complex sequences of data analysis, feature engineering, model selection, and hyperparameter optimization, tasks that require sophisticated planning and iteration. While recent work on building ML agents has explored using large language models (LLMs) for direct code generation, tool-augmented approaches offer greater modularity and reliability. However, existing tool-use benchmarks focus primarily on task-specific tool selection or argument extraction for tool invocation, failing to evaluate the sophisticated planning capabilities required for ML Agents. In this work, we introduce a comprehensive benchmark for evaluating tool-augmented ML agents using a curated set of 61 specialized tools and 15 tabular ML challenges from Kaggle. Our benchmark goes beyond traditional tool-use evaluation by incorporating an in-memory named object management, allowing agents to flexibly name, save, and retrieve intermediate results throughout the workflows. We demonstrate that standard ReAct-style approaches struggle to generate valid tool sequences for complex ML pipelines, and that tree search methods with LLM-based evaluation underperform due to inconsistent state scoring. To address these limitations, we propose two simple approaches: 1) using shaped deterministic rewards with structured textual feedback, and 2) decomposing the original problem into a sequence of sub-tasks, which significantly improves trajectory validity and task performance. Using GPT-4o, our approach improves over ReAct by 16.52 percentile positions, taking the median across all Kaggle challenges. We believe our work provides a foundation for developing more capable tool-augmented planning ML agents.

Figures (5)

• Figure 1: An illustration of our named-object management scheme. Green circles denote tool calls; the blue rectangle denotes the scratchpad (a key–value store). Each tool can read any named object from the scratchpad and write outputs back to it, depending on their read-write behavior. Arrows into a tool indicate inputs; arrows from a tool to the scratchpad indicate outputs. read_csv is a set tool; fill_nans_with, fit_model, and predict are get–set tools; create_feature is an override tool. There are two read_csv tool calls in the figure, one for train data and one for test.
• Figure 2: An example illustration of how rewards are provided in MCTS-Shaped. If a particular stage is judged to be successfully completed at a node, a reward is given, which is used to update the value of all the nodes in this trajectory. It needs to be noted that these stage-wise rewards are only provided once per trajectory and only if the earlier stages were successfully completed.
• Figure 3: A schematic of Hierarchical MCTS. The task is decomposed into an ordered sequence of subtasks. For each subtask, MCTS searches for all solution nodes up to a prescribed maximum subtask depth to avoid locally optimal but globally suboptimal choices. The solution nodes from subtask $t$ are appended to the root of subtask $t{+}1$, and the search resumes. In the example, the solution node from subtask 1, $S^1_{1}$, initializes subtask 2; its solution nodes $S^1_{2}$ and $S^2_{2}$ initialize subtask 3, and so on. The highest-value solution at the final subtask is returned as the overall outcome of Hierarchical MCTS.
• Figure 4: Plots of consistency and median leaderboard percentile across all competitions in ML-Tool-Bench, for different planning algorithms. The top row shows results for GPT-4o, with the left plot showing consistency and the right plot showing the median leaderboard percentile. The bottom row shows results for GPT-4.1-mini. Hierarchical MCTS outperforms LATS and ReAct, followed by MCTS-Shaped, in terms of leaderboard performance, for both LLMs. Also, both Hierarchical MCTS and MCTS-Shaped improve consistency over the other baselines. In the box plots, $\mu$ denotes the mean, $\sigma$ denotes the standard deviation, and $M$ denotes the Median
• Figure 5: A pictorial illustration of Monte Carlo Tree Search