PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Darshan Deshpande

arXiv:2607.16204v1cs.AIcs.LG

TL;DR

Text-based world models must support diverse, steerable, multi-turn agent training, but autoregressive models struggle with globally interdependent state anchors. This paper formalizes steerable transition dynamics and evaluates masked diffusion language models, which produce more coherent, grounded, and diverse rollouts and improve zero-shot agent transfer by up to 47%.

  • Problem

    Multi-turn agent training faces sparse rewards, cross-turn distribution shift, and limited generalization, while autoregressive world models cannot condition bidirectionally on globally interdependent state anchors.

  • Method

    The paper formalizes text-based world modeling as steerable transition dynamics and compares anchor-aware masked diffusion models with autoregressive models across grounded trajectories and transfer tasks.

  • Results

    Across tasks, masked diffusion models produce more coherent, grounded, and diverse rollouts than larger autoregressive models, while MDLM-generated rollouts improve zero-shot agent performance by up to 47%.

  • Takeaways & Limitations

    Bidirectional masked diffusion is a strong, steerable basis for text-based world models and supports downstream agent training without environment-specific fine-tuning.

Abstract

from arXiv · show

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.

1 Introduction

The introduction frames reinforcement learning as having transformed LLMs into decision-making agents for complex, dynamic tool-use environments. These agents have demonstrated strong performance across diverse domains.

  • RL has transformed LLMs from passive sequence generators into decision-making agents.
  • These agents can operate in complex and dynamic tool-use environments.
  • RL-trained agents have demonstrated strong performance across diverse domains.

3) MDLM

MDLMs provide a bidirectional alternative to causal world models by iteratively denoising masked sequences, enabling globally consistent environment-state simulation. Experiments indicate that smaller MDLMs can outperform substantially larger finetuned causal models, implicating causal inductive bias rather than capacity as the primary bottleneck.

  • Anchor-aware generation: AR decoding cannot condition on trailing schema or context anchors, producing prefix-consistent but globally incoherent samples.The figure illustrates a contradiction such as refund_processed: true paired with status: "error".
  • Motivation: MDLMs address scalable environment modeling needs arising from sparse rewards, distributional shift across turns, and limited generalization beyond training environments.These challenges affect multi-turn agents that must maintain dialogue memory, make sequential decisions, and adapt to environmental feedback.
  • Motivation: MDLMs iteratively denoise corrupted token sequences with bidirectional context, conditioning predicted states on preceding and succeeding observations without causal masking.This enables modeling of global consistency, which the paper identifies as essential for faithful environment simulation.
  • Results: 8B-parameter MDLMs, including SDAR, LLaDA-2.1, and WeDLM, consistently outperform finetuned causal models up to 35B mixture-of-expert and 27B dense AR parameters.The compared causal models include Qwen-3.5, Nemotron-3, GLM-4.7-Flash, and GPT-OSS-20B.
  • Results: The results suggest that causal inductive bias, rather than model capacity, is the primary bottleneck in LLM-based world modeling.The conclusion follows from smaller MDLMs outperforming larger finetuned causal models across the reported scales.

2 Relevant Work

Prior work uses LLMs for scalable synthetic data and tool simulation, while MDLMs offer bidirectional, parallel generation for controllable state prediction. World-modeling research addresses the need for specialized, maintainable RL environments and has begun benchmarking autoregressive models.

  • LLMs as synthetic data simulators: LLMs support large-scale synthetic data generation through persona, document or web, workflow, and long-horizon diversification and grounding techniques.These approaches target post-training data generation and scalable tool simulations.
  • MDLMs as controllable models: An ideal world model should be fast, high fidelity, generalizable, and steerable, but autoregressive models impose left-to-right bias and sequence-length-scaled inference costs.These limitations restrict diversity and editability and complicate efficient, high-performance inference.
  • MDLMs as controllable models: MDLM training enables parallel token generation and bidirectional infilling, improving applicability to arbitrary state prediction.Prior work also reports potential for high-quality generations and addresses the reversal curse.
  • MDLMs as controllable models: Converting large pretrained autoregressive models to MDLMs is reported as both feasible and performant, addressing the compute intensity of training MDLMs from scratch.This line of work builds on studies of reversal-curse mitigation and high-quality generation.
  • World Modeling: World modeling is being explored to reduce RL dependence on specialized, robustly curated environments that are difficult to scale and maintain.In text-based world modeling, early work exhaustively benchmarked and ablated autoregressive model performance.

3 Masked Diffusion Language Models for Steerable World Modeling

This section formalizes text-based world modeling as learning environment transition dynamics over structured token sequences while internally tracking prior environment states. It presents MDLMs as steerable alternatives to autoregressive models because bidirectional masked denoising supports arbitrary conditioning and structurally diverse rollouts.

  • Formalizing the World Model Objective: World models approximate environment transition dynamics by representing each state through an initial environment state, task context, tool schemas, reward rules, and domain information.The formalization decomposes each state as st = (et, ct, T, R, D), with T, R, and D fixed across the trajectory.
  • Formalizing the World Model Objective: Text-based states are serialized as token sequences, while prior environment states are tracked internally rather than emitted to avoid context pollution.The fidelity of this learned distribution determines the quality of imagined rollouts for downstream agent training.
  • Masked Diffusion Language Models for World Modeling: MDLMs learn full next-state sequences by iteratively denoising masked tokens conditioned on prior states and actions.Training independently masks tokens at a sampled diffusion time and trains a predictor to recover the masked positions.
  • Steering via Conditional Masked Generation: MDLMs support arbitrary known-position conditioning because their any-order likelihood bound learns every conditional direction, unlike autoregressive models’ fixed ordering.Parallel masked-token prediction from a shared context also avoids within-state error accumulation from conditioning later tokens on earlier hallucinations.
  • Diversity of MDLM Rollouts: MDLM iterative denoising produces greater structural rollout diversity than causal autoregressive sampling by varying both token choices and the positions committed at each step.The passage attributes this diversity to the denoising process rather than temperature alone and notes that empirical evidence appears in Section 5.

4 Experimental Setup

The experiments curate a broad grounded trajectory dataset and compare autoregressive and masked diffusion world models under standardized training and evaluation protocols. Models are assessed for generation fidelity, downstream training utility, and human-judged semantic correctness across in-domain and out-of-domain settings.

  • Environment Selection and Dataset Curation: 239,403 trajectories span software engineering, research automation, customer service, and general-purpose tool-use environments.The environments include SWE-SMITH, CODERFORGE, DEEPRESEARCHQA, OPENRESEARCHER, TAUBENCH, GORILLA/BFCLV4, TOOLATHLON, PANDORA 2, and WEBSHOP.
  • Baselines: The baseline suite includes 8B–35B open-source autoregressive LLMs with dense and mixture-of-experts configurations.Listed baselines include QWEN3.5-27B, GPT-OSS-20B, GLM-4.7-FLASH, NEMOTRON-3-NANO-30B-A3B, and QWEN3.5-35B-A3B.
  • Training: All models use AdamW, gradient clipping, cosine decay with warmup, and a maximum sequence length of 16,384 tokens.Learning rates and batch sizes are swept, and the best scores are reported.
  • Training: Each experiment uses three seeds, with means reported for the generation-quality and diversity tables.The seeds are 42, 7000, and 9000; training uses ms-swift, dFactory, or WEDLM’s official code depending on the model family.
  • Metrics: Evaluation covers generation fidelity, downstream training utility, and human-judged semantic correctness.Generation fidelity uses BLEU-1, ROUGE-L, and MAUVE across in-domain and out-of-domain test sets.

5 Results and Discussion

MDLMs outperform comparable AR models on distributional alignment and rollout diversity, with gains attributable to bidirectional masked diffusion rather than parameter count. Their generated trajectories also support downstream agent training and receive strong human ratings for realism, correctness, utility, and steerability.

  • RQ1: World-modeling comparison: SDAR-30B-A3B achieves the best zero-shot MAUVE on three of four splits, with gaps reaching .697 vs. .532 on API-BANK and .261 vs. .232 in-domain.Three-shot SDAR-30B-A3B further raises in-domain MAUVE to .995.
  • RQ1: World-modeling comparison: SDAR-8B outperforms fine-tuned QWEN3-8B across metrics, including a +.283 in-domain MAUVE gain, showing diffusion—not capacity or training data—drives the improvement.The results identify causal inductive bias as the primary bottleneck in LLM-based world modeling.
  • RQ2: Downstream agent performance: Downstream evaluation measures task success on APPWORLD, SCIENCEWORLD, and ALFWORLD using agent backbones ranging from 1.2B to 7B parameters, with QWEN3.5-27B and SDAR-8B selected for strong performance and efficient dense deployment.QWEN3.5-27B matched QWEN3.5-35B-A3B within ±0.5% overall performance in independent evaluations.
  • RQ3: Human evaluation: SDAR receives mean human ratings of 4.75 for realism, 4.25 for outcome correctness, and 4.50 for training utility, with Krippendorff’s α above 0.89 on every dimension.Experts also observed strong adherence to steering directives, including forced tool failures, while retaining realism.
  • Behavioral and steerability analysis: Behavioral analysis contrasts SDAR with AR baselines across adversarial settings involving incomplete information, verbose tool outputs, trajectory-length constraints, and infeasible tasks.The probe also covers database states with adjacent information and tasks requiring action under information insufficiency.

6 Conclusion · Appendix

The paper presents a systematic study of MDLMs as text-based world simulators for agentic RL, showing that bidirectional, anchor-aware denoising supports coherent, grounded, and diverse rollouts.

  • 6 Conclusion: The work is presented as the first systematic study of MDLMs as text-based world simulators for agentic RL.
  • 6 Conclusion: The authors formalize world modeling as a steerable transition-dynamics problem with five grounded components.
  • 6 Conclusion: The study curates a 239,403-trajectory dataset spanning nine environments and twelve frontier model families.
  • 6 Conclusion: MDLMs produce more coherent, grounded, and diverse environment rollouts than autoregressive LLMs more than 4× their parameter size.The reported advantage is attributed to bidirectional, anchor-aware denoising.

A Human Evaluations … Dimension 3: outcome correctness

The study human-evaluates MDLM world-model outputs across realism, outcome correctness, and training utility using independently scored predicted next states. Annotators report high mean scores and consistency, while identifying concrete API and grounding errors.

  • A Human Evaluations: The evaluation tests whether MDLM gains arise from masked diffusion rather than simply greater data access.The authors conduct a dedicated human evaluation of MDLM world-model outputs.
  • A Human Evaluations: Annotators assess 100 conversation-based predicted next states on three independent dimensions using 1–5 ratings.Ratings must remain independent because tool outputs can legitimately produce different states, including errors or correct responses.
  • A Human Evaluations: Annotators must understand tool-calling LLMs, hold a CS degree, have AI-agent experience, and be at least 18 years old.They also complete five qualification samples before accessing the full dataset.
  • Dimension 1: training_utility: Training utility asks whether a tool output would help an AI model proceed normally and correctly with its task.This dimension evaluates usefulness for model training rather than surface plausibility alone.
  • Dimension 2: realism: Realism asks whether each predicted response resembles a real API response for the relevant domain.The rubric treats realism as a separate dimension from outcome correctness and training utility.
  • Dimension 3: outcome correctness: Outcome correctness asks whether the model correctly predicts the action’s outcome, with independent scoring and domain judgment for plausible but hallucinated outputs.Annotators focus on valid, realistic tool outputs and apply consistent standards across items.
  • Dimension 3: outcome correctness: Mean scores across four annotators were 4.75 for realism, 4.25 for outcome correctness, and 4.5 for training utility.Krippendorff’s alpha indicated high interannotator consistency; reported issues included incorrect numeric types, corrupted API keys, and ungrounded or unrealistic outputs.

B Fairness of Evaluation

The evaluation protocol removes formatting artifacts and task-irrelevant fields that can deflate text-overlap scores. It also standardizes MAUVE computation and selects AR decoding temperatures to balance verbosity and correctness.

  • Fair evaluation preprocessing: JSON canonicalization sorts keys, strips whitespace, and normalizes stringified numeric entries to eliminate spurious prediction-reference mismatches.These steps address nondeterministic outputs and unordered JSON mappings that can deflate precision-based BLEU and ROUGE scores.
  • Fair evaluation preprocessing: API-BANK evaluation omits the shared input field because it does not reflect output quality.The retained schema fields are output, exception, and api_name.
  • Metric standardization: MAUVE uses the GPU implementation with k = max(2, round(min(|p|, |q|)/10)) and the authors’ default GPT-2 implementation.This follows the original MAUVE paper’s auto-sizing recommendation and maintains consistency with Pillutla et al. (2023).
  • AR temperature selection: 0.5 achieves the best verbosity-correctness balance for AR models evaluated at temperatures {0.5, 0.7, 0.9}.Higher temperatures produce hallucinated and more verbose outputs, while temperatures closer to 0 reduce MAUVE scores by up to 0.2 points.

C World Model Training Dataset Curation Details

The dataset was curated from diverse real-environment rollouts because no existing world-modeling datasets were available. Processing diversified trajectory lengths, while synthesized grounding added environment context for steerable state prediction and received high human approval.

  • Motivation: No existing world-modeling datasets were available, motivating a broad-coverage corpus of real environment rollouts that captures unique state-action pairs.The curation procedure included dataset filtering and analysis to support comparisons between causal and masked diffusion language models.
  • Trajectory Generation: Trajectory-generating models were sampled uniformly at random from twelve model families to promote diverse trajectories and broad behavioral coverage.The set included Qwen, GPT, Claude, Gemini, MiniMax, GLM, and DeepSeek variants.
  • Trajectory Length Diversification: Trajectory lengths were diversified through middle truncation and random sub-trajectory extraction to reduce fixed-context overfitting and expose varied rollout behaviors.Middle truncation targeted histories exceeding 16,384 tokens, while sub-trajectory extraction was applied with 15% probability at randomly selected tool-call boundaries.
  • Environment State Grounding: Claude-4.6-Sonnet synthesized hindsight context covering database states, reward structure, behavior patterns, and tool schemas to ground environment-state prediction.The instructions were structured as ground-truth context conditioned on the observed trajectory and next action.
  • Human Evaluation: 87% of evaluated grounded-and-steered instances received human approval, while these augmented instances constituted 17.2% of the final training dataset.Two independent annotators evaluated trajectories both in isolation and with grounding and steering objectives.

D Environment Setup for Downstream RL Agent Training with World Model Backend … D.4 Common Infrastructure

The downstream RL setup uses world-model rollouts with rewards computed from ground-truth state, across ScienceWorld, ALFWorld, and AppWorld, under a shared GRPO training infrastructure. Environment-specific splits, demonstrations, reward structures, and tool or action interfaces support evaluation against real environments.

  • D Environment Setup for Downstream RL Agent Training with World Model Backend: During GRPO, a World Model simulates rollouts while evaluation uses the real environment and rewards are computed from ground-truth environment state.This infrastructure is shared across the three interactive environments.
  • D.1 ScienceWorld: ScienceWorld covers 30 task types across 10 interconnected rooms, using 2 training and 2 evaluation variations per task type from disjoint pools.Agents issue free-text commands and submit answers through focus-on actions.
  • D.1 ScienceWorld: 37 perfect-score demonstrations are retained from GPT-5.5 trajectories capped at 50 turns for ScienceWorld supervised fine-tuning.Prompts include task-specific instructions such as substance/container distinctions, lifespan rankings, and exploration strategy.
  • D.2 ALFWorld: ALFWorld uses fixed-grammar household actions across standard splits, with 300 games sampled for RL from 3,553 training games.Its tasks include picking, placing, cleaning, heating, cooling, and examining objects in receptacle-filled rooms.
  • D.2 ALFWorld: Approximately 50 steps preceded consistent abandonment of exploitation-heavy ALFWorld tasks by small LFM-2.5 and Qwen3 models, motivating action-oriented reward shaping.The shaping combines completion rewards, partial pipeline credit, inactivity penalties, and repetition scaling.
  • D.3 AppWorld: AppWorld agents complete tasks through API calls to simulated applications, with world-model prompts containing full schemas, conversation state, and the active tool call.Rewards check task-completion predicates against the trajectory and evaluation information.
  • D.4 Common Infrastructure: All models use ms-swift with GRPO and vLLM colocate mode, while Mistral-7B-v0.3 uses DeepSpeed ZeRO-2 and all experiments use 8xH100s.Learning-rate, KL-penalty, temperature, and rollout sweeps are performed across environments.

E Inference Harness and Optimizations for MDLMs … COMPUTED QUERY STATE

The paper pairs an optimized MDLM inference harness with structured, deterministic prompts for downstream world-model rollouts. These designs balance denoising speed, enforce environment and tool constraints, and precompute valid query outcomes for exact responses.

  • E Inference Harness and Optimizations for MDLMs: The evaluation serves MDLMs with lmdeploy, uses a 1.3 repetition penalty, and selects 50 denoising steps to balance quality and speed.SDAR models showed no observed quality difference between 25 and 100 steps.
  • E Inference Harness and Optimizations for MDLMs: Inference block-size increases improve latency, but SDAR performance drops considerably from 4 to 8, while LLaDA-2.1-mini tolerates increases to 64 with up to 2× speed-ups.For LLaDA-2.1-mini, author-recommended Joint-Threshold decoding caused ≤5% test-split performance degradation.
  • F Impact of thinking on AR model performance: Test-time thinking improves AR performance through greater chain-of-thought diversity, strengthening MDLMs’ diversity-to-latency advantage despite their smaller size.The passage describes this conclusion as a Table 6 ablation finding.
  • G Prompts used for World Model for Downstream RL experiments: Downstream prompts encode task state, inventories, receptacles, object modifiers, and actions, then require only the environment’s immediate feedback in a fixed textual format.The ALFWorld prompt uses lowercase, instance-numbered entities and forbids planning, JSON, markdown, and action echoing.
  • G Prompts used for World Model for Downstream RL experiments: ALFWorld dynamics are deterministic through explicit action preconditions, state updates, sticky modifiers, and the universal fallback response “Nothing happens.”Rules cover movement, containers, inventory transfers, cleaning, heating, cooling, slicing, and object use.
  • TOOL SCHEMAS: AppWorld prompts validate tool names and argument types, require strict JSON responses, and distinguish valid calls from deterministic error cases.Tool schemas constrain list/search return shapes and require records and IDs to come only from supplied data.
  • COMPUTED QUERY STATE: Computed query state filters and paginates baked records before generation, supplies exact response-page records, and locally handles invalid calls without world-model involvement.The framework also specifies expected validation status, exact error bodies, totals, arrays, and login-token conditions.

H World Model Failure Modes

The section catalogs recurring world-model failures observed during GRPO training and links them to degraded RL signals and category-specific policy collapse. It also notes that these failures affected both MDLM and AR models and were addressed through grounding and response-handling mechanisms.

  • Table 7 catalogs recurring world-model failure modes observed during GRPO training, using trajectory snippets from actual runs to support explainability.
  • These failures degraded the RL training signal and caused agent policy collapse for specific task categories on the held-out test set.
  • The failures were observed in both MDLMs and AR models and addressed through grounding rules, local deterministic responders, and JSON unwrapping.
Loading 2607.16204v1…