Paper: Learning Successor Representations with Distributed Hebbian Temporal Memory (ICLR-2024)

Reviewer comment (Experiments):
The method as presented wouldn’t scale to infinite horizon or large POMDPs. At the very least, I’d like to see a discussion on the challenges presented when scaling DHTM.

Author response:
**Our response:** The main idea of our agent architecture is to combine planning and TD learning, which results in n-step TD learning. It doesn’t need infinite horizon predictions as it learns linear function mapping from hidden space to Successor Features space. We only need to predict the first several steps — to speedup TD learning and reduce its bias — and approximate the tail with the learned mapping (see Eq. (4)).

Scaling to POMDPs with large observation space is also possible due to usage of multiple feature variables, which constitute sparse distributed representation of observations. Sparsity [of observations and connections] is the important key to deal with high dimensionality in our approach. To show that, we also added experiments in AnimalAI environment, which gives more complicated visual input to the agent. 
We should have stated it clearer, as scalability is what we aimed for designing DHTM. Nevertheless, our method still may have scalability issues, so we will add corresponding discussion to the Conclusion section. Thank you for pointing this out.

Reviewer comment (Reproducibility):
The empirical methodology and agent architecture are poorly explained, which makes it hard to evaluate the method. More direct questions can be found below.

Author response:
**Our response:** We agree that the agent’s architecture could be better explained. We took your comment into consideration and will extend the description of the agent and the experimental setup.

Reviewer comment (Experiments):
The empirical results are of limited value; specifically, the authors use a single custom environment with non-standard metrics (see below for more questions about pseudo-surprise).

Author response:
**Our response:** We agree that our custom environment experiments may be not very convincing, so we will add results of experiments in AnimalAI, which is a known benchmark for Reinforcement Learning architectures [2]. Besides, we are going to add standard n-step surprise metrics and qualitative examples of the resulting value functions. We have also tested our algorithm in a 2D maze experiment, similar to what is used in the paper you mentioned, and will add the results to the paper soon. 

[1] Eszter Vertes and Maneesh Sahani. A neurally plausible model learns successor representations in partially observable environments. NeurIPS 2019.  
[2] Crosby, Matthew, Benjamin Beyret, and Marta Halina. 2019. ‘The Animal-AI Olympics’. Nature Machine Intelligence 1 (5): 257–257. https://doi.org/10.1038/s42256-019-0050-3.

Reviewer comment (Novelty):
The paper could be better positioned in the literature regarding successor representations in POMDPs. For example, I expected to see a discussion with [1], which seeks to model an SR in POMDPs.

Author response:
**Our response:** We agree that our custom environment experiments may be not very convincing, so we will add results of experiments in AnimalAI, which is a known benchmark for Reinforcement Learning architectures [2]. Besides, we are going to add standard n-step surprise metrics and qualitative examples of the resulting value functions. We have also tested our algorithm in a 2D maze experiment, similar to what is used in the paper you mentioned, and will add the results to the paper soon. 

[1] Eszter Vertes and Maneesh Sahani. A neurally plausible model learns successor representations in partially observable environments. NeurIPS 2019.  
[2] Crosby, Matthew, Benjamin Beyret, and Marta Halina. 2019. ‘The Animal-AI Olympics’. Nature Machine Intelligence 1 (5): 257–257. https://doi.org/10.1038/s42256-019-0050-3.

Reviewer comment (Theory):
Section 2.1, shouldn't optimal policy be defined as $\text{argmax}_a$ not $\text{softmax}_a$?

Author response:
**Our response:** Thank you noticing that! Our algorithm employs softmax strategy for enabling exploration, however, indeed, the optimal strategy is deterministic.

Reviewer comment (Theory):
The SR in RL is defined with respect to a policy; in DHTM, the target policy is ignored. Is this on purpose? If so, could authors comment on this and the difficulties of adapting their estimate with a changing policy?

Author response:
**Our response:** Thank you for commenting on that, we will add corresponding clarifications to the manuscript soon. We include actions in the model by forcing some of the variables $H^k_{t-1}$ to represent actions. That is, we assume that information about action is included in the hidden state of the model. If we have 4 actions, we set 4 states for one of the hidden variables and set its state from observation of the action. We form on-policy SRs, i.e. relying on policy iteration theorem. DHTM can be used to form SRs for any policy, given it learned dynamics of the environment.

Reviewer comment (Theory):
Section 2.3, I’m quite confused about this section; although the SR could be used to model observations, you’ve discussed the SR up to this point as being a good representation of the hidden state. I’m not sure that the expected discounted observation occupancy, given a state, is the same object. It feels like successor features should be discussed here, as there’s an implicit assumption that the reward function lies in the span of the observation features. You might not be able to linearly predict the value function $V(h)$ given $M^\pi(\cdot | h) \in \mathbb{R}^{|\mathcal{O}|}$.

Author response:
**Our response:** Thank you for noticing that! We have called it Successor Representation as a widespread term among neuroscientists (see [1] for example), and we aimed to bridge knowledge about place cells to our work. Nevertheless, we agree that it’s not common to define Successor Representation in observation space, and this may cause confusion for the reader familiar with the topic from the RL field. We will reformulate our work in terms of the Successor Features framework. In DHTM, variables $O$ are encoded observations, so they are, in effect, already in feature space, not raw observations from the environment. Variables $O$ are observations with respect to the memory itself, but not to the agent. To avoid confusion, further, we call these variables $\Phi$ to distinct them from raw observations.

[1] Samuel J. Gershman. 2018. ‘The Successor Representation: Its Computational Logic and Neural Substrates’. The Journal of Neuroscience 38 (33): 7193. https://doi.org/10.1523/JNEUROSCI.0151-18.2018.

Reviewer comment (Writing):
Section 4, can you describe how the pseudo-surprise is computed? I’m not sure I fully understand what’s going on here.

Author response:
**Our response:** As we mentioned in the text, to measure pseudo-surprise of SF, we do the following:
1. Normalize SF by summing it over corresponding variables $\Phi$ and dividing SF by these sums. The result is SF-induced probability distribution $p(\varphi^k)$ of $\Phi^k$ variables. 
2. Measure average surprise over future observed states $\varphi^k$, according to this distribution: $-\log p(\varphi^k=j)$, where $j$ is the observed state.

Normalized SF represents future observation (feature) profile for the current state. Pseudo-surprise shows whether SF is consistent with the observed feature states or not. For example, if SF doesn’t predict feature $j$ ($p(\varphi^k=j) = 0$), but we observe it, this’ll result in infinite surprise, which means that the SF is of a bad quality. Thank you for your comment, we will devote more space to this explanation in the text. For comparison, we will also put this metric calculated for uniform SFs.

Reviewer comment (Reproducibility):
Can you provide more details about the baseline methods and how they are employed in the architecture? e.g., what’s the input to the RNNs?

Author response:
We have uploaded an updated version of the text. We have added baselines description to the Appendix and formulation of value function decomposition we use.

Reviewer comment (Presentation):
There remains three-quarters of a page that could be employed to break down and refine Section 3.1 on Distributed Hebbian Temporal Memory. Dividing this section into smaller subsections would aid in better organizing the material.

Author response:
Thank you for your advice! We are going to divide Section 3.1 into several subsections and add a glossary at the end of the paper as soon as possible.

Reviewer comment (Writing):
Incorporating a notation glossary, either in the main paper or in an appendix, would be a valuable addition.

Author response:
Thank you for your advice! We are going to divide Section 3.1 into several subsections and add a glossary at the end of the paper as soon as possible.

Reviewer comment (Experiments):
Although the model itself is quite compelling, the experiments fall short of doing it justice. Only a single experimental result is provided, where the performance is marginally inferior to one of the baseline models, the fchmm. There is potential for more comprehensive experiments that can showcase all the intriguing aspects of this model.

Author response:
We agree that one experimental result may be not enough to justify the model. Therefore, we have added an easily interpretable 2D maze experiment to show qualitatively that DHTM value function estimate is adequate. Also, we have added examples of the learned Successor Features for our initial setup, which are easily interpretable too. Besides, we have added results of experiments in AnimalAI environment to show that our method is scalable and environment agnostic. We will upload an updated version of the paper ASAP. 

Although our method is inferior to FCHMM according to some metrics, we should emphasize that, in contrast to FCHMM and other baselines, our method learns in fully online mode and doesn’t employ error backpropagation algorithm. FCHMM needs the whole sequence being available while updating learnable parameters. So, FCHMM provides a lower [surprise] bound that we seek to achieve, using only online Hebbian learning. However, hidden state representations, resulting from FCHMM are, in contrast to our method, dense, and that hinders linear mapping of hidden space to Successor Features. We will discuss this and other limitations of the baselines in the Conclusion section.

Reviewer comment (Writing):
In the factor graph model, what do the superscripts $i$ and $k$ correspond to?

Author response:
Superscripts $i$ and $k$ are used to enumerate different variables as in our model we have multiple hidden and feature variables at the same time, similarly to the Factorial HMM[1].

[1] Ghahramani, Z., and M.I. Jordan. 1997. ‘Factorial Hidden Markov Models’. Machine Learning 29 (2–3): 245–73. https://doi.org/10.1023/a:1007425814087.

Reviewer comment (Experiments):
The authors only consider one simple environment. It will be interesting to study their approach in 2D maze, such as minigrid environments which are commonly used to study successor representations. This environment will also allow more analysis and make better and more accurate comparisons with ground truth successor representations.

Author response:
**Our response:** Thank you for your suggestion! We agree that experiments in 2D maze would be easier to interpret. Therefore, we are going to add them into an updated version of the manuscript.

Reviewer comment (Evaluation):
At the moment, it is hard to understand why their proposed model is performing better than the baseline. The baseline models seem to have more complex architectures but no learning curves were presented to verify that the complexity reduces the learning efficiency of these models.

Author response:
**Our response:** Thank you for pointing this out! We will add learning curves as well.

Reviewer comment (Theory):
In algorithm 1, link it with the equations being optimized during the learning process. At the moment, it is hard to tell how your proposed algorithm contributes to the standard reinforcement learning framework.

Author response:
**Our response:** Algorithm 1 scarcely portrays common training procedure for all methods used, not only for DHTM. Algorithm 1 is not valuable as a contribution, rather it is a clarification of memory testing procedure. There is no strict correspondence between steps described in the Section 3.1 and Algorithm 1, rather it’s a description of the agent's architecture workflow.

Reviewer comment (Presentation):
“The next step after computing p(h_{t}^{k}) distribution parameters is to incorporate information about .." How does this relate to algorithm 1?

Author response:
**Our response:** Algorithm 1 scarcely portrays common training procedure for all methods used, not only for DHTM. Algorithm 1 is not valuable as a contribution, rather it is a clarification of memory testing procedure. There is no strict correspondence between steps described in the Section 3.1 and Algorithm 1, rather it’s a description of the agent's architecture workflow.

Reviewer comment (Reproducibility):
How are the weights w_u_l being learned? Is it through equation 18? If it is, eq 18 should be moved to the main paper.

Author response:
**Our response:** In Section 3.1 we mention that synaptic segment weights $w_{ul}$ are learning through local Hebbian-like plasticity. They are targeted to represent probability $p(s_l=1\ |\ c^u_{t-1}=1)$ that segment $s_l$ is active, given cell $u$ was active at the previous time-step. It can be learned just by counting activation coincidences and mismatches. But in our algorithm it is approximated as exponential moving average of segment’s $s_l$ frequency activation, given $c^u_{t-1}=1$: $\Delta w_{ul} = \alpha * \mathbb{I}[c^u_{t-1}=1] * (\mathbb{I}[s_l=1] - w_{ul})$, where $\alpha \in [0, 1)$ — learning rate. We will add this explanation to the main text to avoid reader confusion.

Reviewer comment (Reproducibility):
Details for the agent’s architecture are missing. Please provide them in the appendix. Ideally, given that you proposed a novel method, a computational graph of your architecture should also be included so that the readers can relate it to figure 1 and algorithm 1.

Author response:
**Our response:** Thank you for your interest! We will add a more detailed description of the algorithm to an Appendix.

Reviewer comment (Experiments):
Figure 4 seems to have very high variance. How many seeds are being used to generate the results in this plot? Looks like more seeds are required.

Author response:
**Our response:** As we mentioned in Section 4, we tested three seeds for every setup. We agree that, perhaps, we need to increase the number of trials to show stable results. However, we should note that high variance may be inevitable, due to usage of non-deterministic softmax strategy. We will update experiments and corresponding figures ASAP.