Reinforcement World Model Learning for LLM-based Agents

• URL: http://arxiv.org/abs/2602.05842v2
• Date: Mon, 09 Feb 2026 04:31:11 GMT
• Title: Reinforcement World Model Learning for LLM-based Agents
• Authors: Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu,
• Abstract summary: Reinforcement World Model Learning (RWML) is a self-conditioned method that learns action-supervised world models for LLM-based agents.<n>Our method aligns simulated next states produced by the model with realized next states observed from the environment.<n>We evaluate our method on ALFWorld and $2$ Bench and observe significant gains over the base model, despite being entirely self-supervised.
• Score: 60.65003139516272
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Grounded in Reality: Learning and Deploying Proactive LLM from Offline Logs [72.08224879435762]
  textttLearn-to-Ask is a simulator-free framework for learning and deploying proactive dialogue agents.<n>Our approach culminates in the successful deployment of LLMs into a live, large-scale online AI service.
  arXiv  Detail & Related papers  (2025-10-29T12:08:07Z)
• Internalizing World Models via Self-Play Finetuning for Agentic RL [65.96875390986655]
  Large Language Models (LLMs) as agents often struggle in out-of-distribution (OOD) scenarios.<n>We show how to encode this world model by decomposing it into two components: state representation and transition modeling.<n>We introduce SPA, a simple reinforcement learning framework that cold-starts the policy via a Self-Play supervised finetuning stage to learn the world model.
  arXiv  Detail & Related papers  (2025-10-16T18:03:39Z)
• Foundation Models as World Models: A Foundational Study in Text-Based GridWorlds [2.9165586612027234]
  Foundation models (FMs) are natural candidates to improve sample efficiency as they possess broad knowledge and reasoning capabilities.<n>We consider the use of foundation world models (FWMs) that exploit the prior knowledge of FMs to enable training and evaluating agents with simulated interactions.<n>Second, we consider the use of foundation agents (FAs) that exploit the reasoning capabilities of FMs for decision-making.
  arXiv  Detail & Related papers  (2025-09-19T12:10:28Z)
• RLAE: Reinforcement Learning-Assisted Ensemble for LLMs [21.77261258691006]
  Large language models (LLMs) can effectively combine diverse strengths of different models, offering a promising approach to enhance performance across various tasks.<n>We propose Reinforcement Learning-Assisted Ensemble for LLMs, a novel framework that reformulates ensemble through the lens of a Markov Decision Process (MDP)<n>Our approach introduces a RL agent that dynamically adjusts ensemble weights by considering both input context and intermediate generation states.
  arXiv  Detail & Related papers  (2025-05-31T07:38:41Z)
• WALL-E 2.0: World Alignment by NeuroSymbolic Learning improves World Model-based LLM Agents [55.64361927346957]
  We propose a training-free "world alignment" that learns an environment's symbolic knowledge complementary to large language models (LLMs)<n>We also propose an RL-free, model-based agent "WALL-E 2.0" through the model-predictive control framework.<n> WALL-E 2.0 significantly outperforms existing methods on open-world challenges in Mars (Minecraft like) and ALFWorld (embodied indoor environments)
  arXiv  Detail & Related papers  (2025-04-22T10:58:27Z)
• WALL-E: World Alignment by Rule Learning Improves World Model-based LLM Agents [55.64361927346957]
  We propose a neurosymbolic approach to learn rules gradient-free through large language models (LLMs) Our embodied LLM agent "WALL-E" is built upon model-predictive control (MPC) On open-world challenges in Minecraft and ALFWorld, WALL-E achieves higher success rates than existing methods.
  arXiv  Detail & Related papers  (2024-10-09T23:37:36Z)
• From Words to Actions: Unveiling the Theoretical Underpinnings of LLM-Driven Autonomous Systems [59.40480894948944]
  Large language model (LLM) empowered agents are able to solve decision-making problems in the physical world. Under this model, the LLM Planner navigates a partially observable Markov decision process (POMDP) by iteratively generating language-based subgoals via prompting. We prove that the pretrained LLM Planner effectively performs Bayesian aggregated imitation learning (BAIL) through in-context learning.
  arXiv  Detail & Related papers  (2024-05-30T09:42:54Z)
• Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models [52.98743860365194]
  We propose a new fine-tuning method called Self-Play fIne-tuNing (SPIN) At the heart of SPIN lies a self-play mechanism, where the LLM refines its capability by playing against instances of itself. This sheds light on the promise of self-play, enabling the achievement of human-level performance in LLMs without the need for expert opponents.
  arXiv  Detail & Related papers  (2024-01-02T18:53:13Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.