DSWorld: A Data Science World Model for Efficient Autonomous Agents

🕒 Published (v1): 2026-07-17 12:14 UTC · Source: Arxiv · link

Why this paper was selected

World model cuts expensive trial-and-error in data-science agents; pattern generalizes to any code-execution agent loop

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

DSWorld is a data science world model that predicts the outcomes of data science operations (feature engineering, model training, evaluation) without executing them, replacing expensive real computation with LLM-based simulation. It achieves ~14× speedup in RL-based agent training and 3–6× speedup in search-based inference, while maintaining competitive downstream performance on MLE-Bench Lite.

Problem

Autonomous data science agents rely on iterative trial-and-error loops that require expensive real execution (e.g., ML-Master spends >86% of runtime on model training). No environment model exists to let agents anticipate the effect of data science operations before committing to costly computation, creating a fundamental scalability bottleneck for both RL training and test-time search.

Method

DSWorld frames the data science environment as a world to be modeled: given state \(S_t\) and action \(A_t\), a transition model predicts \(S_{t+1} = W(S_t, A_t)\).

The framework has four components: 1. State Constructor — rule-based program extracting task description, dataset statistics, execution history, outputs, and error messages into a structured state \(S_t = \text{SC}(E_t)\). 2. Router — encodes actions with Harrier OSS v1 0.6B, then classifies via a 2-layer MLP whether execution is lightweight (execute) or expensive (simulate). 3. Compiler — runs lightweight actions in the real environment for guaranteed accuracy; also serves as a timeout fallback. 4. Simulator — Qwen3-8B fine-tuned to predict next states (execution status, outputs, errors, performance scores) for expensive actions without real execution.

Training uses a two-stage strategy: SFT warm-up on DSWorld-8K (8K real + synthesized transition trajectories with chain-of-thought rationales), then Reflective World Model Optimization (RWMO) — an error-aware RL approach. The Simulator generates \(n\) rollout predictions, computes reflection feedback \(f = S_\theta(\hat{S}', S')\) that identifies discrepancies, produces refined predictions \(\hat{S}_r\), and jointly optimizes both original and refined outputs via GRPO:

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

The reward evaluates correctness of predicted execution status, outputs, error type, and task performance.

DSWorld-8K dataset: combines real agent trajectories from existing benchmarks with a synthetic pipeline that samples tables from MMTU (>60K real tables), constructs states, samples operation/error type pairs, generates executable actions via LLM, executes them in the Compiler for ground-truth verification, and synthesizes CoT trajectories.

Key Contributions

• Formalizes the concept of a Data Science World Model for predicting environment state transitions in autonomous data science workflows.
• DSWorld framework: cost-aware routing (Compiler vs. LLM Simulator) with timeout fallback, enabling hybrid execution-simulation.
• Reflective World Model Optimization (RWMO): error-aware RL that generates reflection feedback on mispredictions and refines outputs via GRPO.
• DSWorld-8K: 8K-scale transition trajectory dataset with real and synthesized samples, including CoT rationales.

Results

Transition prediction (vs. strongest baseline o4-mini, AVG score): - DSWorld: 0.781 vs. o4-mini: 0.576 — +35.6% improvement - Execution success prediction (ESP): 0.950 vs. 0.680 (+39.7%) - Error type prediction (ETP): 0.922 vs. 0.585 (+57.6%) - Execution result similarity (ERS): 0.871 vs. 0.489 (+78.1%)

Agent training (MLE-Bench Lite, 10 RL steps on MLE-Dojo): - DSWorld-trained Qwen3-8B score: 17.67 vs. Compiler-trained: 18.11 (−2.4% performance gap) - Training time: 277 min (DSWorld) vs. 3854 min (Compiler) — ~14× speedup - DSWorld-trained agent outperforms Qwen3-14B baseline (16.50) despite smaller backbone

Agent inference (search-based agents on MLE-Bench Lite): - AIDE + DSWorld: Score 10.58 in 676s vs. Compiler: 10.7 in 4102s (~6× speedup, <1.2% score drop) - ML-Master + DSWorld: Score 10.34 in 371s vs. Compiler: 12.39 in 1421s (~3.8× speedup) - DeepSeek 3.2 as simulator causes severe performance degradation (score drops to 7.0–7.2) despite similar time savings

Ablation: - Qwen3-8B + SFT alone improves avg. by +37.5% over untuned backbone; adding GRPO (DSWorld) further improves to best overall

Limitations

• The Simulator is trained on data science workflows (NumPy/Pandas ecosystems); generalization to other programming environments or domain-specific libraries is untested.
• Performance prediction tasks (PP, PR) remain harder; DSWorld only matches rather than clearly surpasses strong baselines on ML algorithm reasoning.
• DSWorld-8K is 8K samples — scale sensitivity experiments suggest performance still grows with more data, implying the current dataset may be a bottleneck.
• Evaluation is on MLE-Bench Lite (a subset), and high variance in medal-rate metrics (large ±σ) makes fine-grained comparisons unreliable.
• The Router is a learned MLP; routing errors on edge cases fall back to a timeout mechanism rather than a principled recovery.

Relevance to Agentic AI / LLM Agents

DSWorld directly addresses one of the core bottlenecks in agentic systems operating in computationally heavy environments: the cost of environment interaction during both RL training and test-time search. By decoupling agent reasoning from real execution via a learned world model, it enables more exploration steps per unit time — a critical capability for sample-efficient RL and beam/tree-search inference strategies used in state-of-the-art data science agents (AIDE, ML-Master, AutoMLGen). The Reflective World Model Optimization strategy, which uses error-aware reflection to refine predictions within the RL loop, is a transferable design pattern for training LLM-based environment simulators in other agentic domains (web, GUI, code execution). This work extends the world-model paradigm from physical/visual domains into structured, symbolic data science environments, opening a new design axis for agent efficiency beyond prompt engineering and backbone scaling.