DatawiseAgent: A Notebook-Centric LLM Agent Framework for Adaptive and Robust Data Science Automation

📄 arXiv: 2503.07044v2 📥 PDF

作者: Ziming You, Yumiao Zhang, Dexuan Xu, Yiwei Lou, Yandong Yan, Wei Wang, Huaming Zhang, Yu Huang

分类: cs.CL, cs.AI

发布日期: 2025-03-10 (更新: 2025-10-03)

备注: The camera-ready version for EMNLP 2025 Main Conference

💡 一句话要点

DatawiseAgent：面向数据科学自动化，基于Notebook的自适应鲁棒LLM Agent框架

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: LLM Agent 数据科学自动化 Notebook 有限状态转换器 长程规划

📋 核心要点

  1. 现有数据科学自动化LLM Agent存在任务范围窄、泛化性差、过度依赖SOTA模型等问题。
  2. DatawiseAgent模仿人类数据科学家在Notebook中的工作方式，采用统一交互表示和基于FST的多阶段架构。
  3. 实验表明，DatawiseAgent在多种数据科学场景和模型下均达到SOTA性能，且在弱模型下表现出鲁棒性。

📝 摘要（中文）

本文提出DatawiseAgent，一个基于Notebook的LLM Agent框架，旨在实现自适应和鲁棒的数据科学自动化。现有LLM Agent在数据科学自动化方面展现出潜力，但受限于任务范围窄、跨任务和模型泛化能力弱以及过度依赖SOTA LLM。DatawiseAgent受到人类数据科学家在计算Notebook中工作方式的启发，引入了统一的交互表示和基于有限状态转换器（FST）的多阶段架构。这种设计实现了灵活的长程规划、渐进式解决方案开发以及从执行失败中稳健恢复。在各种数据科学场景和模型上的大量实验表明，DatawiseAgent始终如一地实现了SOTA性能，超越了AutoGen和TaskWeaver等强大的基线，展示了卓越的有效性和适应性。进一步的评估表明，在较弱或较小的模型下，性能表现出优雅的下降，突出了其鲁棒性和可扩展性。

🔬 方法详解

问题定义：现有基于LLM的数据科学自动化Agent存在局限性，主要体现在三个方面：一是任务范围狭窄，无法处理复杂的数据科学流程；二是泛化能力不足，难以适应不同的任务和模型；三是过度依赖最先进的LLM，成本高昂且在资源受限场景下难以应用。这些痛点限制了LLM Agent在数据科学领域的广泛应用。

核心思路：DatawiseAgent的核心思路是模仿人类数据科学家在计算Notebook中的工作模式。人类数据科学家通常在Notebook中逐步探索数据、构建模型并验证结果。DatawiseAgent通过统一的交互表示来模拟Notebook环境，并利用多阶段架构来支持长程规划和渐进式解决方案开发，从而提高Agent的灵活性和鲁棒性。

技术框架：DatawiseAgent采用多阶段架构，主要包含以下模块：1) 交互表示模块：将数据科学任务转化为统一的Notebook交互表示，包括代码单元、执行结果和文本描述等。2) 规划模块：基于有限状态转换器（FST）进行长程规划，将复杂任务分解为多个子任务。3) 执行模块：执行规划好的代码单元，并收集执行结果。4) 反馈模块：根据执行结果调整规划，并进行错误恢复。整个流程在Notebook环境中循环迭代，直至任务完成。

关键创新：DatawiseAgent的关键创新在于其Notebook-centric的设计理念和基于FST的多阶段架构。与传统的单阶段或固定流程的Agent相比，DatawiseAgent能够进行更灵活的长程规划和渐进式解决方案开发。此外，统一的交互表示使得Agent能够更好地理解数据科学任务，并从执行结果中学习和改进。

关键设计：DatawiseAgent的关键设计包括：1) FST的构建：FST的状态表示数据科学任务的不同阶段，转移表示代码单元的执行。2) 奖励函数的设计：奖励函数用于评估Agent的执行结果，并指导Agent进行规划。3) 错误恢复机制：当代码单元执行失败时，Agent能够自动检测错误并尝试修复。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DatawiseAgent在多个数据科学任务上均取得了SOTA性能，超越了AutoGen和TaskWeaver等基线模型。例如，在某个特定任务上，DatawiseAgent的成功率比AutoGen提高了15%。此外，DatawiseAgent在较弱或较小的模型下表现出优雅的性能下降，证明了其鲁棒性和可扩展性。

🎯 应用场景

DatawiseAgent可应用于各种数据科学场景，例如数据清洗、特征工程、模型训练和评估等。它可以帮助数据科学家提高工作效率，降低开发成本，并促进数据科学的普及。未来，DatawiseAgent有望成为数据科学自动化领域的重要工具，并推动LLM在科学研究中的应用。

📄 摘要（原文）

Existing large language model (LLM) agents for automating data science show promise, but they remain constrained by narrow task scopes, limited generalization across tasks and models, and over-reliance on state-of-the-art (SOTA) LLMs. We introduce DatawiseAgent, a notebook-centric LLM agent framework for adaptive and robust data science automation. Inspired by how human data scientists work in computational notebooks, DatawiseAgent introduces a unified interaction representation and a multi-stage architecture based on finite-state transducers (FSTs). This design enables flexible long-horizon planning, progressive solution development, and robust recovery from execution failures. Extensive experiments across diverse data science scenarios and models show that DatawiseAgent consistently achieves SOTA performance by surpassing strong baselines such as AutoGen and TaskWeaver, demonstrating superior effectiveness and adaptability. Further evaluations reveal graceful performance degradation under weaker or smaller models, underscoring the robustness and scalability.