Brain-Like Replay Naturally Emerges in Reinforcement Learning Agents

## Abstract

Abstract Replay is a powerful strategy to promote learning in artificial intelligence and the brain. However, the conditions to generate it and its functional advantages have not been fully recognized. In this study, we develop a modular reinforcement learning model that could generate replay. We prove that replay generated in this way helps complete the task. We also analyze the information contained in the representation and provide a mechanism for how replay makes a difference. Our design avoids complex assumptions and enables replay to emerge naturally within a task-optimized paradigm. Our model also reproduces key phenomena observed in biological agents. This research explores the structural biases in modular ANN to generate replay and its potential utility in developing efficient RL.

## 1 Introduction

Experience replay (Lin, ) is an important technique for online reinforcement learning (RL). By storing past experiences in a circular buffer and training the agent on these experiences repeatedly, experience replay improves sample efficiency and training stability (Schaul et al., ), and successfully improves the performance of RL algorithms (Mnih, ; Mnih et al., ). There are several works on how to design more sophisticated experience replay mechanisms, like hyperparameter tuning (Fedus et al., ), prioritizing replay with higher temporal difference (TD) error (Schaul et al., ), or developing a generative replay model (Van de Ven & Tolias, ). However, given the significant efficiency gap between biological and RL agents, we explore whether replay algorithms can be further improved by drawing inspiration from brain structures.

In neuroscience, replay is the reactivation of a sequence of place cells that encode recently experienced locations. It has been observed in many brain areas during rest or sleep, from the hippocampus (Buzsáki, ; Nádasdy et al., ) to the neocortex, including the prefrontal cortex (Peyrache et al., ; Kaefer et al., ), visual cortex (Ji & Wilson, ), and motor cortex (Eichenlaub et al., ). Some people believe it represents the ”virtual trajectory” in the brain that primarily serves two functions: memory consolidation (Foster & Wilson, ; O’Neill et al., ) and planning (Pfeiffer & Foster, ; Widloski & Foster, ). Its dynamic properties and functional diversity have attracted significant attention in neuroscience research.

There are currently two focuses in this field: the emerging condition of replay and the function of replay. Many models have been proposed to explain these two aspects. Some biophysical models suggest that replay arises from the combination of past experiences and specific neuronal network connectivity rules, such as asymmetrical synaptic connections (Tsodyks et al., ) and short-term depression (Romani & Tsodyks, ). In addition to connectivity rules, task optimization provides another way to understand the emergence of replay. (Krishna et al., ) proved that denoising dynamics under task optimization could lead to offline reactivation of past experiences. Levenstein et al. () verified that only the network trained on multi-step predictive learning with recurrent connection and head-direction input can generate offline trajectories resembling replay. These models are illustrative in that they connect conditions like circuit structures, task optimization or predictive learning to the generation of replay. However, they fail to explain the functional importance of replay in the task.

On the other hand, experience replay provides good evidence that replay does have some importance and helps learning. However, this approach is not ideal for exploring the conditions under which replay emerges, as all replay of stored experiences is hard-coded and predefined. The model never finds a way to organize experiences and perform efficient replay without force. The gap between these two research fields leads to the natural question of whether we could find some conditions under which a model could naturally generate adaptive replay sequences with functional importance. Identifying such conditions could provide important insights on how to design new reinforcement learning agents that quickly adapt to the environment, given the present low efficiency of RL algorithms and the unique advantage of biological agents in learning new tasks.

Recently, some studies have attempted to approach this problem within the RL framework. For example, Mattar & Daw () propose that a special variable called the expected value of backup (EVB) determines the replay priority of different experiences, and the replay is accompanied by Bellman updates which are essential for learning the value function and policy for the task. However, because we cannot know EVB until we perform value iteration for the whole environment, and after value iteration it seems unnecessary to use replay at all, the generating condition in there are still some hard-coded constraints in this method.

In this article, we propose a model to circumvent hard-coded settings, provide some conditions as a natural explanation for the emergence of replay, and then show the functional importance of our generated replay sequences. In summary,

- 1.
Based on two straightforward conditions, we develop a model capable of generating offline replay activation.
- 2.
Throughout the entire learning process of the flexible navigation task, the evolution of replay spatial distribution of our model resembles that observed in real animals.
- 3.
Through ablation experiments, we verify the functionality of replay by showing that the exploration efficiency of our model surpasses that of models without replay.
- 4.
We analyze the information that replay flow carries, and explain how replay can promote efficient learning.

## 2 Methods

#### Conditions to generate instrumental replay

What could these conditions possibly be? The main areas where replay occurs are the prefrontal cortex (PFC) and hippocampal formation (HF), usually modeled as policy network and world model respectively due to their unique functions. Replay in these two individual areas is often observed to be coordinated (”phase-lock”) when a new environment learning task is presented (Place et al., ). We hypothesize that communication between these two areas is crucial for the generation of instrumental replay. Therefore, we propose two conditions that could be implemented in RL agents:

$$ Condition 1: Replay serves for reward maximization. (1) Condition 2: Replay is accompanied by communication between PFC and HF. $$

To achieve Condition 1, we employ the end-to-end RL framework without hard-coded design. For Condition 2, our RL agent has a world model and a policy model, representing HF and PFC respectively, and we add an additional information passage between these two models that is trained alongside the policy network. This setting allows the agent to adjust the information communication flow based on its influence on actions and expected rewards. Discussions about the plausibility of these settings can be found in subsection 4.1.

Figure: Figure 1: A-C. Model structure. A The HF GRU module completes the task of path integration and episodic memory. The PFC RNN module is responsible for decision-making. The information passage between the PFC and HF only opens at rest. B. During mobility, the information passage remains closed, and the HF and PFC modules operate independently. The place cell activation reflects the real present locations. C. During immobility, the information passage opens. The HF and PFC modules start to communicate with each other. The activated place cell output constituted a replay sequence. D-F. Task setting. D. The agent should start from S, first get to checkpoint C, consume a small amount of reward (0.5), and then arrive at goal G to get a large reward (1.0). Directly moving to G will cause no rewards. E. A representative trajectory generated by the trained RL agent. First, it reaches C and gets the small reward. The replay happens at this time. F. The replay event in E.
Refer to caption: x1.png

#### Principle of designing different modules

We preface the introduction of model structure with key biological biases that we incorporate into our design.

Biologically, HF has abundant recurrent connections in its CA3 area (Le Duigou et al., ), which is also thought to be a major area generating replay. Thus we implement a gated recurrent unit (GRU) for HF. Besides, there are two main functions for HF. First, the HF is good at encoding the environment structure (i.e. cognitive map) which can serve as a basis for path integration and flexible navigation (Whittington et al., ). In our setting, the HF module learns to do this by predicting the next location (path integration) and reward, based on history (hidden state) and environmental information (input) (Figure 1A). The HF is also thought to store episodic memory (Whittington et al., ; Battaglia et al., ), and in our model, it learns to memorize past reward history.

The PFC is also modeled as a recurrent network (Figure 1A) due to the recurrent nature of the real prefrontal cortex. The prefrontal cortex is thought to be the center for decision-making (Barraclough et al., ) and value encoding (Padoa-Schioppa & Assad, ), and in our task, it learns to estimate values and select actions to maximize the expected reward using the hidden state and environmental information.

To handle high-dimensional visual inputs, we use an ancillary Encoder module modeled by a convolutional neural network (CNN), simulating the sensory cortex (Figure 1A).
Regarding replay, we establish an information passage between HF and PFC. Because replay is only detected when the subject receives a reward and stops to consume it (Igata et al., ), the information passage remains closed during movement (Figure 1B) and opens when the agent receives a reward (Figure 1C). During the agent’s movement, the PFC decides the agent’s next action (i.e., up, down, left, or right) based on the input of environmental representation processed by the Encoder module. At this time, it does not communicate with the HF module. Meanwhile, the HF maintains its own dynamics for correct place cell activation and recent episodic memory (Figure 1B). When the agent receives a reward, the two modules interact for several steps and influence each other through the information passage. The reactivation of place cells in the HF is then probed and analyzed as replay.

#### Design of the HF Module

The update formula of the HF module can be written as:

$$ $\displaystyle h_{t}=f_{\mathrm{HF}}(W_{h}h_{t-1}+(1-\mathbb{I}_{\mathrm{replay }})W_{h,\mathrm{in}}[o_{t},r_{t-1},a_{t-1}]+\mathbb{I}_{\mathrm{replay}}W_{h, \mathrm{r}}\theta_{t-1}),$ (2) $\displaystyle[\hat{r}_{t-k:t},\hat{G}(s=1:m)_{t}]=\sigma(W_{h,\mathrm{out}}h_{ t}).$ (3) $$

$h$ is the hidden state of HF. Input consists of $o_{t}$, $r_{t-1}$, and $a_{t-1}$, which are the current partially observable state, the previous reward and the previous action, respectively. $W_{h}$, $W_{h,\mathrm{r}}$, $W_{h,\mathrm{in}}$, and $W_{h,\mathrm{out}}$ are trainable parameters of neural networks. $\theta$ is the hidden state of the PFC. $\mathbb{I}_{\mathrm{replay}}$ indicates the beginning of the rest period when replay and communication happens and input stops. $f_{\mathrm{HF}}$ is the activation function for the HF module. The HF module then anticipates a Gaussian distribution peaked at the next location as well as rewards for recent steps $\hat{r}_{t-k:t}$. $m=5\times 5$ is the number of locations in this grid world. $k=10$ is the length of the reward history the agent needs to remember.

The loss for structural learning can be written as the cross-entropy between the softmax output, $\hat{G}(s))$, and the real distribution, $G(s)$: $\mathcal{L}_{\mathrm{pred-loc},t}=-\sum_{s=1}^{m}\hat{G}(s)\log G(s)$. To equip the HF module with the reward structure of the environment, it predicts the reward of the next step, $r_{t}$, with the loss written as $\mathcal{L}_{\mathrm{pred-r},t}=|\hat{r}_{t}-r_{t}|$. The loss of memorizing rewards is calculated as the L1 difference between the remembered rewards and the actual rewards: $\mathcal{L}_{\mathrm{memory-r},t}=\sum_{i=1}^{k}|\hat{r}_{t-i}-r_{t-i}|$.

#### Design of the PFC Module

The update formula of PFC could be written as

$$ $\displaystyle\theta_{t}=f_{\mathrm{PFC}}(W_{\theta}\theta_{t-1}+(1-\mathbb{I}_ {\mathrm{replay}})W_{\theta,\mathrm{in}}[o_{t},r_{t-1},a_{t-1}]+\mathbb{I}_{ \mathrm{replay}}W_{\theta,\mathrm{r}}h_{t-1})$ (4) $\displaystyle[a_{t},v_{t}]=W_{\theta,\mathrm{out}}\theta_{t}.$ (5) $$

$\theta$ is the hidden state of the PFC module. Input $[o_{t},r_{t-1},a_{t-1}]$ and the indicator $\mathbb{I}_{\mathrm{replay}}$ are the same as in Equation 2. $W_{\theta}$, $W_{\theta,\mathrm{r}}$, $W_{\theta,\mathrm{in}}$, and $W_{\theta,\mathrm{out}}$ are trainable parameters of neural networks. $f_{\mathrm{PFC}}$ is an activation function. The state value $v_{t}$ is also estimated by the PFC for training purposes. We can see that replay is optimized because it changes the hidden state $\theta_{t}$ of the PFC and thus affects the expected reward.

#### Training and Test Paradigm

We first pre-train the world model HF and the Encoder using random trajectories. Then the weights of the Encoder and HF are frozen, and we start training the PFC module by proximal policy optimization (PPO) to maximize the reward. We conduct the following analysis with all model weights fixed.

The reward settings during pre-training and training are the same. The reward is randomly chosen from one of the nine locations in the center of the room, and the location is reset at a probability = 0.1 so that memorizing reward locations is helpful to finish the task. However, in the test stage, the reward is the same as in the biological experiment.

## 3 Results

### 3.1 Biologically similar replay sequences emerge after training

Figure: Figure 2: A. Performance for the RL agent in the test period. B. Legends for C and E. Different colors represent different parts of the room. C. Change of replay distribution for different segments in the animal data, adapted from (Igata et al., ). D. Distribution of distances between adjacent replay steps. E. Change of replay distribution generated by the RL agent.
Refer to caption: x2.png

#### Task setting

We use the animal experiment from (Igata et al., ) as a reference, in which a rodent was trying to navigate to a dynamically changing reward location in an open area (Figure 1D). First, the animal learned to run from the starting point S to checkpoint 1 (C1) to get a small reward, and then run from C1 to G to get a large reward. Violation of the order would cause no return at all (Figure 1D). When the stage of pre-learning was finished, the small reward was relocated to checkpoint 2 (C2, Figure 2B), requiring the animal to learn a new optimal path (S-C2-G). This was the beginning of the learning stage. Initially, the animal continued to take the path S-C1-G followed by G-C2-G after failing to receive a reward at G. It then realized the existence of a shortcut (S-C2) and gradually settled on it. The aim was to investigate changes in the relative amount of replay sequences representing different parts of the room (”replay distribution” in this paper) during the learning stage when the animal was trying to adapt to the change of reward location. They found that the amount of replay of the original trajectory (S-C1, C1-G) decayed rapidly, while the shortcut (S-C2) increased even before the agent adopted it (Figure 2C). This result could be evidence that replay supports flexible navigation by exploring possible paths and strengthening better ones.

#### Simulation result

We create a similar simulation environment as in (Igata et al., ). The animal is replaced by an RL agent moving in a $5\times 5$ open arena. As in the animal experiment, we put the small reward at C1 and then relocate it to C2. Without adjusting the network parameters, and simply by modifying the hidden states of the RNNs, the RL agent successfully identifies the optimal path to the new checkpoint. (Figure 3A).

The locations of C1 and C2 are the same as in the original experiment (Figure 2B). We first measure the distribution of distances of adjacent replay steps (Figure 2D). Compared with the random level or the model without pretraining, the optimized model shows a significant tendency to generate adjacent replay steps. Therefore, it tends to generate continuous trajectories rather than random skipping points.

We attribute each replay trajectory to one of the four paths in Figure 2B and then calculate the relative fraction of the four paths at each time step, as the ”replay distribution” (See subsection A.1 for calculation details). The change in replay distribution of the RL agent (Figure 2E) closely mirrors that of the rodent (Figure 2C) in two important aspects. First, the amount of replay of C2-G increases and then drops down. Second, the amount of replay of S-C2 increases. Recalling that the animal first took path S-C1-G-C2-G and then S-C2-G, This trend possibly reflects the animal’s focus from ”finding a plausible path” to ”finding an optimal path”. These results prove that the conditions proposed in 1 are sufficient to generate replay.

### 3.2 Ablation study demonstrates that replay helps learning

Figure: Figure 3: A. The total reward decreases when the signal from HF to PFC is replaced by random noise (left) or all-zero vectors (right). B. The different number of replay steps is masked (right), and the performance decreases monotonically as we mask more steps (left). C. The number of steps it takes to find the new reward increases when the multi-step information emission is replaced by one step. D. The performance is impaired only a little when the order of information is shuffled.
Refer to caption: x3.png

Figure: Figure 4: A. Decoding accuracy for the reward location by activities in HF (left), PFC (middle), and information passage (right), respectively. The X-axis is composed of three parts: the initial stage before replay, steps 1-4 during replay, and the output stage after replay. B. Decoding error for correct future actions from PFC activities after replay. Red, orange, yellow, and brown denote the first, second, third, and fourth future action following replay. In A-C, darker colors represent the agent’s first getting the relocated reward, while lighter colors represent the second time. C. Value map produced through the “stop and scan” method. The C1 map represents the value before the checkpoint change, and the subsequent value maps are when the agent gets to C2 repeatedly. D. Value advantages of checkpoint 2 over checkpoint 1, calculated by convolving the value map with a Difference of Gaussian (DoG) filter. E. Left, value advantage calculated as the value of the path S-C2 minus the value of S-C1. Right, the value of the path C2-G minus the value of C1-G.
Refer to caption: x4.png

To prove replay is functionally important, first, we replace the signal in the information passage to be random noise or an all-zero vector. Because the information exchange between HF and PFC is bidirectional, we carry out this replacement in both directions respectively and find that the performance is only impaired when the information from HF to PFC is replaced by meaningless signals. The information in the reverse direction is of no use (Figure 3A). This interesting finding echoes some neuroscientific discoveries that brain wave in HF leads that in PFC in the task encoding stage when phase lock between these two brain areas happens (Jadhav et al., ; Spellman et al., ).

Second, we replace the multi-step information exchange process with one-step information emission and retrain the RL part, i.e. the PFC network and information passage (Figure 3C, left). We measure the number of exploration steps it takes to get to the reward at the time when it’s relocated but the agent continues to follow the old path S-C1 and fails to find it. The original model demonstrates a significantly lower number of exploration steps and higher exploration efficiency compared to the ablation models (Figure 3C, right). This suggests that multi-step information emission is superior to single-step emission. We further mask varying numbers of replay steps from the final decision point and replace them with random noise (Figure 3B, right). Performance decreases monotonically as more replay steps are masked (Figure 3B, left). This proves that the information is incrementally unrolled as a sequence during replay. When we shuffle the order of the messages sent during replay, the performance is only slightly affected (Figure 3D), suggesting that the information may be sent in the form of independent packages rather than a whole sequence.

### 3.3 Information flow during replay entails information about context and action plan

Thanks to the gray-box nature of AI, we can now extract the hidden states of each module and analyze the information inside. We assume that, the PFC initially has no idea that the reward location has switched from C1 to C2, but it changes its context encoding and gets to know this through replay. To test this, we train a Gaussian Naïve Bayes decoder to predict the reward location from the PFC’s hidden states, and observe a gradual increase in decoding accuracy as the agent first encounters the new reward location, which remains high with repeated encounters (Figure 4A, middle). The analysis on HF hidden states has the same result (Figure 4A, left). This shows that replay helps identify the new context quickly. Additionally, another decoder trained on the messages sent from HF to PFC achieves high accuracy in predicting reward location. These results suggest that HF-PFC information emission helps update memory.

We also train a Ridge Classifier to predict actions from PFC hidden states, and find that decoding errors decrease significantly as the agent gets to the new checkpoint repeatedly. This shows that replay helps form future direction intentions (Figure 4B). It’s important to note that the weights are fixed during the entire test period, so the model learns how to use replay to update memory and form plans to complete the flexible navigation task.

In the end, we investigate how replay guides the update of hidden states in the PFC using a ”stop and scan” paradigm. When replay ends, we pause the real-time clock and allow the agent to explore the map randomly for 100 steps, recording the PFC’s value output as its estimation of every location. Finally, we plot the values of different locations in the whole map (Figure 4C). As predicted, the high-value area shifts from S-C1-G to S-C2-G after the change of reward location (Figure 4E, left and right are for S-C and C-G respectively). We then convolve the value map with a Difference of Gaussian (DoG) filter. The filter has a positive peak at C2 and a negative peak at C1, which could measure the relative advantage of the area around C1 to C2. The result that the advantage changes from negative to positive confirms the conclusion above (Figure 4D).

Figure: Figure 5: Manifold analysis reveals the context-switch process in detail. A. The 3D embedding of the “neural” manifold through dimension reduction of PFC activities when the small intermediate reward stays at C1 (left), the agent meets the relocated reward at C2 for the first time (middle), and the agent meets C2 for the second time (right). The highlighted trajectory represents the neural manifold at the corresponding stage, but the other trajectories are also plotted for comparison. B. AEV of different PCs of PFC activities during replay when the agent first finds the new reward. C. Left, the dimension of the PFC activities during replay (virtual experience) when the threshold for AEV is set to 70%. Right, the dimension of PFC activities during both movement and replay. The results are the same. D. The mean square distance of data points to their KNN centroids calculated from hidden states in HF and PFC during movement.
Refer to caption: x5.png

### 3.4 Manifold Analysis Reveals That Replay Flow Bridging Between Contexts

To better visualize the dynamic changes in PFC activities, we perform a 3-dimensional principal component analysis (PCA) on PFC hidden states during the whole trajectory (Figure 5A). Each trajectory is composed of three stages: moving from the starting point to the checkpoint, getting the reward and generating replay, and moving from the checkpoint to the goal. For visualization, we only pick those stages that the animal finishes in four steps. Each step has a different color shown at the bottom. We connect the centers of the data point groups of each step to form an average trajectory. We use two distinct colors to differentiate activities during replay and the real movement. For example, in the left figure, we use red to denote the PFC hidden states during movement and green to denote hidden states during replay. The red continuous dots are separated by the green dots just like how a real trajectory is separated by replay. We use gradually darker colors in each trajectory to represent the gradually increased time steps, shown at the bottom. In the left and right figures, the reward is at a known point, and we can see the PFC hidden states constitute an orbit. In the middle figure, the agent first ”thinks” the reward is at C1, then changes its mind and starts to explore, and finally finds C2, and the PFC hidden states settle on the new orbit. Replay connects two low-dimensional context subspaces. The activities during random exploration when the animal cannot find the reward are too dispersive to see clearly, so we don’t show them here.

Then we analyze the dimensions of different subspaces spanned by PFC hidden states. The dimension is defined as the first index of PC whose accumulated explained variance (AEV) reaches 70%. Figure 5B is a representative AEV figure of PCs of PFC hidden states during replay when the agent finds the relocated reward. The dimension of this subspace is 2. Similarly, we show the subspace dimension of PFC hidden states only during replay (Figure 5C, left), or in the whole trajectory (Figure 5C, right). When the reward location changes, the all-state subspace dimension briefly increases to 3 and then drops back to 2, indicating a more complex subspace during the context switch. If we consider one ”context orbit” as a 2-dimensional circle, then switching between these two orbits would produce a 3-dimensional curve, This is a rough and intuitive understanding.

We also evaluate the trajectory stability during the context switch. We first perform the K-Nearest Neighbour (KNN, $K=20$) algorithm on hidden states during movement to find the centroids for the data points. Then we calculate the mean square distance of data points from their KNN centroids (Figure 5D) as the variance of the hidden states. The variance increases briefly during the context switch and then drops back when the agent gets to C2 for a second time. $K=8$ gives out the same result. This shows that the first encounter of the new reward is accompanied by an unstable and dispersive trajectory. And the trajectory becomes more stable in the following time.

## 4 Discussion

In this article, we have demonstrated that biologically similar replay can emerge from the interaction between the HF module and PFC module within a task-optimization framework. We show that replay in our model is actually a form of multi-step information emission from HF to PFC. Then we show that during replay, the context is updated both in HF and PFC, and the cognitive map and future plan are updated in PFC. Therefore we describe how replay helps learning. At last, we plot the manifold to give a more intuitive understanding. See Appendix A.3 for reproducibility results.

### 4.1 Plausibility

#### Conditions

Our two sufficient conditions (1) are inspired by neuroscientific research. For Condition (1), studies have shown that blocking replay impairs spatial learning (Jadhav et al., ; Girardeau et al., ; Ego-Stengel & Wilson, ). The presence of task-focused replay can predict the accuracy of following decisions (Ólafsdóttir et al., ). Replay is modulated by reward (Ambrose et al., ; Liu et al., ; Michon et al., ), and is associated with reward-guided updates of internal models based on past experiences (Carey et al., ).

For condition (2), the interaction between HF and PFC during replay is supported by a lot of evidence. For example, replay is accompanied by sharp-wave ripples that synchronize the wave phases of HF and PFC (Tang et al., ). Replay is thought to be the substrate of the continuous memory transition from the hippocampus to the neocortex, namely memory consolidation(Kumaran et al., ). Recent computational models also reveal that the communication between the world model and policy network could help decision-making (Jensen et al., ).

However, there is some evidence that does not support this hypothesis as a generating condition for replay. For example, Kaefer et al. () reported PFC replay that is strongly associated with rule-switching performance but is independent of HF replay. As we do not measure PFC replay, it may be that not replay but other PFC activities in our model. Future work could design experiments to make this clear.

#### The segregation of HF and PFC during movement

The notion that the HF and PFC are entirely segregated during movement may seem somewhat implausible. We acknowledge that this highlights a limitation of our model’s fidelity to realistic biological settings. However, we focus on replay which occurs only at rest. During mobility, the HF and PFC could communicate with each other not by replay but by theta waves, which are thought to be mainly driven by external signals and involve only nearby place cells’ firing, rather than a global sequence. Allowing the two modules to communicate during movement might add unnecessary details less relevant to the problem we want to analyze. However, it would also be exciting to see how adding biological details could still produce the phenomena we want and how different factors interact with each other.

#### The difference of PFC from the striatum

Previous RL theories for animal behaviors also emphasize the striatum as the center for model-free learning. The dorsal lateral striatum (DLS) and the ventral lateral striatum (VLS) correspond to the actor and critic parts in the classical Actor-Critic algorithm. The dopamine serving as TD signal is calculated in the ventral tegmental area (VTA) and drives synaptic plasticity in the whole network (Montague, ). Does this role of striatum overlap with PFC? Wang et al. () proposes a meta-reinforcement learning model where PFC is the inner-loop ”fast” learner and the dopamine system is the ”slow” learner. We think it’s also useful to consider the training process involved weight changes as a kind of dopamine learning, and later context identifying and following decision-making involved hidden state changes as PFC learning.

### 4.2 How replay promotes learning

The classical view posits that replay aids value-based reinforcement learning (Liu et al., ; Foster & Wilson, ; Mattar & Daw, ; Ambrose et al., ), and our model provides a concrete implementation supporting this view. In our model, cognitive map modulation in the PFC changes as a result of replay and then stabilizes (Figure 4). A value-modulated cognitive map in the PFC is reminiscent of value and task structure representations in the orbitofrontal cortex (OFC) (paenel2020stable; Zhou et al., ; Padoa-Schioppa & Assad, ), a part of PFC. Our work suggests these representations could combine as a value function (V function) to guide decision-making. Our PFC cognitive map is formed with the help of HF, which also aligns with previous findings (Mızrak et al., ).

Despite those model-free components, our model should still be classified as model-based. The main difference from AI model-based methods is that they sample from the distribution learned by the world model and use these samples to improve the policy. The advantage is that it could provide gradients for the sampling process and have the gradients of the reward loss directly back-propagated to the policy network. Our method, however, doesn’t design how the world model should help the policy network. The replay in our model is similar to a sampling process and maybe gradients are propagated through this process but future work needs to make this clearer. Recently Jensen et al. () explores the model-based explanation of replay, but they also have an explicit design that the samples from the world model are directly taken as input to the policy network.

## Appendix A Appendix

### A.1 Calculation of replay distribution

Figure: Figure 6: Calculation of replay distribution
Refer to caption: extracted/5905634/Figures/distribution_cal_complicated.png

Replay is defined as the sequential activation of place cells. Therefore, a replay trajectory would represent a path in the room. We determine the attribution of replay by calculating the overlapping of our replay with different parts of the room (Figure 6 left). Equal overlapping would cause exclusion of that trajectory. Finally, the numbers of replay trajectories representing each part of the room are normalized to get the spatial distribution of replay, namely ”replay distribution”.

The rule we use is a little more complicated. Each part of the room is divided into two segments with scores that equal to 1 and 0.5 respectively. subsection A.1 provides a concrete example. The red trajectory is one replay.

### A.2 Annotation for Figures

All error bars indicate $\pm 1$ standard error.

In Figure 4, for all decoding tasks, 80% of the data is used for training and 20% for testing. The Gaussian Naïve Bayes and Ridge Regression methods are used to decode the reward location and future directions respectively. Each step has an independent decoder, and they do not share training sets.

In Figure 5C, measurements are repeated for 10 times with different seeds. The results are the same.

### A.3 Reproducibility

We repeat the experiments across different seeds, and the result has been presented in the main text. Then we retrain the RL across different seeds and different numbers of replay steps, the result is shown in Figure 7.

Figure: Figure 7: A-F. Changes of replay distribution across different seeds or different numbers of replay steps, like in Figure 2E. They have the same trends that the replay of green parts first increases and then decreases, and yellow parts keep increasing. A. Seed = 1. B. Seed = 2. C. Seed = 3. D. Seed = 4. E. Number of replay steps = 2. F. Number of replay steps = 6. G. Exploration steps as in Figure 3C. The model without replay suffers from lower exploration efficiency.
Refer to caption: x6.png

### A.4 Training Curve of Multiple Agents

The RL agents are tested every 5000 steps. The training curves are shown in Figure 8. The model without HF clearly receives fewer rewards than the other models.

Figure: Figure 8: Training curve of different models
Refer to caption: extracted/5905634/Figures/train_curve.png

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: - "There are several works on how to design more sophisticated experience replay mechanisms, like hyperparameter tuning (Fedus et al., 2020)". I don't think Fedus et al. 2020 do this. They do investigate what happens under different replay buffers sizes and update rates. But beyond that, there is not much.
- The citation "(Van de Ven & Tolias, 2018)". The "V" in "Van" should be lower case.
soundness: 2
strengths: This paper is interesting, with a large amount of neuroscience inspiration and guidance into developing the model. The specifics of the model do not look easy to implement, so the emergence of replay is certainly interesting. This paper seems exploratory, and in such papers that aim to model some natural phenomenon in the brain it is naturally difficult to ascertain the validity and accuracy of the model. This is fine, as these kinds of works are just steps towards future investigations. The design of the experiments seem well done and I do believe that the claims of the paper are consistent with the experiments.
summary: Replay has an important role in the brain, and experience replay in RL seems to be important and beneficial in RL. However, experience replay in RL is a hardcoded and predefined procedure in the sense that it is explicitly implemented rather than being an emergent phenomenon. This paper is rather exploratory in that it takes inspiration from replay in the brain and develops a model based off of it to investigate the emergence of replay in RL agents without explicitly programming in experience replay. Inspired by the brain, they propose/hypothesize two conditions for replay to motivate their model. They implement a modular neural network with different parts serving as analogs to different parts of the brain. This model exhibits replay and they then study replay within this model.
weaknesses: They do evaluate their model on essentially a single environment, and it does call into question the degree to which the results and the model generalize for studying replay. They do use an environment that is based on prior research, which is good. Nonetheless, the generality of the results would be somewhat a concern of mine.

---

confidence: 2
contribution: 1
initial_rating: 3
presentation: 1
questions: Please provide definitions for the following terms (mathematically, where possible): ‘place cell’, ‘replay’, ‘replay step’, ‘phase-lock’, ‘one-step information emission’. 

The RL task used in all experiments is not sufficiently explained. While the high-level intent is clear, low-level details are lacking. Please include a precise definition of the environment’s transition dynamics and reward structure. 

In Section 3.1, when attributing replay trajectories to one of the four paths, what happens to trajectories that do not clearly pass through / near either of the checkpoints? 

In Section 3.3, the authors write that they “pause the clock and allow the agent to explore the map randomly for 100 seconds”, but it is highly unclear what this actually means. What are the inputs and outputs of the network during this time?
soundness: 2
strengths: Originality: Uncertain. I am not familiar with other work. 

Clarity: Poor. The authors do not define key terms such as ‘replay’ and ‘place cell’. There are also multiple instances where methodology is not described with enough precision to fully understand the experiment that has been conducted. Lastly, the paper uses many jargon terms which are presumably familiar to neuroscience researchers, but which are unfamiliar for general machine learning researchers. Overall, the paper is difficult to understand and would benefit from significant rewriting to make it more accessible to the broader ML community. 

Quality: Uncertain. I am not able to accurately assess the quality of the paper at the moment, due to significant conceptual confusions regarding the research question under investigation, the  experimental methodology used, as well as the broader significance of the result. 

Significance: Poor. The stated impact of the paper seems to be that agents can ‘generate replay’, and this may be significant to neuroscientists, but the relevance to the broader machine learning / RL community is unclear. As a machine learning researcher, the part I found most interesting was the demonstration that the network was capable of adapting to a different reward location at test-time. Adopting this framing, I find the experiments or results extremely lacking, since the task is extremely toyish and there is no comparison to baselines.
summary: The authors propose a biologically-inspired model-based RL architecture for reinforcement learning. They train this network using a two-stage process to learn a world model and policy separately. They show that this architecture is capable of 'generating replay' organically from reward-based learning. They demonstrate using a simple toy setting that this architecture is capable of adaptation to changing reward structure without parameter modification. They also perform limited interpretability studies on the information being passed from the world model to the policy.
weaknesses: The paper’s central result is that the proposed experimental design is capable of ‘generating replay’ organically from reward-based learning. However, it is never concretely defined as to what ‘replay’ is or how replay can be ‘generated’. A clear mathematical definition would be useful. Furthermore, it is not explained what the significance of ‘generating replay’ in this fashion is. Lastly, it is not clear whether these claims generalise to a broader RL setting, since the task is extremely toy-ish.  

The figures are not well explained. In Figure 1A and 1E, what does “C2 checking time” mean? In Figure 1C, how are the boundaries between the different stages determined? In Figure 1D, what is ‘distance’? In Figure 3B, what does it mean to mask a replay step? In Figure 3C, what does it mean to replace the multi-step emission with one step? In Figure 3D, what does it mean to shuffle the order of information? 

In Section 3.4, it is unclear what conclusions should be drawn from the manifold analysis. What are the takeaways and why are they important?

---

confidence: 2
contribution: 3
initial_rating: 3
presentation: 2
questions: - Regarding weakness 2, I am not sure what the authors mean by "modular" in the abstract. I assume this means the proposed model is a general framework on its own and can be integrated with other blocks. If the authors meant something else, please provide clarifications.
  
- The paper classifies the approach as "model-based," primarily due to the role of the HF as a "world model". Could the authors provide empirical evidence showing HF’s accuracy in predicting next states or rewards, demonstrating its role as an explicit model?
  

I am happy to raise the score if the points in the weaknesses and questions can be addressed.
soundness: 3
strengths: - This paper introduces a neuroscience-inspired model that simulates the replay phenomenon without predefined or hardcoded replay rules, which is interesting, innovative and aligns with biological replay mechanisms observed in the brain.
  
- The authors demonstrate the functional importance of replay by drawing parallels with the biological experiment in [1]. The results show that its RL agent’s replay patterns match the replay distribution and adaptive pathfinding seen in rodents. The abolation studies also show how disrupting the replay signal impairs the agent's exploration efficiency.
  

**References**:

[1] Hideyoshi Igata, Yuji Ikegaya, and Takuya Sasaki. Prioritized experience replays on a hippocampal predictive map for learning. Proceedings of the National Academy of Sciences
summary: This paper presents a novel, neuroscience-inspired model that successfully demonstrates the replay phenomenon without relying on hardcoded rules. However, the experiments are limited to a specific setting, and the paper does not provide justifications that the replay mechanism improves exploration across broader RL algorithms and tasks.
weaknesses: - The paper should contain a background section that describes the overall problem formulation along with some preliminary knowledge. This would help improve the readability of the main technical section.

- The paper does not sufficiently demonstrate how the replay phenomenon would enhance exploration in broader RL settings. The experiments are restricted to a specific grid-world task, limiting the ability to assess the model's general applicability. Additional downstream tasks or diverse environments would help clarify the robustness and transferability of the proposed methods.
  
- Although the authors claim the model is modular, there is limited evidence showing how this modular framework could integrate effectively with mainstream model-free (value-based) algorithms. If would be helpful if the authors can discuss how and HF and PFC could improve some standard RL algorithms.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 1
questions: 1) In Figure 1, where does the replay trajectory in F come from? I realize it is a schematic, but it is somewhat confusing how to interpret it. 

2) In equations 2, 3, should I_replay be named something else, such as I_rest, otherwise it is sort of asserting the conclusion in the definition?

3) Does the r in W_{h, r} stand for replay? Or reward? I’m somewhat confused by what is going on there. 

4) What is the axis of 2E? What is the relationship between the x-axis of 2C and 2E? What is the color scheme in 2E? What are the vertical dashed lines in 2D? Why is step count negative in 2A (and why are the y-ticks covered up)?

5) ‘C2 checking times’ are never explicitly introduced as a metric, and it’s unclear what the means as an x-axis in 2E. 

6) I don’t know exactly what it means for the results in Figure 2 to prove that conditions in equation 1 are sufficient to generate replay. What are alternatives outcomes that could have happened that might not have demonstrated this? 

7) In 4A, what is ‘1st’ and ‘2nd’?

8) Is replay actually necessary for the results in Figs 3 and 4, or is it just a quirk of how this model was trained. In particular, this model was trained using RL in a setting where after reward was received, the communication channel opens up, so it will naturally find some way to use that communication channel. So disrupting that channel should disrupt behavior, but I'm not certain that is a particularly deep conclusion. 

9) My understanding is that this model is structurally designed to produce place cell activations. Thus, when it is unlinked from environmental input, and the ‘PFC’ activity yields place cell activations in ‘HF’, is there any alternative? Should the place cells just be silent? This seems like a fairly weak way to assert that replay emerges. Can we see more of an emphasis on the actual spatial trajectories?

10) If the weights are fixed, then what is the role of replay? This is suggesting a role for replay during in-context updating, but I’m a little confused how that relates to the idea of using replay for weight updates and Experience Replay. I was really excited at the beginning of the paper that it would show how replay-for-updating-synaptic-weights could emerge, but I don't see any discussion of that. 

11) Have any animal experiments been published that show that blocking communication between PFC and HF blocks replay? If so, it should be made extremely clear in the text. This should be a highly doable experiment using optogenetics. 

12) When testing the necessity of HF to PFC channel: is signal only replaced during replay periods? It’s not clear to me how disrupting the transmission from HF to PFC actually disrupts PFC, in 3A. Is it just because the state of PFC is disrupted, and because that state is disrupted, then the selected actions are also disrupted? 

13) I don’t follow what the implication is that shuffling the messages during replay has minimal effect. What is an example of an ‘independent package’?
soundness: 2
strengths: - The premise that replay may emerge naturally as a result of a broader reinforcement learning process is interesting. 
- The hypothesis that communication between the world model and policy is important for generating replay is interesting. 
- The connection to animal data is interesting.
summary: This paper proposes that replay emerges naturally during reinforcement learning as a result of communication between a ‘hippocampal’ world model and ‘prefrontal’ policy. To demonstrate this, the authors set up a network in which communication between the world model and policy occur only during ‘rest’ periods after a reward has been received. They test this on a simple animal-inspired task in which the path to acquiring reward changes. They show how the place cell activation changes over time to reflect the updated position, and that this roughly aligns with the dynamics of previously reported rodent replay. They further show that intervening on the communication channel from the world model to the policy disrupts the behavior, and that multiple steps of communication between the world model and policy appears to yield better behavior.
weaknesses: - The presentation, experiments, and analyses make it difficult to discern what to actually conclude about their hypothesis and premise, and I am not convinced to a degree that would be necessary for publication. See detailed questions below. 

- I would want much more detailed and convincing demonstration of the actual replay, perhaps closer to the neuroscientific representation such as in Carr et al. 2011, ‘Hippocampal replay in the awake state: a potential substrate for memory consolidation and retrieval’ (e.g. Fig 1), and in more straightforward replay settings (such as a linear track or W maze).

Meta-review: This paper explores the emergence of replay in reinforcement learning (RL) agents, drawing inspiration from the replay phenomenon observed in the brain. The authors propose a biologically-inspired model-based RL architecture and demonstrate that replay can emerge naturally without explicit programming.

Strengths
-----------
- **Novel perspective:** The paper offers a fresh perspective on replay in RL by investigating its emergence as a natural phenomenon rather than a hardcoded mechanism.
- **Interpretability:** The authors conduct analyses to understand the information flow within the model and the role of replay in adaptation to changing reward structures.

Weaknesses
--------------
- **Limited generalizability:** The experiments are primarily conducted in a single, albeit well-established, environment, raising concerns about the generalizability of the findings to other RL tasks and algorithms.

- **Lack of clarity:** Reviewers found the paper difficult to understand in parts, particularly regarding the definition and significance of "replay" and the explanation of certain figures and experimental procedures.

- **Limited significance for RL:** While the emergence of replay is interesting from a neuroscience perspective, its practical implications for improving RL algorithms remain unclear.

This paper presents an interesting exploration of the emergence of replay in RL agents, drawing valuable insights from neuroscience. However, it needs further development to address concerns about generalizability, clarity, and relevance to the broader RL community.

Decision: Reject