DSWorld: 効率的な自律エージェントのためのデータサイエンス世界モデル

要旨

データ理解と意思決定において強力な能力を持つ一方で、自律的なデータサイエンスエージェントは依然として高コストな計算を伴う試行錯誤のワークフローに大きく依存している。このボトルネックが、実際の実行前にデータサイエンス操作の効果を予測できるモデルの動機となっている。本論文では、データサイエンスワールドモデルの概念を紹介する。これは、現在のワークフロー状態と候補操作に基づいて環境状態遷移を予測することで、データサイエンス実行環境をモデル化するものである。さらに、構造化状態構築、コスト考慮ルーティング、軽量な実際の実行、および高コスト操作のためのLLMベースのシミュレータを組み合わせた実用的フレームワークDSWorldを提案する。訓練を支援するため、8K規模の遷移軌跡データセットを構築し、遷移予測を改善するためのエラーを考慮した強化学習戦略であるリフレクティブワールドモデル最適化を導入する。実験により、DSWorldは競争力のある性能を維持しつつ、RLベースのエージェント訓練を約14倍、探索ベースの推論を約3～6倍高速化し、遷移予測タスクにおいて最強のLLMベースラインを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.