TensorX

TensorX · Paper Blog

Scaling Automatic Research Agents via World Models

Xiyuan Yang (UIUC), Sheikh Sarwar (Amazon), Jingru Cheng (Amazon), Zhan Shi (Amazon), Duanshun Li (Amazon), Huiyuan Chen (Amazon), Haiyang Zhang (Amazon), Xing Fan (Amazon), Chenlei Guo (Amazon), Jingrui He (UIUC), Zhenyu Liao (Amazon) · 2026-09-10

Reinforcement LearningLLM AgentsWorld ModelsAutoResearchGRPOVLA

针对自动化科研智能体强化学习中环境沙盒真实执行带来的算力瓶颈,本文提出 WMRL 框架,利用世界模型模拟运行并结合在线去偏与逆方差降噪机制,实现 3 至 4 倍的训练加速与更优的泛化表现。

💡 核心创新点

自动化科研智能体在给定科研问题后,能够端到端完成构思、代码实验编写、分析日志与多轮迭代。然而,在通过强化学习(RL,通过策略梯度根据奖励信号优化模型参数)对科研智能体进行后训练时,学术界长期默认了一个基准假设:策略探索产生的大量代码必须放入完全隔离的真实执行沙盒中逐一运行以获取真实奖励信号。

本文彻底打破了这一高成本假设,明确指出:大规模强化学习并不需要完全依赖昂贵的物理沙盒运行,而可以通过语言模型构建的世界模型(World Model,以代码和任务为输入、纯静态推理预测执行结果与分数的代理环境)提供密集且高吞吐的模拟信号。相比将环境黑盒化的传统方案,本文提出的 WMRL(World Model RL,基于世界模型的强化学习)具有两项根本性机制创新:

  1. 双重校准的代理奖励框架:世界模型的预测并非完美无瑕,其奖励输出不可避免地伴随系统性偏置与随机噪声。WMRL 仅引入极小比例(约 10%)的真实环境执行作为锚点信号(Anchor Signal),创新性地设计了在线去偏(Online Debiasing,通过保序回归自适应拟合非线性映射校准漂移)与逆方差降噪(Inverse-Variance Denoising,依据残差自适应调控权重融合梯度)两级纠偏机制。
  2. 可证明的收敛性保障:从理论上证明了原生世界模型引入的二次偏置误差与方差膨胀,并在引入两级校正后严格消解了不可消除的偏置误差下界,使代理信号训练在渐近意义下收敛到与纯真实环境训练完全一致的最优解。

🚀 重大突破:WMRL 在训练计算成本大幅削减 3 至 4 倍的前提下,后训练所得的 4B 和 9B 参数规模科研智能体在测试基准上全面超越了真实沙盒训练的基线,甚至击败了规模达 48B 和 120B 的开源大模型!

🔍 背景与动机

自动科研智能体的强化学习困境

利用大语言模型构建自动科研智能体已在机器学习建模、数据科学竞赛以及跨学科实验探索中展现出巨大潜力。为了让智能体掌握在多轮交互中纠错和自我演进的能力,业界普遍采用基于组相对策略优化(GRPO,一种无需独立价值网络、在采样组内对奖励做标准化的策略梯度算法)的后训练机制。

然而,当我们将轨迹采样量推向大规模时,强化学习系统遭遇了严重的计算不对称性瓶颈。一个完整的科研交互轨迹包含两个阶段:智能体生成代码阶段与环境执行评估阶段。

  • 在生成阶段,现代大语言模型推理框架(如 vLLM 和 SGLang)依靠批处理技术使得成百上千条并发轨迹共享矩阵计算,增加单条轨迹的边际算力成本极低;
  • 在执行阶段,每个候选方案必须部署在包含 GPU 资源与完整数据依赖的隔离容器沙盒中,运行真实模型训练、计算损失与评估指标,执行时间长达数分钟乃至数十分钟。

因此,执行成本完全无法通过批处理进行摊销,总成本随轨迹数量线性激增。正如在日常科研中,构思实验假设只需要数秒钟,但每个实验在计算集群排队运行却需要耗费整夜。真实的沙盒执行使得强化学习训练集群的 GPU 绝大部分时间处于空闲等待或重度浪费状态。

自动科研智能体轨迹的不对称扩展瓶颈与 WMRL 解决方案

世界模型的引入与代价

为了破解这一不对称性瓶颈,一个自然的构想是:能否用一个前向推理速度极快的大语言模型充当“代码运行模拟器”?然而天下没有免费的午餐,直接使用世界模型评分会面临严重的奖励黑客与误差累积问题。如果世界模型存在偏置,智能体会学会投机取巧迎合模拟器,导致策略在真实环境中迅速退化。

如何以极轻量的真实执行为锚点,既享受世界模型带来的吞吐量提升,又消除模型瑕疵带来的优化偏差,正是本文研究的核心动机。

📊 数据集与实验环境

基准数据集构建与任务划分

评估工作涵盖了自动科研代码任务以及具身智能连续决策任务两类典型环境:

  • MLE-Dojo:源自真实的 Kaggle 机器学习工程竞赛基准,原作者剔除了原始 MLE-Bench 中下载失效的 2 个任务与环境不兼容任务,筛选出符合沙盒计算约束的 59 个竞赛任务。将其严格划分为 45 个训练任务(包括表格、文本和图像三大类别各 15 个)以及 14 个完全未见过的测试任务(包含 5 个表格、4 个文本和 5 个图像任务),且所有测试任务在数据规模上均大于训练集,专门用于检验从小型任务向大型真实竞赛的外推能力。
  • DSBench:数据科学建模评测基准,包含 60 个端到端运行的 Kaggle 真实任务(19 个二分类、10 个多分类、23 个回归任务以及 8 个非标准评估指标任务),与训练集完全无重叠。
  • LIBERO-Long:具身机器人操控基准,在机械臂长程操作任务中要求从图像观测与语言指令输出动作块。

实验硬件与超参数细节

  • 硬件配置:所有实验均在配置统一的 NVIDIA A100 GPU 集群上运行,真实沙盒中单个任务分配约 40 GB GPU 显存,单次执行设置 1200 秒的硬性超时上限。
  • 基础模型:科研智能体采用 Qwen3.5-4B 与 Qwen3.5-9B。世界模型不引入任何外部更强模型的隐式蒸馏,直接使用与智能体同源同尺寸的骨干模型(同为 Qwen3.5-4B 或 9B),通过结构化提示词进行静态代码审查与评估。
  • 强化学习超参:采用 GRPO 算法,优化器学习率设置为固定值 1e-6,KL 散度惩罚系数为 0.04,熵正则化系数为 0.002,采样温度 temperature=1.0,top_p=1.0。每个优化步骤包含 8 个组,每组包含 8 条独立探索轨迹(即 group size n=8),单条轨迹最多允许 4 轮交互,每轮生成上限 4096 tokens,观测截断为 1024 tokens。
  • 沙盒依赖:沙盒基于 Linux 系统的 Python 3.11 虚拟环境构建,预装 PyTorch 2.12 (CUDA 13.0)、scikit-learn 1.4.2、transformers 5.12.1 等 122 个科学计算与机器学习库。

🛠️ 核心机制与研究方法

整体系统架构与标准 RL 基础

在自动科研任务中,智能体策略在多轮迭代中生成解决方案代码,并在沙盒中运行后获得标量质量评分。标准的强化学习通过最大化预期得分来进行优化。在 GRPO 框架下,对于每一个任务采样一组相互独立的轨迹,其在真实环境中执行的评分为基准,经过组内均值中心化后得到轨迹优势,进而更新策略参数。策略梯度估计量由以下公式给出:

g^:=1n∑i=1nAisi,si:=∇θlog⁡πθ(τi)\hat{g} := \frac{1}{n} \sum_{i=1}^n A_i s_i, \quad s_i := \nabla_\theta \log \pi_\theta (\tau_i)

在公式 (1) 中,符号代表组内包含的样本轨迹总数,为当前轨迹相对于该组平均表现的优势函数值,而代表智能体策略在轨迹上的对数似然梯度向量。

当利用大语言模型构建的世界模型替代真实沙盒执行时,输入智能体生成的代码与任务说明,世界模型通过前向推断预测出执行结果并提取出代理分数。在无校正的纯世界模型设定下,组内优势函数被直接替换为:

A^i:=r^(τi)−1n∑jr^(τj)\hat{A}_i := \hat{r}(\tau_i) - \frac{1}{n} \sum_j \hat{r}(\tau_j)

在公式 (2) 中,代表世界模型对第条轨迹给出的质量打分,其计算过程与完全对齐。然而,世界模型的预测得分不可避免地偏离真实执行得分,可以分解为真实信号、系统偏置与零均值随机残差的三重叠加:

r^(τ)=r(τ)+b(τ)+ξ(τ)\hat{r}(\tau) = r(\tau) + b(\tau) + \xi(\tau)

在公式 (3) 中,为真实执行环境给出的金标准得分,代表无论如何平均都持续存在的系统性偏置函数,而代表均值为零、方差有界的独立随机噪声扰动。

WMRL 双阶段校准流水线:在线去偏与逆方差降噪

为了解决偏置与噪声对强化学习优化造成的不可逆损伤,WMRL 架构引入了以下三个核心模块:

模块一:世界模型多维度静态仿真器

  • 输入是什么:比赛的完整背景描述、真实数据文件目录规格、表格列名类型,以及当前轮次智能体编写的带行号 Python 脚本源码。
  • 它做了什么:世界模型不直接盲猜一个综合得分,而是将评测任务拆解为 8 个并行的单一维度自检专家提示,分别核查依赖库与变量命名、API 版本匹配度、数据路径与字段索引、张量形状与数据类型、计算显存预算、潜在运行时异常、提交文件规范以及代码算法质量。每一个维度输出标准 JSON 格式的结构化诊断信息。
  • 输出是什么:判定结果(通过或失败)、具体的错误代码行号、异常类名称,以及与真实 Linux 沙盒字节级高度相似的报错回溯信息,最终聚合为预测标量分数。
  • 为什么这么设计:大语言模型直接预测浮点数分数极易出现幻觉,但大模型擅长代码静态语法分析和运行时行为推演。通过将执行结果映射为结构化异常与格式校验,世界模型能够快速拦截 80% 以上的代码低级语法与环境接口错误。

模块二:基于保序回归的在线去偏机制

  • 输入是什么:训练过程中积累的锚点样本对集合,其中约 10% 的任务轨迹组同时送入真实沙盒与世界模型中运行,形成预测得分与真实得分的一一映射元组。
  • 它做了什么:世界模型的打分往往呈现非线性的整体虚高或过度保守。在线去偏模块假设偏置在分数尺度上满足单调变换,通过拟合一个单调非减的重校准映射函数来对齐两者的边缘分布:
f^=arg⁡min⁡f∑(r^j,rj)∈P(f(r^j)−rj)2\hat{f} = \arg \min_f \sum_{(\hat{r}_j, r_j) \in \mathcal{P}} (f(\hat{r}_j) - r_j)^2

在公式 (4) 中,代表受约束的单调递增函数空间,目标是通过保序回归(Isotonic Regression)最小化预测输出与真实分数之间的均方误差。系统在前 200 个锚点对积累完成前保持恒等映射,随后首次拟合并在每新增 64 个锚点对时重新动态更新。

  • 输出是什么:经过非线性单调映射矫正后的无偏代理得分。
  • 为什么这么设计:单调性保证了不会逆转同组代码之间的相对优劣排序,而动态保序回归能以非参数化的方式自适应追踪训练过程中策略漂移带来的世界模型误差分布变化。

模块三:基于残差审计的逆方差降噪融合

  • 输入是什么:来自真实执行沙盒锚点组的干净策略梯度估计,以及来自校正后世界模型组的含噪策略梯度估计。
  • 它做了什么:如何将小体量但无偏干净的真实梯度与大体量但存在随机扰动的世界模型梯度进行最优组合?根据最小方差无偏估计理论,两路相互独立的梯度流的最优线性加权系数应严格与各自的方差成反比:
g^=VE−1gE+VWM−1gWMVE−1∣GE∣+VWM−1∣GWM∣,gE:=∑k∈GEg^k,gWM:=∑k∈GWMg^k\hat{g} = \frac{V_E^{-1} g_E + V_{WM}^{-1} g_{WM}}{V_E^{-1} |G_E| + V_{WM}^{-1} |G_{WM}|}, \quad g_E := \sum_{k \in G_E} \hat{g}_k, \quad g_{WM} := \sum_{k \in G_{WM}} \hat{g}_k

在公式 (5) 中,代表真实执行组索引集合,代表世界模型预测组索引集合,与分别代表单组梯度在真实环境与代理模型下的协方差矩阵迹。根据理论推导,方差比例满足关系式,其中常数可通过预热阶段结束时的残差方差倒数直接标定,在线跟踪的均方残差则实时反映当前世界模型的预测离散度。

  • 输出是什么:方差严格低于任一单一流的最优合成梯度,用于更新智能体策略参数。
  • 为什么这么设计:该机制构建了一套动态的自适应回退机制。当智能体探索到世界模型完全无法理解的全新代码模式时,锚点上的预测残差急剧上升,权重自动向真实环境组倾斜,算法平滑退化为标准 GRPO,防止策略被含噪梯度误导。

📈 结果与深度分析

主基准表现与训练效率对比

在 MLE-Dojo 测试集与完全不相交的 DSBench 基准上的全面测试显示,WMRL 在大幅削减算力的同时取得了卓越的性能表现。

主要实验结果:在 MLE-Dojo 与 DSBench 上的排行榜百分位及训练算力对比

从数据对比中可以总结出以下关键结论:

  1. 大幅节省训练算力:在 Qwen3.5-4B 上,纯真实沙盒 GRPO 消耗了 883 GPU 小时,而 WMRL 仅需 286 GPU 小时,实现了 3.1× 的计算加速;在 Qwen3.5-9B 上,真实环境训练耗费 1174 GPU 小时,WMRL 仅消耗 349 GPU 小时,加速比达到 3.4×。
  2. 最终性能不降反升:WMRL 在 4B 尺度上在两个基准的平均百分位得分达到 16.4% 与 28.8%,相较于真实环境 GRPO 分别取得了 +1.2 和 +3.1 个百分点的净提升;在 9B 尺度上更是达到 21.6% 与 32.8%(平均分提升 +2.8 和 +1.6 个百分点)。
  3. 越级超越大参数模型:经过 WMRL 后训练的 4B 智能体在综合性能上全面击败了未经强化学习的 48B 参数开源模型 Kimi-48B-A3B(测试集均分仅 8.1% 和 17.3%),而 9B 智能体则全面超越了 120B 参数的 Nemotron-120B-A12B,充分证明针对任务环境进行高效强化学习后训练的效果远胜于单纯堆砌模型参数量。

跨领域泛化:具身智能操控任务迁移

为检验 WMRL 是否具有跨模态与跨任务的普适性,研究团队将该方法直接迁移到机械臂长程操控基准 LIBERO-Long 上。在这一场景中,策略网络直接输出连续控制动作块,世界模型采用预训练的多模态奖励模型 Robometer-4B 评估轨迹视频帧,真实仿真环境输出离散成功信号充当锚点。

具身智能操控任务 LIBERO-Long 上的 VLA 后训练成功率对比

实验结果表明,在没有做任何特定针对性修改的前提下,WMRL 将机械臂的总体平均成功率提升到了 41.2%,相较于 SFT 基线提升了 3.8 个百分点,相较于耗时极长的真实仿真环境 GRPO 提升了 2.9 个百分点。在完全未见过的测试初始状态下,WMRL 表现出最强的鲁棒性,证明双重校准框架在强化学习后训练中具备坚实的跨领域适用性。

消融实验与机制有效性验证

针对在线去偏(OD)与逆方差降噪(IVD)两项核心机制的消融实验进一步厘清了各模块的贡献度。

在线去偏与逆方差降噪模块的消融实验结果

消融数据揭示:

  • 若仅将真实环境信号与世界模型信号进行简单无校正混合(OD 与 IVD 均关闭),其在两个数据集上的得分均显著劣于纯真实沙盒训练;
  • 仅开启逆方差降噪时,平均分提升了 0.9 至 1.7 个百分点;
  • 仅开启在线去偏时,平均分提升了 2.2 至 2.8 个百分点,这与理论分析中“偏置项是常数阶误差下界、而方差项受学习率阻尼抑制”的数学推导完全吻合;
  • 当两者联合启用时,性能全面跃升 2.9 至 4.8 个百分点,证实了两个校准模块在去除不可控偏差与平抑高方差扰动上的互补特性。

局限性与未来展望

⚠️ 局限性:作者在讨论中坦承,当前世界模型的准确性依然高度依赖于智能体产物本身的信息显式性。如果某个科研任务的关键瓶颈在于极其隐蔽的随机性(例如多线程竞争条件或非确定性随机数种子引发的偶发数值溢出),在代码文本层面无法分析,世界模型在这些情况下的预测残差将持续偏大,导致系统不得不频繁回退至真实沙盒执行,从而削弱训练加速比。

未来的研究方向包括:探索智能体与世界模型的双向对抗演进,使世界模型能在智能体探索过程中同步进行在线参数更新;以及将世界模型扩展到包含真实物理仪器操作与生物湿实验的更广泛科研自动化闭环中。

北京市昌平区探索星信息技术及软件开发工作室

京ICP备2026059466号