DSWorld: A Data Science World Model for Efficient Autonomous Agents
Artificial Intelligence
Summary
The authors address the problem of slow and costly trial-and-error processes in autonomous data science agents by creating a way to predict what will happen before actually running operations. They introduce the Data Science World Model, which forecasts changes in the data environment based on current steps and possible next moves. Their system, DSWorld, uses a smart combination of quick tests and a language model simulator to speed up learning and decision-making. Testing showed that DSWorld trains agents much faster and predicts outcomes better than previous methods. The authors also made a large dataset and a new learning technique to improve prediction accuracy.
Authors
Zherui Yang, Fan Liu, Hao Liu
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 $14\times$ and search-based inference by approximately $3$-$6\times$ 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.