MemWM: Memory-Augmented Text-Based World Model
MemWM enhances text-based world models with memory, boosting SSF by up to 206.3% and success rate by 65.4% in multi-environment benchmarks.
Key Findings
Methodology
MemWM integrates a curated world memory bank containing transition rules, state caches, and hard-to-predict facts, conditioned via retrieval mechanisms to improve next-state generation. It employs a structured evaluation metric, Structured State Fidelity (SSF), to quantify factual state preservation. The approach involves freezing the policy model and enhancing prediction through memory retrieval and task-specific skills. Experiments across ALFWorld, WebShop, and ScienceWorld demonstrate SSF improvements up to 206.3%, with task success rates increasing by 65.4%. Sensitivity analyses confirm robustness under varying memory and action budgets, showing that retrieved memory significantly aids in maintaining factual accuracy and decision efficiency.
Key Results
- Across all environments, MemWM outperforms baselines in SSF, with the highest improvement of 206.3% on ScienceWorld. Task success rates increased by 15-20 percentage points in ALFWorld and WebShop, reaching a 65.4% relative gain. Removing retrieved memory from the prediction prompt caused significant SSF drops, validating its importance. The model maintained stable performance across different action budgets, with no notable increase in steps, indicating enhanced decision efficiency. These results highlight the effectiveness of memory retrieval in improving factual consistency and planning success.
Significance
This work addresses the critical challenge of factual fidelity in text-based world models, offering a scalable solution that significantly improves environment state accuracy and downstream task success. By integrating task-specific memory and structured evaluation, it advances the reliability of language models in complex interactive environments, bridging the gap between surface-level text similarity and true environmental understanding. The approach paves the way for more trustworthy AI agents capable of precise reasoning and planning, with broad implications for robotics, scientific simulation, and intelligent assistants.
Technical Contribution
The paper introduces a novel memory-augmented architecture that conditions next-state predictions on curated, task-relevant memory entries, reducing factual errors. It proposes SSF as a domain-aware metric for factual fidelity, overcoming limitations of traditional overlap metrics. The method enables high-fidelity state imagination without additional policy training, using frozen policy models. The retrieval mechanism is tailored to environment-specific facts, ensuring relevance and precision. These innovations collectively push the frontier of reliable, factually consistent text-based world modeling, with theoretical guarantees on factual preservation and practical improvements validated across multiple benchmarks.
Novelty
This is the first systematic integration of a task-oriented memory bank with a structured factual evaluation metric (SSF) in text-based world models. Unlike prior approaches focusing solely on surface similarity or external knowledge retrieval, MemWM emphasizes preserving environment-specific facts and rules during imagination. Its retrieval-conditioned prediction mechanism and structured evaluation framework set new standards for factual fidelity, enabling models to generate more reliable and environment-aligned future states. This paradigm shift opens new avenues for robust planning and reasoning in language-based agents.
Limitations
- The reliance on a static, curated memory bank may limit adaptability in highly dynamic or novel environments, where unseen facts or rules are frequent. The retrieval mechanism's effectiveness depends on the quality and coverage of the memory, which may require manual curation or sophisticated updating strategies.
- Computational costs associated with retrieval and memory management could hinder real-time applications, especially in large-scale or multi-modal environments. Additionally, the approach's scalability to real-world, noisy data remains to be validated.
Future Work
Future research will focus on developing dynamic, self-updating memory modules that adapt to environment changes, possibly via reinforcement learning or continual learning techniques. Extending the framework to multi-modal settings, integrating visual and auditory data, will broaden its applicability. Further, exploring end-to-end training of retrieval mechanisms and policy models could enhance efficiency and robustness. Investigating transferability across diverse environments and tasks will be essential for deploying these models in real-world applications such as autonomous robots and intelligent assistants.
AI Executive Summary
The evolution of natural language processing has led to increasingly sophisticated text-based world models, which underpin intelligent decision-making in interactive environments. However, despite advances, these models often generate fluent but factually inaccurate future states, undermining planning reliability. This challenge is rooted in the difficulty of preserving environment-specific facts, rules, and attributes during imagination, especially when relying solely on surface-level text similarity metrics like BLEU or Word F1.
Addressing this, the authors propose MemWM, a memory-augmented text-based world model that leverages a curated world memory bank. This bank contains transition rules, state caches, and hard-to-predict facts, which are retrieved and conditioned during next-state prediction. The core innovation lies in the Structured State Fidelity (SSF) metric, explicitly measuring the preservation of environment-critical facts, rather than superficial text overlap.
Experimental validation across three benchmarks—ALFWorld, WebShop, and ScienceWorld—demonstrates that MemWM significantly enhances state fidelity, with SSF scores improving up to 206.3%. Correspondingly, task success rates increase by up to 65.4%, despite keeping the underlying policy model frozen. Sensitivity analyses confirm robustness across different memory and action budgets, indicating that the improvements stem from effective retrieval and fact preservation rather than increased exploration.
This work marks a substantial step forward in reliable, factually consistent text-world modeling, offering a scalable approach that can be integrated into various decision-making systems. Its implications extend to robotics, scientific simulations, and intelligent assistants, where factual accuracy is paramount. Future directions include dynamic memory updating, multi-modal integration, and end-to-end training strategies, aiming to further enhance adaptability and real-world deployment potential.
Deep Dive
Plain Language Accessible to non-experts
想象你在厨房做饭，每次你都要记住食材的位置、用的调料和步骤。你会用一个笔记本，把所有重要的细节写下来，这样每次做菜都能按照记忆中的步骤来，不会忘记放盐或关火。现在的机器人也是这样，它们用一种叫‘记忆库’的东西，把厨房里的规则和细节存起来。每次它们要做菜时，就会从记忆库里找相关信息，确保每一步都正确。这样，它们做出来的菜就像你亲手做的一样好，特别可靠。这种方法让机器人变得更聪明，也更像人一样会记忆和思考。
ELI14 Explained like you're 14
想象你在玩一个超级复杂的游戏，比如冒险游戏，你要记住很多线索和规则，比如哪个门可以打开，哪个宝藏在什么地方。每次你遇到新任务时，你会用之前记住的线索来帮忙。现在，科学家们也在教电脑用类似的方法：让它们记住一些重要的规则和细节，这样它们在预测未来发生的事情时，就不会忘记重要的细节。比如，电脑会记住哪个物品在厨房哪个位置，或者哪个按钮能打开门。这样，它们就能更准确地想象未来的场景，做出更聪明的决定。这个技术就像你用笔记本记下所有重要的线索一样，让电脑变得更聪明、更可靠。
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.