Efficient Integration of External Knowledge to LLM -based World Models via Retrieval-Augmented Generation and Reinforcement Learning (2025)

Abstract: World models achieve remarkable success in predicting future states and planning in complex environments and Large Language Models (LLMs) serve as promising foundation to build general world models. However, their performances are usually constrained by the limited external knowledge to specific environments. Existing research attempts to enhance LLM-based world models through prompting or fine-tuning approaches, which are either requiring human knowledge or computationally extensive. Therefore, we introduce Retrieval-Augmented World Models (RAWM), a novel framework that leverages retrieval-augmented generation to efficiently integrate the external knowledge to LLM-based world models. Our main contributions are threefold: (i) We introduce a memory system and design an embedding model to retrieve relevant experiences as the in-context examples to improve the world model’s predictive accuracy. (ii) We develop a reinforcement learning (RL) training pipeline that fine-tunes a small MLP head on the pre-trained embedding model using Proximal Policy Optimization (PPO), further enhancing prediction performance. (iii) We conduct extensive experiments across three diverse environments, i.e., Game24, BlocksWorld, and BabyAI, demonstrating that RAWM consistently outperforms baseline models and exhibits strong generalizability. By leveraging the retrieval-augmented generation and the efficient RL training pipeline, RAWM dynamically utilizes relevant historical experiences and equips LLMs with environment-specific external knowledge without retraining, enabling more accurate and generalizable predictions.

## introduction why world model is important? the world model @xcite emerges to be an important module in decision making due to * corresponding author the celebrating success of muzero @xcite and dreamer @xcite @xcite . as learned accurate simulators, world models encode rich representations of the complex dynamics of the environment to predict the future states and the rewards. world models are critical for several key capabilities, such as generalization to novel tasks @xcite @xcite , efficient planning @xcite @xcite , and offline learning @xcite @xcite . beyond decision-making tasks, recent works such as genie @xcite and vista @xcite demonstrate that world models can be generalpurpose world simulators and users can directly interact with them for playing and planning. why llm-based world models? the past five years witness the remarkable success of large language models (llms) in enormous text generation and understanding tasks @xcite . llms serve as the world model explicitly in reasoning via planning (rap) @xcite and reason for future, act for now (rafa) @xcite , where the llms predict the next states based on the actions executed at current states, e.g., the states of blocks in the blocksworld @xcite , for planning. llms serve as the world model implicitly in the widely-used tree of thoughts (tot) @xcite , as well as graph of thoughts (got) @xcite , where the llms need to predict the states and evaluate the thoughts for the selection of the thoughts to advance the reasoning. the main advantage of llm-based world models is that llms are pre-trained over internet-scale data and can capture diverse patterns in different environments. more discussion can be found in appendix a.1 why llm-based world models may fail? however, the pre-trained llms may lack the external knowledge of specific environments, which pro- room 1 (current room) room 2 (next room) wm what can i see if i move to the next room? i don't know. rawm i saw a chair in the current room. as both rooms are similar, there may be a chair in the next room. hibits them to be accurate world models. for the example in figure 1 , the llm cannot provide the accurate predictions whether there is a chair in room 2 if room 2 is never been visited. to address this issue, we can carefully design the prompts to add the specific knowledge to help the llms in making predictions, e.g., the rules for objects and actions @xcite . however, the knowledge is even usually not available for humans. alternatively, we can fine-tune the llms on the specific environments @xcite . however, the training of llms brings additional complexities for building the world models with llms and may also hurt the generalizability of llms across different tasks. ## preliminaries in this section, we present the preliminaries of rawm, including the formulation of the decision making, the llms, and the world models. ## retrieval-augmented world models in this section, we introduce retrieval-augmented world models (rawm). we will first introduce the architecture of rawm and then introduce the rl training pipeline for the retrieval process. the architecture of rawm is displayed in figure 2 . we introduce a memory ξ, which stores the pre-collected experiences, an embedding model, which is used to rank and retrieve the relevant experiences. specifically, given the query @xmath0, a) ∈ q, where q is the query dataset, we will use the embedding model to query top k relevant experiences @xmath1, where c @xmath2, a k , s ′ k ), @xmath3, . . . , k. the retrieved experiences c will be concatenated with the query q to form the input to the world model ω. we note that for the environments where the states are not texts, e.g., babyai (chevalier-boisvert et al., 2019a), we need to first transform them into the text representation. prompt design. for the prompt design, any information related to the environments will not be provided to the world model, including the tasks, the object and action rules. we expect that all the environment knowledge is provided by the incontext examples retrieved from the memory. the prompt template is displayed as follows: ## conclusions in this work, we introduce retrieval-augmented world models (rawm), which leverage the retrieval-augmented generation for efficient integration of external knowledge into llm-based world models. we then introduce an efficient rl training pipeline to further improve the performance. extensive experiments demonstrate the effectiveness and the generalizability of rawm. rawm is an efficient method to build the highly capable llmbased world models without fine-tuning llms.