Learning Successor Features with Distributed Hebbian Temporal Memory

## Abstract

Abstract This paper presents a novel approach to address the challenge of online temporal memory learning for decision-making under uncertainty in non-stationary, partially observable environments. The proposed algorithm, Distributed Hebbian Temporal Memory (DHTM), is based on factor graph formalism and a multicomponent neuron model. DHTM aims to capture sequential data relationships and make cumulative predictions about future observations, forming Successor Features (SF). Inspired by neurophysiological models of the neocortex, the algorithm utilizes distributed representations, sparse transition matrices, and local Hebbian-like learning rules to overcome the instability and slow learning process of traditional temporal memory algorithms like RNN and HMM. Experimental results demonstrate that DHTM outperforms LSTM and a biologically inspired HMM-like algorithm, CSCG, in the case of non-stationary datasets. Our findings suggest that DHTM is a promising approach for addressing the challenges of online sequence learning and planning in dynamic environments.

## 1 Introduction

Modelling sequential data is one of the essential tasks in artificial intelligence as it has many applications, including decision-making and learning world models [], natural language processing [], conversational AI [], time series analysis [], and video and music generation []. One of the classical approaches to modelling sequential data is to form a representation that stores and condenses the most relevant information about a sequence, and to find a general transformation rule of this information through the dimension of time [, , ]. We refer to the class of algorithms that use this approach as Temporal Memory (TM) algorithms because they essentially model the cognitive ability of complex living organisms to remember the past experience and make future predictions based on that memory [, , , ].

This paper addresses the problem of online sequence learning for planning and decision-making under uncertainty, which can be formalized as reinforcement learning (RL) for a partially observable Markov decision process (POMDP) [, ]. Inferring the hidden state in a partially observable environment is a sequence modelling problem, since it requires processing a sequence of observations to obtain enough information about the hidden state. One of the most efficient representations of hidden states for discrete POMDPs is the Successor Representation (SR), which disentangles hidden states and goals given by the reward function [, , ]. The Successor Features framework is an extension of SR to continuous POMDP, using the same idea of value function decomposition, but instead for features of a hidden state [, ]. Temporal memory (TM) algorithms can be used to make cumulative predictions about future states and their features, forming SR or SF. This work shows that the proposed algorithm, Distributed Hebbian Temporal Memory (DHTM), can dynamically form SFs for navigation tasks in Gridworld and AnimalAI [] environments.

The most prominent TM algorithms, such as a Recurrent Neural Network (RNN) [] or LSTM [], use backpropagation to capture data relationships, which is known to be unstable due to recurrent nonlinear derivatives. They also require complete data sequences to be available during training. Although the gradient vanishing problem can be partially circumvented in a way that Receptance Weighted Key Value (RWKV) [] or Linear Recurrent Unit (LRU) [] models do, the problem of online learning, i.e. learning on non-stationary data sets, is still a viable issue. Unlike graphical models such as HMM [], RNN models and their descendants also lack a probabilistic theoretical foundation, which is advantageous for modelling sequences captured from stochastic environments [, ]. There is little research on TM models that can be used in fully online adaptive systems interacting with partially observable stochastic environments with access to only one sequence data point at a time, a common case in reinforcement learning [].

We propose a Distributed Hebbian Temporal Memory (DHTM) algorithm based on factor graph formalism and a multi-compartment neuron model inspired by the HTM neuron []. The resulting graphical structure of our model is similar to that of the factorial HMM []. An important feature of our model is that the transition matrices for each factor are stored as different components (segments) of artificial neurons, which makes computation very efficient in the case of sparse transition matrices. Our TM forms sequence representations entirely online and uses only local Hebbian-like learning rules [, , ], which avoid the drawbacks of gradient methods and make the learning process much faster than gradient methods. This comes at the expense of hidden state generalization, since DHTM’s hidden state uniquely encodes each observation sequence. Different trajectories leading to the same position in the partially observable Gridworld can result in different hidden states inferred by DHTM. In essence, this makes DHTM a model of episodic memory [], or an adaptive trajectory buffer, with built-in machinery for making multistep future predictions. In this light, we also show that an agent forming SF with DHTM can be interpreted as a kind of episodic control [, , , ].

The proposed TM is tested as a world model for an RL agent architecture navigating in a Gridworld environment and a more challenging AnimalAI testbed []. The results demonstrate that our algorithm outperforms classical LSTM and a biologically inspired HMM-like world model CSCG [] in changing environments and is in order of magnitude more sample efficient than Dreamer V3 []. Another advantage of our algorithm is that it can be implemented on neuromorphic processors using only local learning rules.

Our contribution to this work includes the following key elements:

- •
A novel episodic memory model, called DHTM, which is based on a factor graph formalism and features an efficient, bio-inspired implementation.
- •
We demonstrate that the DHTM can operate fully online, using only local Hebbian-like learning rules.
- •
We show that DHTM learns at a speed comparable to model-based table episodic control while supporting distributed representations. This capability significantly broadens its applicability.

## 2 Background

This section provides basic information about some of the concepts necessary to follow the paper.

### 2.1 Reinforcement Learning

This paper considers decision-making in a partially observable environment, usually formalized as a partially observable decision process []. A POMDP is defined as a tuple $\mathcal{M}=(S,A,P,R,O,D,\gamma)$, where $S$—state space, $A$—action space, $P(s,a,s`)=Pr(s`\ |\ s,a)$—transition function, $R(s)$–reward function, O—observation space, $D(a,s`,o)=Pr(o\ |\ a,s`)$—sensor model and $\gamma\in[1,0)$—discount factor, given a transition $s,a\to s`$, where $s\in S$, $a\in A$, $o\in O$. If $S,A,O$ are finite, $P,D$ can be viewed as real-valued matrices, otherwise, they are conditional density functions. Here we consider deterministic rewards that depend only on the current state, i.e. $R(s):S\to\mathbb{R}$.

The task of RL is to find a policy $\pi(a\ |\ s):S\times A\to[0,1]$ that maximizes the expected return $G=\mathbb{E}[\sum^{T}_{t=0}\gamma^{l}R_{t}]$, where $T$ is an episode length. Value-based methods usually aim to estimate the Q-function given a policy $\pi$: $Q^{\pi}(s_{t},a_{t})=\mathbb{E}[\sum_{l\geq t}\gamma^{l}R(s_{l+1})\ |\ s_{t},a
_{t},\pi]$. For an optimal value function $Q^{*}$, an optimal policy can be defined as $\pi(a\ |\ s)=\underset{a}{\mathrm{argmax}}\ Q^{*}(s,a)$.

One of the drawbacks of value-based methods is that they require learning a new Q-function when the reward function changes, even if the environment remains the same. To get around this, the dynamics of the environment and the reward function can be decoupled, which leads us to the successor representation framework.

### 2.2 Successor Representation

Successor representations are such representations of hidden states from which we can linearly infer the state value given the reward function []. We assume that the observation and state spaces are discrete.

$$ $V(h_{t}=i)=\mathrm{E}[\sum^{\infty}_{l=0}\gamma^{l}R_{t+l+1}\ |\ h_{t}=i]=\sum _{j}\mathrm{SR}_{ij}R_{j},$ (1) $$

where $\gamma$ is a discount factor, $i$-th row in the matrix $\mathrm{SR}$ is the Successor Representation of a state $i$, and $\mathrm{SR}_{ij}=\sum^{\infty}_{l=0}\gamma^{l}p(h_{t+l+1}=j\ |\ h_{t}=i)$. $R_{j}$ is a reward for observing the state $j$.

An extension of SR for large or continuous state spaces that we use is Successor Features [], which is based on a similar idea of value function decomposition, but in a feature space instead (see details in Appendix A).

It is easy to show that, similarly to the Q-function, both SR and SF can be obtained using dynamic programming. However, this limits its use to relatively stable environments, since changes in the transition function can’t be taken into account immediately.

The next step to improve an agent’s adaptability is to use a world model to approximate the environment’s transition function and use it directly to form SR according to its definition. And one of the most basic such world models is the Hidden Markov Model.

### 2.3 Hidden Markov Model

A partially observable Markov process can be approximated by a Hidden Markov Model (HMM) with hidden state space $H$ and observation space $O$. $O$ is the same as in $\mathcal{M}$, but $H$ is generally not equal to $S$. Variables $H_{t}$ represent an unobservable (hidden) approximate state of the environment that evolves over time, and observable variables $O_{t}$ represent observations that depend on the same time-step state $H_{t}$, and $h_{t},o_{t}$ are corresponding values of these random variables. For simplicity, we assume that actions are fully observable and that information about them is contained in $H_{t}$ variables. For the process of length $T$ with state values $h_{1:T}=(h_{1},\ldots,h_{T})$ and $o_{1:T}=(o_{1},\ldots,o_{T})$, the Markov property yields the following factorization of the generative model:

$$ $p(o_{1:T},h_{1:T})=p(h_{1})\prod_{t=2}^{T}p(h_{t}|h_{t-1})\prod_{t=1}^{T}p(o_{ t}|h_{t}).$ (2) $$

In the case of a discrete hidden state variable, a time-independent stochastic transition matrix can be learned with the Baum-Welch algorithm [], a variant of the expectation maximization (EM) algorithm. To compute the statistics for the expectation step, it uses the forward-backward algorithm, which is a special case of the sum-product algorithm [].

Since the transition matrix grows quadratically with the size of the hidden state space, it may become infeasible to train the HMM in large state spaces. Since it is possible that not all hidden states are dependent on each other, the Factorial HMM can mitigate this problem by partitioning the hidden state spaces into smaller independent subspaces, providing an efficient way to handle distributed representations.

We base DHTM on ideas from the Baum-Welch algorithm and the Factorial HMM, while adapting them for online learning.

### 2.4 Episodic Control

Episodic Control (EC) is a general computational approach to storing and reusing an agent’s experience, inspired by the notion of episodic memory [] from psychology.

The core of EC is a dictionary $\mathcal{D}_{a}=(K_{a},V_{a})$ that stores key-value pairs for each discrete action $a\in A$. $K_{a}$ is a dynamically growing array of the agent’s state representations, and $V_{a}$ is of the same size and can contain any values that are useful for solving a control task, for example, corresponding Q-value estimates []. In this case, given that the current state of the agent is $h_{t}$, the best action can be found by searching the dictionary for corresponding keys for each action $\underset{a}{\mathrm{argmax}}\ \mathcal{D}_{a}(h_{t})$. In the simplest form of EC, learning means adding new key-value pairs to the dictionary $\mathcal{D}_{a}$.

In this paper, we refer to EC to better explain the principle of the proposed algorithm for solving RL tasks.

## 3 Distributed Hebbian Temporal Memory

### 3.1 Factor Graph Model

Distributed Hebbian Temporal Memory is based on the sum-product belief propagation algorithm in a factor graph (see Figure 1A). Analogous to Factorial HMM [], we divide the hidden space $H$ into subspaces $H^{k}$. There are four sets of random variables (RV) in the model: $H^{i}_{t-1}$—latent variables representing hidden states from the previous time step (context), $H^{k}_{t}$—latent variables for the current time step, and $\Phi^{k}_{t}$–feature variables. All random variables have categorical distributions. RV state values are denoted by corresponding lowercase letters: $h^{i}_{t-1}$, $h^{k}_{t}$, $\varphi^{k}_{t}$, $o^{lm}_{t}$.

Each variable $\Phi^{k}_{t}$ is considered independent and has a separate graphical model for increased computational efficiency. However, in practice, hidden variables of the same time step are statistically interdependent. We introduce their interdependence through a segment computation trick that goes beyond the standard sum-product algorithm (see equation 6).

The model also has three types of factors: $M^{i}_{t-1}$—messages from previous time steps, $F^{k}_{c}$—context factor (generalized transition matrix), $F^{k}_{e}$–emission factor. We assume that messages $M^{i}_{t-1}$ contain posterior information from time step $t-1$, so we don’t plot observable feature variables for previous time steps in Figure 1A.

Figure: Figure 1: A. Example of a factor graph for the DHTM. The input to the model is a sequence of observations $o_{t}$. The encoder block forms categorical features $\Phi^{k}_{t}$. Each feature $\Phi$ has its own explaining hidden variable, which may depend on hidden variables of the other features and on itself from the previous time step. $F^{k}_{c}$ and $F^{k}_{e}$ are context and emission factors for the corresponding variables. Unary factors $M^{i}_{t-1}$ called messages represent accumulated information about previous time steps. B. Neural implementation. Random variables are represented by cell clusters (white circles), where each cell corresponds to a state and its spike frequency to the probability of the state $p(h^{k}_{t})$. The dendritic segments of the cell correspond to the context factor values $f_{l}$ for a given combination of states (active presynaptic cells). Segments’ excitations $E_{l}$ are combined to determine the cell’s spike frequency $p(h^{k}_{t})$. Segment synaptic weights reflect the specificity of the combination of presynaptic cells for the segment. Emission factors $F^{k}_{e}$ are fixed and represented by minicolumns within a variable.
Refer to caption: extracted/5907925/whole_graph_combined.png

We will discuss only the upper block of the graph in Figure 1A, which is DHTM itself. The lower block—an encoder—is described in Appendix C. The only requirement for the encoder is that its output should be represented as states of categorical variables (features) for the current observation.

The main routine of the DHTM is to estimate distributions of hidden state variables given by the equation 3, the computational flow of which is schematically shown in Figure 1B:

$$ $p(h_{t}^{k})\propto\sum_{\Omega_{k}}\prod_{i\in\omega_{k}}M_{t-1}^{i}(h_{t-1}^ {i})F_{c}^{k}(h_{t}^{k},\Omega_{k}),$ (3) $$

where $\Omega_{k}=\{h^{i}_{t-1}:i\in\omega_{k}\}$, $\omega_{k}={i_{1},\dots,i_{n}}$—set of previous time step RV indexes included in $F_{c}^{k}$ factor, $(n+1)$—factor size.

### 3.2 Neural Implementation

For computational purposes, we translate the factor graph to the Hebbian learning neural network architecture (for a biological interpretation of the model, see Appendix B). As shown in Figure 1B, each RV can be viewed as a set of spiking neurons representing the states of the RV, i.e. $p(h^{k}_{t})=p(c^{j}_{t}=1)$, where $j$–is the index of a neuron corresponding to the state $h^{k}_{t}$. Cell activity is binary $c^{j}_{t}\in\{0,1\}$ (spike/no spike), and the probability can be interpreted as a spike rate. Factors $F^{k}_{c}$ and $M^{i}_{t-1}$ can be represented as vectors, where elements are factor values for all possible combinations of RV states included in the factor. Let’s denote the elements of the vectors as $f_{l}$ and $m_{u}$ respectively, where $l$ corresponds to a particular combination of $k$ state values $h^{k}_{t},h^{i_{1}}_{t-1},\dots,h^{i_{n_{l}}}_{t-1}$, and $u$ indexes all neurons representing states of RVs from previous time steps.

Inspired by biological neural networks, we group the connections of a neuron into dendritic segments. A segment acts as an independent computational unit that detects a particular input pattern (a contextual state) defined by its receptive field. In our model, a segment connects the factor value $f_{l}$ and the excitation $E_{l}$ induced by the segment $l$ to the cell to which it is connected. The segment is active, i.e. $s_{l}=1$ if all its presynaptic cells are active, otherwise $s_{l}=0$. Computationally, a segment transmits its factor value $f_{l}$ to the cell to which it is attached if the context matches the corresponding state combination.

We can now rewrite equation 3 as the following:

$$ $p(h_{t}^{k})\propto\sum_{l\in\mathrm{seg}(j)}L_{l}f^{k}_{l},$ (4) $$

where $L_{l}=\prod_{u\in\mathrm{rec}(l)}m_{u}$ is segment’s likelihood as long as messages are normalized to probability distributions, $\mathrm{seg}(j)$—indexes of segments that are attached to cell $j$, $\mathrm{rec}(l)$—indexes of presynaptic cells that constitute receptive field of a segment with index $l$.

Initially, cells have no segments. As learning progresses, new non-zero connections, grouped into segments, are created. In equation 4, we benefit from having sparse factor value vectors because their complexity depends linearly on the number of non-zero components. And that’s usually the case in our model due to the one-step Monte Carlo learning and the special form of the emission factors $F^{k}_{e}$:

$$ $F^{k}_{e}(h^{k}_{t},\varphi^{k}_{t})=\mathbb{I}[h^{k}_{t}\in\mathrm{col}( \varphi^{k}_{t})],$ (5) $$

where $\mathbb{I}$—indicator function, $\mathrm{col}(\varphi^{k}_{t})$ is a set of hidden states connected to the feature state $\varphi^{k}_{t}$ forming a column. The shape of the emission factor is inspired by the presumed columnar structure of the neocortex and has been shown to induce a sparse transition matrix in HMM [].

The segment likelihood $L_{l}$ resulting from the sum-product algorithm is computed as if the presynaptic cells were independent. However, this isn’t usually the case for sparse factors. To account for their interdependence, we substitute the following equation for segment log-likelihood:

$$ $\displaystyle\log L_{l}$ $\displaystyle=w_{l}\log\frac{1}{n_{l}}\sum_{u\in\mathrm{rec}(l)}w_{ul}m_{u}+ \sum_{u\in\mathrm{rec}(l)}(1-w_{ul})\log m_{u},$ (6) $$

where $w_{ul}$—synapse efficacy or neuron specificity for segment, such that ${w_{ul}=p(s_{l}=1|c^{u}_{t-1}=1)}$, $w_{l}=\frac{1}{n_{l}}\sum_{u}w_{ul}$—average synapse efficacy for segment $l$, and $n_{l}$-number of cells in segment’s receptive field.

The idea behind the formula is to approximate between two extreme cases. The first is ${p(s_{l}=1|c^{u}_{t-1}=1)}\to 1$ for all $u$, which means that all cells in the receptive field are dependent and part of a cluster, i.e. they fire together. In this case, ${p(s_{l})=m_{u}}$ for all $u$, but we also reduce the prediction variance by averaging across different $u$. In the opposite case, ${p(s_{l}=1|c^{u}_{t-1}=1)\to 0}$ for all $u$, presynaptic cells don’t cluster and the segment activation probability is just a product of the activation probability of each cell.

The resulting equation for belief propagation in DHTM is the following:

$$ $p(h^{k}_{t})=p(c^{j}_{t}=1)=\underset{j\in\mathrm{cells}[H^{k}_{t}]}{\mathrm{ softmax}}(\underset{l\in seg(j)}{\mathrm{max}}(E_{l})),$ (7) $$

where ${E_{l}=\log f_{l}+\log L_{l}}$, $\mathrm{cells}[H^{k}_{t}]$—indexes of cells representing states for the variable $H^{k}_{t}$. Here we also approximate the logarithmic sum with the $\mathrm{max}$ operation, inspired by the neurophysiological model of segment aggregation by cells [].

The next step after computing the $p(h^{k}_{t})$ distribution parameters is to incorporate information about the current feature states ${p(h^{k}_{t}\ |\ \varphi^{k}_{t})\propto p(h^{k}_{t})\mathbb{I}[h^{k}_{t}\in
\mathrm{col}(\varphi^{k}_{t})]}$. After that, the learning step is executed. The loop closing step of our TM algorithm is to assign the posterior of the current step ${p(h^{k}_{t}\ |\ \varphi^{k}_{t})}$ to $M^{i}_{t-1}$.

### 3.3 Learning

DHTM learns $f_{l}$ and $w_{ul}$ weights by Monte-Carlo Hebbian-like updates. First, $h^{i}_{t-1}$ and $h^{k}_{t}$ are sampled from their posterior distributions: ${p(h^{i}_{t-1}\ |\ \varphi^{i}_{t-1})\propto M^{i}_{t-1}}$ and ${p(h^{k}_{t}\ |\ \varphi^{k}_{t})}$ respectively. Then $f_{l}$ is updated according to the activity $s_{l}$ of the segment and the activity $c^{j}_{t}$ of its cell, so that $f_{l}$ is proportional to several coincidences ${s_{l}=c^{j}_{t}=1}$ in the recent past:

$$ $\Delta f_{l}=\alpha(c^{j}_{t}-f_{l})s_{l},$ (8) $$

where $\alpha\in[0,1)$ is the segment’s learning rate.

That is, $f_{l}$ is trained to be equal to the average cell activation frequency when the segment is active.
It’s similar to Baum-Welch’s update rule [] for the transition matrix in HMM, which in effect counts transitions from one state to another, but in our case, the previous state (context) is represented by a group of RVs.

Weights $w_{ul}$ are also updated by the Hebbian rule to reflect the specificity of a presynaptic $u$ for activating a segment $l$. That is, they are targeted to represent the probability ${p(s_{l}=1\ |\ c^{u}_{t-1}=1)}$ that segment $s_{l}$ is active, given that cell $u$ was active at the previous time step. We could learn this by counting activation matches and mismatches. But in our algorithm, it is approximated as an exponential moving average of the frequency activation of segment $s_{l}$, given $c^{u}_{t-1}=1$:

$$ $\Delta w_{ul}=\beta\cdot\mathbb{I}[c^{u}_{t-1}=1]\cdot(\mathbb{I}[s_{l}=1]-w_{ ul}),$ (9) $$

where $\beta\in[0,1)$ — learning rate.

### 3.4 Agent Architecture

We incorporate DHTM as a part of an RL agent. The same agent is used for other memories tested. The agent consists of a memory model and a feature reward function. The memory model generates SF by predicting cumulative future distributions of feature variables $\Phi^{k}$. We assume that the reward function can be decomposed linearly in feature space as ${R_{t}=\frac{1}{n}\sum^{n}_{k=1}{r^{k}_{i}}}$,
where $r^{k}_{i}$ is a reward associated with feature state $\varphi^{k}_{t}=i$ and $n$—number of feature variables. $r^{k}_{i}$ is learned during interaction with the environment and is used in combination with SF representations to estimate the action value function (see Appendix A for details).

Figure: Algorithm 1 General agent training procedure

The agent training procedure is outlined in Algorithm 1. For each episode, the memory state is reset to a fixed initial message with RESET_MEMORY() and the variable action is set to a fixed initial action. An observation returned by an environment (obs) is encoded as a set of categorical variables. In the OBSERVE() routine, memory learns to predict next feature states as described in section 3. An agent learns associations to feature states and rewards in the REINFORCE() function:

$$ $r^{k}_{i}\leftarrow r^{k}_{i}+\alpha\mathbb{I}[\varphi^{k}_{t}=i](R_{t}-r^{k}_ {i})$ (10) $$

where $\alpha$ is a learning rate, $R_{t}$—a reward for the current time step.

We include actions in DHTM by forcing some of the hidden variables $H^{k}_{t}$ to represent actions. That is, we assume that information about action is contained in the hidden state of the model. For example, if we have 4 actions, we set 4 states for one of the hidden variables and set its state from observing the action.

In SAMPLE_ACION(), for each action and current state, TM predicts feature states up to time step $T$ under policy $\pi$, and the predictions are combined to form SF (equation 11). For simplicity, we always build SF under a uniform policy, which doesn’t guarantee optimality, but we found that it doesn’t make a significant difference to iterative policy improvement in our experimental setups. The planning horizon $T$ is determined dynamically from the distribution of $\Phi^{k}_{t+l}$. If the distribution of $\Phi^{k}_{t+l}$ is close to uniform or states with positive reward have high probability, the prediction cycle is stopped. The action value is determined by combining the corresponding SF and feature state rewards as shown in equation 12. Finally, the action is sampled from the softmax distribution over the action values.

$$ $\mathrm{SF}^{\pi}_{t+T}(\varphi^{k}=j\ |\ h_{t})=\sum^{T}_{l=0}\gamma^{l}p_{ \pi}(\varphi^{k}_{t+l+1}=j\ |\ h_{t})$ (11) $$

$$ $Q^{\pi}(h_{t},a_{t})=\sum_{jk}\sum_{h_{t+1}}\mathrm{SF}^{\pi}_{t+T}(\varphi^{k }=j\ |\ h_{t+1})\cdot p(h_{t+1}\ |\ h_{t},a_{t})\cdot r^{k}_{j},$ (12) $$

### 3.5 DHTM and Episodic Control

Figure: Algorithm 2 Episodic memory learning

It can be shown that DHTM combined with the agent architecture described in the previous section is similar to episodic control.

DTHM performs best when its hidden state space is large. Since the hidden state for unpredictable observations is formed randomly, it’s important that there are no accidental state collisions when remembering transitions. That is, DHTM tries to remember new sequences without interfering with old memories. In effect, this makes DHTM a buffer of agent trajectories with a built-in recall engine. Episodic Control is based on a similar idea, and to investigate its difference from DHTM, we implemented an Episodic Control agent, a simplified counterpart of DHTM, reflecting its core property that presumably helps it solve RL tasks.

In our EC version, $K_{a}$ is an array of hidden states $h_{t}$ and $V_{a}$ is an array of corresponding subsequent hidden states $h_{t+1}$. If the key state $h_{t}$ already exists in the dictionary, its value is rewritten. That is, $\mathcal{D}_{a}$ represents the transition matrix of a deterministic HMM for a feature variable $\Phi_{t}$. The hidden state $h_{t}$ and the feature state $\varphi_{t}$ have the same dependence as in DHTM (equation 5), but here we assume infinite cells per column. To predict the next state, we just have to look up $\mathcal{D}_{a}$ for the current state $h_{t}$, but if there is no entry for $h_{t}$ or the prediction does not match the observed state $\varphi_{t+1}$, then the new state is formed. Unlike DHTM, the new state is not completely random, but to avoid collisions, the algorithm makes sure that this state hasn’t been chosen before. The learning procedure of the EC agent’s memory $\mathcal{D}_{a}$ is summarized in the algorithm 2.

Since $\mathcal{D}_{a}$ can be transformed into a transition matrix, the process of forming SF can be reduced to one in Sec. 3.4. However, since the resulting transition matrix is sparse, it’s more efficient to consider $\mathcal{D}_{a}$ as a graph representation and to perform SF formation as a breadth-first search (see Appendix D.3 for details).

## 4 Experiments

We test our model in a reinforcement learning task in the Gridworld environment and in a more challenging AnimalAI 3D environment. This section shows how different memory models affect the adaptability of an RL agent. For each setup, the results are averaged over five experiments with different seed values, and the standard deviation is shown. See Appendix D for details on the baselines, their model parameters, and training regimes. See Appendix E for a more detailed description of the environments and setups. The source code of the experiments and algorithms is at [https://anonymous.4open.science/r/dhtm-FA0E](https://anonymous.4open.science/r/dhtm-FA0E).

### 4.1 Gridworld

Figure: Figure 2: POMDP 5x5 Gridworld foraging task using different types of memory: DHTM, CSCG, LSTM. CSCG and LSTM are trained in two modes. The first mode is with oracle, which tells the agent to reinitialize memory and clear the trajectory buffer during environmental change. The second mode is without oracle, so the agent continues using its current strategy to gather new trajectories and update memory. There is also an episodic control agent (ec) among baselines, which uses a simple dictionary to store trajectories and make predictions. A. Dynamic of steps required for the agent to reach the goal. In the 300th episode, the goal is blocked by a wall, so the agent needs to find a new optimal trajectory. B. Total steps taken during 1000 episodes.
Refer to caption: x1.png

The first test was conducted in partially observable Gridworld in a 5x5 alternating maze with a stationary end state and the agent’s starting position. The agent can observe the colour of the cell it is in and the colour of an obstacle it encounters, while being punished in addition to the constant step penalty. Since the observation can be represented directly as a categorical variable with the number of states equal to the number of possible colours, no encoder is needed. Each episode starts with the agent in the same position and ends when the agent enters the terminal state (i.e., when the reward is collected) or when the maximum number of action steps is reached. In each episode, we measure the number of action steps required for the agent to reach the goal.

Each trial lasts $1000$ episodes and consists of two phases. In the first phase, which lasts $300$ episodes, there are no obstacles on the shortest path to the terminal state. In the second phase, the previously optimal path is blocked by a wall, so the agent has to take a detour to get the reward. The experiment aims to show the ability of a memory to relearn the structure of the maze in the test time. If the agent is able to reach the end state in the second phase, it suggests that the model can forget, or at least suppress, the memory of the old shortest path when it is blocked. Otherwise, the agent will continue to hit the wall and fail to reach the terminal state.

In this experiment, we compare our memory model DHTM with CSCG [] and LSTM [] within the same agent architecture that uses memory to form successor features and evaluate actions (see Sec. 3.4). CSCG and LSTM are trained in two modes. The first mode is with an oracle that tells the agent to reinitialize memory and clear the trajectory buffer just in time for the environmental change. The second mode is without an oracle, so the agent continues to use its current strategy to collect new trajectories and update memory in the second phase. There is also a simple episodic control agent (EC) among the baselines, which uses a dictionary to store trajectories and make predictions (see Section 3.5 for details).

The results in Figure 2 show that CSCG effectively can’t adapt to the environmental change in the online regime, but only if it is retrained from scratch on new trajectories. And, apparently, LSTM is only able to learn a predictive model for a small horizon, since it performs significantly worse than other baselines in the second phase, even with oracle enabled. In contrast, DHTM and EC are able to adapt online without explicit memory reset. Another important observation is that the EC agent performs on par with DHTM, suggesting that its mechanism can explain most of DHTM’s performance in this task.

### 4.2 AnimalAI

Figure: Figure 3: Amount of action steps the agent performs to reach the goal in a 10x10-meter room in AnimalAI environment using DHTM, CSCG, LSTM and EC dictionary. The agent and the goal have the same initial positions throughout the episodes. Each time-step, the agent gets an encoded image of a first-person view of the environment and produces one of three actions: turn left, turn right, or go forward. A. Dynamic of steps required for the agent to reach the goal. B. Total steps taken during 1000 episodes by the agent with DHTM with K-Means encoder and Categorical VAE. For the latter, DHTM’s context factor size is varied (i.e., how many previous time-step hidden variables are considered for prediction).
Refer to caption: x3.png

AnimalAI is a 3D environment capable of providing rich visual input to an agent in the form of RGB images. The goal of the experiment described here is to show that DHTM can handle distributed representations and is scalable to environments with larger state spaces.

The agent’s task is to reach a sphere of radius \qty0.5, which is located in a square area of size $10$ by $10$ meters. The start position of the agent and the position of the ball are the same between episodes. In each episode, we measure the number of actions (steps), which are AnimalAI’s standard left/right turn and forward walk, the agent takes to reach the goal, with the maximum number of actions limited to $40$. Each action is repeated three times and three frames are skipped, so one agent action step corresponds to three steps in the environment.

We test the agent using two different encoders: K-Means and Categorical VAE []. Both encoders are pre-trained on RGB images collected with a uniform strategy during $2000$ episodes in the same room used for agent testing. K-Means represents an image as a single categorical variable with $50$ states, and Categorical VAE is configured for distributed representation with five variables with $50$ states each (see Appendix C for details).

The results are shown in Figure 3A. It can be seen that after about $150$ episodes, the agent with DHTM and the EC agent reach the goal stably, although the DHTM agent with distributed encoding (with a factor size of $4$) converges to a slightly shorter path on average. The agent with CSCG slowly improves its strategy, but the LSTM agent doesn’t show any learning on this time-step scale. To understand a common learning time-step scale for contemporary model-based algorithms, we also tested a small version of Dreamer V3 (18M parameters) [], which successfully solves the same task after about $35\text{\times}{10}^{3}$ action steps ($1000$ episodes) versus $3.5\text{\times}{10}^{3}$ steps ($150$ episodes) for DHTM (see Appendix F for details).

To show that DHTM can’t simply be replaced by multiple parallel EC dictionaries or CSCGs in the case of distributed encoding, we also vary the context factor size of DHTM, that is, how many previous time-step variables are taken into account to predict the current state of a variable. The factor size of 1 corresponds to the graph of the Factorial HMM, with parallel HMMs. It can be seen from Figure 3B that in our case of categorical VAE, apparently knowing only the previous state of the variable is not enough to reliably predict rewarding features, since the agent doesn’t converge to the shortest path. We also tested the variant (factor 1*) where each feature variable is represented by three hidden states, to ensure that the result for factor size $1$ is not due to DHTM’s hidden state collisions. DHTM with a factor that takes into account more than one previous time-step hidden variable (excluding the action variable) performs significantly better. For diagrams depicting factor graphs used in these experiments, see Appendix D.4.

## 5 Conclusion

This paper introduces a novel episodic memory model, DHTM. It combines ideas from belief propagation, episodic control, and neurophysiological models of the neocortex, resulting in a fast probabilistic sequence memory.

The experiments show that our memory model can quickly learn observation sequences in changing environments and effectively reuse them for RL tasks, on par with model-based table episodic control. We also show that, unlike CSCG and table EC, DHTM is able to handle distributed representations and account for statistical dependencies between variables of these representations. Potentially, it extends the application of our algorithm to higher state spaces, exploiting their factorizations for more efficient computations and to reduce the number of stored parameters.

Our results suggest that DHTM can be viewed as an adaptive trajectory buffer for a slower generalizing algorithm, such as a Neural State Space Model-based []. Through prediction error learning, DHTM can selectively discard and replace obsolete transitions in observed trajectories without disrupting relevant experience. It can also provide an initial policy for an agent through episodic control until a slower but more accurate world model takes over. These features of a trajectory buffer are highly desirable in the case of online learning in non-stationary environments.

Further improvements of DHTM can be directed towards the implementation of an automatic feature space factorization discovery to determine an optimal context factor size during learning. To avoid excessive segment growth in noisy environments, it’s also possible to implement an adaptive segment activation threshold. One of the main limitations of this study is that the evaluation of the algorithms is limited to basic RL setups due to the simple agent architecture. More advanced use cases could be demonstrated using better exploration strategies and in combination with other memory models, which is also a matter of future research.

## Appendix A Value Function Decomposition

In our agent model, we approximate the reward function R(s) as a sum:

$$ $R_{t}=\frac{1}{n}\sum^{n}_{k=1}{r(\varphi^{k}_{t})},$ (13) $$

where $r(\varphi^{k}_{t})$ is a reward associated with state $\varphi^{k}_{t}$, $n$–number of feature variables. Then, similarly to the Successor Representation idea (see Section 2.2), the value function can be represented as:

$$ $\displaystyle V(h_{t})$ $\displaystyle=\mathrm{E}[\sum^{\infty}_{l=0}\gamma^{l}R_{t+l+1}\ |\ h_{t}]= \sum^{\infty}_{l=0}\gamma^{l}\mathrm{E}[\frac{1}{n}\sum^{n}_{k=1}{r(\varphi^{k }_{t})}\ |\ h_{t}]$ $\displaystyle=\frac{1}{n}\sum^{n}_{k=1}\sum_{j}\sum^{\infty}_{l=0}\gamma^{l}p( \varphi^{k}_{t+l+1}=j\ |\ h_{t})r^{k}_{j}$ $\displaystyle=\frac{1}{n}\sum^{n}_{k=1}\sum_{j}\mathrm{SF}^{k}_{j}(h_{t})r^{k} _{j},$ (14) $$

where $\mathrm{SF}^{k}_{j}(h_{t})=\sum^{\infty}_{l=0}\gamma^{l}p(\varphi^{k}_{t+l+1}=
j\ |\ h_{t})$, $h_{t}=(h^{1}_{t},...,h^{n}_{t})$—hidden state vector of variables $\{H^{k}_{t}\}_{k}$.

## Appendix B Biological Interpretation

Neural implementation of the DHTM is inspired by neocortical neural networks (see Fig. 4). Hidden variables $H^{k}$ may be considered as populations of excitatory pyramidal neurons in cortical layer L2/3 of somatosensory areas, with lateral inhibition modelled as $\mathrm{softmax}$ function. Staiger and Petersen [] showed that neurons in this layer are responsible for temporal context formation.

The neuronal activity at timestep $t$ can be thought to carry messages $M^{k}_{t-1}$. Messages are propagated through synapses of dendritic segments, which correspond to factors $F^{k}_{c}$. Dendritic segments of biological neurons are known to be coincidence detectors of their synaptic input []. We use the notion of dendritic segment to sparsely represent context factors $F^{k}_{c}$, as each factor value corresponds to a particular combination of states (or active cells).

Feature variables $\Phi^{k}_{t}$ may be considered to represent cells of a granular layer (L4), as they are known to be the main hub for sensory excitation for L2/3. L2/3 cells that have common sensory input from the layer L4 are modelled as columns for particular feature states $\mathrm{col(\varphi^{k}_{t})}$ [].

Figure: Figure 4: Biological view of the neural implementation of the DHTM. Variables $H^{\cdot}_{t-1}$ correspond to populations of neurons that have common sensory input and lateral inhibitory competition. Dendritic segments correspond to factor values $f_{l}$. Spike frequency of a neuron reflects state probability $p(h^{k}_{t})$ of a variable.
Refer to caption: extracted/5907925/bio_intr.png

## Appendix C Encoding Observations

Since DHTM is designed to process categorical random variables, we need a categorical encoder for RGB images to solve AnimalAI tasks.

To compare our model to CSCG and EC in AnimalAI, we use the approach described in Guntupalli et al. [] to quantize observation space. K-Means (implementation from scikit-learn) with $50$ clusters pretrained on $4\text{\times}{10}^{4}$ 64x64 RGB observations from trajectories sampled by uniform policy. To encode an image, the distance from each cluster is calculated and the index of the closest cluster is returned. In that way, AnimalAI’s observation space is represented by one categorical variable.

In order to test DHTM’s ability to deal with distributed encoding, we also trained Categorical VAE [] with the latent space represented by five categorical variables with $50$ states each. We use a slightly modified implementation from Subramanian []. The only difference is that we modify the loss function by adding mean squared reward prediction error, which is linearly calculated from the latent codes to ensure the linear decomposition in feature space (equation 13). The model is trained on $6\text{\times}{10}^{4}$ 64x64 RGB images for $20$ epochs with batch size $64$. An image is encoded by first passing it through VAE’s encoder block, then, to ensure stable encoding, instead of sampling from the latent distribution, the state with the highest logit is chosen for each variable.

## Appendix D Implementation details

As mentioned in Section 3.4, we incorporate DHTM as a part of an RL agent, which has a memory model and a feature reward function. All tested memory models share the same pipeline—conditioned on actions, they learn sequences of categorically encoded observations and are used to form SFs.

In all experiments, the agent uses softmax exploration with temperature equal to $0.04$, except for EC agent we set softmax temperature to $0.01$, since for higher temperatures its policy is unstable. For SF formation the agent uses $\gamma=0.8$ and $\gamma=0.9$ for Gridworld and AnimalAI experiments respectively. Maximum SF horizon is set to $50$, but it is usually shorter since we employ early stop condition based on feature variables distribution. That is, if the prediction is close to uniform or the probability of a state with associated positive reward meets a threshold, the planning is interrupted. We use probability thresholds equal to $0.05$ and $0.01$ for Gridworld and AnimalAI respectively.

### D.1 LSTM

LSTM baseline was implemented with a single LSTMCell from PyTorch library []. It is supported by an additional symexp-layer to encode input before passing it to the LSTM cell and a symexp-layer to decode the LSTM cell’s output from the LSTM’s hidden state back to the input representation, where symexp activation function, $\mathrm{symexp}(x)=\mathrm{sign}(x)e^{\lvert x\rvert-1}$, is a reverse of symlog function: $\mathrm{symlog}=\mathrm{sign}(x)\log{(\lvert x\rvert+1})$.

For both Gridworld and AnimalAI tasks, LSTM with a hidden state size of 100 is trained on a buffer of $1000$ trajectories with learning rate $0.002$ on $50$ randomly sampled batches of size $50$ every $10$ episodes.

We also incorporated some notion of random variables in hidden states by splitting the hidden state of the tested LSTM into groups. In all experiments the hidden state represents $10$ categorical variables with $10$ states. That is, LSTM is forced to learn $10$ categorical distributions with multi-cross-entropy loss to explain the observed sequences, which is a somewhat close to the multi-categorical hidden state representation used in DreamerV2/V3 []. The idea of using symexp activation function, mentioned above, is inspired by Dreamer too, and is used to remedy the problem of learning extreme probability values. Without symexp the neural network has to represent zero probability with high negative logit values and one-probability with high positive logit values, which is hard to reach with low learning rate and may lead to instabilities. Thus, symexp function makes it faster to reach target values in log space.

### D.2 CSCG

CSCG baseline was implemented using code from the repository accompanying the paper ([https://github.com/vicariousinc/naturecomm_cscg](https://github.com/vicariousinc/naturecomm_cscg)). CSCG was trained on buffers of $500$ action steps with $1000$ EM steps until convergence. We iteratively calculated the exponential moving average of transition matrices obtained for different batches with a smoothing coefficient $\alpha=0.05$ as described in the paper []. This smoothed transition matrix was used as an initialization for the next training batch and for inference. The first batch of observations is gathered using the uniform policy, and later batches are gathered using the current best policy.

In so-called oracle mode, the trajectory buffer is continuously growing, and the transition matrix is trained from scratch every $500$ steps. The buffer is discarded at the environmental change and the transition matrix is reinitialised.

### D.3 Episodic Control

Here, we provide an algorithm for SF formation using EC dictionary $\mathcal{D}_{a}$ (see Algorithm 3). In contrast to other memory algorithms, EC dictionary doesn’t estimate state transition probabilities, but it stores connections between them in all-or-nothing faction. Therefore, on each time step, all predicted observations are assumed equally probable. But it’s not a limitation in our set of experiments, since we tested our algorithms in deterministic environments.

Figure: Algorithm 3 SF formation for EC

### D.4 DHTM

In the Gridworld task and AnimalAI with K-Means encoder, the factor-graph for DHTM consists of three hidden variables $H$ with $40$ states per one feature state. Each hidden variable gets the same copied observation, and their predictions are averaged. Using several hidden variables here is essential, as it significantly increases the amount of trajectories memory is able to store by decreasing collision probability for hidden state representations of different trajectories.

In AnimalAI experiments, hidden variables state spaces are of size $40$ states per one feature state. For Categorical VAE encoder, we use one hidden variable per feature variable, except the case (factor 1*) that tests that information from one feature variable is not sufficient for correct prediction, in which three hidden variables per observation variable are used and the factor of size three connecting to only hidden copies of the same feature variable. Examples of typical DHTM’s graphical models used in our experiments are presented on Figure 5. It’s important to note that edge colours denote separate graphical models, but the variables are shared (copied) between them. So there are no loops in these models. For simplicity, we define context factor size as the number of hidden variables from the previous time-step it has connections to, but it also connects to one hidden variable of the current step and the action variable.

The most important DHTM’s parameters, except the factor size, are factor learning rate $\alpha$ and its initial value $f_{0}$. For better adaptability in changing and stochastic environments, the initial factor value should be set closer to zero, but in stationary environments, it can be safely set $f_{0}=1$, allowing faster convergence. In Gridworld we use $\alpha=0.1$ and $f_{0}=0.05$, and in AnimalAI we set $\alpha=0.001$ and $f_{0}=1$. We also set the limit on the number of segments $5\text{\times}{10}^{4}$, but practically never reach it in our experiments.

Figure: Figure 5: Examples of factor-graphs for two main setups. Edge colours show that edges correspond to separate graphical models with copied variables. All context factors are also connected to action variable $A_{t-1}$. A. Graphical model for one feature variable. In that case, feature value is copied for three hidden variables, to increase hidden state space and avoid memory collisions. B. Graphical model for Categorical VAE distributed encoding, example for factor of size $3$ (i.e. takes into account tree hidden variables from the previous time-step).
Refer to caption: extracted/5907925/graph_kmeans.png

## Appendix E Experimental Setups

### E.1 Gridworld

We use a simple homebrew Gridworld implementation. Each Gridworld position maps to a colour, reward, and terminal state indicator. There are two special colours for the world border and obstacles. Each time-step, an agent gets the colour index of the current agent’s position. An agent is able to see obstacles or border colour when trying to go through them, which just replace current position colour for one step without actually changing the agent’s position. We also punish an agent trying to go through a wall or world border in addition to base step reward punishment to facilitate faster convergence to optimal trajectories.

Samples of setups used in our experiments are depicted in Figure 6. Here, the red colour corresponds to obstacles and the light gray colour to the reward. Floor colours are resampled for each trial, but we ensure that there are no more than five floor colours, including the terminal state. In each episode, the agent and the terminal rewarding state have the same initial positions. The episode ends when the terminal state is entered or the limit of action steps is reached.

Figure: Figure 6: Schematic view of POMDP Gridworld setups. The left picture shows the initial setup of the environment, with the reward delivered by the light gray square. The right picture represents the environment after the reward is blocked by obstacles (red squares). Agent’s initial position in the lower left corner.
Refer to caption: extracted/5907925/free.png

### E.2 AnimalAI

AnimalAI is a testbed inspired by experiments with animals []. The environment consists of 3D area surrounded by a wall and many different objects that can be placed using a configuration file, including walls, food, ramps, trees, movable obstacles, and more. We tested our agent in a 10x10-meter room surrounded by walls (see Fig. 7). Each timestep, the agent gets an RGB image of a first-person view of the environment and is able to influence its behaviour by choosing one of three actions: turn left, turn right, or go forward. The agent and food item have the same fixed initial positions for each episode. The episode ends if the food item is collected or the maximum episode time is reached. To make an agent’s actions more discretized, we repeat its choice three times and skip three frames.

Figure: Figure 7: Sample of the agent’s first-person view of our setup in AnimalAI environment. The green ball represents a goal, a food item to be collected.
Refer to caption: extracted/5907925/animalai.png

## Appendix F Additional Experiments

In order to ground our experiments to state-of-the-art model based algorithms, as an additional baseline, we test the smallest version of Dreamer V3 with 18 million parameters. We use the implementation described in Voudouris et al. [], which is based on Hafner et al. [] and adapted for AnimalAI.

We tested it with default parameters in the same setup as described in Section 4.2, except we don’t apply frame skip. The agent successfully solves the task after $105\text{\times}{10}^{3}$ action steps. Taking into account the frame skip used in our experiments with DHTM, it’s approximately $35\text{\times}{10}^{3}$ action steps versus $3.5\text{\times}{10}^{3}$ for the DHTM agent to solve the task.

Figure: Figure 8: Dreamer V3 (18M) score in AnimalAI experiment described in this paper.
Refer to caption: x5.png

## Appendix G Glossary

Categorical Random Variable—a discrete random variable that can take on of finite $K$ possible states.
Cortical Column or Minicolumn—a population of neurons in the neocortex that spans across layers and shares sensory input.
Dendritic segment—a group of synapses (neuron’s connections) that acts as an independent computational unit affecting the resulting neuron’s activity.
Factor Graph—bipartite graph representing the factorization of a probability distribution, with one part representing factor nodes and another—random variables.
Factor size—number of variables connected to the factor. For simplicity, we refer only to the number of previous hidden states, used for prediction.
Hidden Markov Model (HMM)—statistical model of a stochastic process where state probability depends only on previous state of the process.
Multi-compartment neuron model—a model of neuron that divides neuron’s connections into groups (segments) of different types (compartments), where each group may be considered as partly independent computational unit and groups of each compartment may affect the neuron’s activity differently.
Oracle—a mythical creature that gives up to our algorithms insightful information about the environment that is usually hidden.
Successor Representations (SR)—a discounted sum of future [one-hot encoded] observations.
Successor Features (SF)—a generalization of SR, a discounted sum of future latent states.
Temporal Memory (TM)—in this work by this term we mean “memory for sequences”.

## Peer Reviews

confidence: 4
contribution: 4
initial_rating: 6
presentation: 2
questions: Any of the weakness are also questions.
soundness: 3
strengths: This paper is an very interesting contribution. The connection of successor features, factor graphs, episodic control and complex neuron morphologies is very attractive. In particular, it offers a practical computational interpretation of the role of complex morphologies in biological neural network: computing factors/features based on previous experiences. The paper covers a lot of theoretical ground in connecting the different ideas used to build the model. In general each aspect of the model is theoretically sound and the experiments are tested with a good set of baselines. I will try to be comprehensive in the weaknesses and questions with the intentions of making the contribution stronger.
summary: This paper proposes a novel algorithm for learning successor features for reinforcement learning agents in non-stationary partially observable environments. The main feature of the algorithm is a new episodic memory architecture that combines ideas from Factor graphs, Hidden Markov Models and episodic control, along with a neuron concept similar to the Hierarchical Temporal Memory, which is in itself inspired by biological neurons as complex units with spatio-temporal computation. Features of this new type of memory include, efficiency of computation, distributed representations and a local (Hebbian) learning rule. The algorithm is embedded in a simple RL agent and tested, both, in a simple gridworld environment and in a more structure AnimalRL environment.
weaknesses: 1. The basis of comparison with the other baselines in experiment number 1, is the capacity of the model to forget, however, in algorithm 1, the DHTM agent includes a memory_reset() step. This seems to invalidate the whole argument of this experiment.
2. The connection with biology is only superficial, but I think it can be made stronger. It is not clear if some of the things I am going to mention are already in the authors minds, but they are definitely not clear in the text:
- The activity of each of the neurons in the model is associated with a spike only makes sense if there are two different timescales at play, so that you can collect many spikes before making a decision (this is what firing rate means). Do the neurons only spike once per episode/state visited?
- The appendix does not explain very clearly the connection of your model with HTM is not very well explained (and the citation of Hawkins paper is not the appropriate paper, choose one in which they explain the model). For example, it is never clear what you mean by the receptive field of a cell.
- Explain better how the computation of Fe relates to the columnar structure of the cortex? The authors seem to be confusing the morphology of one neuron with the structure of the cortex (this is related to the connection with HTM above). One is related with the sort of computations performed, the other with the origin and information carried by the afferents.

3. The presentation style is probably the weakest point of this paper. It is very difficult to follow. At the start, the definition of a POMDP is very difficult to read, adding a bit more words can make this and other definitions easier to read (if they are in the paper, they will be read by someone, otherwise they can go in the appendix).
- Many things are explained in the wrong place, before or after they appear in any equation. For example rec(l) for the receptive field, or the definition of log Li.
- The calculation of the emission factors E is never shown explicitly! but it appears in equation 7
- Equation 7 is not well justify at all, it is very difficult to see how you go from (4) to (7)
- Many times it is not clear what is part of your algorithm and what is you explaining the old way of doing things, for example: line 191- "the feature variables are considered independent (in this work?), however, the are interdependent (so which?).
- Similar to the previous one, in the exposition of algorithm 2, it seems that you are sometimes talking about DHTM and sometimes about the naive EC agent.

4. The scope of the paper seems a bit misleading. If the focus is the learning of successor representations for non-stationary environments, why not testing it in more standard and more complicated environments than the ones shown in this paper? The experiments are barely non-stationary and not the usual testing ground for comparison with the baselines chosen. As I said, I think this is a very interesting idea but I think the scope needs to be made more realistic or add the necessary tests. This makes me think that there is some problem with scaling the presented ideas; if that is the case what is it?
- Do the number of segments grow fast or exponentially with the complexity of the environment (please show evidence otherwise). The factors are defined in terms of combinations of ALL the RVs involved!
- Is the computation as efficient as suggested? Is having segments more efficient as claimed? How fast? (the results are shown in action steps but no mention about memory/time efficiency)
- What is the capacity of the memory?

5. Definitely needs more figures! figure 1 is infinitely complex, it has a lot of information (more than 3 sections worth of information). This figure can be split to illustrate different aspects of the exposition. For example, the details of the computation of LogLi can be added in a later figure. In relation to the figures, figure 6 also needs more information. Mark what the goal and the start position is (the colors have barely any contrast and someone with visual impairments could not get any information from this figure)

6. Others
- It is not clear to me how are new segments created
- the change from Fe to fl for the emission factors was confusing, please add some warning
- The calculation of the conditional probability on the features in 294 is only tangentially mentioned, but it is very important, here is the only time the features are introduced, and it deserves some discussion
- How is the planning horizon "T" determined. How are you measuring the shape of the distributions to determine this parameter?
- The explanation of the Episodic control agent needs to be rewritten for clarity.
- In equation 10, make clear that Rt is the reward coming from the environment (there are other definitions for Rt before)
- It took me some time to understand that each neuron is a state, and each neuron is also a factor. It is not clear how the population of neurons is set up. It seems that we would need to keep adding neurons with longer experiences? Again, it will be very helpful if the architecture and computations are separated from figure 1.
- 291,292, please explain this step.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - Can you talk about how to incorporate mechanisms that enable the model to generalize across different sequences leading to the same state. Would hierarchical or abstract representations be useful here?
- It would be nice to see a experiment that shows a substantial improvement of the DHTM model over Episodic Control, or an explanation of why it's better to implement the full algorithm. 
- In the AnimalAI task, you only test DHTM with VAE whereas all of the other algorithms are only tests with k-means. Is there a reason why you can only test DHTM with VAE and not the others? If it's possible to test the others with VAE, I would be interested to see that in a comparison.
- It would be worthwhile to see the performance of Modern RNN architectures (such as LRU or RKWV mentioned in the beginning of the paper) instead of LSTM, unless there is a reason those architectures don't apply to these tasks. These algorithms also avoid using BPTT and are quite performant.
- Can you comment on how the architecture compares to other Temporal Memory architectures that are based on Hebbian Learning Rules such as Sequential Hopfield Networks? For example, how does it relate to Chaudhry et al. 2024 (Long Sequence Hopfield Memory) or Tang et al. 2024 (Sequential memory with temporal predictive coding). I think you should add a more thorough literature review of temporal memory.
- I would like to see the effect of an oracle on Episodic Control and DHTM in Figure 4.
soundness: 3
strengths: - DHTM operates fully online using local learning rules due to its local Hebbian learning rules, making it well-suited for non-stationary environments without needing explicit resets or retraining. Its ability to dynamically form and update memory representations allows it to cope with sudden changes in the environment effectively. Furthermore, you can train it fully online since it doesn't require backpropagation through time. 
- The algorithm is computationally efficient due to sparse transition matrices and distributed representations, and could be potentially implemented efficiently in neuromorphic hardware.
- The experiments indicate that DHTM could achieve higher sample efficiency compared to some state-of-the-art models. For example, in the AnimalAI environment, DHTM requires significantly fewer interactions with the environment to learn effective policies than models like Dreamer V3.
summary: This paper introduces a novel algorithm for online temporal memory learning in decision-making tasks with non-stationary / partially-observable environments. Distributed Hebbian Temporal Memory is a biologically-inspired algorithm that uses Hebbian learning rules in order to learn Successor Features, a powerful framework in reinforcement learning to decouple the environment's dynamics from the reward function.
weaknesses: - The unique encoding of each observation sequence limits the model's ability to generalize across different sequences which represent the same underlying state.
- The performance of the algorithm is very close to that of Episodic Control on some tasks, which leads to a question of its advantage over a similar method. 
- The baseline models like LSTM may not be fully optimized for online learning, potentially skewing the comparison results.
- The paper lacks a thorough theoretical analysis of the algorithm's properties such as convergence guarantees and computational complexity.
- I think you should have more exposition explaining what CSCG, Episodic Control, and Dreamer V3 are.
- I am willing to improve my score if you address some of the questions posed in the following section.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - How does the DHTM scale? The experiments applied the model to 5x5 and 10x10 environments. What if the environment had a larger number of states and a larger number of actions? How does this affect the computational complexity and loss curve?
- How well does DHTM handle partial observability? Can it distinguish aliased observations well?
- CSCG has been shown to learn the spatial structure of the environment which allows for generalization during navigation. Can DHTM do the same?
soundness: 3
strengths: - The model proposed by the authors is novel and formulated using the factor graph framework, making it theoretically-grounded, allowing for techniques in that area to be easily applied.
- The model is online and uses local learning rules, which makes it more applicable.
- The model is biologically inspired and the authors make specific connections to its neural implementation/plausibility.
- The authors provide an interpretation of the model in terms of episodic control.
summary: The paper introduces Distributed Hebbian Temporal Memory (DHTM), a model inspired by the HTM neuron (modeled after cortical columns) for spatial-temporal learning. The model takes in a sequence of observations generated from a trajectory through an environment and generates a hidden state representation. These are then used to predict the cumulative distribution of future states under a uniform policy to form a Successor Feature (SF) representation. This is then subsequently used by an RL agent to estimate the Q value (assuming the reward can be decomposed linearly) for action selection.

The authors perform experiments on GridWorld and AnimalAI environments and compare DHTM to LSTM and CSCG baselines, showing that it outperforms both, reaching the goal state in a few number of states in fewer episodes.
weaknesses: - The description of DHTM is difficult to follow. I suggest adding paragraph-level divisions to clarify the logical structure of the sections, or at least outline the general structure of the section at the beginning.
- Since the model uses a factor graph formalism, I think having a section for that in the background would improve clarity.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: * To my knowledge the term "temporal memory" is not in common use and should be briefly defined when it is used in the abstract.
* The distinction between requiring "complete data sequences to be available during training" (lines 53-54) and "access to only one sequence data point at a time" (lines 61-62) is not clear to me.
* "Factor graph formalism" should be defined in the introduction.
* "local Hebbian-like learning rules ... make the learning process much faster than gradient methods" (lines 70-71) should be revised to clarify in which sense "faster" is meant.
* Eq. (3) is unclear because of the summation over $\Omega_k$. I believe this is meant to express summation over all possible values of the "previous time step RV indexes [sic] included in the $F_c^k$ factor" which should be clarified.
* What is meant by "$f_l$ is proportional to several coincidences $s_l = c_t^j = 1$ in the recent past" (line )? 
* Isn't the VAE encoder in section 4.2 significantly more expressive than the k-means encoder? Doesn't this give DHTM-vae an advantage as the only model with this representation?
* It is claimed that the paper shows that "unlike CSCG and table EC, DHTM is able to handle distributed representations and account for statistical dependencies between variables of these representations". I'm not sure where this was shown or how.
* It may be beneficial to revise the storytelling in the introduction to emphasize the points raised in the conclusion that this is a weak but very fast form of learning which may complement a slower, more accurate method over the course of learning.
soundness: 3
strengths: * Exploring how bio-inspired neural networks can implement and perform inference on statistical models is an interesting opportunity for interplay between AI and neuroscience.
* The experimental evaluation of the method is reasonable and with some more detail would support its effectiveness.
* The discussion of limitations of the method in the conclusion is thoughtful and thorough.
summary: This paper proposes a new bio-inspired algorithm, Distributed Hebbian Temporal Memory (DHTM), for online learning to make decisions under uncertainty. This algorithm resembles a model of episodic memory, or a dynamic trajectory buffer, which remembers and replays sequences of actions which were rewarding in the past. It is implemented by compartmentalized neurons, with dendrites that recognize particular previous states and encourage recall of associated actions, and learning is realized by a Hebbian-like rule on these dendrites. The method is evaluated on two navigation tasks (GridWorld and AnimalAI) and compared against a few baselines -- LSTM, CSCG, and a simplified episodic controller which it resembles in performance.
weaknesses: * The method is somewhat misrepresented in the introduction. What exactly it's meant to accomplish is not clear, and it's compared alongside very general learning methods which are meant to generalize effectively. After reading the entire paper, its focus seems to be on learning simple, non-generalizable strategies quickly.
* The description of the model is difficult to follow. A huge amount of notation is introduced, and much of it is not defined precisely and has to be inferred from context. This is most problematic in section 3.2, where the bio-inspired implementation is described. Ideally the neuron model would be laid out, and then the variables of the neuron model would be associated with variables from DHTM.
* Key details of the experiments, particularly how the baseline models were tuned, are missing and make it difficult to evaluate them.
* The paper presents a simple episodic control-based model which is highly competitive with the much more complicated DHTM which it focuses on. It is unclear what the advantages of DHTM are compared to this simple model.
* There is no substantial discussion of related work.

Meta-review: This paper introduces a novel, biologically-inspired algorithm (DHTM) designed for online sequence learning for decision making in non-stationary, partially observable environments. Leveraging factor graph formalism and Hebbian-like local learning rules, DHTM captures Successor Features (SFs) to facilitate decision-making without the need for backpropagation through time. The authors demonstrate the algorithm’s efficacy in tasks requiring episodic memory, such as GridWorld and AnimalAI, showing that it outperforms LSTM, RWKV, and CSCG baselines in terms of sample efficiency and adaptability. Three of the four reviewers voted for acceptance.

Decision: Accept (Poster)