DSWorld: 효율적인 자율 에이전트를 위한 데이터 과학 세계 모델

초록

강력한 데이터 이해 및 의사 결정 능력에도 불구하고, 자율 데이터 사이언스 에이전트는 여전히 값비싼 계산을 수반하는 시행착오 기반 워크플로에 크게 의존합니다. 이러한 병목 현상은 실제 실행 전에 데이터 사이언스 연산의 효과를 예측할 수 있는 모델의 필요성을 제기합니다. 본 논문에서는 데이터 사이언스 실행 환경을 모델링하여 현재 워크플로 상태와 후보 연산에 따른 환경 상태 전이를 예측하는 데이터 사이언스 세계 모델(Data Science World Model) 개념을 소개합니다. 또한 구조화된 상태 구축, 비용 인식 라우팅, 경량 실제 실행 및 고비용 연산을 위한 LLM 기반 시뮬레이터를 결합한 실용적 프레임워크인 DSWorld를 제안합니다. 학습을 지원하기 위해 8K 규모의 전이 궤적 데이터셋을 구축하고, 전이 예측 개선을 위한 오류 인식 강화학습 전략인 반성적 세계 모델 최적화(Reflective World Model Optimization)를 도입합니다. 실험 결과, DSWorld는 경쟁력 있는 성능을 유지하면서 RL 기반 에이전트 학습을 약 14배, 탐색 기반 추론을 약 3~6배 가속화하며, 전이 예측 과제에서 가장 강력한 LLM 기준선보다 35.6% 높은 성능을 보여줍니다. 코드는 https://anonymous.4open.science/r/DSWorld 에서 확인할 수 있습니다.

English

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.