How To Train Your World Model: Fine-tuning vs RAG for LM-based World Modeling
cs.AI
Submitted: 2026-10-01
Updated: 2026-10-01
License: http://creativecommons.org/licenses/by/4.0/
The gist: World models (WMs) simulate the transition dynamics of environments, enabling agents to plan over the consequences of their actions.
Terminology
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.
Sources
- RAG vs Fine-tuning: Pipelines, Tradeoffs, and a Case Study on Agriculture
- Language Models are Few-Shot Learners
- Retrieval-Augmented LLM Agents: Learning to Learn from Experience
- The Llama 3 Herd of Models
- Is Your LLM Secretly a World Model of the Internet? Model-Based Planning for Web Agents
- World Models
- Dream to Control: Learning Behaviors by Latent Imagination
- Mastering Atari with Discrete World Models
- Mistral 7B
- Do LLMs Strategically Reveal, Conceal, and Infer Information? A Theoretical and Empirical Analysis in The Chameleon Game
- R-WoM: Retrieval-augmented World Model For Computer-use Agents
- ALFWorld: Aligning Text and Embodied Environments for Interactive Learning
- Gemma 3 Technical Report
- Omni-Embed-Nemotron: A Unified Multimodal Retrieval Model for Text, Image, Audio, and Video
- Self-Evolving World Models for LLM Agent Planning
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems