Cognitive map formation under uncertainty via local prediction learning
Calvin Yeung, Zhuowen Zou, Nathaniel D. Bastian, Mohsen Imani · Intelligent Systems with Applications · 2025
Cognitive maps are internal world models that enable adaptive behaviour including spatial navigation and planning. The Cognitive Map Learner (CML) has been recently proposed as a model for cognitive map formation and planning. A CML learns high dimensional state and action representations using local prediction learning. While the CML offers a simple and elegant solution to cognitive map learning, it is limited by its simplicity, applying only to fully observable environments. To address this, we introduce the Partially Observable Cognitive Map Learner (POCML), extending the CML to handle partially observable environments. The POCML uses a superposition of states represented via random Fourier features for probabilistic representation and uses the binding operation for parallel state updates. It features an associative memory to enable adaptive behaviour across environments with similar structures. We derive local update rules based on the POCML’s probabilistic state representation and associative memory. We show that a POCML is capable of learning the underlying structure of an environment via local next-observation prediction learning. In addition, we show that a POCML trained on an environment is capable of generalizing to environments with the same underlying structure but with novel observations, achieving good zero-shot next-observation prediction accuracy, significantly outperforming sequence models such as LSTMs and transformers. Finally, we present a case study of navigation in a two-tunnel maze environment with aliased observations, showing that a POCML is capable of effectively using its probabilistic state representations for disambiguation of states and spatial navigation.