Building Spatial World Models from Sparse Transitional Episodic Memories¶

Conference: ICLR2026 arXiv: 2505.13696 Code: To be confirmed Area: Robotics Keywords: world model, episodic memory, spatial reasoning, cognitive map, navigation

TL;DR¶

This paper proposes the Episodic Spatial World Model (ESWM), which constructs spatial world models from sparse, disconnected episodic memories (one-step transitions). The model's latent space spontaneously gives rise to cognitive maps aligned with environmental topology, supporting zero-shot exploration and navigation.

Background & Motivation¶

Background: Existing world models typically require long sequences of continuous trajectories for training, encoding environmental knowledge into model weights. Representative approaches such as TEM and GTM-SM rely on continuous observation sequences and assume a fixed shared structure across environments.

Limitations of Prior Work: (1) In real-world scenarios, an agent's observations are often fragmentary—accessing different parts of an environment at different times without continuous long trajectories; (2) Environments may undergo structural changes (e.g., newly added obstacles), requiring weight-based models to be retrained for adaptation; (3) Sequence models incur prohibitive computational costs when scaling to large environments.

Key Challenge: Existing models encode structural knowledge of the environment into weights, making it (a) impossible to build maps rapidly from fragmented experience, and (b) infeasible to dynamically adapt to environmental changes.

Goal: Can a consistent spatial world model be efficiently constructed from only sparse, disconnected episodic memories?

Key Insight: Inspired by neuroscience—the medial temporal lobe (MTL) is responsible for both spatial representation and episodic memory, constructing relational networks by integrating overlapping episodic memories. The authors hypothesize that a model can infer complete spatial structure from a set of independent one-step transitions, without requiring continuous trajectories.

Core Idea: Reformulating world modeling from sequential learning to set-based reasoning—using a Transformer to infer spatial relationships from a collection of disconnected episodic memories.

Method¶

Overall Architecture¶

ESWM takes as input a memory bank \(M\) (an unordered set of disconnected one-step transitions \((s_s, a, s_e)\)) and a partially masked query transition \(q\) (with one of the start state, action, or end state randomly masked). The model's objective is to predict the masked element. This is essentially a set-to-value reasoning problem: inferring unobserved spatial relationships from fragmented memories.

Training adopts a meta-learning strategy: each sample randomly draws an environment, a memory bank, a query, and a masking scheme, preventing the model from memorizing specific environments and forcing it to learn general spatial reasoning capabilities.

Key Designs¶

Memory Bank Construction:
- Function: Generate a set of transitions for each environment that covers all locations without forming continuous trajectories.
- Mechanism: The memory bank satisfies three properties—disconnectedness (transitions do not form continuous paths), coverage (the graph induced by transitions is connected and covers all locations), and minimality (removing any single transition disconnects the graph).
- Design Motivation: The minimality constraint forces the model to perform multi-step reasoning (inferring unobserved spatial relations from multiple memory fragments) rather than simple table lookup.
Masked Prediction Task:
- Function: Randomly mask one component of the query transition (\(s_s\), \(a\), or \(s_e\)) and predict the masked value.
- Mechanism: \(q^* = f(M, q)\), where \(f\) is a Transformer encoder. The three components of each transition are projected into a shared high-dimensional space and averaged into a single token; tokens from the memory bank and the query are concatenated and fed into the Transformer, with three linear heads predicting \(s_s\), \(a\), and \(s_e\) respectively.
- Design Motivation: The three masking types test the model's capabilities for "forward prediction" (predicting the next state given state and action), "action inference" (inferring the action given start and end states), and "backward inference" (inferring the start state given action and end state).
Uncertainty Classification (I don't know):
- Function: When parts of the memory are missing, the model must determine whether a query is answerable.
- Mechanism: During training, a random subset of memories is deleted, creating unobserved regions in the environment. For queries involving unobserved regions, the model is trained to output an additional "I don't know" class.
- Design Motivation: This serves as the foundation for the exploration algorithm—the agent can leverage high "I don't know" probability to identify actions with maximum information gain.
Architecture Comparison:
- Transformer (ESWM-T), LSTM (ESWM-LSTM), and Mamba (ESWM-MAMBA) are compared.
- Key Findings: Only the Transformer succeeds in Open Arena (which requires compositional generalization); LSTM and Mamba overfit. This demonstrates that the attention mechanism—analogous to classical content-addressable memory—is critical for learning generalizable world models from episodic memories.

Loss & Training¶

Cross-entropy loss is used with equal weighting across the three prediction heads (\(s_s\), \(a\), \(s_e\)). Training runs for 460K iterations with a batch size of 128 and cosine learning rate scheduling. The meta-learning setup ensures that the environment, memory bank, and query are randomly generated for each sample.

Key Experimental Results¶

Main Results¶

Environment	Model	State Prediction Accuracy	Action Prediction Accuracy	vs. TEM-T
Open Arena	ESWM-T-2L	~85% (\(s_s\)), ~85% (\(s_e\))	~95% (\(a\))	TEM-T significantly lower
Random Wall	ESWM-T-14L	High accuracy	High accuracy	TEM-T completely fails (cannot handle structural changes)
MiniGrid 9×9	ESWM-T-12L	Successful prediction	Successful prediction	—
ProcThor 3D	ESWM-T-12L	High cosine similarity	Accurate \(\Delta xy\), \(\Delta\theta\) prediction	—

Downstream Task Performance¶

Task	Metric	ESWM	EPN (baseline)	Best Oracle
Exploration (15 steps)	Unique states visited	+16.8% vs. EPN	—	96.48% of Oracle
Navigation	Success rate	96.8%	78.8% (+18%)	—
Navigation	Path optimality	99.2%	78.2% (+21%)	—
Adaptation (navigation with added obstacles)	Success rate	93%	72%	baseline drops to 56%

Key Findings¶

The Transformer's attention mechanism is critical for learning spatial reasoning from episodic memory sets; LSTM and Mamba fail in Open Arena where compositional generalization is required.
ESWM's latent space spontaneously gives rise to spatial maps consistent with environmental topology (ISOMAP projections reveal smooth manifolds with local discontinuities corresponding to obstacle regions).
Path length is highly correlated between latent and physical space (\(R^2 = 0.89\)).
The model's predictive uncertainty (output entropy) increases monotonically with the length of the memory integration path required by the query, demonstrating that the model genuinely performs multi-step reasoning.
ESWM achieves superior navigation using only 1/4 of EPN's memory, reflecting high sample efficiency.

Highlights & Insights¶

From Set Reasoning to Spatial Maps: Reformulating world modeling from sequential processing to set-based reasoning is the core innovation. The model requires no continuous trajectories—only a set of independent transition memories—substantially reducing data requirements while naturally supporting dynamic environments.
Decoupling Memory and Reasoning: Environmental knowledge is stored in an external memory bank rather than model weights, enabling true "plug-and-play" adaptation—modifying a few memories suffices to adapt to environmental changes without retraining. This design principle transfers to any scenario requiring rapid adaptation.
Spontaneous Emergence of Cognitive Maps: Although the model is never explicitly supervised to learn spatial structure, its latent space naturally forms a geometric map consistent with environmental topology after training—closely paralleling findings on hippocampal place cells in neuroscience.
Zero-Shot Downstream Capabilities: Both exploration and navigation require no additional training; near-optimal policies are achieved directly by leveraging the world model's predictions and uncertainty estimates.

Limitations & Future Work¶

Experiments are primarily conducted in controlled discrete or simple continuous environments; validation in real robotic scenarios has yet to be demonstrated.
The ProcThor experiments demonstrate feasibility only, without comparison against strong baselines.
The minimality constraint on the memory bank is difficult to satisfy in practice—real agent memories typically contain redundancy and noise.
The current framework models only spatial structure and does not incorporate semantic information (e.g., object categories or functional attributes).
Meta-learning incurs high training costs (460K iterations), and the distribution of pretraining environments may affect generalization.

vs. TEM (Whittington et al.): TEM assumes all environments share a unified structural template encoded into RNN weights; ESWM makes no such assumption, dynamically inferring structure from external memory and handling structurally diverse environments (e.g., random mazes). TEM completely fails on Random Wall.
vs. GTM-SM (Fraccaro et al.): GTM-SM similarly relies on sequential trajectories and assumes shared structure; ESWM operates on disconnected episodic memories and is substantially more sample-efficient.
vs. Ha & Schmidhuber (2018) World Models: Traditional world models encode knowledge into weights and cannot rapidly adapt to environmental changes; ESWM's external memory mechanism enables immediate adaptation.
This work introduces a new paradigm for embodied AI and robot navigation: rather than requiring extensive training in the target environment, a usable spatial model can be built from only a small number of exploratory memories.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ — Reformulating world modeling from sequential learning to set-based reasoning represents a conceptually significant breakthrough.
Experimental Thoroughness: ⭐⭐⭐⭐ — Progressive validation from simple grids to 3D environments with thorough analysis, though real-world validation remains insufficient.
Writing Quality: ⭐⭐⭐⭐⭐ — Logically clear, visually polished, with a natural integration of neuroscientific motivation and computational methodology.
Value: ⭐⭐⭐⭐⭐ — Proposes a broadly influential new paradigm and makes important contributions at the intersection of cognitive science and AI.