Masked diffusion models beat LLMs 4x their size as agent world models
239K trajectories, 9 environments, 12 model families — MDLMs deliver 47% absolute gains.
Darshan Deshpande's new paper formalizes text-based world modeling as a steerable transition-dynamics problem, breaking it into initial state, task context, tool schemas, domain rules, and steering directives. The team curated 239,403 grounded state-action trajectories spanning nine open-source environments and twelve frontier model families. They compared autoregressive LMs (AR LMs) against masked diffusion language models (MDLMs), finding that MDLMs—via bidirectional anchor-aware denoising—achieve better coherence, groundedness, and empirically validated rollout diversity than LLMs four times their parameter size, all at comparable inference latency.
To make these world models practical for agentic RL, the authors introduced a plug-and-play GRPO training framework with deterministic state checks. In zero-shot transfer ablations across three out-of-distribution environments (ScienceWorld, ALFWorld, AppWorld) and three 1.2B–7B agent backbones (LFM2.5, Qwen3, Mistral), the approach delivered up to 47% absolute gains over baselines without any environment-specific fine-tuning. Behavioral analysis of failure modes under adversarial scenarios and human evaluations on realism, outcome correctness, and training utility further validated the approach. The entire dataset and training code are open-sourced.
- MDLMs outperform autoregressive LLMs 4x their size in coherence, groundedness, and rollout diversity at comparable latency.
- GRPO training framework with deterministic state checks achieves up to 47% absolute gains on OOD environments without fine-tuning.
- Dataset includes 239,403 trajectories from 9 environments and 12 model families, all open-sourced.
Why It Matters
More scalable, diverse training environments for RL agents without expensive hand-curation or fine-tuning.