MemWM: Memory-Augmented Text-Based World Model

2026-08-07

AI Summary

This work addresses the limitations of existing text-based world models, which often fail in planning due to omitting critical facts, altering entity attributes, or misapplying transition rules. To mitigate these issues, we propose the first text world model augmented with a structured external memory mechanism, comprising a world memory bank that stores transition rules, state caches, and hard-to-predict facts. This memory is leveraged through retrieval and task-level skill injection to guide state generation and policy correction. We further introduce Structured State Fidelity (SSF), a novel evaluation metric, and demonstrate efficient memory-guided planning under policy freezing. Experiments on ALFWorld, WebShop, and ScienceWorld show that our approach improves downstream task success rates by up to 65.4% over supervised fine-tuning baselines and achieves a 206.3% gain in SSF.

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.

Problem

• world models
• systematic prediction errors
• factual state preservation
• text-based planning
• state transition accuracy

Innovation

• memory-augmented world model
• world memory
• Structured State Fidelity
• text-based planning
• retrieval-augmented reasoning