Modelling the control of offline processing with reinforcement learning
Authors: Eleanor Spens
Neil Burgess
Timothy Edward John Behrens
Venue: NeurIPS 2025 poster ()
Keywords: computational neuroscience
offline learning
reinforcement learning
hippocampal replay
generative replay
Primary area: neuroscience_and_cognitive_science
TL;DR: We propose a framework in which a meta-controller learns to coordinate offline learning in 'sleep' phases to maximise reward in an 'awake' phase, choosing between different actions which correspond to types of offline process in the brain.

Abstract: Brains reorganise knowledge offline to improve future behaviour, with 'replay' involved in consolidating memories, abstracting patterns from experience, and simulating new scenarios. However, there are few models of how the brain might orchestrate these processes, and of when different types of replay might be useful. Here we propose a framework in which a meta-controller learns to coordinate offline learning of a lower-level agent or model in 'sleep' phases to maximise reward in an 'awake' phase. The meta-controller selects among several actions, such as learning from recent memories in a hippocampal store, abstracting patterns from memories into a 'world model', and learning from generated data. In addition, the meta-controller learns to estimate the value of each episode, enabling the prioritisation of past events in memory replay, or of new simulations in generative replay. Using image classification, maze solving, and relational inference tasks, we show that the meta-controller learns an adaptive curriculum for offline learning. This lays the groundwork for normative predictions about replay in a range of experimental neuroscience tasks.

Modelling the Control of Offline Processing with Reinforcement Learning
Eleanor Spens' ellie. spensenden. ox. ac. uk
Neil Burgess?
n. burgess@ucl ac uk
Timothy Behrens
timothy. behrens ox ac.uk
University of Oxford
University College London
Abstract
Brains reorganise knowledge offline to improve future behaviour, with 'replay' involved in consolidating memories, abstracting patterns from experience, and simulating new scenarios. However, there are few models of how the brain might orchestrate these processes, and of when different types of replay might be useful. Here we propose a framework in which a meta-controller learns to coordinate offline learning of a lower-level agent model in 'sleep' phases to maximise reward in 'wake" phase. The ceta-controller selects among several actions, such learning from recent memories in hippocampal store, abstracting patterns from memories into a 'world model', and learning from generated data. In addition, the meta-controller learns to estimate the value of each episode, enabling the prioritisation of past events in memory replay, or of new simulations in generative replay. Using image classification, maze solving, and relational inference tasks, we show that the meta-controller learns an adaptive curriculum for offline learning. This lays the groundwork for normative predictions about replay in a range experimental neuroscience tasks.
1
Introduction
Brains are busy during rest, reorganising knowledge to improve future behaviour. In particular, rest improves the ability generalise, infer 'missing links', and learn statistical patterns 2 This is linked to replay, the hippocampal reactivation of experiences in temporally compressed form [3] Offline replay has several proposed functions: Firstly, it is thought reactivate real memories to consolidate them into neocortex [4]
3 or to abstract general patterns from episodes [5]. Secondly, replay can generate new experiences (6), sometimes referred as generative replay, which may help generalisation (7) and lessen catastrophic forgetting [ In addition, replay is thought to recombine conceptual 'building blocks' extrapolate beyond direct observation
77 Excitingly, recent studies show that the content of replay depends on the stage of learning [10 and the stage of cognitive development (II), suggesting variation in the brain's 'strategy" for offline learning over both short and long timescales. (See Appendix A.2 for additional neuroscience context.)
One might expect the optimal balance of these offline processes vary based on task demands, the uncertainty the environment, and the reliability of internal models. Machine learning involves similar trade-offs: collecting real-world data may be expensive, while synthetic data may not be reliable. Whilst 'data augmentation" is often used to expand datasets by specifying label-invariant transformations (e.g. adding noise, rotating, or cropping image in the case of image classification), identifying the right transformations can be process of trial and error. As a result, training data tends
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

to be curated and ordered based on human judgement. The fields of meta reinforcement learning [12] curriculum learning [3] and data valuation [4 15] explore solutions to some of these issues.
Drawing together these ideas, we propose a framework in which a meta-controller learns to select between different offline actions in a 'sleep' state to maximise reward in a 'wake' state. In the simulations of image classification and maze solving, these actions are i) learning from recent memories in a hippocampal store, ii) abstracting patterns from memories into a 'world model', and iii learning from data generated with this 'world model'. In the case of relational inference, the model has a more complex range of actions, including the abstraction of relational patterns, and the application of those patterns to infer new relationships. This framework could be used to make normative predictions about the types of offline processing and content of replay observed in neuroscience data. It also suggests how access to metacognitive actions', with which an agent can process own knowledge offlin,, can improve the performance of a reinforcement learning agent, as can the ability to construct custom curriculum for this learning.
1.1 Previous work
Reinforcement learning (RL), which has long used experience replay to improve performance, is a common framework in which to consider offline learning. In particular, Dyna [16] is an approach that uses real experiences update both a policy and learned model of the environment, which is then used to generate simulated experiences for additional updates, improving the sample efficiency of learning. However, Dyna and subsequent algorithms [17] rely on fixed heuristics to schedule real-world and synthetic updates. Meta RL refers to approaches in which a higher-level agent guides the learning of a lower-level agent [2] Our work extends these ideas to metacognitive control over offline processes in neuroscience context, and is related to the suggestion that the prefrontal cortex
functions as meta-controller
18
Many studies in the RL literature explore ways to prioritise experience replay, either at the level of episodes or individual transitions [19], e.g. based on reward, uncertainty, or prediction error (20). This often requires balancing selective replay with random sampling to maintain some diversity in the training data. Other studies suggest different prioritisation criteria for experience replay at different stages of learning [2I]. Mcanwhile neuroscience research suggests that memories are replayed in order of utility, i.e. the amount they would improve decision- making [22], or based on their longer-term value for 'map-building' (23] Note that methods for prioritising experiences are not only useful for replaying past data they can also guide which new experiences an agent should
gather next in online learning, e.g. to enable efficient exploration (24).
Curriculum learning refers to the arrangement of training examples into an order that benefits learning, e.g. in which the model is exposed to data of an appropriate difficulty level at each stage [K] Meanwhile data valuation is a growing field of applied machine learning which explores how to select the most useful samples with which to train a model, e.g. by estimating the marginal contribution of a given example [15], or by training a network select the best examples via RL [14). The data valuation aspect of our framework can be thought of in terms of learning a curriculum for offline learning.
2
The model
In each simulation, we consider alternating sleep and wake states (Figure [p). During the wake state, the meta-controller receives new observations (images, maze navigation episodes, or observed relationships) from the environment which are stored in its model hippocampus. In addition, reward is obtained reflecting performance on some task (image classification, maze navigation, or relational inference). During the sleep state, the meta-controller takes actions which correspond to different ways to process offline knowledge, gradually learning to act offline in a way that maximises the awake reward.
More formally, consider a meta-controller which co-ordinates the offline learning of a lower-level RL agent. For each sleep/wake cycle = I....,N:
1. Sleep phase: A meta-controller takes T offline actions 4, using its recurrent proximal policy optimisation (PPO; (25)) policy To(a 2),, where x is an observation capturing the state of

---

b)
Wake state: encode new experiences
Updase
on simulations
Update world model
Train on roal examppee
Train on imagiration
episndea
D
Select epiendes to reply
Select epleades simulate
Select images to raplay
Select imgge to imagiiee
Sleep state: procees experiences
Figure I: The model. a) In all simulations, a meta-controller learns to select from a range of possible offline actions during sleep, receiving reward only at the end of the episode based on performance in the wake state. b) The maze task involves learning to navigate to a goal in a gradually changing maze. The meta-controller chooses between training lower-level agent on real episodes from the hippocampus, updating its 'world model' of the maze's 's transition statistics from the hippocampus, and training the lower-level agent on episodes simulated with the world model. In addition, the system learns which episodes are most valuable to replay or to simulate. c) In the image task, images encoded in the hippocampus during the wake state, where the distribution of classes varies, and the goal is to learn to classify them accurately. The meta-controller chooses between training a classifier real images from the hippocampus, updating simple generative (Gaussian mixture) model on images from the hippocampus, and training classifier on generated images. In addition, the system learns which images are most valuable to replay or generate.
the system. These actions involve offline processing of current memories and knowledge, without interacting with the real environment.
2. Wake phase: The agent tested on the task for frest episodes, recording the evaluation reward Rn, L Diee I'm where T' is the total reward in episode m. The meta-controller's policy is Tipu updated based on this reward.
The agent interacts with the environment- which may have changed since the previous wake phase for K steps, encoding memories H, {(sk-1, 1, 4t-1, 8k, rk) rk)) k= in the model hippocampus (here simply list).
The lower-level task can be an arbitrary supervised learning task rather than an RL one by replacing the references to episodes above with observations and targets, and storing memories as H, {(observation, target,))
The system also learns which items are most valuable to replay or generate, as follows. Truth data for the value is obtained for a small subset of candidate items by calculating their utility with a Shapley value approach, which estimates the change in performance if an example were omitted [5] Then small model is trained on this subset to predict the value of items (with the model fine-tuned each time data is selected rather than trained from scratch so that value estimates improve over time). Crucially, the data valuation procedure estimates the value of an example for subsequent learning contingent on the state of the current model, allowing to capture the value of examples at different stages of training; this differs from some existing approaches which compute Shapley values with respect to a surrogate model trained from scratch, c.g. (26). Given predicted values for all candidates, items are selected using maximal marginal relevance (MMR; (27), which balances selecting the most valuable examples with promoting variety. Note that data valuation can be used to either undersample or oversample from set of candidates. The data valuation component can be trained at the same time as the meta-controller. or separately for efficiency as shown here, in which case random sampling is used during meta-controller training. Whilst the meta-controller is frozen during inference, the value estimator network continues to be updated each time memory or generative replay is selected.
Our experiments used Stable-Baselines3 28 for the RI. algorithms, Gymnasium 29 to create the
custom environments, and PyTorc 30 for other neural network training. See Appendix A. for further details.
3

---

3 Results
3.1
Image task
At the start of each episode 200 new images are stored in a buffer representing the hippocampus. The reward in the wake state, after three offiine steps in the asleep state, is the accuracy of a classifier. Fashion-MNIST (a more challenging variant of MNIST featuring images of ten items of clothing) is used toy dataset [3]]
The meta-controller agent is a recurrent PPO agent, enabling it to learn an optimal sequence of offline actions. The agent's inputs are the number of steps since the sleep phase began, the fraction of observations per class in the buffer, and the most recent action. The action space consists of the following actions (see Figure T
Update the image generator with stored images from the model hippocampus. Specifi- cally, the image generator is a Gaussian mixture model from which samples can be drawn for particular class.
Action 0:
Action I: Train the classifier on stored images from the model hippocampus, selecting the 50 most useful 'memories' to replay with the value estimator model.
Action 2: Train the classifier on generated images. As above, the 50 most useful items for 'generative replay' are selected using the value estimator model.
Note that we assume a small validation set exists for the data valuation element. To describe this more concretely, if the hippocampus contained many images, the marginal contribution of each of a subset (here, 25%) of images to the classification accuracy on small validation set would be obtained. Then the value estimator (a linear regression model) would learn to estimate these values. Predicting the value of the many candidate items could then performed efficiently with the value estimator, with these estimates used to prioritise which items to replay. (But in principle this validation set could be generated by the generative model, as we demonstrate for the maze task.)
3.1.1 Image results
Catastrophic forgetting or interference occurs when newly learned information overwrites previously learned information neural network, leading to poor performance on the previously learned tasks
33.. Continual learning is the ability to learn series of tasks sequentially without the occurrence of catastrophic forgetting 33), and poses a general challenge for connectionist models of the brain. The training data for machine learning problems is carefully curated, e.g. to balance the number of examples of different categories, and shuffled before training. Reality is far messier, so biological learning must be robust to changing distributions over time.
To explore whether the agent can learn strategy for continual learning, we compare a balanced and an imbalanced experience case. In the balanced experience case, examples from all ten Fashion- MNIST classes are encoded in the hippocampus in each wake state, and the model's accuracy on all ten classes is returned as feedback. In the imbalanced experience case, examples from two classes are encoded in the hippocampus in each wake state, and the model's accuracy on all ten classes is returned as feedback. This is sometimes referred to as incremennllllttal learning in the continual learning literature [34]. (Note that the image classification, generation, and valuation models are re-initialised when all categories of the dataset have been encoded, i.e. every five episodes in the imbalanced case.) We find that the balanced experience case, replaying from the hippocampus is optimal (Figure P)), but in the imbalanced experience case, training the generative model then learning from simulations is (Figure 24. Unlike the basclines in Table the meta-controller approach can do well in both settings, using generative replay reduce catastrophic forgetting where necessary. (Note that the Gaussian mixture model learns a distribution per class, so in generative replay an equal number of samples are taken per class, from all classes added to the model so far.)
The behaviour of the data valuation component of the model was explored with i) the original Fashion-MNIST dataset (Figure 21 and g), and ii) a noisy version of the same dataset in which a randomly chosen fraction of pixels (25% 50%, 75% or 100%) was set to zero in each image (Figure 2 and f), intended as simple check that data valuation improved performance. Both variants were tested in the class-balanced case, hard-coding the actions to three stages of either memory replay or generative replay. That is, we enforced three consecutive actions and recorded the estimated values

---

of each image at each stage of memory or generative replay. Figure and f show that selecting the images based on value estimates plus MMR results in significantly better performance than random selection, for memory replay and generative replay respectively. Table 2 shows it performs better than several other baselines too.
Figure 2 and g show that novelty (the distance from the mean of the class) is negatively correlated with estimated value in the first step of learning, but positively correlated by the third, for both memory replay and generative replay. One way to interpret this is that novel examples are less useful at first because they are too challenging, as observed in curriculum learning [3]. but become more useful over time. We visualise value in the UMAP projected datasets in Figure and also corroborate this finding with MNIST (Figure
D and h).
6. Action probabilities over CL setting
b)
probabiltties over setting
Step eptiode
Step ectoode
Step eptoode
Shep of aphante
Shep of sphasde
Shep of sphant
1.0
4003
2031 4991 Teeey spade
4990 thringg spodd
400 wruik rytadd
400
66
200
420
(6)
Tuetin gtodd
wwain relocd
wwaing re-ooe
Nemory replay
Update WM
Generative replay
replay
Generative raplay
Learning curriculum memory raplay
Harder examples become more useful
memory replay
MNIST memory replay
Early episode: examples more examples useful
oplande: harder
High
High
value
2
Low value
Low
Step
Steo
Step
Stea
0
Learning curriculum for generative raplay
Harder examples become more useful over time in generative roplay
h)
MNIST generative raplay
episode: easier
Late episode:
examples useful
examples
High value
I
High
value
Low value
Low value
Bardonn
0
Data value by step
UMAP projection of labelled memory replay value
UMAP projection generated data by generative replay
Step
Step
Figure 2: Image simulations. a) Action probabilities for each position in the sleep sequence over the course of training in the continual learning (CL) setting, in which the distribution of the agent's s experiences is imbalanced, with only two classes observed in each wake phase. Shading indicates the SD across a sliding window of 100 episodes. (In other words, step 0 is the first action chosen in the sleep phase, step is the second, and so on.) b) Action probabilities for each position in the sleep sequence over the course of training, but in the balanced experience baseline case. c) Learning to select the images for memory replay results in significantly better performance than random selection. The means across 10 repetitions with the addition of random noise (setting a variable fraction of pixels to zero), with a balanced class distribution (i.e. the non-CL setting above). d) Here we take the original Pashion-MNIST dataset (no noise), and look at the relationship between novelty (distance from the class mean) and value for three consecutive stages of memory replay. At first the correlation is negative, indicating that easier examples are more valuable. But by stage three, the correlation is positive, indicating that harder examples are more valuable. e) Repeating the experiment part d but for the MNIST dataset displays a similar pattern. f-h) As in parts c-e but for generative replay. Note that three steps of generative replay, following the training of the 'world model', are shown for consistency with the memory replay results. i) A projection into 2D with UMAP of the Fashion-MNIS" data and their value estimates for memory replay (left), and of the generated data and their value estimates for generative replay (right).
shown, and error bars give the SEM. Pashion-MNIST is used
5

---

Table 1: Image task results by approach. The meta meta-controller learns use generative replay in the continual learning (CL) setting but memory replay in the non- mn-CL setting. Results are averaged across 25 meta-episodes in each setting.
Approach
Acc. (CL)
Acc. (non-CL)
Mean acc.
SEM
Meta-controller
0.307
0,411
0.359
0.010
Memory replay only
0.171
0.411
0.291
0.003
Generative replay only
0.307
0.355
0.331
0.011
Table 2: Data valuation network for the image task compared to other data selection baselines, for choosing data to replay from the model hippocampus (memory replay; MR) simulate (generative replay: GR). Ten meta-episodes are tested in each case. Diverse random selection uses the MMR method but with random data values. For the most and least challenging image baselines, the difficulty of an image is quantified as the fraction of an ensemble of scikit-learn classifiers that labelled the image correctly.
Approach
Mean acc. (MR)
SEM (MR)
Mean acc. (GR)
SEM (GR)
Data valuation network Random selection
0.438
0.007
0.373
0.007
0.307
0.017
0.290
0.013
Diverse random selection
0.378
0,012
0.273
0.018
Most challenging
0.145
0,014
0.300
0.009
Least challenging
0.364
0.009
0,300
0.009
3.2 Maze task
Next we model a maze solving task which the maze gradually changes over time, based on rodent neuroscience task 335 Here the higher-level agent is recurrent PPO agent as above, but the lower-level agent for navigating from random start to random reward location in the environment is Deep Q-Network (DQN; (36) agent. The inputs to the DQN agent one-hot encoded representations of the current and goal locations. The maze is simply a six-by-six grid where subset of squares are barriers (but all non-barrier squares are connected). In the wake phase, the maze is updated by adding one barrier square and removing another (see Figure oa. Appendix (A.3). This means that the lower-level agent must update its strategy by replaying or simulating trajectories in the new maze.
In the wake phase the hippocampus receives 50 trajectories in the maze, before the meta-controller takes three actions the sleep phase. The reward in the next wake phase is the performance of the lower-level maze solver agent on the same maze. The agent's inputs are the number of steps since the sleep phase began, the current accuracy of the world model, and the most recent action. The action space consists of three actions (see Figure [D), closely analogous to those for the image classification task:
Action 0: Update the world model from memories. This is simple network for predicting the next state given a state and action, which sufficient to allow the simulation of new episodes. (In deterministic environment, the world model could simply cache the result of the most recent observation of a state-action pair.)
Action Train the maze solver on stored episodes from the model hippocampus (50 episodes
oversampled to 200).
Action 2: Train the maze solver on 200 generated episodes, in which the agent acts within the world model, i.e. epsilon-greedy rollouts are simulated.
In addition, the system learns which episodes are most valuable to replay to simulate, similarly to the approach in the image task described above. Specifically, the Shapley values of a subset of the candidates (50 out of 576 possible start-goal combinations, in the case of generative replay) are obtained, using the world model to estimate the improvement to the agent. Then a value estimator predicts the utility of particular start-goal pair. Unlike in the image case, this does not require

---

real validation set, as the world model can be used to quantify the benefit. That is, the system can simulate how well variant of the model trained on subset of the training data would do.
3.2.1
Maze results
Figure 3h shows the learned schedule with which the meta-controller improves the solver; if only a few sequences are stored, it is more effective to learn the transition structure the maze then simulate episodes than it is to learn from the limited data in the hippocampus. The meta-controller's "success rate' of reaching the goal is higher than action selection and data selection baselines (Tables aandd
Even though the agent is tested only on the current maze, its learned strategy of training a generative model then learning from simulations resembles that of the image agent in the imbalanced case. This is because RL agents perform poorly with very limited training data, so learning from the hippocampus alone is not feasible. In addition, the world model in the quality samples than the world model in the image case, so there is less of a trade-off associated with generative replay.
task generates higher
Start-goal pairs for which the shortest path was changed in the most recent maze update are higher value in both stages of generative replay (Figure 3p) but the relationship between shortest path length and value varies with learning (Figure 3d.. We can also inspect the average estimated value maze trajectories passing through each point in the maze. In some cases, locations nearby recent change are more valuable at first, but learning saturates in later stages (Figure Bb), though these effects are not significant (Figure 3
Action probabiltties over time eploode
b) value
Step up sode
Step spasode
1.8
X
0.0
Tre code
ring sode
Tox ring spodd
Memoory replay
Update WM
Generative replayy
Addectias
Removed
Value paths
Value length
Value by praximity
Step
Step
Step
Step
0.0044
00925
$.800
00082
04832
5.000
00804
areax
2.000
-E4031
831
-0.000
Changed
Unchwnnget
Ohanged
Unchinged
0,0000
Helr
FM
Figure 3: Maze simulations. a) Action probabilities for each position in the sleep sequence over the course of training. (In other words, step 0 the first action chosen in the sleep phase, and so on.) The agent learns to update the world model then simulate new episodes with it. Shading indicates the SD across a sliding window of 100 episodes. b) Visualising the learned values of different locations in the maze. Each square shows the mean estimated value of all candidate episodes which pass through that point. The most recently added and removed squares are indicated. e) The mean-centered value of start-goal pairs for which the shortest path is changed vs. unchanged, for two stages of generative replay. Error bars give the SEM. d) The mean-centered value of start-goal pairs by their shortest path length, for two stages of generative replay. The mean value of locations near vs. far from maze changes.
3.3 Relational knowledge model
Biological intelligence excels at generalising from relatively few examples. In our framework, can the agent learn how to infer as much information as possible offline from small number of online observations? We trained and tested the model on two relational inference tasks: spatial inference, and family tree inference, as in 377, although the approach is applicable to any scenario in which observations can be expressed edges in graph with an underlying structure. In the family tree
7

---

Table 3: Maze task results for the meta- controller's s action selection compared to memory replay only baseline (with random data selection for both), averaged across ten incremental maze updates for three different random seeds.
Table 4: Maze task results for the data valu- ation method compared to baselines (all with meta-controller action selection), averaged across mazes and seeds as in Table
Approach
Success rate
SEM
Approach
Success rate SEM
Data valuation
0.871
0.010
Meta-controller
0.822
0.007
Random selection
0.822
0.007
Memory replay
0.161
0.005
Longest paths
0.736
0.013
inference task, the goal is predict many relationships as possible from subset of family relationships (see Figure 4b). while in the spatial inference task, the goal is to predict the relative directions of locations in 2D grid (i.e. whether each location is east / west / north / south of every other location).
In the wake phase the hippocampus receives 20 observed relationships, corresponding to a few of the edges of a true underlying graph. The agent interacts with a relational graph each environment, iteratively updating its knowledge by choosing various meta-actions. In the previous simulations, the world models were used to generate new items in isolation with which to train a model to perform a given task. Here the model trained offline has a different function, inferring new facts to flesh out partial observations. Specifically,, the world model is graph autoencoder, with graph convolutional layers which allow it to learn shared patterns across graphs. Its decoder predicts missing edges in graph, allowing the agent to infer new relationships. (Note that for training efficiency the model is cached, rather than retrained repeatedly.)
The meta-controller, recurrent PPO agent as in the previous simulations, receives a three-dimensional observation vector consisting of the number of relationships learned so far, the number of edges in the current graph, and the accuracy of the world model. The action space consists of six discrete meta-actions:
Action O: Learn the task (i.e. add relationships to list of known facts) from the model hip- pocampus. Here the task is trivial, but this could be some other operation the knowledge, e.g. identifying the youngest members of the family, or finding the shortest path to a goal location.
Action 1: Transfer (or 'consolidate') observed relationships in the model hippocampus into the knowledge graph.
Action 2: Learn the task (i.e. add relationships to list of known facts) from the knowledge graph.
Action Train the world model on all instances of graphs that have been built so far during meta-training.
Action 4: Expand the knowledge graph using predictions from the world model. For example, the world model has learned that parent of a parent of X a grandparent of X. applying the world model to the graph would add inferred grandparent edges to a family tree.
Action 5: Do nothing. This avoids a penalty of 0.2 for all other actions, to incentivise efficient behaviour, and reflects the cognitive cost of taking the other actions.
The episode ends when either 10 meta-steps have been reached or random stopping probability (0.05) triggers termination. At this point, the reward computed based on the number of correct relationships inferred by the model compared to the true graph. The data valuation aspect is omitted in this simulation for simplicity.
3.3.1 Relational knowledge results
The meta-controlle learns schedule for offline actions to maximise online reward (Figure and d). the spatial task, the agent first learns from memories in the hippocampus, which is desirable because the episode can stop at any point, and this ensures at least some reward is obtained. Then the observed relationships are transferred (or 'consolidated") into the knowledge graph, as this prerequisite for relational inference to take place. After this, the world model is updated. (Note
8

---

Diagram of model
b)
Action space
Learn from observations
Consolidate abservations into graph
Update world model
Expand graph with world model
Learn from Inforred graph
ANAEO of
Encode observed relationships
Process these observations
Evaluate inference performance
or
Spatial task
d)
Family tree task
Troiing rnllade
Thaiing neteee
Thaning myneee
Taainng rehaod
Triing rabule
Vine
Ming Pippnalle
Mapp Tyoudds
Nirp Partace
Wia Faeact
L9G0
1990 1400
500 2008
5R0 ctice 2018
100 1997 1980
infened
model
Learn infend
moce Sovnating
Convotiaat
prch
Oreard pash
nerting
Cravitass dieeeddo nio groc
Courd pich
BOOKE
Figure 4: Relational knowledge simulations. a) Diagram of the model. At the start of each episode, twenty observations from new graph added to the hippocampus. The meta-controller learns to choose between several actions for expanding this knowledge, taking up to ten steps. The reward is the number of edges in the true graph which be inferred by the end of the episode. b) The possible actions are: learning a given task from the observations in the model hippocampus, 'consolidating' observations from the hippocampus into graph, updating a world model that captures relational structure, expanding the graph with new edges based on this world model, and learning a given task from the inferred graph. (Here the task simply build up a list of facts, but this could be some other operation on the knowledge, e.g. identifying the youngest members of the family, or finding the shortest path to goal.) c) Above: Mean reward and action probabilities over the course of training for the spatial model, where the probabilities are plotted separately for each position in the sleep sequence. Shading indicates the SD across sliding window of 100 episodes. Below: The final frequencies of actions chosen at each position, calculated across 20 episodes. Error bars give the SEM. Note that only the first five steps are shown. d) Likewise but for the family tree task.
that in this simple demonstration, the world model must be relearned in each episode, but one would expect this choice to be modulated by the accuracy of the world model if it were not reset.)
Next, the meta-controller uses the trained world model to infer additional relationships in the graph, e.g. inferring the fact that X is west of Z from the facts that X is west of Y and Y is west of Z. This action can be applied multiple times, iteratively expanding the graph. With each application of the graph convolutional autoencoder, links are predicted with a given probability, and those with probability greater than 0.5 are added (see Figure Appendix A.3 for an example of the iterative expansion of family tree graph from a few edges). The chance of Tearning from the inferred graph ramps up later in the episode, as does the chance of doing nothing, as this avoids the small penalty for all other actions. The family tree task displays a similar pattern.
Discussion
We trained RL agents to perform metacognitive actions in an offline 'sleep' state in order to improve the "awake" capabilities of lower-level agent / model. The actions correspond to different processes

---

observed in biological offline learning, namely learning directly from memories, the training of predictive models, and the simulation of new events.
We found that the agents learned an effective curriculum for offline learning. In the maze navigation task, the agent learned first update world model capturing transition statistics, then simulate episodes with this model. The learned strategy depended on the demands of the task. In an image classification scenario requiring continual learning, the agent updated its generative model after receiving new observations, then learned from generated samples. Whilst these samples were lower quality, generative replay improved reward in the wake state by reducing catastrophic forgetting. But in the case that the agent observed all relevant categories in each wake state, this trade-off off was no longer worthwhile, with the agent reverting to learning from real images in the model hippocampus. This could make predictions about how the brain regulates the replay of recent and remote examples to enable lifelong learning. The framework flexible that other metacognitive actions can easily be added, as demonstrated in the relational inference simulations, in which the agent learned to extract statistical patterns and then apply these to iteratively infer new facts.
In our simulations, the meta-controller also learns to choose which events to replay or simulate. For example, in the image task, easier examples (those closer the mean of the class) were selected earlier in learning, whereas harder examples were selected later in learning. This was observed for both memory and generative replay. One might expect this be true over short timescales (within a single block of sleep or rest), and over longer timescales while learning is gradually acquired. This lays the groundwork for making normative predictions about the optimal schedule of replays for learning, which could be compared to experimental neuroscience data. Previous work suggested that offline replay of memories regulated by their noise level, so that noisy memories do not impair neocortical generalisation (38]. On the other hand, novelty is proposed to promote memory encoding [3). Our approach could reconcile these different proposals about the relationship between schema congruence and replay by explaining how the optimal curriculum varies over the course of offline learning.
This work has several limitations, which reflect the fact that this is a very simplified model. Firstly, there are many variables that affect the behaviour of the data valuation component, so thorough tests of the sensitivity of the results to these parameters should be performed before we draw any general conclusions about which episodes are optimal to replay or simulate. Secondly, the data valuation aspect is computationally expensive, especially when fine-tuned in each stage of memory or generative replay. It would be more powerful for the data valuation network to learn to predict value based partly information about the current progress of learning. Then the network could predict different values for the same items of data at different stages of training, without the need to fine-tune the model. Thirdly, a drawback of the current Shapley value-based data valuation mechanism in the image task is that it assumes a small validation set available to the brain. The model could be refined to avoid this assumption, c.g. by using a subset of items in the hippocampus or generated data as a validation set (the latter is demonstrated in the maze task).
There are several other directions for future research. Here we only consider how offline processes are orchestrated, but an agent could also learn the optimal way to encode and retrieve memories. Such an agent might learn to only encode memories when it is worth the cognitive cost to do so (see also (40)) ). Similarly, one could integrate a mechanism for learning abstractions, e-g. 'options' in the maze navigation task. If the feedback in the wake state included penalty for the cognitive cost of behaviour, so that solutions that combined a few simple components were advantageous, this might incentivise the meta-controller to develop 'primitives'.
In conclusion, we present a novel framework in which cont-conr chooses between different metacognitive actions during 'sleep', inspired by the processes observed in resting brains. The meta-controller orchestrates offline learning to optimise 'wake" performance, and learns a curriculum for what to replay or simulate.
10

---

Acknowledgements
Thanks Rachel Swanson, Kris Jensen, Xiao Qin, Diksha Gupta, Mathias Sablé-Meyer, and Thomas Akam for helpful discussions.
References
[1] Jeffrey M Ellenbogen, Peter T Hu, Jessica D Payne, Debra Titone, and Matthew P Walker. Human relational memory requires time and sleep. Proceedings of the National Academy of Sciences, 104(18):7233-778, 2007.
[2] Simon Durrant, Charlotte Taylor, Scott Cairney, and Penelope A Lewis. Sleep-dependent consolidation of statistical learning. Neuropsychologia, 49(5):1322-1311, 2011.
[3] David Foster. Replay comes of age. Annual review of neuroscience, 40:581-600 1-282, 2017.
[4] Matthew A Wilson and Bruce I. McNaughton. Reactivation of hippocampal ensemble memories during sleep. Science, (55(5572):666-666 1994.
[5] James L McClelland, Bruce McNaughton, and Randall C O'Reellly Why there are comple- mentary learning systems in the hippocampus and neocortex: insights from the successes and failures connectionist models of learning and memory. Psychological review, 102(3):419, 1995.
[6] Anoopum Gupta, Matthijs AA Van Der Meer, David Touretzky, and A David Redish. Hippocampal replay is not a simple function of experience. Neuron, 65(5):695-705, 2010.
[7] Philipp Schwartenbeck, Alon Baram, Yunzhe Liu, Shirley Mark, Timothy Muller, Raymond Dolan, Matthew Botvinick, Zeb Kurth-Nelson, and Timothy Behrens. Generative replay underlies compositional inference the paapocampal-prffroonn circuit. Cell, 186(22):4885- 4897, 2023.
[8] Gido M van de Ven, Hava Siegelmann, and Andreas S Tolias. Brain-inspired replay for continual learning with artificial neural networks. Nature communications, 11(1):1-14, 2020.
[9] Jacob JW Bakermans, Joseph Warren, James CR Whittington, and Timothy EJ Behrens. Con- structing future behaviour in the hippocampal formation through composition and replay. bioRxiv, pages 2023-04, 2023.
[10] Jianxin Ou, Yukun Qu, Yue Xu, Zhibing Xiao, Tim Behrens, and Yunzhe Liu. Replay builds an efficient cognitive map offline to avoid computation online. bioRxiv, pages 2025 01, 2025.
[11] Luoyao Pang, Yukun Qu, Jianxin Ou, Shuqi Wu, Yuejia Luo, Chetan Gohil, Mark Woolrich, Tim Behrens, and Yunzhe Liu. Development of replay-dmn coordination predicts entorhinal grid codes. bioRxiv, pages 2025-04, 2025.
[12] Jane X Wang, Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel 7. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint /6 /.05/66 2016.
[13] Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 41-48, 2009.
[14] Jinsung Yoon, Sercan Arik, and Tomas Pfister. Data valuation using reinforcement learning. In International Conference on Machine Learning, pages 10842-10851. PMLR, 2020.
[15] Amirata Ghorbani and James Zou. Data shapley: Equitable valuation data for machine learning. In International Conference Machine Learning, pages 1242-2251, PMLR, 2019.
[16] Richard Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 166-16,, 1991.

---

[17] Jessica B Hamrick. Analogues of mental simulation and imagination in deep learning. Current Opinion in Behavioral Sciences, 29:8-16, 2019.
[18] Jane X Wang, Zeb Kurth-Nelson, Dharshan Kumaran, Dhruva Tirumala, Hubert Soyer, Joel Leibo, Demis Hassabis, and Matthew Botvinick. Prefrontal cortex retf-reiiforrement learning system. Nature neuroscience, 21(6):860 868,, 2018.
[19] Tom Schaul, John Quan, loannis Antonoglou, and David Silver. Prioritized experience replay, arXiv preprint arXiv: 151.05552 2015.
[20] Andrew W Moore and Christopher Atkeson. Prioritized sweeping: Reinforcement learning with less data and less time. Machine learning, 13:103-130, 1993.
[21] Hu Li, Xuezhong Qian, and Wei Song. Prioritized experience replay based on dynamics priority. Scientific Reports, 14(1): 16614, 2024.
[22] Marcelo Mattar and Nathaniel Daw. Prioritized memory access explains planning and hip- pocampal replay. Nature neuroscience. 21(11):1609-1617, 2018.
[23] Yotam Sagiv, Thomas Akam, Ilana Witten, and Nathaniel D Daw. Prioritizing replay when future goals unknown. bioRxiv, 2024.
[24] Changmin Yu, Neil Burgess, Maneesh Sahani, and Samuel Gershman. Successoo-predeccssor intrinsic exploration. Advances in neural information processing systems, 36:73021-73038, 2023.
[25] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint wrXiv:1707.0634/4 2017.
[26] Bojan Karlaš, David Dao, Matteo Interlandi, Bo Li, Sebastian Schelter, Wentao Wu, and Ce Zhang. Data debugging with shapley importance over end-to-end machine learning pipelines. arXiv preprint arXiv: 2022.
[27] Jaime Carbonell and Jade Goldstein. The use of mmn, diversity-based reranking for reordering documents and producing summaries. In Proceedings of the 21st annual international ACM SIGIR conference on Research and development in information retrieval, pages 335-336, 1998.
[28] Antonin Raffin, Ashley Hill, Maximilian Emestus, Adam Gleave, Anssi Kanervisto, and Noah Dormann. Stable baselines3, htttss://githhii www.coom//////
www.bbleeebaeeeeesee 2019.
[29] Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint arXiv:2407.17032. 2024.
[30] A Paszke. Pytorch: An imperative style, high-performance deep learning library. arXiv preprint 19Xiv:1912.01703, 2019.
[31] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist. a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv: 1000::77700000 1708.07777 2017.
[32] Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. In Psychology of learning and motivation, volume 24, pages 165. Elsevier, 1989.
[33] Raia Hadsell, Dushyant Rao, Andrei A Rusu, and Razvan Pascanu. Embracing change: Con- tinual learning in deep neural networks. Trends in cognitive sciences, 24(12):1028-100 1040, 2020.
[34] Gido M Van de Ven, Tinne Tuytelaars, and Andreas S Tolias. Three types of incremental learning. Nature Machine Intelligence, 4(12):185-11177 2022.
[35] Simões Godinho. Hippocampal representations a complex route planning task. PhD thesis, University of Oxford, 2024.
12

---

[36] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: 131.5500, 2013.
[37] James CR Whittington, Timothy H Muller, Shirley Mark, Guifen Chen, Caswell Barry, Neil Burgess, and Timothy EJ Behrens. The nolman-eichenbaum machine: Unifying space and relational memory through generalization in the hippocampal formation. Cell, 183(5):1249- 1263, 2020.
[38] Weinan Sun, Madhu Advani, Nelson Spruston, Andrew Saxe, and James E Fitzgerald. Organiz- ing memories for generalization in complementary learning systems. bioRxiv, 2021.
[39] Michael Hasselmo, Bradley P Wyble, and Gene Wallenstein. Encoding and retrieval of episodic memories: role of cholinergic and gabaergic modulation in the hippocampus. Hippocampus, 6(6):693-708, 1996.
[40] Qihong Lu, Uri Hasson, and Kenneth Norman. A neural network model of when to retrieve and encode episodic memories. elife, 11:e74445, 2022.
[41] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay. Scikit-learn: Machine learning in Python, Journal of Machine Learning Research, 12:2825-2830, 2011.
[42] David Marr. theory for cerebral neocortex. Proceedings of the Royal society of London. Series B. Biological sciences, 176(1043):161-234, 1970.
[43] David Marr. Simple memory: a theory for archicortex. Philosophical Transactions of the Royal Society of London. B, Biological Sciences, 262((41):2338)) 1971.
[44] Timothy J Teyler and Pascal DiScenna. The hippocampal memory indexing theory. Behavioral neuroscience, 100(2):147, 1986.
[45] Gabrielle Girardeau, Karim Benchenane, Sidney I Wiener, György Buzsáki, and Michaël B Zugaro. Selective suppression of hippocampal ripples impairs spatial memory. Nature neuro- science, 12(10):1222-1223, 2009.
[46] Philippe Peigneux, Steven Laureys, Sonia Fuchs, Fabienne Collette, Fabien Perrin, Jean Reggers, Christophe Phillips, Christian Degueldre, Guy Del Fiore, Joel Aerts, al. Are spatial memories strengthened in the human hippocampus during slow wave sleep? Neuron, 44(3):535-545, 2004.
[47] Freyja Ólafsdóttir, Caswell Barry, Aman B Saleem, Demis Hassabis, and Hugo Spiers. Hippocampal place cells construct reward related sequences through unexplored space. Elife, 4:c06063, 2015.
[48] Federico Stella, Peter Baracskay, Joseph O' O'Neill, and Jozsef Csicsvari. Hippocampal reactiva- tion of random trajectories resembling brownian diffusion. Neuron, 102(2)450-441, 2019.
[49] Yunzhe Liu, Raymond Dolan, Zeb Kurth-Nelson, and Timothy EJ Behrens. Human replay spontaneously reorganizes experience. Cell, 17463):640-652, 2019.
13

---

A Appendix
Code for all simulations can be found
at https://githhi: cttpb//coclellllee as/ Simulations were on Linux virtual machines with NVIDIA A100 GPUs, and on MacOS with the MPS backend for GPU support.
mt-wiab-metacoonittvee-eo actioonss
Diagrams were created using BioRender.com, with some icons created using ChatGPT.
A.1 Further model details
A.1.1 Proximal policy optimisation
Actor-critic methods are family of reinforcement learning algorithms which jointly learn both a policy (the 'actor") ) and a value function (the 'critic'), where the critic's value predictions guide the learning of the actor.
Proximal policy optimisation 22 is one such actor-critic algorithm, which adjusts the policy in small steps to avoid destabilising training.
In PPO, batches of experience collected with the current policy Toaa are used to update the parameters 0 of the policy by maximising clipped objective:
L CLIP (0) Er min (r)(0) Å., -ip(r:(0),1,,11123) An)
where:
no (a: 8t)
r+())=
the ratio of the new and old action probabilities at time t.
Ta (az st)
A, is an estimate of the advantage. i.e. how much better the chosen action a, compared the average.
is a small clipping parameter (e.g., 0.2) that limits how far rx can deviate from 1.
The clipping prevents large policy updates that could reduce performance.
To learn the right sequences of offline actions, we use a recurrent PPO network so the agent can condition its choices the history of previous steps within the same 'sleep' phase. We use the default LSTM network (but with one hidden layer of 64 units) provided for recurrent PPO agent in Hable-Baselines,, (8
A.1.2 Estimating value via Shapley values
To teach the agent which images trajectories are most useful to replay, we use small validation set and compute approximate Shapley values [15] for subset candidates. The Shapley value &; for example measures its average contribution to performance across all subsets of examples. Formally,
|S!! - - 1)! (c(SU x(S)),
|N!
SEN
where v(S) is the classifier accuracy (or maze solver performance) trained subset S. The second term therefore measures how much adding example after subset increases performance. The first term gives the probability that, if you shuffle all N examples into random order, the set of items that appear before is This sum is therefore the expected marginal contribution item i when you add it into randomly shuffled training set.
Because calculating this exactly is very expensive, Shapley values can be approximated by sampling random subsets of the candidates, with and without particular pair. In the image experiments, we use standard data Shapley: for each training image, we estimate value by sampling coalitions of images, training the classifier each coalition with and without the target image, and averaging the
14

---

marginal accuracy gains. In the maze experiments, we value all candidate start-go -goal pairs in single pass per permutation for efficiency. For each sampled ordering, we sequentially train on each pair and record the marginal performance gain it produces, then we average across all permutations as in standard data Shapley.
These estimated Shapley values for a subset of the candidates are used to train value estimator model which quickly predicts value for every candidate. In the image task, the input the value estimator is the image, whilst in the maze task, the input to the value estimator is the start-goal combination, represented as co-ordinates of the start concatenated with co-ordinates of the goal, and the distance between the start and goal.
A.1.3 Selecting examples with maximal marginal relevance
To pick a set of items (real or generated) that are both informative and diverse, we use maximal marginal relevance [27], a common approach for information retrieval. We add examples one at time to a set at each step scoring candidate item by:
MMR(d) Relevance(d) (1 A) max Similarity(d,d))
d'ES
where:
Relevance(d) is how valuable the example is predicted to be (from our value estimator), Similarity(d, d') measures how similar two items are (e.g. cosine distance between the pixel vectors in the image task, or the sum of distances between the start locations and goal locations in the maze task),
x E [0, 1] balances relevance versus diversity. This is set to 0.9 by default in our simulations.
We then add the item with the highest MMR score, and repeat until the batch is full.
A.1.4 Gaussian mixture model for image generation
To model generative replay in the image classification task, we use Gaussian mixture model (GMM) as simple generative model of images, implemented with seikit-learn [1.. This assumes that images from each class come from mixture of K Gaussian components. For each class the probability of an image x is:
K
p(x
Fek A(x Meck: Ee
Here Te.k are mixing weights that sum to 1, and each Gaussian component has mean Ma, and covariance Ee,-- The parameters are found by the expectation maximisation algorithm.
Images can then be generated by sampling from the learned distribution for a particular class.
A.1.5 Deep Q-network maze solving agent
The DQN agent uses Stable- Baselines3's DQN with MLP Q-network: feed-forward network mapping the observation vector to action-values, with the default architecture of two fully connected hidden layers (64 units each) and ReLU activations, followed by linear output actions.
Training uses a learning rate of Ic-4, with epsilon fixed at 0.2 (higher than typical because of the frequent changes to the maze, which require significant exploration).
A.2 Further neuroscience background
Recent memories are encoded the hippocampus, which capable 'one-shot learning' of particular events. During sleep and rest, the brain conducts offline processing to better support future behaviour. particular, replay is a phenomenon which hippocampal neurons reactivate memories in 'fssr-forward' [3]. In machine learning terms, hippocampal sequences reactivated
15

---

to train other brain regions, with the neocortex thought to gradually learn statistical patterns across memories through replay [42
This view is supported by the fact that interventions which interfere with replay lead to increased forgetting [45] and replay after learning a task is correlated with subsequent performance [46]. (Note that much of this research relates to rodents performing
spatial tasks like navigating to reward in maze, motivating our use of a maze one scenario in the paper.)
A newer view is that replay may be generative rather than, or as well as, veridical. In rodents, 'replayed' sequences can join together paths that were experienced on separate occasions [ cross regions that have been seen but not visited [47], and even diffuse throughout an open environment [88] In human neuroimaging studies, 'replayed" sequences do not always correspond real memories either [4]. In addition, conceptual "building blocks' are thought to be learned offline which can then be flexibly recombined.
These findings have led to body of research about what kinds of replay are observed when, and which events are replayed simulated during rest. This is thought depend several factors. For example, consider the case of gradually changing maze. Simulating trajectories offline without first updating the world model is detrimental the world has changed, but updating the world model is costly if the world has not changed. So to optimise offline activity for future behaviour, there must be some mechanism for assessing which offline action take based on the current cognitive state (e.g. whether recent experiences diverge from the world model's predictions).
In summary, during rest, the brain is thought perform some combination of i) replaying real memories, ii) learning a general world model, and iii) simulating events with this world model. There evidence for all of these processes, but there is little understanding of how they regulated, and how this depends on the environment and/or current learning. The 'metacognitive actions' chosen in the paper are intended to reflect these observed offline processes in the neuroscience literature.
A.3 Supplementary results
intaal graph
expansion
expansions
Three expanstons
NERC CALD
DUN:
On
Figure 5: A toy example iterative graph expansion. From left to right: The initial family tree graph shown. When the graph autoencoder model for link prediction is applied this graph, the links shown orange, which have probability greater than 0.5, are added. The graph autoencoder is then applied to this expanded graph.
maze update example
goal selection
c) Example
Mazer
Selected starts
Selected goals
Figure Additional maze visualisations. a) Incremental maze updates break a high centrality link and add another, whilst keeping the maze fully connected. b) Example heatmaps showing the frequency of selected start (left) and goal (right) locations. MMR applied to select start-goal combinations which balance high value and high diversity. The selected start and goal locations may very different. c) An example task solved by the DQN agent.
16

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The abstract and introduction accurately reflect the contributions and scope as specified below.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contribu- tions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: The Discussion covers limitations of the paper (see paragraph beginning "This work has several limitations'), and how future work could address them.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results to vi- olations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by review- grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
17

---

The paper presents empirical rather than purely theoretical results. Where the chosen approach depends on existing theoretical results (e.g. Shapley-bar based data valuation), these are cited and explained in the Appendix where possible.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experi- mental results of the paper the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: A detailed description of all experiments included, as well as the code to reproduce the results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others replicate the model with the same dataset, provide access to the model. general, releasing code and data often one good way to accomplish also be provided via detailed instructions for how to replicate the
this, but reproducibility results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (c.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, described in supplemental material?
18

---

Answer: [Yes]
Justification: Code is provided alongside the paper to reproduce the results. It will be made available GitHub repo upon publication. The only datasets used are open-source ones, which the code downloads automatically.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please the NeurIPS code and data submission guidelines (https://niis..cc/ppil Telide//ooeSupa wwi.omoliiiipoppyoooocmm.. for more details.
While we encourage the release of code and data, we understand that this might not be possible,
so "No" is an acceptable answer. Papers camnot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
to
/nips..c/puulic/guudeee http:///wwwo/eaaaoooooooo///oaaa
The authors should provide instructions data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental sefting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: These details are provided in the Appendix and the code to reproduce the results. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code. in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Results in figures are accompanied by error bars where appropriate, and tables include the SEM.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
19

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 96% CL. if the hypothesis of Normality of emors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type compute workers, memory, time of execution) needed to reproduce the experi- ments?
Answer: [Yes]
Justification:
This information is provided in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experi- mental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https::/neeripp::: http://earrrp..ccc/aa/allccccaaa ines?
Answer: [Yes]
Justification: This complies fully with the ethics guidelines, as there harms caused by the research process or potential harmful consequences.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consideration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
The purpose of this work is to better understand the brain through computational neuroscience, which has direct socictal impact (although the field of neuroscience as a whole has positive societal impacts).
Justification:
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no socictal impact or why the paper does not address socictal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g.,
disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploy- ment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
20

---

The conference expects that many papers will be foundational research and not tied to par- ticular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML..
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Creators of the code and data used clearly eited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, soperswithcooe com /atassts has curated licenses for some datasets. Their licensing guide can help determine the license of dataset,
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
21

---

Answer: [Yes]
Justification: A draft of the documentation for code to reproduce the results is provided alongside the paper. More detailed documentation will accompany public GitHub repo released upon publication.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL include anonymized zip file.
14. Crowdsourcing and research with human subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, as well as details about compensation (if any)?
Question:
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing or research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
22

---

Justification: LLMs were only used for coding assistance / debugging (with human verification and oversight).
Guidelines:
The answer NA means that the core method development in this research does not involve
LLMs as any important, original, or non-standard components.
Please refer to our LLM policy should or should not be described.
http:///uurpps.ccc//ffeeeeeeeeee /LLM) for what
23


Reviews:
[{"review_id": "r0fvupc8pu", "reviewer": "Reviewer_mBgr", "summary": "The paper proposes an RL agent operating in two states, “sleep” and “awake,” inspired by work in cognitive science. Learning occurs only in the “sleep” state, and data collection occurs in the “awake” state.", "questions": "1. What is the proposed approach in formal mathematical notation?\n2. What is the hippocampus model in the experiments?\n3. What is the world model in the experiments?\n4. Why were these particular experimental setups chosen?\n\nThe main drawback of the work is the lack of a general picture after reading the paper due to poorly constructed narration and positioning. I believe that the work requires significant revision in this part and is not ready for publication at the current stage.", "limitations": "The limitations are described in sufficient detail.", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 1, "originality": 1, "strengths_and_weaknesses": "**Strengths:**\n\n1. A large number of experiments, which nevertheless do not add up to a general picture.\n\n**Weaknesses:**\n\n1. The paper is poorly constructed at the narrative level, making it difficult to draw conclusions about the significance and originality of the work.\n\n**Quality**\n\nThe paper presents a large number of experiments, but they do not fit into the overall picture and, due to the lack of a clearly formulated contributions, cannot confirm them.\n\n**Clarity**\n\nThe work is very difficult to read and perceive, the narrative jumps between sections without linking the individual parts together.\n1. There is no clear indication of the contribution of the paper in the introduction.\n2. There is no clear formal description of the approach and problem statement, including through formulas and algorithms. Figure 1 gives only a general idea of ​​​​the method.\n3. There is no motivation why these particular experimental setups were chosen.\n4. The part from the Discussion section describing the limitations of the work would be better allocated to a separate section.\n\n**Significance and Originality**\n\nIt is really hard to evaluate these aspects of the work because on the one hand there is no clear, well-articulated contribution of the work, on the other hand there is no clear positioning of the work relative to others in the Previous work section. Authors should clearly position themselves relative to the works mentioned.\n\nHowever, I should note that I am not an expert in neuroscience and cognitive science and may be missing important results for these areas.", "was_revised": true, "final_justification": "The paper suffers from a lack of a clear formal description of the proposed approach and technical details. To eliminate this shortcoming, a significant revision of the paper is required. In this regard, I believe that in its current form the paper cannot be presented at such a top-tier conference as NeurIPS.", "rebuttal": "Thank you for your comments. We’re sorry that the paper was unclear to those without a neuroscience background, and hope we can clarify the ‘general picture’ below. \n\nHere is a revised framing of the paper with some more neuroscience context:\n* Recent memories are encoded in the hippocampus, which is capable of ‘one-shot learning’ of particular events; as well as capturing our recent experiences, it is often thought of as a store of training data for other networks in the brain. \n* During sleep and rest, a lot is happening to update neural networks to better support future behaviour. In particular, replay is a phenomenon in which hippocampal neurons reactivate memories in ‘fast-forward’ (Foster, 2017). In ML terms, data are reactivated to train other brain regions, with the neocortex thought to gradually learn statistical patterns across memories through replay (Marr, 1970, 1971; McClelland et al., 1995; Teyler & DiScenna, 1986). Interventions that interfere with replay lead to increased forgetting (Girardeau et al., 2009), and replay after learning a task is correlated with subsequent performance (Peigneux et al., 2004). (Note that much of this research relates to rodents performing spatial tasks like navigating to a reward in a maze, motivating our use of a maze as one scenario in the paper.)\n* A newer view is that replay may be generative rather than, or as well as, veridical. In rodents, ‘replayed’ sequences can join together paths that were experienced on separate occasions (Gupta et al., 2010), cross regions that have been seen but not visited (Olafsdottir et al., 2015; Pfeiffer & Foster, 2015), and even ‘diffuse’ throughout an open environment (Stella et al., 2019). In human neuroimaging studies, ‘replayed’ sequences do not always correspond to real memories either (Liu et al., 2019). In addition, conceptual 'building blocks' are thought to be learned offline which can then be flexibly recombined.\n* These findings have led to a body of research about what kinds of replay are observed when, and which events are replayed / simulated during rest. This is thought to depend on the situation. To illustrate this, consider the case of adding and removing links gradually in a maze. One might expect that simulating trajectories offline without first updating the world model would be detrimental to learning. But updating the world model is costly to do when the world has not changed. So if the brain is optimising its offline activity for future behaviour, there must be some mechanism for assessing which offline action to take based on the current cognitive state (e.g. whether recent experiences diverge from the world model’s predictions). \n* So, during rest, the brain is doing some combination of 1) replaying real memories, 2) learning a general world model, 3) simulating events with this world model, and 4) learning specific tasks. There is evidence for all of these processes, but there is little understanding of how they are regulated, and how this depends on the environment and/or the animal’s current learning. The ‘metacognitive actions’ chosen in the paper are intended to reflect these observed offline processes in the neuroscience literature. \n* We want to develop an approach to understand what sequence of offline processing is optimal, and which events should be replayed / simulated to optimise future behaviour. This could make predictions about the pattern on offline activity that could be tested in future experiments.\n* Whilst the training data in modern ML tends to be carefully curated, e.g. by balancing and interleaving different categories, brains excel at data-efficient learning in a changing environment, in which generalisable knowledge must be extracted from limited, noisy experience. The offline processing of memories and beliefs is thought to support this, so we hope these mechanisms are of wider interest to the ML community, and could lead to brain-inspired innovations.\n\nI hope this provides some more insight into why the orchestration of offline processing is an important unsolved problem in neuroscience, motivating our study of whether a meta-controller could learn when to perform different ‘metacognitive actions’. We have updated our draft to make the paper assume less prior knowledge.\n\nIn response to your specific questions above:\n1. The approach is described more mathematically in the supplementary information of the initial version, but we have extended this to provide more details. As a brief summary:\n\nFor each sleep/awake cycle $n = 1,\\\\ldots,N$: \n\n*Awake phase:* Agent $\\\\theta_n$ interacts with the environment for $K$ steps (i.e. ${n\\_{episodes}}$), encoding memories $H\\_n = \\{(s_k,a_k,r_k)\\}\\_{k=1}^{K} $ in the model hippocampus and recording an overall reward for the awake phase $R\\_n = \\tfrac{1}{n\\_{episodes}}\\\\sum\\_{k=1}^{K} r\\_k$. \n\n*Sleep phase:* A meta-controller takes $T$ offline actions $a_{n,0:T-1}$, corresponding to the kinds of offline processing observed in the neuroscience literature described above.\n\n*Objective:* The recurrent PPO policy $\\\\pi_{\\\\phi}(a \\\\mid x)$ updates its parameters to maximise the expected cumulative return $\\mathbb{E}\\_{\\pi\\_{\\phi}}\\left[\\sum\\_{n=1}^{N} R_n\\right]$.\n\nIn the maze task, the reward $R_n$ is the mean reward per episode of the lower-level agent in the awake state, i.e. in the real maze environment. In the image task the reward is the mean classification accuracy for new images encountered in the awake state.\n\n2. The hippocampus in our simulations is simply a list of episodes the agent has experienced in the awake state. However in reality we would envisage the hippocampus to be a network capable of one-shot learning of sequential experiences, e.g. an asymmetric modern Hopfield network (see Millidge et al., 2022).\n3. The world model is a model that can simulate experience with which to train the model for the task. In the image case, the world model is the image generator. In the maze case, the world model is the transition model. However the ideas in this paper can be applied to more complex kinds of world model, e.g. consider an agent learning to play a video game from simulations with a world like PlaNet (Hafner et al., 2019) or DreamerV2 (Hafner et al., 2021). In a setting in which the video game environment changes over time, encoding and replay of new experiences to update the world model need to be balanced with learning from the world model, and our approach would be applicable in this context.\n4. The image case is chosen because it is a classic setting for considering learning from generated data, and because it enables the data valuation network’s selections to be easily visualised. The maze case is chosen because a lot of neuroscience relates to rodents performing spatial tasks like navigating to a reward in a maze, and a maze is a natural task in which to consider changes to an environment over time (by adding or removing links). However, many kinds of low-level agent and environment can work within this framework. (If accepted the final version of the code will demonstrate this.)\n\nThanks again for your feedback - we hope you’ll consider updating your evaluation if we have addressed your concerns.\n\n**References:**\n\n* Foster, D. J. (2017). Replay comes of age. Annual Review of Neuroscience, 40, 581–602. \n* Marr, D. (1970). A theory for cerebral neocortex. Proceedings of the Royal Society of London. Series B: Biological Sciences, 176(1043), 161–234. \n* Marr, D. (1971). Simple memory: A theory for archicortex. Philosophical Transactions of the Royal Society of London. B, Biological Sciences, 262(841), 23–81. \n* McClelland, J. L., McNaughton, B. L., & O’Reilly, R. C. (1995). Why there are complementary learning systems in the hippocampus and neocortex: Insights from the successes and failures of connectionist models of learning and memory. Psychological Review, 102(3), 419–457.\n* Teyler, T. J., & DiScenna, P. (1986). The hippocampal memory indexing theory. Behavioral Neuroscience, 100(2), 147–154. \n* Girardeau, G., Benchenane, K., Wiener, S. I., Buzsáki, G., & Zugaro, M. B. (2009). Selective suppression of hippocampal ripples impairs spatial memory. Nature Neuroscience, 12(10), 1222–1223. \n* Peigneux, P., Laureys, S., Fuchs, S., Collette, F., Phillips, C., Aerts, J., … Maquet, P. (2004). Are spatial memories strengthened in the human hippocampus during slow wave sleep? Neuron, 44(3), 535–545. \n* Gupta, A. S., van der Meer, M. A. A., Touretzky, D. S., & Redish, A. D. (2010). Hippocampal replay is not a simple function of experience. Neuron, 65(5), 695–705. \n* Ólafsdóttir, H. F., Barry, C., Saleem, A. B., Hassabis, D., & Spiers, H. J. (2015). Hippocampal place cells construct reward related sequences through unexplored space. eLife, 4, e06063.\n* Pfeiffer, B. E., & Foster, D. J. (2015). Autoassociative dynamics in the generation of sequences of hippocampal place cells. Science, 349(6244), 180–183.\n* Stella, F., Baracskay, P., O’Neill, J., & Csicsvari, J. (2019). Hippocampal reactivation of random trajectories resembling Brownian diffusion. Neuron, 102(2), 450–461.e7. \n* Liu, Y., Dolan, R. J., Kurth-Nelson, Z., & Behrens, T. E. J. (2019). Human replay spontaneously reorganizes experience. Cell, 178(3), 640–652.e14. \n* Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., & Davidson, J. (2019). Learning latent dynamics for planning from pixels. arXiv preprint arXiv:1811.04551.\n* Hafner, D., Lillicrap, T., Norouzi, M., & Ba, J. (2021). Mastering Atari with discrete world models. arXiv preprint arXiv:2010.02193."}, {"review_id": "FhbTtg074z", "reviewer": "Reviewer_1iZc", "summary": "This paper prpopses a new framework to better use the acquired experiences for training. Specifically, a meta controller is used to decide whether the agent should choose from 1) training the model using real data, 2) build a world model from the data and 3) train the model using data generated by the world model. The method addresses some issues with previous methods such as the reliance on human judgement on the data augmentation/curriculum design process, and the low quality of synthetic data. Empirical study shows the method can effectively learn an adaptive curriculum for offline learning.", "questions": "My concerns listed in \"weaknesses\" are also my questions to the author(s). I will consider updating my evaluation if my concerns are addressed.", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The proposed method is novel, it provides a new view for curriculum design or data augmentation.\n\n2. The idea is conceptually sound and easy to follow. The method is well described.\n\n\n\nWeaknesses:\n\n1. This paper lacks comparative study. While the existing experiments show that the framework can learn an effective curriculum for offline learning, comparisons to previous image classification, navigation and relational inference algorithms are needed to better demonstrate the framework's performance and potential advantages over these related work. For example, the author(s) may want to compare the classification accuracy with traditional image classifiers in the image task or compare the cumulative reward / success rate with offline RL algorithms in the navigation task.\n\n2. Some ablation studies are needed to better understand the effect of each component in the proposed framework. An example to do this is to try (world model constructuion + learn from data generated only) to see the importance of training on real data. Another example is sample training data with some simple distribution instead of a shapley value estimator to show the effect of data valuation.", "was_revised": true, "final_justification": "I raised my score as the authors addressed my concerns.", "rebuttal": "Thank you for your comments. We are glad you found the general idea of the paper novel and conceptually sound.\n\nWe agree that the paper would benefit from more comparison and ablation studies, so we have run additional simulations to show that metacognitive control of offline processing and custom curriculum design improve performance compared to sensible baselines. (Due to restrictions on the rebuttal format, we summarise the results below, but we have prepared improved figures to add if the paper is accepted, as well as additional text. Results below give the mean across several repeats, but the improved figures have error bars.)\n\n**Metacontroller action selection baselines**\n\nIn the image task, we wanted to compare the metacontroller’s adaptive strategy to i) a fixed memory replay strategy and ii) a fixed world model update + generative replay strategy. The table below shows that the two baselines each do well in only one setting, with the highest average accuracy across the two settings obtained by the metacontroller. Note that we have set the data valuation to be random to isolate the effect of the actions, and give results for Fashion-MNIST only:\n\n| Approach                | Accuracy&nbsp;(CL) | Accuracy&nbsp;(non‑CL) | Mean accuracy |\n|-------------------------|:------------------:|:----------------------:|:-------------:|\n| Metacontroller          | 0.36              | 0.41                   | 0.39          |\n| Memory replay only      | 0.14              | 0.41                   | 0.28          |\n| Generative replay only  | 0.34              | 0.30                   | 0.32          |\n\nIn the maze task, we wanted to compare the metacontroller’s adaptive strategy to a fixed memory replay strategy (which performs very poorly as there are not enough real episodes to learn a good policy from, meaning that the metacontroller is better at generalising to new tasks). As above, we have set the data valuation to be random to isolate the effect of the actions:\n\n| Approach               |  Mean reward |\n|------------------------|:-----------:|\n| Metacontroller         | 0.90        |\n| Memory replay only     |  0.47        |\n\n**Data valuation baselines**\n\nWe also tested new baselines for the data valuation component of the model. In the image task, we previously tested random selection as the only alternative to the learned data valuation network. To add to this, we tested training on the easiest and hardest images only (where the difficulty of an image was quantified as the fraction of an ensemble of scikit-learn classifiers that labelled the image correctly). Both perform worse than regulating the difficulty by the stage of learning (here we give results for just the non-CL setting with Fashion-MNIST):\n\n| Data‑valuation approach | Mean accuracy |\n|-------------------------|:-------------:|\n| Data valuation network  | 0.58          |\n| Random selection        | 0.41          |\n| Most challenging     | 0.44          |\n| Least challenging      | 0.34          |\n\nThe fact that none of i) most challenging, ii) least challenging, or iii) random data alternatives do better provides support for the claim that the system is choosing the optimal examples for its current stage of learning, i.e. that it is learning a curriculum from less challenging to more challenging data.\n\nIn the maze task, the only previous baseline was random selection, so we tried two baselines that correspond to simple heuristics for which episodes are most useful to the agent: i) training on the longest trajectories, and ii) training on only paths that visit the new link. Both performed worse than the data valuation network, suggesting the data valuation networks learns something more sophisticated than these simple heuristics.\n\n| Data‑valuation approach | Mean accuracy |\n|-------------------------|:-------------:|\n| Data valuation network  | 0.94          |\n| Random selection        | 0.90          |\n| Longest paths           | 0.85          |\n| Paths through new link  | 0.87          |\n\nThanks again for your feedback- we hope you’ll consider updating your evaluation if we have addressed your concerns. If accepted, the additional results will obviously be included in the final version, as will the updated code to replicate them."}, {"review_id": "WxBfUDj4AB", "reviewer": "Reviewer_9NGj", "summary": "This paper proposes a computational framework to explain how the brain might effectively balance different offline learning strategies, namely consolidating direct experiences, abstracting a world model, and learning from imagined scenarios. The authors introduce a meta-reinforcement learning model where a meta-controller learns to guide the offline learning process of a lower-level agent. Specifically, during simulated sleep periods, this meta-controller selects from a set of metacognitive actions: training on real experiences, updating a world model, or training on experiences simulated from that model. The framework also learns to prioritize experiences by assigning them value. The paper's core contributions are to (1) provide a normative model that makes testable predictions about the content and timing of neural replay and offline processing, which the authors validate in three experiments, and (2) demonstrate how this method of constructing a bespoke learning curriculum can improve the overall performance of a reinforcement learning agent.", "questions": "- What is the performance of these models in comparison to heuristic baselines?\n- How are metacognitive actions designed?\n- For the maze experiment, are the start and end locations fixed? \n- If the reward is the fraction of edges in the true graph which can be inferred by the end of the episode, why is the y-axis of 4(c) and 4(d) (top-left) showing values greater than 1? Is this actually showing a percentage.\n- Is it a realistic assumption that the meta-controller always observes how accurate its world model is? What does this mean? Where does it get this information from?\n- “In our simulations, the meta-controller also learns to choose which events to replay or simulate.”\n    - Does the meta-controller do that, or the separate valuation network?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Generally speaking, the paper is clearly written and well organized. It is clear what the contribution of the paper is (the framework), how it works (in the model section), as well as the two key claims it makes regarding that contribution.\n- The paper does provide an interesting framework that one could see validated in a range of other tasks (for instance, in an RLHF domain). \n- As the paper points out, it also lays the groundwork for fruitful comparisons to experimental neuroscience data on the optimal scheduling of replays.\n- This paper novelly employs meta-reinforcement learning in a neuroscience context when looking at the idea of metacognitive control.\n\nWeaknesses:\n- One of the central claims--that metacognitive actions and custom curriculum design improve performance--is not well-supported. To wit:\n    - In Experiment 2, performance is only evaluated in the narrow domain of evaluating the data valuation component of the model, showing that this improves accuracy over random replay using both real and simulated examples.\n    - In Experiment 3, no measure of performance (that is, the agent’s realized ability to efficiently navigate the maze) is reported at all.\n    - In all three experiments, no baselines are evaluated showing how learning a policy for meta-cognitive actions would improve performance over a fixed-heuristic model like Dyna, Prioritized Sweeping, or its deep learning counterparts. These should be easy baselines to show in order to buttress this claim. Visualizing action policies is not enough.\n- Fig. 3(c) is hard to interpret the values. They’re all too small to be seen in the annotation, but the colors are all different.\n- The normative neuroscience predictions could be better elaborated upon.", "was_revised": true, "final_justification": "The updates to the results made by the authors in their rebuttal have addressed well my central critique of the paper, namely that one of the central claims initially lacked any supportive evidence in the quantitative results. Given that, I have decided to revise my score upwards accordingly.", "rebuttal": "Thank you for your comments. We are glad you found the proposed framework interesting, and the application of meta-RL in the context of metacognitive control novel. \n\nWe agree that comparison to baselines would strengthen the paper, so we have run additional simulations to show that metacognitive control of offline processing and custom curriculum design improve performance. (Due to restrictions on the rebuttal format, we summarise the results below, but we have prepared improved figures to add if the paper is accepted, as well as additional text. Results below give the mean across several repeats, but the improved figures have error bars.)\n\n**Metacontroller action selection baselines**\n\nIn the image task, we wanted to compare the metacontroller’s adaptive strategy to i) a fixed memory replay strategy and ii) a fixed world model update + generative replay strategy. The table below shows that the two baselines each do well in only one setting, with the highest average accuracy across the two settings obtained by the metacontroller. Note that we have set the data valuation to be random to isolate the effect of the actions, and give results for Fashion-MNIST only:\n\n| Approach                | Accuracy&nbsp;(CL) | Accuracy&nbsp;(non‑CL) | Mean accuracy |\n|-------------------------|:------------------:|:----------------------:|:-------------:|\n| Metacontroller          | 0.36              | 0.41                   | 0.39          |\n| Memory replay only      | 0.14              | 0.41                   | 0.28          |\n| Generative replay only  | 0.34              | 0.30                   | 0.32          |\n\nIn the maze task, we wanted to compare the metacontroller’s adaptive strategy to a fixed memory replay strategy (which performs very poorly as there are not enough real episodes to learn a good policy from, meaning that the metacontroller is better at generalising to new tasks). As above, we have set the data valuation to be random to isolate the effect of the actions:\n\n| Approach               |  Mean reward |\n|------------------------|:-----------:|\n| Metacontroller         | 0.90        |\n| Memory replay only     |  0.47        |\n\n**Data valuation baselines**\n\nWe also tested new baselines for the data valuation component of the model. In the image task, we previously tested random selection as the only alternative to the learned data valuation network. To add to this, we tested training on the easiest and hardest images only (where the difficulty of an image was quantified as the fraction of an ensemble of scikit-learn classifiers that labelled the image correctly). Both perform worse than regulating the difficulty by the stage of learning (here we give results for just the non-CL setting with Fashion-MNIST):\n\n| Data‑valuation approach | Mean accuracy |\n|-------------------------|:-------------:|\n| Data valuation network  | 0.58          |\n| Random selection        | 0.41          |\n| Most challenging     | 0.44          |\n| Least challenging      | 0.34          |\n\nThe fact that none of i) most challenging, ii) least challenging, or iii) random data alternatives do better provides support for the claim that the system is choosing the optimal examples for its current stage of learning, i.e. that it is learning a curriculum from less challenging to more challenging data.\n\nIn the maze task, the only previous baseline was random selection, so we tried two baselines that correspond to simple heuristics for which episodes are most useful to the agent: i) training on the longest trajectories, and ii) training on only paths that visit the new link. Both performed worse than the data valuation network, suggesting the data valuation network learns something more sophisticated than these simple heuristics.\n\n| Data‑valuation approach | Mean accuracy |\n|-------------------------|:-------------:|\n| Data valuation network  | 0.94          |\n| Random selection        | 0.90          |\n| Longest paths           | 0.85          |\n| Paths through new link  | 0.87          |\n\nWe hope this addresses the first question, and the first weakness identified in your review.\n\nAnswers to your other questions are as follows.\n\n* *How are metacognitive actions designed?* The metacognitive actions are chosen to reflect a highly simplified version of metacognitive processes discussed in the neuroscience literature. See the response to reviewer mBgr for ore explanation of evidence for these actions.\n* *For the maze experiment, are the start and end locations fixed?* No, the start and end locations can be anywhere in the maze.\n* *If the reward is the fraction of edges in the true graph which can be inferred by the end of the episode, why is the y-axis of 4(c) and 4(d) (top-left) showing values greater than 1? Is this actually showing a percentage?* Apologies, this shows the number of edges, not the fraction of edges. We have corrected this error.\n* *Is it a realistic assumption that the meta-controller always observes how accurate its world model is? What does this mean? Where does it get this information from?* We assume a small test set exists, which could be a subset of hippocampal memories. If this subset is not used to update the world model, it can instead be used as a test set to evaluate it, by tracking the prediction error of the world model on the test set. \n* *“In our simulations, the meta-controller also learns to choose which events to replay or simulate.” Does the meta-controller do that, or the separate valuation network?* We were using meta-controller to refer to the combined system, but you are correct that we should be more precise, as this claim could be confused with the recurrent PPO agent itself choosing events to replay or simulate. We have clarified this.\n\nIn addition, you pointed out that ‘In Experiment 3, no measure of performance (that is, the agent’s realized ability to efficiently navigate the maze) is reported at all’. We have added a plot to show that the agent’s mean reward per episode in the awake state increases over time as the meta-controller learns, which shows that once the meta-controller has learned to sequence its meta-actions correctly, reward is near optimal (with optimal performance plotted for comparison). We have also shown trajectories for a random subset of (start, goal) pairs to illustrate this qualitatively.\n\nThanks again for your feedback- we hope you’ll consider updating your evaluation if we have addressed your concerns. If accepted, the additional results will obviously be included in the final version, as will the updated code to replicate them."}]

Author rebuttal:

Meta-review: The paper presents a framework in which an agent learns to control states of learning and how to balance inputs.  Although one reviewer pointed out several weaknesses as the reason for lower scores, the authors seemed to have fully addressed these weaknesses and can easily integrate these findings into the paper.  
The approach is novel, sound and the paper is well written, with some expected interesting impact.
Decision: Accept (poster)