Skip to content

Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL

🕒 Published (v1): 2026-05-07 00:00 UTC · Source: HuggingFace · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Masked diffusion language models (MDLMs) at 8B parameters outperform fine-tuned autoregressive LLMs up to 35B parameters as text-based world models for agent RL, because bidirectional anchor-aware denoising avoids the left-to-right causal bias that causes globally incoherent state predictions. The authors curate 239,403 grounded state–action trajectories, then show that GRPO training on MDLM-generated rollouts yields up to 47% absolute task-success gains on three out-of-distribution environments without any environment-specific fine-tuning.

Problem

AR LLMs used as world models are architecturally constrained to condition on preceding tokens only. In structured tool-call outputs, anchors such as trailing status fields, tool schemas, and expected outcomes are globally interdependent with earlier fields, so causal decoding produces prefix-consistent but globally incoherent states (e.g., refund_processed: true with status: "error"). Additionally, AR sampling suffers from prefix mode collapse—high-probability opening tokens are committed early and constrain all subsequent decisions—reducing the diversity of imagined rollouts precisely when agents need the most varied training signal.

Method

The world model is formalized as learning the transition distribution:

\[p_\theta(c_{t+1} \mid e_0,\, c_{\leq t},\, a_{\leq t};\; \mathcal{T}, \mathcal{R}, \mathcal{D})\]

where each state is decomposed into initial environment state \(e_0\), task context \(c_t\), tool schemas \(\mathcal{T}\), domain rules \(\mathcal{R}\), and steering directives \(\mathcal{D}\) (fixed across a trajectory). MDLMs learn this distribution via a masked diffusion objective: the forward process independently masks each next-state token with probability \(\tau \sim \mathcal{U}(0,1]\), and the model recovers masked tokens from full bidirectional context:

\[\mathcal{L}_{\text{MDLM}}(\theta) = -\mathbb{E}\!\left[\frac{1}{L}\sum_{k=1}^{L} \mathbf{1}\!\left[x_{k,\tau}^{(t+1)} = [\text{M}]\right] \log p_\theta\!\left(x_k^{(t+1)} \mid x_\tau^{(t+1)}, x^{(\leq t)}, a^{(\leq t)}\right)\right]\]

Steering is implemented by leaving known anchor positions unmasked at inference time; the model fills in remaining positions while attending to anchors on both sides. A dataset of 239,403 trajectories is curated from nine open-source environments (SWE-Smith, CoderForge, TAU-Bench, WebShop, Gorilla/BFCLV4, etc.) using 12 frontier model families. For downstream RL, a plug-and-play GRPO training framework with deterministic state checks is used; MDLM-generated rollouts substitute for real environment interaction.

Key Contributions

  • Formal decomposition of text-based world modeling into five steerable components (\(e_0, c_t, \mathcal{T}, \mathcal{R}, \mathcal{D}\)) enabling fine-grained guidance of imagined rollouts.
  • World model corpus of 239,403 grounded state–action trajectories across nine environments and twelve frontier model families (open-sourced on HuggingFace).
  • Comparative study demonstrating that the causal inductive bias, not model capacity, is the primary bottleneck: 8B MDLMs (SDAR-8B, WeDLM-8B, LLaDA-2.1-mini) consistently outperform fine-tuned AR models up to 35B dense / MoE parameters.
  • Plug-and-play GRPO framework integrating MDLM world models with deterministic state verification for stable RL training.
  • Zero-shot OOD transfer ablations on ScienceWorld, ALFWorld, and AppWorld across LFM2.5 (1.2B), Qwen3, and Mistral (up to 7B) agent backbones.
  • Behavioral failure-mode analysis under adversarial scenarios and a human evaluation of realism, outcome correctness, and training utility.

Results

  • Generation fidelity (MAUVE, in-domain): SDAR-30B-A3B achieves best zero-shot MAUVE on 3 of 4 evaluation splits; gap most pronounced on API-Bank (.697 MDLM vs. .532 best AR baseline GLM-4.7-Flash zero-shot).
  • 3-shot performance: SDAR-8B reaches MAUVE 0.990 (in-domain) and 0.981 (Intercode SQL), and SDAR-30B-A3B 0.995 / 0.983 respectively.
  • Diversity (Table 2): SDAR-8B Self-BLEU 0.601 vs. best AR (Qwen3.5-27B 0.670; GPT-OSS-20B 0.659); Distinct-N 0.385 vs. 0.321 / 0.228 for AR baselines; MAUVE 0.902 vs. 0.889/0.867.
  • Downstream RL (GRPO, OOD): Up to 47% absolute task-success-rate improvement over raw baselines on ScienceWorld/ALFWorld/AppWorld, zero-shot, without environment-specific fine-tuning, across 1.2B–7B agent backbones.
  • Human evaluation (1–5 Likert, 4 annotators): Training utility 4.75, realism 4.25, outcome correctness 4.50; inter-annotator agreement \(\alpha \geq 0.89\).

Limitations

  • MDLMs are computationally expensive to train; the paper relies on existing pre-trained checkpoints (SDAR, WeDLM, LLaDA) rather than training from scratch.
  • The human evaluation uses only four annotators, limiting statistical power despite high inter-annotator agreement.
  • The 47% gain is the upper bound across settings; per-environment breakdown and the variance across seeds is not shown in the truncated text.
  • Failure-mode analysis is confined to adversarial scenarios described in an appendix; systematic coverage of in-the-wild failure types is limited.
  • The world model represents the full environment state \(e_{\leq t}\) implicitly (not emitted as output), which may limit interpretability and auditability of latent state tracking.
  • OOD environments (ScienceWorld, ALFWorld, AppWorld) cover text-game and household tasks; generalization to GUI/browser agents or code execution environments is not established.

Relevance to Agentic AI / LLM Agents

This work directly addresses the scalability and rollout-diversity bottleneck in agentic RL training—a foundational open problem as LLM agents improve faster than hand-curated environments can be updated. The anchor-aware generation paradigm is particularly important for tool-use agents whose structured outputs (JSON tool responses, API schemas) encode interdependent fields that AR decoding corrupts. The plug-and-play GRPO integration means any agent backbone can be trained on MDLM-generated imagined rollouts without environment-specific infrastructure, lowering the barrier to curriculum generation at scale. More broadly, the finding that architectural inductive bias (causal vs. bidirectional) matters more than model scale challenges the prevailing assumption that larger AR LLMs are always better world models.