Latent Dynamics · 2026-10-05
How To Train Your World Model: Fine-tuning vs RAG for LM-based World Modeling
Dhananjay Ashok, Shantanu Agarwal, Vivek Datla, Jonathan May, Alfy Samuel
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors systematically compare fine-tuned vs RAG-based LM world models across five text environments. They report fine-tuning yields higher rewards on 15/20 model-environment combinations, while RAG is more data-efficient and outperforms in Chameleon. They find retrieval often returns suboptimal transitions (recall@5 < 0.7) and propose hierarchical retrieval plus a hybrid approach that combines parametric dynamics with retrieval, outperforming other methods on 17/20 combinations.
Why it matters
The paper evaluates how to build LM-based world models from self-collected experience, focusing on whether RAG can replace fine-tuning. It also diagnoses a concrete failure mode of RAG—retrieval surfacing suboptimal transitions—and tests query reformulation via hierarchical retrieval. The hybrid method provides a practical alternative that the authors report improves utility across environments.
Method
- Constructs fine-tuned LM world models (FT) by fine-tuning on verbalised transitions from an experience replay buffer; and retrieval-augmented generation world models (RAG) by indexing transitions and using top-5 retrieved in-context demonstrations at inference.
- Evaluates world models for fidelity (LM-as-a-judge on a 1–10 scale; plus BLEU and BertScore in the appendix) and for utility via total reward scaled by an oracle-access upper bound.
- For RAG, introduces counterfactual intervention to estimate retrieval-stage error; evaluates query reformulation and finds hierarchical retrieval (observation-similarity then action-similarity reranking) outperforms traditional retrieval; proposes a hybrid (parametric + retrieved memory) system.
Limitation
The paper reports that retrieval for LM-based world modelling “remains underexplored” and, for fidelity metrics, notes that BLEU and BertScore exhibit low variance and are likely poor measures of WM fidelity.
Abstract (from arXiv)
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.
Related papers
- Model What Matters: Decision-Guided Selective World Modeling for Language Agents
- MemoWM: How World Models Change What Agents Need to Remember
- What Matters for World Modeling in Vision-Language-Action Models?
- RoleWorld: Knowledge Transfer in World Modeling with Functional Role Abstractions
- Steering Black-Box Large Language Model Agents with Pitfall Anticipation World Model
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-01 · Indexed by Pub-AI: 2026-10-05 UTC
How To Train Your World Model: Fine-tuning vs RAG for LM-based World Modeling
Preprint posted: 2026-10-01 · Original record ↗
- arxiv
- 2610.02542
Explicit source versions (1)
- v1 · 2026-10-01
- Preprint available
Source date: 2026-10-01 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
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.