Back to AI Research

AI Research

DSWorld: A Data Science World Model for Efficient A... | AI Research

Key Takeaways

  • DSWorld: A Data Science World Model for Efficient Autonomous Agents Autonomous data science agents are powerful tools for tasks like predictive modeling and...
  • 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.
  • 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.
Paper AbstractExpand

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 this https URL .

DSWorld: A Data Science World Model for Efficient Autonomous Agents
Autonomous data science agents are powerful tools for tasks like predictive modeling and data analysis, but they currently rely on a "trial-and-error" approach. This means they must constantly run code, train models, and evaluate results in real-time, which is computationally expensive and slow. This paper introduces the Data Science World Model, a framework designed to help agents "imagine" the outcome of their actions before they actually execute them. By predicting how a data science environment will change, the system avoids unnecessary, resource-heavy computations.

How DSWorld Works

The DSWorld framework acts as a bridge between an agent’s decision and the actual execution environment. It consists of four main components:

  • State Constructor: Converts raw environment data—such as dataset statistics, task descriptions, and execution logs—into a structured format that the model can easily process.

  • Router: A smart decision-maker that evaluates whether an agent’s proposed action is simple or complex. Lightweight tasks are sent to a compiler for direct execution, while computationally heavy tasks are sent to the simulator.

  • Simulator: An LLM-based model that predicts the next state of the environment, including potential errors and performance metrics, without needing to run the actual code.

  • Reflective Optimization: A training strategy that uses reinforcement learning to help the model identify its own prediction errors, reflect on them, and refine its future accuracy.

Building the Foundation

Because there was no existing large-scale dataset for data science transitions, the authors created "DSWorld-8K." This dataset contains 8,000 high-quality trajectories of data science workflows. The team combined real-world agent logs with a synthetic data pipeline that uses advanced LLMs to generate diverse scenarios, including both successful operations and common execution errors. Each sample includes "Chain-of-Thought" reasoning, which helps the model understand the logic behind why a specific action leads to a specific outcome.

Key Results

The researchers tested DSWorld in both training and inference scenarios, finding significant improvements in efficiency:

  • Speed: DSWorld accelerates the reinforcement learning training of agents by approximately 14 times and speeds up inference by 3 to 6 times.

  • Accuracy: On transition prediction tasks, DSWorld outperformed the strongest LLM baseline by 35.6%.

  • Performance: Despite the speed gains, the agents maintained competitive performance levels compared to traditional, slower methods.

Important Considerations

The effectiveness of DSWorld relies on its ability to accurately predict complex data science outcomes. While the framework significantly reduces the need for real-world execution, it is designed to work alongside a compiler that handles lightweight tasks, ensuring that the system remains grounded in reality. The authors note that this approach is particularly useful for scaling autonomous data science systems, as it shifts the bottleneck from raw computational power to more efficient, model-based reasoning.

Comments (0)

No comments yet

Be the first to share your thoughts!