DSWorld: A Data Science World Model for Efficient Autonomous Agents

2026-07-17

AI Summary

This work addresses the high cost and inefficiency of trial-and-error in existing autonomous data science agents, which lack the ability to predict the outcomes of their actions beforehand. To overcome this limitation, the authors propose a Data Science World Model that integrates structured state representations, cost-aware routing, lightweight real-world execution, and an LLM-based simulator into an end-to-end transition prediction framework. They further introduce Reflective World Model Optimization, a reinforcement learning strategy to refine agent behavior. The study contributes the first large-scale dataset of 8K state-transition trajectories and demonstrates substantial performance gains: the model improves state-transition prediction accuracy by 35.6% over the strongest LLM baseline, accelerates agent training by 14×, and achieves 3–6× speedup in search and reasoning.

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.

Problem

• autonomous data science agents
• trial-and-error workflows
• expensive computation
• world model
• transition prediction

Innovation

• Data Science World Model
• DSWorld
• State Transition Prediction
• Reflective World Model Optimization
• Cost-aware Simulation