Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL
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.
Community
Can masked diffusion LMs, by denoising bidirectionally, fix the left-to-right bias that keeps autoregressive world models from staying consistent with global anchors (tool schemas, prior turns, expected outcomes)? This paper shows that MDLMs beat autoregressive models over 4x their size on coherence, groundedness, and rollout diversity, at comparable latency. Agents trained in these simulated worlds gained better performance improvements than being trained on AR models on three out-of-distribution environments with no environment-specific fine-tuning, showing that MDLMs are more steerable simulators
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- dOPSD: On-Policy Self-Distillation for Diffusion Language Models (2026)
- Trace-Based On-Policy Distillation for Masked Diffusion Language Models (2026)
- Back on Track: Aligning Rewards and States for Reasoning in Diffusion Large Language Models (2026)
- Data-Efficient Autoregressive-to-Diffusion Language Models via On-Policy Distillation (2026)
- Self-Augmenting Retrieval for Diffusion Language Models (2026)
- Learning Agent-Compatible Context Management for Long-Horizon Tasks (2026)
- Diffusion-GR2: Diffusion Generative Reasoning Re-ranker (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Get this paper in your agent:
hf papers read 2607.16204 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 1
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper