DSWorld: A Data Science World Model for Efficient Autonomous Agents

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 14times and search-based inference by approximately 3-6times 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.

한 줄 요약

DSWorld는 데이터 과학 작업의 효율성을 높이기 위해 작업 흐름의 상태 전이를 예측하는 모델로, RL 기반 학습을 14배, 추론을 3~6배 가속화한다.

핵심 기여도

데이터 과학 월드 모델(Data Science World Model) 개념 도입: 실제 실행 없이 데이터 과학 작업의 상태 전이를 예측.
DSWorld 프레임워크 제안: State Constructor, Router, Compiler, LLM-based Simulator 모듈을 결합.
Reflective World Model Optimization 제안: 오류 인식 기반 강화 학습 전략으로 전이 예측 정확도 35.6% 개선.
DSWorld-8K 데이터셋 구축: 8천 개의 고품질 데이터 과학 작업 흐름 전이 데이터.

핵심 아이디어

기존 자율 데이터 과학 에이전트는 반복적인 시도-오류 작업 흐름에 의존하며, 이는 높은 계산 비용을 유발한다. 이를 해결하기 위해 DSWorld는 **작업 흐름의 상태와 후보 작업을 조건으로 하여 다음 상태를 예측**하는 월드 모델을 도입한다. 이는 실제 실행 없이 작업 효과를 사전에 예측함으로써 학습 및 추론 과정을 가속화한다. 핵심 아이디어는 **LLM 기반 시뮬레이터와 경량 실행 컴파일러를 결합**하여 정확성과 효율성을 균형 있게 유지하는 것이다. 특히, **Reflective World Model Optimization**은 오류를 반영하여 예측을 반복적으로 개선하는 강화 학습 전략으로, 전이 예측 정확도를 높인다.

기술적 접근법

**State Constructor**: 원시 데이터 과학 환경을 작업, 데이터셋, 실행 이력, 출력, 상태 등을 포함한 구조화된 상태 표현으로 변환.
**Router**: 에이전트의 액션에 따라 경량 작업은 **Compiler**를 통해 실행, 계산 비용이 높은 작업은 **LLM-based Simulator**로 예측.
**Reflective World Model Optimization (RWMo)**: 오류 인식 기반 반복적 강화 학습 전략.
**DSWorld-8K**: 8천 개의 고품질 데이터 과학 작업 흐름 전이 데이터. 실제 데이터와 CoT 추론을 결합한 합성 파이프라인 사용.
**모델**: Qwen3-8B 기반, SFT 및 GRPO(Reflective World Model Optimization)를 사용한 학습.

주요 결과

RL 기반 에이전트 학습 가속: 약 **14×**.
검색 기반 추론 가속: 약 **3~6×**.
전이 예측 성능: 최고 LLM 베이스라인 대비 **35.6% 개선**.
Qwen3-8B-sft 모델: SFT만 적용 시 평균 성능 **37.5% 향상**.
Qwen3-8B-grpo 모델: GRPO 적용 시 전체 성능 **1.05% 향상**.
DSWorld: RWMo 적용 시 Qwen3-8B-grpo 대비 **1.3%**, Qwen3-8B-sft 대비 **2.36% 향상**.

의의 및 한계

DSWorld는 자율 데이터 과학 시스템의 효율성과 확장성을 크게 향상시키며, **작업 흐름의 상태 전이를 사전 예측**함으로써 실제 실행 비용을 줄인다. 특히, **Reflective World Model Optimization**은 오류 인식을 통해 예측 정확도를 반복적으로 개선하는 점에서 학술적 기여가 크다. 그러나, **작업 흐름의 복잡성과 다변성**에 따라 예측 오류가 발생할 수 있으며, **모든 작업 흐름에 대한 전이 데이터를 확보하는 것이 여전히 어려움**이다. 또한, **LLM 기반 시뮬레이터의 예측 한계**가 있을 수 있어, 실제 실행과의 일관성을 지속적으로 검증해야 한다.

실용적 활용

DSWorld는 **자동화된 데이터 분석, 예측 모델링, 머신 러닝 경진대회 자동화** 등 다양한 데이터 과학 작업에 적용 가능하다. 특히, **시간과 계산 자원이 제한된 환경**에서 효율적인 작업 흐름 최적화를 지원하며, **대규모 데이터셋 처리나 복잡한 모델 학습 과정**에서 유용하게 사용될 수 있다.