Building spatial world models from sparse transitional episodic memories
Authors: Zizhan He
Maxime Daigle
Pouya Bashivan
Venue: ICLR 2026 Poster (2026)
Keywords: Spatial Representations
World Models
Episodic Memory Models
Transformers
Navigation
Primary area: applications to neuroscience & cognitive science

Abstract: Many animals possess a remarkable capacity to rapidly construct flexible cognitive maps of their environments. These maps are crucial for ethologically relevant behaviors such as navigation, exploration, and planning. Existing computational models typically require long sequential trajectories to build accurate maps, but neuroscience evidence suggests maps can also arise from integrating disjoint experiences governed by consistent spatial rules. We introduce the Episodic Spatial World Model (ESWM), a novel framework that constructs spatial maps from sparse, disjoint episodic memories. Across environments of varying complexity, ESWM predicts unobserved transitions from minimal experience, and the geometry of its latent space aligns with that of the environment. Because it operates on episodic memories that can be independently stored and updated, ESWM is inherently adaptive, enabling rapid adjustment to environmental changes. Furthermore, we demonstrate that ESWM readily enables near-optimal strategies for exploring novel environments and navigating between arbitrary points, all without the need for additional training. Our work demonstrates how neuroscience-inspired principles of episodic memory can advance the development of more flexible and generalizable world models.

Published as a conference paper at ICLR 2026
BUILDING SPATIAL WORLD MODELS FROM SPARSE TRANSITIONAL EPISODIC MEMORIES
Zizhan Het.3 Maxime Daigle Pouya Bashivan 1.2.3
1 Department of Computer Science, McGill University
2 Department of Physiology, McGill University
3 Mila, Université de Montréal
ABSTRACT
Many animals possess remarkable capacity to rapidly construct flexible cogni- tive maps of their environments. These maps are crucial for ethologically relevant behaviors such navigation, exploration, and planning. Existing computational models typically require long sequential trajectories to build accurate maps, but neuroscience evidence suggests maps can also arise from integrating disjoint ex- periences governed by consistent spatial rules. We introduce the Episodic Spa- tial World Model (ESWM), novel framework that constructs spatial maps from sparse, disjoint episodic memories. Across environments of varying complexity, ESWM predicts unobserved transitions from minimal experience, and the geom- etry of its latent space aligns with that of the environment. Because it operates on episodic memories that can be independently stored and updated, ESWM is inher- ently adaptive, enabling rapid adjustment environmental changes. Furthermore, we demonstrate that ESWM readily enables near-optimal strategies for explor- ing novel environments and navigating between arbitrary points, all without the need for additional training. Our work demonstrates how neuroscience-inspiee principles of episodic memory can advance the development of more flexible and generalizable world models.
INTRODUCTION
When visiting a new city, we quickly piece together scattered memories an alley passed on a walk, glimpse of landmark from taxi-intoo mental map that lets navigate, explore, and plan. This ability to build coherent model of the world from fragmented experiences, and to use it flexibly for imagining new routes outcomes, is a core feature of human intelligence that is fundamental to planning, problem-solving and decision-making all of which are crucial for survival (Bennett, 2023).
How do humans achieve this extraordinary feat? Insights from neuroscience and psychology provide compelling evidence for the neural mechanisms underpinning this capacity. One pivotal discovery the dual role of a key brain area called the medial temporal lobe (MTL) in both spatial represen- tations and episodic memories. Individual neurons within the MTL exhibit spatial selectivity, firing in response to specific locations "eefee 1976 Hafting et al. 2005), while collectively forming population code that represents animal's instantaneous position in space (Wilson & McNaughton 1993]. In parallel, MTL plays a critical role in the formation of episodic memories, enabling the rapid acquisition of associations between arbitrary stimuli (Howard & Eichenbaum 2015). Strik- ingly, lesions to this area result in profound deficits: individuals lose not only their ability to explore and navigate effectively (Kimble) 1963) but also their capacity to form episodic memories (Scoville) & Milner 1957) and imagine novel scenarios (Hassabis al. (007). These converging findings strongly suggest that the MTL constructs structured relational networks, integrating overlapping episodic memories into cohesive framework (Eichenbaum 2004 McKenzie et al., 2014) (See Ap- pendix B for more a detailed review). This mechanism may underlie a flexible and adaptive world model, allowing for inferences about spatial relationships beyond immediate perception.
Despite these insights, most prior computational approaches to world modeling have relied on learn- ing from continuous sequences of recent observations and actions (Ha & Schmidhuber 2018 Hatner

---

Published as a conference paper at ICLR 2026
et al.
2019 Whittington et al. 2020 Levenstein et al. 2024 Bar et al.
2025 Fraccaro al.
2018
fixed hand designed circuitry for representing and updating spatial information (Wan et al.
2023
2024), allocentric observations Khan et al. (2018), and itera- tive update procedures that hinders their ability to quickly adapt to environmental changes( Wang et al., 2023 Chandra et al. 2025 Stachenfeld al.) 2017). These models are typically trained on large datasets collected from a given environment, with knowledge encoded in their weights, While effective in certain contexts, such approaches face inherent limitations. For instance, agents may encounter different parts of an environment at widely separated times, visit states along arbitrary and unrelated trajectories that do not form continuous sequences, or experience environments that undergo structural changes (A). These challenges hinder the abilitty of sequence- models to generalize effectively environments where observations are disjoint and episodic.
Chandra et al. 2025: Kymn et al.
This raises a key question: Can a neural network (NN) rapidly construct coherent spatial map of the environment using only sparse and disjoint episodic memories? Drawing inspiration from the neuroscience of memory and navigation, we introduce the Episodic Spatial World Model (ESWM), a framework designed to infer spatial structure from independently acquired episodic experiences.
ESWM meta-trains the model to predict unseen transitions from a sparse set of one-step transitions (referred to as the memory bank) sampled across diverse environments. This is in contrast with conventional world models in two ways: 1) ESWM does not require sequential observations for predictions-the memory bank can contain transitions from completely different episodes. This is crucial in large environments where single walk spanning all locations might be long and therefore computationally expensive to process. In comparison, ESWM's memory bank can be constructed to contain only the useful transitions. 2) Each transition in the memory bank can be stored and updated independently. As result, ESWM can rapidly adapt to environmental changes (e.g, addition of obstacles) by simply editing a handful of transitions that are related to that change (see Section (4.6).
Our main contributions are: 1) We introduce ESWM, a framework that models spatial environments from sparse, episodic transitions. 2) We find that ESWM's latent space forms a geometric map reflecting the environment's s topology, which dynamically adapts to new memories and structural changes like obstacles. 3) We demonstrate that this learned model enables near-optimal, zero-shot exploration and navigation in novel environments without any task-specific training. 4) Finally, show that ESWM is scalable to realistic continuous environments with high-dimensional inputs.
Nemory Bank
Eatoodlc Spatial World Model
192
Tensturner Blcc
Tonsoormn Block
Trirdwess Bleck
$2
150
Cuery Chotcee
Presiction
Figure I: Episodic Spatial World Model. a) Three common scenarios that hinder the ability of typical world models to generalize effectively. (top) Observation of the full environment may take many time steps leading to long sequences; (middle) Environment structure may change across tri- als; (bottom) Information about a particular environment may be collected across separate exposures to the environment and not within a single one. b) Memory bank and query selection in square grid environment. c) Architecture of ESWM and training procedure. Model input consists of a bank of transitional memories (corresponding to the black arrows in (b)) and single query (q arrow in (b)), with either start-state, action, or end-state randomly masked with equal chances. The sequence of transitions is processed by a sequence model (e.g, Transformer encoder block) and the model parameters are updated to output the correct value for the masked component.
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORKS
2.1 WORLD MODELS
Humans excel at forming mental models of the world from egocentric observations and action, enabling efficient interaction with complex, dynamic environments. Inspired by this cognitive abil- ity, prior works have used NN to forecast future outcomes based on past and present states (Ha & Schmidhuber 2018) using methods such as image reconstruction (Watter al.
2015:
Alonso al.
2024) in domains like video games (Hafner et al. 2020 Schrittwieser et al.,
2020 Kaiser et al., 2019) and robotic control tasks (Gupta al., 2022; Ebert et al.] 2018). spatial settings, agents
have been trained to predict future sensory observations following arbitrary actions, typically using
recurrenll-cooneeteed networks (Stachenfeld et al. 2017 Whittington et al. 2020 Levenstein et al,
2024
Fraccaro et al. 2018).
More recently, the emergence of transformer architectures (Vaswani et al. 2017) has opened new possibilities for improving world models. Recent works such as TWM (Robine et al. 2023 STORM (Zhang et al. 2024). and IRIS (Micheli et al.] (022) have leveraged transformers based world models to improve performance unlike our proposed model, these models train extensively on a single environment and absorb the knowledge about that environment in the model's weight.
sample-limited benchmarks like Atari 100k. However,
Most notably, our work is closely related to Generative Temporal Models with Spatial Memory (GTM-SM) Fraccaro al. (2018) and the Folman-EEchenbel Machine (TEM) (Whittinttn al.) 2020 2022). the latter of which has been proposed as model of medial temporal lobe function. These models factorize sensory and structural information. The shared structure encoded in recurrent network weights, while environment-speccffc observations are stored in Hebbian-slot- memory, self-attention module, or differentiable memory dictionary. While both of these model classes include content addressable memory modules, important distinctions separate them from our proposed model. First, GTM-SM and TEM both presuppose a common structural template across environments (e.g, open 2D grid), whereas ESWM does not. Second, unlike these models, ESWM does not embed structural knowledge in its weights but instead infers it dynamically from external memories, enabling it capture distributions of environments with varying structures (e.g, a set of 2D mazes) where these models fail (Fig. ( Finally, while GTM-SM and TEM operate over sequential trajectories, ESWM integrates disjoint experiences, making it more sample-effccien and better suited to rapid adaptation, such as handling newly introduced obstacles. See Appendix for a more comprehensive comparison with other MTL models.
2.2 NAVIGATION AGENTS
The concept of cognitive map, an internal allocentric representation of the space, O'keete & Nadel (1978) has been thought enable navigation and spatial reasoning. Several studies have demon- strated that similar spatial representations emerge in NN trained on navigation tasks (Gornet & Thomson 2024). For example, [Banino et al. (2088) and Cueva & Weil (2018) showed that spatial representations arise in recurrent neural networks trained for path integration, while in parallel, NN lacking explicit mapping modules have achieved high performance on navigation tasks in unseen environments (Wijmans et al.) 2019; Reed et al. 2022; Khandelwal et al. 2022] (haah et al.) 2023). Wijmans et al. (2023) recently showed that spatial representations can emerge in navigation-agnose NN architectures trained for navigation, emphasizing memory's 's role in shaping these spatial repre- sentations. This raises a key question: Can spatial maps also emerge in
purpose models not explicitly trained for navigation and operating entirely on finite set of episodic memories?
METHODS
3.1 EPISODIC SPATIAL WORLD MODEL
Formally, our Episodic Spatial World Model (ESWM) function f that takes as input a memory bank M and a partially masked transition g. and predicts the complete transition q":
s(M,q)
3

---

Published as a conference paper at ICLR 2026
We define an episodic memory one-step transition in a spatial environment, represented by the tuple ($,, a, denoting the source state, action, and end state, respectively. For each environ- ment, we construct memory bank M as unordered set of such transitions, constructed to be: Disjoint, transitions do not form a continuous trajectory; Spanning, the transitions collectively form connected graph covering all locations; Minimal, removing any transition would disconnect the graph (Fig. [p). However, Minimality is not required for training ESWM but is imposed in the grid environment to prove that the absolute minimum information is sufficient to construct functional spatial map (see Fig. [9 for results on non non-minimal memory banks), We algorithmically generate such memory banks to sparsely cover each environment (Extended Methods C.3). Importantly, a single room may yield multiple valid memory banks, enabling our meta-training setup.
We formulate spatial modeling as the ability to infer any missing component of an unseen but plau- sible transition g-that is, a query transition not present in M (Fig. [D). The task is to predict the masked element (either 8g. a, or 82) of q. given the other two and the memory bank M (Fig.Te). For example: predicting the next state from s, = 5 and a move right; inferring the action connecting states 5 and 10; or recovering the start state given a = move right and se = 10.
To train ESWM, we adopt meta-learning approach where on each trial (i.e. sample), we randomly sample: (1) An enviromment e (though random assignment of states to locations); (2) A disjoint memory bank M from e; (3) A plausible query transition q from (e MM; (4) A masking choice (randomly masking 8s, a, or 8z). This process prevents the model from memorizing any specific environment since the state-location mapping is randomized across trials (i.e, the model cannot simply memorize that state "5" is at top-left because "5" could be at top-right the next trial), and our environments (Open Arena, Random Wall; detailed below) contain an intractable amount of configurations (> 10t3, see Extended Methods C.2 Fig.
3.2 MODEL SELECTION AND TRAINING
We considered three NN architectures for f. 1) Encoder-only transformer with no positional encod- ing Vaswani et al., 2017), ESWM-T, 2) long short-term memory (LSTM) (Hochreiter & Schmid- huber 1997). M-LTM,,,M, and Mamba (Gu & Dao 2024), ESWM-MAMBA, Each transition component (s,,, a, se) is projected separately into shared high-dimensional space and averaged to form a single input token.
The transformer model receives as input the concatenated sequence of memory bank and query tokens [b). The LSTM and Mamba models receive a randomly permuted sequence of memory tokens, each with the query token added. Three separate linear heads read out the model's 's prediction on the three parts (ss, a, Sc), from the final layer's 's last token. We use cross-entropy loss between predicted and actual states and actions. Equal weight is assigned to each prediction (8x, a, 8c)- Training is done over 460k iterations with a batch size of 128 and cosine learning rate schedule.
A spatial world model that assumes complete memory coverage is rather constrained. Instead, should reason over partial observations. To achieve this, we create trials where some memories are deliberately removed. The model is expected to solve queries within observed regions and classify unsolvable ones (e.g., transitions involving unobserved regions) under an additional "I don't 't know" category (Fig.
3.3 MODEL EVALUATION
We evaluate ESWM under three settings of varying complexity:
1. 2D grid environments: We evaluate ESWM in family of discrete, hexagonally structured grid environments (Extended C.22 Each location in the grid environment is assigned an integer-valued observation, drawn from a predefined set of states. An environment is uniquely defined by its mapping from states to spatial locations, allowing us to generate a vast number of distinct rooms by randomly shuffling this mapping-critie The hexagonal layout allows for six possible actions per location, offering richer connectivity than traditional square grids. We design two types of environments to test generalization: (1) Open Arena to evaluate generalization to unseen states and; (2) Random Wall to evaluate generalization unseen structures.
for our meta-learning objective.
4

---

Published as a conference paper at ICLR 2026
bj
Product SCUEEM
Prodect ACI ON
Pradict Tergat
Dw ooa Prodict
Thin
Tramoon varr
Yout OR UnSeET S3809
Traniin
Teall On LETSOOT
Traiiig Data Portt Soon
Figure 2: Evaluation accuracy in a) Open Arena and b) Random Wall. In both schematics, blue arrows are transitions in the memory bank and the q-labeled arrows are examples of query transitions q = (s,, 8c). ) Open Arena, the states are represented as 6-bit binaries and the model is tested in environments entirely filled with unseen states, while in Random Wall, the states are integers and the model is tested on its generalizability to unseen wall patterns. In addition, a random subregion is masked in Random Wall and the query can be either unsolvable, unseen, or seen, depicted as q q-labeled arrows with different colors. The Random Wall displayed here is a scaled down version of what the model is trained on (19 v.s 36 locations).
Each room consists of 19 locations and 64 possible state values. These states are represented compositionally as 6 binary bits (Fig.2p) enabling us to train and test ESWM on mutually exclusive sets of states (e.g, train on odd states and test on even states). For action prediction, the model performs 6-way classification (one for each action). For source and end state prediction, the task is broken into six independent binary classification problems one for each bit. This setup forces it to build representations that support spatial reasoning even in environments filled with previously unseen but compositionally generated state observations.
Open Arena:
Random Wall: To test for generalization unseen structures, this environment features larger rooms (36 locations) with unique integer states and a randomly shaped wall of ob- stacles (Fig. 2)). During testing, the model must generalize to novel wall configurations and state-location mappings. The prediction task is 6-way classification for actions and 36- way for states. We simplified the observations to integers to focus the task on inferring the environment's underlying structures.
2. MiniGrid environments: We used MiniGrid Shevalier-Boisverr et al., 2023) to create environ- ment layouts of size Each state consist of egocentric 5 x 5 view of the environment containing non-unique col- ored objects. The agent takes one of three egocentric actions: turn left, turn right, go straight. Each procedurally generated environment is assigned random color layout and wall patter.. The model has to solf-localize based on partial egocentric views of the environment experienced from varying head directions.
with egocentric multi-dimensional observations and actions (Fig. (6)
3. Simulated 3D indoor environments: We use ProcThor (Deitke et al. 2022), large set of proce- durally generated 3D indoor scenes. Each state consists of a high-resolution RGB first-person observation (224 x 224). The agent takes action defined by egocentric displacements Ary (along left/right and front/back axes) and rotation A6, both from a continuous range.
For MiniGrid and ProcThor environments, the memory bank consists of episodic memories that collectively cover the room, with partial overlap across observations enable memory integration. Importantly, the memory banks in these environments remain compact.
4 RESULTS
4.1 MODELING SPATIAL ENVIRONMENTS FROM EPISODIC MEMORIES
We first assessed ESWM on its core ability to predict missing elements of unseen transitions given a set of memories from environment and partially observable query. In both settings, we consider a closely related sequence-based based model TEM-T as baseline (Whittintton et al.) 2022). Instead of a set of disjoint memories, TEM-T receives a spanning trajectory the memory bank and is queried
5

---

Published as a conference paper at ICLR 2026
to predict unseen transition the end of the trajectory. To enable fair comparison, the spanning
trajectory is generated to be minimal (Extended Methods
We found stark contrast between the model types in Open Arena setting (Fig. 2a). The ESWM- models overfitted (see Fig. 8 and struggled reach above chance the validation set. This suggests that the attention mechanism in Transformer architectures sntiitectures-rrmiiiseeen
model excelled at learning this task while ESWM-LSTM and
classic content-addressalle key role in learning generalizable world models from episodic mem- ories, However, simpler, hand-engineered memory-lookup mechanisms like those in Coda-Forno al. (2022a) fail on this inference task (Fig. [7). highlighting the need for learnable attention- based mechanisms. Moreover, we also show that ESWM-T can also model spatial environments when trained on inputs distributions that lack the Minimality constraint (Fig. 9 see section [D.2 for com- paring ESWM-T trained with/without the Minimality constraint). In addition, the sequence- -ased transformer model TEM-T, while outperforming con-transformer ESWM models, still significantly underperforms relative to the transformer ESWM-T, highlighting the advantage of directly inferring from episodic memories, rather than segregated encoding of structure and memories, Furthermore, while ESWM-T clearly improves with larger model size, the performance in larger TEM-T models significantly degrades (Fig. b). For further discussions in this setting, see section
memory mod-
els (Kanerva 1988)
In the second scenario, the Random Wall, found that all ESWM models are capable of learn- ing the task (Fig. D). Although, there was a substantial gap between transformer-basee models and other models, even with equal number of parameters (LSTM-4L, T-2L, and MAMBA-5L). This showed that ESWM learned general strategy for modeling an environment from sparse and disjoint episodic memories, enabling to generalize across environments with different structures. Consis- tent with our prediction, TEM-T fails at modeling the set of environments with varying structures. This underscores the importance of direct structural inference from memories rather than encoding it in the model's weights.
Finally, we also considered environments with duplicate (non-unique) states. In the absence of historical context, one-step transition prediction prevents the model from disambiguating between duplicate states. However, this limitation can easily overeome by prepending a trajectory, serving as context, to the query. To show this, we train ESWM-T in environments with duplicated states and find that it capable solving the task in both Open Arena and Random Wall (see Fig.
a)
Dim
b)
d)
Figure 3: Spatial map emerges in ESWM's latent space. a), b) ISOMAP projections of ESWM- T-2L's activations for action prediction. Different columns the same projection from different viewing angles. a) Spatial map in the absence of obstacles and b) in the presence of obstacles (a straight wall). c), ISOMAP projections of ESWM-T-14L: activations. e) From left right: A sample room with two disconnected regions whose shapes match boundaries' shape; ESWM's latent space when memory bank observing either the top or bottom region is given as input; ESWM's latent space when memory bank observing both regions is given input. d) From left to right: A sample room containing two disconnected regions whose shapes give no cues about their relative position; ESWM's 's latent space when a memory bank of the room is given as input; Updated room with a new wall, the regions remain disconnected; ESWM's latent space when a memory bank of the updated room is given input. See Fig. 10 for more examples.
4.2 SPATIAL MAPS EMERGE ROBUSTLY WITHIN ESWM'S LATENT SPACE
To simulate unseen transition, ESWM must integrate disjoint memory fragments into coher- ent internal map that reflects both observed states and obstacles. Inspired by low-dimensional, environment-like neural manifolds in rodents' MTI Nakaia al. (2024). we applied ISOMAP

---

Published as a conference paper at ICLR 2026
baum et al. (2000) (Extended Methods C.6) to ESWM's latent space. The resulting manifolds vealed that ESWM consistently organizes its internal representations to mirror the spatial layout of each specific environment.
In obstacle-frec environments, the manifold adopts a smooth, saddle-like shape (Fig. 3d. When obstacles are introduced, they create localized discontinuities that correspond precisely to blocked regions in physical space (Fig. BD)- This structure remarkably robust and persists under out- of-distribution obstacle shapes (Fig. I0 partial observation of environments (Fig. 12 across architectures (LSTM, Transformer, Mamba; Figs TT and throughout all transformer layers, which become progressively more continuous with depth (Fig. 14). All prediction tasks ($g. a, se) yield similar maps (Fig. [2p). In the presence of disconnected memory clusters with unclear spatial relations, ESWM pieces its fragmented internal representation into coherent map by aligning them along the inferred wall shapes or boundaries' shapes (Fig. Br Additional Analysis (.3) Quan- titatively, path lengths in latent and physical spaces correlate strongly (R² 0.89, N path 1500,
Nenv 75; Fig. [6p). and a linear classifier on ESWM activations can accurately discriminate which of two states is farther from a reference state with 93.42 + 0.006% accuracy (chance 50%, Noeed 10; Extended Methods C.7). Furthermore, the sRSA (Levenstein al. 2024) of a trained ESWM significantly higher than that of an untrained one (trained = 0.78 = 0.02,
untrained 0.5 + 0.04, Neny 1000; paired (-test, /=172.17, p < 10-3; Extended Methods C.8).
Notably, such structured spatial representations do emerge in episodic agents trained for navi- gation, such as Episodic Planning Network (EPN; Ritter et al. (2020)), despite being trained with meta-RL on the same environments and data as ESWM (Fig. 10 Appendix H.
4.3 ESWM INTEGRATES OVERLAPPING MEMORIES
We ran several experiments test whether ESWM truly retrieves the information scattered across its memory bank. First, we find that ESWM's prediction mncertainty-mmasure by entropy of its output distribution probability-increede increases with the length of the integration path needed to solve a query (Fig. 4p). Second, we found that introducing shortcut memories that reduce this path length, significantly alters the prediction distribution (independent two-sample t-test; n 5000, p < 0.001; Fig.Ep) Fp). Third, we find that ESWM trained in size- 19 environments can generalize to substantially larger environments
size 37 in Open Arena, achieving lower, but still well-above-chance pre- diction accuracy (s,: 59%, 90%, se: 55%) unlikely if the model had simply memorized patterns. Fourth, ESWM's performance improves with denser, non non-minimal memory banks, despite being out-of-distribution in Random Wall (Fig. 4D). Lastly, despite being trained only clean data, ESWM develops some robustness to noisy memory banks. Its prediction accuracy only dropped by 1% and 10,, respectively, when tested memory banks containing contradictory state observa- tions and transition outcomes (n=5000; Extended Methods C.9). Altogether, these results suggest that ESWM indeed builds coherent spatial representations by integrating overlapping episodes.
4.4 EXPLORATION IN NOVEL ENVIRONMENTS
So far, we have assumed that ESWM has access to informative episodic memories. However, an- imals actively acquire such memories through exploration. The ability to rapidly map out novel environments- wwirooments--ideniffiing obstacles, rewards, and threats without prior knowledge is critical for survival. We show that ESWM enables autonomous exploration in unfamiliar spaces without any additional training.
We designed simple exploration algorithm based on ESWM's "I don't know" predictions (Alg Starting with an empty memory bank, the agent evaluates all possible actions at each step by query- ing ESWM. It selects the action with the highest predicted uncertainty (i.c., highest "I don't 't know" probability) and adds the resulting transition to memory. When ESWM is confident about all im- mediate actions, the agent performs multi-step lookahead by unrolling the model to identify frontier states with low prediction confidence (Yamauchi) 1997). A prediction deemed unconfident ESWM either outputs "I don't 't know" or assigns 80% probability to the predicted outcome.
Interestingly, the agent adopts an efficient zig-zag zag exploration strategy, only revisiting states when separated by long paths. In the Random Wa environment, ESWM explores 16.8% more unique
7

---

Published as a conference paper at ICLR 2026
(b)
(c)
Prediction Entropy
2.0
Predict Action
Predict Source
D.Overgencc
(Minimal) [Non-Minimal) ID G00D
Predict Target
Accuracy
Brodec Soarce
predict Actin
Predkt Target
Predict Start
Preditt Action
Informative
Predict End
Path Length
20 5) Memory Bank Size
Figure 4: ESWM integrates memories. a) X-axis is the shortest path length between the source and end states in the query, with edges corresponding to episodic memories in the memory bank, and the Y-axis is the entropy of ESWM's prediction probabilities (n=5000). The shaded region is 95% confidence interval computed over all predictions for each path length. b) KL divergence between ESWM prediction distributions before and after adding an episodic memory the memory bank. An informative episodic memory shortens the integration path required for the model to solve the prediction task, while non-informative episodic memory does not (n=5000). e) Prediction accuracy across memory bank sizes. In the gray region, memory banks minimally observe the environment; larger banks correspond to larger observed areas. In the blue region, the observed area is fixed (full environment), so larger memory banks correspond to more densely observed area. ID and OOD denote in-distribution and out-of-distribution settings, (n=2000), ESWM-T- 14L used for a) and b) while ESWM-T-12L. is used for c). All experiments in Random Wall.
1000, F 2803.63, p 0.05; Tukey's s HSD post hoc test: ESWM > EPN, < 10-3), , and achieves 96.48% of the performance of an oracle agent with full obstacle knowledge (Fig. 5h. Fig. [5p. When exploration continues until the memory bank reaches the max length seen during training, ESWM visits on average 91 + 2% of all unique states (Nrriols 1000). Moreover, the collected memory banks allow ESWM to predict 8,, a, S. with 97%, 93%, 98% of the accuracy of those achieved using procedurally generated memory banks (n 1000). Together, these results highlight ESWM's ability to support data-efficient, near-oracle exploration and mapping in novel environments.
states than EPN (one-way ANOVA;
(a)
(b)
(c)
25WM
Ravalire
FRWN
-a-- FPM
Succcso Bate 195
Optemalty
Succeess
0.
Exploration Oouerver
Path Length
Number Costacles
Figure 5: Exploration, navigation, and adaptability. a) Comparison of exploration strategies based on the number of unique states visited over 15 time steps in Random Wall. The optimal agent explores along the path found by the Traveling Salesman Algorithm over known free space. N=1000. b) Comparison of navigation success rate and path optimality over path lengths between EPN and FSWM, n=2400. с) Comparison of navigation success rates with increasing number of unexpected obstacles. The baseline agent is trained on the original environment and tested after structural changes, while ESWM navigates with memory banks from the original environment and needs to autonomously adapt to changes (n=100). The best seed for ESWM and EPN is used. We compare their performance across seeds in Fig. 15] Random Wall Experiments include 19 locations.
4.5 NAVIGATING WITH MINIMAL SPANNING EPISODIC MEMORIES
A key benefit of world model is the ability to plan by simulating trajectories. We test this capabil- ity by using the memory-banke eank-ee ESWM for planning in the Random Wall environment. The task is to find a path from source state Ssouce to goal state Sgaal using only the current ob- servation and the episodic memories the memory bank, without access to global information like coordinates complete map of obstacles.
8

---

Published as a conference paper at ICLR 2026
Our planning strategy relies ESWM's ability to infer the end state of any given transition. Ini- tialized at a state, an agent can imagine multiple steps into the future by querying ESWM on the consequences of all applicable actions, bootstrapping ESWM's own predicted end states, and tracking visited states until the goal state is reached (Alg TT. Then the agent can act out the imag- ined sequence of actions in reality.
Using this strategy, an ESWM-based navigational agent can navigate between arbitrary states in Random Wall with 96.8% success rate and 99.2% path optimality, outperforming EPN al. (2020) by 18% and 21 % respectively (Fig- 5p) (one-sided independent (-fest; n=2400: ESWM>>PN on optimality: p < 10-3), 3). Navigation is successful the agent reaches the target within 15 time steps, and path optimality is defined as the ratio between the shortest and actual path length. We also include a variant of such an algorithm in the appendix (Additional Analysis D.4) that biasedly considers the action that shortens the geodesic distance to the target on ESWM's latent spatial map, demonstrating how structured latent spaces in world models can serve as efficient heuristics for planning.
is important to emphasize that the navigation success primarily stems from the high-fidelity world models constructed by ESWM using minimal memory (requiring 4x fewer memories than EPN), rather than from the planning algorithm itself. Thus, the key finding is not that a classic search algorithm can outperform RL agent, but that world model learned from such sparse, disjoint data is accurate enough to enable it to do so.
A0ky B) = (move left 0 5mm, forward 1 1.5m, rotate right 272 degrens)
Start Stato
End State
Weldttin
(a)
(b)
Figure 6: ESWM in scaled-up environments a) (Top-Left) Examples of mooedurally-geneee MiniGrid environments. (Top-Right) Accuracies for Sa. and se predictions for ESWM-T-12L. (Bottom) ISOMAP analysis on ESWM's latent space shows map-like representation. Inset shows the color scheme across spatial locations. b) (Top-Left) Examples of procedurally-generate Thor environments. (Top-Right) A sample transition. States are high-dimensional images, and actions continuous 3-dimensional vectors. (Bottom) ESWM-T-I2L's performance in unseen test environments for action (measured by accuracy on sub-components Ary and Af) and end-state prediction (measured by cosine similarity to the frozen vision model's 's representation).
Proc-
4.6 ESWM ENABLES FAST ADAPTATION TO CHANGES
Unlike traditional world models that embed environment knowledge in their weights, ESWM de- couples memory from reasoning. This separation makes ESWM inherently adaptable to changes: when the environment changes, no retraining is needed-onle the external memory bank is updated. Because ESWM operates on sparse, independent one-step transitions, targeted edits to memory are straightforward We demonstrate this flexibility in dynamic navigation task. After introducing sub- stantial changes- suchngeesschhhaee as new obstacles an ESWM agent continues to navigate between arbitrary states with 93% success rate (Fig. 58)

---

Published as a conference paper at ICLR 2026
The adaptive agent works as follows: Given a (possibly outdated) memory bank M, it plans path from Sstart to 8gaal (Section 4.5). During execution, it compares predicted observations with actual ones. When mismatches occur (indicating environmental changes), the agent removes outdated transitions (where 8, or Se match the predicted observation from M), prompting exploration (Section 4.4 to fill the gaps. After acquiring
transitions, the agent replans with the updated memory.
ESWM's adaptability significantly surpasses both EPN and a baseline RL agent trained specifically for the original environment (see Extended Methods C0 for baseline details). For example, after adding a new wall consisting of multiple contiguous obstacles, ESWM maintains a 93% naviga- tion success rate. whereas EPN and the baseline drop to 72% and 56% respectively (see Fig. (p). Unsurprisingly, the baseline agent fails to form any structured latent spatial map.
4.7 SCALING ESWM
In the previous sections, we focused on simple grid environments, where episodic memories are tran- sitions between integer observations via allocentric actions (e.g.. go north), to examine how ESWM behaves under various scenarios. To test ESWM's scalability, we next evaluate it in Minigrid and ProcTThor (see Se. two more complex scenarios with egocentric high-ddmensional observa- tions and actions. To adapt ESWM to MiniGrid and ProcTThor, its operation on discrete bank of transitions remains unchanged, and we only adjusted the encoding function for states (trained from scratch for MiniGrid, pretrained vision model for ProcThor) and actions and the readout function for generating the predictions to match those in this environment (Extended Methods C.T.
In MiniGrid, ESWM continues to exhibit strong predictive performance (Fig. (a) Top- Top-Right) while (6ig.6l)
constructing coherent internal map representations (Fig. (6)) Bottom) that adapt to the environment's structure. In ProcThor, our results show that ESWM solving this challenging task by learning to predict both state vectors and action parameters (Fig. 6p).
DISCUSSION
In this work, we introduced ESWM, a NN model that can rapidly construct coherent spatial world model from sparse and fragmented one-step transitions. ESWM exhibits strong spatial reasoning and supports downstream tasks like exploration and navigation without explicit training for them. Unlike prior models that rely on fixed circuitry for spatial representation (Wang al.) 2023 Chandra et al. 2025 Kymn et al. 2024), ESWM learns both representation and update mechanisms directly from experience. It encodes new memories in a single shot, avoiding the iterative parameter tuning required by earlier approaches (Wang et al., 2023 Chandra et al., 2025 Stachenfeld et al., 2017 By storing transitional memories instead of state-observation pairs (Whitington al. [2020] Coda-Forno et al. 2022bb Whittington et al.) 2022), ESWM also adapts quickly to environmental changes, such as the addition or removal of obstacles an area where many existing models struggle.
A limitation is the use of highly controllable environments, which allowed to scale up the train- ing and control different aspects of the environment and data freely. Future work could explore extending this approach more realistic settings such as robotic navigation with natural, complex sensory information. Although theoretical, our work has potential societal impacts autonomous navigation and exploration, highlighting the need for safe, ethical, and responsible deployment (see Section l).
REPRODUCIBILITY STATEMENT
We provide detailed descriptions to ensure the reproducibility of our work. These include training details (Extended Methods C.I data generation pipeline (Extended Methods C.2 C.4), bascline implementations (Extended Methods C.3 C.10, Section[H), methods of analysis (Extended Methods
C.6 C.7) and algorithm pseudo codes (Section G
10

---

Published as a conference paper at ICLR 2026
7 DECLARATION OF THE USE OF LLM
We recognize the importance of LLM usage disclosure and declare that our use of LLM is strictly limited word editing and polishing.
ACKNOWLEDGEMENT
M.D. was supported by the UNIQUE PhD Fellowship. This research was supported by the NSERC Discovery grant RGPIN-2021-03035, and CIHR Project Grant PJT-191957. P.B. was supported by FRQ-S Research Scholars Junior grant 310924, FRONT-NSERC NOVA grant 2024-NOVA- 346823, and the William Dawson Scholar award. All analyses were executed using resources pro- vided by the Digital Research Alliance of Canada (Compute Canada) and funding from Canada Foundation for Innovation project number 42730.
REFERENCES
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos Storkey, Tim Pearce, and François Fleuret. Diffusion for world modeling: Visual details matter in atari. arXiv preprint arXiv:2405. 2024.
Jacob W. Bakermans, Joseph Warren, James C. R. Whittington, and Timothy E. Behrens. Con- structing future behavior in the hippocampal formation through composition and replay. Nature Neuroscience, pp. 1-12, March 2025. ISSN 1546-1726. doi: 10. 111010/0//1153--5--0000
nwwuratuueee.e com ture Publishing Group.
BtRL:ttps:://// htt.://ee....mmaaatcccc///////// arttcceeess419933 020- 01988----00--11
Publisher: Na-
Andrea Banino, Caswell Barry, Benigno Uria, Charles Blundell, Timothy Lillicrap, Piotr Mirowski, Alexander Pritzel, Martin Chadwick, Thomas Degris, Joseph Modayil, et al. Vector-basee navigation using grid-like representations in artificial agents. Nature, 557(7705):429 433, 2018.
Amir Bar, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun. Navigation world models. 2025. hRtps::t::///rr wwtp:///gggabb///..aa///// 03572
Max Bennett. A brief history of intelligence: evolution, AI, and the five breakthroughs that made our brains. HarperCollins, 2023.
Iva K. Brunec, Jessica Robin, Rosanna K. Olsen, Morris Moscovitch, and Morgan D. Barense. Integration and differentiation of hippocampal memory traces. Neuroscience & Biobehavioral Reviews, 118:196-200 208, November 2020. ISSN 01497634. doi: 10.1016/
jmeubiorew.2020.007.00 .244. URL htttps://linkigghht www.uiie..lleeee....m..... wtt.:/oo///reeeee wttp:e..cw/reeeeeeeeeee eis/8014976642203422
Sarthak Chandra, Sugandha Sharma, Rishidev Chaudhuri, and ll Fiete. Episodic and associative
memory from spatial scaffolds in the hippocampus. Nature, (381)))))7739 February 2025.
ISSN 1476-4687.
doi: 10. R..1008//41150 10.10388/411586-02--0000
0860054-8399... URL https: / www. naturat....... com/ Publisher: Nature Publishing Group.
articles/s41586- 020-1-588--4----9999 020-2-88922-1
Maxime Deovalie--Boivver Bolun Dai, Mark Towers, Rodrigo Vicente Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and Jordan Terry. Minigrid & miniworld: Mod- ular customizable reinforcement learning environments for goal-oriented tasks. Advances in Neural Information Processing Systems 36, New Orleans, LA, USA, December 2023.
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation, 2014. URL htll:/ttp///xxxi/// www.:i.i..orr.oooooo abs///sss/100 1078
Julian Coda-Forno, Changmin Yu, Qinghai Guo, Zafeirios Fountas, and Neil Burgess. Lever- aging Episodic Memory to Improve World Models for Reinforcement Learning. https: themari-workshop..o github. io/ htoppapess/pappr. 3. 3.fdf 2022a. [Accessed 09-05- 2025].

---

Published as a conference paper at ICLR 2026
Julian Coda-Forno, Changmin Yu, Qinghai Guo, Zafeirios Fountas, and Neil Burgess. Leveraging Episodic Memory to Improve World Models for Reinforcement Learning. NeurlPS MemARI Workshop, 2022b.
Christopher J Cueva and Xue-Xin Wei. Emergence of grid-like representations by training recurrent neural networks perform spatial localization. arXiv preprint
2018
Matt Deitke, Eli VanderBilt, Alvaro Herrasti, Luca Weihs, Kiana Ehsani, Jordi Salvador, Winson Han, Eric Kolve, Aniruddha Kembhavi, and Roozbeh Mottaghi. Procthor: Large-scale embodied ai using procedural generation, 2022.
Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual fore- sight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint arXiv:1812.00568. 2018.
Howard Eichenbaum. Hippocampus: Cognitive Processes and Neural Representations that Underlie Declarative Memory. Neuron, 44(1):109-120, September 2004. ISSN 0896-6273, doi: 10.10166 2ouron,2004.08. 028. URL htt.ss//www..eeeeeeeeeeee...o.
www www.::/ccenneeirec.... cww htt.://awc...occccccccaarraai.aa
p11750896627304005555
Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dumning, et Impala: Scalable distributed deep-rl with im- portance weighted actor-learner architectures. In International conference on machine learning, pp. 1407-1416. PMLR, 2018.
Dave Ferguson, Maxim Likhachev, and Anthony Stentz. A guide to heuristic-basee path planning. In Proceedings of the international workshop on planning under uncertainty for autonomous sys- tems, international conference on automated planning and scheduling (ICAPS), pp. 9-18, 2005.
Corey Fernandez, Jiefeng Jiang, Shao-Fang Wang, Hannah Lee Choi, and Anthony D Wagner. Rep- resentational integration and differentiation the human hippocampus following goal-directed navigation. eLife, 12:880281, February 2023. ISSN 2050-084X. doi: 10.7554/eLife, 10288 URL tions, Ltd.
2t1.111.555511111 80281
Publisher: eLife Sciences Publica-
Marco Fraccaro, Danilo Jimenez Rezende, Yori Zwols, Alexander Pritzel, M Ali Eslami, and Fabio Viola. Generative Temporal Models with Spatial Memory for Partially Observed Environments. ICMI,, 2018.
James Gornet and Matt Thomson. Automated construction of cognitive maps with visual predictive coding. Nature Machine Intelligence, 6(7):820-88) -333, 2024.
Albert Gu and Tri Dao. Mamba: Lincar- time sequence modeling with selective state spaces, 2024.
URL httttttt////iiii www. abs/ 2212. 00752
Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martin-Martin, and Li Fei-Fei. Maskvit: Masked visual pre-training for video prediction. arXiv preprint arXiv:2206 1894, 2022.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint 1rXiv:1803.10122 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1912.0169 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02193, 2020.
Torkel Hafting, Marianne Fyhn, Sturla Molden, May-Britt Moser, and Edvard Moser. Microstruc- ture of a spatial map in the entorhinal cortex. Nature, 436(7052):801 8060 August 2005. ISSN 1476-4687. doi: 10.1033./ URL https:/ www. nature com/ articles/ nature03721. Number: 7052 Publisher: Nature Publishing Group.
12

---

Published as a conference paper at ICLR 2026
Peter E Hart, Nils J Nilsson, and Bertram Raphael. A formal basis for the heuristic determination of minimum cost paths. IEEE transactions on Systems Science and Cybernetics, 4(2):100-107, 1968.
D. Hassabis, D. Kumaran, S. D. Vann, and E. A. Maguire. Patients with hippocampal amnesia cannot imagine new experiences. Proceedings of the National Academy of Sciences, 104(5):122-177) 2007. ISSN 0027-8424. doi: 10. 1073//0a.s....mm htt.ppaas..00000000000 0610561 104. URL http: /ttp:://wwwww.eo pnas org/
0rw.maas.0115661111 Number: 5.
Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory, Neural Computation, 9(8): 1997.9...1......... Conference
1735-1780, November 1997. ISSN 0899-7667.
doi: 10.
Name: Neural Computation.
Mare W. Howard and Howard Fichenbaum. Time and space in the hippocampus. Brain
Research, 1621:345-355 September 2015. ISSN 0006-8993, doi: 10.10166/.brainee
2014.10.069. URL https://www. www.:wsseeecdirec.oomm 500068993 4014917
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based reinforcement learning for atari. arXiv preprint urXiv:1903 003400374 2019.
Pentti Kanerva. Sparse distributed memory. MIT press, 1988.
Arbaaz Khan, Clark Zhang, Nikolay Atanasov, Konstantinos Karydis, Vijay Kumar, and Daniel D. Lee. Memory Augmented Control Networks, February 2018. URL http:: org/
abs/1700 05706 1rXiv:1109.. 05706 [cs].
Apoorv Khandelwal, I Luca Weihs, Roozbeh Mottaghi, and Aniruddha Kembhavi. Simple but effec- tive: Clip embeddings for embodied ai. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 14829-14838, 2022.
Daniel P. Kimble. The effects of bilateral hippocampal lesions in rats. Journal of Comparative and Physiological Psychology, 56(2):273-28)))) 1963. ISSN 0021-9940. doi: 10.1037/h0048903.
URL RtlLhttp::// http:://doo..pa.o.oo..... httpp/pp..........11111100///000 /11.//0..11170/000000000
Christopher J. Kymn, Sonia Mazelet, Anthony Thomas, Denis Kleyko, E. P. Frady, Friedrich Som- mer, and Bruno A. Olshausen. Binding in hippocampall-entorhin circuits enables composition- ality in cognitive maps. Advances in Neural Information Processing Systems, 37:39128-39157, December 2024. URL https::tt////oooootote .eg.riippppi.n................
nwwripuu.....m cc http:/cccppppppepeeeeeeeeeeppeee
html
Daniel Levenstein, Aleksei Efremov, Roy Henha Eyono, Adrien Peyrache, and Blake A. Richards. Sequential predictive learning is a unifying theory for hippocampal representation and replay, April 2024. URL https: www biorxiv.....v. org/ http:/ooteet/0-110//04.....22
591528v1
Pages: .2..04.28.591528 Section: New Results.
Sam McKenzie, Andrea J. Frank, Nathaniel R. Kinsky, Blake Porter, Pamela D. Rivière, and Howard
Eichenbaum. Hippocampal Representation of Related and Opposing Memories Develop within Distinct, Hierarchically Organized Neural Schemas. Neuron, 83(1):202-215, July 2014. ISSN 0896-6273. doi: 10. 10.10116j.neuron...o ot5t:///...rrro221.......0.000 001..05.019. URL URL/ttpss:/ www. www..cell.. com neuron/ abst
S0896 /S0t 62 73 (14) 00405-X
Publisher: Elsevier.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod-
els. arXiv preprint arXiv:2209.00588, 2022.
Volodymyr Mnih. Asynchronous methods for deep reinforcement learning. 783, 2016.
arXiv preprint
Neal W Morton, Katherine Sherrill, and Alison R Preston. Memory integration constructs maps of space, time, and concepts. Current Opinion in Behavioral Sciences, 17:161-168, October 2017. ISSN 23521546. doi:
000.0.0.......... 007. URL https: /tts/////nnginbbbb..
elsevi com otte://r.......eeeeeeeeeee/////a
13

---

Published as a conference paper at ICLR 2026
Shinya Nakai, Takuma Kitanishi, and Kenji Mizuseki. Distinct manifold encoding of naviga-
tional information in the subiculum and hippocampus. Science Advances, 10())cadi4441, 2024. doi: 10.1126sscidd.adi4444
URL https: www.w.cciece..om wwt./oigggooooeeeee
sci aiiidd... adv.adddd adi4444 4471
John O'Keefe. Place units in the hippocampus of the freely moving rat. Experimental
Neurology, 51(1)788-09, January 1976. ISSN 0014-4886.
doi: 10.10166001444886(66)
90055-8.
URL https: http://wwwwcieneeereeeeee www.:/cience//rrriicll//o///////
0014483676905558
John O' O'keefe and Lynn Nadel. The hippocampus as cognitive map. Oxford: Clarendon Press, 1978.
Seongmin A. Park. Douglas S. Miller, Hamed Nili, Charan Ranganath, and Erie D. Boor- man. Map Making: Constructing, Combining, and Inferring on Abstract Cognitive Maps. Neuron, (2266):222-11233 2020.06.030. URL https: www. www.scienceiirec..........mccee. com/ http://cieee/e//tteeee//////////
September 2020. ISSN 0896-6273. doi: 10.1016/j.neuron.
5089662732 0304840
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06 2022.
Sam Ritter, Ryan Faulkner, Laurent Sartran, Adam Santoro, Matt Botvinick, and David Raposo. Rapid task in novel environments. arXiv preprint S1Ki:::000.036666 2020.
William A. Roberts, Catherine Cruz, and Joseph Tremblay. Rats take correct novel routes and shortcuts in an enclosed maze. Journal of Experimental Psychology: Animal Behavior Pro- cesses, 33(2);99-91 2007. ISSN 1939-2184, 0097-7403. doi: 10.1037/0097-74033 doii
URL
https: /ooi apa. 0ttg:///66601.1113///000000.3333
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-bast world mod-
els are happy with 100k interactions. arXiv preprint arXiv:2303.07/1 109, 2023.
Margaret L. Schlichting, Jeanette A. Mumford, and Alison R. Preston. Learning-related represen- tational changes reveal dissociable integration and separation signatures in the hippocampus and prefrontal cortex. Nature Communications, 6(1):8151, August 2015. ISSN 2041-1723. doi: 10. Pub-
1038/ncomms9151, URL ltttps://wwtt http:/wwwwwtttttaa...m...m.a...a ure com wttp:///rtteeeee//ooeesss. lisher: Nature Publishing Group.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839) 2020.
W. B. Scoville and B. Milner. Loss of recent memory after bilateral hippocampal lesions. Journal of neurology, neurosurgery, and psychiatry, 20(1):11-21, 1957. ISSN 00223050. doi: 10.1136/
20.1.20.1.11. Number:
Dhruv Shah, Ajay Sridhar, Nitish Dashora, Kyle Stachowicz, Kevin Black, Noriaki Hirose, and Sergey Levine. ViNT: A Foundation Model for Visual Navigation, June 2023. URL http:
wwrxiv..rr//bs//23000 14846 arXiv:2306. 14846 [cs].
Oriane Siméoni, Huy V. Vo, Maximilian Seitzer, Federico Baldassarre, Maxime Oquab, Cijo Jose, Vasil Khalidov, Mare Szafraniec, Seungeun Yi, Michaël Ramamonjisoa, Francisco Massa, Daniel Haziza, Luca Wehrstedt, Jianyuan Wang. Timothée Darcet, Théo Moutakanni, Leonel Sentana, Claire Roberts, Andrea Vedaldi, Jamie Tolan, John Brandt, Camille Couprie, Julien Mairal, Hervé Jégou, Patrick Labatut, and Piotr Bojanowski. Dinov3, 2025. URL
abs 2aas/2508.0 10104
Kimberly L. Stachenfeld, Matthew M. Botvinick, and Samuel J. Gershman. The hippocampus
predictive map. Nature Neuroscience, 20(11): 1653, 2017. ISSN 15461726. doi: 10. no.1038/mn.450. URL http: ///: doi //c.//1.... 1038 10w8/nn. 4650) Number: 11 arXiv:
1011..660 111696
Publisher: Nature Publishing Group ISBN: 9788578110796.
14

---

Published as a conference paper at ICLR 2026
Richard S Sutton. Reinforcement learning; An introduction. A Bradford Book, 2018.
Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction, science, 290(5500):2319 2323, 2000.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need, December 2017. URL arXiv:1706.03762 [es].
http::/arxiv..a/ org www.a//117..0077.
Jianyuan Wang. Minghao Chen, Nikita Karaev, Andrea Vedaldi, Christian Rupprecht, and David Novotny. Vggt: Visual geometry grounded transformer, 2025. bttttsss://axxiv.ooo
abs/2503. 11651
Raymond Wang, Jacdong Hwang, Akhilan Boopathy, and Ila R. Ficte. Rapid Learning without Catastrophic Forgetting in the Morris Water Maze. 2023. URL htttps://ppeneevlle
het//orum??ld-MuuNNee
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances in neural information processing systems, 28, 2015.
James C. R. Whittington, Timothy H. Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy E. Behrens. The Tolman- Eichenbaum Machine: Unifying Space and Relational Memory through Generalization in the Hippocampal Formation. Cell, (83(5):1249 1263.e23, November 2020. ISSN 0092-8674, 1097-4172, doi: 10. 10.1016/j.cell.2200 201. 10.024. URI.
https://www.c cell. com/ cww./eelleccee http///ssttttttteeeoeeeeeeeaae SO092 74 (20) 31388-388 - Elsevier.
Publisher:
James C. R. Whittington, Joseph Warren, and Timothy J. Behrens. Relating transformers to models and neural representations of the hippocampal formation, 2022. URL htttp:////axx
org/abs/211 2212/ 2112. 04035 arXiv:21 12.04035 [cs, q-bio).
Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames. arXiv preprint arXiv:191..00357 2019.
Erik Wijmans, Abhishek Kadian, Ari Morcos, Stefan Lee, Irfan Essa, Devi Parikh, Manolis Savva, and Dhruv Batra. Dd-ppo: Learning near-perfect pointgoal navigators from 2.5 billion frames, 2020. (ntts::/: http://aaavvvvvv.. org
wwtp:/s.g//aaaa....o.o aaa/abb//111 00357
Erik Wijmans, Manolis Savva, Irfan Essa, Stefan Lee, Ari S Morcos, and Dhruv Batra. Emergence of maps in the memories of blind navigation agents. Al Matters, 9(2):8-14, 2023.
Matthew A. Wilson and Bruce L. McNaughton. Dynamics of the Hippocampal Ensemble Code for Space. Science, 261(5124): 1055-1058. August 1993, doi: 10.1126/science.835 $351520. URL hattps://www... science org doi/ lisher: American Association for the Advancement of Science.
wtt.//scincee... 8351520 Pub-
Brian Yamauchi. A frontier-based approach for autonomous exploration. In Proceedings 1997 IEEE International Symposium on Computational Intelligence in Robotics and Automation CIRA "URA"97."Tovii 'Towards New Computational Principles for Robotics and Automation', pp. 146-151.
IEEE, 1997.
Jun Yokose, Reiko Okubo-Suzuki, Masanori Nomoto, Noriaki Ohkawa, Hirofumi Nishizono, Aki- nobu Suzuki, Mina Matsuo, Shuhei Tsujimura, Yukari Takahashi, Masashi Nagase, Ayako M. Watabe, Masakiyo Sasahara, Fusao Kato, and Kaoru Inokuchi. Overlapping memory trace dispensable for linking, but not recalling, individual memories. Science, 355(6323) January 2017. doi: 260.11c2ccee
URL httpp:://ww..t.... http://wwwcccenc..e.mm org/doi// Publisher: American Association for the Advancement of
10.1122/ stt//sciencee wat.cce.aall50 aa12690 Science.
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang. Storm: Efficient stochastic transformer based world models for reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A RELATION TO OTHER MODELS OF MEDIAL TEMPORAL LOBE
Successor Representations (SR). In the context of reinforcement learning, successor representa- tion has been shown to explain many non-spatial aspects of hippocampus responses (Stachenfeld al. 2017). SR and ESWM are similar in that both rely on transition information, but they differ in several important ways. First, classical SR methods require either a full transition matrix or incre- mental experience sufficient to estimate one. They are not designed to infer missing transitions or to complete a partially observed graph. ESWM, by contrast, is built specifically for that setting: it takes a sparse set of one-step episodic transitions and infers the missing structural relations directly at inference time. Conceptually, ESWM is estimating the entire family of transition distributions from incomplete data, whereas SR presupposes access to the full underlying dynamics.
Second, SRs are learned through iterative, policy-depender counted future occupancy. ESWM does not rely on iterative updates and is not tied to a particular policy. The transitions in its memory bank are single, policy-free episodic observations, and the model performs transitive reasoning over these episodes on demand to answer arbitrary transition queries.
updates and encode multi-step, dis-
In addition, SRs collapse the environment into a future-occupancy representation, which mixes structural and predictive information. ESWM keeps the structure explicit and manipulable, allowing to reconstruct the transition graph itself rather than a discounted predictive summary (e.g. state values). ESWM is also symmetric in its inference capabilities: it jointly estimates forward transi- tions, backward transitions, and actions from the same memory bank, whereas SRs typically encode only forward predictive relationships.
Finally, ESWM supports rapid structural updates by directly editing or adding episodic transitions, without needing to relearn transition matrix. This contrasts with SRs and related models, which must re-estimate transition dynamics when the environment changes.
-ichen-ai Machine. Tolman -ichennba Machine (TEM) Whittingtnn et 2020
2022 Bakermans et al., 2025) has been proposed as model of medial temporal lobe function and meta-trained across diverse environments to predict unseen transitions from episodic experience. TEM factorizes sensory and structural information: the shared structural scaffold is encoded in the weights of a recurrent network, while environment-speciffe observations are bound to that scaffold through Hebbian-slot memory or self-attention module. The recurrent state therefore encodes generic structural prior-ttpicca typically a Euclidean, grid-like topology- and the attractor dynamics al- low for content-specifie retrieval. This makes TEM well-suited for environments that share this underlying family of structures and allows it to tolerate large changes in sensory input at familiar locations. However, this same design makes the model less flexible when the underlying topology itself changes or when the environment obeys different structural rules (e.g., Random Wall environ- ments where walls can appear in arbitrary shapes and configurations).
In contrast, ESWM does not factorize structural and sensory information. Both the structural layout and the expected sensory observations are inferred jointly and directly from episodic transitions. Because ESWM does not rely on fixed structural template encoded in its weights, it can accom- modate changes to the environment by directly modifying its memory bank-allowing predictions immediately without retraining. This makes ESWM particularly effective in settings where the topology can vary widely across environments or can change abruptly within an environ- ment.
update
Finally, whereas TEM operates on sequential trajectories and depends on an explicit dynamical update of its recurrent state, ESWM integrates disjoint and non-sequential experiences. This enables more sample-efficient learning from sparse observations and supports rapid adaptation to newly introduced obstacles or structural changes that were never part of trajectory.
Recurrent Predictive Models. Recurrent predictive models have also been proposed as models of neural computations within the hippocampus (Levenstein al. 2024). These models consist of re- current neural networks trained to predict future observations from a sequence of prior observations along particular trajectory. In these models, both the structural and sensory information em-
16

---

Published as a conference paper at ICLR 2026
bedded directly in the recurrent network's 's weights. This architecture makes them even less flexible in the face of changes, more so than the TEM class of models. Whereas TEM includes an episodic module that allows for rapid adaptation to sensory changes within fixed structural template, recur- rent predictive models must update their recurrent weights to incorporate both sensory and structural changes, which is considerably harder achieve quickly.
In addition, recurrent predictive models depend on having continuous, sequential experience to grad- ually shape the internal dynamics into coherent predictive structure. This reliance on dense ex- perience makes them poorly suited to integrating disjoint or fragmentary observations. ESWM, by contrast, designed specifically to operate on sparse. non-sequential episodic transitions, allowing it to infer the structural relationships among states without needing long trajectories.
B EVIDENCE FOR BUILDING SPATIAL MAPS FROM DISJOINT EXPERIENCES IN ANIMALS
Several behavioral studies have demonstrated that animals and humans integrate information collected across disjoint episodes into coherent spatial or relational maps. (Roberts et al.] (007) showed that rats inferred a shortcut between two locations in a maze after exploring its subcompo- nents in separate sessions on different days. From the perspective of our framework, the animals observed disconnected edges of the underlying spatial graph- never experiencing them as a contin- uous trajectory- -yet were still able to infer the globally shortest route during the test. (Fernandez) al., 2023) similarly demonstrated that humans trained to navigate locally within three separate en- vironments on different days could integrate these fragmented experiences to navigate across them globally. (Park et al.] 2020) found that human subjects who learned each dimension of a rela- tional space in separate episodes nonetheless formed joint map of the full space, again illustrating cross-episode integration.
There is also substantial neural evidence that hippocampal and amygdala circuits integrate memo- ries according to shared structure rather than temporal contiguity. (Yokose et al. 2017) showed that manggdaaa-dependeen memories from distinct tasks form overlapping engrams when they share crit- ical stimulus. (Schlichting et al.) 2015) found that the anterior hippocampus integrates across paired associates (A-B B C)) to support the inferred A-cC relationship. (McKenzie et al., 2014) ported that hippocampal population codes become similar for objects that share spatial context, even when those objects were experienced in separate episodes. Related reviews (Morton et al. 2017 Brunec et al., 2020) synthesize a broad literature showing that the hippocampus organizes memories into integrated representations that reflect spatial, temporal, and conceptual structure often beyond what was directly and continuously experienced.
Altogether, these findings provide strong evidence that biological systems do face and solve form of the "disjoint memory" problem. They routinely combine episodic fragments acquired at different times, in different contexts, and along different task dimensions, to construct coherent internal maps that support flexible behavior.
C EXTENDED METHODS
C.1 TRAINING DETAILS
C.1.1 OPEN ARENA
In the Open Arena environment, each integer state value is encoded as a 6-bit binary. Each bit is independently projected-via either the source-state or end-state linear projector-into 128-dimensional embedding. The six embeddings are concatenated into a single 168-dimensional state vector. Actions are similarly mapped into 768 768-dimensional embeddings using dedi- cated linear projector. The final transition embedding is then obtained by averaging the three 768-dimensional vectors corresponding (1) the start-state, (2) the action, and (3) the end-state.
17

---

Published as a conference paper at ICLR 2026
C.1.2 RANDOM WALL
In the Random Wall environment, the same procedure applies, except that each state is represented by an integer and projected into 1,024-dimensional embedding, yielding 1.024-dimensional transition vector.
C.1.3 MINIGRID
In the Minigrid environment, the states are 5 x 5 egocentric view, observing 25 locations all at once. Each location in the environment assigned one of the colored objects (e.g, red floor) among a pre-defined set of 9 possible colored objects. Each of the objects among the 25 observed objects in a state is encoded into a 64 dimensional vector separately. The encodings of the 25 objects are then concatenated Open Arena and Random Wall, 1, the 8s, a, and 8, embeddings are averaged to produce the final emhedding for the transition. The prediction for the state is 25 independent 9-way classifications.
produce single 1600 dimensional representation for the state. Similar to
C.1.4 PROCTHOR EXPERIMENTS
In ProcThor (Deitke et al. 2022), an agent can be placed anywhere in the environment to collect 224 224 3 RGB observations (i.e. states; FOV=90). A transition consists of start state 8 a continuous action vector AAr, Ay, A/] specifying how much to move right, forward, and rotate clockwise (negative values mean moving along the opposite axis) to arrive the end state, and an end state Se- To sample transition, we first spawn the agent at a random, unobstructed location with a random rotation (start state). Then we randomly teleport the agent to new location that is within the 3m x 3m box of the start state, followed by a random rotation (end state). Lastly, compute the egocentric action vector from the start state to the end state's 's poses.
The states are embedded by DINOV3 (Siéénni et al. 2025 into W x H x D = 14 x 14 x 1024 feature map, padded and pooled into W x x D = X X 1024, then flattened into 16, 1024- dimensional tokens. The action is projected by a I-layer MLP into standalone 1024-dimensional token. The start state, end state, and action tokens are concatenated into a sequence 33 tokens to represent a single transition. Suppose memory bank has N memories, then ESWM receives the concatenation the memory bank and query, totaling (N + 1) x 33 tokens, as input.
The action and end state of the query transition masked with equal chances (the start state is left out for our preliminary experiment). For state prediction, the model is tasked to maximize the cosine similarity between its output and DINOV3's represennaaion. For action prediction, the model performs categorical prediction over a discretized action space. To discretize action prediction, the 3m 3m movement box discretized into grids, with each grid representing a class. ESWM needs to essentially predict how many steps to move right/left, forward/backward until the end state's location reached, with the step size (or discretization granularity) determined so that there are 64 classes. The rotation is discretized by 10-degree intervals, resulting in 36 classes.
The concatenation of the memory bank and query fed into a multi-layer transformer. After all transformer blocks, we linearly read out the action prediction from the action token, and directly match the state tokens with their target representations without any further transformation. We apply cross-entropy loss to action predictions and cosine distance loss states.
To train ESWM in GPU-based environments like ProcThor, de- vise a memory-effiiient fast, and scalable training pipeline that largely follows DD-PPO (Wijmans) et al. 2020). The training is distributed over N GPUs, with each GPU alternating between online data generation across K scenes and training. In our preliminary experiment, we train ESWM with 4 A100 GPUs for 19 days over 80k iterations with an effective batch size of 144, using only 250 GB of memory.
ProcThor Training Pipeline
C.1.5 MODEL. ARCHITECTURES.
Transformer: attention heads per layer; feed-forwart hidden dimension of 2,048, unless otherwise specified.
LSSM: hidden state size of 1,024.
18

---

Published as a conference paper at ICLR 2026
Mamba: model dimension of 768 in Open Arena and 1,024 in Random Wall.
For ProcThor experiment, replaced the standard transformer block with the Alterating Atten- tion block from VGGT (Wang et al. 2025). Each transformer block consists of a memory-wiss attention layer, where the attention scope is limited the tokens within transition consolidating the relation between the start and end state, and a global attention layer, where all tokens can attend each other enabling memory integration. We use 16 heads in our ProcThor runs.
C.1.6 TRAINING HYPERPARAMETERS
All models were trained on a single NVIDIA A100 GPU for 480,000 iterations using:
Optimizer: A with initial learning rate x 10 and cosine decay schedule.
Dropout: 0.1.
Batch size: 128 memory bank and query pairs.
Training time: 19-28 hours, depending model capacity.
Memory usage: 2GB. The low memory consumption is attributed to the online generation data.
C.2 BUILDING RANDOM 2D GRID ENVIRONMENTS
To generate input (memory bank, query) pair to ESWM, first generate environment & Each is represented as a 5-tuple (G, g, W, th A), where:
G = (L, E): An undirected graph, with L representing unique locations and E represent- ing undirected edges connecting adjacent locations. The graph adopts fixed hexagonal structure across all environments. It models grid-like movement dynamics and serves as the canvas for creating complex environments.
EL,
g: A filtering function that generates subgraph Gobs S G, containing only observable locations Lobs C L and transitions Eons C E (i.e. part of the environment that will be observable to the agent in that sample).
W: A subset of contiguous locations in L, representing wall that restrict agent's motion in the environment.
db: mapping from integer observations, assigning each location to a unique obser- vation from S C z. This allows random association of locations and observed values (or states) each sample.
A: Maps undirected edge = ly) E, which is a pair of locations, an action- induced transition either starting from 4 or ly with equal chances. c.g. A(4,, (4) returns 3-tuple transition (4, go West, 4j) or (ly, go East, 4) with equal probability.
In the Open Arena setting (Section 3.3 the entire environment is observable and there are no obstacles (i.e, g is the identity function, W 0). In the Random Wall setting (Section (3.3), removes random subset V Vunobs S V along with their attached edges, and W o 6. The components W, and s are randomly generated for each environment e, where applicable.
C.3 GENERATING MEMORY BANKS IN 2D GRID ENVIRONMENTS
The minimal spanning episodic memories for an environment e = (G. W. is generated follows:
1. Apply the filtering function g to G produce the partially observed subgraph Gats (Laba Easx).
2. Run a minimal spanning tree (MST) algorithm on Goss- with random weights assigned to edges ensure diversity. Tree effectively models set of fragmented transitions that do not naturally integrate into continuous trajectory. This produces set of undirected edges
T (44,1 4 E T.
19

---

Published as a conference paper at ICLR 2026
3. Use A to map edges in T to directed 3-tuple transitions between locations and ti map locations to observations. This produces an array of transitions between observations [[s,, s,)].
4. Modify transitions ending at locations W to form self-loops, effectively modeling blocked movement. Then, randomly permute the array of transitions to produce the final array of transitions M.
In Open Arena and Random Wall, the memory banks M observe at most 18 and 36 transitions out of the total 84 and 180 transitions, respectively.
C.4 QUERY SELECTION
The resulting memory bank M is an unordered set of one-step transitional episodic memories that span all observable locations in the environment (Fig. (). Given M, the model is tasked to infer the missing component of query transition q = (s,, a, Se)-
For Random Wa 1 1, the query q is selected from 1) the set of transitions that are in Gobs but not in the memory bank (unseen transitions; 68%); 2) from the memory bank (seen transitions; 15%) and; 3) from set of unsolvable transitions where one end is in unobs (17%). In Open Arena, the query is always selected from unseen transitions.
The model receives the M and q = (ss, a $e), with either 8s, a, or 8e masked randomly with equal probability, as input, and asked to infer the masked component.
C.5 TEM-T IMPLEMENTATION
Our goal is to maximally equalize the training of TEM-T and ESWM to eliminate potential con- founds such as task differences. Thus, TEM-T also receive the concatenation of memory bank and unseen query transition as input, except that the disjoint memory bank is replaced with a sequential trajectory that spans the entire environment. The task is predict the missing end-state of the final query transition, analogous to inferring the masked end state in ESWM. The trajectory is kept mini- mal, generated by Networkk travel inn salesman problemt em, match the sparsity of the memories we provide for ESWM. In Random Wall, 1, we explicitly ensure that the agent runs into each obstacle at least once throughout the traversal so that the model has sufficient knowledge to infer the underlying structure of the environment.
To our best knowledge, no publicly available code exists for TEM-T. Therefore, we resort custom implementation that closely follows the authors' recipe. We use the same dimensions to encode the positions é and observations x, (384 in Open Arena and 512 in Random Wall) and ensure the combined dimension of h, [xt, ex] matches ESWM's embedding dimension. Sigmoid used as the non-linear activation a in positional encoding generation, following the observation that linear and ReLU activation produce exploding gradients. We use 8 heads and 2048 feed-forwart dimensions, matching the configuration of ESWM. The prediction is read out linearly from the embedded query transition after singlc/multiple transformer blocks. We train TEM-T over 460k iterations with batch size of 128 and a cosine learning rate schedule, same as ESWM.
C.6 ISOMAP
We consider ESWM's population neuron activity at a particular location as the ESWM's activation vector as it infers the missing component of a query transition that either starts or ends at that location. For 8, and a predictions (i.e, 8, is visible to ESWM), the activation vector is associated with the s. 's location, while for sa prediction (i.e. 8, is not visible), associate the activation with the se location. For transformer models, we used the query token after each transformer encoder block the activation vector, while for multilayer LSTM and Mamba, we used the final hidden state (i.e, the last token from the processed sequence) from the last layer.
To collect ESWM's activation across all locations in environment, task ESWM do either 8, o, or s, prediction for all transitions the environment conditioned on a memory bank. This process is repeated with multiple memory banks from the same environment until approximately
20

---

Published as a conference paper at ICLR 2026
1,000 activation vectors collected. We then apply ISOMAP using cosine distance and 20-60 neighbors to generate the visualizations.
C.7 EUCLIDEAN DISTANCE ESTIMATION FROM ACTIVATION
A linear layer receives as input the concatenation of three vectors: the activations from Transformer- 2L ESWM corresponding to action predictions starting from states A, B, and C. It outputs scalar representing the probability that the Euclidean distance between and exceeds that between A and C. The classifier is trained on 3,000 state triplets sampled all from different environments and evaluated on 1,000 held-out triplets. Reported accuracies are averaged over 10 runs with different randomly initialized weights for the linear layer.
C.8 SRSA
We use the formula from (Levenstein et al. 2024) to compute sRSA This metric measures the Spearman correlation between latent and physical distances. To get the latent representation particular location, we task the model to solve all queries that start in that location and average the activation. We repeat this procedure across all locations in an environment and across 5 memory banks from each environment. We used ESWM-T-2L's first layer's 's query token's 's activation vector during the action prediction task for this analysis.
C.9 NOISY MEMORY BANKS IMPLEMENTION
To create memory banks with noisy, contradictory, state observations, we pick location and change the corresponding memories in the memory bank to observe different states at that location.
To create memory banks with noisy, contradictory, transitions, randomly pick transition from the memory bank, edit its end state, and append the modified transition to the memory bank.
This analysis is done with ESWM-T-4L. in Random Wa 1.
C.10 BASELINE NAVIGATIONAL AGENT
The baseline agent is trained via A2C (Sutonn 2018 Mnih 2016) to navigate between any two states in single Random Wall environment, which differs from ESWM and EPN, which are meta-trained across many environments. At each time step f, the agent receives as input a sequence of historical interactions with the environment, each as a tuple (5t-1, 0r-1, St, 5goal), and is tasked to output action towards the goal state 8goal- Each part of the tuple is distinctly embedded onto a 128D vector and the vectors are averaged to produce the final observation each time step. The embedded sequence is fed into ony-l LSTM (with a hidden dimension of 256) and the predicted action and value are read out from the final hidden state via distinct linear heads.
The agent receives a penalty of -0.01 every time step except when it reaches the goal state, where receives a reward of 1. The training is done over 5000 episodes with a discount factor of 0.99 and learning rate of 10-3
21

---

Published as a conference paper at ICLR 2026
D ADDITIONAL ANALYSIS
D.1 COMPARING MODELS IN OPEN ARENA
Different models exhibit drastically different performance in Open Arena, This section aims to provide insights into the factors that contribute to each model's performance.
The ESWM-T models perform noticeably well in the Open Arena. We investigate which hy- perparameter affect its performance the most. In 2 we fix the number of heads and observe a positive correlation between model depth and performance. vary the number of heads. We observe that while increasing depth leads monotonic increase in performance, increasing the number of heads initially improves performance, but exceeding heads leads to quick drop in performance. In addition, the model depth impacts all three prediction tasks, while the number of heads selectively impacts state prediction tasks more than the action prediction task.
Table. ييل we fix the depth and
As shown in Fig 8 ESWM-MAMBA and ESWM-LSTM both overfit on the training states, explain- ing their poor performance in completely unseen test states. However, we note an additional obser- vation: LSTM struggles to perform even on the training set in Open Arena, unlike in Random Wall where it succeeds on both train and evaluation splits. We suspect another task difference between Open Arena and Random Wall contributes to this discrepancy: states in Open Arena consist of 6-bit value, each represented by 12-diimensional vector while states Random Wall consist of single continuous vector. The latter may allow more compact latent representation that better fits within LSTM's limited memory capacity.
TEM-T models, which also incorporate transformer into its architecture, underperform compared with ESWM-T models. TEM-T is architecturally distinct (See Section 2 and is trained on a different distribution (minimal spanning trajectories rather than minimal spanning trees) compared to ESWM T. To determine which of these two factors contributes to the performance gap, we train ESWM-IT T- 10L on TEM-T's training distribution and find that it can reach similar performance compared with training on ESWM's, predicting unseen 8s, a, 8. with accuracies 99%, 99%, 99%. The similar is true in Random Wa Wall, where ESWM-T-6L. reaches 8g, 0, 8, accuracies of 81%, 89%, 78% when trained with TEM-T's training distribution. These results indicate that TEM-T's architecture most likely explains its performance gap with ESWM-T.
D.2 COMPARING ESWM TRAINED WITH OR WITHOUT THE MINIMALITY CONSTRAINT ON MEMORY BANKS
We show that ESWM-T performs well when trained with minimal memory banks (Fig. 2 and non- minimal memory banks (Fig. 9.. However, ESWM-T trained on non-mine non-minimal memory banks, where more environmental transitions are observed, underpeforms the one trained on minimal memory bank, where less environmental transitions are observed.
To investigate why the ESWM-T trained on non-minimal memory banks underperforms while hav- ing access to more experience in an environment, we break down the model's performance across memory bank sizes and observe substantial gap between denser and sparser memory banks, even though they are equally likely to be in the training data (Fig. D. This suggests that, when trained on a mixture of densities, the model focuses more on solving the easier tasks where is given an over- complete memory bank, and only partially succeeds in learning the more challenging tasks when minimal information is given. This leads to weaker spatial reasoning skills and poorer performance when the observation pattern changes.
In contrast, the ESWM-T that trained under minimality constraints is forced to solve the harder spatial reasoning problem consistently. As shown in the (Fig.Ap) this produces representations that transfer robustly to denser, caatoo--Ditributiion (OOD) memory banks. We also note that within the training distribution (ID), increasing memory bank size corresponds area, which requires the model to coherently build a larger map which is more challenging and likely explains the performance drop.
expansion of the observed
Collectively, these findings suggest that applying sparsity constraints during training leads to better test-time flexibility in terms of information availability.
22

---

Published as a conference paper at ICLR 2026
D.3 HANDLING DISCONNECTED MEMORY CLUSTERS
During Random Wall training, memory banks can contain disconnected memory clusters, with no paths between observations in distinct clusters. This leads to uncertainty spatial relationships as path integration provides no avail. In rare cases, the model is asked to solve queries that bridge disconnected clusters. This raises two questions: (1) Does the model infer the relative positions of disconnected memory clusters (e.g., which part of the room each cluster maps onto)? (2) Which cues enable these inferences?
D.3.1 INFERENCE FROM ENVIRONMENTAL STRUCTURE
We observe that when memory clusters are disconnected, the model leverages the underlying struc- ture of the room infer their relative positions. For example, the model aligns memory clusters with the room's 's boundary when their shape matches. This allows the model to piece together dis- connected memory clusters, like solving puzzle, by identifying the configuration that satisfies the invariant structural constraints. Evidence for this behaviour is twofold. First, in the model's la- tent space, the disconnected clusters are stitched together, rather than overlapping or disconnecting, to form a smooth surface reflecting the room's 's structure (Fig. 3t) Second, the model's end state prediction distributions, starting from the boundary states of one cluster, are biased towards the boundary states of the inferred neighbouring cluster, instead of exclusively outputting the "I don't know" option.
D.3.2 GENERALIZATION FROM OBSERVATIONS OF OBSTACLES
In cases where disconnected memory clusters' shapes do not provide clear cues about their relative locations, the model relies observations of obstacles deduce their positions. For instance, given two disconnected clusters of memories, the model initially maps them to different locations in its latent space and outputs predictions exclusively reflecting uncertainty ("I don't know"). However, if both clusters observe obstacles, even though the clusters remain disconnected, the model applies its knowledge of common wall shapes (e.g, walls are usually straight) make inferences about the relative locations of the clusters. In the latent space, the two previously disconnected clusters are connected and aligned along the inferred wall shape to form a continuous representation (Fig. [3d, 13 The model's predictions are also biased towards the inferred neighbours. Moreover, a single observation of obstacles from both clusters is sufficient for the model to infer the shape of the wall, and therefore infer the relative positions of the clusters (Fig. 13) These results demonstrate that ESWM can leverage the invariant structure of the room and combine observations of obstacles with prior knowledge of wall shapes to construct coherent spatial representation.
D.4 PLANNING BY GUIDED IMAGINATION
While the strategy introduced in the main paper enables planning, it uniformly searches through the space, unbiasedly imagining the consequences of all actions, even if some move further from the goal. In contrast, animals with internal maps encoding geometric constraints plan biasedly, prioritizing actions that shorten the geodesic distance (i.e. the shortest path length) to the goal. We show that the spatial maps in ESWM's latent space, which accurately encode obstacles and geodesic distance between states (Fig. 166), can be used for guided imagination towards the goal.
The previous strategy is akin Dijkstra's algorithm, or A with a null heuristic function (Hart et al. 1968 Ferguson et al., 2005). The key to efficient planning, which prunes undesirable actions early, then, a heuristic function containing rich information on the geodesic distances between states in the environment. Inspired by ISOMAP TTeenbaum et al. (2000). we construct such heuristic function Masen as follows (See Alg for pseudocode): 1) Query the model solve the end state prediction task for all possible (s, a) pairs, collecting activations. 2) Construct graph Giatent where nodes are activations and edges connect nodes within radius Rigent (Appendix GG.. Edges are weighted by the cosine distance between their two endpoints (Fig. 16p). Intuitively, each node Giment is the activation of predicting the outcome of action a starting at state s. 3) Compute pairwise shortest paths this graph to create a geodesic distance table Matent-
Maten captures ESWM's latent spatial map into a look-up table containing information on which action takes you closer the goal state at every step of imagination. A guided imagination with
23

---

Published as a conference paper at ICLR 2026
Matent. improves planning efficiency by 57% (Fig. 16a), measured by the number of states visited, over Dijkstra's 's while maintaining path optimality (Fig. IGp). Furthermore, greedy policy based on Maten achieves 70% success in navigating around obstacles (Fig. [1 These results suggest that structured latent spaces world models can serve as efficient heuristics for planning.
24

---

Published as a conference paper at ICLR 2026
E ADDITIONAL TABLES
Heads
4 Heads
8 Heads
16 Heads
Predict Source
40.4
78.9
83.3
42.0
Predict Action
84.6
91.8
92.6
91.8
Predict End
37.6
78.1
79.6
40.4
ESWM-T number of heads ablation We report the accuracies for the start state, action, and end state prediction task in Open Arena with ESWM-T-6L.
Table 1:
Predict Source
Predict Action
Predict End
Seen
0.9050
0.8700
1.0000
Unseen
0.8750
0.8880
0.9030
Unsolvable
0.8270
0.9050
0.8180
Table 2: ESWM-T-4LL performance Random Wall Accuracies for source, action, and end prediction.
Predict Source
Predict Action
Predict End
Seen
0.8780
0.869()
1,0000
Unseen
0.6660
0.7530
0.7030
Unsolvable
0.7620
0,8830
0.7560
Table 3: ESWM-LSTM-4L performance in Random Wall Accuracies for source, action, and end prediction.
F ADDITIONAL FIGURES
25

---

Published as a conference paper at ICLR 2026
Predict Source
Predict Action
Predict End
Scen
0,9000
0.8810
1,0000
Unseen
0.7710
0.8110
0.8230
Unsolvable
0.8110
0.8990
0.8010
Table 4: ESWM-Mamba-SL. performance in Random Wall Accuracies for source, action, and end prediction.
Prodict Action
Predictt End State
Predict Actioon
Predict End State
Myr pu
0.0
M/
|M11
2.6
307
367
Datapoint
Datapoints Seen
(b)
Figure 7: ESWM's evaluation accuracy in grid environments with duplicated states. A trajec- tory is prepended to the query transition to serve as context for disambiguating two locations with the same state. The context trajectory begins at random state and ends at 8, of the query transition. The query transition's ' a and & are masked with equal chances. a) Accuracy in Open Arena for ESWM-T-8L. and b) Random Wall for ESWM-T-4L.
ESWM Model
Curve
MAMBA-12L
Train
Eval
Data Points Seen
Figure 8: Training and evaluation loss for various ESWM models in Open Arena.
15/M61-12
Models CSMM MAM-A--
Kandom
100.0
ESIM LSTN
TEMT22L
Non-Minimal (ID)
Predict
w
Accuracy
Predict Start
Accuracy
80.0
Predict Action
87.5
Predict End
85.0
82.3
80.0
Training
Memory Bank Size
(a)
(b)
Figure 9: Training with non-minimal memory banks. a) Evaluation accuracy on unseen queries for various models trained with non-minimal memory banks. ESWM receives memories that are disjoint, spanning, but not minimal (constructed by adding extra episodic memories to the minimal memory banks) and is tasked to predict an unseen query. TEM receives a minimal spanning trajectory padded with random walks as input. b) Prediction accuracy over memory bank sizes for ESWM-T-12L. from a). The shaded area emphasizes that this model is trained on non-minimal memory banks. means In Distribution.
26

---

Published as a conference paper at ICLR 2026
Transforme-21
LSTM-4L
EPN
Random network
Figure 10: ISOMAP projection of various ESWM-trained models' latent space. The last two environments out-of-distribution ESWM has never seen multiple walls in the a cluster of obstacles that does not form a long, straight wall during training. For Transformer- 2L, LSTM-4L, and Random network (LSTM-4L initialized with random weights), the activations for end-state prediction are used while for Mamba, the activations for action prediction are used. EPN (Section H) is episodic-memory-base navigational and exploration agent trained the same environment as ESWM.
27

---

Published as a conference paper at ICLR 2026
Transformer21
LSTM-4L
Mamba
EPN
Random network
Figure 11: Same ISOMAP projections as Fig. 10 but in a different angle.
28

---

Published as a conference paper at ICLR 2026
Predict Start State
Predict Action
Predict End State
Iransformer-21
LSTM-4L
(a)
(b)
Figure 12: ESWM's latent spatial map persists for all prediction tasks and partial observations. a) ESWM's 's latent spatial map for s.. a, and s, predictions. Mamba is excluded as the spatial map only found for the action prediction task. For Transformer-2L. the spatial map is found in the first layer for prediction, and the second layer for and s. prediction. b) Spatial map when the memory bank observes an irregular subsection of the room. Trass first layer's activation for action prediction is used.
29

---

Published as a conference paper at ICLR 2026
(a)
(b)
Figure 13: ESWM's latent spatial map adjusts to structural changes. ISOMAP projections of the 14-layer transformer model's seventh layer's activations. a) ESWM's latent space when the input memory bank M consists two disjoint memory clusters. b) ESWM's latent space when both memory clusters observe wall their left; despite remaining separate in the memory bank, ESWM inferred them to be connected as reflected in its latent space. c) ESWM's latent space when both memory clusters observe single obstacle their left. ESWM, again, inferred the clusters connected but with sparser observations of obstacles than in b). 900 activations and 45 neighbours are used to the ISOMAP.
0000 0000
Figure 14: Representation becomes more continuous with depth. Left to right are 30-neighbours ISOMAP projections from layer 9, and 14 of Transformer 14L (2000 activations). The model was trained without random masking of subregions. Representation in the latent space becomes progressively more continuous.
30

---

Published as a conference paper at ICLR 2026
Number of Unique Staters
Number of Unique States
p= 2.57e-02
Random
FSWM Exploration Strategies
Ideal
ESWM
EPN
Exploration Strategies
(a)
ESWM
EPN
Success Rate
0.9
Number of Obstacles Added
(d)
Figure 15: Exploration, navigation, and adaptability for ESWM and EPN. a) Number of unique states visited by different exploration strategies a Random Wall environment with 37 locations over 25 time steps (n=1000). On average, ESWM explores a comparable number of states to an oracle explorer with access to obstacle locations. b, c, d) Same as Fig. 5 exxcep the performance of ESWM and EPN is averaged over random seeds. Error bars show standard deviation across seeds. Welch's two-sample t-test was used to compare ESWM and EPN (a 0.05). A dashed line with sge indicates statistically significant difference (p <0.05)
31

---

Published as a conference paper at ICLR 2026
Weth Lergh
Nhhrrh
(a)
Hose Ayal
Contur: pilANA
Gecdek Eitarc on Cryar
(b)
H+- Djaoza +Bo A" Withhwww
Sauccaii hate
Optirality
Path
Figure 16: The latent space of ESWM can be harnessed as a heuristic for more efficient plan- ning. All the results are on Random Wall environment with locations. a) (Left) Average number of nodes expanded for different heuristics over 1000 paths. (Middle) Success rates of greedy nav- igation using different heuristics over 2400 paths. (Right) Success rates for greedy navigation on paths requiring detours around obstacles 800 paths. b) (Left) Correlation between geodesic distances on Gueett (see D.4) and corresponding distances in the environment, using n=1500 pairs of states (Sstarts Sgaal) over 75 environments. The linear fit shows a strong positive correlation with R² 0.89. (Mid-Left) sample grid environment. (Mid-Right) Giatent for the same environment. (Right) Agent navigates greedily on Maten on a path that requires detours around obstacles. c) Com- parison of navigation success rate and path optimality between Dijkstra's 's and A* different path lengths (n=1500). Although not guaranteed to be admissible, /asen empirically preserves path opti- mality.
32

---

Published as a conference paper at ICLR 2026
1.0
0.5
0.5
0.4
0.2
0.0 Seen
Query
Figure 17: Comparison of end-state prediction accuracy between EMWM Coda-Forno et al. (2022a) and Transformer-2L ESWM in RandomWall. EMWM predicts using weighted aver- age over the top-K matches (K=5), but we found K=1 performs best, likely due to the sparsity of memory banks (i.e, transitions unique, averaging more than one transition offers no ben- efit). To adapt the author's s implementation to our task, we embed each token (transitional memory tokens or the end-state masked query) by averaging the embedding of each item (st, 0t, St+1), each produced using a distinct MLP layer identical to ESWM. All tokens are sequentially fed into I-layer Cho al. (2014) obtain the sequence of hidden states, which become the keys with the corresponding end-state embeddings as values for the memory table. The network is optimized predict the s1+1 for each token.
33

---

Published as a conference paper at ICLR 2026
G ALGORITHMS
Algorithm Dijkstra- Pathfinding with ESWM
procedure FINDPATH(ESWM, memoryBank, Syaats Sgaals T maxx)
2:
const ACTTON COST
3: Initialize priority queue OpenSet.
4: Initialize dictionary Cost
5:
Initialize dictionary Parent
6: Cost [sttr]] = 0, Parent [saun] nil
7: Push (Suart: cost = 0, depth = 0) into OpenSet
8:
while OpenSet 2 do
9:
(u, d) + PoMMinCossNodee(yyennee)))
if goal then
11:
return ReconstnucPath(PPrentt u) end
13:
if d 2 Tmax then continue
remove path that exceeds horizon
15:
end if
for all action a in PossibleActions do PredictNextSaate(ESS Weeeee reenorynnnn u, a)
17:
18:
c+ + ACTION.COST d' d+1 e not Cost or d' < [ust(u) then
19:
20:
21:
Push (u, c', d') into OpenSet
22:
Cost(o] =
23:
Parent [u]]
24:
end if end for
26:
end while
return FAILO)
no path within horizon Tmax
28:
end procedure
Algorithm Explore Function
1: procedure EXPLORE(ESWM. memoryBank, actions)
s current observation
2: Initialize q max heap
3: for all E actions do
4: s', logits <- W.W. PREDICCTEEDDOwmmrryyBbkkkk 8, a)
5:
probs + SOFTMAX(logits)
6:
if s' 'I don't 't know' then
7:
Enqueue into with cost = 1 (ENTROPY(probsj) 3 =02 in our experiments
8:
else probs[s'] < 0.8 then
9:
Enqueue into with cost = {ENTROPY(probs)) end
by= 0.1 in our experiments
10:
11:
end for
12: return q. R..EMOVEMAX))
13:
end procedure
34

---

Published as a conference paper at ICLR 2026
Algorithm 3 Get all State-Action Pairs
1: procedure GETSAS(memoryyBan
2: S' + GETUUNIQQESSAATESS//mmmmmyyy
3: Initialize as an empty list
4: for all E S do
5: for all E Actions do L.pppend (8, a)) 6: end for end for
7:
8:
9: return /
10: end procedure
Algorithm 4 Compute Geodesic Distance Table /ttenntt
I: procedure GETHEURISTTCSSEESSWW memoryBank, Rmment)
2: SAs + GETSAs(memoryBannk
3: N + SSAA
4: Initialize modesToSAs an empty dictionary
5: Initialize activations as an empty list
6: Initialize hasteur as a matrix of shape N N
7:
for all (s, a) E sa saPairs do activation ESWM's activation for end state prediction task starting at taking action
8:
a nodes + (s, a) uecivationsaappenloullaaalio end for
9:
10:
11:
12:
Glanen <- NEAAESSNNEACHHUURGGLLAAAACCIIIIT Raaten
for all nı, E nloeen do for all n2 E Gaatem modes do
14:
D T4y My are activations
Sa,, sa, + nodesToSAs nodesToSAs(n.] soae [sa sa,1 + SHORTESTTAATLLENNHHHHHH 7, n,))
16:
17:
end for end for
18:
19:
return Maateen
20:
end procedure
Note that Matent approximates the geodesic distance between any two state-action pairs rather than between any two states. Geodesic distances between states can be obtained by averaging over the action space.
Selection of Riatent
In essence, Riaeen the radius among a set radii in which the corresponding hjssnt obtained from Alg 4 best captures the geodesic distance information in the enviromment. More specifically, Reaen is the E that maximizes both the success rate and the optimality of the agent's path when it navigates greedily to the corresponding Amsen- We consider navigation as successful if the the agent reaches the goal state within time steps. Path optimality is measured by the ratio between the shortest path and the agent's path. To determine Ruatent, we conducted search over bounded range R, which includes those whose ISOMAP projections accurately reflect the structure of the environment.
Limitations Alg4
One drawback is that /tttent can overestimate the true geodesic distance, i.e. it is non-admissible. This is because Giatent might be under-connected for the given Rmment and memory bank (i.e. nodes aren't connected by edge when they should). This can be mitigated by using a larger radius R + to build G (i.e prefer over-connection over under-connection). However, we ac- knowledge that there is way guarantee admissibility, hence, the optimality of the path. This
35

---

Published as a conference paper at ICLR 2026
trade-off for faster search. In future work, we will explore adaptive radius where radius, hence Matent, can be adjusted dynamically as agents interact with the environment.
H EPISODIC PLANNING NETWORKS (EPN)
The Episodic Planning Networks (EPN) model Ritter et al. (2020) specifically designed and trained to rapidly learn to navigate in novel environments. Comparatively, ESWM is a general- purpose model not explicitly trained for navigation. By comparing with EPN, we investigate whether ESWM's ability construct internal models of spatial structures enables navigational capabilities. Similarly to ESWM, EPN is designed to leverage episodic memory to facilitate rapid adaptation in unfamiliar environments, and both models are episodic memory-driven transformer-based models.
Hyperparameters- We use the same hyperparameters as those reported by Ritter et al. 20200)
environment structure variations and amount of data, approximately 1.14e9 steps, which is restrictive data regime for EPN. EPN is trained to navigate through reinforcement learning, using IMPALA (Espeholt al. 2018 ESWM trained to build internal model of environments from sparse disjointed episodic experiences through self- supervision.
Training- Both models trained the
Memory bank - EPN has memory bank containing 200 transitions, whereas ESWM only stores the minimal spanning tree, with 19 transitions for smaller environments and 36 for larger ones. When evaluating EPN ability to navigate, we follow the methodology from Ritter et al. 2020). EPN explores to populate its memory bank for the first 2/3 of episodes and the performance only measured on the last third.
I IMPACT STATEMENTS
This paper aims to advance the foundations of machine learning by developing model for struc- tured memory and spatial understanding. Although our primary contributions are theoretical and methodological, the ideas explored here connect directly to real-world applications in autonomous exploration, navigation, and decision- makiin.. particular, the compact latent representations and rollout capabilities of ESWM make at promising component for future embodied systems operating in unstructured physical environments.
Potential positive societal impacts include safer and more efficient navigation technologies, im- proved exploration strategies for robotics, and progress toward Al systems that reason about space ways inspired by human cognition. At the same time, recognize the risks associated with deploying autonomous agents in real-world settings, including the possibility of collecting sensitive data or influencing high-stakes decisions. These considerations underscore the importance of sponsible development, careful evaluation in physical contexts, and adherence to ethical guidelines when integrating such models into embodied platforms.
36


Reviews:
[{"review_id": "VSh6zWlZSH", "reviewer": "Reviewer_d6hX", "summary": "The paper proposes a novel framework Episodic Spatial World Model (ESWM) for constructing cognitive maps from disjoint episodic memories, akin to human learning behavior. Instead of using long consecutive trajectories over the environment, ESWM uses a meta-learning algorithm across many environments to infer unseen transitions from a sparse memory bank of one-step experiences, learning spatial regularities on the way. Experiments demonstrate useful properties of the framework, including the ability to learn an internal representation that closely mirrors the environment layouts with minimal observations, and being sample-efficient and adaptive.", "questions": "1. Since ESWM essentially learns to infer spatial structure, have you considered comparing it against a non-learned mapping baseline?\n2. How crucial is it for the memory bank to form a perfect spanning tree of the environment? In cases where the episodic memories are incomplete, how might ESWM perform or adapt?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The empirical evaluation is thorough, rigorous, and insightful. The authors test ESWM on multiple tasks and environment variants (open arenas vs. maps with random walls/obstacles) and over multiple base models (transformers, LSTM, and Mamba) to assess its prediction accuracy, representation quality, and control performance. The set of experiments covers many insightful explorations of the ESWM behavior, including memory integration, navigation, adaptation to environment change, and scaling. Open Arena training also reveals nontrivial performance difference between the transformer ESWM and the other two variants. \n2. This work has substantial implications for both reinforcement learning and cognitive modeling. By showing that an agent can rapidly form a useful world model from a handful of potentially disjoint experiences, it highlights a pathway toward more generalizable and lifelong learning in AI. In practical terms, the ability to plan efficient exploration and navigation with less training could be impactful for robotics or embodied AI, where an agent might get only sparse observations of a new environment before needing to act.", "weaknesses": "1. The innovation is a bit limited, and there is no baseline outside of the ESWM model to compare with, making it hard to ground the model properties in context. Given the nature of the base models and the environment, the key innovation seems to be the training scheme over multiple memory banks. On the grounding, it would be interesting to see how the model performs against ordinary baselines or training schemes. For example, one may compare ESWM-T with ordinary transformers (with a single memory bank). Comparing ESWM-T with a model trained with longer trajectories would also provide insight into how sparse episodic memory effects learning. \n2. While the results are strong, the experiments are confined to relatively simple, synthetic environments. Additionally, the environments considered all follow consistent spatial rules. If an environment had more complex dynamics or non-grid connectivity, it’s unclear if the current approach would generalize as well (beyond scale).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5vKaPAhSYg", "reviewer": "Reviewer_3Ryj", "summary": "This paper proposes the Episodic Spatial World Model (ESWM), a framework for building coherent spatial world models from **sparse and disjoint one-step transitional memories**. Inspired by the dual role of the Medial Temporal Lobe (MTL) in processing both episodic memory and spatial information, ESWM infers the underlying structure of an environment from a 'memory bank'—a set of independent transitions ($s_s​,a,s_e$​)—unlike prior models that rely on continuous trajectories.\n\nThe authors demonstrate the efficacy of ESWM through experiments in various grid-world environments (Open Arena, Random Wall, and MiniGrid). The results show that ESWM achieves higher prediction accuracy than the sequence-based model, TEM-T, and exhibits strong generalization to environments with unseen structures. Through ISOMAP visualizations, the authors reveal that ESWM's latent space forms a geometric map that accurately reflects the physical topology of the environment, including obstacles. This learned model enables high performance on downstream tasks such as zero-shot exploration and navigation. Furthermore, by decoupling memory from reasoning, the proposed architecture allows for rapid adaptation to environmental changes by simply updating the memory bank without retraining the model.\n\nThe paper's claims are well-supported by a detailed appendix that enhances reproducibility. It includes specifics on the model architecture, data generation pipeline, and additional results, such as the extension to the high-dimensional, vision-based MiniGrid environment (Figure 8), providing substantial evidence for the method's effectiveness.", "questions": "- In the Figure 2 results, the performance of Mamba and LSTM architectures is notably poor in Open Arena but significantly better in Random Wall. This is counter-intuitive, as the Random Wall task seems more complex due to obstacle wall. Could the authors provide an explanation for this performance discrepancy?\n    \n- Regarding the experiments on non-minimal memory banks (lines 293-295), you show that ESWM-T performs well. Does this finding hold for other architectures like ESWM-Mamba and ESWM-LSTM, as well as for the TEM-T baseline? Additionally, could you elaborate on why a duplicated memory item might be problematic? Is it because it reduces the diversity of transitions sampled within a finite memory capacity?\n    \n- In Appendix A.4, you describe three different query types used for the Random Wall experiments (unseen, seen, and unsolvable). Could you please provide a performance breakdown for each query type? I suspect this data might help explain why ESWM-Mamba and LSTM perform better in the Random Wall setting.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Important Research Question and Interdisciplinary Connection:** This paper addresses a significant limitation of prior world models: their difficulty in inferring robust spatial knowledge from continuous sequences and adapting to environmental changes. The authors propose the Episodic Spatial World Model (ESWM), a novel approach inspired by the dual role of the Medial Temporal Lobe (MTL) in processing both episodic and spatial memory. This represents an excellent example of leveraging insights from neuroscience to advance machine learning research.\n    \n- **Comprehensive and In-depth Experiments:** The authors validate their claims through a diverse and rigorous set of experiments. By evaluating prediction accuracy in various environments (Open Arena, Random Wall), visualizing the latent space with ISOMAP, and demonstrating strong performance in zero-shot exploration and navigation, the paper not only shows that ESWM is effective for spatial reasoning tasks but also provides intuitive evidence for why it works.\n    \n- **High Reproducibility and Detailed Analysis:** The manuscript enhances the credibility and reproducibility of its results by providing meticulous details on model architectures and experimental setups in both the main paper and the appendix. The authors offer specific descriptions of the query generation (Appendix A.4), the ISOMAP visualization process (Appendix A.6), and provide extensive additional analyses (Appendix B) and visualizations (Appendix C), which lend strong support to their findings.", "weaknesses": "- **Insufficient Discussion of Closely Related Work and Unclear Novelty:** The core mechanism of the proposed method shares significant similarities with Generative Temporal Models with Spatial Memory (GTM-SM) [1], particularly the idea of leveraging one-step transitions to build a spatially-aware model. GTM-SM also demonstrated that its inferred state representations align with the true geometry and enable high-quality long-term predictions (e.g., Figure 2 in [1]). However, this paper lacks any discussion comparing ESWM with GTM-SM. I believe a thorough comparison is crucial for contextualizing the contributions of this work and clarifying its novelty. I strongly encourage the authors to incorporate this discussion.\n    \n- **Limited Evaluation Outside of Grid-World Environments:** While the authors test their model in several environments, all of them are grid-based. The discrete actions and cyclic nature of grid-worlds may provide implicit structural cues that simplify the task of learning spatial relationships. It is unclear whether ESWM's strong performance would generalize to more realistic, continuous environments (e.g., DMLab [2]) where such cues are absent. The lack of experiments in these settings makes it difficult to assess the true robustness and scalability of the proposed method.\n    \n- **Compliance with LLM Usage Policy:** The authors do not explicitly state whether or how they used Large Language Models (LLMs) in preparing their submission, as required by the ICLR 2026 Author Guide.\n    \n- **Minor Presentation Issues:**    \n    - In Figure 1a (middle), the caption \"Specific parts of the environment may be changed dynamically\" could be misinterpreted as changes occurring mid-trajectory. It would be clearer to state that wall configurations are randomized across different episodes.\n    - In Figure 2, the font size is too small, and the x-axis labels and ticks are occluded. In the caption, `Open Arena` should be $\\texttt{Open Arena}$, and the formatting for ”-T” is inconsistent.\n    - In Figure 3, the font sizes are too small for easy reading.\n\n[1] Fraccaro, Marco, et al. \"Generative temporal models with spatial memory for partially observed environments.\" _International conference on machine learning_. PMLR, 2018.\n\n[2] Beattie, C., Leibo, J. Z., Teplyashin, D., Ward, T., Wainwright, M., Kuttler, H., Lefrancq, A., Green, S., Vald ¨ es, ´ V., Sadik, A., Schrittwieser, J., Anderson, K., York, S., Cant, M., Cain, A., Bolton, A., Gaffney, S., King, H., Hassabis, D., Legg, S., and Petersen, S. DeepMind Lab. CoRR, abs/1612.03801, 2016.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yH2Jo4EGaA", "reviewer": "Reviewer_xMiP", "summary": "This paper introduces the Episodic Spatial World Model (ESWM), a framework that\nconstructs spatial representations from sparse, disjoint episodic memories. The\nmodel is meta-trained to predict unseen transitions from minimal one-step\ntransition memories across diverse environments. The authors demonstrate that\nESWM forms geometric latent representations, enables zero-shot exploration and\nnavigation, and adapts rapidly to environmental changes.", "questions": "1. Can you provide evidence that biological systems construct spatial maps from genuinely disjoint (not just sparsely sampled) experiences?\n2. How does your approach differ fundamentally from learning successor representations or other graph-based transition models? What unique capability does ESWM provide?\n3. Have you tested on substantially larger environments (100+ states)? What is the scaling behavior?\n4. Why is transformer attention specifically required? Have you tried other relational reasoning architectures (e.g., Graph Neural Networks, which seem natural for transition graphs)?\n5. Can you clarify the relationship between your work and existing transition system models, particularly for grid cells and hippocampal representations?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Well-documented implementation: The paper provides comprehensive technical details including pseudocode (Algorithms 1-4), extensive appendices, and thorough experimental setup descriptions. This supports reproducibility.\n2. Strong downstream performance: ESWM demonstrates impressive zero-shot capabilities on exploration (96.48% of oracle performance) and navigation tasks (96.8% success rate, 99.2% path optimality), outperforming the task-specific EPN baseline by substantial margins.\n3. Interesting architectural insights: The finding that transformer attention mechanisms are critical for this task (while LSTM and Mamba struggle in Open Arena) provides valuable insights about the role of content-addressable memory in spatial reasoning.\n4. Adaptability demonstration: The ability to handle environmental changes by simply updating memory banks (93% success rate with new obstacles) is a genuine advantage over weight-based world models.\n5. Emergence of spatial structure: The spontaneous emergence of environment-like manifolds in latent space (Figures 3, 9-13) is compelling, especially the model's ability to infer spatial relationships from obstacle observations and boundary shapes.", "weaknesses": "## Major Concerns\n1. Insufficient Justification of Core Premise.\nThe abstract and introduction claim that animals/humans build spatial maps from \"disjoint experiences governed by consistent spatial rules,\" but no neuroscience evidence is provided to support this specific claim. The cited disruption studies (lesions, amnesia) show that MTL is important for spatial cognition and episodic memory, but don't demonstrate that humans actually construct maps from genuinely disjoint, fragmented experiences.\nIn reality, human and animal spatial experience is largely continuous and spatiotemporal. We don't teleport randomly through environments. The authors should either:\n  * Provide evidence that biological systems actually face and solve this \"disjoint memory\" problem, or\n  * Reframe their contribution as addressing a computational/AI challenge rather than claiming biological inspiration for this specific aspect\n\n2. Unclear Relationship to Prior Work on Transition Systems.\nThe memory bank structure (graph of one-step transitions) closely resembles successor representations and transition graphs that have been extensively studied in both:\n  * Hippocampal/entorhinal modeling: Stachenfeld et al. (2017), and related work on spatial transition representations in grid cells\n  * General cognitive maps: Work from Behrens' lab on relational structure\n\nThe paper mentions Stachenfeld et al. only briefly and doesn't adequately distinguish ESWM's contribution from this and other related substantial body of work. How does ESWM's approach differ conceptually from learning successor representations or other transition models? This needs clear articulation.\n\n3. Limited Scope and Generalization Concerns.\nEnvironment Scale: All tested environments are very small (19-37 locations). With such limited state spaces and the model seeing many different configurations during meta-training, it's unclear whether ESWM is:\n  * Learning generalizable spatial reasoning principles, or\n  * Memorizing patterns across the finite set of possible small-scale configurations\n\nThe one generalization test (size-19 -> size-37) shows substantial performance drops (source->action->end accuracy: 59%->90%->55%), suggesting the model may not scale well.\nEgocentric Setting: While MiniGrid experiments (Figure 8) show promise, they're presented quite briefly. Given that real-world spatial cognition is fundamentally egocentric, this deserves more thorough investigation.\nWhy hexagonal grids? The choice of hexagonal structure is never justified. This seems arbitrary and limits comparability with standard benchmarks.\n\n4. Missing Statistical Rigor.\nMany key results lack error bars or significance tests:\n  * Figure 2: Training curves show variance bands, but test results don't\n  * Figure 4a: No confidence intervals on entropy vs. path length\n  * Figure 5: Some comparisons report significance, others don't\n  * Exploration and navigation comparisons need statistical testing across multiple seeds\nFor a model claiming to learn generalizable principles, demonstrating statistical reliability is essential.\n\n5. Insufficient Ablation Studies.\nThe paper would benefit from:\n  * Ablating the \"minimality\" constraint more thoroughly: Figure 7 shows non-minimal banks work, but how does performance scale with memory bank size? Is there an optimal density?\n  * Analyzing what makes transformers work here: Beyond showing LSTM/Mamba fail, what properties of self-attention are critical? Number of heads? Depth? Context length?\n  * Testing robustness to noise: Real-world transitions are noisy—how sensitive is ESWM to observation or transition noise?\n\n\n## Minor Concerns\n1. Conceptual and Framing Issues.\nSection 2.2's key question: \"Can spatial maps also emerge in general-purpose models not explicitly trained for navigation?\"\nThis has essentially been answered affirmatively by prior work (TEM itself, Behrens lab work, emergence in language models). The authors should clarify what new dimension their work adds to this question.\n\nSection 2.1: The first paragraph of Related Work feels unfocused, jumping between human cognition, RL, and vision without clear connections.\n\n2. Technical Clarifications Needed.\nGraph theory terminology (Section 3.1): Please verify that \"disjoint\" is used consistently with standard graph theory definitions. In particular, \"disjoint\" typically refers to vertex or edge-disjoint subgraphs, but here seems to mean \"non-sequential transitions\". If that's the case, please state this somewhere.\n\n\"I don't know\" classification: The handling of unobservable regions (17% of queries in Random Wall) is clever, but how is the threshold for outputting \"I don't know\" determined? Is it learned or hand-tuned?\n\n\n\n## Minor Comments\n\n* Line 124-125: \"TEM factorizes sensory and structural information\". Clarify how ESWM's approach differs\n* Figure 1c: The masking procedure is clear, but consider showing an example with actual states/actions\n* Section 4.4: The \"zig-zag exploration strategy\" is interesting. Can you characterize this more formally?\n* Appendix B.2: The guided imagination approach using latent space geometry is valuable and could be elevated to the main paper", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CMmbjMwLnm", "reviewer": "Reviewer_o4ni", "summary": "Paper proposes a novel representation of spatial memory for mapping and navigation. A transformer-based model is trained via mask-prediction self-supervised method to predict the unseen region/path. Visualization of the neural network with ISOMAP projection shows that it is able to learn to infer the relative position of disjoint regions. The framework is further extended to the exploration task and goal-seeking task. It can also effectively handle changing environment for the goal-seeking task.", "questions": "1. Does the authors have specific plans for generalizing their methods for real-world environment, i.e. physical embodiment, for their future work?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The inspiration from neuroscience has been a common thread in the major breakthroughs in computational science. In particular, the neural network is proven to be instructional in how computational theories are advanced through such inspiration. While this paper does not directly model brain episodic memory, the concept being borrowed/inspired shows good potential in the application of mapping and navigation.\n\nOne of the major challenges facing machine learning and AI is the effective and efficient representation and modelling of memory within/coupled with neural network training. This work may lead to some interesting ideas for other researchers. S1. Claims are supported by the strong experimental results for the exploration tasks and the goal-seeking task (without global map).\n\nS2. Proposed methods are sensible and applicable for the problem.\n\nS3. Experimental designs and analyses are sound and valid. There are sufficient experiments and ablation studies to support the central claim of the paper. There are additional experiments to compare transformer variants of the model, ESWM-T, against transformer-based such as TEM-T.", "weaknesses": "Minor: The proposed solution may not generalized to continuous environments. However, given the highly theorical framing of the paper, this is not a major concern.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)