T HE E MERGENCE OF A BSTRACT AND E PISODIC N EURONS IN E PISODIC M ETA -RL
A BSTRACT
In this work, we analyze the reinstatement mechanism introduced by Ritter et al.
(2018) to reveal two classes of neurons that emerge in the agent’s working memory (an epLSTM cell) when trained using episodic meta-RL on an episodic variant of the Harlow visual ﬁxation task.
Speciﬁcally, Abstract neurons encode knowledge shared across tasks, while Episodic neurons carry information relevant for a speciﬁc episode’s task.
I NTRODUCTION
Starting as a method to study animal conditioning in psychology (Pavlov, 1927; Rescorla & Wagner, 1972), Reinforcement Learning (RL) has become an efﬁcient way to train artiﬁcial agents in solving complex tasks such as Go or StarCraft (Silver et al., 2016; Vinyals et al., 2019).
Despite such successes important problems remain.
State-of-the-art RL algorithms require enormous amount of training data and do not easily adapt to new tasks.
One research strand trying to address these issues is meta-reinforcement learning (meta-RL) (Wang et al., 2016) - in which agents have to learn to deal with a number of different tasks.
Typically, in this work recurrent neural networks - speciﬁcally LSTMs - are used to learn representations that encode an RL algorithm.
Ritter et al.
(2018) proposed to extend these LSTMs with neural memory - so the agents are able to remember and retrieve knowledge gained over past tasks when re-encountering them.
The memory proposed by Ritter et al.
(2018) relies on a gating mechanism that decides which memory activations are retrieved and reinstated into the LSTM.
This gating mechanism is trained as part of the overall optimization problem and constitutes a key artifact of learning.
We study how the gating mechanism interacts with LSTM neurons, and show that they can be roughly categorized.
We identify two classes of neurons: episodic and abstract neurons - that differ in their characteristics w.r.t how information is restored as well as their impact on the performance of the system.
Abstract neurons encode structural task knowledge relevant across different episodes, while Episodic neurons carry episode-speciﬁc environmental information (potentially reoccurring in later episodes).
This paper proceeds as follows: (1) we introduce a simpliﬁed version of the Harlow Task (a standard
Meta-RL environment) with episodic cues, (2) we introduce the model implementation, followed by (3) an analysis, deﬁnition and tests for abstract and episodic neurons.
Lastly, we brieﬂy discuss our results in the wider context of meta-RL.
M ETHODS 2. T ASK F ORMULATION
[Figure]
Let M i P D be a distribution of tasks each characterized as a Markov Decision Process (MDP):
M i “ p S , A , T i , R i q .
The agent learns over a sequence of MDPs by taking an action a P A to transition from state s to s 1 (where s, s 1 P S ) and receiving a scalar reward r , using some transition probability distribution T : S ˆ A Ñ S and reward function R : S ˆ A Ñ R .
To introduce the concept of identiﬁably reoccurring tasks, Ritter et al.
(2018) extend the previous formulation by associating a context k
i with each M i , sampled as p M i , k i q „ D uniformly with replacement.
With each new task, the agent can use the context k i to identify if the task had been seen before, and hence leverage previously discovered policies to avoid redundant exploration.
The (One-dimensional) Symbolic Episodic Harlow Task To study and analyze episodic metaRL agents, we develop a simpliﬁed symbolic version with exact parallels to the task structure of the Harlow visual ﬁxation task found in the PsychLab environment (Leibo et al., 2018) but which factors out the visual and spatial modeling of the environment.
Further details can be found in Appendix A.
2. M ODEL D ESCRIPTION
The agent is trained 1 using the Advantage Actor-Critic (A2C) RL algorithm on a single thread (Mnih et al., 2015).
The architecture follows that of the LSTM A3C model from Wang et al.
but uses an epLSTM instead.
The encoder is a stack of 2 afﬁne layers with 64 and 128 units, respectively, and a ReLU non-linearity in-between.
The epLSTM takes a concatenation of: (a) The encoding of the receptive ﬁeld, (b) the reward at t ´ 1 , and (c) the action at t ´ 1 .
The epLSTM is a one layer LSTM with 256 2 hidden units plus the reinstatement mechanism.
The memory module maps the context associated with the current task k
i (as key) to the cell state c T (as value) at the end of each episode (time T ).
This memory is updated at the same key each time the task M i reoccurs.
The experiment is repeated 50 times with different initializations.
We analyze the top 30 models (ﬁltered by a threshold on the reward calculated on 100 randomly generated episodes).
Each instance is trained for 25 ,episodes and tested for 1 , 000 episodes (with different objects).
The code is made open-source 3 .
2. R EINSTATEMENT M ECHANISM
We use the reinstatement mechanism from (Ritter et al., 2018)—deﬁned as:
[Equation]
where r t controls the ﬂow of information from the retrieved memory m t into the epLSTM cell state:
[Equation]
The vector c t can be seen as encoding the agent’s working memory state up to time t .
Therefore, at the end of each episode the agent commits the accumulated knowledge learned about the current
MDP M i to an external long-term memory module with the associated context vector k i as key.
In a later episode where M i reoccurs, the context k i is used to query the memory to retrieve the corresponding cell state as m t (equals the zero vector if M i is novel).
The retrieval occurs when the agent ﬁrst observes the pair of objects (after ﬁxation at the ﬁrst trial).
The vector m t is then interpolated into the current working memory using the reinstatement gate r t .
A NALYSIS 3. T ASK P ERFORMANCE
Figure 1d shows the performance over successive stages of training as a function of the trial number.
The performance on the ﬁrst trial improves and is not stuck at random as in the classical Harlow experiment because the agent is able to reinstate relevant information when it re-encounters a speciﬁc task.
Figure 1e shows the testing performance as a function of the number of times an agent is exposed to a particular task.
It can be seen that when a task reoccurs the agent immediately identiﬁes it and is able to solve it from the ﬁrst trial.
3. R ECURRENT N EURONS AND THE R EINSTATEMENT G ATE
[Figure]
Noting that r t r j s P p 0 , 1 q (by the σ function) modulates the reinstatement of individual neurons from m t , we may interpret r t r j s as the importance of neuron j for recurring episodic information.
Figure 2a shows the values of r r j s at ﬁxation (as heat) for each neuron j across training episodes.
Notice that the r r j s activations converge to stable values as training progresses, independent of changes in input or hidden state among the different episodes.
We calculate from the lasttraining episodes a static value for r ﬁx ( r t at t “ ﬁxation) and use it for all further gate analysis.
Formally ( e is the episode index and N e is the training episode count):
[Equation]
Figure 2b is a histogram of r ˚ r j s values.
About 25% of the neurons have become biased to be open ( r ˚ r j s ě 0 . 9 ), while „ 30% are biased to be closed ( r ˚ r j s ă 0 . 1 ).
Figure 2c shows the cosine similarity of certain regions (indicated by hue) of c t to c ﬁx (averaged over 1000 testing episodes).
This suggests that some speciﬁc neurons consistently hold the information needed across episodes to identify the winning object, and so change values less often within each episode (see Figure 3).
3. T ESTING FOR A BSTRACT AND E PISODIC N EURONS
[Figure] To clarify the roles of individual neurons in c t , we test while gradually masking out the neurons c t r j s based on t j : r ˚ r j s ž θ u then analyze the behavioral change in two signals: (a) number of steps before ﬁxation, and (b) ﬁrst trial performance.
Figure 3a shows the average time-steps before ﬁxation when zeroing out c t based on the upper/ episodic (orange) or lower/ abstract (blue) regions of r ˚ .
Dropping more of the “abstract” region leads to worse performance (more steps to ﬁxation).
Performance remains largely stable as we drop “episodic” neurons ﬁrst, up until before the extreme where all neurons are dropped ( θ “ 0 . 0 ).
Figure 3b shows objective performance (choosing the rewarding object) at the ﬁrst trial during testing where the MDP M i has occurred before.
Dropping neurons from the “abstract” region (blue curve) shows a smooth drop in performance, while dropping “episodic” neurons shows a steep drop from as early as r ˚ ą 0 .
9 , suggesting a strong correlation between the neurons where r ˚ r j s ą θ (for some reasonable θ ) and task performance on object re-occurrence.
This seems to indicate that this region holds most of the object-based reward information.
D ISCUSSION
In this work, we have shown the existence of two classes of neurons that emerge in the memory reinstatement-based episodic meta-RL paradigm.
Each neuron may belong to a class that encodes episodic information, or a class that encodes abstract knowledge that is shared across episodes.
This ﬁnding implies that one does not need to store the whole cell-state when committing it to the long-term memory module since only a fraction of the activations are actually going to be reinstated.
Therefore, one optimization method is to store a sparse representation of c T while storing the required indices only once.
In the case of the experiments conducted in this paper, this method can save up to 75% of the storage cost for the memory module while maintaining close to optimal performance after deployment once r ˚ is computed.
Wang et al.
(2018) had shown that the meta-RL framework has direct connections with structures and functions in the brain.
Speciﬁcally, they conceptualize the prefrontal cortex (PFC) along with the subcortical structures to which it connects as forming a homogeneous recurrent neural network that is trained using striatal dopamine reward prediction error signals.
Inline with this theory and the work presented in this paper, previous work have shown that the PFC contain single neurons that encodes abstract rules (Wallis et al., 2001).
Future work may extend the analysis to different episodic tasks, and utilize the ﬁndings for incorporating stronger inductive biases.
We hope this work meaningfully furthers the sharing of insights between the neuroscience and machine learning ﬁelds.
A T HE S YMBOLIC E PISODIC H ARLOW T ASK
The task consists of a one-dimensional circular state-space with 16 discrete cells, of which the agent can observe 8 cells at any time step (see Figure 1 for an illustration of the task).
It starts with a central ﬁxation cross placed in the initial observable space (i.e. receptive ﬁeld) of the agent; similar to the PsychLab version.
The agent can then select one of two actions | A | “ 2 : move one cell to the left or to the right .
After the ﬁxation cross appears in the center of the agent’s receptive ﬁeld, it is removed and two objects are introduced to the left and right of the center, one of which is randomly assigned to be the rewarding object throughout the episode.
The objects are uniformly sampled from a distribution of n “ 100 objects split into 80 for training and 20 for testing, resulting in n p n ´ 1 q “ 9 , 900 possible combinations of object-reward pairs, because in each task either of the object pair may be rewarding.
The agent must then choose one of the objects by orienting it towards the center of its receptive ﬁeld.
Following that, the ﬁxation cross reappears initiating the next trial.
Similar to the PsychLab version, one episode consists of 6 trials, but here it is terminated after a maximum of 120 total time steps.
We use the same reward values as used in Wang et al.
(2018):and ´ 1 for the rewarding and non-rewarding objects, respectively, and 0 . 2 for arriving at the ﬁxation cross.
The episodic structure comes from the fact that the same objects with their associated rewards can be sampled more than once.
In order for the agent to identify the current task, a context vector is randomly generated the ﬁrst time the agent encounters a particular task.
This creates a unique mapping between each possible MDP and its corresponding context.