cs.AIOct 1, 2026

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

Authors: Dhananjay Ashok, Shantanu Agarwal, Vivek Datla, Jonathan May, Alfy Samuel

Organizations: Information Sciences Institute University of Southern California · CapitalOne · Capital One

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.

Figures & tables

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Explore similar work

CardsList
  1. Qwen-AgentWorld: Language World Models for General Agents

    Jun 23, 2026Yuxin Zuo, Zikai Xiao, Li Sheng +30LLM World ModelsAgentic RL

  2. Bridging the Agent-World Gap: Text World Models for LLM-based Agents

    Jun 8, 2026Yixia Li, Hongru Wang, Peng Lai +13LLM World ModelsWorld Models

  3. From Traces to Agentic Worlds: Agentic Language World Models for Interactive Environment Simulation

    Oct 5, 2026Quanyu Long, Xiao Chen, Jianda Chen +4LLM World ModelsInteractive World Models