Masked Diffusion LMs as Text World Models for Agentic RL

Lead

Agentic reinforcement learning needs more than clever reward design. As agents become stronger, fixed hand-built environments quickly stop providing useful signal, while sparse rewards over long horizons can push models toward narrow workflows or repeated tool-use patterns. This is why world models—models that simulate environment states and transitions—are becoming attractive as scalable training environments.

The paper asks whether text-based world models should remain primarily autoregressive. Autoregressive language models generate left to right, which is convenient for text completion but less ideal when a simulated environment must remain consistent with many global anchors at once: tool schemas, previous turns, domain rules, initial conditions, and desired outcomes. The authors propose that Masked Diffusion Language Models, or MDLMs, are a better fit because they can denoise bidirectionally while conditioning on these anchors throughout generation.

Key points

• A new formulation: The paper frames text world modeling as a steerable transition-dynamics problem. It decomposes the simulated state into initial state, task context, tool schemas, domain rules, and steering directives.
• Grounded trajectory data: The authors curate 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families.
• MDLMs versus autoregressive LMs: The comparison shows MDLMs achieving stronger coherence, groundedness, and empirically validated rollout diversity than autoregressive LLMs more than 4x their parameter size, while maintaining comparable inference latency.
• Training integration: The work introduces a plug-and-play GRPO training framework with deterministic state checks, aiming to reduce the damage caused by inconsistent simulated states.
• OOD transfer: In zero-shot ablations on ScienceWorld, ALFWorld, and AppWorld, across 1.2B-7B agent backbones including LFM2.5, Qwen3, and Mistral, agents trained with MDLM-simulated worlds show better gains than those trained with autoregressive world models. The paper reports up to 47% absolute gains over baselines without environment-specific fine-tuning.

Why it matters

The main contribution is not simply that a diffusion language model performs well on another benchmark. The larger claim is architectural: world modeling for agents may benefit from generation methods that are not locked into a left-to-right order. In an agent environment, the next state is not just a sentence; it is a constrained object that must respect tools, history, rules, and outcomes simultaneously.

Bidirectional denoising gives MDLMs a natural advantage for this kind of global consistency. Instead of committing early to local tokens and repairing contradictions later, the model can iteratively fill in the state under shared constraints. If this direction holds up, future agent RL pipelines may rely more on steerable simulated worlds that can generate diverse, validated training rollouts on demand.

There are still open questions. The findings are tied to the curated environments and trajectories in the study, and more work is needed to test whether the same benefits extend to messier real-world tool ecosystems, long-running applications, and higher-stakes domains. Still, the paper offers a concrete alternative to the default autoregressive world-modeling path.

Source: Hugging Face Daily Papers