DSWorld: A Data Science World Model for Efficient Autonomous Agents

DSWorld: A Data Science World Model for Efficient Autonomous Agents

Paper: DSWorld: A Data Science World Model for Efficient Autonomous Agents | Authors: Zherui Yang, Fan Liu, Hao Liu | Institution: HKUST (Guangzhou) | Link: https://arxiv.org/abs/2607.15901 | Code: https://anonymous.4open.science/r/DSWorld

---

One-Sentence Summary

DSWorld introduces the concept of a Data Science World Model, predicting the effects of data science operations without real execution through structured state construction, cost-aware routing, and an LLM-based simulator, achieving approximately 14× RL training speedup and 3-6× inference acceleration while outperforming the strongest LLM baseline by 35.6% on transition prediction.

---

Background and Motivation

Autonomous data science agents have recently been proposed to automate a wide range of data science tasks, from exploratory data analysis to predictive modeling. ML-Master 2 achieves medal-level performance on 56.4% of Kaggle competition tasks on MLE-Bench. Existing methods typically leverage test-time scaling strategies, exploring numerous candidate solutions through iterative trial-and-error workflows. However, these strategies heavily rely on expensive analytical computation—including data processing, model training, evaluation, and workflow updating—causing the majority of execution time to be spent on computation rather than agent reasoning. ML-Master spends over 86% of its execution time on model training in MLE-Bench.

This computational overhead fundamentally limits the efficiency and scalability of autonomous data science systems. A critical question arises: Can we develop a transition prediction model for data science workflows, enabling agents to anticipate the effects of operations before performing costly computation?

To this end, the paper introduces the concept of a Data Science World Model. As illustrated in Figure 1, similar to vision world models that imagine future states of the physical world, a data science world model treats the data science execution environment as the "world" to be modeled. Given a workflow state and a candidate operation, the model predicts the next environment state—including dataset and model changes, execution feedback, errors, and performance signals. This capability enables agents to anticipate operation effects without expensive real-world execution, substantially accelerating both training and inference.

Figure 1: DSWorld simulates data science environments and accelerates agent training and inference.

---

Method

DSWorld consists of four components $\mathcal{W} = \{\mathcal{SC}, \mathcal{R}, \mathcal{C}, \mathcal{S}\}$, where $\mathcal{SC}$ is the State Constructor, $\mathcal{R}$ is the Router, $\mathcal{C}$ is the Compiler, and $\mathcal{S}$ is the LLM-based Simulator.

Data Science Workflow State Definition

The workflow state at time step $t$ is represented as $S_t = \{T_t, D_t, P_t, L_t\}$, where $T_t$ denotes the task, $D_t$ denotes the data state (e.g., dataset statistics and previews), $P_t$ denotes the execution environment (e.g., libraries and runtime configurations), and $L_t$ denotes execution logs, intermediate outputs, and task progress. The agent produces an action $A_t$ conditioned on state $S_t$, where $A_t$ denotes data science operations such as feature engineering, model training, and evaluation. The Data Science World Model is defined as a transition model that predicts operation effects:

$$S_{t+1} = \mathcal{W}(S_t, A_t)$$

where $\mathcal{W}$ denotes the world model. This formula captures the core idea: predicting the consequences of an action before execution.

State Constructor

The State Constructor transforms the raw execution environment into a structured state representation: $S_t = \mathcal{SC}(E_t)$, where $E_t$ is the data science environment at time step $t$. It is a rule-based program that extracts and organizes key information from the environment, including task descriptions, dataset statistics, data previews, execution environments, execution histories, intermediate outputs, and error messages. This structured representation enables DSWorld to model environment transitions in a unified manner.

Cost-Aware Router

Given the current state, the agent generates an action $A_t = \pi(S_t)$. To enable efficient routing, the action is first encoded into a dense embedding, which together with the current state is fed into the Router for decision making: $m_t = \mathcal{R}(S_t, A_t)$, where $m_t \in \{\texttt{execute}, \texttt{simulate}\}$. Intuitively, lightweight operations (e.g., simple data manipulation or environment inspection) are routed to direct execution, while computationally expensive operations (e.g., large-scale model training) are routed to simulation. To improve robustness against routing errors, a time limit is imposed on Compiler execution—if it exceeds the threshold, the action is redirected to the Simulator.

Compiler and Simulator

When the action is inexpensive, the Compiler executes it directly: $\hat{S}_{t+1} = \mathcal{C}(S_t, A_t)$, interacting with the actual execution environment and returning the resulting state. When the action is expensive, the Simulator predicts the next state without real execution: $\hat{S}_{t+1} = \mathcal{S}(S_t, A_t)$. The Simulator is an LLM-based transition model that predicts execution outcomes and potential errors directly from the current state and action. The overall transition process is formulated as:

$$\hat{S}_{t+1} = \begin{cases} \mathcal{C}(S_t, A_t), & m_t = \texttt{execute} \\ \mathcal{S}(S_t, A_t), & m_t = \texttt{simulate} \text{ or Timeout} \end{cases}$$

Through this hybrid execution-simulation mechanism, DSWorld balances efficiency and accuracy, enabling scalable environment interaction for autonomous data science agents.

Figure 2: DSWorld overview. (a) Predicts operation effects; (b) Reflective World Model Optimization; (c) LLM synthesizes and verifies state transitions.

Reflective World Model Optimization

Training adopts a two-stage post-training strategy: SFT warm-up followed by Reflective World Model Optimization. The SFT objective is:

$$\mathcal{L}_{\text{SFT}} = -\log \mathcal{S}_\theta(S' \mid S, A)$$

In the RL stage, the Simulator first predicts the next state $\hat{S}' \sim \mathcal{S}_\theta(\cdot \mid S, A)$, compares it with the ground-truth $S'$ to generate reflection feedback $f = \mathcal{S}_\theta(\hat{S}', S')$ identifying missing, incorrect, or inconsistent predictions. Conditioned on the feedback, the Simulator refines its prediction: $\hat{S}'_r = \mathcal{S}_\theta(S, A, f)$. For each sample, $n$ rollouts are performed to obtain both original and refined predictions $\mathcal{P} = \{\hat{S}'_i, \hat{S}'_{r,i}\}_{i=1}^n$, jointly optimized using GRPO:

$$\mathcal{L}(\theta) = \mathbb{E}\left[\frac{1}{2n}\sum_{i=1}^{n}\left(\mathcal{L}_{\text{clip}}(\hat{S}'_i, A_i) + \mathcal{L}_{\text{clip}}(\hat{S}'_{r,i}, A_{r,i})\right)\right] - \beta_{\text{KL}} \mathbb{D}_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})$$

where the advantage $\mathcal{A}$ is computed as a group-relative advantage: $\mathcal{A}_i = \frac{R_i - \mu(R)}{\sigma(R) + \epsilon}$, with $R_i$ being the reward of the $i$-th rollout and $\mu(R)$, $\sigma(R)$ the mean and standard deviation of rewards within the rollout group.

flowchart TB
    A["Raw Environment E_t"] --> B["State Constructor SC
extract structured state S_t"]
    B --> C["Agent generates action A_t"]
    C --> D["Action Encoder"]
    D --> E["Router R
decide execute or simulate"]
    E -->|lightweight op| F["Compiler C
real execution"]
    E -->|expensive op or timeout| G["Simulator S
LLM predicts next state"]
    F --> H["Next state S_t+1"]
    G --> H
    G --> I["Reflective optimization
compare with ground truth"]
    I --> J["Refined prediction S_r'"]
    J --> H

Training Data Construction

Due to the lack of state transition data for data science workflows, the authors construct DSWorld-8K, containing both real and synthesized transition trajectories. Real trajectories are collected by running existing agents on real tasks to record $(S, A, S')$, then using an LLM to synthesize CoT reasoning trajectories $\tau = p_\eta(S, A, S')$. Synthetic trajectories leverage MMTU's 60K real tables to construct diverse environment states, randomly sampling data operations $o$, error types $e$, and execution status $r$, with the LLM generating executable actions $A \sim p_\eta(A \mid S, o, e, r)$. These are then executed via the Compiler to obtain real next states $S' = \mathcal{C}(S, A)$, and verified $\text{Verify}(S, A, S', e, r)$ to retain only valid samples.

---

Experimental Results

DSWorld is evaluated on five transition prediction tasks: Execution Success Prediction (ESP), Error Type Prediction (ETP), Execution Result Similarity (ERS), Execution Keyword Matching (EKM), and Performance Prediction (PP), plus Performance Ranking (PR). DSWorld uses Qwen3-8B as the simulator backbone, Harrier OSS v1 0.6B as the encoder, and a two-layer MLP as the Router.

Method	ESP↑	ETP↑	ERS↑	EKM↑	PP↑	PR↑	Avg↑
Llama-3.1-8B	0.480	0.322	0.318	0.043	0.622	0.492	0.379
GPT-4o	0.712	0.502	0.472	0.173	0.757	0.492	0.518
o4-mini	0.680	0.585	0.489	0.382	0.789	0.514	0.576
Qwen3-8B-sft	0.917	0.885	0.843	0.574	0.849	0.509	0.763
Qwen3-8B-grpo	0.937	0.912	0.859	0.556	0.848	0.513	0.771
DSWorld	0.950	0.922	0.871	0.575	0.856	0.518	0.781

Table 1: Transition prediction performance. DSWorld achieves best results on nearly all dimensions.

DSWorld outperforms the strongest baseline o4-mini by 35.6% on average. On execution-related tasks, it improves by 33.4%, 57.6%, 71.5%, and 50.5% respectively over the strongest baseline, demonstrating more accurate modeling of execution dynamics. Performance-related tasks (PP, PR) require stronger reasoning about ML algorithms, where DSWorld achieves best PP and second-best PR.

Backbone	Simulator	Any Medal↑	Score↑	Time (min)↓
Qwen3-8B	Compiler	11.11	18.11	335
Qwen3-8B	DeepSeek 3.2	1.59	10.86	3854
Qwen3-8B	DSWorld	9.52	17.67	277

Table 2: Agent training comparison. DSWorld achieves ~14× training speedup while maintaining competitive performance.

For agent training, the DSWorld-trained agent achieves competitive performance with the Compiler-trained agent on MLE-Bench Lite, but with training time reduced from 335 to 277 minutes (~14× speedup). Using DeepSeek 3.2 as simulator also reduces time but causes severe performance degradation due to inaccurate feedback and hallucinated transitions. For inference acceleration, DSWorld achieves approximately 3-6× speedup over Compiler while preserving downstream performance.

Figure 1b: DSWorld accelerates agent RL training by approximately 14×.

The ablation study shows Qwen3-8B-sft improves average performance by 37.5% over the base backbone, validating the synthetic data pipeline. Qwen3-8B-grpo further improves by 1.05%, and DSWorld adds another 1.3% on top, demonstrating the effectiveness of Reflective World Model Optimization. Further analysis shows DSWorld consistently benefits from more training data (0.1k→6.4k) and larger model scale (0.6B→14B).

---

State representation

$$ S_{t}=\{T_{t},D_{t},P_{t},L_{t}\} $$

Limitations

1. Performance prediction tasks remain difficult: execution-related tasks (ESP, ETP, ERS) are easier to model due to explicit execution patterns; however, performance prediction and ranking require stronger reasoning about ML algorithms, task characteristics, and evaluation metrics, showing limited improvement. The authors note these tasks "require stronger reasoning about machine learning algorithms, task characteristics, and evaluation metrics."
2. Routing error risk: cost-aware routing depends on accurate judgment of operation computational cost. Misrouting expensive operations to the Compiler may cause timeouts, while misrouting inexpensive operations to the Simulator may introduce unnecessary prediction errors. Although a timeout fallback mechanism exists, routing precision still affects overall efficiency.
3. Limited synthetic data coverage: synthetic trajectories are based on NumPy/Pandas ecosystem operation and error libraries, which may not cover all operation patterns in real-world data science scenarios, limiting generalization.

---

Conclusion and Outlook

DSWorld extends the world model concept from the physical world to the data science execution environment, demonstrating that predicting operation effects before execution is feasible. The cost-aware routing mechanism elegantly balances precision and efficiency—lightweight operations are directly executed for accuracy, expensive operations are predicted by the LLM simulator for efficiency. Reflective World Model Optimization further improves prediction quality through error-aware iterative refinement. The 14× training speedup and 3-6× inference acceleration demonstrate that world models can serve as efficient environment simulators for autonomous data science agents.

This work opens a new direction for scalable training and inference of data science AI agents. As stronger LLM backbones and larger-scale transition data become available, data science world models are poised to become infrastructure components for accelerating data science automation.

Golden insight: Rather than letting agents waste 86% of their time waiting for computation through trial and error, let the world model tell them "what would happen if you did this" first—DSWorld proves that prediction itself is the most efficient execution.