LLM 기반 에이전트를 위한 강화 세계 모델 학습
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 기반 에이전트에 세계 모델링 능력이 필요함을 시사합니다. 본 연구에서는 시뮬레이션-현실 간격 보상을 활용하여 텍스트 상태에서 LLM 기반 에이전트의 행동 조건부 세계 모델을 학습하는 자기 지도 방법인 Reinforcement World Model Learning(RWML)을 제안합니다. 우리의 방법은 모델이 생성한 시뮬레이션 다음 상태와 환경에서 관찰된 실제 다음 상태를 사전 학습된 임베딩 공간에서 정렬함으로써, 내부 세계 시뮬레이션과 실제 환경 역학 간의 일관성을 촉진합니다. 토큰 수준 충실도(즉, 정확한 어휘 재현)를 의미적 등가성보다 우선시하여 모델 붕괴를 초래할 수 있는 다음 상태 토큰 예측과 달리, 우리의 방법은 더 강력한 학습 신호를 제공하며 실증적으로 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.