DSGym: A holistic framework for evaluating and training data science agents

DSGym: A holistic framework for evaluating and training data science agents

⚡️ FlashAttention-4: up to 1.3× faster than cuDNN on NVIDIA Blackwell →

Introducing Together AI's new look →

🔎 ATLAS: runtime-learning accelerators delivering up to 4x faster LLM inference →

⚡ Together GPU Clusters: self-service NVIDIA GPUs, now generally available →

📦 Batch Inference API: Process billions of tokens at 50% lower cost for most models →

🪛 Fine-Tuning Platform Upgrades: Larger Models, Longer Contexts →

All blog posts

Research

Published 1/26/2026

DSGym: A holistic framework for evaluating and training data science agents

Authors

Fan Nie, Junlin Wang, Harper Hua, Federico Bianchi, Yongchan Kwon, Zhenting Qi, Owen Queen, Shang Zhu, James Zou

Table of contents

40+ Models Chosen for Production...40+ Models Chosen for Production...40+ Models Chosen for Production...

Links in this article

arXiv paper Github repo ‍

Summary

Current data science benchmarks rely on incompatible evaluation interfaces. Moreover, many tasks can be solved without using the underlying data. We address these limitations by introducing DSGym, an integrated framework for evaluating and training data science agents in self-contained execution environments. Using DSGym, we trained a state-of-the-art open-source data science agent. arXiv paper:  https://arxiv.org/abs/2601.16344 Github repo: https://github.com/fannie1208/DSGym

Data science serves as the computational engine of modern scientific discovery. However, evaluating and training LLM-based data science agents remains challenging because existing benchmarks assess isolated skills in heterogeneous execution environments, making integration costly and fair comparisons difficult. We introduce DSGym , a unified framework that integrates diverse data science evaluation suites behind a single API with standardized abstractions for datasets, agents, and metrics. DSGym unifies and refines existing benchmarks while expanding the scope with novel scientific analysis tasks (90 bioinformatics tasks from academic literature) and challenging end-to-end modeling competitions (92 Kaggle competitions). Beyond evaluation, DSGym provides trajectory generation and synthetic query pipelines for agent training—we demonstrate this by training a 4B model on 2k generated examples, achieving state-of-the-art performance among open-source models.

(a) Typical Scientific Discovery Process. DSGym focuses on the data-driven investigation phase with tasks spanning 10+ domains and file types. (b) DSGym Overview. The framework supports both evaluation (data analysis and prediction tasks) and training through synthetic data generation. Framework and datasets One of the main contributions of DSGym is that it abstracts the complexity of code execution behind containers that can be allocated in real time to execute code safely; these containers come with pre-installed dependencies and data available for processing. DSGym provides a unified JSON interface for all benchmarks, where each task is expressed as: data files, query prompt, evaluation metric, and metadata. We strive to make the design modular and straightforward. In this way adding new tasks, agent scaffolds, tools, and evaluation scripts should be simpler for users. The tasks in DSGym are categorized into two primary tracks: Data Analysis (query-answering via programmatic analysis). Data Prediction (end-to-end ML pipeline development).

In addition to integrating established benchmarks like MLEBench and QRData, DSGym introduces original datasets. Specifically, we expand the general scope by creating two novel suites: DSBio (90 bioinformatics tasks from academic literature probing domain-specific workflows) and DSPredict (92 Kaggle competitions spanning time series, computer vision, molecular property prediction, and single-cell perturbation). The next figure summarizes our creation process for these two suites:

DSGym construction pipeline. Left: Scientific tasks derived from bioinformatics papers via paper-reported and expert-derived queries (90 tasks). Right: Kaggle tasks from archived competitions, filtered and split by difficulty into Easy (38) and Hard (54) subsets. To support task execution and data generation, DSGym provides a data generation pipeline to execute queries and generating trajectories, turning the framework into a data factory that can effectively train models. Using this pipeline, we generated 3,700 synthetic queries. After applying LLM-based quality filtering, we obtained 2,000 high-quality query-trajectory pairs for supervised finetuning. Our results (presented next) demonstrate that these data can be an effective way to improve model performance on data science tasks, even for small models. Results We present here our main findings. Additional results are available in the paper. Addressing the memorization gap A first and important result concerns memorization. We observe that many existing benchmark queries provide weak signals: a non-trivial fraction remains solvable even without data file access, suggesting LLMs might have learned about these tasks during training.

We show that most models can answer data science questions without examining the data, instead relying on information memorized during training. This is problematic because it means models aren't actually performing data analysis. Thus, we made sure to flag and exclude these examples that are likely in the training sets of the models. DSGym applies quality filtering and prompt-only shortcut filtering to remove such tasks, producing refined datasets: DAEval-Verified, QRData-Verified, DABStep, and MLEBench-Lite. Benchmark performance & failure Mmodes After creating these new benchmarks, we test frontier proprietary and open-weight LLMs across general-purpose data science and domain-specific scientific tasks. Our trained 4B model ( Qwen3-4B-DSGym-SFT-2k ) achieves competitive performance with much larger models on general analysis benchmarks. 77.78 33.07 86.19

Interestingly, most models are still far from getting perfect scores on these benchmarks. To understand why models fail on these tasks, we conducted a manual error analysis of 50 randomly sampled failed trajectories per model and task family. This analysis reveals an interesting pattern: while general analysis tasks show diverse failure modes, with statistical knowledge gaps and planning errors being most common, scientific analysis tasks are dominated by a single failure…