Cognitive map formation under uncertainty via local prediction learning

We thank the reviewer for their insightful feedback.

W1. Observations in POCML are assumed to be one-hot as a simplifying assumption for theoretical development and experimental validation. Computationally, the tokenization of state and action can be generalized to account for more general state/action representation. In future work, we are interested in how POCML may model continuous actions and their interpretations.

W2/Q2. While TEM and CSCG are alternative probabilistic models that learn in partially observable environments, their complex computation typically renders the model much more complex than POCML, and the resulting representation of state and action is hard to interpret. In comparison, POCML leverages representation with a simple and unified geometric interpretation of state and action: applying action to a state is equivalent to binding the action vectors to state vectors (at the standard representation). This is further supported by the PCA visualization, where the geometry learned by the state and action representation matches directly with the underlying environment structure. In comparison, TEM uses recurrent neural networks (RNNs) to predict the next abstract state, and POCML uses the binding operation. Our setup with POCML enables representations in which the underlying geometry is apparent. This holds even in environments of different structures. CSCG is interpretable in the sense that it learns the transition probabilities between different states. However, compared to TEM and POCML, it cannot generalize to larger environments of the same structure as there is no notion of action representations.

W3. Thank you for pointing it out! We have made the changes in the latest draft of the paper.

Q1. Unfortunately, due to time constraints, we were unable to perform experimentation on more standardized environments such as minigrid. However, we would like to point out that we have evaluated Transformer and LSTM models with many more (up to 100x) parameters on the same tasks and they yielded similar results as before. We point this out in the revised paper. We also discussed the issue with the scale of experimentation in a newly added Limitations section in the appendix.

Thank you for the insightful feedback.

Simplistic experimental setup: We acknowledge the need for more complex environments. While the experimental environments presented in this work are relatively simple in comparison with the common setup in TEM and CSGS, they are purposefully chosen to evaluate the foundational capabilities of POCML as an extension of CML. These setups enable rigorous testing of the model's ability to learn structure and disambiguate observations in partially observable settings, a critical challenge for cognitive map formation. In particular, we followed the practice in CML [1] to test the POCML in the grid and tree structures, which are representative abstractions of more complex environments. Because we assume a partially observable setting, the problem complexity increases significantly from that with the same underlying topology in a CML experiment, leading to reduced-size grid and tree environments for our experimentation. For future work, we aim to validate POCML in larger-scale and more diverse settings, such as 3D mazes or complex graph environments.

Transformers/LSTM fewer than 500 parameters: We agree with the reviewer that the main results (Table 1) of our comparison with Transformers and LSTM use a small number of parameters, though we would like to point out that we have evaluated Transformer and LSTM models with many more (up to 100x) parameters on the same tasks and they yielded similar results as before. The parameter choice in Table 1 is intended to investigate the model performances with a similar number of trainable parameters. We have emphasized this in the revised paper.

[1] Stöckl, Christoph, Yukun Yang, and Wolfgang Maass. "Local prediction-learning in high-dimensional spaces enables neural networks to plan." Nature Communications 15.1 (2024): 2344.

Thank you for the positive review and the insightful feedback.

Unclear Motivation: Thank you for pointing out the ambiguity in our motivation. The main goal of our paper is to extend the existing Cognitive Map Learner (CML) algorithm to handle partially observable environments, a setting adopted by concurrent theoretical models of cognitive map formation, resulting in the proposed Partially Observable Cognitive Map Learner (POCML).

Lack of Limitation: Thank you for pointing out the lack of limitation sections in our work. We have followed your advice and added a limitations section to the appendix. This work has several limitations. First, evaluation was performed on simplistic environments. These experiments illustrate the properties and functionality of the POCML model but further work is needed to demonstrate the practicality of the model in larger scale applications. However, we would like to note that the purpose of this work is to introduce the POCML model and the novel representational paradigm based on RFF for encoding uncertainty; we leave extending the model to larger-scale environments for future work. Second, the POCML model here uses discrete observations. Third, as in the original CML work, action affordances have to be provided by the environment.

Experimental setup: We acknowledge the need for more complex environments. While the experimental environments presented in this work are relatively simple in comparison with the common setup in TEM and CSGS, they are purposefully chosen to evaluate the foundational capabilities of POCML as an extension of CML. These setups enable rigorous testing of the model's ability to learn structure and disambiguate observations in partially observable settings, a critical challenge for cognitive map formation. In particular, we followed the practice in CML [1] to test the POCML in the grid and tree structures, which are representative abstractions of more complex environments. Because we assume a partially observable setting, the problem complexity increases significantly from that with the same underlying topology in a CML experiment, leading to reduced-size grid and tree environments for our experimentation. For future work, we aim to validate POCML in larger-scale and more diverse settings, such as 3D mazes or complex graph environments.

We sincerely thank the reviewer for the detailed constructive feedback. It has been immensely helpful to us. We have tried to follow your feedback in the revised copy of the paper and for what we could not immediately apply we will definitely consider for our future work.

Unfortunately, due to time constraints, we could not perform experimentation on more standardized environments such as minigrid. However, we would like to point out that we have evaluated Transformer and LSTM models with many more (up to 100x) parameters on the same tasks and yielded similar results.

Regarding weakness (2), POCML works best in environments with regular structures (e.g. grids and trees), as action and state representations are decoupled. This is because, taking after TEM, the motivating example used in developing the model is to learn spatial structure. So we are mostly concerned with environments of that type here, explaining our choice of environments, which are also used in the TEM paper. To improve model performance in more complex environments, we plan to extend it in future work.

Based on your feedback, to address weakness (5), we added a related work section in the appendix comparing in detail POCML against TEM and CSCG.

Specifically, TEM and POCML are similar in that both have a separation between abstract state and observation, but TEM uses the conjunctive representation $p$ in its memory, a Hopfield network. In contrast, POCML uses a hetero-associative memory that associates state and observation implemented as a table of expected counts. TEM uses recurrent neural networks (RNNs) to predict the next abstract state while POCML uses the binding operation. While RNNs are more expressive, properties of binding enable a geometric interpretation of the learned representations: next state prediction can be performed by adding the action representation to the state representation in this space. That said, TEM does learn grid-cell representations that reflect biology and explain the remapping phenomenon; thus, it is valuable from a neuroscience perspective.

From a computational perspective, in both learning and inference, TEM requires iterative processing to perform memory retrieval and state and observation inference, which produces a sizable computational overhead. POCML does not have this limitation.

Algorithmically, there is a comparatively greater difference between CSCG and POCML. Both models can be considered action-augmented hidden Markov models (HMMs) though actions are treated slightly differently in both models. Moreover, CSCG uses a variant of an HMM called a cloned HMM (CHMM). While CSCG represents state-transition probabilities via a state-transition matrix, POCML does this through kernel density estimation (based on RFF). Given that POCML uses vector representations for states and actions, it can easily be extended to larger environments with the same regularity structure simply by keeping the action representations fixed.

We also discuss the relation to RL here to address weakness (7), which we also include in the appendix. It is important to note that POCML learns the structure of the environment in a reward-free manner and the resulting state and action representation is independent of policy, which sets it apart from value-based reinforcement learning (RL) techniques. That said, the structure of the environment learned by the POCML encodes environment dynamics $p(s'|s,a)\propto \delta(\phi(\mathbf{s}'),\phi(\mathbf{s})\odot\phi(\mathbf{a}))$, which can subsequently be used for model-based RL. The tunnel-maze experiment in the results section demonstrates a way in which this model of the environment can be used given a policy. It is of interest to investigate the consequence of a tighter coupling between model and policy in future work. As you point out, Successor Representations (SR) have been used in RL as a reward-free representation of states. However, SRs consider distributions over future states $s’$ at any given state $s$ via $M(s,s’)$, given a policy so they are still policy-dependent. It has been pointed out that SRs need to be recomputed for dynamic planning [A]. It is also unclear how it can generalize beyond known states as it lacks action representations. POCML representations, conversely, can be used to generate a policy. One could potentially connect SRs and POCML representations by considering the relation between $M(s,s’)$ and $\sum_a p(s'|s,a)p(a)$ (from above), but this still needs further exploration.

To address weakness (4), we have added an explanation of the binding operation and why this property holds in the background section.

We have also clarified both PCA figure in the results section and that $\Delta$ does not refer to an operator; $\Delta Q$ simply refers to the matrix that we add to $Q$ when applying the update rule.

[A] Stachenfeld, K. L., Botvinick, M. M. & Gershman, S. J. The hippocampus as a predictive map. Nat. Neurosci. 20, 1643–1653 (2017).