World Modeling for Agentic RL

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

A language model that stands in for an environment while an agent trains.

Darshan Deshpande
“Earth’s Aquarium”, Bering Sea
01

In brief

02

Key figure

AppWorld, ScienceWorld and ALFWorld, for three policy models trained four ways. RL with rollouts from the diffusion world model gives the highest success in all nine model and environment pairs, with the largest gains on the partially observable environments.
AppWorld, ScienceWorld and ALFWorld, for three policy models trained four ways. RL with rollouts from the diffusion world model gives the highest success in all nine model and environment pairs, with the largest gains on the partially observable environments.
03

Runs

03

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.

Cite this
TextBibTeXCopy
Darshan Deshpande. Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL. In NeurIPS 2026. arXiv:2607.16204.
@inproceedings{deshpande2026worldmodelin, title={Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL}, author={Darshan Deshpande}, booktitle={NeurIPS 2026}, year={2026}, eprint={2607.16204}, archivePrefix={arXiv} }
arXiv 2607.16204
May 7, 2026
NeurIPS 2026
Main
Main conference
world-modeling-for-agentic-rl
More research
All research →
“Luminescence”, Suwannee River plume, Gulf of Mexico
COLM 2026 Workshop
meta

Prompt Infilling for Diffusion LMs

Teaching diffusion language models to fill in prompts, not only responses.

April 4, 2026
arXiv 2604.03677
Workshop
prompt-infilling-for-diffusion-lms
COLM 2026 Workshop
“A Study in Color”, Kenya
meta

Generative Simulators

Environments that generate their own tasks, world dynamics and rewards.

December 17, 2025
Technical report (PDF)
generative-simulators
“Himalayas”, Eastern Himalaya, southwestern China
Nature
meta

Humanity’s Last Exam

A multi-modal benchmark at the frontier of human knowledge.

January 24, 2025
arXiv 2501.14249
Main
humanitys-last-exam
Nature
From the blog
All posts →
No items found.