Home | Publications | WZB+26

MemWM: Memory-Augmented Text-Based World Model

MCML Authors

Link to Profile Hinrich Schütze

Hinrich Schütze

Prof. Dr.

Core PI

Computational Linguistics

Abstract

World models are increasingly used to support planning in agents by predicting how environment states evolve in response to agent actions. Yet fluent next-state predictions can still omit task-critical facts, corrupt product attributes, or apply incorrect transition rules. To address such systematic prediction errors, we introduce MemWM, a memory-augmented text-based world model. MemWM uses world memory, a curated memory bank of transition rules, state caches, and hard-to-predict facts, to condition next-state imagination. We evaluate factual state preservation with Structured State Fidelity (SSF), which scores predicted states through benchmark-specific facts and fields. Compared with SFT, memory-augmented training improves SSF by up to 206.3%. In the full planning setting, we keep the policy model frozen and provide policy-side world skill: retrieved task-level skills and step-wise corrective guidance for action selection. Across ALFWorld, WebShop, and ScienceWorld, memory-augmented agents improve downstream success over an SFT-trained world-model agent, with up to a 65.4% relative gain. Sensitivity analyses further show that retrieved memory improves task success and efficiency under different memory and action-budget settings.

inproceedings WZB+26

EMNLP 2026

Conference on Empirical Methods in Natural Language Processing. Budapest, Hungary, Oct 24-29, 2026. To be published. Preprint available.
Conference logo
A* Conference

Authors

Y. Wang • T. Zhang • J. Bi • A. Aniri • W. Ye • B. Liu • S. Yan • S. Wang • X. Zhou • S. Pirk • H. Schütze • Y. Ma

Links

arXiv

Research Areas

 A3 | Computational Models

 B2 | Natural Language Processing

BibTeXKey: WZB+26

#p-schuetze #p-tresp
Back to Top