论文

Masked Diffusion Language Models 是用于智能体 RL 的强且可操控的基于文本的世界模型

Masked Diffusion Language Models 是用于智能体 RL 的强且可操控的基于文本的世界模型

近期强化学习 (RL) 的发展催生了对多样化、专门化训练环境的需求。手工定制的固定难度环境随模型性能提升而失效,长跨度稀疏奖励则导致特定工作流或工具结构上的模式崩溃。能模拟环境状态的世界模型在纯 rollout 性能上与传统方法持平,有望按需扩展多样性。然而,自回归 (AR) 世界模型存在从左到右的偏差,无法对全局相互依赖的状态锚点(如工具架构、先前回合、预期结果)进行条件建模。 本文 (i) 将基于文本的世界建模形式化为可操控的转换动力学问题,分解为初始状态、任务上下文、工具架构、领域规则和操控指令;(ii) 构建包含 239,403 条 轨迹的数据集,覆盖九个开源环境和十二个前沿模型族。我们比较了 AR 语言模型和掩码扩散语言模型 (MDLMs),发现 MDLMs 通过双向锚点感知去噪,在连贯性、扎根性和 rollout 多样性上优于 4 倍参数量的 LLM,且推理延迟相当。 我们引入即插即用的 GRPO 训练框架,包含确定性状态检查,并在三个 OOD 环境 (ScienceWorld, ALFWorld, AppWorld) 上对三个 1.2B-7B 智能体主干 (LFM2.5, Qwen3, Mistral) 进行零样本迁移消融实验,取得高达 47% 的绝对增益(无需环境特定微调)。此外,我们进行了对抗场景下的失败模式行为分析和人类评估(真实感、结果正确性、训练效用)。项目已开源。

论文精读

TL;DR 掩码扩散语言模型(MDLM)作为文本世界模型,通过双向去噪克服自回归模型的左-右偏见,实现更连贯、多样且可操控的环境模拟,在智能体强化学习中零样本迁移性能提升高达47%。

问题

在智能体强化学习(RL)领域,训练环境多样性不足成为主要瓶颈。手工设计的环境任务和奖励难度固定,当智能体性能提升后,奖励信号变得稀疏,导致策略过度拟合特定工作流或工具结构,即模式崩溃(mode collapse)。

现有方法的局限

基于自回归语言模型(AR LM)的世界模型被用于模拟环境状态,但其单向生成机制存在根本性局限:从左到右的生成顺序强制一种因果偏置,使其无法有效条件化于全局相互依赖的状态锚点(如工具schema、对话历史、预期结果)。这种偏置会导致生成的轨迹在逻辑上断裂、与现实工具定义不一致,或多样性受限。即使增大模型参数量,AR世界模型在需要全局上下文一致性约束的场景中仍表现不佳。

技术挑战与行业重要性

文本世界建模的核心难题在于:如何在严格遵循领域规则、工具模式、任务上下文等全局锚点的同时,生成连贯、多样化且可操控的交互轨迹。这需要模型具备双向条件生成能力,但传统AR模型难以满足。随着LLM智能体应用(如代码助手、多步API调用智能体)快速增长,RL训练需要大规模、高保真且可定制的模拟环境,工业界对可扩展的世界模型需求迫切。能够按需生成多样化场景的世界模型能大幅降低环境构建成本,提升智能体的泛化能力。

类似自动驾驶仿真中需要生成丰富多样的交通场景来训练规划模型,文本世界模型若能可控地生成多种交互分支,将成为智能体RL训练的关键基础设施。

核心洞察

  • **MDLM 双向锚点感知** 解决了自回归世界模型的左到右偏置,使生成过程可被全局状态(工具模式、历史回合、预期结果)操控,从而提升长序列状态推演的连贯性与多样性。与 AR world model 相比,MDLM 能同时关注完整上下文中的多个锚点,在 sparse reward、长 horizon 的 agentic RL 场景下有效避免模式坍塌,并生成更接近真实环境的 rollout 变体。
  • **小参数 MDLM 超越大 4 倍的 LLM**,在参数效率与推理成本上为 world modeling 提供了新选择。实验表明,1.2B-7B 规模的 MDLM 在 groundedness 和 rollout 多样性指标上优于同规模甚至 4 倍参数量的 AR 模型,且推理延迟相当。这一发现挑战了“world modeling 必须依赖大规模自回归模型”的惯常假设,为资源受限下的高质量环境模拟提供了可行路径,同时借助条件掩码生成轻松实现按需操控,降低了构建可定制训练环境的门槛。

方法

形式化世界建模目标

本文将基于文本的世界建模定义为可操控的转移动态问题,输入被显式分解为五个组件:

  • 初始状态 s0:环境的起始文本描述;
  • 任务上下文 c:目标、指令或对话历史;
  • 工具模式 T:可用工具的函数签名与约束;
  • 领域规则 R:环境隐含的逻辑规则;
  • 操控指令 z:希望世界模型遵循的高层指引(如多样性、场景偏差)。

世界模型的任务是,在给定当前状态 st、动作 at 及上述全局信息后,预测下一状态 st+1

模型架构:掩码扩散语言模型

核心生成器采用掩码扩散语言模型 (MDLM),而非自回归 (AR) 模型。MDLM 通过迭代去噪方式生成文本:

  • 输入序列以一定比例随机掩码,模型学习恢复被掩码的 token;
  • 训练目标为 masked token 的交叉熵损失,训练时掩码比例动态调度(如余弦调度);
  • 推理时从全掩码状态开始,逐步替换低置信度 token,经过若干步去噪得到最终状态。

MDLM 的双向上下文表示能力使其能够同时关注全局锚点(如工具 schema、历史回合、期望结果),避免 AR 模型“从左到右”的偏差,这在需要前后文相互依赖的状态生成中至关重要。

操控性生成

通过条件掩码生成实现可操控性:将操控指令 z 编码为文本前缀或特殊 token,与状态、动作拼接后送入模型。模型在去噪过程中受该指令引导,从而控制生成风格、多样性或特定场景倾向。实验表明,MDLM 能在不牺牲连贯性的前提下大幅提高 rollout 的多样性。

训练框架:GRPO + 确定性状态检查

为进一步提高模型在序列决策中的稳定性,引入插件式 GRPO(Group Relative Policy Optimization)训练框架

  • 基于 MDLM 的 rollout 被送入 RL 智能体,在模拟环境中执行并收集实际状态;
  • 通过确定性状态检查(如关键实体匹配、约束满足)对生成状态进行验证,提供可解释的奖励信号;
  • 使用 GRPO 微调世界模型,使其生成的状态不仅语言连贯,且符合环境真实动态。

与同类方法的差异

与主流的 AR 世界模型不同,MDLM 凭借双向注意力和迭代去噪,天然适合捕捉全局依赖,无需复杂的 prompt 工程或后处理来修正左到右偏差;同时 GRPO 训练和状态检查机制直接优化了世界模型在下游 RL 训练中的效用,而非仅关注生成文本的困惑度。

实验

实验设计

实验围绕文本世界模型的转向转移动力学形式化问题展开,将建模任务分解为初始状态、任务上下文、工具模式、领域规则与转向指令。基于此框架,收集了涵盖 9 个开源环境、12 个前沿模型家族的 239,403 条有状态‑动作轨迹,构建大规模训练集。

对比目标为自回归语言模型 (AR LM)掩码扩散语言模型 (MDLM)。训练采用即插即用的 GRPO 训练框架,内置确定性状态检查,以增强生成轨迹的可靠性。下游评估在 ScienceWorldALFWorldAppWorld 三个分布外环境上执行零样本迁移,测试 LFM2.5Qwen3Mistral 三种 1.2B–7B 规模智能体骨干。评估指标涵盖连贯性、接地性、多样性,并辅以人类评估与对抗场景下的失败模式分析。

关键发现

  • MDLM 通过双向锚感知去噪,在文本世界建模的连贯性、接地性与 rollout 多样性上全面超越 AR LM,即使后者参数量为其 4 倍,且推理延迟相当。
  • 零样本迁移结果显示,MDLM 世界模型使下游智能体在三个 OOD 环境上获得最高 47% 的绝对成功率提升,无需任何环境特定微调。
  • 定性分析表明,AR 模型的左到右生成偏差导致其对工具模式、先前轮次等全局状态锚点的条件化能力不足,而 MDLM 的 mask‑and‑recover 机制可自然融入此类约束。

与基线对比的深度解读

AR 世界模型受限于从左到右的自回归范式,难以在生成过程中回溯并协调全局依赖信息,尤其在长轨迹、多工具交互场景下容易发生模式坍塌。MDLM 通过将状态‑动作序列的部分位置 mask 并由双向上下文重建,实现了任意位置的条件化与全局一致性建模。实验证实,这种机制不仅提升了单条轨迹的准确性与多样性,还使训练出的世界模型具备更强的鲁棒性与泛化能力,能够跨环境、跨智能体骨干迁移,且计算效率不逊于 4 倍参数量的 AR 模型。GRPO 框架引入的确定性状态检查进一步降低了幻觉累积,为 RL 训练提供了更可靠的模拟信号。

行业影响

落地场景

基于 Masked Diffusion Language Models (MDLMs) 的 steerable 文本世界模型,可直接用于强化学习 agent 的训练环境生成,尤其适合需要长程、稀疏奖励的任务(如电商导购对话、企业流程自动化、游戏 NPC 行为学习)。它能够按需生成多样化、可干预的状态轨迹,缓解传统手写环境导致的模式坍塌过拟合,并支持零样本迁移到新领域(如 ScienceWorld、ALFWorld、AppWorld),显著降低环境适配成本。

商业价值

  • 降本:替代大量人工构建训练场景的工作量,一次训练即可生成无限变体,减少数据采集与标注支出。
  • 增效:通过更丰富的状态分布覆盖,提升 agent 在真实部署中的成功率与泛化能力,缩短从研发到上线的迭代周期。
  • 体验提升:使对话或操作型 agent 在面对突发请求(如用户临时更改意图、工具异常)时表现更稳定,直接改善客户满意度。

与现有工作流的集成

该方案以 plug-and-play GRPO 框架提供,可无缝嵌入现有 RL 管线:在 agent 与环境交互阶段,用 MDLM world model 替代真实环境或作为辅助采样器;训练时结合确定性状态校验保证轨迹可靠性。部署上可通过轻量 API 调用,与 LangChain、AutoGPT 等 agent 框架协同,用户只需提供初始状态 schema 和 steering 指令即可生成定制化 rollout。

具体落地用例

  • 电商智能客服:world model 动态生成不同用户画像(意图摇摆、模糊描述)与库存/物流变化,训练客服 agent 处理退换货、多轮议价等复杂流程。
  • 业务流程自动化(RPA):模拟 ERP 系统状态转移,包括权限变更、异常报错等,让 agent 学会自适应地调用数百个 API 组合,而无须维护庞大的灰盒测试环境。

局限

  • **训练数据规模与领域覆盖有限**。数据集包含 239,403 条轨迹,覆盖 9 个开源环境和 12 个模型家族,但轨迹均由现有语言模型生成,可能继承其内在偏差,且未充分包含真实环境中的噪声、歧义与长尾场景。在 ScienceWorld、ALFWorld、AppWorld 等 OOD 环境上的 zero-shot 评估虽取得绝对值提升,但整体成功率仍较低,表明世界模型在完全未见任务类型上的泛化性能受限于训练数据的多样性。扩大数据规模和纳入更多样化的真实交互轨迹将是提升鲁棒性的关键方向。
  • **对显式 steering 组件的依赖**。方法将世界模型分解为初始状态、任务上下文、工具模式、领域规则和操控指令等组件,并通过条件掩码生成实现操控。这要求为每个新环境预先定义结构化的 steering 组件,引入领域专家成本,且在环境动态变化或缺乏明确结构化描述(如开放域对话或弱规则游戏)时,操控能力可能衰减,限制了方法在高度非结构化场景中的直接应用。
  • **扩散模型训练与采样效率**。尽管论文指出推理延迟与 AR 模型可比,但 Masked Diffusion 模型的多步去噪过程在迭代数较大时计算开销较高,尤其在线训练中需要大批量快速生成世界模型 rollout 时,可能成为瓶颈。此外,扩散调度器和混合顺序采样器的调优对最终生成质量敏感,未充分讨论长序列降噪步数的选择、与 GRPO 在线微调的计算开销折衷,以及更大模型(如 13B+)的扩展效率。
论文Darshan Deshpande2026-05-07原文

相关内容