DSWorld:面向高效自主智能体的数据科学世界模型 DSWorld: A Data Science World Model for Efficient Autonomous Agents
用世界模型预测数据操作效果,智能体训练加速14倍、推理加速3-6倍
前置知识
世界模型(World Model)
世界模型通过学习环境的动力学,在给定当前状态和动作的条件下预测下一个状态,从而支持预测、控制和仿真。视觉/物理世界模型用于生成未来视频或 3D 场景,数字世界模型用于网页/GUI 智能体的状态转移。其核心价值是让智能体在脑中想象动作后果,避免昂贵的真实试错。形式化记作 $S_{t+1} = W(S_t, A_t)$,其中 $W$ 即世界模型。
本文的核心贡献是把世界模型从物理/数字世界推广到数据科学执行环境,理解世界模型的基本范式才能判断其创新点和局限。
自主数据科学智能体(Autonomous Data Science Agent)
这类智能体自动完成数据科学任务,从探索性数据分析到预测建模,自主进行特征工程、模型训练、评估与工作流更新。代表系统如 ML-Master、AIDE、AutoMLGen、AutoMind,它们多用测试时扩展策略(搜索大量候选解)或强化学习训练骨干模型来提升表现。ML-Master 在 MLE-Bench 上有 56.4% 的任务达到奖牌级别。
本文要解决的就是这类智能体执行成本过高的痛点,理解它们的工作方式才能明白世界模型为什么能加速。
测试时扩展与搜索式推理(Test-time Scaling)
测试时扩展指在推理阶段通过探索大量候选方案(如树搜索、进化算法、采样多条轨迹)来提升答案质量。AIDE、ML-Master、AutoMLGen 等智能体在代码解空间中搜索并选优。代价是每个候选都要真实执行代码、训练模型、跑评估,导致推理时间随搜索规模急剧增长,绝大多数时间耗在计算而非推理上。
这正是 DSWorld 推理加速(3-6×)要切入的场景:用世界模型模拟代替昂贵真实执行,从而大幅压缩搜索成本。
GRPO 强化学习(Group Relative Policy Optimization)
GRPO 是 DeepSeek 提出的策略优化算法,用一组 rollout 的相对优势替代独立价值函数来降低估计方差。优势定义为组内归一化 $A_i = (R_i - \mu(R))/(\sigma(R) + \epsilon)$,配合带裁剪的策略比 $L_{clip}$ 和对参考策略的 KL 惩罚 $-\beta_{KL}D_{KL}(\pi_\theta\|\pi_{ref})$ 来稳定训练。本文用它来优化世界模型的转移预测质量。
理解 GRPO 才能看懂反思式世界模型优化如何用 RL 信号逐步逼近正确的下一状态预测。
状态转移预测(State Transition Prediction)
状态转移预测指给定当前状态 $S_t$ 和动作 $A_t$,预测下一状态 $S_{t+1}$。在数据科学语境下,状态包括任务、数据集统计与预览、执行环境、执行历史与中间输出、错误信息;动作是特征工程、模型训练、评估等操作。预测难点在于既要判断执行成败、错误类型,又要预测输出文本与下游 ML 性能。
转移预测是数据科学世界模型的直接建模目标,也是论文评测(ESP/ETP/ERS/EKM/PP/PR)的对象。
研究动机
自主数据科学智能体近年来在探索性数据分析、预测建模等任务上展现出强大能力,例如 ML-Master 在 MLE-Bench 上有 56.4% 的竞赛任务达到奖牌级别。然而现有方法普遍依赖测试时扩展策略,通过迭代试错探索大量候选方案。这种范式导致执行时间几乎全部耗费在昂贵的数据处理、模型训练、评估与工作流更新上,而非智能体推理本身。以 ML-Master 为例,它在 MLE-Bench 中超过 86% 的执行时间都花在模型训练上。如此沉重的计算开销严重制约了自主数据科学系统的效率与可扩展性,使得大规模强化学习训练和高频搜索式推理都变得代价高昂乃至不可行。
本文的目标是本文旨在为数据科学工作流构建一个状态转移预测模型(即数据科学世界模型),使智能体能够在执行昂贵的真实计算之前,就预判候选操作(特征工程、模型训练、评估等)对环境状态的影响,包括数据集变化、模型变化、执行反馈、错误信息和性能信号。具体可量化的目标是:在不显著损失下游任务性能的前提下,把基于强化学习的智能体训练加速约 14 倍,把基于搜索的智能体推理加速约 3-6 倍,并在数据科学转移预测任务上以较大优势(约 35.6%)超越最强的通用大模型基线。
与已有工作不同的是,本文的独特切入角度在于把'世界模型'这一概念从视觉/物理世界(想象未来画面)迁移到数据科学执行环境:把代码执行环境本身视为待建模的世界,给定当前工作流状态与候选操作,直接预测下一个环境状态。在此之前的同类工作主要面向网页或 GUI 智能体的状态转移,而数据科学环境的世界建模几乎是空白。此外,作者没有走'用一个 LLM 预测一切'的极端,而是设计了代价感知的路由机制——廉价操作真实执行、昂贵操作交给模拟器——从而在精度与效率之间找到了关键平衡点。
核心方法
DSWorld 的整体思路可以类比下棋引擎:在真正落子前先在脑中模拟若干步,估算每步的后果,从而避免昂贵的真实试错。技术上,给定一个数据科学工作流的原始执行环境,State Constructor 先用规则程序把它抽取成结构化状态表示 $S_t = SC(E_t)$,包含任务描述、数据集统计与预览、执行环境、执行历史、中间输出和错误信息。智能体据此生成动作 $A_t = \pi(S_t)$,动作被编码成稠密向量后送入 Router 判断该操作是否昂贵:轻量操作(简单数据查看、环境检查)直接交给 Compiler 真实执行以保证精度,重型操作(大规模模型训练)则交给 LLM 模拟器预测下一状态以节省算力。为防止路由失误,Compiler 设有超时阈值,超时则回退到模拟器,整体转移可写成分段公式 $\hat{S}_{t+1} = C(S_t, A_t)$(执行)或 $\hat{S}_{t+1} = S(S_t, A_t)$(模拟/超时)。
核心创新点是'代价感知的混合执行-模拟'框架与'反思式世界模型优化'的结合。与已有方法最本质的区别在于:现有数据科学智能体要么全部依赖真实执行(慢但准,如 Compiler),要么用通用大模型直接预测一切(快但极易幻觉,例如直接把 DeepSeek 3.2 当模拟器会导致性能严重退化),而 DSWorld 通过 Router 把两类操作的优点拼到一起——该跑真的就跑真的,该想象就交给微调过的 LLM 模拟器。在此之上,Reflective World Model Optimization 不再只做一次预测,而是先预测 $\hat{S}' \sim S_\theta(\cdot|S,A)$,再与真实下一状态对比生成错误反思 $f = S_\theta(\hat{S}', S')$,并据此自我修正 $\hat{S}_r = S_\theta(S, A, f)$,让模型在迭代中逼近正确转移。
方法步骤详情
DSWorld 分数据构造、SFT 预热、反思式 RL 三步。第一步构造 DSWorld-8K:让现成智能体在真实任务上运行收集 $(S, A, S')$,并从 MMTU 的 6 万余张真实表中采样状态 $S=SC(E)$,再基于 NumPy/Pandas 操作库与错误库采样操作 $o$、错误 $e$、成败 $r$,由 LLM 生成动作 $A \sim p_\eta(A|S,o,e,r)$,真实执行得 $S'=C(S,A)$,仅保留通过校验并补写 CoT,约 8K 条样本。第二步 SFT 预热目标 $L_{SFT}=-\log S_\theta(S'|S,A)$,5 个 epoch、batch 32、学习率 $1\times10^{-5}$。第三步反思式优化:对每个样本做 $n$ 次 rollout 得原始与修正预测,用 GRPO 联合优化(含 $L_{clip}$ 与 KL 惩罚),rollout 8、学习率 $1\times10^{-6}$、200 步,4 张 A800 上用 VeRL 完成。
技术新颖性
技术新颖性体现在三方面。其一,首次提出'数据科学世界模型'这一概念,把代码执行环境当作世界来建模,与物理/视觉世界模型(生成视频、3D 场景)和数字世界模型(网页/GUI 转移)形成互补,填补了该方向的研究空白。其二,代价感知路由加超时回退的混合执行机制,避免了'全真执行'的低效和'全模拟'的幻觉,是精度与效率的关键折中。其三,反思式世界模型优化把'预测-对比-反思-修正'引入 GRPO,让 LLM 在错误信号中迭代逼近真实转移。配合可扩展的数据合成流水线(MMTU 大规模真实表加操作/错误库再真实执行校验),系统性地解决了真实转移数据稀缺的瓶颈。
实验结果
RQ1 转移预测:Table 1 中 DSWorld 平均 0.781,比最强通用基线 o4-mini(0.576)高 35.6%;执行类 ESP/ETP/ERS/EKM 较最强基线分别提升 33.4%/57.6%/71.5%/50.5%(ESP 0.950、ETP 0.922、ERS 0.871、EKM 0.575),PP 0.856 取最佳、PR 0.518 偏弱。RQ2 训练:Table 2 中 DSWorld 训练的智能体在 MLE-Bench Lite 得 17.67/277 分钟,接近 Compiler(18.11/335),远优于 DeepSeek 3.2 模拟器(10.86/3854),约 14× 训练加速。RQ2 推理:Table 3 中 DSWorld 作执行器比 Compiler 加速约 3-6×(AIDE+Qwen3-8B 从 4102s 降到 676s)且性能基本保持,用 DeepSeek 3.2 则性能退化。RQ3 消融:SFT 较骨干提升 37.5%,GRPO 加 1.05%,反思优化再加 1.3%;RQ4 显示随数据与模型规模增大性能稳步上升。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数据科学转移预测综合 | AVG(六项任务平均) | 0.781 | o4-mini 0.576(最强通用 LLM) | 平均提升 35.6% |
| 执行错误类型预测(ETP) | 准确率 | 0.922 | Qwen3-8B 原始 0.573 | 提升 57.6%(相对最强基线) |
| RL 智能体训练效率(MLE-Bench Lite) | 训练时间(min) | 277 | Compiler 335 / DeepSeek 3.2 模拟器 3854 | 较 DeepSeek 3.2 模拟器约 14× 加速 |
| 搜索式推理效率(AIDE, Qwen3-8B) | 推理时间(s) | 676 | Compiler 4102 | 约 6× 加速且性能持平 |
| DACode 智能体训练 | DACode 得分 | 0.232(Qwen3-8B 骨干) | Qwen3-14B 基线 0.214 | 用更小骨干反超 8.4% |
局限与改进
作者承认三点局限:其一,DSWorld 当前只建模数据科学内部的转移,并未显式建模智能体工作流中对外部工具调用的转移;其二,预测质量受底层 LLM 模拟器能力限制,在复杂工作流中偶尔会产生不准预测;其三,合成轨迹与真实自主工作流之间仍存在分布差距,可能影响泛化。此外我观察到,PR(性能排序)和 PP(性能预测)的提升远小于执行类任务(PR 仅 0.518,仍是难点),说明对'给定代码预测最终 ML 指标'这类强推理任务,世界模型还力不从心;同时 MLE-Bench Lite 仅含 21 个任务,奖牌率绝对值很低(Gold 约 1.59%),方差较大,统计意义有限。
独立分析的弱点
第一个弱点是性能/排序类预测仍是短板。PR 与 PP 需要对 ML 算法、任务特征与评估指标的深度推理,单纯学执行轨迹难以胜任。改进方向是引入更强的反思信号,或把性能预测解耦为专门的子模型,借助历史竞赛结果做对比学习。第二个弱点是合成数据的分布偏差:操作/错误库主要覆盖 NumPy/Pandas 生态,真实数据科学还涉及 Spark、SQL、深度学习框架等,合成轨迹未必覆盖。改进方向是扩展操作库到更多框架,并引入真实 Kaggle 解题轨迹做风格校准。第三个弱点是路由决策依赖启发式阈值与小型 MLP,在边界操作(中等开销)上可能误判,导致要么白白执行、要么错误地模拟。改进方向是把路由建模为带开销感知的强化学习子问题,或用学习到的开销估计器替代固定超时。第四个弱点是评估规模偏小,21 个 MLE-Bench 任务难以支撑强结论,改进方向是扩展到完整 MLE-Bench 或更大数据科学基准。
未来方向
作者提出的未来方向包括:显式建模外部工具调用转移、用更强的底层 LLM 提升模拟器上限、缩小合成与真实轨迹的分布差距。基于本成果可延伸的方向包括:把数据科学世界模型与其他数字世界模型(网页/GUI)统一,构建跨环境的通用智能体世界模型;将代价感知路由思想推广到更多昂贵执行场景,如机器人规划、科学实验模拟;研究世界模型自身的自我改进闭环——用模拟器生成更多轨迹再反哺训练;探索如何把性能预测的反思信号用作奖励模型,直接优化智能体策略。
复现评估
论文提供了匿名代码仓库(https://anonymous.4open.science/r/DSWorld),但尚未正式开源。训练用 Qwen3-8B 作骨干、Harrier OSS v1 0.6B 作动作编码器、Router 为两层 MLP(隐层 256/64),数据合成用 DeepSeek 3.2;SFT 5 个 epoch、batch 32、学习率 $1\times10^{-5}$,RL rollout 8、学习率 $1\times10^{-6}$、200 步,4 张 A800 上用 VeRL 完成,硬件门槛中等偏上。DSWorld-8K 数据构造依赖 MMTU(6 万+真实表)、DABench、MLE-Dojo 等外部数据源,部分需自行获取;Router 训练所用的'代码-执行时间'对的收集细节披露较少。整体复现可行但工程量较大,主要难点在数据合成流水线与多基准评测体系的搭建。
论文图表
该图用三张子图给出全文的直觉。(a) 把视觉世界模型(想象物理世界未来画面)与数据科学世界模型(不执行昂贵计算就预测数据科学操作的效果)做类比;(b) 展示 DSWorld 把基于强化学习的智能体训练加速约 14×;(c) 展示它把基于搜索的智能体推理加速约 3-6×,三条时间/得分曲线对比了 Compiler、DSWorld 与 DS3.2 的代价。
这张图一图概括了论文的核心动机和三大卖点(概念、训练加速、推理加速),是快速建立直觉的关键入口。