大規模言語モデルベースエージェントのための強化学習世界モデル学習
Reinforcement World Model Learning for LLM-based Agents
February 5, 2026
著者: Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu
cs.AI
要旨
大規模言語モデル（LLM）は、言語中心のタスクにおいて強力な性能を発揮している。しかし、エージェント環境では、LLMは行動の結果を予測し、環境のダイナミクスに適応することにしばしば困難をきたすため、LLMベースのエージェントにおける世界モデリング能力の必要性が浮き彫りになっている。本研究では、強化学習世界モデル学習（RWML）を提案する。これは、シミュレーションtoリアルギャップ報酬を用いて、テキスト状態におけるLLMベースのエージェントのための行動条件付き世界モデルを学習する自己教師あり手法である。本手法は、モデルが生成するシミュレートされた次の状態と、環境から観測された実際の次の状態とを、事前学習済みの埋め込み空間内で整合させ、内部的な世界シミュレーションと実際の環境ダイナミクスとの一貫性を促進する。トークンレベルの忠実度（すなわち、正確な文言の再現）を意味的等価性よりも優先し、モデル崩壊を引き起こしうる次の状態トークン予測とは異なり、本手法はよりロバストな訓練信号を提供し、LLM-as-a-judgeよりも報酬ハッキングの影響を実証的に受けにくい。ALFWorldおよびτ^2 Benchにおいて本手法を評価した結果、完全に自己教師ありでありながら、ベースモデルを大幅に上回る性能向上が確認された。タスク成功報酬と組み合わせた場合、本手法はALFWorldおよびτ^2 Benchにおいて、直接タスク成功報酬を用いた強化学習をそれぞれ6.9ポイント、5.7ポイント上回り、専門家データを用いた訓練と同等の性能を発揮した。
English
Large language models (LLMs) have achieved strong performance in language-centric tasks. However, in agentic settings, LLMs often struggle to anticipate action consequences and adapt to environment dynamics, highlighting the need for world-modeling capabilities in LLM-based agents. We propose Reinforcement World Model Learning (RWML), a self-supervised method that learns action-conditioned world models for LLM-based agents on textual states using sim-to-real gap rewards. Our method aligns simulated next states produced by the model with realized next states observed from the environment, encouraging consistency between internal world simulations and actual environment dynamics in a pre-trained embedding space. Unlike next-state token prediction, which prioritizes token-level fidelity (i.e., reproducing exact wording) over semantic equivalence and can lead to model collapse, our method provides a more robust training signal and is empirically less susceptible to reward hacking than LLM-as-a-judge. We evaluate our method on ALFWorld and τ^2 Bench and observe significant gains over the base model, despite being entirely self-supervised. When combined with task-success rewards, our method outperforms direct task-success reward RL by 6.9 and 5.7 points on ALFWorld and τ^2 Bench respectively, while matching the performance of expert-data training.