Technical Summary: GRACE-DS

Problem Statement

Organizations are increasingly evaluating Large Language Model (LLM) agents not merely as coding assistants but as potential participants in end-to-end data science and machine learning workflows. However, existing benchmarks for ML agents primarily focus on final task performance (e.g., leaderboard scores) or general code correctness. These metrics fail to capture critical enterprise deployment requirements: whether an agent can solve specific organizational tasks using internal data, adhere to governance and privacy constraints, avoid data leakage, ensure reproducibility, and recover from execution errors.

In production environments, a model that performs well on public benchmarks may fail due to messy internal schemas, unstable data distributions, or methodological flaws like target leakage and non-reproducible pipelines. Furthermore, a single-shot generation approach does not reflect the iterative nature of real-world ML work, where analysts inspect data, revise assumptions, and debug failures. There is a lack of a controlled evaluation harness that measures an agent's ability to execute ML workflows under realistic, organization-specific constraints while distinguishing between a good final score and a trustworthy process.

Methodology: GRACE-DS

GRACE-DS (Guarded Reward-guided Agent Correction Environment in Data Science) is a controlled evaluation platform designed for pre-deployment assessment of LLM-powered AutoML agents on tabular ML tasks. It operates within an isolated, evaluator-owned sandbox environment.

Core Design Principles

1. Iterative Workflow: Unlike single-shot benchmarks, GRACE-DS structures the agent's interaction as a multi-stage workflow. Agents must navigate specific states: Planning, Data Inspection (EDA), Feature Engineering, Model Development, Validation, Code Repair, and Final Submission.
2. Guarded Evaluation: The environment employs hidden, executable validators that operate independently of the agent's visible data. These validators check for:
   • Integrity & Leakage: Ensuring no target leakage, no access to evaluator-private namespaces, and no tampering with protected snapshots.
   • Process Reliability: Verifying reproducibility, valid validation practices (e.g., proper stratification), and adherence to protocol.
   • Critical Errors: Detecting fatal methodological flaws (e.g., training on validation labels) which immediately zero out the reward.
3. Decomposed Reward System: The system provides a reward signal ($w$) composed of three terms:
   • Performance ($r_{perf}$, 55%): Normalized hidden-test quality.
   • Plan Coverage ($r_{plan}$, 15%): Adherence to the planning stage.
   • Code Quality ($r_{code}$, 30%): Based on validators for reproducibility, pipeline structure, and checklist compliance.
   • Note: The reward is a proxy for progress but is explicitly separated from the final hidden-test score to prevent reward hacking.
4. Hidden Test & Reproducibility: The final submission must be a reproducible artifact (e.g., a predict_fn or fitted Pipeline) that can be re-run on raw input data. The final score is computed once on a hidden test set by the evaluator, ensuring the agent cannot tune to the test labels.

Experimental Setup

The evaluation spans 10 production-like tabular tasks (including post-cutoff Kaggle competitions, industry datasets, and synthetic data) and 8 frontier LLMs (ranging from open-weight MoE models to proprietary systems). The study compares several regimes:

• Baselines: Single-shot generation, unstructured interaction, and restart-based approaches.
• Structured Iteration: Flexible iterative interaction (agent chooses stage order) and fixed-stage iteration (prescribed order).
• Stress Tests: Reward-optimization prompts (instructing agents to maximize process reward) and red-team adversarial prompts (instructing agents to ignore task solving and focus solely on reward).

Key Contributions

1. Enterprise-Oriented Evaluation Harness: GRACE-DS introduces a framework for testing AutoML agents under realistic workflow, validation, reproducibility, and data-governance constraints, moving beyond public leaderboard scores.
2. Guarded Protocol: A novel protocol combining hidden executable validators, evaluator-owned validation, reproducible final submissions, and process-level reward decomposition.
3. Demonstration of Structured Iteration: The paper provides empirical evidence that flexible structured iteration significantly outperforms single-shot, unstructured, and restart-based baselines in both end-to-end effectiveness and protocol-valid completion rates.
4. Deployment-Relevant Diagnostics: The platform generates detailed diagnostics including runtime, iteration counts, quality deltas, reward alignment, error statistics, and recovery behavior, offering insights into the "why" behind an agent's performance.
5. Safety and Reliability Validation: Through reward-optimization and red-team stress tests, the study demonstrates that process rewards are useful for monitoring but cannot replace isolated hidden-test evaluation, validating GRACE-DS as a safety and reliability harness.

Results

• Performance: The flexible iterative regime achieved the highest end-to-end normalized hidden-test quality (0.754), significantly outperforming single-shot generation (0.536), unstructured agents (0.527), and restart-based baselines (0.672–0.686).
• Reliability: Structured iteration improved protocol-valid completion rates to 96.9%, compared to 88.8% for single-shot and 69.2% for unstructured agents. The structured feedback loop effectively converted potential failures into valid submissions.
• Error Correction: Structured regimes reduced critical methodological errors (e.g., target leakage) to near zero (0.2% for flexible iteration vs. 5.2–7.1% for restart baselines). The CODE_FIX loop allowed agents to recover from execution errors with high success rates (93.4%).
• Reward vs. Quality: While process rewards correlated with task quality within specific tasks, they could not replace hidden-test evaluation. Adversarial agents attempting to maximize process rewards failed to improve hidden-test quality and were blocked by validators from committing critical errors.
• Model Scale Effects: Larger models benefited more from structured environments, but the flexible approach acted as a capability equalizer, narrowing the performance gap between models compared to single-shot approaches. Rigid, fixed-stage protocols proved fragile for smaller models, whereas flexible iteration remained robust across model scales.
• Efficiency: Flexible iteration was more efficient than restart-based baselines, requiring fewer LLM calls and less wall-clock time to achieve higher quality.

Significance

The paper positions GRACE-DS as a necessary evolution in AI agent evaluation for enterprise settings. It argues that for organizations to trust LLM agents with internal data and production workflows, evaluation must shift from "what score did the agent get?" to "how did the agent get there, and can we trust the process?"

The study concludes that structured, flexible iteration is the most effective regime for deploying AutoML agents, offering a balance of high predictive performance, methodological safety, and robust error recovery. By combining hidden validators with process-level feedback, GRACE-DS provides a mechanism to detect and prevent deployment-critical failures (such as data leakage) that traditional benchmarks miss, thereby serving as a critical pre-deployment safety and reliability harness for organizations adopting AI-driven data science.