How To Train Your World Model: Fine-tuning vs RAG for LM-based World Modeling

2026-10-01

AI Summary

This study addresses the efficiency of dynamics modeling in language model-based world models for text environment simulation by systematically comparing fine-tuning with retrieval-augmented generation (RAG). While fine-tuning achieves superior performance at high data cost, RAG proves more data-efficient. To bridge this gap, we introduce counterfactual interventions to estimate retrieval errors and design a hierarchical query reformulation strategy. Building on these insights, we construct a hybrid world model that integrates parameterized core dynamics with a non-parametric memory bank. Experiments demonstrate that the proposed system consistently outperforms existing methods across diverse environments and model configurations, validating the effectiveness of combining parameterized and non-parametric approaches to enhance agent planning capabilities.

Abstract

World models (WMs) simulate the transition dynamics of environments, enabling agents to plan over the consequences of their actions. In text-based environments, fine-tuning a Language Model (LM) to serve as a WM has emerged as a dominant paradigm. However, despite the widespread success of non-parametric approaches such as Retrieval Augmented Generation (RAG), retrieval for LM-based world modelling remains underexplored. We conduct a systematic evaluation across five diverse environments spanning embodied, web navigation and social settings, comparing fine-tuning and RAG-based approaches for LM-based world modelling. Our study reveals that fine-tuning often outperforms RAG, with fine-tuned WMs enabling agents to obtain higher rewards on 15/20 settings. While both construction paradigms benefit from additional and more diverse exploration, RAG-based approaches prove more data-efficient, and fine-tuning approaches disproportionately benefit from scaling the amount of experience collected. With a focus on RAG-based WMs, we devise a procedure that uses counterfactual intervention to estimate the error rate of the retrieval stage, and show that retrievers consistently surface suboptimal transitions from the experience buffer. Hoping to address this failing, we study a variety of query reformulation strategies, demonstrating that a hierarchical approach outperforms the traditional retrieval pipeline. Finally, we compose our findings into a hybrid world modelling system that parametrically captures core environment dynamics, while learning to rely on retrieval from an actively maintained memory store. Our hybrid system consistently outperforms other methods across multiple environments and models, showcasing the robustness of the approach and the applicability of our findings.

Problem

• World Models
• Language Models
• Fine-tuning
• Retrieval Augmented Generation (RAG)
• Text-based Environments

Innovation

• World Models
• Retrieval Augmented Generation
• Counterfactual Intervention
• Query Reformulation
• Hybrid World Modelling