DSWorld: A Data Science World Model for Efficient Autonomous Agents Zherui Yang1 , Fan Liu1 , Hao Liu1 * 1
The Hong Kong University of Science and Technology (Guangzhou) [email protected]; [email protected]; [email protected]
... State Action
Despite strong capabilities in data understanding and decision-making, autonomous data science agents still heavily rely on trial-anderror 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× and search-based inference by approximately 3-6× 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.
1
World Model
...
DS State
DSWorld Next State DS Operation Action
Next State
(a) Vision World Model imagines future states of the physical world, while Data Science World Model predicts the effects of data science operations without costly execution.
2k
14 ×
Speedup
1k 0 Compiler DSWorld
20
6k
15
4k
Time w/ Compiler Time w/ DSWorld
10 5
DS3.2 0
2k 0 ML-Master
Score w/ Compiler Score w/ DSWorld 60
3-6 × Speedup
40
Score
3k
25
Score Time (min)
Time Score
4k
Time (min)
arXiv:2607.15901v1 [cs.AI] 17 Jul 2026
Abstract
20
AIDE AutoMLGen0
(b) DSWorld accelerates agent (c) DSWorld accelerates agent RL training by ∼ 14×. inference by ∼ 3-6×.
Figure 1: DSWorld simulates data science environments and accelerates agent training and inference.
workflows (Jiang et al., 2025b; Du et al., 2025). However, these strategies heavily rely on expensive analytical computation, including data processing, model training, evaluation, and workflow updating. As a result, the majority of execution time is spent on computation rather than agent reasoning. For example, ML-Master (Liu et al., 2025a) spends over 86% of its execution time on model training in MLE-Bench. Such heavy computational overhead substantially limits the efficiency and scalability of autonomous data science systems. This raises a critical question: Can we develop a transition prediction model for data science workflows, enabling agents to anticipate the effects of data science operations before performing costly computation? To this end, we introduce the concept of Data Science World Model. As illustrated in Figure 1a, similar to vision world models that imagine future states of the physical world (Chu et al., 2026; Ding et al., 2026), data science world models treat the data science execution environment as the world to be modeled. Given a workflow state and a candidate data science operation, the model predicts the next environment state, including dataset
Introduction
Autonomous data science agents have recently been proposed to automate a wide range of data science tasks, ranging from exploratory data analysis to predictive modeling (Tang et al., 2025; Zhu et al., 2025). For example, ML-Master 2 (Zhu et al., 2026) achieves medal-level performance on 56.4% of Kaggle competition tasks on MLE-Bench (Chan et al., 2025). Existing methods typically leverage test-time scaling strategies, exploring numerous candidate solutions through iterative trial-and-error * Corresponding author.
1
and model changes, execution feedback, errors, and performance signals. Such a capability enables agents to anticipate the effects of data science operations without performing expensive real-world execution, thereby substantially accelerating both autonomous agent training and inference. To realize this vision, we propose DSWorld, a practical data science world model framework for autonomous data science workflows. Specifically, a State Constructor transforms raw data science environments into structured state representations containing tasks, datasets, execution histories, outputs, and environment status. To balance prediction accuracy and efficiency, we further introduce a cost-aware routing mechanism. Given an agent action, a Router determines whether the action requires heavy computation. Lightweight actions are directly executed through a Compiler to ensure accuracy, while computationally expensive actions are handled by an LLM-based Simulator, which predicts the next environment state without actual execution to improve efficiency. Building upon this framework, we first perform supervised finetuning (SFT) for initialization and further propose Reflective World Model Optimization, a reflective reinforcement learning strategy that improves nextstate prediction through error-aware reflection and iterative refinement. Due to the lack of large-scale transition data for autonomous data science workflows, we construct DSWorld-8K, a training dataset containing 8 thousand high quality data science agent trajectories. Specifically, we first collect authentic transition trajectories from real-world autonomous data science tasks and further synthesize corresponding Chainof-Thought (CoT) trajectories (Wei et al., 2022) to explain transition logic. However, real-world trajectories are expensive and inherently limited in scale. Therefore, we further develop a scalable data synthesis pipeline that leverages advanced LLMs and large-scale data sources from MMTU (Xing et al., 2025) to synthesize diverse workflow states and data science operations. The generated actions are executed in real environments to obtain next states, and only verified samples are retained for training. Additional CoT trajectories are synthesized to further explain transition dynamics. We evaluate DSWorld in two practical downstream settings: (1) as a transition model for training autonomous data science agents, and (2) as a transition model for agent inference. Experimental results demonstrate that DSWorld substantially ac-
celerates both agent training and inference while maintaining strong agent performance, achieving approximately 14× RL training speedup and 36× inference acceleration. In addition, we evaluate DSWorld across multiple transition prediction tasks. Experimental results show that DSWorld achieves strong predictive performance and outperforms the strongest LLM baseline by 35.6% on average on transition prediction tasks. In summary, our contributions are as follows: • We introduce the concept of Data Science World Models, which aim to predict the effects of data science operations in autonomous data science workflows without performing costly real-world computation. • We propose DSWorld, a practical data science world model framework for modeling state transitions, and further introduce a reflective reinforcement learning strategy that improves transition prediction through erroraware reflection and iterative refinement. • We develop a scalable data synthesis pipeline for constructing large-scale transition data for data science world model training. • Extensive experiments demonstrate the effectiveness of DSWorld, achieving approximately 14× RL training speedup and 3-6× inference acceleration while maintaining strong downstream agent performance.
2
Related works
2.1
World Models
World models aim to model environment dynamics for prediction, control, and simulation, typically by predicting future states conditioned on current states and actions (Chu et al., 2026; Ding et al., 2026; Ha and Schmidhuber, 2018). Existing world model research spans multiple domains. In the Physical World, world models are used for video generation (Bruce et al., 2024; Chen et al., 2026), 3D world generation (Kong et al., 2025; World Labs, 2025), and latent-space dynamics modeling (Ha and Schmidhuber, 2018; Assran et al., 2023). In the Digital World, recent works simulate webpage or software state transitions for web and GUI agents (Xiao et al., 2026; Chae et al., 2025). Despite these advances, world modeling for autonomous data science environments remains 2
(b) Reflective World Model Optimization
...
Action
DS Environment
Router
State Constructor State
State and Action
Timeout
Simulator
Next Turn
Policy Model
Policy Model
... Sn ′
Preditions
′ Sn
...
′ Sr,n
Refined Preditions
′ Sr,1
...
Reference Model
R1
Reward Model
Sr,n ′
... R2n
Advantage Computation
Rewards
A1
... A2n
Advantages
Verify
... Next State
...
′ Sr,1
Compiler
Agent
S1 ′
′ S1
...
State
Sample LLM
Data and Conditions
Compiler
Next State
Code Action (c) Training Data Construction
(a) The DSWorld Framework
LLM
CoT
Figure 2: Overview of DSWorld. (a) DSWorld predicts the effects of data science operations before performing costly computation. (b) Reflective World Model Optimization refines next-state prediction through error-aware reflection. (c) LLMs synthesize and verify data science state transitions and generate CoT trajectories.
configurations), and Lt denotes execution logs, intermediate outputs, and task progress.
largely unexplored. Different from prior work, we study data science world models that model environment transitions in data science workflows. 2.2
Definition 2 (Action) At time step t, a data science agent produces an action At conditioned on the current state St , where At denotes data science operations such as feature engineering, model training, and evaluation.
Autonomous Data Science Agents
Recently, autonomous data science agents have been proposed to automate a wide range of data science tasks (Tang et al., 2025; Zhu et al., 2025). Existing methods mainly improve agent performance in two ways. One line of work enhances the backbone model through training. For example, ML-Agent (Liu et al., 2025b) trains agents with reinforcement learning for machine learning tasks. Another line of work improves performance through test-time scaling strategies. For example, AIDE (Jiang et al., 2025b), ML-Master (Liu et al., 2025a), AutoMLGen (Du et al., 2025), and AutoMind (Ou et al., 2025) use search algorithms to explore candidate solutions and select the optimal one. Despite their effectiveness, these approaches heavily rely on time-consuming computation during both training and inference. Therefore, developing data science world models that can model environment transitions and enable agents to anticipate the effects of data science operations before costly computation is critical.
3
Definition 3 (Data Science World Model) A Data Science World Model is defined as a transition model that predicts the effects of data science operations before costly computation: St+1 = W(St , At ),
(1)
where W denotes the data science world model.
4
Methodology
Figure 2 illustrates the overall framework of DSWorld. Given a data science workflow state, a State Constructor converts it into a structured representation. Based on the current state, a data science agent generates an action. The action is encoded into a dense embedding and routed to either a Compiler for real execution or an LLM-based Simulator for next-state prediction. 4.1
Preliminaries
DSWorld
The proposed data science world model consists of four components W = {SC, R, C, S}, where SC denotes the State Constructor, R denotes the Router, C denotes the Compiler, and S denotes the LLM-based Simulator. State Constructor. We first introduce the State Constructor, which transforms the raw execution
Definition 1 (Data Science Workflow State) The data science workflow state at time step t is represented as St = {Tt , Dt , Pt , Lt }, where Tt denotes the task, Dt denotes the data state (e.g., dataset statistics and previews), Pt denotes the execution environment (e.g., libraries and runtime 3
environment into a structured state representation, formulated as St = SC(Et ), where Et denotes the data science environment at time step t. Specifically, the State Constructor 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. The resulting structured state representation enables DSWorld to model environment transitions in a unified manner. Router. Given the current state, a data science agent generates an action as At = π(St ), where π denotes the agent policy. To enable efficient routing, we first encode the generated action into a dense embedding representation using an action encoder. The resulting action embedding, together with the current state, is then fed into the Router for decision making. The Router determines whether the generated action can be executed efficiently as mt = R(St , At ), where mt ∈ {execute, simulate}. Intuitively, lightweight operations such as simple data manipulation or environment inspection are routed to direct execution, while computationally expensive operations, such as large-scale model training, are routed to simulation. Compiler. If the action is inexpensive, the Compiler executes it directly: Ŝt+1 = C(St , At ).
Through this hybrid execution-simulation mechanism, DSWorld balances efficiency and accuracy, enabling scalable environment interaction for autonomous data science agents. 4.2
To effectively train DSWorld and improve transition prediction quality, we adopt a two-stage posttraining strategy consisting of SFT followed by Reflective World Model Optimization, a reflective reinforcement learning strategy that explicitly analyzes prediction errors and refines subsequent predictions to provide higher-quality training signals. 4.2.1
SFT Warm-Up
We first perform supervised fine-tuning on data science transition trajectories to initialize the Simulator with basic transition modeling capabilities. The SFT objective is defined as: ′
LSFT = − log Sθ (S | S, A), ′
4.2.2
Reflective World Model Optimization
After initializing DSWorld through SFT, we further optimize it with RL. Given the current state and action, the Simulator first predicts the next state as ′ Ŝ ∼ Sθ (· | S, A). We then compare the prediction with the ground-truth next state and generate reflection feedback:
(2)
′
C(St , At ), S(St , At ),
mt = execute, mt = simulate or Timeout.
(6)
where f identifies missing, incorrect, or inconsistent predictions. Conditioned on the reflection feedback, the Simulator refines its prediction as
(3)
′
Ŝr = Sθ (S, A, f ).
(7)
For each training sample, we perform n rollouts to obtain both original and refined predictions as ′ ′ P = {Ŝi , Ŝr,i }ni=1 . We jointly optimize all trajectories using Group Relative Policy Optimization (GRPO) (Shao et al., 2024). The overall objective is defined as: "
Ŝt+1 =
′
f = Sθ (Ŝ , S ),
The Simulator is an LLM-based transition model that predicts execution outcomes and potential errors directly from the current state and action, thereby avoiding expensive computation. To improve robustness against routing errors, we further impose a time limit on Compiler execution. If execution exceeds the predefined timeout threshold, the action is redirected to the Simulator for transition prediction. Thus, the overall transition process can be formulated as: (
(5)
where S, A, and S denote the current state, action, and next state, respectively.
The Compiler interacts with the actual execution environment and returns the resulting next state, including updated data states, execution outputs, and runtime feedback. Simulator. Otherwise, the Simulator predicts the next state without real execution: Ŝt+1 = S(St , At ).
World Model Optimization
# n ′ ′ 1 X L(θ) =E Lclip (Ŝi , Ai ) + Lclip (Ŝr,i , Ar,i ) 2n i=1 (8)
(4)
− βKL DKL (πθ ∥ πref ),
4
State Synthesis. To construct diverse environment states, we leverage datasets from MMTU (Xing et al., 2025), which contains over 60K real tables across diverse domains, providing rich and heterogeneous data science environments. For each sample, we randomly select a dataset and construct the corresponding state using the State Constructor as S = SC(E), where E denotes the sampled data science environment. The constructed state includes task descriptions, dataset statistics, and data previews. Action Synthesis. To generate diverse actions, we construct data operation and execution error libraries based on the NumPy and Pandas ecosystems. Each entry contains an operation or error type, textual descriptions, and code examples. For each synthesized sample, we randomly sample a data operation o, an error type e, and an execution status r ∈ {success, failure}. Conditioned on the current state and sampled attributes, the LLM synthesizes executable actions:
where ′ " |Ŝ | πθ (yt |x, y<t ) 1 X min A, Lclip (Ŝ , A) = ′ π θold (yt |x, y<t ) |Ŝ | t=1 # πθ (yt |x, y<t ) clip , 1 − ε, 1 + ε A . πθold (yt |x, y<t ) (9) ′
Here, ε and βKL are hyperparameters, πref denotes a reference policy, and DKL denotes the Kullback–Leibler divergence. The advantage A is computed as a group-relative advantage: Ai =
Ri − µ(R) , σ(R) + ϵ
(10)
where Ri denotes the reward of the i-th rollout, µ(R) and σ(R) denote the mean and standard deviation of rewards within the rollout group, and ϵ is a small constant for numerical stability. The reward function evaluates whether the Simulator correctly predicts execution status, outputs, error information, and task performance. 4.3
A ∼ pη (A | S, o, e, r),
Training Data Construction
(12)
Due to the lack of state transition data for data science workflows, we construct a high-quality training dataset, DSWorld-8K, consisting of both real and synthesized state transitions. To support transition reasoning learning, we additionally synthesize reasoning explanations for each transition. Each ′ sample is represented as (S, A, S , τ ), where S is ′ the current state, A is the action, S is the next state, and τ is the reasoning trajectory explaining the transition process.
where A denotes the generated action. This strategy enables the training data to cover both successful executions and diverse execution failures in data science workflows. Ground-Truth Construction and Verification. After generating executable actions, we execute the synthesized action using the Compiler to obtain the corresponding next state:
4.3.1 Real-World Transition Collection We first leverage existing autonomous data science agents and real-world data science benchmarks to collect authentic environment transition trajectories. Specifically, we run agents on real-world tasks and record the resulting environment transitions as ′ (S, A, S ). Based on the collected transitions, we further prompt an advanced LLM to synthesize corresponding CoT reasoning trajectories:
The resulting transition captures actual execution outcomes, including updated data states, execution outputs, runtime logs, and error messages. To ensure data quality, we further verify whether the execution results satisfy the intended execution status and error constraints:
′
τ = pη (S, A, S ),
′
S = C(S, A).
′
Verify(S, A, S , e, r).
(13)
(14)
Only valid and consistent samples are retained. ′ For each verified transition tuple (S, A, S ), we further prompt the LLM to synthesize the corresponding CoT reasoning trajectory:
(11)
where pη denotes the LLM.
′
τ = pη (S, A, S )
4.3.2 Synthetic Transition Construction However, collecting real-world transition data is expensive and inherently limited in scale. To enable scalable training, we further design a synthetic pipeline for generating transition trajectories.
(15)
Combining both real-world and synthesized trajectories, we construct a final training dataset containing approximately 8K transition samples with corresponding reasoning trajectories. 5
Table 1: Performance comparison on data science transition prediction tasks. Execution Prediction
Methods
Performance Prediction
AVG. ↑
ESP ↑
ETP ↑
ERS ↑
EKM ↑
PP ↑
PR ↑
Training-free
Llama-3.1-8B Qwen3-8B DeepSeek-3.2 GPT-4o o4-mini
0.480±0.112 0.710±0.009 0.628±0.023 0.712±0.018 0.680±0.010
0.322±0.047 0.573±0.023 0.403±0.028 0.502±0.033 0.585±0.023
0.318±0.061 0.508±0.010 0.420±0.022 0.472±0.020 0.489±0.015
0.043±0.021 0.193±0.020 0.250±0.010 0.173±0.009 0.382±0.031
0.622±0.02 0.840±0.012 0.851±0.005 0.757±0.010 0.789±0.005
0.492±0.021 0.507±0.009 0.539±0.011 0.492±0.013 0.514±0.028
0.379 0.555 0.516 0.518 0.576
Trained on DSWorld-8K
Llama-3.1-8B-sft Llama-3.1-8B-grpo Qwen3-8B-sft Qwen3-8B-grpo DSWorld (Ours)
0.872±0.028 0.910±0.017 0.917±0.010 0.937±0.013 0.950±0.005
0.852±0.028 0.892±0.01 0.885±0.015 0.912±0.008 0.922±0.003
0.798±0.029 0.829±0.014 0.843±0.008 0.859±0.010 0.871±0.005
0.529±0.017 0.533±0.011 0.574±0.020 0.556±0.005 0.575±0.018
0.809±0.011 0.819±0.011 0.849±0.017 0.848±0.007 0.856±0.001
0.498±0.012 0.502±0.016 0.509±0.011 0.513±0.007 0.518±0.008
0.726 0.747 0.763 0.771 0.781
5
Experiments
8B-grpo, Qwen3-8B-sft, and Qwen3-8B-grpo. Implementation Details. DSWorld employs Qwen3-8B as the simulator backbone. The encoder is implemented using Harrier OSS v1 0.6B (Microsoft, 2026), while the Router is implemented as a two-layer MLP. Detailed implementation settings are provided in the Appendix A.3.
This section aims to answer the following research questions: RQ1: How effective is DSWorld on data science transition prediction tasks? RQ2: Can DSWorld effectively support the training and inference of agents? RQ3: Is the proposed optimization strategy effective for training DSWorld? RQ4: How do training data scale and model scale affect DSWorld performance? 5.1
5.2
Transition Prediction Performance (RQ1)
Table 1 presents the overall performance comparison between DSWorld and strong LLM baselines across multiple prediction tasks. Overall, DSWorld consistently achieves the best performance on nearly all evaluation dimensions, improving the average performance by 35.6% over the strongest baseline, o4-mini. These results demonstrate the effectiveness of DSWorld for modeling data science environment transitions. Specifically, compared with general-purpose LLMs, DSWorld substantially improves prediction accuracy on execution-related tasks. Compared with the strongest baseline, DSWorld achieves improvements of 33.4%, 57.6%, 71.5%, and 50.5% on these tasks, respectively. These results indicate that DSWorld can more accurately model execution dynamics and environment transitions in data science workflows. In addition, both SFT- and GRPOtrained models significantly outperform their untuned backbones, validating the effectiveness of the proposed synthetic transition data for training. For performance-related tasks, DSWorld achieves competitive results with the bestperforming methods. DSWorld achieves the best performance prediction score and the second-best performance ranking result. We attribute this to the fact that these tasks require stronger reasoning about machine learning algorithms, task characteristics, and evaluation metrics, making them more
Experimental Setup
Benchmarks. We evaluate DSWorld on the Predict-before-Execute (Zheng et al., 2026) benchmark to assess its performance ranking (PR) capability. In addition, since no benchmark exists for data science world models, we construct 540 evaluation tasks to measure execution success prediction (ESP), error type prediction (ETP), execution result similarity (ERS), execution keyword matching (EKM), and performance prediction (PP). Furthermore, we evaluate DSWorld as an environment simulator for autonomous data science agents on MLEBench Lite (Chan et al., 2025). Detailed benchmark descriptions are provided in Appendix A.1. Metrics. For PR, ESP, ETP, and EKM, we use accuracy as the evaluation metric. For ERS, we use embedding cosine similarity, while for PP, we use 1 − RMSE. For MLE-Bench Lite, we report the medal rate and overall score. Detailed metric descriptions are provided in Appendix A.2. Baselines. Since no existing methods are specifically designed for data science world modeling, we compare DSWorld with two categories of baselines: (1) advanced LLMs, including Llama-3.18B (Team, 2024), Qwen3-8B (Yang et al., 2025), DeepSeek 3.2 (DeepSeek-AI, 2025), GPT-4o (OpenAI, 2023), and o4-mini; and (2) trained LLM baselines, including Llama-3.1-8B-sft, Llama-3.16
Table 2: Performance and training time comparison of agents trained with different environment simulators. Simulators
Gold ↑
Silver ↑
Bronze ↑
Any ↑
Median ↑
Score ↑
Time (min) ↓
Qwen3-8B Qwen3-14B Qwen3-8B Qwen3-8B Qwen3-8B
DeepSeek 3.2 Compiler DSWorld
1.59±2.75 1.59±2.75 1.59±2.75 1.59±2.75 1.59±2.75
1.59±2.75 0.00±0.00 0.00±0.00 1.59±2.75 1.59±2.75
1.59±2.75 6.35±2.75 0.00±0.00 6.35±2.75 6.35±2.75
4.76±4.76 7.94±2.75 1.59±2.75 11.11±2.75 9.52±0.00
7.94±2.75 7.94±2.75 6.35±2.75 12.70±2.75 11.11±2.75
13.80±0.53 16.50±2.75 10.86±1.07 18.11±1.93 17.67±2.27
3854 335 277
with DeepSeek 3.2
with DSWorld
Cumulative Time (min)
Backbones
with Compiler
4000
DSWorld. These results further demonstrate the efficiency and scalability advantages of DSWorld.
3000
5.4 14 × Speedup
2000
We further investigate whether DSWorld can serve as an efficient inference environment for searchbased autonomous data science agents. Specifically, we use the Compiler, DeepSeek 3.2, and DSWorld as the executers of search-based agents, including AIDE, ML-Master, and AutoMLGen. Table 3 presents the performance and efficiency results on MLE-Bench Lite. Overall, DSWorld consistently achieves competitive downstream performance while substantially reducing inference time across different agents and backbone models. Specifically, compared with Compiler-based execution, DSWorld achieves approximately 3-6× acceleration while largely preserving downstream performance. In contrast, although directly using DeepSeek 3.2 also reduces execution time, it leads to severe performance degradation due to inaccurate execution feedback and hallucinated environment transitions. These results indicate that DSWorld can effectively simulate real environment transitions and provide sufficiently accurate execution feedback. Since search-based agents require evaluating many candidate solutions, replacing expensive real execution with efficient world model simulation can significantly reduce inference overhead while preserving search quality.
1000 0
1
2
3
4
5
6
Step
7
8
9 10
Figure 3: RL training time across different simulators.
difficult to learn. In contrast, execution-related prediction tasks exhibit more explicit execution patterns and environment dynamics, which are easier for DSWorld to model effectively. 5.3
Agent Inference with DSWorld (RQ2)
Training Agents with DSWorld (RQ2)
To investigate whether DSWorld can effectively support agent training, we use 105 machine learning tasks from MLE-Dojo (Qiang et al., 2025) to train agents under the same ReAct (Yao et al., 2023) framework for 10 RL steps and evaluate them on MLE-Bench Lite. As shown in Table 2, the agent trained with DSWorld achieves strong performance and remains competitive with Compiler-based training, while outperforming training with DeepSeek 3.2 as the simulator. Notably, the DSWorld-trained agent also outperforms the stronger Qwen3-14B baseline despite using Qwen3-8B as the backbone. These results demonstrate that DSWorld can serve as an effective training environment for autonomous data science agents. More importantly, Figure 3 also compares the training efficiency of different simulators. Compiler-based training with real execution is significantly more time-consuming, whereas DSWorld achieves approximately 14× acceleration while maintaining comparable downstream performance. Although Compiler-based training achieves slightly better final performance, the gap is relatively small compared with the substantial efficiency gain of
5.5
Ablation Study (RQ3)
To investigate the effectiveness of the proposed training strategy, we compare different variants of DSWorld. The results are shown in Table 1, where Qwen3-8B-sft and Qwen3-8B-grpo denote models trained with only SFT and with additional GRPObased RL, respectively. Compared with the original Qwen3-8B backbone, Qwen3-8B-sft achieves substantial improvements across all tasks, improving the average performance by 37.5%. These results demonstrate the effectiveness of the proposed data synthesis pipeline and supervised world model 7
Table 3: Performance and inference efficiency comparison using different executers on MLE-Bench Lite. Methods AIDE
Qwen3-8B
ML-Master
AutoMLGen
AIDE
DeepSeek-3.2
ML-Master
AutoMLGen
Executer
Gold ↑
Silver ↑
Bronze ↑
Any ↑
Median ↑
Score ↑
Time (s) ↓
Compiler DeepSeek 3.2 DSWorld Compiler DeepSeek 3.2 DSWorld Compiler DeepSeek 3.2 DSWorld
3.17±2.75 1.59±2.75 3.17±2.75 1.59±2.75 1.59±2.75 1.59±2.75 0.00±0.00 0.00±0.00 0.00±0.00
0.00±0.00 0.00±0.00 0.00±0.00 3.17±2.75 0.00±0.00 4.76±4.76 0.00±0.00 0.00±0.00 0.00±0.00
0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00 0.00±0.00
3.17±2.75 1.59±2.75 3.17±2.75 4.76±4.76 1.59±2.75 6.35±5.5 0.00±0.00 0.00±0.00 0.00±0.00
6.35±2.75 3.17±2.75 7.94±2.75 9.52±0.00 4.76±4.76 9.52±0.00 0.00±0.00 0.00±0.00 0.00±0.00
10.7±2.27 7.21±3.46 10.58±2.32 12.39±3.51 7.00±4.08 10.34±1.56 0.33±0.56 0.00±0.00 0.40±0.69
4102 806 676 1421 653 371 1470 476 312
Compiler DeepSeek 3.2 DSWorld Compiler DeepSeek 3.2 DSWorld Compiler DeepSeek 3.2 DSWorld
4.76±4.76 3.17±2.75 6.35±2.75 15.88±2.75 0.00±0.00 11.11±5.50 3.17±2.75 0.00±0.00 3.17±2.75
3.17±2.75 3.17±5.50 3.17±5.50 6.35±2.75 1.59±2.75 7.94±2.75 0.00±0.00 0.00±0.00 0.00±0.00
1.59±2.75 4.76±4.76 3.17±5.50 3.17±5.5 7.94±2.75 3.17±2.75 3.17±5.5 3.17±2.75 3.17±5.50
9.52±4.76 11.11±5.50 12.70±9.91 25.40±5.5 9.52±0.00 22.23±7.28 6.35±7.27 3.17±2.75 6.35±7.27
14.29±4.76 14.29±9.52 17.46±7.27 30.16±2.75 19.05±4.76 26.99±2.75 15.87±7.27 9.52±2.75 15.87±7.27
17.41±4.13 15.12±5.04 20.49±3.27 31.61±1.59 22.32±3.57 29.26±1.74 17.29±4.36 12.46±2.63 17.62±2.82
4344 1362 1199 3232 936 1065 5836 1273 989
training. In particular, Qwen3-8B-sft significantly improves execution-related prediction tasks, indicating that supervised fine-tuning on large-scale transition trajectories enables the model to effectively learn data science environment dynamics. Building upon Qwen3-8B-sft, Qwen3-8B-grpo further improves the overall performance by 1.05%, demonstrating the effectiveness of RL for improving prediction quality through iterative refinement. Furthermore, DSWorld achieves the best overall performance across nearly all evaluation dimensions. Compared with Qwen3-8B-grpo, DSWorld further improves the overall performance by 1.3%, and improves by 2.36% over Qwen3-8Bsft, demonstrating the effectiveness of Reflective World Model Optimization. By introducing reflective error-aware optimization, DSWorld can iteratively refine transition reasoning and improve nextstate prediction quality, enabling more accurate simulation of data science environment transitions. 5.6
Llama3.1-8B
Qwen3-8B
Average Score
0.8
0.8
Average Score
Backbones
0.7
0.7
0.6
0.6 0.5 0.1k 0.2k 0.4k 0.8k 1.6k 3.2k 6.4k
Training Data Size
0.5 0.6B
1.7B
4B
8B
Model Scale
14B
(a) Performance under differ- (b) Performance under different training data scales. ent Qwen3 model scales.
Figure 4: Scale effects on DSWorld.
accurate environment transition modeling and reasoning capabilities. Overall, these results indicate that DSWorld scales favorably with both data and model size, highlighting its potential for further improvement with larger-scale training resources.
6
Conclusion
In this paper, we present DSWorld, a learned transition model for data science workflows. By modeling environment transitions conditioned on the current state and actions, DSWorld enables agents to anticipate the effects of data science operations before performing costly computation during both training and inference. To support effective world model learning, we construct a large-scale trajectory dataset combining real and synthetic transition data, and further propose Reflective World Model Optimization to improve prediction quality through reflective reinforcement learning. Extensive experiments demonstrate that DSWorld achieves strong predictive performance across diverse transition prediction tasks while substantially accelerating autonomous agent training and inference.
Further Analysis (RQ4)
In this section, we further investigate the effects of training data size and model scale on DSWorld. Figure 4 shows that DSWorld consistently benefits from more training data and larger backbone models. Specifically, as the training data size increases from 0.1k to 6.4k samples, the performance of both Llama3.1-8B and Qwen3-8B steadily improves, demonstrating the importance of large-scale transition data for world model learning. Moreover, scaling the Qwen3 backbone from 0.6B to 14B parameters also leads to significant performance gains, suggesting that stronger LLMs provide more 8
Limitations
DeepSeek-AI. 2025. Deepseek-v3.2: Pushing the frontier of open large language models. CoRR, abs/2512.02556.
Despite the promising results, this work still has several limitations. First, DSWorld currently focuses on modeling data science transitions and does not explicitly model external tool-call transitions in autonomous agent workflows. Second, the quality of transition prediction remains constrained by the capability of the underlying LLM simulator, which may occasionally produce inaccurate predictions in complex workflow scenarios. Third, synthesized trajectories may still exhibit distribution gaps compared with real-world autonomous workflows, potentially affecting generalization performance. We leave these limitations for future work.
Jingtao Ding, Yunke Zhang, Yu Shang, Yuheng Zhang, Zefang Zong, Jie Feng, Yuan Yuan, Hongyuan Su, Nian Li, Nicholas Sukiennik, Fengli Xu, and Yong Li. 2026. Understanding world or predicting future? A comprehensive survey of world models. ACM Comput. Surv., 58(3):57:1–57:38. Shangheng Du, Xiangchao Yan, Dengyang Jiang, Jiakang Yuan, Yusong Hu, Xin Li, Liang He, Bo Zhang, and Lei Bai. 2025. Automlgen: Navigating finegrained optimization for coding agents. CoRR, abs/2510.08511. David Ha and Jürgen Schmidhuber. 2018. Recurrent world models facilitate policy evolution. In NeurIPS, pages 2455–2467.
References
Xueyu Hu, Ziyu Zhao, Shuang Wei, Ziwei Chai, Qianli Ma, Guoyin Wang, Xuwu Wang, Jing Su, Jingjing Xu, Ming Zhu, Yao Cheng, Jianbo Yuan, Jiwei Li, Kun Kuang, Yang Yang, Hongxia Yang, and Fei Wu. 2024. Infiagent-dabench: Evaluating agents on data analysis tasks. In ICML, pages 19544–19572.
Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael G. Rabbat, Yann LeCun, and Nicolas Ballas. 2023. Selfsupervised learning from images with a jointembedding predictive architecture. In CVPR, pages 15619–15629.
Yiming Huang, Jianwen Luo, Yan Yu, Yitong Zhang, Fangyu Lei, Yifan Wei, Shizhu He, Lifu Huang, Xiao Liu, Jun Zhao, and Kang Liu. 2024. Da-code: Agent data science code generation benchmark for large language models. In EMNLP, pages 13487–13521.
Jake Bruce, Michael D. Dennis, Ashley Edwards, Jack Parker-Holder, Yuge Shi, Edward Hughes, Matthew Lai, Aditi Mavalankar, Richie Steigerwald, Chris Apps, Yusuf Aytar, Sarah Bechtle, Feryal M. P. Behbahani, Stephanie C. Y. Chan, Nicolas Heess, Lucy Gonzalez, Simon Osindero, Sherjil Ozair, Scott E. Reed, and 6 others. 2024. Genie: Generative interactive environments. In ICML, pages 4603–4623.
Dongfu Jiang, Yi Lu, Zhuofeng Li, Zhiheng Lyu, Ping Nie, Haozhe Wang, Alex Su, Hui Chen, Kai Zou, Chao Du, Tianyu Pang, and Wenhu Chen. 2025a. Verltool: Towards holistic agentic reinforcement learning with tool use. CoRR, abs/2509.01055.
Hyungjoo Chae, Namyoung Kim, Kai Ong, Minju Gwak, Gwanwoo Song, Jihoon Kim, Sunghwan Kim, Dongha Lee, and Jinyoung Yeo. 2025. Web agents with world models: Learning and leveraging environment dynamics in web navigation. In ICLR, pages 63707–63738.
Zhengyao Jiang, Dominik Schmidt, Dhruv Srikanth, Dixing Xu, Ian Kaplan, Deniss Jacenko, and Yuxiang Wu. 2025b. AIDE: ai-driven exploration in the space of code. CoRR, abs/2502.13138. Lingdong Kong, Wesley Yang, Jianbiao Mei, Youquan Liu, Ao Liang, Dekai Zhu, Dongyue Lu, Wei Yin, Xiaotao Hu, Mingkai Jia, Junyuan Deng, Kaiwen Zhang, Yang Wu, Tianyi Yan, Shenyuan Gao, Song Wang, Linfeng Li, Liang Pan, Yong Liu, and 4 others. 2025. 3d and 4d world modeling: A survey. CoRR, abs/2509.07996.
Jun Shern Chan, Neil Chowdhury, Oliver Jaffe, James Aung, Dane Sherburn, Evan Mays, Giulio Starace, Kevin Liu, Leon Maksin, Tejal Patwardhan, Aleksander Madry, and Lilian Weng. 2025. Mle-bench: Evaluating machine learning agents on machine learning engineering. In ICLR. Taiye Chen, Xun Hu, Zihan Ding, and Chi Jin. 2026. Learning world models for interactive video generation. In NeurIPS, pages 154456–154483.
Zexi Liu, Yuzhu Cai, Xinyu Zhu, Yujie Zheng, Runkun Chen, Ying Wen, Yanfeng Wang, Weinan E, and Siheng Chen. 2025a. Ml-master: Towards ai-for-ai via integration of exploration and reasoning. CoRR, abs/2506.16499.
Meng Chu, Xuan Billy Zhang, Kevin Qinghong Lin, Lingdong Kong, Jize Zhang, Teng Tu, Weijian Ma, Ziqi Huang, Senqiao Yang, Wei Huang, Yeying Jin, Zhefan Rao, Jinhui Ye, Xinyu Lin, Xichen Zhang, Qisheng Hu, Shuai Yang, Leyang Shen, Wei Chow, and 23 others. 2026. Agentic world modeling: Foundations, capabilities, laws, and beyond. CoRR, abs/2604.22748.
Zexi Liu, Jingyi Chai, Xinyu Zhu, Shuo Tang, Rui Ye, Bolun Zhang, Lei Bai, and Siheng Chen. 2025b. Mlagent: Reinforcing LLM agents for autonomous machine learning engineering. CoRR, abs/2505.23723. Microsoft. 2026. harrier-oss-v1.
9
OpenAI. 2023. Hello gpt-4.
Jingsheng Zheng, Jintian Zhang, Yujie Luo, Yuren Mao, Yunjun Gao, Lun Du, Huajun Chen, and Ningyu Zhang. 2026. Can we predict before executing machine learning agents? CoRR, abs/2601.05930.
Yixin Ou, Yujie Luo, Jingsheng Zheng, Lanning Wei, Shuofei Qiao, Jintian Zhang, Da Zheng, Huajun Chen, and Ningyu Zhang. 2025. Automind: Adaptive knowledgeable agent for automated data science. CoRR, abs/2506.10974.
Xinyu Zhu, Yuzhu Cai, Zexi Liu, Bingyang Zheng, Cheng Wang, Rui Ye, Jiaao Chen, Hanrui Wang, WeiChen Wang, Yuzhi Zhang, Linfeng Zhang, Weinan E, Di Jin, Siheng Chen, and Yanfeng Wang. 2026. Toward ultra-long-horizon agentic science: Cognitive accumulation for machine learning engineering. CoRR, abs/2601.10402.
Rushi Qiang, Yuchen Zhuang, Yinghao Li, Dingu Sagar V. K, Rongzhi Zhang, Changhao Li, Ian ShuHei Wong, Sherry Yang, Percy Liang, Chao Zhang, and Bo Dai. 2025. Mle-dojo: Interactive environments for empowering LLM agents in machine learning engineering. CoRR, abs/2505.07782.
Yizhang Zhu, Liangwei Wang, Chenyu Yang, Xiaotian Lin, Boyan Li, Wei Zhou, Xinyu Liu, Zhangyang Peng, Tianqi Luo, Yu Li, Chengliang Chai, Chong Chen, Shimin Di, Ju Fan, Ji Sun, Nan Tang, Fugee Tsung, Jiannan Wang, Chenglin Wu, and 6 others. 2025. A survey of data agents: Emerging paradigm or overstated hype? CoRR, abs/2510.23587.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Mingchuan Zhang, Y. K. Li, Y. Wu, and Daya Guo. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. CoRR, abs/2402.03300. Zirui Tang, Weizheng Wang, Zihang Zhou, Yang Jiao, Bangrui Xu, Boyu Niu, Xuanhe Zhou, Guoliang Li, Yeye He, Wei Zhou, Yitong Song, Cheng Tan, Bin Wang, Conghui He, Xiaoyang Wang, and Fan Wu. 2025. Llm/agent-as-data-analyst: A survey. CoRR, abs/2509.23988.
A
Experimental Setup
A.1
Benchmarks
• Predict-before-Execute (Zheng et al., 2026). Predict-before-Execute is a performance ranking benchmark containing 28 machine learning problems and 18,438 ranking tasks in total. Each task consists of a machine learning problem together with several candidate code solutions, and the model is required to predict which solution achieves better performance without executing the code. Since evaluating all 18,438 tasks is computationally expensive, we randomly sample 20 performance ranking tasks for each machine learning problem, resulting in 471 evaluation tasks in total.
Llama Team. 2024. The llama 3 herd of models. CoRR, abs/2407.21783. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Ed H. Chi, Quoc V. Le, and Denny Zhou. 2022. Chain-of-thought prompting elicits reasoning in large language models. In NeurIPS. World Labs. 2025. Marble: A multimodal world model. Zikai Xiao, Jianhong Tu, Chuhang Zou, Yuxin Zuo, Zhi Li, Peng Wang, Bowen Yu, Fei Huang, Junyang Lin, and Zuozhu Liu. 2026. Webworld: A largescale world model for web agent training. CoRR, abs/2602.14721.
• Synthetic Evaluation Tasks. Since there are currently no dedicated benchmarks for data science world modeling, we construct 540 evaluation tasks using our training data construction pipeline. To avoid data leakage, we use additional data sources from DABench (Hu et al., 2024) and MLEDojo (Qiang et al., 2025), which are different from the data sources used for synthesized training transitions. The tasks evaluate five core prediction capabilities:
Junjie Xing, Yeye He, Mengyu Zhou, Haoyu Dong, Shi Han, Lingjiao Chen, Dongmei Zhang, Surajit Chaudhuri, and H. V. Jagadish. 2025. MMTU: A massive multi-task table understanding and reasoning benchmark. CoRR, abs/2506.05587. An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, Chujie Zheng, Dayiheng Liu, Fan Zhou, Fei Huang, Feng Hu, Hao Ge, Haoran Wei, Huan Lin, Jialong Tang, and 40 others. 2025. Qwen3 technical report. CoRR, abs/2505.09388.
Execution Success Prediction (ESP). Predict whether a code action can execute successfully under the current environment state.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R. Narasimhan, and Yuan Cao. 2023. React: Synergizing reasoning and acting in language models. In ICLR.
Error Type Prediction (ETP). Predict the execution error category for failed code actions, such as syntax, runtime, or data-related errors. 10
For PP, we use 1 − RMSE: v u N u1 X t Score = 1 − (ŝi − si )2 , N
Execution Result Similarity (ERS). Evaluate the semantic similarity between predicted execution outputs and ground-truth execution results. Execution Keyword Matching (EKM). We use DeepSeek 3.2 to extract several keywords from the ground-truth execution results and evaluate whether the predicted execution outputs contain these keywords.
where ŝi and si denote the predicted and groundtruth task performance scores, respectively. For MLE-Bench Lite, we follow the official evaluation protocol and report the percentages of Gold, Silver, Bronze, and Any medals, as well as the percentage of tasks achieving above-median leaderboard performance. In addition, following MLEDojo, we compute a normalized leaderboard score ratio and use 1 − ratio as the final score metric.
Performance Prediction (PP). Given a machine learning task and a code solution, predict the resulting task performance. • MLE-Bench Lite (Chan et al., 2025). MLEBench is a machine learning benchmark that requires agents to autonomously solve realworld machine learning tasks. Due to the high computational cost, we use MLE-Bench Lite for evaluation and further remove one task with a dataset exceeding 100GB. The final evaluation set contains 21 tasks. A.2
A.3
For PR, ESP, and ETP, we use accuracy as the evaluation metric: N
(16)
i=1
where ŷi and yi denote the predicted and groundtruth labels, respectively. For EKM, we use keyword matching accuracy to measure the proportion of correctly predicted keywords: N
Score =
1 X |Kip ∩ Kig | , N |Kig |
(17)
i=1
where Kip and Kig denote the predicted and groundtruth keyword sets, respectively. For ERS, we use embedding cosine similarity between predicted outputs and ground-truth execution results: Score =
e⊤ p eg , ∥ep ∥∥eg ∥
Implementation Details.
DSWorld employs Qwen3-8B as the simulator backbone. The encoder is implemented using Harrier OSS v1 0.6B (Microsoft, 2026), while the Router is implemented as a two-layer MLP with hidden dimensions of 256 and 64, trained on collected code-execution time pairs. Data synthesis is performed using DeepSeek 3.2. For realworld transition collection, we use tasks from DACode (Huang et al., 2024) together with the ReAct framework. For synthesized transitions, we use data sources from MMTU (Xing et al., 2025), a large-scale multi-task table understanding and reasoning benchmark containing 28,136 table-centric questions over 61,763 real tables across 25 task categories, providing diverse data science environments and analytical workflows. For all experiments, each task is evaluated over three independent runs, and we report the mean and variance of the results. For SFT, we train the model for 5 epochs with a batch size of 32 and a learning rate of 1×10−5 . For RL, the reward function for execution prediction tasks is defined as the average score across the four execution prediction objectives:
Metrics
1 X Score = 1(ŷi = yi ), N
(19)
i=1
1 Rexec = (RESP + RETP + RERS + REKM ). 4 (20) For performance prediction tasks, the reward function is defined as:
(18)
Rperf = 1 − (ŝ − s)2 ,
where ep and eg denote the embeddings of the predicted and ground-truth execution outputs, respectively. The embeddings are extracted using Harrier OSS v1 0.6B.
(21)
where ŝ and s denote the predicted and groundtruth performance scores, respectively. 11
Table 4: Performance of DSWorld under different execution environments. Execution Prediction
Environment Ubuntu CentOS Windows
Performance Prediction
ESP ↑
ETP ↑
ERS ↑
EKM ↑
PP ↑
PR ↑
0.950±0.005 0.943±0.006 0.941±0.009
0.922±0.003 0.923±0.005 0.916±0.007
0.871±0.005 0.864±0.007 0.872±0.005
0.575±0.018 0.562±0.005 0.569±0.017
0.856±0.001 0.856±0.005 0.853±0.007
0.518±0.008 0.513±0.008 0.516±0.009
Table 5: Performance comparison of agents trained with different simulators.
Backbone Qwen3-8B Qwen3-14B Qwen3-8B Qwen3-8B Qwen3-8B
Simulator DeepSeek 3.2 Compiler DSWorld
We use a rollout size of 8 and a learning rate of 1 × 10−6 for 200 training steps. The maximum response length is set to 16K tokens. All experiments are conducted using VeRL (Jiang et al., 2025a) on 4 NVIDIA A800 GPUs.
Additional Experiments
B.1
Cross-Environment Generalization
C
Task Examples
In this section, we provide two representative examples from our constructed evaluation tasks. The first example evaluates execution-level transition prediction, where the model predicts detailed execution outputs generated by data analysis code. The second example evaluates performance prediction, where the model estimates the downstream task performance of a machine learning solution without actual execution.
In this section, we evaluate the robustness of DSWorld across different execution environments, including Ubuntu, CentOS, and Windows systems. As shown in Table 4, DSWorld achieves consistently strong performance across all environments, with only minor variations in both execution prediction and performance prediction metrics. In particular, Ubuntu achieves the best overall average performance, while CentOS and Windows remain highly competitive. These results demonstrate that DSWorld generalizes well across heterogeneous operating environments and is not overly dependent on a specific execution platform. B.2
0.781 0.777 0.778
Table 5 presents the results. Overall, agents trained with DSWorld achieve the best average performance among all settings. Consistent with the findings in Section 5.3, DSWorld-trained agents outperform the stronger Qwen3-14B baseline while only using Qwen3-8B as the backbone model. Compared with training using DeepSeek 3.2 as the simulator, DSWorld substantially improves downstream task performance, demonstrating stronger environment transition modeling capabilities. In addition, DSWorld achieves performance comparable to Compiler-based training while avoiding expensive real execution during training. These results further demonstrate that DSWorld can serve as an effective and scalable training environment for autonomous data science agents.
DACode ↑ 0.200±0.018 0.214±0.006 0.158±0.012 0.231±0.009 0.232±0.010
B
AVG. ↑
Example 1: Execution Prediction Task. Please preprocess and analyze Microsoft stock market data to compute volatility, log returns, VWAP, momentum indicators, risk-adjusted returns, trend signals, and cumulative growth metrics.
Evaluating DSWorld as a Training Environment on Additional Benchmark
To further evaluate the effectiveness of DSWorld as a training environment, we conduct additional experiments on 100 machine learning tasks from DACode (Huang et al., 2024). Specifically, we train autonomous data science agents under the same ReAct framework using different simulators, including DeepSeek 3.2, real execution through the Compiler, and DSWorld.
Data. | Unnamed: 0 | Date | Open | High | Low | Close | Volume | |————-:|:———-|——-:|——-:|—— :|——–:|———:| | 0 | 19-Jan-18 | 90.14 | 90.61 | 89.66 | 90 | 36875013 | | 1 | 18-Jan-18 | 89.8 | 90.67 | 89.66 | 90.1 |
12
24159683 | ... omit here ... | 81 | 22-Sep-17 | 73.99 | 74.51 | 73.85 | 74.41 | 14111365 | | 82 | 21-Sep-17 | 75.11 | 75.24 | 74.11 | 74.21 | 19186140 |
|– train.csv
Action. import pandas as pd import numpy as np df = pd.read_csv(’microsoft.csv’) df[’Date’] = pd.to_datetime(df[’Date’], format=’%d-%b-%y’) ... omit here ... final_result = df[[’Date’, ’Close’, ’VWAP’, ’Risk_Adjusted_Return’]].tail(10) print(final_result)
Action. import pandas as pd import numpy as np import lightgbm as lgb # Load data train = pd.read_csv(’train.csv’) test = pd.read_csv(’test.csv’) ... omit here ... submission = pd.DataFrame(’id’: test_ids, ’ACTION’: test_pred) submission.to_csv(’submission.csv’, index=False)
## Data Preview ... omit here ... [END Preview]
Next State. Date Close VWAP Risk_Adjusted_Return 73 2017-10-04 74.69 83.796359 -2.640024 74 2017-10-03 74.26 83.725570 -0.914520 75 2017-10-02 74.61 83.641401 0.685921 76 2017-09-29 74.49 83.548067 -0.234191 77 2017-09-28 73.87 83.485573 -1.188027 78 2017-09-27 73.85 83.376069 -0.039253 79 2017-09-26 73.26 83.270265 -1.140838 80 2017-09-25 73.26 83.131894 0.000000 81 2017-09-22 74.41 83.062009 1.608727 82 2017-09-21 74.21 82.966613 -0.278535
Score. 0.813
D
Case Studies
In this section, we present a representative case study to illustrate the prediction capabilities of DSWorld on real-world machine learning tasks. The example shows that DSWorld can effectively predict both execution outcomes and downstream task performance without actual code execution. Specifically, DSWorld correctly predicts that the generated solution can execute successfully and produce a valid submission file. In addition, DSWorld accurately estimates the performance, achieving a predicted score of 0.528 compared with the ground-truth score of 0.543. The generated reasoning trajectory further demonstrates that DSWorld can analyze multiple aspects of the machine learning pipeline, including dataset characteristics, preprocessing strategies, feature representations, model selection, and potential performance limitations. For example, DSWorld recognizes that the TF-IDF and Logistic Regression pipeline provides a strong text classification baseline while also identifying its limitations in handling complex linguistic phenomena such as negation and sarcasm. These results demonstrate that DSWorld can capture high-level relationships between code implementations and downstream task performance, which is important for efficient autonomous data science agent planning and search.
Example 2: Performance Prediction Task. Machine Learning task: ### Description In this competition, participants are tasked with developing a model that predicts employee access needs based on historical data. ... omit here ... Participants should focus on creating a robust model that minimizes manual access transactions as employee attributes change over time. Data. "— COMPREHENSIVE DATA REPORT — ## Data Directory Structure |– description.txt |– sample_submission.csv |– test.csv
13
Case of DSWorld
... omit here ... [END Preview]
Task. Machine Learning task:
Action. import pandas as pd from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.linear_model import LogisticRegression
### Description Dive into the world of sentiment analysis with the Rotten Tomatoes movie review dataset, a rich corpus originally collected by Pang and Lee. This competition invites participants to label phrases on a five-point sentiment scale: negative, somewhat negative, neutral, somewhat positive, and positive. The challenge lies in navigating obstacles such as sentence negation, sarcasm, terseness, and language ambiguity. Hosted by Kaggle, this competition is designed for the machine learning community to engage in fun and practice, inspired by the pioneering work of Socher et al. Explore the accompanying website for source code, a live demo, and an online interface to aid in model training: http://nlp.stanford.edu/sentiment/.
# Load data train_df = pd.read_csv(’train.tsv’, sep=’’) test_df = pd.read_csv(’test.tsv’, sep=’’) # Preprocess X_train_text train_df[’Phrase’].fillna(”).str.lower() y_train = train_df[’Sentiment’] X_test_text test_df[’Phrase’].fillna(”).str.lower()
=
=
# TF-IDF vectorization vectorizer = TfidfVectorizer(max_features=10000, ngram_range=(1, 2), stop_words=’english’) X_train_vec = vectorizer.fit_transform(X_train_text) X_test_vec = vectorizer.transform(X_test_text)
### Evaluation Submissions are evaluated on classification accuracy (the percent of labels that are predicted correctly) for every parsed phrase. The sentiment labels are:
# Train logistic regression on full data model = LogisticRegression(max_iter=1000, random_state=42, class_weight=’balanced’) model.fit(X_train_vec, y_train)
0 - negative 1 - somewhat negative 2 - neutral 3 - somewhat positive 4 - positive
# Predict on test test_predictions model.predict(X_test_vec)
... omit here ...
# Create submission submission = pd.DataFrame( ’PhraseId’: test_df[’PhraseId’], ’Sentiment’: test_predictions )
Data. "— COMPREHENSIVE DATA REPORT — ## Data Directory Structure |– description.txt |– sample_submission.csv |– test.csv |– train.csv
# Save submission.to_csv(’submission.csv’, index=False) Ground Truth. 0.543
## Data Preview
Prediction of DSWorld.
14
=
Reasoning. <think>First, I assess the code’s execution feasibility. The code imports standard libraries (pandas, sklearn) that are widely available in data science environments. The syntax appears correct with no obvious typos or indentation errors. The data loading uses pd.read_csv with sep=’’, which is appropriate for TSV files as indicated in the dataset description. The column names referenced (’Phrase’, ’Sentiment’, ’PhraseId’) align with typical expectations for this dataset, though I note that the data preview shows errors loading the TSV files, which could indicate potential issues with file paths or formatting. However, the code structure itself is sound and likely to run without syntax or immediate runtime errors.
Given the above, the code is highly likely to execute successfully without errors, assuming the TSV files exist at the expected paths and are correctly formatted. The absence of error handling for missing files or columns is a minor risk, but the dataset description provides clear file names and structure. Now, for performance estimation: The approach uses a simple but standard method (TF-IDF + Logistic Regression) that has been shown to achieve moderate performance on sentiment analysis tasks. However, several factors may limit the score: (1) The model does not incorporate advanced text preprocessing (e.g., removing punctuation, handling negations) or feature engineering (e.g., sentiment lexicons). (2) No hyperparameter tuning or cross-validation is performed. (3) The dataset includes challenges like negation and sarcasm, which TF-IDF features may not capture well. (4) The class weighting is applied but may not fully address underlying distribution shifts.
Next, I examine the preprocessing and modeling pipeline. The use of TF-IDF vectorization with max_features=10000, ngram_range=(1,2), and stop_words=’english’ is a reasonable baseline for text classification. The LogisticRegression configuration with max_iter=1000 and class_weight=’balanced’ addresses potential class imbalance, which is relevant given the five-class sentiment task. The pipeline follows a standard supervised learning workflow: load data, preprocess text, vectorize, train model, predict, and save results. There are no apparent logical inconsistencies, such as data leakage or incorrect handling of train/test splits.
For a five-class classification task with balanced class weighting, a reasonable baseline accuracy might be in the range of 0.40.6, assuming random guessing would yield 0.2 (if classes were highly imbalanced) to 0.25 (if balanced). Given the simplicity of the pipeline and the inherent difficulty of the task, I expect the accuracy to be modest but above random. The lack of advanced techniques (e.g., deep learning, BERT) suggests the score will not be state-of-the-art, but the TF-IDF + Logistic Regression combination often achieves solid performance on clean text data. Considering the evaluation context (Kaggle competition), a score around 0.5-0.6 is plausible, with the upper bound limited by the model’s inability to handle complex linguistic patterns effectively. The final estimated score would reflect this baseline performance, acknowledging both the method’s strengths and its limitations. </think> Score. 0.528
However, I must consider potential hidden issues. The dataset description mentions that phrases are parsed from sentences, and some phrases may be repeated. The code does not explicitly handle duplicate PhraseIds, but this is unlikely to cause execution failure. The evaluation metric is classification accuracy, and the submission format requires PhraseId and Sentiment columns, which the code correctly generates and saves. The submission file name matches the requirement (’submission.csv’).
15