Learning Successor Features with Distributed Hebbian Temporal Memory
Authors: Evgenii Aleksandrovich Dzhivelikian
Petr Kuderov
Aleksandr Panov
Venue: ICLR 2025 Poster (2025)
Keywords: temporal memory
successor features
online learning
Hebbian learning
Primary area: learning on time series and dynamical systems
TL;DR: We present a new temporal memory algorithm with online learning for decision-making in changing, partially observable environments.

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

Published as a conference paper at ICLR 2025
LEARNING SUCCESSOR FEATURES WITH DISTRIBUTED HEBBIAN TEMPORAL MEMORY
Evgenii Dzhivelikian
Petr Kuderov AIRI, MIPT Moscow, Russia
Aleksandr Panov
MIPT, AIRI
AIRI, MIPT
Moscow, Russia
Moscow, Russia panoveairi.c net
dzhi vel ikian. wewehhstthhhhh..m.. .du
kuderov@airi.nee
ABSTRACT
This paper presents a novel approach to address the challenge of online sequence learning for decision making under uncertainty non-stationary, partially ob- servable environments. The proposed algorithm, Distributed Hebbian Temporal Memory (DHTM), is based on the factor graph formalism and a multi-component model. DHTM aims to capture sequential data relationships and make mulative predictions about future observations, forming Successor Features (SFs). Inspired by neurophysiological models of the neocortex, the algorithm uses dis- tributed representations, sparse transition matrices, and local Hebbian-like learn- ing rules to overcome the instability and slow learning of traditional temporal memory algorithms such as RNN and HMM. Experimental results show that DHTM outperforms LSTM, RWKV and a biologically inspired HMM-like al- gorithm, CSCG, on non-stationary data sets. Our results suggest that DHTM is a promising approach address the challenges of online sequence learning and planning in dynamic environments.
INTRODUCTION
Modelling sequential data is one of the essential tasks in artificial intelligence as it has many ap- plications, including decision- making and learning world models (Moerland et al.] (023). natural language processing (Min et al.) (021), conversational Al (Dwivedi et al. 2023), time series anal- ysis (Hraslan al., 2019). and video and music generation (Ji et al. 2020). 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 this infor- mation through the dimension of time (Lipton et al. 2015 Harshvardhan et al. 2020 Mathys al.) 2011]. We refer to the class of algorithms that use this approach as temporal memory (TM) algo- rithms because they essentially model the cognitive ability of complex living organisms remember the past experience and make future predictions based on that memory (Hochreiter & Schmidhuber)
1997 Friston al. 2016 2018 Parr & Friston 2017).
Temporal memories differ in their ability to capture data dependencies and the extent to which they can generalise their predictions to unseen sequences. There are two extremes known in psychology as cognitive maps] Tolman 1948), which extract commonalities or schemas generalising experience, and episodic memory Tulving al. (1722), which stores fragments of experience separately. They both can be viewed as temporal memories, however, the former is slower and require lot of data to built up useful schema, according to contemporary computational models based on HMM (George et al. 202T) and Transformers (Dedieu et al. 2024), while the latter is instant and is usually viewed as data buffer (Zhang et al. 2021).
In this paper, we explore advantages of an episodic-like memory in addressing the problem of online sequence learning for planning and decision-makie under uncertainty, which can be formalized as reinforcement learning (RL) for partially observable Markov decision process (POMDP) (Poupart 2003] Singh et al. 2021). Inferring the hidden state in partially observable environment is a se- quence modelling problem, since it requires processing 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

---

Published as a conference paper at ICLR 2025
given by the reward function (Dayan 1993, SamuelJ. J. Gershman 2018 Millidge & Buckley 2022
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 (Barreto al. 2017 Touati et al. 2022). Temporal memory (TM) algorithms can be used to make cumulative predictions about future states and their features, forming SR SF. This work shows that the proposed algo- rithm, Distributed Hebbian Temporal Memory (DHTM), can dynamically form SFs for navigation tasks in Gridworld and AnimalAI (Fountas et al. 2020) environments.
The most prominent TM algorithms, such as Recurrent Neural Network (RNN) (Boden 2002) or LSTM (Hochreiter Schmidhuber, 1997), use backpropagation to capture data relationships. which is known to be unstable due to recurrent nonlinear derivatives. They also require complete data se- quences to be available during training. Although the gradient vanishing problem can be partially circumvented in a way that Receptance Weighted Key Value (RWKV) (Peng et al. 2023) or Linear Recurrent Unit (LRU) (Orvieto al., (023) models do, the problem of online learning. i.e. learning on non-stationary data sets, is still viable issue. Unlike graphical models such as HMM (Eddy) 2004). RNN models and their descendants also lack a probabilistic theoretical foundation, which is advantageous for modelling sequences captured from stochastic environments (Salaün et al. 2019] Zhao et al. 2020). There is little research on TM models that can be used in fully online adap- tive systems interacting with partially observable stochastic environments with access to only one sequence data point at time, a common case in reinforcement learning (Jahromi et al. (022))
We propose a Distributed Hebbian Temporal Memory (DHTM) algorithm based on factor graph for- malism and a nulti-compartment neuron model inspired by the HTM neuron (Hawkins & Ahmad, 2016). Factor graph a flexible graphical representation of probabilistic distributions, which is com- monly used in combination with sum-product algorithm (Loeliger] (004). We found this represen- tation particularly convenient in our attempt of probabilistic interpretation of the HTM framework. The resulting graphical structure of our model is similar to that of the factorial HMM (Ghahramani & Jordan 1995). An important feature of our model 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 (Hebb) 2005| Churchland & Sejnowski 1992] Lillicrap et al. 2020), which avoid the drawbacks of gradient methods and make the learning pro- cess much more sample efficient than gradient methods. This comes at the expense of hidden state generalization, since DHTM's '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 (EC) (Lengyel& & Dayan 2007 Pritzel et al. 2017 Ritter et al. 2018, Emukkere et al. 2023).
Though our approach is similar to EC, we are not aware of other works, demonstrating its application to partially observable environments. There are other temporal memories that, similarly to DHTM, use Hebbian rules and learn online, like Asymmetric Hopfield Network (AHN) (Chaudhry et al.) 2023), Temporal Predictive Coding (IPC) (Tang et al., (023) and Hierarchical Temporal Memory (Hawkins & Ahmad, 2016). Compared to DHTM, AHN is a first-order memory. That is, it is not able to disentangle identical observations that correspond to different underlying environmental states. Therefore, it is not suitable for partially observable environments. In contrast, iPC has hidden layer that, potentially, makes it high-order memory. Despite that it has been tested only in MDP environments (Tang al.) (024) far, this method paves promising direction, which aim to compare with in future work. Finally, the algorithm that inspired our work. HTM, is also Hebbian high-order temporal memory, which, opposed to DHTM, doesn't 't allow for probabilistic predictions due to its all-or-nothing computations.
The proposed TM is tested as an episodic memory for an RL agent architecture navigating in a Gridworld environment and a more challenging AnimalAI testbed (Crosby et al. 2020). The results demonstrate that our algorithm outperforms classical LSTM, RWKV and biologically inspired HMM-like world model CSCG (George et al., 202T) in changing environments and is in order of magnitude more sample efficient than Dreamer V3 (Hafner et al.) 2023). 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:
2

---

Published as a conference paper at ICLR 2025
A novel episodic memory model, called DHTM, which is based on 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 speed comparable to model-based table episodic trol while supporting distributed representations. This capability significantl broadens its applicability.
BACKGROUND
2.1 REINFORCEMENT LEARNING
This paper considers decision-making in partially observable environment (Poupart 2005). The environment is defined by its state transition function P(s, a, s') = Pr(s') s, a) determining its dynamics, where s' E S is called state space, A-action space. Partial observability means that an agent has access to observations E O, which only partially inform about the environment's state each moment.
In reinforcement learning (RL.,, the goal for the agent set through the reward function R(s) : S R. The task of RL is find the agent's policy (a s) 5 x A [0, 1] that maximizes the expected return EI Re]. where an episode length and discount factor y sets the importance of remote rewards. Value-based methods usually aim estimate the Q-function given policy m: Q (sc. a²) = E B| 7 R+((())+) 8t, a,, æ... For an optimal value function Q", an optimal policy can be defined as (a 8) argmax Q" (s,
One of the drawbacks of value-based methods is that they require learning 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 (Dayan 1993). For discrete state space, such representa- tion can be derived just by rearranging sums the equation for value function:
V(h, = F
Res1+1 h = =]= = Σ SR
t=0
where y is a discount factor, i-th row in the matrix SR is the Successor Representation of state i, h i). Rg is reward for observing the state j.
and SR,y Ei= p'p(e+444
extension of SR for large or continuous state spaces that we use is Successor Features (Barreto et al. 2017), which is based 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 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 s adaptability to use world model to approximate the environ- ment's transition function and use it directly to form SR according its definition. And one of the most basic such world models is the Hidden Markov Model.
2.3 HIDDEN MARKOV MODEL.
A partially observable environment can be approximated by Hidden Markov Model (HMM) with hidden state space H and observation space O. Generally, H is not equal environemnts state
3

---

Published as a conference paper at ICLR 2025
space S since the environemnt's true structure is usually unknown. For simplicity, assume that actions are fully observable and that information about them is contained in hidden state. For the process of length T with state values 1 (hi..... hr)) and observations 01.7 (O1,...,,7) 0r), the Markov property yields the following factorization of the generative model:
T
T
plour, = p(h1) P(hhaa II [p(o))))
1=2
In the case of a discrete hidden state variable, a time-independent stochastic transition matrix can be learned with the Baum-Welch algorithm (Baum et al.] 1970), a variant of the expectation maximiza- tion (EM) algorithm. To compute the statisties for the expectation step, uses the greward-backwan algorithm, which is a special case of the cum-product algorithm (Kschischang ct al. 2001).
Factor graph is convenient graphical representation (Loeliger 2004) of probabilistic models like HMM (see F for an example). Formally, factor graph is bipartite graph, where one part con- sists of factors (squares), which, replacing conditional probabilities, define dependencies between random variables (circles) constituting the other part. Such representation also allows extension of forward-backe backward algorithm to sum-product algorithm. As shown in section 3.1 we base DHTM on ideas from the sum-product algorithm and the Factorial HMM, while adapting them for online learning.
DISTRIBUTED HEBBIAN TEMPORAL MEMORY
3.1 FACTOR GRAPH MODEL.
Distributed Hebbian Temporal Memory is based the sum-product belief propagation algorithm in a factor graph (see Figure TA). Analogous to Factorial HMM (Ghahramani & Jordan 1997). we divide the hidden space H into subspaces H There are four sets of random variables (RV) in the model: latent variables representing hidden states from the previous time step (context), H latent variables for the current time step, and 4 f variables. All random variables have categorical distributions. RV state values are denoted by corresponding lowercase letters: hi-1, Ni spt o.
In this work, each variable & is considered independent and has a separate graphical model for increased computational efficiency. However, in practice, hidden variables of the same time step statistically interdependent. We introduce their interdependence through segment computation trick that goes beyond the standard sum- sum- nun-product algorithm (see Appendix
B for details).
The model also has three types of factors: Miss-messgee from previous time steps, F- - factor (generalized transition matrix), F factor. We assume that messages M, contain posterior information from time step - 1, so we don't plot observable feature variables for previous time steps in Figure [.A
We will discuss only the upper block of the graph in Figure DA which is DHTM itself. The lower block-an encoder-is described in Appendix [ The only requirement for the encoder that its output should be represented as states of categorical variables (features) for the current observation.
The main routine of the DHTM to estimate distributions of hidden state variables given by the equation 3 the computational flow of which is schematically shown in Figure IB:
p(h+) x ΣI M (6;- (h₂ 24).
(3)
Ol LEWE
where summation is over all possible hidden RV states in Ω4 {hi_- i E wh]. i....., set of previous time step RV indexes included in ph factor, (n + 1) factor size.
3.2 NEURAL IMPLEMENTATION
For computational purposes, we translate the factor graph shown in Figure TA to the Hebbian learn- ing neural network architecture shown in Figure [ (for biological interpretation of the model,
4

---

Published as a conference paper at ICLR 2025
Chardlaaed 1e6i
MIT
MET
MIT
MF
H
log
ML
M
E!
ML
28L
MEN
Σ
maa
softmax
H
H4-1
segiment
A
cargette rocide
Figure 1: A. Example of factor graph for the DHTM. The input the model is a sequence of observations The encoder block forms categorical features Φ Each feature 4 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 and F are context and emission factors for the corresponding variables. Unary factors Mi-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)). The dendritic segments of the cell correspond to the context factor values for given combination states (active presynaptic cells). Segments' excitations Eg are combined to determine the cell's spike frequency p(h(). Segment synaptic weights reflect the specificity of the combination of presynaptic cells for the segment. Emission factors F& are fixed and represented by minicolumns within a variable.
Appendix C Each RV can be viewed set of neurons representing the states of the RV, i.e. p(h+) p(c) = 1), where j-is the index of neuron corresponding to the state hy and of {0, cell's activity. Since RVs are categorical, only one cell per variable can be active at a time and their activity is defined through sampling of RV's distribution.
Since factors 1% represent dependencies between variables, in neural implementation view them as collections of connections between cell groups corresponding to RVs that F takes as arguments.
Each factor's entry maps a particular combination of states ht -.......................... h, to a factor value fi- We found it useful to represent each such entry separate computational unit, which we call segment (alluding to dendritic segments of a biological neuron). As shown in Figure [B, segment / connects its presynaptic cells hi' h," (one per each RV), constituting its receptive field, to the cell it excites ht. The factor fi value can be thought of segment's excitation efficiency, i.e. how strong its presynaptic pattern influences the cell's activation.
Reformulation of belief propagation step in equation 3 in terms of segmental computations results in the following equations:
p(bt)) softmax max (Ei)), [HN (Eseglj)
(4)
1
E = log Si + wi log
+
Σ
(1 waa) log mu:
(5)
wErec(!)
wErec(!)
where fells[/[]]-indee)) of cells corresponding to states of the variable H.S. seg(i)-all segments for cell j, Wal synapse efficacy or neuron's specificity for segment, such that Wal = p(s = 1 1|c"-1
Ni 1), wi 1 E. average synapse efficacy for segment I, and ny- number of cells in segment's 'ss receptive field rec((). my is presynaptic cell's activation probability p(c²_, 1), which carries messages from the previous time-step and corresponds entry in Mi For detailed derivation and rationale behind these equations, see Appendix B
Initially, cells have segments. As learning progresses, segments, are created. In equation we benefit from having sparse factors because their complexity depends linearly on the number of non-zero entries or segments. And that's usually the case in
non-zero connections, grouped into
5

---

Published as a conference paper at ICLR 2025
model due to the one-step Monte Carlo learning and the special form of the emission factors F
(h s0f) T[ E col())"),
where 1-indicator function, col()p}) is a set of hidden states connected the feature state of 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 (George et al. 2021).
The next step after computing the p(h()) distribution is to incorporate information about the current feature states p(hg vf) x p(ht)T(i E col()))). After that, the learning step is executed. The loop closing step of our TM algorithm to assign the posterior of the current step p(ht (e) to M
3.3 LEARNING
DHTM learns fi and Wat weights by Monte-Carlo Hebbian-like updates. First, hi-1 and hy are sampled from their posterior distributions: p(hi11 (pi-1) x Mi and p(hy of)) respectively. Cells that correspond to sampled states become active. For cell j corresponding to hū we check if there are segments matching the pattern of sampled states (active cells) from the previous time-step (ht_)): If there are no such segments, the new segment connecting sampled cells is created. For all newly created segments and segments with receptive fields matching active cells from the previous time-step, its activity 57 is set to 1. Then fi is updated according to the activity of the segment 8t and the activity of its cell c;, that fi is proportional to several coincidences si cf = 1 in the recent past (exponential moving average of coincidence):
Afi = aad f1)8
where a E [0, 1) is the segment's ' learning rate.
That is, ft trained be equal the cell's average activation frequency when the segment is active. It's similar to Baum- Welch's update rule (Baum et al. 1970) for the transition matrix in HMM, which, in effect, counts transitions from one state another, but our case, the previous state (context) is represented by a group of RVs.
Weights wha are also updated by the Hebbian rule to reflect the specificity of a presynaptic u for activating segment I. That is, they are targeted to represent the probability p(st = ci_ 1) that segment 81 is active, given that cell u was active at the previous time step. In our algorithm, is approximated oi-1-
exponential moving average of the frequency activation of segment st, given
Awal B -I-I 1] (I[s wal):
where B E (0, 1) - learning rate.
3.4 AGENT ARCHITECTURE
We incorporate DHTM 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 ohk We assume that the reward function can be decomposed linearly in feature space R3 n - rit where r is reward associated with feature state sp i and n-nummer of feature variables. rf is learned during interaction with the environment and is used combination with SF representations to estimate the action value function (see Appendix A for details).
The agent training procedure is outlined in Algorithm T For each episode, the memory state is reset fixed initial message with RESET_ RESET_MEMORY () and the variable act on is set fixed initial action. An observation returned by an environment (obs) is encoded as set of categorical variables. In the OBSERVE () routine, memory learns to predict next feature states described in section 3 An agent learns associations to map feature states and rewards in the REINF ORCE () function:
N < r + allge i) +
(9)
where a is a learning rate, R reward for the current time step.

---

end for Published as a conference paper at ICLR 2025
Algorithm General agent training procedure
We include actions in DHTM by forcing some the hidden variables H. to rep- resent actions. That is, assume that information about action is contained in the hidden state of the model. For exam- ple, if we have actions, we set 4 states for one of the hidden variables and set its state from observing the action.
1: for episode= I... do
2: RESET MEMORY)
3: action initial_action _action
4: while (not terminal) and (steps < mtx_steps) do
5: obs, reward STEPO
6:
features <- ENCODEEPREPROCEESSOO)
7:
8: OfetRVEsfeatures, action)
In SAMPLE (), for each action and current state, TM predicts feature states up to time step T under policy n, and the predictions are combined to form SF (equation [Tp. For simplicity, we al- ways build SF under a uniform policy, which doesn't guarantee optimality, but found that it doesn't tt make significant difference to iterative policy improvement in our perimental setups. The planning horizon T is determined dynamically from the distribution of 4/+4- For each prediction step 1, we check whether the distribution of t is close to uniform (according Kullback-Leible distance) or whether any of states with positive reward has probability above threshold: in that case, the prediction cycle is stopped. The action value is determined by combin- ing the corresponding SF and feature state rewards as shown in equation II Finally, the action sampled from the softmax distribution over the action values.
features)
9: action +- SAMPLE
10: ACT(action)
11: end while
12:
SF = he) I'pa =j
(10)
(Te, a.) a=)== ΣΣ ST (
= = let1) ) p(he+1 | he, 0z)
(11)
kk he-1
3.5
DHTM AND EPISODIC CONTROL
It can be shown that DHTM combined with the agent architecture described in the previous sec- tion similar episodic control.
Algorithm Episodic memory learning
Input: Pr41, Oz
DTHM performs best when its hidden state space is large. Since the hidden state for unpredicted ob- servations is formed randomly, it's important that there are no accidental state collisions when re- membering transitions. That is, DHTM tries to remember new sequences without interfering with old memories. In effect, this makes DHTM buffer of agent trajectories with a built-in recall en- gine. Episodic Control based on similar idea, and to investigate its difference from DHTM, we implemented an Episodic Control agent, simpli- fied counterpart of DHTM, reflecting its core property that presumably helps it solve RL tasks.
1: h411 + (hr)
2: Pi+1 t argmax Fe(he+1,4p)
3: if 1 is null or s441 is not 44+1 then
4: h+11 + --rgmaax Fe (h, po+1) # K is the
hgg
set of keys for all actions
5: Do. (h+) + h4t1
6: end if
7: h4 + h441
Episodic Control (EC) a general computational approach storing and reusing an agent's expe- rience, inspired by the notion of episodic memory from psychology (Tulving et al. 1972).. The core of EC is a dictionary D. = (Ka, Va) that stores key-value pairs for each discrete action a E A. In our EC version, K, is an array of hidden states ht and V, is an array of corresponding subsequent hidden states hx+1- If the key state hú already exists the dictionary, its value is rewritten. That is, D, represents the transition matrix of deterministic HMM for a feature variable Φ,, The hidden state h1 and the feature state yi have the same dependence as in DHTM (equation 6. but EC we assume that there are infinite hidden states per feature state. To predict the next state, just have to look up p, for the current state ht but there is no entry for h, or the prediction does not match the observed state 11 then the new state is formed. Unlike DHTM, the state is not completely

---

Published as a conference paper at ICLR 2025
c
(onell)
cocg (onac
Strop
(orack)
c5og
(eracie)
rMi (crack)
(oradie)
tereckor
200 400
600
1000
1.0
155
2.0
2.5
3.0
3.5
4.0
4.5
Episode
Iotal Stags
K10*
Figure 2: POMDP 5x5 Gridworld foraging task using different types of memory: DHTM, CSCG, LSTM, RWKV. All agents, except for DHTM, are trained two modes. The first mode with oracle, which tells the agent reinitialize memory and clear the trajectory buffer during environ- mental change. The second mode is without oracle, so the agent continues using its current strategy gather new trajectories and update memory. There is also an episodic control agent (ec) among baselines, which uses simple dictionary to store trajectories and make predictions. A. Dynamic of steps required for the agent reach the goal. In the 300th episode, the goal blocked by a wall, so the agent needs to find a new optimal trajectory. B. Total steps taken during 1000 episodes.
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 's memory D2 is summarized in the algorithm
Since Da can be transformed into a transition matrix, the process of forming SF can be reduced to one in Sec. [.4] However, since the resulting transition matrix is sparse, it's more efficient to consider D, as a graph representation and perform SF formation as breadth-first search (see
Appendix E.3 for details).
4 EXPERIMENTS
We test our model in reinforcement learning task in the Gridworld enviromment 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 one standard deviation is shown everywhere, except for boxplots. See Appendix D for details on the baselines, their model parameters, and training regimes. See Appendix For more detailed description of the environments and setups"
4.1 GRIDWORLD
The first test was conducted in partially observable Gridworld in a 5x5 alternating maze with stationary end state and the agent's starting position. The agent can observe only the colour of the cell it is in and the colour of obstacle it encounters, while being punished in addition to the constant step penalty. Each episode ends when the reward is collected or when the maximum number of action steps reached. In each episode, 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, there the shortest path to the terminal state. In the second phase, after 300 episodes, the previously optimal path is blocked by a wall, so the agent has to take 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 reach the end state in the second phase, it suggests that the model can forget, at least suppress,
obstacles
'The source code of the experiments and algorithms is at httpa://github.com Cognitive Aariiiee-A---ysee -tw.reee---syyttttt/eeeeeeeeeeee
8

---

Published as a conference paper at ICLR 2025
kimcans
factor
25
30
factor
25
factor
20
factor
Dreamer V3 (18M) after 1k apisodes
factor1
10
factor
A 100 200
300
400
500
0.50
0.75
1.00
1.25
1.50
1.75
2.00
Episade
Total Stops
x10
Figure 3: Amount of action steps the agent performs reach the goal in 10x room in AnimalAI environment using DHTM, CSCG, LSTM, RWKV 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 left, turn right, or go forward. A. Dynamic of steps required for the agent to reach the goal. Total steps taken during 1000 episodes by the agent with DHTM with K-Means encoder and Categorical VAE. For the latter, DHTM's s context factor size is varied (i.e., how many previous time-step hidden variables are considered for prediction).
view of the environment and produces one of three actions: turn
the memory of the old shortest path when is blocked. Otherwise, the agent will continue to hit the wall and fail reach the terminal state.
In this experiment, we compare our memory model DHTM with CSCG (George et al. 2021), LSTM (Hochreiter & Schmidhuber 1997) and RWKV (Peng al. 2023) within the same agent architecture that uses memory to form successor features and evaluate actions (see See. 3.4.. All memories, except DHTM, are trained in two modes. The first mode is with 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 oracle, the agent continues trajectories and update memory in the second phase. There is also a simple episodic control agent (EC) among the baselines, which uses dictionary store trajectories and make predictions (see Section 3.5 for details).
its current strategy collect
The results in Figure 2 show that CSCG effectively can't adapt the environmental change in the regime, but only if it is retrained from scratch on new trajectories. And, apparently, LSTM and RWKV are only able to learn a predictive model for small horizon, since they perform significantly worse than other baselines in the second phase, even with oracle enabled. In contrast, DHTM and EC are able adapt online even without explicit memory's trainable parameters reinitialisation. Another important observation is that the EC agent performs par with DHTM, suggesting that its mechanism can explain most of DHTM's performance in this task. However, see Appendix H.2 for experiments that show their difference in scalability.
4.2 ANIMALAI
AnimalAI is 3D environment capable of providing rich visual input to 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 environments with larger state spaces.
The agent's task is to reach a sphere of radius 0.5 m, 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 AnimalAl'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.

---

Published as a conference paper at ICLR 2025
We test DHTM and EC agents using two different encoders: K-Means and Categorical VAE (Jang ct al. 2017). ) 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 VA is configured for distributed representation with five variables with 50 states each (see Appendix D details). In case of EC agent, since it doesn't support distributed representations, we map each unique VAE encoding to one feature state using simple dictionary. Since the resulting number of states for such feature variable is large (typically more than 1000), it's infeasible to use the same mapping for CSCG. We also didn't 't find much difference for LSTM and RWKV using K-Means or VAE, so we only keep the results for VAE encoder.
The results are shown in Figure AA. 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 LSTM and RWKV agents don'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 small version of Dreamer V3 (18M parameters) (Hafner et al. 2023), which success- fully solves the same task after about 35 x 10 action steps (1000 episodes) versus 3.5 X 10 steps (150 episodes) for DHTM (see Appendix H.. for details).
To show that DHTM can't simply be replaced by multiple parallel EC dictionaries or CSCGs in the case of distributed encoding, 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 the graph of the Factorial HMM, with parallel HMMs. It can be seen from Figure B that in our case of categorical VAE, knowing only the previous state of the variable not enough to reliably predict rewarding features, since the agent doesn't converge 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 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 E..4 For further experiments showing the importance of distributed encoding in DHTM, see Appendix H.3.
5 CONCLUSION
This paper introduces novel episodic memory model, DHTM. It combines ideas from belief prop- agation, episodic control, and neurophysiological models of the ncocortex, resulting in a fast proba- bilistic 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- table episodic con- trol. We also show that, unlike CSCG and table EC, DHTM is able to handle distributed representa- tions 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 gen- eralizing algorithm, such Neural State Space Model (Tiang et al., 2023). Through prediction error learning, DHTM can selectively discard and replace obsolete transitions in observed trajecto- ries 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 trajec- tory buffer are highly desirable in the case of online learning non-stationary environments.
Further improvements of DHTM can be directed towards the implementation of an automatic fea- ture 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 matter of future research.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. Advances in neural information processing systems, 30, 2017.
Leonard E Baum, Ted Petrie, George Soules. and Norman Weiss. A maximization technique occur- ring in the statistical analysis of probahilistic functions of markov chains. The annals of mathe- matical statistics, 41(1):164-11) 171, 1970.
PENG Bo. Blinkdl/rwkv-Imm 0.01, August 2021.
URL
012o65..199577
Mikael Boden. A guide to recurrent neural networks and backpropagation. the Dallas project, 2(2):
1-10, 2002.
Hamza Chaudhry, Jacob Zavatone- Veth, Dmitry Krotov, and Cengiz Pehlevan. Long Sequence Hop- field Memory. Advances in Neural Information Processing Systems, 36:54300-54340, Decem- ber 2023.
cc httpppppeer_iilleppprr///22222 hash/aa haas:///222ccc2eeeeeeeeefffffff. - Abstract--ooffeeecee.
Patricia Smith Churchland and Terrence Joseph Sejnowski. The computational brain. MIT press, 1992.
Matthew Crosby, Benjamin Beyret, Murray Shanahan, José Hernández-Orallo Lucy Cheke, and Marta Halina. The animal-ai testbed and competition. In Hugo Jair Escalante and Raia Hadsell (eds.), Proceedings of the NeurIPS 2019 Competition and Demonstration Track, volume 123 of Proceedings of Machine Learning Research, pp. 164-176. PMLR, 08-14 Dec 2020.
Peter Dayan. Improving generalization for temporal difference learning: The successor representa- tion. Neural computation, 5(4):613-624, 1993.
Antoine Dedieu, Nishad Gothoskar, Scott Swingle, Wolfgang Lehrach, Miguel Lázaro- Gredilla, and Dileep George. Learning higher-order sequential structure with cloned hmms. MeXXiv1995.005000 May 2019. URL 1xXiv:1905.00507 [cs, stat].
Antoine Dedieu, Wolfgang Lehrach, Guangyao Zhou, Dileep George, and Miguel Learning Cognitive Maps from Transformer Representations for Efficient Planning in Partially Observed Environments, January 2024. arXiv:2401 .05946 [es].
Yogesh K Dwivedi, Nir Kshetri, Laurie Hughes, Emma Louise Slade, Anand Jeyaraj, Arpan Kumar Kar, Abdullah M Baabdullah, Alex Koohang, Vishnupriya Raghavan, Manju Ahuja, et al. "so what if chatgpt wrote it?" multidisciplinary perspectives on opportunities. challenges and impli- cations of generative conversational ai for research, practice and policy. International Journal of Information Management, 71:10264:. 2023.
Sean R. Eddy. What is a hidden markov model? Nature Biotechnology, 22(10):1315-1316 Oct 2004. ISSN 1546-1696. doi: 10. 10.10388/bc1004413155 URL
David Emukpere, Xavier Alameda-Pineda, and Chris Reinke. Successor Feature Neural Episodic Control, August 2023. URL [cs].
arXiv:2111.03110
Gökcen Eraslan, Ziga Avsec, Julien Gagneur, and Fabian Theis. Deep learning: new computational modelling techniques for genomics. Nature Reviews Genetics, 20(7):399-03,, 2019.
Zafcirios Fountas, Noor Sajid, Pedro Mediano, and Karl Friston. Deep active inference agents using monte-carlo methods. In Advances in Neural Information Processing Systems, volume 33, pp. 11662-11675. Curran Associates, Inc., 2020.
Karl Friston, Thomas FitzGerald. Francesco Rigoli, Philipp Schwartenbeck, Giovanni Pezzulo, et al. Active inference and learning. Neuroxcience & Biobehavioral Reviews, 68:862-879, 2016.
11

---

Published as a conference paper at ICLR 2025
Karl J Friston, Richard Rosch, Thomas Parr, Cathy Price, and Howard Bowman. Deep temporal models and active inference. Neuroscience & Biobehavioral Reviews, 90:486-501, 2018.
Dileep George, Rajeev V. Rikhye, Nishad Gothoskar, J. Swaroop Guntupalli, Antoine Dedieu, and Miguel Eizaro-Gredilla. Clone-structured graph representations enable flexible learning and vi- carious evaluation of cognitive maps. Nature Communications, 12(1 22319), Apr 2021. ISSN 2041-1723. doi: 10.10099/41147--22-22555
Z. Ghahramani and M.I. Jordan. Factorial Hidden Markov Models. Machine Learning, 29(2-3): 245-273, 1997. ISSN 0885-6125, doi: 10. 1023/a:10074255:
Zoubin Ghahramani and Michael Jordan. Factorial hidden markov models. Advances in neural information processing systems, 8, 1995.
Swaroop Guntupalli, Rajkumar Vasudeva Raju, Shrinu Kushagra, Carter Wendelken, Danny Sawyer, Ishan Deshpande, Guangyao Zhou, Miguel I Lázaro-Gredilla, and Dileep George. Graph schemas as abstractions for transfer learning, inference, and planning. December 2023. URL
http://arxiv.cr wttp://v.rrg/bbb////ooo/o 2332 .03.5002033300 arXiv:2302.07:50 [cs, q-bio].
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04044 2023.
GM Harshvardhan, Mahendra Kumar Gourisaria, Manjusha Pandey, and Siddharth Swarup Rautaray. A comprehensive survey and analysis of generative models in machine learning. Com- puler Science Review, 38: 100285, 2020.
Jeff Hawkins and Subutai Ahmad.
Why neurons have thousands of synapses, a theory of se- Frontiers in Neural Circuits, 10, 2016. ISSN 1662-5110. URL https: / /ww. frontiersin org/ jrrr/journals//
quence memory in neocortex.
doi: 10.3389/fncir22216.0002
heural-cireultts/attiiiees
3389/ fneir. 2016.0002.
Donald Olding Hebb. The organization of behavior: neuropsychological theory. Psychology press, 2005.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9:1735- 80, 12 1997. doi: 10. 00.1162/neco.1199 wwt.eec....999999. 1997.9.8.1735.
Mehdi Jafarnia Jahromi, Rahul Jain, and Ashutosh Nayyar. Online learning for unknown partially observable mdps. In International Conference on Artificial Intelligence and Statistics, pp. 1712- 1732. PMLR, 2022.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax, 2017. URL httttttss///aiiii wrg. wtargaab//1611.0 01144
Shulei Ji, Jing Luo, and Xinyu Yang. A comprehensive survey deep music generation: Multi-level representations, algorithms, evaluations, and future directions. arXiv preprint arXiv:2011.0680, 2020.
Xiajun Jiang Ryan Missel, Zhiyuan Li, and Linwei Wang. Sequential latent variable models for few-shot high-dimensional time-series forecasting. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: /wwp//oeeewww.oco//// net wtwp://wwwee//orumt.e.mm forum?id 7C9aRX2nBf2.
F.R. Kschischang, B.J. Frey, and H.-A Loeliger. Factor graphs and the sum-product algorithm. IEEE Transactions on Information Theory, 47(2):988-51, 2001. doi: 10.1109/18.910777
Máté Lengyel and Peter Dayan. Hippocampal Contributions to Control: The Third Way. In Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc., 2007. URL /tttps://papers..t. wwwp:ppeeeerssss.peeeeeeeeeeeeee 1E4477badaaf361ooo L14773aaaEEA111133222aaaaaeeeee Abstrrctt .tml
Timothy P. Lillicrap, Adam Santoro, Luke Marris, Colin J. Akerman, and Geoffrey Hinton. Back- propagation and the brain. Nature Reviews Neuroscience, 21(6):335-344 Jun 2020. ISSN 1471- 003X, 1471-0048.
doi: 10. 0000008//411553-00007777777
12

---

Published as a conference paper at ICLR 2025
Zachary C Lipton, John Berkowitz, and Charles Elkan. A critical review of recurrent neural networks for sequence learning. arXiv preprint arXiv:1506.00019, 2015.
H.-A. Loeliger. An introduction to factor graphs. IEEE Signal Processing Magazine, 21(1):28-41, 2004. doi: N0.110//MSP.2044.22000
Christoph Mathys, Jean Daunizeau, Karl J Friston, and Klaas E Stephan. A bayesian foundation for individual learning under uncertainty. Frontiers in human neuroscience, 5:39, 2011.
Beren Millidge and Christopher Buckley. Successor Representation Active Inference, July 2022.
URL http:://xxxxi
09897 arXiv:2207.0989) [cs].
Bonan Min, Hayley Ross, Elior Sulem, Amir Pouran Ben Veyseh, Thien Huu Nguyen, Oscar Sainz,
Eneko Agirre, Ilana Heintz, and Dan Roth. Recent advances in natural language processing via large pre-trained language models: A survey, ACM Computing Surveys, 2021.
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-basee rein- forcement learning: A survey. Foundations and Trends@ in Machine Learning, 16(1):1-118, 2023.
V. Mountcastle. The columnar organization of the neocortex. Brain, 120(4):701-722, April 1997. ISSN 14602156. doi: 10. 10.1093//rrai/1120.47000 URL https::/accaee::: oup. com/ 4.tp:.1000///aa/112.....7...
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint arXiv:23033063449 2023.
Thomas Parr and Karl J Friston. Working memory, attention, and salience in active inference. Scientific reports, 7(1):14678, 2017.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Viti,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, al. Rwkv: Reinventing rans for the transformer era. arXiv preprint arXiv:2305.13048 2023.
Pascal Poupart. Exploiting structure to efficiently solve large scale partially observable Markov decision processes. Citeseer, 2005.
Alexander Pritzel, Benigno Uria, Sriram Srinivasan, Adrià Puigdomènech Badia, Oriol Vinyals, Demis Hassabis, Daan Wierstra, and Charles Blundell. Neural Episodic Control. In Proceed- ings of the 34th International Conference on Machine Learning, pp. 2827-2836. PMLR, July 2017. URL https://prooeeenngg.. mlr. http:///pw.p.e//////////aaaa2//. 2640-3498.
ISSN:
S. Ritter, J. X. Wang, Z. Kurth-Nelson, and M. Botvinick. Episodic Control as Meta- Reinforcement Learning, July 2018. URL htR.//ttts///wwwiooooo I1C1/350537v2 Pages: 360537
http///wwbborxii...o//otttttttte Section: New Results.
Achille Salaün, Yohan Petetin, and François Desbouvries.
Comparing the modeling powers of mn and hmm. In 2019 18th ieee international conference on machine learning and applications (iemla). pp. 1496-1499, IEEE, 2019.
Samuel J. Gershman. The Successor Representation: Its Computational Logic and Neu- ral Substrates. The Journal of Neuroscience, 38(33):7193, August 2018. doi: 10.1523/ INEUROSCLO151118.200
URL :ttt:://wwr ineurosci, wrg./oorrg/eoo oeen:oel..og///ctttttt///////// conttettttttt/33/
7193 933 abst.ract.
13

---

Published as a conference paper at ICLR 2025
Gautam Singh, Skand Peri, Junghyun Kim, Hyunseok Kim, and Sungjin Ahn. Structured world belief for reinforcement learning in pomdp. In Marina Meila and Tong Zhang (eds.), Pro- Machine Learning. volume 139 of Proceed-
ceedings of the 38th International Conference ings of Machine Learning Research, pp. 7745-9755. PMLR, 18-24 Jul 2021. URL https: //proceeding ml www.pprssssvv3. 8 mghh1a.......m htmml
Jochen Staiger and Carl C. H. Petersen. Neuronal eircuits in barrel cortex for whisker sensory perception. Physiological Reviews, 101(1):353-415, 2021. doi: 10.1 0tt2/phsseev.000192200
URL https: //oi: wwwp:oo.rg1...0 11t.0.11112phhhs..
0001 9. 2019 PMID: 32816652.
Greg J. Stuart and Nelson Spruston. Dendritic integration: 60 years of progress. Nature Neu-
roscience, 18(12):1713-1721, Dec 2015. ISSN 1546-1726. doi: 10.1038/nm.4177.
URL
https://doi. http://abi..o......003........
A.K Subramanian. Pytorch-vae. httppp///ggchha.....o/t/tt//////
VAE 2020.
Mufeng Tang, Helen Barron, and Rafal Bogacz. Sequential Memory with Temporal Predic- tive Coding. Advances Neural Information Processing Systems, 36:44341 Decem- ber 2023. / heerr neurips cc/paper http:/aprrfflessseeeeee
files/papeer20023 hatt://88bbcff118888888888666662 aostsssstattacccooCeoeeeeeeeeeee
Mufeng Tang, Helen Barron, and Rafal Bogacz. Learning grid cells by predictive coding, 2024.
abs/ 01022]
Edward C. Tolman. Cognitive maps in rats and men. Psychological Review, 55(4):189-208,
1948. ISSN 1939-1471, 0033-295X, doi: 10.1037/h0061620 StRthtp::11 /tttt:://doi.aggoe www.dppi.ppppppoooooooooooo....a apa. org/ getdoi cfm3:..mmmoo.... 10.1:0.100001111111111111111
Ahmed Touati, Jérémy Rapin, and Yann Ollivier. Does Shro-Shot Reinforcement Learning Exist? September 2022. URR.htttps://ppenrvve net.
Endel Tulving, Wayne Donaldson, Gordon H. Bower, and United States (eds.). Organization of memory. Academic Press, New York, 1972. ISBN oe
Konstantinos Voudouris, Ibrahim Alhas, Wout Schellaert, Matthew Crosby, Joel Holmes, John Burden, Niharika Chaubey, Niall Donnelly, Matishalin Patel, Marta Halina, José Hernández- Orallo, and Lucy G. Cheke. Animal-ai 3: What's new & why you should care, 2023. URL attps::/arxxe
/2312.2211.... 11414
Xiaohan Zhang, Lu Liu, Guodong Long, Jing Jiang, and Shenquan Liu. Episodic mem- ory governs choices: An RNN-based reinforcement learning model for mecision-making task. Neural Networks, 134:1-10, February 2021. ISSN 0893-6080. doi: 10.1016/j.neunet.
2020.11.003. URL https://ww..c....... /ttp://ww.eeeeeeee....ccceeeeeee
5089360802C33388
Jingyu Zhao, Feiqing Huang, Jia Lv, Yanjie Duan, Zhen Qin, Guodong Li, and Guangjian Tian. Do mn and Istm have long memory? In International Conference on Machine Learning, pp. 11365-1133 1i15--1375.. PMLR, 2020.
14

---

Published as a conference paper at ICLR 2025
A VALUE FUNCTION DECOMPOSITION
In our agent model, we approximate the reward function R(s) as a sum:
R
r(p)),
(12)
where r(v)) is reward associated with state sp), n--ummer of feature variables. Then, similarly to the Successor Representation idea (see Section 2.2). the value function can be represented as:
V(h,) = El
Rett+1
[4]
Σ E 1=0
E
r(o))) mx]
1=0
Y Y = he)r
SF (hr)r)
(13)
where ST (lx) = Eiso 7 2'p(p/tet> =j he), he = (li),.., - hidden state vector of variables {11)b.
B DETAILS OF NEURONAL IMPLEMENTATION
As shown in Figure BB, each RV can be viewed as set of neurons representing the states of the RV, i.e. p(ht) = p(c] = 1), where j-is the index of a neuron corresponding to the state ht. Cell activity is binary c. 1), and the probability can be interpreted 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 and m. respectively, where I corresponds to a particular combination of state values h. h h and u indexes all neurons representing states of RV's from previous time steps.
spike rate. Factors F and Mi
Inspired by biological neural networks, we group the connections of a neuron into dendritic seg- ments. segment acts as an independent computational unit that detects a particular input pattern (a contextual state) defined by its receptive field. In our model, segment connects the factor value fi and the excitation E; induced by the segment / the cell which it is connected. The segment is active, i.e. 81 if all its presynaptic cells are active, otherwise = 0. Computationally. segment transmits its factor value ft the cell to which attached if the context matches the corresponding state combination.
We can now rewrite equation as the following:
p(ht)
Σ
List.
(14)
LEseg(s)
where L1 = II my is segment's likelihood as long as messages are normalized to probability distributions, seg(i)-indexen of segments that are attached to cell j. rec(I)-indexes of presynaptic cells that constitute receptive field of segment with index 1.
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 Fik
(1 wtt TM E co)())
(15)
where I- indicator function, col(yp!) is a set of hidden states connected to the feature state of forming column. The shape of the emission factor is inspired by the presumed columnar structure
15

---

Published as a conference paper at ICLR 2025
of the neocortex and has been shown to induce sparse transition matrix in HMM GGorge al. 2021).
The segment likelihood L, resulting from the sum oum-product algorithm is computed as if the presy- naptic 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 -igliiklihoood:
log L = wy log ny
Σ wEree(1)
Wathh +
Σ
(1 (1-wigh Wat log may
(16)
aErec())
where Wal Synpsee efficacy or neuron specificity for segment, such that Wal p(s 1 ci- - 1), wt E. Whi averaee synapse efficacy for segment I, and
ny-number of cells segment's receptive field rec(!).
The idea behind the formula is to approximate between two extreme cases. The first is p(5 = 1|c__11 = 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:) My for all but we also reduce the prediction variance by averaging across different u. In the opposite case, p(s; = 1|c"-1 = 1) - for all u. presynaptic cells don't I cluster and the segment activation probability is just product of the activation probability of each cell.
The resulting equation for belief propagation in DHTM is the following:
p(ht)) = p(c 1)
softmax max (Ei)), jEcells[H[]] (Eseg())
(17)
where E = log fi + log Lt, [ of cells representing states for the variable H Here also approximate the logarithmic sum with the max operation, inspired by the neurophysi- ological model of segment aggregation by cells (Stuart & Spruston 2015).
The next step after computing the p( (ht) distribution parameters is to incorporate information about the current feature states p(h, y xc p(h, ([(k E col(pp))] After that, the learning step is ex- ecuted. The loop closing step of TM algorithm is to assign the posterior of the current step p(h v to M
C BIOLOGICAL INTERPRETATION
Neural implementation of the DHTM is inspired by neocortical neural networks (see Fig. Hid- den variables H may be considered as populations of excitatory pyramidal neurons in cortical layer L2/3 of somatosensory areas, with lateral inhibition modelled as softmax function. Staiger & Petersen 2021) showed that neurons in this layer are responsible for temporal context formation.
The neuronal activity timestep & can be thought carry messages Mf_-- Messages are propa- gated through synapses of dendritic segments, which correspond to factors Fk. Dendritic segments of biological neurons are known to be coincidence detectors of their synaptic input (Stuart & Sprus- ton 2015). That is, each segment can be thought of as detector of particular combination of active cells (states in our interpretation). If we associate segment's 's parameter with factor value, collection of dendritic segments represent a mapping from state space factor values, representing context factor Fe.
Our computational model of neuron and its dendrites is based on HTM neuron model (Hawkins & Ahmad 2016). In this model, each segment has receptive field and activates when number of active cells in this receptive field reaches a threshold. The neuron active if any of its segments active. We use the same activation rule, but the threshold is fixed to the receptive field size. Our algorithm of segment's 's growth on demand is also similar to HTM's. In principle, DHTM can be considired as probabilistic interpretation of HTM, allowing belief propagation.
Similarly to HTM framework, our temporal memory can be interpreted as a model of cortical layers. Feature variables & may be considered to represent cells of 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 1.4 are modelled columns for particular feature states cok( (Mountcastle. 1997).
16

---

Published as a conference paper at ICLR 2025
pstamidal nEuEOnssare
10
segment / factor value
log, f
H
max
prdvshilit
Hitl
p(I)
pepuilation of cellevaribble
Figure 4: Biological view of the neural implementation of the DHTM. Variables H;- correspond to populations of neurons that have common sensory input and lateral inhibitory competition. Dendritic segments correspond to factor values fi Spike frequency of a neuron reflects state probability p(hā) of variable.
D ENCODING OBSERVATIONS
Since DHTM designed to process categorical random variables, we need categorical encoder for RGB images to solve AnimalAI tasks.
To compare our model to CSCG and EC AnimalAI, we use the approach described Guntupalli ct al. (2023) to quantize observation space. K-Means (implementation from scikit -laarn) with 50 clusters pretrained on x 10a 64x64 RGB observations from trajectories sampled by uniform policy. To encode image, the distance from each cluster is calculated and the index of the closest cluster is returned. In that way, AnimalAl'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 (Jang al. 2017) with the latent space represented by five categorical variables with 50 states each. We use slightly modified implementation from Subramanian (2020). The only difference is that 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 12 The model is trained on 1 64x64 RGB images for 20 epochs with batch size 64. An image 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.
E IMPLEMENTATION DETAILS
As mentioned in Section 3.4 we incorporate DHTM as a part of an RL agent, which has a mem- ory model and a feature reward function. All tested memory models share the same pipeline- conditioned on actions, they learn sequences categorically encoded observations and are used to form SFs.
all experiments Section 4.1 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 unstable. For SF formation the agent uses 7 = 0.8 and Y = 0.9 for Gridworld and AnimalAI experiments respectively. Maximum SF horizon is set to 50, but it is usually shorter since employ early stop condition based on feature variables distribution. That is, the prediction is close to uniform or the probability of state with associated positive reward meets threshold, the
17

---

Published as a conference paper at ICLR 2025
planning is interrupted. We probability thresholds equal to 0.05 and 0.01 for Gridworld and AnimalAI respectively.
In scaling Gridworld experiments (Appendix H.2
we increase maximum planning horizon up to 100 steps, and we also had to decrease agent's 's temperature to 0.02 for DHTM and to 0.002 for EC agent. Discount factor is 0.7 for DHTM and 0.9 for EC. We found that, in 10x10 maze, DHTM's strategy becomes more stable we increase the positive reward threshold up to 0.4. EC agent con- verges faster if we increase its softmax temperature, however, it converges to less optimal trajectory on average.
In the noise tolerance AnimalAT experiment (Appendix H.3)) for DHTM we use the same set of agent parameters, except we set the goal probability threshold to 0.5. We also set the goal detection threshold equal to the mean positive reward since we found that noise variable states mostly get reward assignment slightly above zero, in contrast to signal variables, in which only several states get positive reward assignment.
E.I LSTM AND RWKV
LSTM baseline was implemented with a single LSTMCell from PyTorch library Paszke et al. 2019). 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 the input representation, where symexp activation function, symexp(x) sign(r)e is reverse of symlog function: symlog = sign(x) log (|æ) + 1).
The similar way we implemented RWKV baseline: a single RWKV layer supported by single-layer encoder and decoder. Current public RWKV implementation is fast evolving framework (Bo) 202T], and for the increased performance for the transformer architectures. In case we needed so-called sequential mode for online learning similar to LSTM. Thus, we adapted another public implementation mentioned in the official documentation (RWKV in 150 lines of code).
tightly bound the offline batch training common
For both Gridworld and AnimalAI tasks, LSTM and RWKV with a hidden state size of 100 are 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 RNNs into groups. In all experiments the hidden state represents 10 categorical vari- ables with 10 states. That is, RNN 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 Dreamer V2/V3 (Hafner al., 2023). The idea of using symexp activation function, mentioned above, is inspired by Dreamer too, and is used to remedy the prob- lem of learning extreme probability values. Without symexp the neural network has to represent zero probability with high negative logit values and one- 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.
E.2 CSCG
CSCG baseline was implemented using code from the repository accompanying the paper (https:
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 a = 0.05 described in the paper (Dedieu al. 2019). This smoothed transition matrix was used as an initialization for the next training batch and for inference. The first batch observations gathered using the uniform policy, and later batches are gathered using the current best policy.
//ithhb.. com/vi ca i ous cec
In so-called oracle mode, the transition matrix is trained from scratch every 500 steps. The buffer is discarded at the environmental change and the transition matrix is reinitialised.
18

---

Published as a conference paper at ICLR 2025
E.3 EPISODIC CONTROL
Here, we provide an algorithm for SF formation using EC dictionary Da (see Algorithm 3) con- trast to other memory algorithms, EC dictionary doesn't It estimate state transition probabilities, but stores connections between them all-or-nothing fation. Therefore, on each time step, all predicted observations are assumed equally probable. But it's not limitation in our set of experiments, since we tested our algorithms in deterministic environments.
Algorithm SF formation for EC
Input: ho. E (0,
Output: SF
I: SF t array of zeros
2: PS { # PS is the set of all states (nodes) on the current BFS depth
3: goal found < false
4: for 1..7 do
5: PS < 4 (h){h)) # get next depth nodes assuming the policy is uniform
6: counts <- array of zeros
7:
for all h E PS do
4 argmax (e) & counts, counts + 10: goal_found < rp >> 0 # check if the feature state is rewarding
8:
9:
11: end for
12:
SF = SF + + NORMALIZE(counts)
13: if goal found is true then
14: break
15: end if
end for
E.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, it significantly increases the amount of trajectories memory is able to store by decreasing collision probability for hidden state representations of different trajectories.
In AnimalAl experiments, hidden variables state spaces of size 40 states per feature state. For Categorical VAE encoder, we use one hidden variable per feature variable, except the case (factoo 1*) that tests that information from one feature variable is sufficient for correct prediction, in which three hidden variables per observation variable are used and the factor of size three connecting only hidden copies of the same feature variable. Examples of typical DHTM's graphical models used in our experiments are presented 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 has connections to, but also connects one hidden variable of the current step and the action variable.
Figure 5 It's important to note that edge colours denote
The most important DHTM's parameters, except the factor size, are factor learning rate a and its initial value fo. 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 fu = 1, 0.1 and fo 0.05, and in AnimalAI we set
allowing faster convergence. In Gridworld a = 0.001 and fo = We also set the limit on the number of segments x 10 but practically never reach it in our experiments.
19

---

Published as a conference paper at ICLR 2025
A
M
MC
A-1
H
M
M
H
2
K- K-Moann Landyarrd
M
M
8F
H
M
M
H
H
H
A
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 1, A. Graphical model for 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).
F EXPERIMENTAL SETUPS
F.1 GRIDWORLD
We use simple homebrew Gridworld implementation. Each Gridworld position maps to colour, reward, and terminal state indicator. There are two special colours for the world border and obsta- cles. Each time-step, 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 through wall or world border in addition to base step reward punishment to facilitate faster convergence optimal trajectories.
Samples of setups used in our experiments are depicted in Figure 6 Additional setups for scaling experiments are shown Figure Here, the red colour corresponds obstacles (negative numbers) and the light gray colour to the reward (number 4). Letters A and G correspond to the agent's 's initial position and goal, correspondingly. 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.
F.2 ANIMALAT
AnimalAI is a testbed inspired by experiments with animals (Crosby et al. 2020). The environment consists of 3D area surrounded by a wall and many different objects that can placed using configuration file, including walls, food, ramps, trees, movable obstacles, and more. We tested our agent in a 10x 10-meter room surrounded by walls (see Fig. 8). Each timestep, the agent gets an RGB image of 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.
20

---

Published as a conference paper at ICLR 2025
2
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 (number 4). The right picture represents the environment after the reward is blocked by obstacles (red squares negative num- bers). The letter A depicts the agent's ' initial position and the letter G-the goal.
a
1
2
3
5
3
3
1
Figure 7: Schematic view of POMDP Gridworld setups for scale experiments. From left to right:
5x5, 7x7, 10x10. The letter A depicts the agent's initial position and the letter G-the goal.
Prev reward: 0.000
Reward: -0.556
Health: 44.4
Figure 8: Sample of the agent's first-person view of our setup in AnimalAI environment. The green ball represents goal, a food item to be collected.
21

---

Published as a conference paper at ICLR 2025
Score
0.5
-0.5
Step
120k
Figure 9: Dreamer V3 (18M) score in Animal experiment described in this paper.
G CAPACITY
In this section, we provide basic theoretical analysis DHTM's capacity. Since the hidden state for unpredicted feature state transition is generated randomly, the main source of errors is accidental collision of hidden states. Due to special form the emission factor (see equation 6), the collision possible only for the same feature sate. Here analise the capacity of DHTM for factor grpah depicted in Figure [ Let's assume that the stream of feature states is sampled from uniformly distributed categorical random variable. Therefore, the probability have collision Pe ".,......"
where numbe of hidden variables, nhe; fl are the number of states feature and hidden variables correspondingly. Therefore, the probability of having no collisions for a l transitions is equal to (1 - Pe)'. Let's set a threshold с E (0, 1), which determines our confidence in that the sequence of transitions is collision free. Solving the equation 1 - (1 Pe) with respect to we get:
log(1 c)
(18)
log( Rhp Ph
Using the Taylor expansion for Ro + inf, we can see that grows linearly with the total number of states n Therefore, DHTM's capacity scales linearly with the number of rons. This, however, can be easily improved if we set additional prior on hidden state activation probability, which would take into account hidden state usage, but we leave it for future research.
H ADDITIONAL EXPERIMENTS
H.1 DREAMER
In order ground our experiments state-of-the-art model based algorithms, an additional baseline, we test the smallest version of Dreamer V3 with 18 million parameters. We use the imple- mentation described in Voudouris et (2023), which is based on Halner al. (2023) 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 x 10 action steps. Taking into account the frame skip used in our experiments with DHTM, it's approximately 35 x 10 action steps versus 3.5 x 10 for the DHTM agent to solve the task.
H.2 SCALABILITY
In this section, we explore DHTM's scalability properties and compare them to the EC agent. In RL experiments, like presented in section the number of segments largely depends on the RL algorithm used, its exploration strategy, and how fast converges to stable policy. We tested DHTM and EC agents in empty Gridworld of different sizes: 5x5, 7x7 and 10x10, where the agent and the goal are placed in opposite corners (see Fig. 7 As shown in Figure TOA, the DHTM-baased
22

---

Published as a conference paper at ICLR 2025
agent requires fewer episodes to find the goal; however, the policy becomes less stable as the size of the Gridworld grows. As a consequence of faster convergence, DHTM produces fewer segments than the number of EC's dictionary entries throughout experiments. It also can be seen that for both DHTM and EC agents, the number of memory units grows non-linearly with the state space size (see Fig. TOB).
X10*
100
Jtem 5c5
1.75
dham
(dictionary sizel
Jiam 7x7
3htm 10x10
5.5
Number of Segments
1.25
Stape
0.75
0.25
0.00
30
50
60
70
90
100
Episode
Gridworid States
Figure 10: Comparison of BC and DHTM-based agents foraging task in POMDP Gridworld of different sizes: 5x5, 7x7 and 10x10. A. Dynamic of steps required for the agent to reach the goal. B. Number of segments grown (or dictionary entries for EC agent) after 300 episodes with respect to state space size.
principle, the number of segments reflects number of transitions stored, which proportional to the amount of experience provided to the memory. To show that, tested DHTM and EC agent in 10x10 empty Gridworld using the same memory parameters as for experiments in section[4.1] but under uniform strategy. That is, agents freely explored the environment during 100 episodes and at the end the number of segments for DHTM and dictionary size for EC was recorded for each episode length. Figure 11 shows that DHTM's number of segments and EC's dictionary size grow linearly with the maximum episode length. In other words, since DHTM lacks generalisation abilities it largely serves as a trajectory buffer, in the same sense dictionary in EC.
DHTM's computational complexity (time and memory) of the prediction step grows linearly with respect the number of segments. Therefore, to use it for huge environments, there is need to prune unnecessary segments or remember only selected episodes. This is a common problem among episodic-like memories. The current approach form SF also should be adjusted to be scalable, since it requires maximum planning horizon to be equal to the episode length. Among the improved approaches could be to employ n-step TD learning, hierarchical planning SF caching, to recaleu- late it less frequently.
H.3 NOISE TOLERANCE
In order show the key difference between DHTM and its simplified EC version, tested agents in the same AnimalAI setup as described in 4.2 but we add one noise variable to the VAE encoder. It's easy to see that the naive mapping from distributed VAE encoding to one categorical variable used for EC doesn't t work in this case, since even identical observations will likely be encoded differently due to the noise variable. contrast, since DHTM is able to handle distributed representations directly (as the name implies), it still can use non-noise variables independently to learn transitions.
In this experiment, we vary DHTM's factor size (see Fig. 12 random, probability to include the noisy variable increases with factor size and this consequently worsen the performance the agent. If the factor size is set to the total number of variables, DHTM becomes effectively unfactorised, similarly to its EC counterpart.
Since factors formed completely
23

---

Published as a conference paper at ICLR 2025
x10*
(dictionary size]
2.0
Number of Segments
1.5
Episade Length
Figure 11: Comparison of DHTM's number of segments with EC agent's 's dictionary size depending on episode length in Gridworld after uniform exploration during 100 episodes.
40
25
20
dHtm factor
factor
15
dHtm factor
dhtm factor
100
200
300
400
500
Episode
Figure 12: Number of action steps the agent takes to reach the goal in a 10x 10-meter room in the AnimalAI environment, using DHTM and EC models with a distributed Categorical VAE encoder that includes one additional noisy variable (six variables in total). DHTM's factor size is varied, determining the number of variables sampled for each factor.
I GLOSSARY
Categorical Random Variable-a discrete random variable that take on of finite K possible states.
Cortical Column or Minicolumn- 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 's activity.
Factor Graph-bipartin graph representing the factorization of probability distribution, with one part representing factor nodes and another-randon variables.
Factor size number variables connected to the factor. For simplicity, we refer only to the number of previous hidden states, used for prediction. Hidden Markov Model (AMMM)-sttissie model of a stochastic process where state probability depends only previous state of the process.
24

---

Published as a conference paper at ICLR 2025
Multi-compartment neuron model-a a model of neuron that divides neuron's '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 environment that is usually hidden.
algorithms insightful information about the
Successor Representations (SR)-a discounted sum of future [one-hot encoded] observations. Successor Features (SF) a generalization of SR, discounted sum of future latent states.
Temporal Memory (TM) in this work by this term we mean "memory for sequences".
25


Reviews:
[{"review_id": "HRqf5ChIQY", "reviewer": "Reviewer_JqYM", "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.", "questions": "Any of the weakness are also questions.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "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.", "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.\n2. 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:\n- 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?\n- 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.\n- 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.\n\n3. 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).\n- 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.\n- The calculation of the emission factors E is never shown explicitly! but it appears in equation 7\n- Equation 7 is not well justify at all, it is very difficult to see how you go from (4) to (7)\n- 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?).\n- 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.\n\n4. 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?\n- 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!\n- 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)\n- What is the capacity of the memory?\n\n5. 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)\n\n6. Others\n- It is not clear to me how are new segments created\n- the change from Fe to fl for the emission factors was confusing, please add some warning\n- 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\n- How is the planning horizon \"T\" determined. How are you measuring the shape of the distributions to determine this parameter?\n- The explanation of the Episodic control agent needs to be rewritten for clarity.\n- In equation 10, make clear that Rt is the reward coming from the environment (there are other definitions for Rt before)\n- 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.\n- 291,292, please explain this step.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7bPn6RNCxN", "reviewer": "Reviewer_uBbj", "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.", "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?\n- 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. \n- 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.\n- 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.\n- 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.\n- I would like to see the effect of an oracle on Episodic Control and DHTM in Figure 4.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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. \n- The algorithm is computationally efficient due to sparse transition matrices and distributed representations, and could be potentially implemented efficiently in neuromorphic hardware.\n- 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.", "weaknesses": "- The unique encoding of each observation sequence limits the model's ability to generalize across different sequences which represent the same underlying state.\n- 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. \n- The baseline models like LSTM may not be fully optimized for online learning, potentially skewing the comparison results.\n- The paper lacks a thorough theoretical analysis of the algorithm's properties such as convergence guarantees and computational complexity.\n- I think you should have more exposition explaining what CSCG, Episodic Control, and Dreamer V3 are.\n- I am willing to improve my score if you address some of the questions posed in the following section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BvlmzxfkEw", "reviewer": "Reviewer_UWu5", "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.\n\nThe 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.", "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?\n- How well does DHTM handle partial observability? Can it distinguish aliased observations well?\n- CSCG has been shown to learn the spatial structure of the environment which allows for generalization during navigation. Can DHTM do the same?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 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.\n- The model is online and uses local learning rules, which makes it more applicable.\n- The model is biologically inspired and the authors make specific connections to its neural implementation/plausibility.\n- The authors provide an interpretation of the model in terms of episodic control.", "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.\n- Since the model uses a factor graph formalism, I think having a section for that in the background would improve clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b70JCMNz8T", "reviewer": "Reviewer_cQzX", "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.", "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.\n* 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.\n* \"Factor graph formalism\" should be defined in the introduction.\n* \"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.\n* 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.\n* What is meant by \"$f_l$ is proportional to several coincidences $s_l = c_t^j = 1$ in the recent past\" (line )? \n* 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?\n* 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.\n* 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.", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "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.\n* The experimental evaluation of the method is reasonable and with some more detail would support its effectiveness.\n* The discussion of limitations of the method in the conclusion is thoughtful and thorough.", "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.\n* 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.\n* Key details of the experiments, particularly how the baseline models were tuned, are missing and make it difficult to evaluate them.\n* 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.\n* There is no substantial discussion of related work.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)