Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL
cs.AI, cs.LG
Submitted: 2026-05-07
Updated: 2026-10-07
Comments: Accepted to NeurIPS 2026 Dataset: https://huggingface.co/PatronusAI/world_model_corpus Training code: https://github.com/patronus-ai/mdlm_world_modeling
Code: https://github.com/patronus-ai/mdlm_world_modeling
License: http://creativecommons.org/licenses/by/4.0/
The gist: Recent growth in reinforcement learning (RL) has surfaced a need for diverse, specialized training environments.
Terminology
Abstract
Recent growth in reinforcement learning (RL) has surfaced a need for diverse, specialized training environments. Hand-curated environments with fixed task and reward difficulties become ineffective signals as model performance improves, and sparse rewards over long horizons induce mode collapse on specific workflows or tool structures. World models that simulate environment states have matched pure rollout performance, making them promising for scaling diversity on-demand. However, autoregressive (AR) world models suffer from a left-to-right bias preventing conditioning on globally interdependent state anchors such as tool schemas, prior turns, and expected outcomes. We (i) formalize text-based world modeling as a steerable transition-dynamics problem decomposed into initial state, task context, tool schemas, domain rules, and steering directives, and (ii) curate 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families. We compare AR LMs and masked diffusion language models (MDLMs), showing MDLMs, via bidirectional anchor-aware denoising, achieve better coherence, groundedness, and empirically validated rollout diversity than LLMs over 4x their parameter size, at comparable inference latency. We introduce a plug-and-play GRPO training framework with deterministic state checks, and perform zero-shot transfer ablations on three OOD environments (ScienceWorld, ALFWorld, AppWorld) across three 1.2B-7B agent backbones (LFM2.5, Qwen3, Mistral), achieving up to 47% absolute gains over baselines without environment-specific fine-tuning. We further conduct behavioral analysis of failure modes under adversarial scenarios and human evaluation on realism, outcome correctness, and training utility. We open-source our work to encourage research in this direction.
Sources
- LFM2 Technical Report
- Large Language Models Hallucination: A Comprehensive Survey
- Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models
- AbideGym: Turning Static RL Worlds into Adaptive Challenges
- From Masks to Worlds: A Hitchhiker's Guide to World Models
- The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A"
- LLaDA2.0: Scaling Up Diffusion Language Models to 100B
- LLaDA2.1: Speeding Up Text Diffusion via Token Editing
- Scaling Agent Learning via Experience Synthesis
- SDAR: A Synergistic Diffusion-AutoRegression Paradigm for Scalable Sequence Generation
- Revisiting Reinforcement Learning for LLM Reasoning from A Cross-Domain Perspective
- SFT Memorizes, RL Generalizes: A Comparative Study of Foundation Model Post-training
- Process Reinforcement through Implicit Rewards
- Agent-RLVR: Training Software Engineering Agents via Guidance and Environment Rewards
- Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning
- A Survey on Code Generation with LLM-based Agents
- Efficient-DLM: From Autoregressive to Diffusion Language Models, and Beyond in Speed
- Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding
- NExtLong: Toward Effective Long-Context Training without Long Documents
- Scaling Synthetic Data Creation with 1,000,000,000 Personas
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems