Characterizing neural representation of cognitively-inspired deep RL agents during an evidence accumulation task
Authors: James Mochizuki-Freeman
Sahaj Singh Maini
Zoran Tiganj
Venue: Submitted to ICLR 2023 (2023)

Abstract: Evidence accumulation is thought to be fundamental for decision-making in humans and other mammals. It has been extensively studied in neuroscience and cognitive science with the goal of explaining how sensory information is sequentially sampled until sufficient evidence has accumulated to favor one decision over others. Neuroscience studies suggest that the hippocampus encodes a low-dimensional ordered representation of evidence through sequential neural activity. Cognitive modelers have proposed a mechanism by which such sequential activity could emerge through the modulation of recurrent weights with a change in the amount of evidence. This gives rise to neurons tuned to a specific magnitude of evidence which resemble neurons recorded in the hippocampus. Here we integrated a cognitive science model inside a deep Reinforcement Learning (RL) agent and trained the agent to perform a simple evidence accumulation task inspired by the behavioral experiments on animals. We compared the agent's performance with the performance of agents equipped with GRUs and RNNs. We found that the agent based on a cognitive model was able to learn much faster and generalize better while having significantly fewer parameters. We also compared the emergent neural activity across agents and found that in some cases, GRU-based agents developed similar neural representations to agents based on a cognitive model. This study illustrates how integrating cognitive models and deep learning systems can lead to brain-like neural representations that can improve learning.

Under review as a conference paper at ICLR 2023
CHARACTERIZING NEURAL REPRESENTATION OF COGNITIVELY-INPPIRED RL AGENTS DURING AN EVI- DENCE ACCUMULATION TASK
Anonymous authors Paper under double-bling blind review
ABSTRACT
Evidence accumulation is thought to be fundamental for decision meassion-maiing in hu- mans and other mammals. It has been extensively studied in neuroscience and cognitive science with the goal of explaining how sensory information is sequen- tially sampled until sufficient evidence has accumulated to favor one decision over others. Neuroscience studies suggest that the hippocampus encodes low- dimensional ordered representation of evidence through sequential neural activity. Cognitive modelers have proposed mechanism by which such sequential activ- ity could emerge through the modulation recurrent weights with a change in the amount of evidence. This gives rise neurons tuned to a specific magnitude of evidence which resemble neurons recorded in the hippocampus. Here we in- tegrated a cognitive science model inside a Reinforcement Learning (RL) agent and trained the agent to perform a simple evidence accumulation tasks inspired by the behavioral experiments on animals. We compared the agent's 's performance with the performance of agents equipped with GRUs and RNNs. We found that the agent based on a cognitive model was able learn faster and generalize bet- ter while having significantly fewer parameters. We also compared the emergent neural activity across agents and found that in some cases, GRU-based agents de- veloped similar neural representations to agents based on cognitive model. This study illustrates how integrating cogmitive models and artificial neural networks can lead to brain-like neural representations that can improve learning.
INTRODUCTION
Converging evidence from cognitive science and neuroscience suggests that the brain represents physical and abstract variables structured form, as mental or cognitive maps. These maps are thought play essential role in learning and reasoning (Tolman 1948 Ekstrom & Ranganath,
2018 Behrens et al. 2018). Cogmitive maps are characterized by neurons that activate sequentially function of the magnitude of the variable they encode. For instance, neurons called place cells activate sequentially as a function of spatial distance from some landmark (Moser et al. 2015] Muller 1996. Sheehan et al. 22021). Similarly, time cells activate sequentially function of
elapsed time from some event (Pastalkova et al. 2008 MacDonald et al. 2011 Cruzado et al. 2020
Salz et al. 2016).
Similar sequential activity has also been observed for sound frequency (Aronov et al. 2017). prob- ability (Knudsen & Wallis (02T] and accumulated evidence (Nich et [202 Morcos & Harvey 2016b). For example. in the "accumulating towers task" Nich et al. (2021) trained mice to move along virtual track and observe objects (towers) on the left- and right-hand sides. When mice arrived at the end of the track, to receive a reward they had to turn left or right, depending on which side had more towers. The difference in the number of towers here corresponds to the amount evi- dence for turning left vs. turning right. Nieh et al. (2021) recorded activity of hundreds of individual neurons from mice hippocampus, part of the brain commonly thought to play key role in naviga- tion in physical and abstract spaces (Bures et al. 1997 Eichenbaum 2014 Moser al. 2015). The results indicated the existence of cells tuned to a particular difference in the number of towers, such that a population of neurons tiles the entire evidence axis (Nieh al. 2021) (see also Morcos &

---

Under review as a conference paper at ICLR 2023
Harvey (0016b)) This provides valuable insight into how abstract variables represented in the brain.
Cognitive scientists have developed elaborate models of evidence accumulation to explain the sponse time a variety of behavioral tasks (Laming 1968 Link 1975 Ratchif 1978). These models hypothesize that the brain contains internal variable that represents the progress towards the decision. A neural-level cognitive model proposed that the brain could implement this process using a framework based the Laplace transform (Howard al.) 2018). The Laplace framework gives rise to map-like representations and it has been successful in describing the emergence of sequentially activated time cells (Shankar & Howard 2012) and place cells (Howard et al. 2014
Howard & Hasselmo (020).
Artificial neural networks (ANNs) are commonly thought to have a distributed representation that does not have map-like structure. While ANNs excel in many domains, they still struggle at many tasks that humans find relatively simple. Unlike humans, ANNs typically require large number of training examples and fail generalize to examples that are outside the training distribution (Bengio) 2017 Vene 2017; Marcus 2020). Using cognitive models informed by neural data as an inductive bias for ANNs is important direction that can help not only advance the current AI systems but also improve our understanding of cognitive mechanisms in the brain.
Here integrate the Laplace framework into reinforcement learning (RL) agents. The Laplace framework based on recurrent neurons with analytically computed weights. We use the Laplace domain to generate map-like representation of the amount of evidence. This representation is then fed into trainable RL module based on the A2C architecture (Mnih et al.) 2016). We compare map-based agents standard RL agents that use simple recurrent neural networks (RNNs) and Gated Recurrent Units (GRUs) (Chung et al. (2014) in terms of performance and similarity of the neural activity to neural activity recorded the brain.
Contributions of this work are as follows:
We integrated cognitive model for evidence accumulation based on the Laplace transform into an RL agent.
We showed that symbolic operations in the Laplace domain give rise to individual neurons that are tuned to the magnitude of the evidence, just like neurons neuroscience studies (Nieh et al. 2021
Morcos & Harvey 2016a).
We found that agents based on the Laplace framework learn faster and generalize better than agents based on commonly used RNNs. This indicates that RL agents were able to efficiently use the brain-like sequential representation of evidence.
We found that GRUs performed much better than RNNs, suggesting that gating plays an important role in constructing neural representation of time- varying latent variables. This is consistent with the cognitive modeling work, which uses gating to convert representa- tion of elapsed time into representation of accumulated evidence.
Episode
Tower appears the right
Tower appears
Tower
End of the Agent needs to turn right
Agent aceds move foeward
en the right
Time
Figure 1: Schematic of the accumulating towers environment. In this simple example, two towers appeared on the right, and one tower appeared on the left, so the agent has to turn right once reaches the end of the track. Each tower is encoded with single pixel value.
2

---

Under review as a conference paper at ICLR 2023
2 METHODS
2.1 ENVIRONMENTS
Inspired by the neuroscience studies from Nieh et al. (2021) Morcos Harvey (201b)), we de- signed simple version of the accumulating towers task. To evaluate robustness and generality of the proposed approach, aside from default version of the task, we also designed two other tasks, namely range count task and exact mentary material, and they will be made publicly available and made open source together with the entire code used implement the artificial agents and generate all of the results in the manuscript.
task. The environments are provided a part of the supple-
2.1.1 ACCUMULATING TOWERS TASK
Agents had to navigate down a virtual track composed of only three inputs: left, right and middle. In each episode, agents start from the beginning of the virtual track and observe towers (represented by the input value changing from 0 to 1) on each side of the environment (Fig. [ Agents had four available actions: left, right, forward, and backward. Positions of towers were decided randomly in each episode. Similar to the neuroscience studies, the maximum number of towers on one side was 14. Once the agent arrived the end of the track, the middle input changed from signifying that hit the wall. In order to receive the reward, the agent had to turn left or right, depending on which side had more towers. We the magnitude of the reward 10, penalize the agents for hitting the wall at the end of the track or going backward with -1 negative reward, and penalize the agent for hitting the side walls with a -0.1 negative reward.
2.1.2 RANGE COUNT TASK
In this task, agents observed the same environment as in the accumulating towers task, but in order to obtain reward, agents had to turn left only if the number of towers the left-hand side was larger than five and right in every other case (regardless of the number of towers on the right-hand side).
2.1.3 EXACT COUNT TASK
In the counting task, agents again observed the same environment as in the accumulating towers task, but the maximum number of towers each side was 5 instead of 14, and agents had to turn left only if the number of towers on the left-hand side was exactly and right in every other case.
2.2 RNN IMPLEMENTATION OF EVIDENCE ACCUMULATION USING THE LAPLACE DOMAIN
Evidence accumulation implemented through the Laplace domain a key component the cognitively-inspired RL agent (Fig. (2) We will first describe the Laplace framework for functions of time and then convert representation time into a representation of evidence and show how can be implemented an RNN.
h t
Recurrent
π
Layer
V
1
h
Figure 2: The agent architecture. We compare simple RNN, GRU and Laplace-based RNN de- scribed here.
3

---

Under review as a conference paper at ICLR 2023
b
500
50
100
150
200
30
30
100
150
200
50
100
150
200
Time
Time
150
210
50
150
50
30
100
150
200
Time
T Time
Time
14
50
200
500
100
160
200
50
301
100
150
2001
T iim
T Time
Time
501
150
200
50
30
100
100
200
300
1000
150
2001
T
T
Time
50
100 150
200
50
100
150
200
Dislance
Figure 3: Example of the Laplace and inverse Laplace transform with and without modulatory input. (a) In the absence of modulatory input (a = 1) the impulse response of the Laplace transform decays exponentially with decay rate s. The impulse response of the inverse Laplace transform has unimodal shape. Note that if time t was shown on the log-scale, the unimodal curves would be equally wide and equidistant. (b, c) a modulates the decay rate F and it is proportional to change the count (b) or distance (c). This makes units in f develop unimodal basis functions that tuned to count or distance rather than to time and peak at n² or x respectively.
2.2. LAPLACE FRAMEWORK
We define the Laplace transform of function [() from -o to the present l:
F(s:t) e (t=(t-
(1)
We restrict variable to real positive values.
The above equation can be expressed in a differential form where appears as a rate constant of leaky integrator:
dF - (s; s(t).
(2)
dl
Fig. 34 shows the impulse response of the above equation for several values of
To convert representation of time into a representation of numerosity n(t) (how many times some input was observed) extend Eq. and modulate with a rate of change a expressed as time derivative of numerosity (a = dn/dt))
dF(stt dn
(-sF(s;1)+ +()))
di
By reorganizing terms in the above equation and applying the chain rule we can rewrite the equation as function of instead of / (Fig. 30):
dF (s;n)) (s; + f(n),
(4)
dn
and we set f(n) = 5(0).
Note that the same approach can convert function of time into function of any variable, which derivative can be learned from the environment. For instance, a = dz/dt velocity, then the network would represent traveled distance (Fig. BD).
'The Laplace transform defines complex variable. This choice would result in expomentially growing and oscillatory neural activity, causing numerical instahilities when computing the inverse Laplace transform.
4

---

Under review as a conference paper at ICLR 2023
Inverting the Laplace transform reconstructs the input function of the internal variable n", which corresponds to n. The inverse, which we denote as f(n"; can be computed using the Post inversion formula (Post 1930):
f(n;;)) = LE ((;;n) = (-1)* k!
k+1 dh
F(;;),
(5)
ds
where n² := k/s and & + 30. As we show below, the reconstruction gives rise to units tuned to O we see that [(n":)) peaks at n" For s being
particular By solving 0f continuous variable and & + oc, the width the peak is infinitesimally small, providing a perfect reconstruction of the observed quantity.
2.2.2 DISCRETE IMPLEMENTATION
For neural network implementation, we discretize the Laplace and inverse Laplace transform for both and L. To select values in an informed way, compute the impulse response of sar
kk+1 u(t) K!
((()
k+1
Fn
(6)
where u = Eizoa(t)) att)) When is discrete and & finite, furm is a set of unimodal basis functions (when is continuous and & - co, those unimodal basis functions turn into delta functions with spacing + 0). The coefficient of variation of fn-:. is independent of n² and n:: This implies that the width of the unimodal basis functions increases linearly with their peak time. When observed as function of log(n), the width of the unimodal basis functions is constant. This property the Post inversion formula is relevant for modeling human perception due to the Weber- Fechner law (Fechner 1860/1912) Portugal & Svaiter, 201T). This law states that the relationship between the perceived magnitude of the stimulus and its true magnitude is logarithmic, motivating the use of logarithmic units such decibel and candela. To ensure equidistant spacing of unimodal basis functions along the log-axis we space n² logarithmically. This results in dramatic conservation of resources, especially when representing large quantities since the number of units in fn+: grows as a finctiin of (og(n) rather than n. Note that fixing the values of n² and choosing & also fixes values of since = k/n".
1/vk + 1.
We now write a discrete-time approximation of Eq. equation an RNN with diagonal connec- tivity matrix and linear activation function:
Fst = WF F fts
where W diag(e -(())=A)). A discrete approximation of the inverse Laplace transform, fara, can be implemented by multiplying Faa with derivative matrix L computed for some finite value of k.
2.2.3 SUBTRACTION OF FUNCTIONS USING THE LAPLACE DOMAIN
In the accumulating towers task, Eq. can enable the agent to learn to represent the number of towers each side. However, the latent variable that should determine the agent's 's decision is not the number of towers on each side but the difference between those numbers (the agent needs to turn towards the side which had more towers). This is non-trivial problem since the number of towers is not represented as a scalar but as a function over n. Fortuitously, the Laplace domain enables access number of useful operations, including subtraction of two functions (Howard al. 2015). To show this, let us define (a) and g(a) as functions representing two distributions of possible values for the number a in the range 0 amaz- Outside this range, the functions are assumed to vanish. We define the operation of subtraction of these two distributions [f - (a) to be the cross-correlation of the two functions:
If gla) = 1 f(=')g(a+ x' d')d''
To illustrate that the above operation results in subtraction of two functions, consider simple case where each of the functions is a delta function: = 8(a,) and g = 8(a2). Then [f g] is delta function @1 d2- To implement cross-correlation in the Laplace domain turn Eq.
5

---

Under review as a conference paper at ICLR 2023
GRU
SnU
RNN
Fau
b
10
Reward
9
Reward
Reward
8-
25k
50k
0
25k Episode
50k
25k
50k
Episode
Episode
Figure 4: Agent performance on (a) accumulating towers task, (b) range count task, and (c) exact count task. In each task proposed architectures (either feub or Feut) learned the task faster than GRU despite having almost three orders of magnitude fewer parameters.
into convolution by reflecting g(a) around amaz: 9r = g(amar - a). Point-wise product of the Laplace transforms of two functions, f(a) and g(a), corresponds their convolution in the time domain. Point-wise multiplication of the Laplace transform of (a) and gr(a) corresponds to cross- correlation of (a) and ((a) in the time domain, which is equivalent their subtraction [f (a). Note that for subtraction we need to consider both positive and negative values. Since we only use positive values of 8, we
able directly represent the negative axis. To work around this, compute both [f (a) and وg fl(a)
2.3 AGENT ARCHITECTURE
The agents received three inputs from the environment that were fed into the recurrent layer (Fig. figg22 The recurrent layer either RNN, GRU RNN based on the Laplace framework as described above. When the Laplace framework was used we had three independent f modules. Each module had 20 n² values spaced logarithmically from 5 to 100. The value of parameter & was set to 8. Parameter controls the sharpness of the unimodal basis functions and this value was chosen to ensure no gaps between them. Note also that the input into ff was delivered to a. which controls recurrent weights of the population of units F. This is conceptual difference in comparison to other RNNs, where recurrent weights are tuned separately for each unit. The strength of the proposed approach is that the population of neurons encodes a function over a latent variable that the network needs to learn such that / directly represents the count of the objects. In addition to computing f. also computed the subtraction frus of each pair of f. This was done by computing the product in the Laplace domain between each pair F as described in the previous section. The total number of units was 180 (20 units per module, 3 independent modules and subtraction modules). When other RNNs were used, the dense layer was mapped to 180 recurrent units.
The output of the recurrent layer passed to actor network and critic network. Both actor and critic consist of a single layer fully connected neural network. We set the discount rate to y (since in this task, the reward immediately available agents after they made correct turn). For all agents, we explored two different learning rates 0.001 and 0.0001.
3 RESULTS
3.1 PERFORMANCE OF RL AGENTS
We trained and evaluated agents in three different RL environments: accumulating towers, range count and exact count. We compared several agents based on the Laplace framework: the proposed agent using the inverse Laplace transform and the subtraction (fsub), without the subtraction (f),

---

Under review as a conference paper at ICLR 2023
50
50
50
50
Neuron
Neuron
Neuron
Neuron
100
100
100
100
Evidence
Evidence
Evidence
Evidence
Figure 5: Neural activity of (fub) agents after 100k episodes of accumulating towers task. Similar to plots in Nich et al. (202)) MMrcos & Harvey (2016b), neurons are sorted by peak activity. Each row is normalized such that the activity ranges from to 1.
a
d
50
Neuron
Neuron
Neuron
Neuron
100
100
100
100
150
150
150
150
10
Evidence
Evidence
Evidence
Evidence
Figure 6: Same Fig.
but for GRU agents.
with the subtraction but without the inverse Laplace transform (Fsu)), and without both subtraction and the inverse Laplace transform (F). This was done to evaluate the importance of different com- ponents of the Laplace framework. We also compared agents based existing RNNs, including simple RNN and GRU, as well as versions those agents but with frozen recurrent weights. The control experiment with frozen recurrent weights was conducted because the Laplace agents have analytically computed recurrent weights s which are modulated on a population level with a as shared weight.
d=300
d=3000
d=10000
Parameters
faub
10.000 + 0.000
10.000
+ 0.000
10.000 + 0.000
724
9.903 + 0.066
9.925 + 0.065
9.925 + 0.041
244
Foub
10.000 0.000
10.000
= 0.000
10,000 = 0.000
724
F
8.995 + 0.287
8.675
0.361
8.900
0.272
244
RNN
4.475 - 0.456
4.600 + 0.281
4.700 + 0.576
34024
GRU
9.98 + 0.000
9.600 = 0.146
8.425 = 0.504
100624
RNN FROZEN
-21.0 + 0.000
-0...0 0.000
-00100 0.000
724
GRUFROZES
14.473 = 5.653
-149.3 + 44.77
449.4 + 131.3
724
Table 1: Results of agents trained of the accumulating towers task. The table shows mean reward +/- standard error across four runs after 100k episodes of training d=300 steps long environment. Validation was done in 300, 3000 and 10000 steps long environments.
The agents were trained and evaluated 300 steps long environment. We trained four different agents for each of the eight models and performed 100 validation runs every 1000 episodes.

---

Under review as a conference paper at ICLR 2023
d=300
d=3000
d=10000
Parameters
Sanb
10.000 + 0.000
10.000
0.000
10.000 + 0.000
724
10.000 + 0.000
10.000
0.000
10.000 + 0.000
244
F Faub
9.325 + 0.074
9.525
0.096
9.400 + 0.122
724
F
9.150 + 0.075
9.350 + 0.115
9.375 + 0.147
244
RNN
-3.225 + 8.033
45.550 + 44.875
-15.22 + 131.6
34024
GRU
9.125 + 0.758
4.15 0.557
4.300 = 0.696
100624
FROFRE
-31.00 + 0.000
-201.0 + 0.000
-00.0 = 0.000
724
FROZEN
-31.00 + 0.000
-201.0 0.000
601.0 = 0.000
724
Table Results of agents trained of the range count task. The table shows mean reward +/- standard emoor across four runs after 100k episodes of training in d=300 steps long environment. Validation was done 300, 3000 and 10000 steps long environments.
d=300
d=3000
d=10000
# Parameters
fsub
10.000
0.000
10.000
0.000
10.000
0.000
724
10.000
0.000
10.000
0.000
10.000 0.000
244
F sub
10.000 0.000
10.000
0.000
10.000 0.000
724
F
7.975 + 0.129
8.250 + 0.075
8.275 + 0.246
244
RNN
7.975 + 0.222
8.000 + 0.146
8.225 + 0.222
34024
GRU
9.375 1 0.375
8,050 + 0.545
8.175 + 0.219
100624
R
21.10 + 8.574
148.7 1: 45.27
448.8 + 131.8
724
GRUFROZER
21.23 + 8.465
148.7 = 45.27
448.8 = 131.8
724
Table Results of agents trained of the exact count task. The table shows mean reward +/- standard error across four runs after 100k episodes of training in d=300 steps long environment. Validation was done 300, 3000 and 10000 steps long environments.
3.2 AGENTS BASED ON THE COGNITIVE MODEL LEARNED THE TASKS FASTER THAN GRU-BASED AGENTS
While in some cases other agents had a better start, Faus agents were the first to learn the tasks and converge to a reward value 10 (see supplemental video of feut agent performing the accumulating towers task). This indicates that the cognitive model provided good representation, and once the A2C algorithm learned to use it, it was able to perform perfectly (Fig. 4 and the first column in Table 1 Table 2 and Table 3 range). In accumulating towers task and exact count task, Faub agents converged next. While the difference between faub and Faub agents is only in the inverse Laplace transform which is just a linear projection, this result suggests that the sequential map-like activation in feub was more useful in learning to perform the task with perfect accuracy than the Laplace representation with exponentially decaying traces. f and F agents performed above chance (agents reached the end of the track and made a correct decision in more than 50% of cases), but did not reach the performance of 10 during 100k episodes on accumulating towers task suggesting that the subtraction operation which constructed map-like representation for evidence was important for learning.
also Fig. 13 Fig. 14 and Fig. 15 for results
wider training
GRU agents managed to reach performance close to 10, indicating that they can learn the tasks as well (see supplemental video of GRU agent performing the accumulating towers task). On the other hand, the RNN agents did not learn the tasks 100k episodes indicating that gating was important for correct performance. It is important to note that the cognitive model also constructed the representation of evidence by a gating mechanism. Similarly, GRU can learn to modulate the range of change in neural activity by the amount of change in the evidence. Frozen models were not able to learn the task, failing to even reach the end of the track and make a random decision as indicated by the total reward being negative.
8

---

Under review as a conference paper at ICLR 2023
3.3 AGENTS BASED ON THE COGNITIVE MODEL WERE ROBUST TO CHANGES IN THE ENVIRONMENT
To test the ability of agents to generalize, we also evaluated them on 3000 and 10000 steps long tracks without ever training them on tracks of that length (second and third column in Table [ Table 2 and Table 3 Agents based on the cognitive model showed great resilience to this kind of rescaling. This is not surprising since the representation was designed to change as a function of change in the amount of evidence, so rescaling the environment did not have any impact. On the other hand, the performance of GRU agents dropped at unseen track lengths but remained well above chance. This suggests that GRU agents were able to learn to modulate their activity by the change in the amount of evidence: when there was no new evidence, there was little to no change in the activity, making the impact of track rescaling relatively small.
We also evaluated agents in environments that had to 30 towers on each side instead of 1 to 14 towers (Table 4). While all agents suffered some loss in reward, GRU and feus seem most resilient to this change, indicating potential similarity in the form of those two representations.
3.4 NEURAL ACTIVITY INSIDE THE RECURRENT LAYER RESEMBLES ACTIVITY IN MICE IIРPOCAMPUS
We visualized the neural activity of each of the four agents for each of 8 models after 100k episodes of training the accumulating towers task. As expected, neurons in feu agents activated sequen- tially as a function of evidence resembling the activity in neural recordings from the hippocampus (Nich et al. 2021 Morcos& & Harvey 2016b) (Fig. 5). f neurons also showed some tuning to the amount of evidence, but since they were able to only count objects on each side (and not subtract them), the tuning is blurry (Fig. (7) Neurons in F and Feut agents showed gradual changes as func- tion of evidence rather than tuning to particular magnitude of evidence reflecting the exponential dynamics of the Laplace transform (Fig. and Fig. prespectively)
Some of the neurons in GRU agents showed tuning to the magnitude of evidence with often promi- nent asymmetry between positive and negative amount of evidence (Fig. (). (See supplemental video of changes in the neural representation during training of GRU agent.) In comparison to neurons in RNN agents (Fig. [ GRUs had firing pattern significantly more dependent on the amount of evidence. Neurons in frozen GRU agents also showed some tuning to the magnitude of evidence (Fig. 10 although less than neurons in trained GRU agents.
CONCLUSIONS
We evaluated different artificial agents on three tasks 1) evidence accumulation task that mimicked the procedure of a recent neuroscience study (Nich et al.] 2021 Morcos & Harvey] (016b)). 2) range count task, and 3) exact count task. We used simple setup with only three inputs and compared agents based on cognitive model with agents based on simple RNN and GRU,
Agents based on a cognitive model were able to learn faster and generalize better despite having almost three orders of magnitude fewer parameters than GRU-based agent (244 vs 100624). This an indicator that the A2C algorithm was able to use the neural representation from the cognitive model. This representation also resembled data from neural recordings in Nieh et al.) (2021); Mor- cos & Harvey (2016b) characterized with the sequential activation as a function of the amount of evidence.
Agents based on the cognitive model represented perceptual information on a logarithmic scale, consistent with the Weber-Fechner law. This type of representation saves resources and can explain human behavioral outputs in perceptual tasks.
While we focused experiments that involve numerosity, the proposed architecture can represent any latent variables (e.g., size, distance or luminosity) as a log-compressed number line. The only condition is that the latent variables change over time such that their time derivative can be extracted from the input signal (see Fig. B for examples of compeeseede number lines for discrete and continuous signals).

---

Under review as a conference paper at ICLR 2023
REFERENCES
Dmitriy Aronov, Rhino Nevers, and David W Tank. Mapping of non-spatial dimension by the hippocampal-entoriine eircuit. Nature, 543(7647)779--22. 2017.
Timothy EJ Behrens, Timothy H Muller, James CR Whittington, Shirley Mark, Alon B Baram, Kimberly L Stachenfeld, and Zeb Kurth-Nelson. What is cognitive map? organizing knowledge for flexible behavior. Neuron, 100(2):90-500, 2018.
Yoshua Bengio. The consciousness prior. arXiv preprint
2017.
J Bures, AA Fenton, Yu Kaminsky, and 1. Zinyuk. Place cells and place navigation. Proceedings of the National Academy of Sciences, 94(1):343-350, 1997.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks sequence modeling. arXiv preprint arXiv:1412.35555, 2014.
Nathanael A Cruzado, Zoran Tiganj. Scott L Brincat, Earl K Miller, and Marc W Howard. Conjunc- tive representation of what and when in monkey hippocampus and lateral prefrontal cortex during an associative memory task. Hippocampus, 30(12):1332-11) 2020.
Howard Eichenbaum.
Time cells in the hippocampus: Nature Reviews Neuroscience, 15(1 15(11):732-744 2014.
dimension for mapping memories.
Arne D Ekstrom and Charan Ranganath. Space, time, and episodic memory: The hippocampus is all over the cognitive map. Hippocampux, 28(9):680-687, 2018.
Gustav Fechner. Elements Psychophysics. Vol. Houghton Miffin,, 1860/1912.
Mare W Howard and Michael E Hasselmo. Cognitive computation using neural representations of time and space in the laplace domain. arXiv preprint arXiv:2003. /668, 2020.
Mare Howard, Christopher J MacDonald, Zoran Tiganj, Karthik H Shankar, Qian Du, Michael Hasselmo, and Howard Fichenbaum. A unified mathematical framework for coding time, space, and sequences in the hippocampal region. Journal of Neuroscience, 34(13):4692-707 2014. doi: 10. S
Mare W Howard, Karthik H Shankar, and Zoran Tiganj. Efficient neural computation in the laplace domain. In Cognitive computations workshop at Advances in neural information processing sys- tems, 2015.
Mare W Howard, Andre Luzardo, and Zoran Tiganj. Evidence accumulation in laplace domain decision space. Computational brain & behavior, 1(3):237-251, 2018.
Eric B Knudsen and Joni D Wallis. Hippocampal neurons construct map of an abstract value space. Cell, 184(18):460-4655, 2021.
Donald Richard John Laming. Information theory of choice-reaction times. 1968.
Steve Le Artificial intelligence pioneer says need to start Arlington, VA: Axios, 2017.
Steve W Link. The relative judgment theory of two choice response time. Journal of Mathematical Psychology, 12(1):114-133 1975.
C. J. MacDonald, K. Q. Lepage, U. T. Eden, and H. Eichenbaum. Hippocampal "time cells" bridge the gap in memory for discontiguous events. Neuron, 7144):337-74, 2011.
Gary Mareus. The next decade in ai: four steps towards robust artificial intelligence. arXiv preprint arXiv:2002.061177. 2020.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International conference on machine learning, pp. 1928-1937. PMLR, 2016.
Ari S Morcos and Christopher D Harvey. History-dependent variability in population dynamics during evidence accumulation in cortex. Nature Neuroscience, 19(12):162-11681 2016a.
10

---

Under review as a conference paper at ICLR 2023
Ari S Morcos and Christopher D Harvey. History- History-dependent variabilityy in population dynamics during evidence accumulation in cortex. Nature neuroscience, 19(12):662--161,, 2016b.
May-Britt Moser, David Rowland, and Edvard Moser. Place cells, grid cells, and memory. Cold
Spring Harbor perspectives biology, 7(2)=021888, 2015.
Robert Muller. quarter of a century of place cells. Neuron, 17(5):813888) 822, 1996.
Edward Nieh, Manuel Schottdorf, Nicolas Freeman, Ryan Low, Sam Lewallen, Sue Ann Koay, Lucas Pinto, Jeffrey L Gauthier, Carlos D Brody, and David W Tank. Geometry of abstract learned knowledge in the hippocampus. Nature, pp. 1021.
Eva Pastalkova, Vladimir Itskov, Asohan Amarasingham, and Gyorgy Buzsaki. Internally generated cell assembly sequences in the hippocampus. Science, (5224)):3322-13333) 2008.
RD Portugal and Benar Fux Svaiter. Weber-fechner law and the optimality of the logarithmic scale. Minds and Machines, 21(1):73-81, 2011.
E. Post. Generalized differentiation. Transactions of the American Mathematical Society, 32:723- 781, 1930.
R. Ratcliff. A theory memory retrieval. Psychological Review, 85:59-108, 1978.
Daniel M Salz, Zoran Tiganj, Srijesa Khasnabish, Annalyse Kohley, Daniel Sheehan, Mare W Howard, and Howard Eichenbaum. Time cells hippocampal area ca3. Journal of Neuroscience, 36(28):7476-7484 2016.
K. H. Shankar and M. W. Howard. A scale-invariant internal representation of time. Neural Com- putation, 24(1):134-193, 2012.
Daniel J Sheehan, Stephen Charczynski, Blake A Fordyce, Michael E Hasselmo, and Mare W Howard. A compressed representation of spatial distance in the rodent hippocampus. bioRxiv, 2021.
Edward C Tolman. Cognitive maps in rats and men. Psychological review, 55(4):189, 1948.

---

Under review as a conference paper at ICLR 2023
A APPENDIX
d=300
d=3000
d=10000
# Parameters
Laub
9.892 + 0.067
9,850 +: 0.056
9.925 + 0.041
724
7.060 + 0.220
7.175 +: 0.343
6.800 + 0.242
244
Faub
8.845 0.151
8.775 4- 0.137
8.525 0.248
724
F
7.293 0.257
7.625 + 0.410
7.350 0.109
244
RNN
4.538 + 0.449
4.800
4.775 + 0.461
34024
GRU
9.895 0.016
9.725 + 0.054
9.025 0.219
100624
RNN FROZEN
21.000 0.000
201.. 0.000
-601.0 = 0.000
724
GRU FROZEN
14.555 + 5.582
1149.6 + 44.5
-449.7 + 131.0
724
Table 4: Mean reward +/- standard error across four runs on accumulating towers task after 100k episodes of training d=300 steps long environment and to 14 towers shown agents. Validation was done in 300, 3000 and 10000 steps long environments and with I to 30 towers shown to agents.
d
Neuron
Neuron
Neuron
50
50
50
50
Evidence
Evidence
Evidence
Evidence
100
100
100
100
Tumedd right (9)
80
Tumedd right (9)
80
80
Turned right (%)
80
60
Turned right
60
40
40
10
Evidence
Evidence
Evidence
Evidence
Figure 7: Neural activity (top row) and psychometric curves (bottom row) of f agents after 100k episodes of accumulating towers task. Top conveys the same information Fig. 5 Since f did not have perfect accuracy, we also show the psychometric curves. Dots in the psychometric curves indicate actions taken by the agent different amounts of evidence.
12

---

Under review as a conference paper at ICLR 2023
Neuron
Neuron
50
50
50
50
Evidence
Evidence
Evidence
Evidence
100
100
100
100
Turned right (%)
Turned right (%)
80
80
Turned right (%)
60
Turned rightt
Evidence
Evidence
Evidence
Evidence
Figure 8:
Same Fig. but for F agents
d
50
50
Neuron
Neuron
Neuron
100
100
100
100
150
150
150
150
-10
Evidence
Evidence
Evidence
Evidence
100
100
100
100
Tumed right 1%
Tumned right 1%)
80
80
Turned right 196)
Turned right
20
-10
-10
10
-10
10
Evidence
Evidence
Evidence
Evidence
Figure 9: Same Fig. but for Foub agents.
13

---

Under review as a conference paper at ICLR 2023
d
50
50
Neuron
Neuron
100
100
100
100
150
150
150
150
10
Evidence
Evidence
Evidence
Evidence
100
100
100
100
Turned night (%)
80
Tumeed right (9)
80
80
Turned right (%)
Turned right
60
60
40
40
10
Evidence
Evidence
Evidence
Evidence
Figure 10: Same Fig. but for GRU frozen agents. Note that some psychometric curves empty, indicating that those agents did not learn to reach the end of the environment.
d
Neuron
Neuron
10
10
100
10
-10
-10
0
10
Evidence
Evidence
Evidence
Evidence
100
100
100
Turned right (%
Turned right (%)
Turned right (%)]
60
Turned right
40
40
40
20
20
20
0
0
-10
0
10
-10
10
-10
10
-10
10
Evidence
Evidence
Evidence
Evidence
Figure 11: Same as Fig
but for RNN agents.
14

---

Under review as a conference paper at ICLR 2023
d
Neuron
Newron
150
-00
10
-10
10
-10
-10
10
Evidence
Evidence
Evidence
Evidence
100
100
100
100
Tureed night (%)
Turned right (%)
Turedd right (9)
right
40
20
20
20
0-
-10
10
-10
-10
10
-10
10
Evidence
Evidence
Evidence
Evidence
Figure 12: Same as Fig. but for RNN frozen agents. Note that psychometric curves are empty, indicating the agents not learn to reach the end of the environment.
GRU
F
Fub
RNN
F sub
10
Reward
4
0
20k
40k
60k
80k
100k
Episode
Figure 13: Agent performance on the accumulating towers task (same data as in Fig. 4h but different scale).
15

---

Under review as a conference paper at ICLR 2023
GRU
F
Jub
RNN
F
10
Reward
8
4
0
20k
40k
60k
80k
100k
Episode
Figure 14: Agent performance on the range count task (same data as in Fig. 4p but different scale).
GRU
F
RNN
F
10
Reward
8
4
0
20k
40k
60k
80k
100k
Episode
Figure 15: Agent performance the exact count task (same data as in Fig. 4e but different scale).
16


Reviews:
[{"review_id": "M8oNgLB-t_3", "reviewer": "Reviewer_tPG3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper presents a novel framework for reinforcement learning, specifically for evidence accumulation tasks that involve counting quantities and deciding which is greater. The general approach is to use cognitive models as inductive biases by learning a function representing the evidence and giving the resulting representations to an RL algorithm. The paper includes a discretized version of the Laplace transform-based framework to use with a neural network.\n\nThe framework is as follows: the Laplace transform is rewritten as a differential equation in terms of the numerosity of the stimulus and discretized; an inverse is also derived. This is meant to enable learning of a counted quantity. To subtract one quantity from another and determine which one is higher, the paper presents a cross-correlation mechanism. \n\nThis framework integrates into an RL algorithm (here A2C) by having environment inputs go into a recurrent layer, have that layer learn a direct representation of the count of the objects and their difference (in theory), and then pass the output of this layer to the actor and critic networks. \n\nExperiments are done on an \"accumulating towers\" task. There are several ablations removing various parts of the Laplace framework, as well as RNN and GRU layers as baselines. Task reward statistics show that the Laplace and inverse Laplace transforms with subtraction do very well, as does GRU. Other ablations have mid-range performance, and RNN does quite poorly. Frozen versions of RNN and GRU are unsurprisingly very poor.\n\nWhen evaluating on harder versions of the task, Laplace with subtraction continues to do well, as does GRU. The paper attributes this to the gating mechanisms present in both methods.\n\nFinally, there is some analysis of neural activity and its connections to human cognitive maps. Results show that Laplace with subtraction appears to activate sequentially as a function of evidence. GRU and other versions of the Laplace framework also shows some tuning to \"magnitude\"/\"amount of evidence\".", "summary_of_the_review": "I am recommending a weak acceptance for this paper. It is a good paper and I enjoyed reading it, but I am not convinced of its significance from a general modeling perspective given the limitations I think I see in its domain. I also am not totally sure I understood all the math, so I may be missing limitations; I may also be missing strengths. \n\nAll that said, the quality of idea and execution is far above a workshop paper and it is a very compelling read, so I wouldn't consider it fair to recommend rejection just because it's an interesting and well-demonstrated idea rather than the next big SOTA system like a lot of papers. I would raise my score if I were convinced that this opens a lot of future work directly (not just in high-level concept) or enables improvement on new tasks out of the box, and if I were convinced that it provides meaningful improvement over existing methods.", "strength_and_weaknesses": "**Strengths**\n- Introduction is strong and motivates the paper well. The idea of cognitive models providing inductive biases is not novel but this is a very interesting instantiation of it\n- The connections to human cognitive maps are more convincing than a lot of \"human cognitive construct + our cool new method\" connections that ML papers attempt to make, especially because of follow-up evidence in the paper such as the logarithmic representation. While even this one is ultimately conceptual (and the logarithmic representation appears to be a design choice rather than an emergent property, if I understand correctly), I find it compelling.\n- The paper is well-packaged - it makes several claims and demonstrates them, and the methods solve exactly the problem motivated and outlined. To some degree, the final contribution about gating could use more expansion, but I realize it's more a validation of prior work and I appreciate the acknowledgement of the similarities/power of existing work\n\n**Weaknesses**\n- The Laplace framework presented here seems very tuned to this specific task, or at least counting-based tasks. The paper acknowledges in \"Future Work\" that this is by no means the only type of latent variable and encourages work on others, but as of right now the framework is engineered in terms of counting - or at least, that's the sense I get reading it. It's still an interesting method, but GRUs and RNNs are much more general modules; it would be helpful to have (brief!) discussion of either 1) the domain limitations of these equations or 2) the potential to easily adapt this method to other tasks. I see how this formulation may extend to some other quantities, but I don't know for sure. \n- The discussion I described above would make the comparisons in this paper make more sense, but for actual acceptance I'm not sure that would be enough. This is an interesting idea and novel to my knowledge, but without seeing experiments in other domains, it's not clear how much value it adds. A more general approach covering a larger domain (potentially through multiple frameworks like this one), and experimental results on those, would be more convincing.\n- From a pure utility perspective, it's unclear how much value this method provides over GRU (not to mention, it has a narrower domain if I understand correctly). Because the performances are so similar, some statistics differentiating them would help.\n- The word \"deep\" is used in the title and throughout the paper, but the version of A2C experimented on uses single-layer FCNs for actor and critic. The paper isn't making claims about scale or ability to handle large real-world tasks or complicated environments, so this claim seems not only false but unnecessary to the story.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LC5m2PCa70", "reviewer": "Reviewer_BAao", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper embeds a learnable recurrent network based on Laplacian transformation into a reinforcement learning agent and claims that the Laplacian based recurrent network outperforms the others.", "summary_of_the_review": "I have gone through the math derivations and read the manuscript for several times and I believe I try my best to understand the paper but there are still some places which are not clear.", "strength_and_weaknesses": "It seems to be novel to use Laplacian transformation to derive a neural representation of evidence or numerosity. I believe this framework may suggest some general principle of neural coding. Nevertheless, I feel some parts of the paper are not well written and can be improved further.\n\n### Major concerns\n\n- Based on the text it is unclear to me that how the authors embed the Laplacian transformation into a recurrent network, even if I read the manuscript several times. Do authors firstly perform a Laplacian transformation on the input, and then feed the transformed inputs to a recurrent net? Or do authors use the Laplacian transformation to propose or constrain the recurrent net structure? \n\n- The connection between Eqs. 3 and 7 implicitly suggests that the firing rate of neurons represents the Laplacian transformation of an input. It is better to have more detailed explanations at here.\n\n### Minor concerns\n\n- If I understood correctly, Eq. 3 is wrong because it is inconsistent with Eq. 2, while Eq. 4 is correct. There must be some typos.\n\n- It is not clear of the meaning of n* in Eq. 5. Is it the preferred evidence of a neuron? Also, is $\\tilde{f}(n*, n)$ the reconstruction of $f(t)$ by only using a Laplacian basis?\n\n- $\\alpha(t_i)$ is not defined and I need to guess. Is it the instantaneous change of numerosity as mentioned right above Eq. 3? Besides, is the modulatory input $(\\alpha=1)$ also the change of numerosity? If so, I don't understand why $\\alpha=1$ represents no modulatory input.\n\n- The definition of $W_L$ right below Eq. 7 is contradictory with each other and quite confusing.\n\n- The math definition in Eq. 8 is no problem but I don't understand why author call the cross-correlation as subtraction between functions.\n\n- The 4th line to the 7th line below Eq. 8 is repetitive.\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "5MlTjpmJVD", "reviewer": "Reviewer_ojuW", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "The paper describes a network that combines a Laplace framework into a reinforcement learning architecture achieve superior performance over GRU and RNN in an evidence accumulation task both in terms of reward gain and generalization. Moreover, the activity of units in the proposed architecture resembles map-like neural activity in the brain for accumulated evidence. ", "summary_of_the_review": "I think because of lack of significant modeling/theory contribution as well lack of test-domains this paper is not ready to be presented in the ICLR venue yet. \n\nUpdate after rebuttal: \nI appreciate the authors' response and adding more experiments. The added experiments, however, are very much in the same domain of previous one and only a threshold has been changed (the accumulation mechanism is the same). In the evidence accumulation task (e.g. in mentioned labs), a sufficiently task has either 1) different number of options (e.g. more than 2) 2) different decision time criterion (reaction time v.s. fixed) or 3) dynamic weigh of evidence through the task. Moreover, even if the proposed approach is new, we do not know whether it is works better than a simple Bayesian update, which is in fact shown to be able to solve evidence accumulation task.  ", "strength_and_weaknesses": "Strength: The paper looks into an interesting and commonly used task in decision making and RL experiments in the field of Neuroscience. \nThe paper is clearly written and different components are very well explained step by step. It also provides a virtual environment, useful for the community to simulate evidence accumulation experiments. \n\nWeakness: This work only covers one specific task from one family of experiments (Evidence accumulation). The paper can be significantly improved if the method is tested on broader range of tasks in the face of uncertainty (e.g. see \"Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs\" by Ni et al, ICML 2022), and/or works on hippocampus and cognitive maps (many references  presented in the paper) and/or different evidence accumulation tasks in neuroscience (e.g. see works of Michael Shadlen lab for human and monkey, and Carlos Brody lab for rodent). In terms of architecture, comparison of an architecture with a Bayes filter (e.g. see implementations of \"End-to-end learnable histogram filters\", by Bonschkowski & Brock, NeurIPS deep RL workshop 2016, or \"QMDP-Net: Deep Learning for Planning under Partial Observability\" by Karkus et al, NeurIPS 2017) instead of Laplace framework could be illuminating. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject