マスク拡散言語モデルは、エージェンティック強化学習のための強力で制御可能なテキストベースの世界モデルである。

要旨

近年の強化学習(RL)の発展に伴い、多様で特殊な訓練環境への需要が高まっている。固定されたタスクと報酬難易度を持つ手作業で設計された環境は、モデルの性能が向上するにつれて効果的なシグナルとして機能しなくなり、長期的なスパース報酬は特定のワークフローやツール構造へのモード崩壊を引き起こす。環境状態をシミュレートするワールドモデルは純粋なロールアウト性能に匹敵する成果を上げており、オンデマンドで多様性を拡張する有望な手段となっている。しかし、自己回帰(AR)ワールドモデルは左から右へのバイアスに悩まされ、ツールスキーマ、先行ターン、期待される結果などの大域的に相互依存する状態アンカーへの条件付けが困難である。本研究では、(i)テキストベースのワールドモデリングを、初期状態、タスクコンテキスト、ツールスキーマ、ドメインルール、誘導指示に分解された操作可能な遷移ダイナミクス問題として定式化し、(ii)9つのオープンソース環境と12のフロンティアモデルファミリーにわたる239,403の grounded な状態行動軌跡をキュレーションする。AR言語モデルとマスク拡散言語モデル(MDLM)を比較し、MDLMが双方向のアンカー認識型デノイジングを通じて、4倍のパラメータサイズを持つLLMよりも優れたコヒーレンス、grounded性、実験的に検証されたロールアウト多様性を達成し、推論レイテンシは同等であることを示す。さらに、決定論的状態チェックを備えたプラグアンドプレイのGRPO訓練フレームワークを導入し、3つのOOD環境(ScienceWorld、ALFWorld、AppWorld)において、3つの1.2B-7Bエージェントバックボーン(LFM2.5、Qwen3、Mistral)を用いたゼロショット転移アブレーションを実施し、環境固有の微調整なしでベースラインを最大47%絶対的に上回る成果を達成した。さらに、敵対的シナリオ下での障害モードの行動解析と、現実感、結果の正しさ、訓練有用性に関する人間評価を実施する。本研究成果はオープンソースとして公開し、この方向性の研究を促進する。

English

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.