DSWorld：面向高效自主智能体的数据科学世界模型

摘要

儘管在資料理解與決策方面能力強大，自主資料科學智能體仍高度依賴試錯工作流程，其中涉及昂貴的計算成本。此瓶頸促使研究人員開發能在實際執行前預測資料科學操作效果的模型。本文提出「資料科學世界模型」概念，該模型透過預測當前工作流程狀態與候選操作條件下的環境狀態轉移，來模擬資料科學執行環境。我們進一步提出實用框架 DSWorld，該框架結合了結構化狀態建構、成本感知路由、輕量級實際執行，以及針對昂貴操作的基於大語言模型的模擬器。為支援訓練，我們建構了一個 8000 規模的轉移軌跡資料集，並引入「反思世界模型優化」，這是一種用於改善轉移預測的錯誤感知強化學習策略。實驗結果顯示，DSWorld 能將基於強化學習的智能體訓練加速約 14 倍，並將基於搜索的推理加速約 3 至 6 倍，同時保持具競爭力的性能；在轉移預測任務上，DSWorld 比最強的大語言模型基準高出 35.6%。程式碼請見 https://anonymous.4open.science/r/DSWorld。

English

Despite strong capabilities in data understanding and decision-making, autonomous data science agents still heavily rely on trial-and-error workflows that involve expensive computation. This bottleneck motivates models that can anticipate the effects of data science operations before real execution. In this paper, we introduce the concept of Data Science World Model, which model the data science execution environment by predicting environment state transitions conditioned on current workflow states and candidate operations. We further propose DSWorld, a practical framework that combines structured state construction, cost-aware routing, lightweight real execution, and an LLM-based simulator for expensive operations. To support training, we construct an 8K-scale transition trajectory dataset and introduce Reflective World Model Optimization, an error-aware reinforcement learning strategy for improving transition prediction. Experiments show that DSWorld accelerates RL-based agent training by approximately 14times and search-based inference by approximately 3-6times while maintaining competitive performance, and outperforms the strongest LLM baseline by 35.6% on transition prediction tasks. The code is available at https://anonymous.4open.science/r/DSWorld.