Masked Diffusion LMはAgent強化学習向けの新しいテキスト世界モデルになるか

導入

Agent向けの強化学習では、環境そのものがボトルネックになりつつある。手作業で作られた環境は数も難度も限られ、モデルが強くなるほど訓練信号として弱くなる。さらに、長いタスクで報酬が疎になると、Agentは特定の手順やツール利用パターンに偏りやすい。そこで注目されるのが、環境状態をシミュレートして多様な訓練ロールアウトを生成する世界モデルだ。

この論文の問いは、テキスト世界モデルに自回帰言語モデルを使い続けるべきか、という点にある。自回帰モデルは左から右へ順番に生成するため、通常の文章生成には適している。しかしAgent環境では、次の状態が単なる文ではなく、ツールスキーマ、過去の会話、ドメインルール、期待される結果と同時に整合していなければならない。著者らは、双方向にノイズ除去を行うMasked Diffusion Language Models（MDLMs）が、この種のグローバルな制約により適していると見る。

核心ポイント

• 問題設定の整理：テキストベースの世界モデリングを、制御可能な状態遷移ダイナミクスとして定式化。初期状態、タスク文脈、ツールスキーマ、ドメインルール、steering directivesに分解して扱う。
• データ構築：9つのオープンソース環境と12のフロンティアモデルファミリーにまたがる、239,403件のgrounded state-action軌跡を整備した。
• モデル比較：MDLMは、双方向かつアンカーを意識したdenoisingにより、一貫性、groundedness、実証的に確認されたrollout diversityで、自身より4倍以上大きい自回帰LLMを上回ったとされる。推論レイテンシは同程度と報告されている。
• 訓練フレームワーク：決定的な状態チェックを組み込んだプラグイン型のGRPO訓練フレームワークを導入し、シミュレーション中の状態崩れを抑える設計にしている。
• 転移評価：ScienceWorld、ALFWorld、AppWorldの3つのOOD環境で、LFM2.5、Qwen3、Mistralを含む1.2B〜7B規模のAgentバックボーンを評価。環境固有の微調整なしで、最大47%の絶対改善が報告された。

意義と影響

この研究の意義は、拡散型言語モデルが単に別の指標で強かったという点にとどまらない。重要なのは、Agent用の世界モデルでは、生成順序そのものが性能を左右しうるという示唆だ。環境状態は複数の制約を同時に満たす必要があり、逐次生成では長いロールアウトほど矛盾が蓄積しやすい。

MDLMの双方向denoisingは、まず全体の制約を保ちながら状態を補完するという発想に近い。そのため、ツール利用や過去履歴との整合性が重要なAgent訓練では、自回帰モデルより扱いやすい可能性がある。

一方で、結果は特定のオープン環境と構築された軌跡データに基づく。より複雑な実ツール環境、長期運用タスク、高リスク領域でも同じ利点が出るかは今後の検証が必要だ。それでも、テキスト世界モデルの設計を見直す有力な方向性を示した研究と言える。

出典：Hugging Face Daily Papers