Retrieval-Augmented World Models Enhanced with Reinforcement Learning

ACL ARR 2025 February Submission3429 Authors

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 knowledge to the specific environments. Existing research attempts to enhance LLM-based world models through prompting or fine-tuning approaches, which are either requiring human knowledge or computational extensive. Therefore, we introduce Retrieval-Augmented World Models (Rawm), a novel framework that leverages retrieval-augmented generation to improve the LLM-based world models.
Our main contributions are threefold:
(i) We introduce a memory system and design an embedding model to retrieve and incorporate relevant experiences, significantly improving 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 RL environments, i.e., Game24, BlocksWorld, and BabyAI, demonstrating that Rawm consistently outperforms baseline models and exhibits strong generalizability.
By leveraging external memory and retrieval techniques and training embedding with RL pipeline, Rawm dynamically utilizes relevant historical experiences and equips LLMs with environment-specific knowledge without retraining, enabling more accurate and generalizable predictions.

Paper Type: Long

Research Area: Generation

Research Area Keywords: retrieval-augmented generation,

Contribution Types: Approaches to low-resource settings

Languages Studied: english

Submission Number: 3429