DSWorld:一种用于高效自主智能体的数据科学世界模型
DSWorld: A Data Science World Model for Efficient Autonomous Agents
摘要
尽管自主数据科学智能体在数据理解和决策方面具备强大的能力,但它们仍然严重依赖需要大量计算资源的试错式工作流。这种限制促使人们开发能够提前预测数据科学操作效果的模型。在本文中,我们提出了“数据科学世界模型”的概念——该模型通过预测当前工作流状态和候选操作的条件下环境状态的变化来模拟数据科学执行环境。我们还提出了DSWorld这一实用框架,它结合了结构化状态构建、成本意识化的路由机制、轻量级实际执行方式,以及用于处理复杂操作的基于LLM的模拟器。为了支持训练过程,我们构建了8K规模的转换轨迹数据集,并提出了“反思性世界模型优化”策略,这是一种能够识别误差的强化学习方法,有助于提升转换预测的准确性。实验表明,DSWorld使得基于强化学习的智能体训练速度提高了约14倍,基于搜索的推理速度则提高了3-6倍,同时仍能保持良好的性能水平;在转换预测任务上,DSWorld的表现比最先进的LLM基线模型高出35.6%。相关代码可访问于https://anonymous.4open.science/r/DSWorld。
English Abstract
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.