PaperlayerView exact revision on arXiv

Source-linked AI summary

DSWorld: A Data Science World Model for Efficient Autonomous Agents

Zherui Yang, Fan Liu, Hao Liu

arXiv:2607.15901v1cs.AI

TL;DR

Autonomous data science agents spend substantial time on costly computation, limiting workflow efficiency. DSWorld predicts workflow state transitions before execution, and outperforms the strongest LLM baseline by 35.6% on transition prediction tasks.

  • Problem

    Autonomous data science agents rely on computationally expensive execution, motivating transition models that anticipate operation effects before costly computation.

  • Method

    DSWorld models workflow state transitions using structured state construction, routing, compilation, and an LLM-based simulator, with reflective reinforcement learning for prediction improvement.

  • Results

    35.6% average improvement over the strongest LLM baseline on transition prediction tasks, while DSWorld maintains strong predictive performance across diverse tasks.

  • Takeaways & Limitations

    DSWorld substantially accelerates autonomous agent training and inference while maintaining strong downstream agent performance.

  • Takeaways & Limitations

    Transition prediction quality remains constrained by the underlying LLM simulator, which may produce inaccurate predictions in complex workflow scenarios.

Abstract

from arXiv · show

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.

1 Introduction

DSWorld introduces Data Science World Models that predict workflow state transitions and operation effects before costly execution, addressing the computational overhead of trial-and-error autonomous data science. The framework combines structured state construction, cost-aware routing, scalable transition-data generation, and reflective refinement to accelerate agent training and inference while preserving strong performance.

  • Motivation: Autonomous data science agents rely on iterative trial-and-error workflows whose analytical computation consumes most execution time, with ML-Master spending over 86% on model training in MLE-Bench.Existing workflows explore numerous candidate solutions while incurring costs from data processing, model training, evaluation, and workflow updates.
  • Data Science World Models: Data Science World Models predict next environment states from workflow states and candidate operations, including dataset and model changes, execution feedback, errors, and performance signals.This enables agents to anticipate operation effects without performing expensive real-world execution.
  • DSWorld Framework: DSWorld constructs structured states containing tasks, datasets, execution histories, outputs, and environment status, then applies cost-aware routing to balance prediction accuracy and efficiency.Its Router determines whether an agent action requires heavy computation before selecting how to handle it.
  • Training Data and Optimization: DSWorld-8K contains 8 thousand high quality data science agent trajectories collected from authentic tasks and augmented with synthesized Chain-of-Thought trajectories explaining transition logic.The framework also introduces reflective reinforcement learning that improves transition prediction through error-aware reflection and iterative refinement.
  • Evaluation: 14× RL training speedup and 3-6× inference acceleration are achieved while maintaining strong downstream agent performance; transition prediction outperforms the strongest LLM baseline by 35.6% on average.These results cover DSWorld’s use in autonomous-agent training, inference, and transition prediction tasks.

2 Related works

Prior world models predict environment dynamics across physical and digital domains, while world modeling for autonomous data science environments remains largely unexplored. Related autonomous data science agents mainly improve performance through backbone training or test-time scaling.

  • World models: World models predict future states conditioned on current states and actions for prediction, control, and simulation.The paper situates world models as environment-dynamics models spanning multiple application domains.
  • World models: Physical-world models support video generation, 3D world generation, and latent-space dynamics modeling.The cited examples include video generation, 3D world generation, and latent-space dynamics modeling.
  • World models: Digital-world models simulate webpage or software state transitions for web and GUI agents, but data science environments remain largely unexplored.DSWorld differs from prior work by modeling environment transitions in data science workflows.
  • Autonomous data science agents: Autonomous data science agents mainly improve performance through backbone training or test-time scaling strategies.ML-Agent exemplifies reinforcement-learning-based training for machine learning tasks, while AIDE exemplifies test-time scaling.

3 Preliminaries

The preliminaries define a data science workflow state as a structured representation of the task, data, execution environment, and logs or progress. They then define agent actions and a Data Science World Model that predicts operation effects before costly computation.

  • Workflow State: The workflow state is S_t = {T_t, D_t, P_t, L_t}, capturing the task, data state, execution environment, and execution logs, intermediate outputs, or task progress.The data state includes dataset statistics and previews, while the execution environment includes libraries and runtime configurations.
  • Action: At time step t, a data science agent produces an action A_t conditioned on the current state S_t.Actions include feature engineering, model training, and evaluation.
  • Data Science World Model: A Data Science World Model predicts the effects of data science operations before costly computation.The world model is denoted by W.

4 Methodology

DSWorld models data science environment transitions with four components: structured state construction, cost-aware routing, direct compilation, and LLM-based simulation. It is trained through supervised fine-tuning and reflective reinforcement learning on real and synthesized transition trajectories.

  • Framework Overview: DSWorld comprises a State Constructor, Router, Compiler, and LLM-based Simulator for modeling data science environment transitions.The Router selects direct execution or simulation based on the current state and generated action.
  • State Constructor: The State Constructor extracts task descriptions, dataset statistics, previews, execution environments, histories, intermediate outputs, and errors into a unified structured state.It represents the transformation as S_t = SC(E_t), where E_t is the environment at time step t.
  • Router: The Router encodes each agent action and determines whether to execute it directly or simulate its transition.Lightweight manipulation and inspection operations are executed, whereas expensive operations such as large-scale model training are simulated.
  • Compiler and Simulator: The Simulator predicts execution outcomes and potential errors without real execution, while timed-out Compiler actions are redirected to simulation.This hybrid mechanism balances efficiency and accuracy for scalable environment interaction.
  • Posttraining: DSWorld uses two-stage posttraining: supervised fine-tuning initializes transition modeling, followed by Reflective World Model Optimization that analyzes errors and refines predictions.The reinforcement-learning reward evaluates predicted execution status, outputs, errors, and task performance, with trajectories jointly optimized using GRPO.
  • DSWorld-8K Dataset: Approximately 8K transition samples with reasoning trajectories are constructed from both authentic agent executions and verified synthesized transitions.Synthetic states draw on diverse datasets, while synthesized actions cover successful executions and diverse execution failures before Compiler-based verification.

5 Experiments

Experiments show that DSWorld achieves leading transition-prediction performance while accelerating agent training and search-based inference. Ablations further support Reflective World Model Optimization and favorable scaling with more data and larger backbones.

  • Transition prediction: 35.6% average improvement over o4-mini: DSWorld achieves the best performance on nearly all transition-prediction dimensions.It substantially improves execution-related prediction accuracy while remaining competitive on performance-related tasks.
  • Agent training: 14× acceleration: DSWorld-trained agents remain competitive with Compiler-based training and outperform training with DeepSeek 3.2 as simulator.Compiler-based training achieves slightly better final performance, but the downstream-performance gap is relatively small compared with DSWorld’s efficiency gain.
  • Search-based inference: 3-6× acceleration: DSWorld largely preserves downstream performance for search-based agents while reducing inference time versus Compiler-based execution.Using DeepSeek 3.2 reduces execution time but causes severe performance degradation from inaccurate feedback and hallucinated transitions.
  • Optimization ablation: 37.5% average improvement: Qwen3-8B-sft outperforms the original Qwen3-8B backbone across all tasks, while Qwen3-8B-grpo adds 1.05%.DSWorld further improves overall performance by 1.3% over Qwen3-8B-grpo and by 2.36% over Qwen3-8B-sft through reflective error-aware optimization.
  • Scaling analysis: 0.1k to 6.4k samples: Performance steadily improves with more transition data, and scaling the Qwen3 backbone from 0.6B to 14B parameters also benefits DSWorld.These results indicate favorable scaling with both training-data size and model size.

6 Conclusion

DSWorld is a learned transition model for data science workflows that predicts the effects of operations from current states and actions before costly computation. Its training is supported by a large-scale trajectory dataset combining real and synthetic transitions and by Reflective World Model Optimization.

  • DSWorld models data science workflow transitions conditioned on the current state and actions.This enables agents to anticipate operation effects before performing costly computation during training and inference.
  • The framework supports anticipatory data science agents during both training and inference by avoiding costly computation when effects can be predicted.
  • DSWorld training uses a large-scale trajectory dataset that combines real and synthetic transition data.
  • The paper further proposes Reflective World Model Optimization to support effective world model learning.

Limitations … A.2 Metrics

The paper evaluates DSWorld on synthetic data-science prediction tasks and MLE-Bench Lite using metrics tailored to execution, error, similarity, keyword, and performance prediction. It also acknowledges limitations in external tool-call modeling, simulator accuracy, and synthesized trajectory distributions.

  • Limitations: DSWorld does not explicitly model external tool-call transitions, and its simulator may produce inaccurate predictions in complex workflows.The paper also notes that synthesized trajectories may still exhibit distribution gaps, although the supplied passage is truncated.
  • A.1 Benchmarks: The evaluation uses 540 synthetic tasks built from additional DABench and MLE-Dojo data sources to avoid training-data leakage.These tasks assess five core prediction capabilities.
  • A.1 Benchmarks: The five synthetic-task capabilities are execution success, error type, execution-result similarity, execution-keyword matching, and performance prediction.ESP predicts executability; ETP predicts error categories; ERS compares semantic output similarity; EKM checks extracted keywords; PP predicts task performance.
  • A.1 Benchmarks: MLE-Bench Lite contributes 21 real-world machine-learning tasks after removing a task whose dataset exceeded 100GB because of computational cost.The evaluation follows the lightweight version of MLE-Bench.
  • A.2 Metrics: Accuracy evaluates PR, ESP, and ETP, with predicted labels compared against ground-truth labels.The metric definition identifies the predicted and ground-truth labels as ŷ_i and y_i.
  • A.2 Metrics: EKM uses keyword matching accuracy, whereas ERS uses embedding cosine similarity between predicted and ground-truth execution outputs.ERS embeddings are extracted using Harrier OSS v1 0.6B.
  • A.2 Metrics: PP uses 1 − RMSE to compare predicted and ground-truth task performance scores.The supplied metric passages define the compared scores as ŝ_i and s_i.
  • A.2 Metrics: MLE-Bench Lite reports Gold, Silver, Bronze, Any, and above-median leaderboard percentages, plus 1 − normalized leaderboard score ratio.The latter is used as the final score metric.

A.3 Implementation Details. · B Additional Experiments · B.1 Cross-Environment Generalization

DSWorld combines specified model components, data sources, and training procedures with evaluations across execution environments and simulator settings. Cross-environment results show consistently strong performance with minor metric variation, with Ubuntu achieving the best overall average.

  • A.3 Implementation Details.: DSWorld uses Qwen3-8B as its simulator backbone, Harrier OSS v1 0.6B as encoder, and a two-layer MLP Router with hidden dimensions 256 and 64.The Router is trained on collected code-execution time pairs.
  • A.3 Implementation Details.: DeepSeek 3.2 synthesizes data, while real-world transitions use DACode tasks with ReAct and synthesized transitions use MMTU data sources.The supplied passage truncates the full list of MMTU-based sources.
  • A.3 Implementation Details.: SFT runs for 5 epochs with batch size 32 and learning rate 1×10−5, while execution-prediction RL reward averages RESP, RETP, RERS, and REKM.The reward is defined as 1/4(RESP + RETP + RERS + REKM).
  • B Additional Experiments: Additional experiments compare DSWorld under different execution environments and compare agents trained with different simulators.These comparisons are presented in Tables 4 and 5.
  • A.3 Implementation Details.: RL uses rollout size 8, learning rate 1 × 10−6, and 200 training steps, with maximum response length 16K tokens on 4 NVIDIA A800 GPUs.Experiments are conducted using VeRL.
  • B.1 Cross-Environment Generalization: Across Ubuntu, CentOS, and Windows, DSWorld maintains consistently strong execution-prediction and performance-prediction metrics with only minor variations.Ubuntu achieves the best overall average performance, while CentOS and Windows remain highly competitive.

B.2 Evaluating DSWorld as a Training Environment on Additional Benchmark

On 100 DACode machine-learning tasks, DSWorld is evaluated as a ReAct training environment against DeepSeek 3.2 and real Compiler execution. DSWorld-trained agents achieve the best average performance, surpassing the stronger Qwen3-14B baseline with a Qwen3-8B backbone and outperforming DeepSeek 3.2 simulation.

  • Experimental setup: 100 DACode machine-learning tasks are used to compare ReAct agents trained with DeepSeek 3.2, real Compiler execution, and DSWorld.The experiments evaluate DSWorld as a training environment under the same ReAct framework.
  • Results: DSWorld-trained agents achieve the best average performance among all evaluated training environments.This result is reported in Table 5.
  • Results: DSWorld-trained agents outperform the stronger Qwen3-14B baseline while using Qwen3-8B as the backbone model.This finding is consistent with the paper’s results in Section 5.3.
  • Results: Compared with DeepSeek 3.2 as the simulator, DSWorld substantially improves downstream task performance.The comparison demonstrates stronger environment transition modeling capabilities for DSWorld.

C Task Examples

The section presents two representative evaluation examples: execution-level transition prediction and performance prediction. They cover forecasting detailed data-analysis outputs and estimating downstream machine-learning performance without actual execution.

  • C Task Examples: Two evaluation examples assess execution-level transition prediction and performance prediction.The first predicts detailed execution outputs from data-analysis code, while the second estimates downstream machine-learning performance without execution.
  • Example 1: Execution Prediction: Execution Prediction uses Microsoft stock data to compute volatility, log returns, VWAP, momentum, risk-adjusted returns, trends, and cumulative growth.The workflow parses dates and produces a final result containing Date, Close, VWAP, and Risk_Adjusted_Return.
  • Example 2: Performance Prediction: Performance Prediction targets employee access-need prediction from historical data while minimizing manual access transactions as attributes change.The example loads train.csv and test.csv with LightGBM and generates an ACTION submission.

D Case Studies

DSWorld’s case study shows it can predict successful execution, valid submission generation, and downstream machine-learning performance without running the code. For Rotten Tomatoes sentiment classification, it evaluates a TF-IDF plus Logistic Regression pipeline and predicts a score close to ground truth.

  • Case Study: DSWorld correctly predicts successful execution and valid submission-file generation without actual code execution.The case study assesses both execution outcomes and downstream task performance.
  • Pipeline Analysis: The evaluated pipeline lowercases text, applies TF-IDF features with 10,000 maximum features and 1–2-gram ranges, then trains balanced Logistic Regression.The model uses English stop-word removal, max_iter=1000, and random_state=42.
  • Execution Analysis: The reasoning identifies standard libraries, correct syntax, expected TSV structure, and matching submission columns as supporting successful execution.It notes missing-file or missing-column handling as a minor residual risk when expected files are absent or malformed.
  • Performance Analysis: The predicted performance is limited by basic preprocessing, absent feature engineering, no hyperparameter tuning or cross-validation, and the task’s inherent difficulty.The reasoning characterizes TF-IDF plus Logistic Regression as a moderate, non-state-of-the-art baseline compared with advanced methods such as BERT.
Loading 2607.15901v1…