Predictive auxiliary objectives in deep RL mimic learning in the brain
Authors: Ching Fang
Kim Stachenfeld
Venue: ICLR 2024 oral (2024)
Keywords: hippocampus
neuroscience
cognitive science
deep reinforcement learning
representation learning
prediction
Primary area: applications to neuroscience & cognitive science

Abstract: The ability to predict upcoming events has been hypothesized to comprise a key aspect of natural and machine cognition. This is supported by trends in deep reinforcement learning (RL), where self-supervised auxiliary objectives such as prediction are widely used to support representation learning and improve task performance. Here, we study the effects predictive auxiliary objectives have on representation learning across different modules of an RL system and how these mimic representational changes observed in the brain. We find that predictive objectives improve and stabilize learning particularly in resource-limited architectures, and we identify settings where longer predictive horizons better support representational transfer. Furthermore, we find that representational changes in this RL system bear a striking resemblance to changes in neural activity observed in the brain across various experiments. Specifically, we draw a connection between the auxiliary predictive model of the RL system and hippocampus, an area thought to learn a predictive model to support memory-guided behavior. We also connect the encoder network and the value learning network of the RL system to visual cortex and striatum in the brain, respectively. This work demonstrates how representation learning in deep RL systems can provide an interpretable framework for modeling multi-region interactions in the brain. The deep RL perspective taken here also suggests an additional role of the hippocampus in the brain-- that of an auxiliary learning system that benefits representation learning in other regions.

Under review as a conference paper at ICLR 2023
PREDICTIVE AUXILIARY OBJECTIVES IN DEEP RL MIMIC LEARNING IN THE BRAIN
Anonymous authors Paper under double-blind review
ABSTRACT
The ability to predict upcoming events has been hypothesized to comprise a key aspect of natural and machine cognition. This is supported by trends in deep reinforcement learning (RL), where self-supervised auxiliary objectives such as prediction are widely used support representation learning and improve task performance. Here, we study the effects predictive auxiliary objectives have on representation learning across different modules of an RL system and how these mimic representational changes observed in the brain. We find that predictive objectives improve and stabilize learning particularly tectures. We identify settings where longer predictive horizons better support rep- resentational transfer. Furthermore, we find that representational changes in this RL system bear striking resemblance to changes in neural activity observed in the brain across various experiments. Specifically, we draw a connection between the auxiliary predictive model of the system and hippocampus, to learn predictive model to support memory-guider behavior. We also connect the encoder network and the value learning network of the RL system to visual cortex and striatum in the brain, respectively. This work demonstrates how repre- sentation learning in deep RL systems can provide an interpretable framework for modeling multi-region interactions in the brain. The deep RL perspective taken here also suggests an additional role of the hippocampus in the brain- that of an auxiliary learning system that benefits representation learning in other regions.
resource-limite archi-
thought
INTRODUCTION
Deep reinforcement learning (RI) models have shown remarkable success solving challenging prob- lems (Sutton & Barto, 2018; Mnih et al., 2013; Silver et al., 2016; Schulman et al., 2017). These models use neural networks to learn state representations that support complex value functions. A key challenge in this setting is to avoid degenerate representations that support only subpar policies or fail to transfer to related tasks. Self-supervised auxiliary objectives, particularly predictive ob- jectives, have been shown regularize learning in neural networks prevent overfit or collapsed representations (Lyle et al., 2021; Dabney et al., 2021; François-Lavet et al., 2019). As such, is common to combine deep RL objectives with auxiliary objectives. The modular structure of these multi-objective models can function as metaphor for how different regions of the brain combine to comprise expressive, generalizable learning system.
Analogies can readily be drawn between the components of a deep RL system augmented with pre- dictive objectives and neural counterparts. For instance, the striatum has been identified as a RL-like value learning system (Schultz et al., 1997). Hippocampus has been linked to learning predictive models and cognitive maps (Mehta al., 1997; O' Keefe & Nadel, 1978; Koene al., 2003). Fi- nally, sensory cortex has been suggested to undergo unsupervised or self-supervised learning akin to feature learning (Zhuang et al., 2021), although reward-selective tuning also been observed (Poort et al., 2015). It is unclear how value learning, predictive objectives, and feature learning mutually interact to shape representations. Comparing representations across artificial and biological neu- ral networks can provide useful frame of reference for understanding the extent artificial models resemble the brain's mechanisms for robust and flexible learning.
These comparisons also provide useful insights into neuroscience, where little is known about how learning in one region might drive representational changes across the brain. For instance, the hippocampus is likely candidate for predictive objectives, as ample experimental evidence has

---

Under review as a conference paper at ICLR 2023
shown that activity in this region is predictive of the upcoming experience of animal (Skaggs & McNaughton, 1996; Lisman & Redish, 2009; Mehta et al., 1997: Payne et al.. 2021; Muller & Ku- bie, 1989; Pfeiffer & Foster, 2013; Schapiro et al., 2016; Blum & Abbott, 1996; Mchta et al., 2000). These observations are often accounted for in theoretical work as hippocampus computing a predic- tive model map (Lisman & Redish, 2009; Mehta al., 2000; Russek al., 2017; Whittington et al., 2020; Momennejad, 2020; George et al., 2021; Stachenfeld et al., 2017). Much has been writ- ten about how learned predictive models may be used by the brain to simulate different outcomes or support planning (Vikbladh et al., 2019; Geerts et al., 2020; Mattar & Daw, 2018; Miller et al., 2017: Olafsdóttir et al., 2018; Redish, 2016; Koene et al., 2003; Foster & Knierim, 2012; McNamee et al., 2021). However, the context of deep RL. the mere act of learning to make predictions in one region confers substantial benefits to other interconnected regions by shaping representations to in- corporate predictive information (Hamrick et al., 2020; Oord et al., 2018; Bengio, 2012). One of the key insights of this work is to propose that an additional role of predictive learning in hippocampus is to drive representation learning that supports deep RL in the brain.
The main contribution of this paper is to quantify how representations in deep RL model change with predictive auxiliary objectives, and to identify how these changes mimic representational changes in the brain. We first characterize functional benefits this auxiliary system confers learn- ing. We evaluate the effects of predictive auxiliary objectives in a simple gridworld foraging task, and confirm that these objectives help prevent representational collapse, particularly in resource- limited networks. We also observe that longer-horizon predictive objectives are more useful than shorter ones for transfer learning. We further demonstrate that a deep RL model with multiple ob- jectives undergo a variety of representational phenomena also observed in neural populations in the brain. Downstream objectives can alter activity in the encoder, which is mirrored in various results that show how visual cortical activity is altered by both predictive and value learning. Additionally, learning in the prediction module drives activity patterns consistent with activity measured in hip- pocampus. Overall we find that interacting objectives explain diverse effects in the neural data not well modeled by considering learning systems isolation. Moreover, this suggests that deep RI, with predictive objectives appears to many ways mirror the brain's approach to learning.
2 RELATED WORK
In deep RL, auxiliary objectives have emerged as crucial tool for representation learning. These additional objectives require internal representations to support other learning goals besides the pri- mary task of value learning. Auxiliary objectives thus regularize internal representations to preserve information that may be relevant for learning. They are thought to address challenges that may arise in sparse reward environments, such as representation collapse and value overfitting (Lyle et al., 2021). Many auxiliary objectives used in machine learning are predictive in flavor. Prior work has found success in defining objectives to predict reward (Jaderberg al., 2016; Shelhamer et al., 2016) or to predict future states (Shelhamer et al., 2016; Oord et al., 2018; Wayne et al., 2018) from his- tory. Predictive objectives may be useful for additional functions well. Intrinsic rewards based on the agent's 's ability to predict the next state can be used to guide curiosity-driven exploration (Pathak al., 2017: Tao et al., 2020). These objectives may also aid with transfer learning (Walker al., 2023). by learning representations that capture features that generalize across diverse domains. The incorporation of auxiliary objectives has greatly enhanced the efficiency and robustness of deep RL models in machine learning applications.
In neuroscience, much theoretical work has sought to characterize brain regions by the computa- tional objective they may be responsible for. Hippocampus in particular has been suggested to learn predictions of an animal's upcoming experience. This has been formalized as learning a transition model similar to model-based reinforcement learning (Fang et al., 2022) to learning long-horizon predictions as in the successor representation (Gershman et al., 2012; Stachenfeld et al., 2017). Sep- arately, the striatum has long been suggested to support model-free (MF) reinforcement learning like actor-critic models (Joel et al., 2002), with more recent work connecting these hypotheses to deep RL settings (Dabney et al., 2021: Lindsey & Litwin- Litwin-Kumar, 2022).
Less work has been done to understand how the computational objectives of multiple brain regions interact, although this has been suggested as a framework for neuroscience (Marblestone et al., 2016; Yamins & DiCarlo, 2016; Botvinick et al., 2020). Prior work has used multi-region recurrent neural networks (Pinto et al., 2019; Andalman et al., 2019; Kleinman et al., 2021) switching nonlinear
2

---

Under review as a conference paper at ICLR 2023
A.
Deaeart
obsecoation
Environmett
Figure 1: A deep RL framework model multi- region computation. A. In the deep RL model we use, reward is provided as a scalar input r. Observations are 2D visual inputs fed into an encoder (green) that learns low-dimensional state space representations The encoder is a convolu- tional neural network. Representations 2 are used to learn Q values via a MLP (blue); these Q values are used to select actions a. A predictive auxiliary objective (orange) is enforced by separate MLP learning predictions from 2.
+1
Enooder
atllon
+ Weet state
senetry cotle
Nepocarpua
Modei-Free Q Learner
Auxiliary Prediclive Model
dynamical systems (Semedo et al., 2014; Glaser et al., 2020; Tammbor al., 2022) to model the interactions of different regions. However, much of this work focuses more on fitting recorded neural activity than taking a normative perspective on brain function. A growing body of work considers modular and multi-objective approaches to building integrative models of brain function. One approach has been to construct multi-region models by combining modules performing inde- pendent computations and comparing representations in these models to neural activity (Frank & Claus, 2006; O Reilly & Frank, 2006; Geerts et al., 2020; Russo et al., 2020; Liu et al., 2023; Jensen al., 2023). On the behavioral end, there has also been prior work discussing how the addition of biologically
regularizers or auxiliary objectives can result in performance more consistent with humans (Kumar et al., 2022; Binz& & Schulz, 2022; Jensen et al., 2023).
Our work differs in that the entire system consists of a neural network that is trained end-to-end, allowing us the opportunity specifically study the effects on representation learning. In this paper, we show how deep RL networks can be testbed for studying representational changes and serve as multi-region model for neuroscience.
3 EXPERIMENTAL METHODS
Network architecture We implement double deep Q-learning network (Van Hasselt al., 2016) with a predictive auxiliary objective, similar to François-Lavet et al. (2019) (Fig IA). A deep con- volutional neural network E' encodes observation o, at time into latent state z, (o, will be a 2D image depicting the agent state in tabular grid world). The state 2 is used by two network heads: a Q-learning network Q(z, a) that will be used to select action a and a prediction network T(z, a) that predicts future latent states. Both Q and T are multi-layer perceptrons with one hidden layer.
Network training procedure The agent trained on transitions (ot, at, Oc+1, at+1) sampled from a random replay buffer. We will also let o; and 0j denote any two observations randomly sam- pled from the replay buffer that may not have occurred in sequence. The weights of E, Q. T are trained end-to-end to minimize the standard double Q-learning temporal difference loss function Lo Van Hasselt et al. (2016) and a predictive auxiliary loss Lpred. The predictive auxiliary loss is simi- lar to that of contrastive predictive coding (Oord et al., 2018). That is. Lpred L+ L_ where 4+ is a positive sampling loss and C_ is a negative sampling loss. The positive sample loss is defined as L4 = - 0x) 2++1 7T (2++1, (1+++111,++111 1++1)]12 where 2; F((o)) and T(=,, a.) = + T T(2,, a:). That is, in the y = O case, the network is learning the difference between current and future latent states such that 7 T(z,, ar) = 2₁ + T T(2,, az) x 2411. This encourages the learned representations to be structured to be structured so as to be consistent with predictable transitions (François-Lavet al., 2019). Additionally, 7 modulates the predictive horizon.
The negative sample loss is defined as L_ = - exp || - =1yl-- We emphasize that and zj are randomly sampled from the buffer and thus may represent states that are spatially far from another. This loss drives temporally distant observations to be represented differently, thereby preventing the trivial solution from being learned (mapping all latent states to a single point). The use of two contrasting terms (C_ and L+) is not just useful for optimization reasons- it also mirrors the hypoth- esized pattern separation and pattern completion within the hippocampus (O (O'Reilly & McClelland, 1994; Schapiro et al., 2017). However, we note that negative sampling elements are not always needed to support seef-predictive learning if certain conditions are satisfied (Tang et al., 2023). Ex- cept where indicated, the agent learns off-policy via random policy during learning, only using its policy during test time. The weights over loss terms Lo L L_ are chosen through a small grid search over the final episode score.
3

---

Under review as a conference paper at ICLR 2023
A. Gridworld Foraging
B.
Seed
Seed
Seed
Episode Score
1.0
X
Observation
0.5
MF only
Het
C_
Seed 2
Seed
Epred
0.0
250
500
Training Steps
Latent Size
Episode Score
1.0
Fraction Silent
1.0
Inner Product
1.0
0.5
0.5
G-Network Depth
Deoth
0.0
0.0
0.0
5
10
10
500
Latent Size
Latent Size
Training Steps
Figure 2: Gridworld performance with predictive auxiliary tasks. A. The model is tested on grid- world task in 8x8 arena. The agent must navigate hidden reward given random initial starting locations. B. Average episode score across training steps for models without auxiliary losses (blue), with only the negative sampling loss C_ (green), and with the full predictive loss Lpred (orange). The maximum score is and |=| = 10 (i.e. contains 10 units). In each step, the network is trained on one batch replayed transitions (batch size is 64). All error bars are standard error mean over 45 random seeds. C. 3D PCA representations of latent states for the models in (B) (two random seeds). The latent states are colored by the quadrant the arena they lie in. The quadrants (in order) purple, pink, gray, brown. The goal location state is colored red. Gray lines represent the true connectivity between states. D. Diagram of the encoder network (red), learned latent state (gray), and value-learning network (blue). We vary |2| (see E, F), as well the encoder/decoder depths (Appendix A.3AB). E. Average episode score at the end of learning (600 training steps) across |z]- F. Fraction of units in that are silent during the task, across |z]. G. Cosine similarity of two ran- domly sampled states throughout learning, |2| = 10.
Experimental comparisons and modifications We will treat the encoder network as sensory cortex analog, the Q-learning network as striatum analog, and the prediction network as hip- pocampus analog (Fig A.IAB). In our analyses, vary several parameters of interest. We vary the size of z to test the effects of the information bottleneck the encoder. We will also modulate the strength of r in the auxiliary loss to test the effects of different timescales of prediction. Finally, also test how the depths of the decoder and encoder networks affect learning.
RESULTS
4.1 PREDICTIVE OBJECTIVES HELP PREVENT REPRESENTATIONA, COLLAPSE.
We first understand the effect predictive auxiliary objectives have on learning system. We test the RL model in simple gridworld foraging task, where an agent must navigate to a hidden reward from any point in 2D arena. The observation received by the agent a 2D image depicting birds-eye view of the agent's ' location. Further details and examples are provided in Figure A.2 A-D. We compare model without auxiliary objectives (MF-only) models with the negative sampling objective C_ only and with the full predictive objective Lpred- Here, the predictive model is trained with one-step prediction (7- = 0).
Given sufficient capacity in the encoder, decoder, and latent layer all models learn the foraging task (Fig 2B). However, the model with prediction reaches maximum performance with fewer train- ing steps than both the negative-sampling model and the MF-only agent (Fig 2B). Additionally, the latent representation in the predictive model appears to capture the global structure the environ- ment better than the other models (Fig 2C). The model without any auxiliary tasks tends to expand the representation space around rewarding states, while the model with negative sampling (Fig 2C) evenly spaces apart state representations without regard for environment structure.
4

---

Under review as a conference paper at ICLR 2023
We next tested how the effects of auxiliary tasks change with the size of the model components (Fig 2D). We first varied the size of 2, and thus the representational capacity of the encoder. We find that, although all models can perform well given large enough latent dimension |2|, supplying the model with a predictive auxiliary objective allows the model to learn the task even with a smaller bottleneck (Fig 2E). This benefit is not conveyed by the negative sampling loss alone, suggesting that learning the environment structure confers its own unique benefit (Fig 2E). We find similar results by varying the encoder network depth and the decoder network depth (Fig A.3AB), showing that the cases.
benefits of predictive auxiliary objectives are more salient in resource-limine
This difference may be because representational collapse is a greater danger in lower-dimensional settings. To test this, we measure how many units in the output of the encoder are involved in supporting the state representation. We find that greater proportion of units are completely silent in the MF-only encoder (Fig 2F), suggesting a less distributed representation. To more directly test for collapse, we measure how the cosine similarity between state representations change across learning. Although all models start with highly similar states, the models with auxiliary losses separate state representations across training more than the MF-only model does (Fig 2G).
Finally, we test more complex versions of this gridworld task to see how performance is affected. We find consistent results in a CIFAR version of this task (Fig A.2D), where models equipped with predictive auxiliary objective outperform the other two models tested (Fig A.3C). We also test a version of gridworld where the environment is less predictable- that is, transitions are no longer determinstic. We find that, the probability of stochastic transitions increse, the benefit of predictive auxiliary objectives vanish (Fig A.3D).
4.2 LONG-IIORIZON PREDICTIVE AUXILIARY TASKS ARE MORE EFFECTIVE AT SUPPORTING REPRESENTATIONAL TRANSFER THAN SHORT- HORIZON PREDICTIVE TASKS.
Thus far, we have tested the predictive auxiliary objective with one-step prediction. However, long horizon predictions are often used as auxiliary objectives (Oord et al., 2018; Hansen et al., 2019), and many neural systems, including hippocampus, have been hypothesized to perform long-horizon predictions (Brunec & Momennejad, 2022; Lee al., 2021). We next sought to understand under what conditions longer horizons of prediction in auxiliary objectives would be useful. In particular, we were interested in exploring how well learned representations could transfer to new tasks. We hypothesize that oong-horizon predictions (larger y in C+) can better capture global environment structure and thus learn representations that transfer better to tasks in similar environments.
We first test representation transfer to new reward locations in gridworld (Fig 3A). After the agent learns an initial goal location in task A, we freeze the encoder, move the goal to new state, and fine- tune the value network for task B. This allows to test how well the learned representation structure can support new value functions. We test models with Epred loss and y E (0.0, 0.25, 0.5, 0.8). We find that, although all models learn task A quickly, models with larger y learn task B more efficiently (Fig 3B). We test how this effect scales with latent sizes. Just having a predictive horizon longer than one timestep appears sufficient improve learning efficiency, with the effect stronger at larger latent sizes. (Fig 3C). The selective benefit of longer time horizons for transfer may explain the observation that regions of hippocampus with larger spatial scales appear to be preferentially active in novel environments (Fredes et al., 2021; Köhler et al., 2002; Poppenk et al., 2010).
We hypothesize that the difference in efficient transfer performance across the models may result from learning a latent structure that better reflects global structure. Long-horizon prediction may be better at smoothing across experience over many timesteps, thus capturing global environment structure better than short-horizon prediction and providing a larger benefit when latent representa- tions are higher dimensional. Indeed, models with smaller y values tend to learn more curved maps that preserve local, but not global, structure (Fig 3D). To quantify this effect, we measured the inner product between the states representing the corners of the environment. These are states that are maximally far from each other, and as such, representations that capture the environment structure accurately should separate these states from each other. We see that, across learning, models with larger y learn to separate corner states better (Fig 3E).
Predictive auxiliary objectives can also be disadvantageous under certain regimes. Predictive objec- tives shape latent representations to reflect transition structure. However, these learned representa- tions might not generalize well to new tasks where the transition structure or the policy changes. We
5

---

Under review as a conference paper at ICLR 2023
A.
B.
encoder frozen
C.
Earlyy Learning
New goal location
goa
goal 2
Episode Score
1.0
Episode Score
1.0
x
0.5
0.5
0.0
0
250
500
0
*
250
500
10
20
Training Steps
Training Steps
Latent Size
7=0
7/= 0.25
Y= 0.50
Corner States Inner Product
1.0
0.5
0.0
0
0.250.5
0.8
New transition structure
G.
encoder
H.
"Early Learning
Episode Score
1.0
Episode Score
1.0
0.8
0.5
+ C-
Epred
0.6
500
1000
0
*200
400
0.4 0.6 0.8 1.0
Training Steps Figure 3: Effects of predictive auxiliary objectives across transfer learning scenarios. A. We test goal transfer by moving the goal location to new state in task B. After training on task A, encoder weights are frozen and the value function is fine-tuned on task B. B. Average episode score across task A, then task B. All models shown use the predictive auxiliary loss, with the shade of each line corresponding to the magnitude of y in Lpred ( E {0.0, 0.25, 0.5, 0.8). |= 17). C. The episode score after 100 training steps for each of the models in (B), |2| is increased. All models achieve maximum performance in task A. 30 random seeds are run for each latent size. D. 3D PCA plots, for three models (y 0.0, 0.25, 0.5) with the same random seed. E. Pairwise cosine similarity values between the corner states the arena for the model shown in (B). F. We test transition transfer by shuffling the connectivity between all states in task B. Freezing and fine-tuning are the same in (A). G. Average episode score across task A, then task B. Here, |z| 17 and e = 0.4-greedy policy during learning. In green is the model with only L_ an auxiliary loss. H. Episode score after 150 training steps for the model with only L_ (green) versus the model with Cpred for 0.8. On the x-axis, the policy e used during training is varied, with e 1.0 corresponding to a fully random policy ((2 = 17, all models achieve maximum performance on task A).
Training Steps
Policy E
test this in a different transfer task, where reward location remains the same in task B, but the envi- ronment transition structure is scrambled (Fig 3F). Additionally, to test for effects of policy change across task A and B, we vary the portion of random actions taken in our c-greedy agent. Under this new transfer task with e = 0.4, we find performance in task B decreases for models with the predictive objective compared to a model with just the negative sampling loss. (Fig. 3G).
Indeed, as e gets smaller and the agent learns more from biased on-policy transition statistics, trans- fer performance task B accordingly suffers (Fig 3G.H). All models with predictive objectives do not perform as well in task B as a model with only negative sampling loss (Fig 3G.H).
4.3 EFFECTS OF VALUE LEARNING AND DISTORY-DEPENDENCE IN PREDICTION NETWORK RESEMBLE HIPPOCAMPAL ACTIVITY.
We next ask how well representations developed in the network can model representations found in neural activity. The output of our T network serves analog the hippocampus, region

---

Under review as a conference paper at ICLR 2023
A.
D.
H.
Trials
E
F
% Units
Epred
10
0.5
0.0
0.5
1.0
Place Field Shift
Cross-Condite Correlation
G
0.8
% Units
10
mendom
MF only
0.6
MF only
-0.5
0
0.5
Start
Middle
Decision Point
Place Field Location
(Relative to Reward)
Central Stem Location
Figure 4: Representational changes in the predictive model are similar to those observed in the hippocampus. A. 2D foraging experiments are simulated in the gridworld task from Fig 1-2. B. 2D receptive fields from top four T units (columns) sorted by spatial information score (Skaggs al., 1992). Three random seeds are shown (rows). The model uses Lpred and |= 10. White asterisk depicts reward. C. As in (B), but the model has no auxiliary objectives. D. Circular track experiments are simulated in a circular gridworld with 28 states. Reward is in a random state for each seed and the agent is rewarded for running clockwise to the reward. E. Receptive fields of two example units in the 7' network before (gray) and after (orange) learning. Histogram over the shift in receptive field peaks for units in 7 over 15 random seeds, where | 24. Positive values indicate shifts forward, and vice-versa for negative values. Black dotted line O. Median of the histogram is -0.034. G. Histogram over the location of receptive field peaks for units in (F), with location centered around the reward site. Random shuffle (gray) control was made by randomly shuffling the weights of the T network. Black dotted line at O. The model median is -0.06, while the random shuffle median is -0.02. H. We simulate a 5x5 alternating-T maze (see Appendix); center corridor in pink. I. Cosine similarity of T population vector responses in the center corridor under left-turn versus right-turn conditions. X-axis depicts location in the center corridor. Data from 20 random seeds. Shown is the model without auxiliary objectives (blue) and the model with Lpred (orange). T's is randomly initialized for the model without an auxiliary objective.
implicated in self- self-predictive learning. We first test whether the T network activity can capture a classic result in the hippocampal literature: formation of spatially local activity patterns, called place fields. We plot the spatial firing fields of individual T units our model trained on gridworld, and find 2D place fields expected (Fig 4B). We also find that the prevalence of these place fields is greatly reduced in models without predictive auxiliary tasks (Fig 4C).
Hippocampal place fields also undergo sxperience-dependen changes. We test for these effects in our model through 1D circular track experiments (Fig 4D). We find that place fields developed on the ID track will skew and shift backwards from the movement of the animal (Fig 4E.F). This is consistent with phenomena in rodent hippocampal data that have been attributed to predictive learning (Mehta al., 2000). We also find that the number of place fields across the linear track is more abundant close to the reward site (Fig 4G), another widely observed phenomena that is considered be a result of reward learning in hippocampus. Our results suggest that value learning in shared representations with other systems can result in reward-related effects in the hippocampus.
Finally, we test more complex form of pnperience-dependeen alternating T-maze task. In this task, animals alternate between two trial types: one where they run down a center corridor and turn left for reward, and another where they run down the same center corridor but turn right for reward. In these tasks, neural activity has been observed to "split" neu- rons fire differently in the center corridor across trial types despite the spatial details of the corridor remaining the same (Duvelle et al., 2023). Interestingly, the degree of splitting is greatest in the beginning of the corridor and also high at the end of the corridor, splitting least in the middle of the corridor (Duvelle al., 2023). To enable the agent perform this task, which requires remembering
neural activity by simulating an
7

---

Under review as a conference paper at ICLR 2023
A. Preferred
Nonrall
Stw aqposay
Fulty stend IT sao tile ancr
stev idfeance
Vatte
revard
Non- preferred
m
A
Snop
meprocues from U bicats 2010)
Response Response
1.0
0.0
Before Loariing
G.
0.5
Wftrr Looming
0.0
1
2 3
3
A(P-N)
-0.2
Proportion
0.5
Proportion
0.5
1.0
0.5
-0.4
MF only
0.0
0.0
0.0
Epred
-2
0
2
V
A
N/A
1
2
3
1
3
S.I.
Selectivity
Figure 5: Representational changes the encoder model resemble recordings from visual cortex. A. Example sequence structure in the preference swap task of Li & DiCarlo (2008; 2010), images num- bered by seqeunce location. B. Example changes in IT neuron response to preferred images (red) and non-preferred images (blue) across exposure to new image transitions. C. Responses of two example units from the model with ppred. Arrows indicate response profile before and after experi- encing swapped transitions. Red indicates the response to P1, P2, P3 states that were selected from the gridworld environment, while blue indicates the response to N1, N2, N3 states selected from the environment. D. Change in response difference between (PI,NI), NN), (P2, N2), and (P3, N3) over 10 units. Each unit is separate transition swap experiment. Shown is the model without any auxiliary objectives (blue) and the model with Lpred (orange). Asterisks indicate significance from a t-test comparing the means from both models. We additionally note that the means of both models significantly different from O. E. Linear track VR experiment used in Poort al. (2015). Vertical stripe corridors were rewarded but angled corridors were not. Animals experienced either condition random following approach corridor. F. Selectivity across the population before learning (gray) and after learning (orange). Selectivity was calculated as in Poort et al. (2015), with negative and positive values corresponding to angled and vertical corridor preference, respectively. Asterisks in- dicate significance from one-tailed t-test (t = - 12.43, 9 x 10c-36) 36) G. Selectivity of individual units before and after learning for vertical condition (V), angled condition (A), or neither (N/A). Units are pooled across 15 experiments.
the previous trial type, we introduce a memory component to the agent so that temporally graded trace of previous observations are made available. That is, the input into the encoder at time is Or + -::-1 + a²o,-2 + ... for some < 1. This decaying sum of recent observations captures information about the recent past in a simple way, and inspired by representations hypothesized by temporal context model (Howard & Kahana, 2002). We measure cosine similarity between pop- ulation activity in the left turn condition and the right turn condition. Lower similarity corresponds to greater splitting. The representations in both MF-only model and the model with the predictive objective show increased splitting in the beginning of the corridor due to the memory component (Fig 4F). However, only the model with the predictive objective shows increased splitting at the end of the corridor (Fig 4F). This shows that the pattern of splitting seen in data can be captured by model using both memory and prediction.
We also test the effects of recurrency in the model by simulating a partially observable version of the alternating-T maze (Fig A.SC). To solve this version of the task, recurrency must be used to infer the current latent state with the model's previous latent state. We find consistent results where the inclusion of a predictive auxiliary objective greatly improves the model's s ability to learn the task (Fig A.5DE) and where only the model with the predictive objective shows splitting pattern consistent with data (Fig A.5F).
4.4 EFFECTS OF VALUE LEARNING AND TRANSITION LEARNING IN THE ENCODER NETWORK RESEMBLE ACTIVITY IN VISUAL CORTEX.
As another example of representational effects arising from mutually interacting regions, we com- pare the activity of our encoder network experimental results in sensory cortices. Neurons in
8

---

Under review as a conference paper at ICLR 2023
visual cortex (even those in primary regions) have been observed to change their tuning result of learning Poort et al. (2015); Li & DiCarlo (2008; 2010); Wilmes & Clopath (2019); Pakan et al. (2018). Our model provides a simple system to look for such effects.
First, we test for effects of prediction and temporal statistics that have been seen in visual cortex. Specifically. Li & DiCarlo (2008) found that object selectivity in macaque IT neurons could be al- tered by exposing animals to sequences of images where preferred stimuli and non- non-preferred stimuli became linked (Fig 5A). The images in the preferred and non-preferred that are linked together are referred to as the "swap position" within a sequence (Fig 5A). An analogous experiment can be run in our gridworld task from Fig 2. We first identify spatially contiguous preferred and non-preferred states of neurons in the encoder network. We then expose the model to sequences where preferred states and non-preferred states became connected at arbitrarily chosen swap positions (Fig 5B). We find neurons in the output of the encoder that, after exposure, decrease their firing rate for the pre- ferred stimulus at the swap location and increase their firing rate for the non-pre non-preferred stimulus at the swap position (Fig 5C). This is consistent with observations in data as well (Li & DiCarlo, 2008; 2010). We quantify this change in firing rate at different sequence locations. We find a similar trend as in data, where tuning for stimuli closer to the swap position is increasingly altered away from the original preferences (Fig 5D). Importantly, this effect is not present without the predictive auxiliary objective, similar to lesion studies carried out in Finnie et al. (2021).
The downstream Q- objective also have an effect on representations in the encoder. We simulate value learning effects in visual cortical activity through linear track experiments used in Poort et al. (2015) (Fig SE). In this experiment, authors found that VI neurons in mice increased selectivity for visual cues in the environment after learning the task. Furthermore, the authors noted a slight selectivity increase for more rewarding cues (vertical gratings) compared to nonrewarding cues (angled gratings). We find a similar effect in units in early layers of the encoder network: a small, but statistically significant increase in proportion of units encoding the rewarded stimulus (Fig 5F). As in Poort et al. (2015), selectivity increases across learning, but with a greater preference for the vertical grating environment (Fig 5G).
5 CONCLUSION
In this work, we explore the representational effects induced by predictive auxiliary objectives. We show how such objectives are useful in resource- settings and in certain transfer learning settings. We also investigate how prediction and predictive horizons affect learned representation structure. Furthermore, we describe how such deep RL models can function as a multi-region model for neuroscience. We show how representation learning in the prediction model recapitulates experi- mental observations made in the hippocampus. We make similar connections between representation learning in the encoder model and learning in visual cortex.
Our results point to a new perspective on the role of the hippocampus in learning. That is, a predic- tive system like the hippocampus can be useful for learning without being used to generate sequences or support planning. Learning predictions is sufficient to induce useful structure into representations used by other regions. This view also connects to trends seen in machine learning literature. In deep RL, predictive models need not be used for forward planning (Hamrick et al., 2020) to be useful for representation learning. Additionally, the contrastive prediction objective used in this work is drawn from machine learning literature but bears interesting similarities to classic descriptions of hippocampal computation. CA3 and CAI in the hippocampus have been implicated in predictive learning similar to the positive sampling loss. Meanwhile, the dentate gyrus in the hippocampus has been proposed to perform pattern separation similar to the contrastive negative sampling loss.
Our results are limited in the complexity of tasks and the diversity of auxiliary objectives tested. Future work can improve on current understanding by more systematically comparing effects across objectives over more complex tasks. We also did not examine representations in the value learning network, which ripe for comparison with striatum data. Future work can also explore the effects of recurrence across modules, which can be both functionally useful and more biologically realistic.
Overall, this work points to the utility of a modeling approach that considers the effect of multiple objectives in deep learning system. The deep network setting reveals new aspects of neuroscience modeling that are less apparent in tabular settings or in simpler models.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Aaron S Andalman, Vanessa M Burns, Matthew Lovett-Barron, Michael Broxton, Ben Poole, Samuel J Yang. Logan Grosenick, Talia Lerner, Ritchie Chen, Tyler Benster, al. Neuronal dynamics regulating brain and behavioral state transitions. Cell, 177(4):970-96) 2019.
Yoshua Bengio. Deep learning of representations for unsupervised and transfer learning. In Pro- ceedings of ICML workshop unsupervised and transfer learning, pp. 17-36. JMLR Workshop and Conference Proceedings, 2012.
Marcel Binz and Eric Schulz. Modeling human exploration through resource-rational reinforcement learning. Advances in Neural Information Processing Systems, 35:31755-31 2022.
Kenneth I Blum and Larry F Abbott. A model of spatial map formation in the hippocampus of the rat. Neural computation, 8(1):85-93, 1996.
Matthew Botvinick, Jane X Wang, Will Dabney, Kevin Miller, and Zeb Kurth-Nelson. Deep reinforcement learning and its neuroscientific implications. Neuron, 107(4):603-611 616, 2020.
Iva K. Brunec and Ida Momennejad. Predictive representations in hippocampal and prefrontal hierarchies. Journal of Neuroscience, 42(2):299-312, 2022. ISSN 0270-6474. doi: 10. 152/UNEUUROSC..132-2110200000 URL https: //wp:///wwwwww// www.:/jjeurssi.commm htw.rrocc...cgg//ooee/t///// 2/299.
Cathrin B Canto, Floris G Wouterlood, Menno P Witter, et al. What does the anatomical organization of the entorhinal cortex tell us? Neural plasticity, 2008, 2008.
Will Dabney, André Barreto, Mark Rowland, Robert Dadashi, John Quan, Mare Bellemare, and David Silver. The value-improvement path: Towards better representations for reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 7160- 7168, 2021.
éléónore Duvelle, Roddy M Grieves, and Matthijs van der Meer. Temporal context and latent state inference in the hippocampal splitter signal. Elife, 12:c82357, 2023.
Ching Fang, Dmitriy Aronov, Larry Abbott, and Emily Mackevicius. Neural learning rules for generating flexible predictions and computing the successor representation. bioRxiv, pp. 2022-05, 2022.
Peter SB Finnie, Robert W Komorowski, and Mark Bear. The spatiotemporal organization of ex- perience dictates hippocampal involvement in primary visual cortical plasticity. Current Biology, 31(18):3996-4008 2021.
David J Foster and James J Knierim. Sequence learning and the role of the hippocampus in rodent navigation. Current opinion in neurobiology, 22(2):294-300, 2012.
Vincent François-Lavet, Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 3582-3589, 2019.
Michael J Frank and Eric D Claus. Anatomy of decision: striato-orbitofrontall interactions in reinforcement learning, decision making, and reversal. Psychological review, 113(2):300, 2006.
Felipe Fredes, Maria Alejandra Silva, Peter Koppensteiner, Kenta Kobayashi, Maximilian Joesch, and Ryuichi Shigemoto. Ventro-dorsal hippocampal pathway gates novelty-induced contextual memory formation. Current Biology, 31(1):25-38.e5 2021. ISSN 0960-9822. doi: https://doi. mtg/l0...01//.ccu22200.....2.4 URL https: : / //wp://wwwww//... www.ncedireete.....ccccceeeeee.e com scinnceeeee////// aar.ccle/pi1/00009922033.33
Jesse P Geerts, Fabian Chersi, Kimberly L Stachenfeld, and Neil Burgess. A general model of hip- pocampal and dorsal striatal learning and decision making. Proceedings of the National Academy of Sciences, 1 11749))314273314333 2020.
10

---

Under review as a conference paper at ICLR 2023
Dileep George, Rajeev V Rikhye, Nishad Gothoskar, Swaroop Guntupalli, Antoine Dedieu, and Miguel dizaro-Gredille Clone-structured graph representations enable flexible learning and vi- carious evaluation of cognitive maps. Nature communications, 12(1):1-17, 2021.
Samuel Gershman, Christopher D Moore, Michael Todd, Kenneth A Norman, and Per B Seder- berg. The successor representation and temporal context. Neural Computation, 24(6):1553-1565 2012.
Joshua Glaser, Matthew Whiteway, John P Cunningham, Liam Paninski, and Scott Linderman. Re- current switching dynamical systems models for multiple interacting neural populations. Ad- vances in neural information processing systems, 33:1433:111 18677-1488, 2020.
Sarah Goodroe, Jon Starnes, and Thackery Brown. The complex nature of hippocampal-striitall interactions in spatial navigation. Frontiers in human neuroscience, 12:250, 2018.
Jessica Hamrick, Abram Friesen, Feryal Behbahani, Arthur Guez, Fabio Viola, Sims Wither- spoon, Thomas Anthony, Lars Buesing, Petar Veličkovié, and Théophane Weber. On the role of planning in model-based deep reinforcement learning. arXiv preprint arXiv:2011.04022 2020.
Steven Hansen, Will Dabney, André Barreto, Tom Van de Wiele, David Warde-Farley and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. CoRR, 00//190600003, 2019. URL /ttttpp//aaxte wwwpr/iii..oo.vooooooo.oo www.ooggg/////////o///o.o/..ee/e 906. 05030.
Mare W. Howard and Michael J. Kahana. A distributed representation of temporal context. Journal of Mathematical Psychology, 46(3):269-299 2002. ISSN 0022-2496. doi: https://
1ttg://0.1000/mpp22001....111 URL / /wwp//wwwww.//////////////. www.eceeieeccct.....c....... www...oooooooooo...oo....... com/science// articlle/p11/$00224461111
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05333 2016.
Kristopher T Jensen, Guillaume Hennequin, and Marcelo G Mattar. A recurrent network model of planning explains hippocampal replay and human behavior. bioRxiv, pp. 2023-01, 2023.
Daphna Joel, Yael Niv, and Eytan Ruppin. Actor-critic models of the basal ganglia: New anatomical and computational perspectives. Neural networks, 15(4-6):535-55) 5473 2002.
Orren Karniol-Tanin Tambour, David M Zoltowski. E Mika Diamanti, Lucas Pinto, David W Tank, Car- los Brody, and Jonathan Pillow. Modeling communication and switching nonlinear dynam- ics in multi-region neural activity. bioRxiv, pp. 2022 09, 2022.
Michael Kleinman, Chandramouli Chandrasekaran, and Jonathan Kao. mechanistic multi-area recurrent network model of decision- decision-making making. Advances in neural information processing systems, 14:23152-23111 2021.
Randal A Koene, Anatoli Gorchetchnikov, Robert C Cannon, and Michael E Hasselmo. Model- ing goal-directed spatial navigation the rat based on physiological data from the hippocampal formation. Neural Networks, 16(5-6):577-588 2003.
Stefan Köhler, Joelle Crane, and Brenda Milner. Differential contributions of the parahippocampal place area and the anterior hippocampus to human memory for scenes. Hippocampus, 12(6): 718-723, 2002.
Sreejan Kumar, Carlos G Correa, Ishita Dasgupta, Raja Marjieh, Michael Y Hu, Robert Hawkins, Jonathan D Cohen, Karthik Narasimhan, Tom Griffiths, et al. Using natural language and pro- gram abstractions instill human inductive biases in machines. Advances in Neural Information Processing Systems, 15:167, 2022.
Caroline Lee, Mariam Aly, and Christopher Baldassano. Anticipation of temporally structured events in the brain. eLife, 10:e64972, apr 2021. ISSN 2050-084X. doi: 10.7554/eLife.6497. URL https::/p//. /tw///o.../i www.ooor///..... 7500...5554.... fe. 64972.

---

Under review as a conference paper at ICLR 2023
Nuo Li and James DiCarlo. Unsupervised natural experience rapidly alters invariant object repre- sentation in visual cortex. science, 321 1216 2008.
Nuo Li and James J DiCarlo. Unsupervised natural visual experience rapidly reshapes size-invanne invariint object representation inferior temporal cortex. Neuron, 6766):1062-1775 2010.
Jack Lindsey and Ashok Kumar. Action-modulas midbrain dopamine activity arises from distributed control policies. Advances in Neural Information Processing Systems, 35:5535-5548, 2022.
John Lisman and A David Redish. Prediction, sequences and the hippocampus. Philosophical Transactions of the Royal Society B: Biological Sciences, (64((621):) 193-1111, 2009..
Ziming Liu, Mikail Khona, Ila R Fiete, and Max Tegmark. Growing brains: Co-emergence of anatomical and functional modularity in recurrent neural networks. arXXiv preprint arXiv:2310.0771, 2023.
Clare Lyle, Mark Rowland, Georg Ostrovski, and Will Dabney. On the effect of auxiliary tasks representation dynamics. In International Conference on Artificial Intelligence and Statistics, pp. 1-9. PMLR, 2021.
Adam H Marblestone, Greg Wayne, and Konrad P Kording. Toward an integration of deep learning and neuroscience. Frontiers in computational neuroscience, 10:94, 2016.
Marcelo G Mattar and Nathaniel D Daw. Prioritized memory access explains planning and hip-
pocampal replay. Nature neuroscience, 21(11):1609-1611 2018.
Daniel C McNamee, Kimberly L Stachenfeld, Matthew M Botvinick, and Samuel Gershman. Flexible modulation of sequence generation in the pooorhinal-iipoccamppp system. Nature neu- doi: 10. 10.10088611592-2--00003 URL http:::tp////c....t.m /europepmc..eceee org/ aataaatiicc http//rrccceeee/Mc 610914,
roscience, 24(6):851 862, June 2021. ISSN 1097-6256.
Mayank R Mehta, Carol A Barnes, and Bruce L. McNaughton. Experience-dependenn asymmetric expansion of hippocampal place fields. Proceedings of the National Academy of Sciences, 94(16): 8918-8921, 1997.
Mayank R Mehta, Michael C Quirk, and Matthew A Wilson. Experience-dependent asymmetric shape of hippocampal receptive fields. Neuron, 25(3):707-715, 2000.
Kevin J Miller, Matthew Botvinick, and Carlos Brody, Dorsal hippocampus contributes to model-based planning, Nature neuroscience, 1299-1266 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Ida Momennejad. Learning structures: predictive representations, replay, and generalization. Cur- rent Opinion in Behavioral Sciences, 32:155-166, 2020.
Nicolás A Morgenstern, Ana Filipa Isidro, Inbal Israely, and Rui M Costa. Pyramidal tract neurons drive amplification of excitatory inputs to striatum through cholinergic interneurons. Science Advances, 8()):eabh4315 2022.
Robert U. Muller and John L Kubie. The firing of hippocampal place cells predicts the future position of freely moving rats. In The Journal of neuroscience the official journal of the Society for Neuroxcience, 1989.
O'Keeee and L. Nadel. The hippocampus as a cognitive map. Clarendon Press, Oxford, United Kingdom, 1978.
Freyja Ólafsdöttir, Daniel Bush, and Caswell Barry. The role of hippocampal replay in memory and planning. Current Biology, 28(1):R37-,50 2018.
12

---

Under review as a conference paper at ICLR 2023
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv:1807.03744 2018.
Randall O'Reellly and Michael Frank. Making working memory work: a computational model of learning in the prefrontal cortex and basal ganglia. Neural computation, 18(2):283-328 2006.
Randall C O'Reilly and James I. McClelland. Hippocampal conjunctive encoding, storage, and recall: Avoiding a trade-off. Hippocampus, 4(6):661-682, 1994.
Janelle MP Pakan, Valerio Francioni, and Nathalie L Rochefort. Action and learning shape the activity of neuronal circuits in the visual contex. Current opinion in neurobiology, 52:88-97, 2018.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
HL Payne, GF Lynch, and D) Aronov. Neural representations of space in the hippocampus of food-caching bird. Science, 373(6552:343-344, 2021.
Brad Pfeiffer and David J Foster. Hippocampal place-cell sequences depict future paths to remem- bered goals. Nature, 497(7447):74-7, 2013.
Lucas Pinto, Kanaka Rajan, Brian DePasquale, Stephan Y Thiberge, David W Tank, and Carlos D Brody. Task deeeddee changes in the large-scale dynamics and necessity of cortical regions. Neuron, 104(4):810-85 824, 2019.
Jasper Poort, Adil G Khan, Marius Pachitariu, Abdellatif Nemri, Ivana Orsolic, Julija Krupic, Mar- ius Bauza, Maneesh Sahani, Georg B Keller, Thomas Mrsic-Flogel, et al. Learning enhances sensory and multiple non-sensory representations in primary visual cortex. Neuron, 86(6):1478- 1490, 2015.
Jordan Poppenk, Anthony R McIntosh, Fergus IM Craik, and Morris Moscovitch. Past experience modulates the neural mechanisms of episodic memory formation. Journal of Neuroscience, 30 (13):4707-4716, 2010.
David Redish. Vicarious trial and error. Nature Reviews Neuroscience, 17())147-159, 2016.
EM Russek, Momennejad MM Botvinick, SJ Gershman, and ND Daw. Predictive representations can link model-based reinforcement learning to model-free mechanisms. PLoS Comput Biol, 2017. doi: 10. l/journal.pcbi. 1005768.
100.50
Abigail A Russo, Ramin Khajeh, Sean R Bittner, Sean M Perkins, John P Cunningham, Laurence Abbott, and Mark M Churchland. Neural trajectories in the supplementary motor area and motor cortex exhibit distinct geometries, compatible with different classes of computation. Neuron, 107 (4):745-758, 2020.
Anna C Schapiro, Nicholas B Browne, Kenneth A Norman, and Matthew M Botvinick. Sta- tistical learning of temporal community structure in the hippocampus. Hippocampus, 26(1):3-8, 2016.
Anna C Schapiro, Nicholas Turk-Browne, Matthew M Botvinick, and Kenneth A Norman. Com- plementary learning systems within the hippocampus: neural network modelling approach to reconciling episodic memory with statistical learning. Philosophical Transactions of the Royal Society B: Biological Sciences, 372((711):201606))) 2017.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Wolfram Schultz, Peter Dayan, and P Read Montague. A neural substrate of prediction and reward. Science, 15935000))-133-15999 1997.
Joao Semedo, Amin Zandvakili, Adam Kohn, Christian K Machens, and Byron M Yu. Extracting latent structure from multiple interacting neural populations. Advances in neural information processing systems, 27, 2014.
13

---

Under review as a conference paper at ICLR 2023
Evan Shelhamer, Parsa Mahmoudich, Max Argus, and Trevor Darrell. Loss is its reward: Self- supervision for reinforcement learning. arXiv preprint arXiv: 2016.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittweeser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484448) 2016.
WE Skaggs and BL McNaughton. Replay of neuronal firing sequences in rat hippocampus during sleep following spatial experience. Science, 1996. doi: 10. 0t6o/science22155255..15
William Skaggs, Bruce Menaughton, and Katalin Gothard. An information-theomm approach to deciphering the hippocampal code. Advances in neural information processing systems, 5, 1992.
Kimberly Stachenfeld, Matthew Botvinick, and Samuel Gershman. The hippocampus predictive map. Nature Neuroscience, 2017. doi: 10.1038/mn. 4650..
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chan- dak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, et al. Understanding self- oeffprredictvv learning for reinforcement learning. In International Confer- ence on Machine Learning, pp- 13632-33656. PMLR, 2023.
Ruo Yu Tao, Vincent François-Lavet, and Joelle Pineau. Novelty search in representational space for sample efficient exploration. Advances in Neural Information Processing Systems, 33:8114-8122 2020.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
Oliver M Vikbladh, Michael R Meager, John King, Karen Blackmon, Orrin Devinsky, Daphna Shohamy, Neil Burgess, and Nathaniel Daw. Hippocampal contributions model-based plan- ning and spatial memory. Neuron, 102(())683-66) 2019.
Jacob C Walker, Eszter Vértes, Yazhe Li, Gabriel Dulac- Amnol, Ankesh Anand, Théophane Weber, and Jessica B Hamrick. Investigating the role of model-based learning exploration and transfer. In International Conference on Machine Learning, pp. 35368-35388 PMLR, 2023.
Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Leibo, Adam Santoro, al. Unsupervised predictive memory goal-directed agent. arXiv preprint arXiv:1803.10700 2018.
James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The tolman-eichenbaum machine: Unifying space and relational mem- ory through generalization in the hippocampal formation. Cell, 183(5):1249-1263 2020.
Katharina Anna Wilmes and Claudia Clopath. Inhibitory microcircuits for top-down plasticity of sensory representations. Nature communications, 10(1):5055 2019.
Daniel LK Yamins and James J DiCarlo. Using goal-driven deep learning models to understand sensory cortex. Nature neuroscience, 19(3):35-335,, 2016.
Chengxu Zhuang, Siming Yan, Aran Nayebi, Martin Schrimpf, Michael C. Frank, James J. Di- Carlo, and Daniel L. K. Yamins. Unsupervised neural network models of the ventral visual stream. Proceedings of the National Academy of Sciences, 118(3)::00141961118 2021. doi: 20073//nas.2014199 URL https: / /ttp///wwww/////////////////tt wtwppppsss...///aa//ooooe 2014196118.
1073 /paas.
14


Reviews:
[{"review_id": "0eeoZqmOGZ", "reviewer": "Reviewer_f2kS", "summary": "The authors use a multi-objective RL model that combines a Q-learning module with an auxiliary predictive objective to demonstrate advantages of predictive learning on representations for multi-regional interactions in neural networks.  Using foraging tasks on gridworld scenarios as an example, they demonstrate that using an auxiliary predictive module results faster training and smoother representations of the task environment.  They also demonstrate that increasing the time horizon for predictive learning produces networks that retrain faster on new tasks in the same environment.  Interestingly, the learned representation is crucial for this as scrambling the transition structure results in much slower learning. They then demonstrate that some of the representational changes observed in the model reflect similar changes observed in real neural networks.", "questions": "- Can you provide example environments?\n\n- What is the effect of varying the complexity of the environment?  \n\n\nSimilar to the above, the encoder seems to have all available information about the environment (in principle), so the predictive task is in some sense simpler than it otherwise might be for a real organism, which only has clues to its environment.  Do you have thoughts about how partial observability might affect the predictive module?  (This is beyond the scope of the paper, but might be worth speculating on.)", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The description of the methods and approach is fairly clear and the overall goals of the paper are clear, with some room for improvement.  The numerical experiments provided demonstrate how a predictive loss benefits the learned representations available in a downstream area (not necessarily directly related to the area responsible for prediction), without the \"predictive area\" necessarily providing any direct information to the area responsible for valuation and action selection.", "weaknesses": "One important feature that isn't clear from what is presented in the paper is how the environment itself may or may not affect these results.  I don't find any examples of what gridworld environments are being solved by these models, let alone how complex they are.  Surely the predictability of the environment itself has some bearing on the rate of learning and retrainability for novel tasks, not to mention the quality of representations?  Please provide examples of these environments.  If possible, please consider varying the complexity of the environments.  \n\nminor:\n\n- Bottom of page 3 \"the standard double deep Q-learning temporal difference loss function\":  Even though standard, please either provide the form of this loss or provide a reference.  \n\n- page 4, just below figure caption, definition of positive sample loss:  should o_{t+1} be z_{t+1}?\n\n- last sentence, first paragraph page 5:  \"the predictive model is trained with...\"  (No \"be\")\n\n- first sentence, section 4.2:  remove \"is used\" at the end of the sentence.\n\n- first sentence, last paragraph on page 7:  remove \"to\",I.e. \"undergo experience-dependent changes\".\n\n- last paragraph, page 8:  \"remembering the previous trial type\". (Remove \"whether\"), also empty reference at the end of that sentence.", "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": "gPBU2tRW95", "reviewer": "Reviewer_CL1C", "summary": "This paper implements a deep RL framework with predictive auxiliary objectives for representation learning. The authors demonstrate in a gridworld setting that predictive objectives improve representation learning by preventing representation collapse and enhancing transfer learning when transition structure remains unchanged. The paper also relates the components of the deep RL model to different brain regions, including the sensory cortex, the hippocampus, and the striatum. They show that representation learning in the predictive model resembles the neural activity observed in the brain, and learning in the encoder model resembles neural observations in the visual cortex.", "questions": "- I'm curious to learn if there is a similarity between the action selection network and the neural activity observed in the striatum\n- Figure 2: Were cells that didn't show place-like activities filtered out in these analyses?\n- Were there also place-like activities in the encoder model?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Overall I really enjoy reading this work, due to its clear presentation both in the text and in the figures, experiments testing different perspectives of the model, and the strong link to the brain\n- This work introduces a multi-region model that is developed from a normative perspective, instead of fitting to recorded data, which can be extended to other tasks and to test against new biological evidence\n- I appreciate the discussion of the limitation that predictive auxiliary objectives may be less helpful when the transition structure or policy changes", "weaknesses": "- It's interesting to see in section 4.4 where the authors describe the effects of value learning in the encoder network, but this part feels somewhat disconnected from the rest of the paper, as the primary focus is to demonstrate how predictive objectives can lead to representation changes similar to those seen in the brain", "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": "ijCuOYYyr4", "reviewer": "Reviewer_T1AU", "summary": "In this work, the authors explore modeling the effect of auxiliary predictive losses on a deep RL agent trained to navigate a spatial gridworld environment. The architecture consists a convolutional encoder that encodes the visual observation, a prediction layer that generates predictions for the auxiliary losses, and a Q-learning agent that takes the encoded state and returns Q-values/actions for the main reward objective. There are two auxiliary losses: a positive sampling loss, which is predicting future states, and negative sampling loss, which encourages representations of non-consecutive states to be distinct from each other. Multiple analyses inspired by neuroscience experiments were employed to show the value of these predictive losses: visualizing the latent state space showing that predictive losses help with representational collapse, showing how long horizon state predictions help with learning new goal locations, looking at how the predictive auxiliary objectives help in producing the \"splitting\" phenomenon, and showing how the encoder (through virtue of prediction) will adapt to novel transition sequences. \n\nOverall, I think this work is cool and exciting. But I do think portions of the submission have clarity that is below the caliber of an iclr paper. I have initially rated marginally below acceptance (5), but I want to assure the authors that most of my issues are with regards to providing clarity/details. I hope to see the incorporation of more clarity and details in the writing during the rebuttal and would be more than happy to re-evaluate my score after such changes.", "questions": "1. Why use an off-policy learner (Q-learning) rather than on-policy? An on-policy learner would seem to be more biologically realistic. Also I think an actor-critic approach (e.g. A2C, PPO, etc) may reflect what striatum is doing more than Q-learning? \n\n2. Is there a reason why there is no recurrence in the model? Striatum will indirectly project back to sensory cortex via thalamus. Also it may be possible that hippocampus projects back to sensory visual cortex. Hippocampal activity early in a trial can be predictive of information in visual cortex at a later time in the trial (see Hindy et al. 2016 Nature Neuroscience). Regardless of the biological realism of these recurrent connections, I think there could be useful normative principles in having recurrent loops between modules in this architecture. \n\n3. Fig 5c: The change in firing rates seem pretty small between before/after exposure. Are these changes statistically significant?\n\nReferences:\nHindy, N. C., Ng, F. Y., & Turk-Browne, N. B. (2016). Linking pattern completion in the hippocampus to predictive coding in visual cortex. [Nature neuroscience](https://www.nature.com/articles/nn.4284), 19(5), 665-667.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The experiments seem to be done soundly and rigorously. \n\n* The authors do an excellent job introducing various relevant neuroscience experiments and grounding the phenomenon into specific predictions for their model. \n\n* The authors use an exciting emerging framework of auxiliary losses to tackle an important problem, which is modeling multiple regions of the brain simultaneously while performing a difficult task.", "weaknesses": "The related works section is a little small/sparse. The authors do a good job in highlighting works on auxiliary predictive losses in RL within the machine learning realm, but I think there is also a growing body of work that is using this framework to produce various behavioral phenomenon in cognitive science/neuroscience. These are complimentary works to the current submission and would be good to include. Here are some examples that I feel should definitely be included:\n\n1. Kumar et al. 2022 NeurIPS use auxiliary predictive losses in RL agents to predict abstractions of the observation, operationalized through language and symbolic programs, in order to reproduce abstract human biases. \n2. Jensen et al. 2023 bioRxiv introduce a predictive auxiliary loss which helps the agent learn when to plan and reproduces replay patterns seen in rodent hippocampal work. \n3. This is not exactly auxiliary predictive loss, but I think it is having the same effect, Binz & Schulz 2022 NeurIPS show adding a regularization term to the loss on the number of bits required to compress the agent's weights reproduces quirks in human exploration.\n\nI think the clarity in some portions in this submission can be improved. \n\n1. On page 4, the description of the auxiliary losses can be improved. Its an important part of the work so it'd be good to have this be as clear as possible. A sentence describing what exactly $\\tau$ is representing and why  $\\mathcal{L_{+}}$ is a loss term that enforces transition structure in the state representations would be very helpful. Also, for the third term in $\\mathcal{L_{+}}$, shouldn't it be $\\tau (z_{t+1},a_{t+1})$ and not $\\tau (o_{t+1},a_{t+1})$? For the negative sampling loss term, making it explicit $z_{i}$ and $z_{j}$ are latent representations of states that are not consecutive in the same area where the loss term is introduced would be helpful. It'd also be good to have a sentence explaining the motivation for choosing these two specific auxiliary losses. I suspect its loosely inspired by pattern completion vs pattern seperation in hippocampus (respectively) but it would be good to confirm that in this section. \n\n2. It would be helpful to state what the colors in Figure 2D mean. It wasn't clear to me upon first read. \n\n3. The memory component in page 8 is not explained at all. It seems important to describe what this is to put Fig. 4's results in context. \n\nThere are a couple of design decisions that left me a little confused (details in the questions section). It would be nice for the authors to clarify the motivations behind them.\n\nLast, I don't think there was any section in the paper explicitly discussing the limitations of the work. I think this is an important part of any iclr paper so it'd be good to include one. \n\nReferences:\n1. Kumar, S., Correa, C. G., Dasgupta, I., Marjieh, R., Hu, M. Y., Hawkins, R., ... & Griffiths, T. (2022). Using natural language and program abstractions to instill human inductive biases in machines. [Advances in Neural Information Processing Systems](https://arxiv.org/abs/2205.11558), 35, 167-180.\n2. Jensen, K. T., Hennequin, G., & Mattar, M. G. (2023). A recurrent network model of planning explains hippocampal replay and human behavior. [bioRxiv](https://www.biorxiv.org/content/10.1101/2023.01.16.523429v1.abstract), 2023-01.\n3. Binz, M., & Schulz, E. (2022). Modeling human exploration through resource-rational reinforcement learning. [Advances in Neural Information Processing Systems](https://arxiv.org/abs/2201.11817), 35, 31755-31768.", "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 (oral)