Paper: Predictive auxiliary objectives in deep RL mimic learning in the brain (ICLR-2024)

Reviewer comment (Novelty):
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 there is also a growing body of work using this framework to produce various behavioral phenomena in cognitive science/neuroscience. These are complementary works to the current submission and would be good to include.

Author response:
Thank you for all the suggestions! We have updated the related works section with the reviewer’s suggestions, as well as additional references that help contextualize the work.

“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 independent 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 et al. 2023). On the behavioral end, there has also been prior work discussing how the addition of biologically-realistic 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 to specifically study the effects on representation learning. In this paper, we show how deep RL networks can be a testbed for studying representational changes and serve as a multi-region model for neuroscience.”

Reviewer comment (Novelty):
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.

Author response:
Thank you for all the suggestions! We have updated the related works section with the reviewer’s suggestions, as well as additional references that help contextualize the work.

“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 independent 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 et al. 2023). On the behavioral end, there has also been prior work discussing how the addition of biologically-realistic 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 to specifically study the effects on representation learning. In this paper, we show how deep RL networks can be a testbed for studying representational changes and serve as a multi-region model for neuroscience.”

Reviewer comment (Writing):
On page 4, the description of the auxiliary losses can be improved. It is an important part of the work, so it would be beneficial to make it 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_{+}}$, it should perhaps be $\tau (z_{t+1},a_{t+1})$ instead of $\tau (o_{t+1},a_{t+1})$. For the negative sampling loss term, explicitly stating that $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. Additionally, a sentence explaining the motivation for choosing these two specific auxiliary losses would be good. The reviewer suspects it is loosely inspired by pattern completion vs. pattern separation in hippocampus, but confirmation in this section would be beneficial.

Author response:
Thank you for the suggestion! We’ve added the following clarification in section 3 where we describe the role of $L_+$:

“The positive sample loss is defined as $L_+ = ||\tau(z_t, a_t) - z_{t+1} - \gamma \tau(z_{t+1}, a_{t+1})||^2$, where $z_{t} = E(o_{t})$ and $\tau(z_{t}, a_{t}) = z_{t} + T(z_{t}, a_{t})$. That is, in the $\gamma=0$ case, the network $T$ is learning the difference between current and future latent states such that $\tau(z_t, a_t)=z_t + T(z_t,a_t) \approx z_{t+1}$. This encourages the learned representations $z$ to be structured so as to be consistent with predictable transitions (Francois-Lavet et al, 2019).”

Yes, thanks for the catch! It is now fixed.

Thanks for the suggestion! We’ve added the following in the beginning of paragraph 3 of section 3:
“We emphasize that $z_i$ and $z_j$ are randomly sampled from the buffer and thus may represent states that are spatially far from another.“

Thanks for the suggestion! That is part of the motivation. We have expanded the explanation in section 3 to the following:

“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 ($L_-$ and $L_+$) is not just useful for optimization reasons -- it also mirrors the hypothesized pattern separation and pattern completion within the hippocampus (O’Reilly & McClelland, 1994; Schapiro et al., 2017). However, we note that negative sampling elements are not always needed to support self-predictive learning if certain conditions are satisfied (Tang et al 2023).”

Reviewer comment (Presentation):
It would be helpful to state what the colors in Figure 2D mean, as it was not clear upon first read.

Author response:
We add the following sentence in the caption for Figure 2D:
“Diagram of the encoder network (red), learned latent state (gray), and value-learning network (blue).”

Reviewer comment (Writing):
The memory component in page 8 is not explained at all. It seems important to describe what this is in order to put Fig. 4's results in context.

Author response:
We have modified the description to clarify this:
“That is, the input into the encoder at time $t$ is $o_t + \alpha o_{t-1} + \alpha^2 o_{t-2}+\dots$ for some $\alpha<1$. This decaying sum of recent observations captures information about the recent past in a simple way, and is inspired by representations hypothesized by temporal context model (Howard & Kahana, 2002). “

Reviewer comment (Writing):
There are a couple of design decisions that are confusing (details in the questions section). The authors should clarify the motivations behind them.

Author response:
This is a really great question! It would be interesting to explore how different value learning networks affect the robustness of the results and the representations found in the system. We unfortunately did not have time to test these different simulations. There were a few factors that motivated our initial  design choice to use a Q-learning system:
1. There is some precedent in that previous papers model the striatum with Q-learning and discuss experiments that indicate the striatum may be explained by both actor-critic and Q-learning setups (Geerts & Burgess 2020, Averbeck & O’Doherty 2022, Blackwell & Doya 2023, ). For our purposes, using a deep Q network was a simple approximation for a general model-free reinforcement learner. Interestingly, a recent theoretical paper has suggested the presence of action-surprise signals in striatal dopamine combined with an actor-critic architecture results in a model that approximates Q-learning (Lindsey & Litwin-Kumar 2022).
2. Freely behaving animals have been observed to learn in an offline manner. For instance, Rosenberg & Meister 2021 show that freely behaving mice in a maze will behave largely randomly in initial environment explorations. These mice will then make directed and efficient trajectories to rewards when not exploring.
3. Finally, the offline learning setting is relevant because it connects to the hippocampus in many ways. Hippocampal replay has been suggested to contribute to offline RL as the supplier of the replay buffer used in the RL updates (Mattar & Daw 2018). The negative sampling we used in the predictive model (which relies on random replay) is also supported by hypotheses of pattern separation in the hippocampus. Thus, there are interesting biological similarities between the use of an offline replay buffer and activity in the hippocampus. In future experiments, we are interested in using this structure to test replay-related hypotheses.

Reviewer comment (Presentation):
There is no section in the paper explicitly discussing the limitations of the work. This is considered an important part of any ICLR paper and should be included.

Author response:
Thank you for this suggestion – we agree it is a good practice. We add the following sentence to the conclusion:
“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 developed in the value learning network, which is ripe for comparison with striatum data.”

Reviewer comment (Presentation):
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.

Author response:
Thanks for noting the lack of clarity! We were not just interested in the predictive network but also wanted to explore how the other modules in the system mutually interacted. We make a few edits to the introduction to emphasize that we are interested in the roles of both learning heads.

In paragraph 2 of the introduction we add:
“It is unclear how value learning, predictive objectives, and feature learning mutually interact to shape representations.”

In the final paragraph of the introduction we make a few modifications:
“We further demonstrate that a deep RL model with multiple objectives 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 hippocampus.”

Reviewer comment (Evaluation):
I'm curious to learn if there is a similarity between the action selection network and the neural activity observed in the striatum.

Author response:
This is a great question! There is prior work that has explored the similarities between action selection networks in deep RL models and striatal neural activity. For instance, Dabney & Botivinick 2020 show how signals in the ventral tegmental area, a main projector to the striatum, can be explained by value encoding in the distributional RL framework. Lowet & Uchida, in an upcoming 2023 NeurIPs UniReps workshop paper, further this line of work by recording from striatum and demonstrating consistency with a distributional RL model as well. As another example, Lindsey & Litwin-Kumar 2022 show how action-modulated striatal dopamine responses can be explained by value encoding in a deep actor-critic model that is approximating deep Q-learning.

However, there is still much to explore in comparing the value network with striatal experiments. In particular, there are likely interesting effects resulting from the other modules of the system, and our model would be a great testbed for these effects. Although we did not have the time to explore this, we added a sentence in the discussion paragraph discussing limitations of our work: “We also did not examine representations in the value learning network, which is ripe for comparison with striatum data.”

Reviewer comment (Experiments):
Figure 2: Were cells that didn't show place-like activities filtered out in these analyses?

Author response:
All cells were included in population level analyses, without any exclusion criteria. However, to display example place fields for Figure 4B, cells are sorted by spatial modulation and the top four cells are shown for each cell. All additional cells are pictured in the Appendix, however.

Reviewer comment (Experiments):
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.

Author response:
> I don't find any examples of what gridworld environments are being solved by these models, let alone how complex they are. 

To make this more clear, we have now included examples of the environment and corresponding visual inputs into Appendix Figure 2. We use a gridworld environment where the underlying state space is a 8x8 grid with walls at the edges of the arena. Four discrete actions are possible (up, down left, right) and the 2D  visual inputs to the agent are either (A) a top-down view of the agent’s location in the arena or (B) the same top-down view with pixels randomly shuffled. In recently introduced experiments (described below) we also experiment with (C) CIFAR images as visual input for added visual complexity. In case B and C, this removes spatial correlations between visual inputs from neighboring states.

> 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. 

Yes, this is an important point. The predictive auxiliary objective assumes that the effects of actions in the environment are predictable. Specifically, the predictive network is trained to predict $z_{t+1}-z_{t}$ given $(z_t, a_t)$.

To further test the effects of environment predictability, we ran additional experiments in a stochastic gridworld environment and included these results in Appendix Figure 3. In this environment, transitions are not deterministic. There is a probability $p$ that the agent’s action $a$ is not obeyed and instead the environment transitions according to an unselected action instead. We find that, as $p$ increases, the relative benefit of the predictive objective vanishes and all models perform similarly (Figure A3D). We discuss these results in the bottom of section 4.1.

Reviewer comment (Experiments):
Can you provide example environments?

Author response:
> I don't find any examples of what gridworld environments are being solved by these models, let alone how complex they are. 

To make this more clear, we have now included examples of the environment and corresponding visual inputs into Appendix Figure 2. We use a gridworld environment where the underlying state space is a 8x8 grid with walls at the edges of the arena. Four discrete actions are possible (up, down left, right) and the 2D  visual inputs to the agent are either (A) a top-down view of the agent’s location in the arena or (B) the same top-down view with pixels randomly shuffled. In recently introduced experiments (described below) we also experiment with (C) CIFAR images as visual input for added visual complexity. In case B and C, this removes spatial correlations between visual inputs from neighboring states.

Reviewer comment (Experiments):
What is the effect of varying the complexity of the environment?

Author response:
> If possible, please consider varying the complexity of the environments... What is the effect of varying the complexity of the environment?

We agree that varying the complexity of the environment would be a helpful way to test the robustness of the observed effects.

To further test the effects of environment complexity, we ran a new set of experiments in which the observation at each state in the gridworld environment consisted of a different CIFAR image, and included these results in Appendix Figure 2-3. We find the results to be consistent with the simpler gridworld inputs we use in the main figures (Figure A3C). That is, predictive auxiliary objectives enable faster learning of the gridworld task compared to other models (Figure A3C). Additionally, the model equipped with a predictive objective can learn the task for small latent sizes whereas the model without auxiliary objectives cannot (Figure A3C). We update section 4.1 to discuss these results.

We also note that prior work has tested predictive auxiliary objectives in single-task settings with greater task complexity. In these findings, the inclusion of these objectives consistently improves performance over models without auxiliary objectives (Jaderberg et al., 2016; Shelhamer et al., 2016; Shelhamer et al., 2016; Oord et al., 2018; Wayne et al., 2018).

Reviewer comment (Theory):
Similar 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.)

Author response:
> Do you have thoughts about how partial observability might affect the predictive module? 

We predict that a predictive auxiliary objective may provide additional benefits in a partially observable task. For instance, the additional predictive information induced by the transition module can help the network infer hidden state dynamics better. In cases of observation noise, predictive information can also help network representations ignore irrelevant or noisy parts of the observation. 

To further test this, we ran new experiments with a partially observable version of the alternating-T maze task in Figure 4H. These new results are added to Appendix Figure 5. In this version, the model is only provided with a 2D visual observation that indicates its current location ($o_t$). The model is modified with recurrency such that its previous latent state ($z_{t-1}$) is fed back into the encoder and can be used to infer the current latent state. Thus, $z_t = E(o_t, z_{t-1})$. To solve the task, the model must use recurrency to maintain an internal memory of its current context. We find that the predictive auxiliary objective greatly improves the model’s ability to learn this partially observable task (Figure A5D). The splitting across population representations are still similar as in the fully observable case ( Figure A5F vs Figure 4I). We discuss these results in the bottom of section 4.3.