# GRACE-DS: a Guarded Reward-guided Agent Correction Environment in Data Science

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Web, search, and research agents, one of 16 directions in LLM agents and tool use.
- Builds on 17 of its 37 references that are in the map.
- First posted on arXiv 14 Jun 2026.

## Placed in

- LLM agents and tool use / [Web, search, and research agents](https://mlsanity.com/agents/web-search-and-research-agents/)

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

- [MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering](https://mlsanity.com/paper/2410.07095/) (2024)
- [DS-1000: A Natural and Reliable Benchmark for Data Science Code Generation](https://mlsanity.com/paper/2211.11501/) (2022)
- [MLAgentBench: Evaluating Language Agents on Machine Learning Experimentation](https://mlsanity.com/paper/2310.03302/) (2023)
- [ScienceAgentBench: Toward Rigorous Assessment of Language Agents for Data-Driven Scientific Discovery](https://mlsanity.com/paper/2410.05080/) (2024)
- [AIDE: AI-Driven Exploration in the Space of Code](https://mlsanity.com/paper/2502.13138/) (2025)
- [InfiAgent-DABench: Evaluating Agents on Data Analysis Tasks](https://mlsanity.com/paper/2401.05507/) (2024)
- [Machine Learning Operations (MLOps): Overview, Definition, and Architecture](https://mlsanity.com/paper/2205.02302/) (2022)
- [DSBench: How Far Are Data Science Agents from Becoming Data Science Experts?](https://mlsanity.com/paper/2409.07703/) (2024)
- [MLGym: A New Framework and Benchmark for Advancing AI Research Agents](https://mlsanity.com/paper/2502.14499/) (2025)
- [BLADE: Benchmarking Language Model Agents for Data-Driven Science](https://mlsanity.com/paper/2408.09667/) (2024)
- [DataSciBench: An LLM Agent Benchmark for Data Science](https://mlsanity.com/paper/2502.13897/) (2025)
- [DABstep: Data Agent Benchmark for Multi-step Reasoning](https://mlsanity.com/paper/2506.23719/) (2025)
- [MLE-Dojo: Interactive Environments for Empowering LLM Agents in Machine Learning Engineering](https://mlsanity.com/paper/2505.07782/) (2025)
- [InsightBench: Evaluating Business Analytics Agents Through Multi-Step Insight Generation](https://mlsanity.com/paper/2407.06423/) (2024)
- [DSGym: A Holistic Framework for Evaluating and Training Data Science Agents](https://mlsanity.com/paper/2601.16344/) (2026)
- [IDA-Bench: Evaluating LLMs on Interactive Guided Data Analysis](https://mlsanity.com/paper/2505.18223/) (2025)
- [TML-Bench: Benchmark for Data Science Agents on Tabular ML Tasks](https://mlsanity.com/paper/2603.05764/) (2026)

## Abstract (arXiv)

We introduce GRACE-DS, a Guarded Reward-guided Agent Correction Environment in Data Science for pre-deployment evaluation of LLM-powered AutoML agents. GRACE-DS is a set of evaluation metrics in an isolated environment that can be applied to tabular ML tasks specific to a particular organization. It exposes agents to realistic workflow stages, from planning and data inspection through feature engineering, model development, validation, and code repair to final submission, while hidden executable validators measure not only final predictive performance but also leakage avoidance, reproducibility, protocol validity, correction behavior, and reward alignment. The strongest structured regime, flexible iterative interaction (our approach), achieves higher end-to-end normalized hidden-test quality than single-shot generation, unstructured interaction, and restart-based baselines, while also improving protocol-valid completion. Validated across more than 7,000 episodes, these results establish GRACE-DS as a robust platform for assessing the capacity of LLM-based AutoML agents to execute machine learning workflows under production-like conditions and in accordance with organization-specific requirements.
