Fast And Slow Learning Of Recurrent Independent Mechanisms
Authors: Kanika Madan
Nan Rosemary Ke
Anirudh Goyal
Bernhard Schölkopf
Yoshua Bengio
Venue: ICLR 2021 Poster (2021)
Keywords: modular representations
better generalization
learning mechanisms

Abstract: Decomposing knowledge into interchangeable pieces promises a generalization advantage when there are changes in distribution. A learning agent interacting with its environment is likely to be faced with situations requiring novel combinations of existing pieces of knowledge. We hypothesize that such a decomposition of knowledge is particularly relevant for being able to generalize in a systematic way to out-of-distribution changes. To study these ideas, we propose a particular training framework in which we assume that the pieces of knowledge an agent needs and its reward function are stationary and can be re-used across tasks. An attention mechanism dynamically selects which modules can be adapted to the current task, and the parameters of the \textit{selected} modules are allowed to change quickly as the learner is confronted with variations in what it experiences, while the parameters of the attention mechanisms act as stable, slowly changing, meta-parameters. We focus on pieces of knowledge captured by an ensemble of  modules sparsely communicating with each other via a bottleneck of attention. We find that meta-learning the  modular aspects of the proposed system greatly helps in achieving faster adaptation in a reinforcement learning setup involving navigation in a partially observed grid world with image-level input.  We also find that reversing the role of parameters and meta-parameters does not work nearly as well, suggesting a particular role for fast adaptation of the dynamically selected modules.

Published as a conference paper at ICLR 2021
FAST AND SLOW LEARNING OF RECURRENT INDEPEN- DENT MECHANISMS
Kanika Madan', Rosemary Nan Ke AAniudh Goyal Bernhard Schölkopf 3 Yoshua Bengio
ABSTRACT
Decomposing knowledge into interchangeable pieces promises generalization advantage when there are changes in distribution. A learning agent interacting with its environment is likely be faced with situations requiring novel combinations of existing pieces of knowledge. We hypothesize that such a decomposition of knowledge particularly relevant for being able to generalize in a systematic way to out-of-distribution changes. To study these ideas, we propose a particular training framework in which we assume that the pieces of knowledge an agent needs and its reward function are stationary and can be re-used across tasks. An attention mechanism dynamically selects which modules can be adapted to the current task, and the parameters of the selected modules are allowed to change quickly as the learner is confronted with variations in what it experiences, while the parameters of the attention mechanisms act stable, slowly changing, meta- parameters. We focus on pieces of knowledge captured by an ensemble of modules sparsely communicating with each other via bottleneck of attention. We find that meta-learning the modular aspects of the proposed system greatly helps in achieving faster adaptation reinforcement learning setup involving navigation in a partially observed grid world with image-level input. We also find that reversing the role of parameters and meta petamparsmeters does not work nearly as well, suggesting a particular role for fast adaptation of the dynamically selected modules.
INTRODUCTION
The classical statistical framework of machine learning is focused on the assumption of independent and identically distributed (i.i.d) data, implying that the test data comes from the same distribution as the training data. However, a learning agent interacting with the world is often faced with non- stationarities and changes in distribution which could be caused by the actions of the agent itself, or by other agents in the environment. A long-standing goal of machine learning has been to build models that can better handle such changes in distribution and hence achieve better generalization Lake & Baroni 2018: Bahdanau et al., 2018). At the same time, deep learning systems built in the form of single big network, consisting of a layered but otherwise monolithic architecture, tend to co-adapt across different components of the network. Due to monolithic structure, when the task or the distribution changes, majority of the components of the network are likely to adapt in response to these changes, potentially leading to catastrophic interferences between different tasks or pieces of knowledge (Andreas et al. 2016 Fernando et al.) 2017: Shazeer et al. 2017 Jo et al.
2018; Rosenbaum et al. 2019 Alet et al. 2018 Kirsch et al. 2018 Goyal et al. 2019 2020 Goyal
& Bengio 2020).
An interesting challenge of current machine learning research is thus out-of-distributio adaptation and generalization. Humans seem to be able to learn a new task quickly by re- re-using relevant prior knowledge, raising two fundamental questions which we explore here: (1) how separate knowledge into easily recomposable pieces (which we call modules), and (2) how to do this so as to achieve fast adaptation to new tasks changes in distribution when module may need to be modified or when different modules may need to be combined in new ways. For the former objective, instead of representing knowledge with homogeneous architecture as standard neural networks. we adopt
recently proposed approaches (Goyal et al. 2019 Mittal et al. 2020 Goyal et al. 2020 Rahaman
Mila, University of Montreal, Mila, Polytechnique Montréal, Max Planck Institute for Intelligent Systems. Corresponding author: madankan ka s@gmai com

---

Published as a conference paper at ICLR 2021
Dynamica
update
M1
M2
M2
Dampaqqea dry Mannnn
Value Head
Mi
M9
Image RGB Observation
Encoder
M2
M4
Ma;
Dynamics update
M3
M4
Independent modules query on input using input Attention "IA"
Top-k modules are activated updata their
Policy
Sparse Communication between modulles occurs through communication attention "CA"
Head
Mission Instruction
dymamics
Recurrent Independent Mechanisms
Recurrent Independent Mechanisms
M
Ingut
Communication
Attention
Figure 1: Proposed Model architecture: Input images processed through encoder, and the em- bedded mission instruction are passed through a set of recurrent independent modules or RIMs (Goyal et al. 2019) which compete using attention mechanisms to capture the dynamics of the system. The network's state is divided into N modules, such that at every time step, conditioned on a given input 3i, only subset out of total N modules are dynamically activated and updated fast in the inner learning loop, while the parameters of the two attention mechanisms, Input Attention and Communication Attention, are learnt slowly in the outer loop as meta parameters.
et al. 2020) to learn a modular architecture consisting of a set of independent modules which compete with each other to attend to input and sparsely communicate using a key-value attention mechanism (Bahdanau et al., 2014, Vaswani et al., 2017, Santoro al., 2018). For the second objective of fast transfer, we adopt meta-learning approach on the two sets of parameters (those of the modules and those of the attention mechanisms) with the goal of achieving fast adaptation to changes in distribution
agent (Duan et al. [2016] Mishra al. 2017 Wang et al. 2016 Nichol et al., 2018 Xu et al. 2018 Houthooft et al.) 2018 Kirsch et al. 2019).
task in reinforcement-leuuee
In this paper, study the generalization ability of the proposed modular architecture tasks not seen during training. We conduct experiments in the domain of grounded language learning, in which poor data efficiency one of the major limitations for agents to learn efficiently and generalize well
(Hermann et al.) 2017] Chaplot et al., 2017] Wu et al. 2018 Yu et al., 2018 "heevalie-Booiver al.] (018). We show, empirically, how the proposed learning agent only on the seen data, but also is more sample efficient, faster to train and adapt, and has better transfer capabilities the face of changes in distributions. We thus present evidence that combining meta-learning with modular architectures can help in building agents that learn and leverage the compositional properties of the environment to generalize better on novel domains and achieve better transferability and more systematic generalization.
generalize better not
META LEARNING OF RECURRENT INDEPENDENT MECHANISMS
We intend to assay whether modular architecture, combined with learning different parts of the model on different timescales, can help in decomposing knowledge into re-usable pieces such that the resulting model is not only more sample efficient, but also generalizes well across changes in task distributions. We first give high-level overview, and then describe the components the model and the two learning phases in more detail: Section 2 contains an overview of the modular architecture consisting an ensemble recurrent modules, and Section (..iiionn222pppp explains the meta-learning approach to learn the parameters of the modular network at different time scales.
Modular Network. The proposed method based on RIMs architecture (Goyal al. 2019) which consists of set of competing modules, such that each module acts independently and interacts with other modules sparingly through attention.
Attention Networks to modulate information. The soft-attention mechanisms control the flow of information in the model such that different modules attend to different parts of the input via input attention, and module queries relevant contextual information from other modules via communication attention. The two attention mechanisms can use multiple attention heads and have their own set of parameters, and the different modules have their own independent parameters.
2

---

Published as a conference paper at ICLR 2021
Meta Sequence S1
Meta Sequence 52
Imput
Individual Episodes
Individual Episodes
Fast Inner Updates
Inner loop for last learning module parameters
Fast Update
Fast Update
Slow Outer
Updates
Outer loop for meta leamiing of attention parameters
Slow Update
Slow Update
ASS
Inner Update:
subset of "K" modules dynamically updated every timestep
Input Communication Attention parameters updited
MA!
The two learning loops (fast and slow) of the meta-learning setup learn different parameters of the model at different timescales. Several episodes are concatenated to form long sequences, $1, 52,... The learning happens over these meta sequences such that at every time-step, conditioned on the current input It, a dynamic subset of k (out of total N) modules learnt over shorter time-spans to capture the quickly changing dynamics, while the two sets of attention mechanisms, which define the connectivity structure and sparse communication between the modules, are learnt over much longer time-spans to capture more stable connectivity structures.
Figure 2: Meta-Attention Setup:
Learning at different time scales. In the proposed model, the entire network is trained end-to-ent using meta-learning setup such that different components of the network are trained at different time scales and some parameters adapt quickly. while others are updated much less frequently.
Fast learning of parameters of relevant modules. In order to capture the quickly changing aspects of the underlying task distribution, the parameters of the activated modules (most relevant to the current input) updated in fast manner by learning over multiple short spans of the input sequence.
Slow learning of stationary meta-par tata-apaameeess of attention mechanisms. In order to capture the more stable aspects of the task distribution, the parameters of the two sets of attention mechanisms are updated less frequently over relatively longer spans the input sequence. These attention mechanisms are responsible for the activation of the relevant modules and for the modulation of an information exchange between different modules. The parameters of the individual modules are not updated in this phase.
2.1 RECURRENT INDEPENDENT MECHANISMS
In this section, we deseribe the parameterization of the modular network, as well as how attention is used to select relevant modules at a particular time-step.
Recurrent Independent Modules. In Meta-Attention Networks, follow similar modular architecture as in (Goyal et al., (019)). such that the network consists of an ensemble of independently parameterized modules in which each module operates with its individual dynamics and these modules sparingly interact with each other through attention. At given time-step t, for each of the j = L....... ., N modules, let ht.j and D, represent the hidden state vector and the recurrent dynamics, respectively, of the module j. Given an input 2, at time t, the hidden states of the independent modules follow the following three-step step process: First, based on their relevance to the current input re a subset of i (out of N) modules is dynamically activated to form an active set A, consisting of the most relevant modules. Second, these activated modules (in active set As) independently process the input using their default transition dynamics, D, for 1,..., k. while the hidden state remains the same for the other - non-activated modules. In the third step, the modules sparsely communicate with the other modules through attention mechanism.
Two sets of attention mechanisms, (1) Input Attention and (2) Communication Attention, labeled as "iA" and "cA" respectively in Fig. I used selectively activate modules and to modulate sparse communication between them. We explain these two attention mechanisms below.
Soft attention mechanisms. Attention mechanisms used for selective activation of different mod- ules (Input Attention) and for communication between the modules (Communication Attention) based on multi-headed soft attention. The attention is computed as Attention K, V)
3

---

Published as a conference paper at ICLR 2021
softmax QK² Vds
V. where Q, K and V represent the query, key and value, and dy is the size of the features K. The query is obtained as function of where the information is written to, while keys and values are obtained as a function of where the information is read from.
Input Attention for selective activation of competing modules. At every time step 1, each module generates a query. and keys/values are obtained as a function of the input. Then a soft key-value based attention mechanism (Bahdanau al.) 2014; Vaswani al.) 2017 Santoro al.) 2018 is used to select the top-k modules that pay most attention to input dy in a differentiable way using a
sparse softmax (Ke et al. 2018 Goyal et al., 2019). Using this style of competition between different modules has shown to improve specialization among different modules (Goyal et al. 2019 2020).
Communication Attention for sparse communication between modules. In order share informa- tion between different modules, modules communicate with each other through second attention mechanism, called Communication Attention. During this process, only the activated modules in A, can read relevant contextual information from all the other modules, while the non-activated modules remain unchanged. This process is also based on soft key-value based attention (Bahdanau et al. 2014, Vaswani et al., 2017). Here, each module in the active set A, makes query, while keys and values are a function of other modules (irrespective of if they are active or not).
Hyperparameters. In the proposed method, the values of lc and N are hyper- spper-parameters that stay constant throughout the learning of a particular task as well as across different environments.
2.2 META-LEARNING ATTENTION TO MODULATE INFORMATION BETWEEN MODULES
In order to appropriately capture the quickly vs. slowly changing aspects of the underlying distribution, we use a meta-learning setup such that not all the parameters are trained at the same pace: some parameters are adapted more quickly than others and the updates happen in their respective loops only, Throughout the meta-learning process, the input observation sequence from several concatenated episodes is divided into shorter and (relatively) longer horizons, as visualized in Fig. 2 and the learning of the full set of network parameters happens over the following two update loops:
Inner / Faster loop. Conditioned on the input 2, at time f, the parameters of the & most relevant modules in active set A, are updated over several short spans of the input sequence, while keeping the parameters of the attention mechanisms constant. Given the dynamic selection of the k most relevant modules, different sets of modules may get updated over different iterations of this update loop.
Outer / Slower loop. In the outer loop, longer spans of input observations are considered to slowly update the parameters of the attention mechanisms, while the parameters of the independent modules are kept unaltered. This captures the more stable aspects and connectivity patterns across the network.
Loss Function. For our experiments, we used the Proximal Policy Optimization algorithm (Schulman et al.] (017)) using parallel data collection. During the training of the parameters of the network, we use separate output heads for the policy and value function. These two heads are updated separately on different learning time scales, such that the value function head is updated slowly in the outer meta loop along with the attention parameters to capture regularities the environment, while the policy head is updated in a fast way in the inner loop along with the parameters of the selected recurrent modules. In each loop, only the respective parameters specified for that loop are updated while keeping others unaltered, and the outer loop gradient is not obtained by backpropagating through the inner loop updates. For this work, the number of steps through which the gradients are back-pru poocppropagated in time in the outer loop is four times longer than in the inner loop, making the outer loop sweep across four times longer time-steps as compared to the faster inner loop.
3 RELATED WORK
Meta-Learning:
Meta-learning (Bengio et al. 1990 Schmidhuberr 1987)
methods gives the flex-
ibility to adapt to new environments rapidly with a few training examples, and has demonstrated success in both supervised learning such as few shot image classification (Ravi & Larochelle] (016)) and reinforcement learning (Wang et al., 2016 Santoro et al., 2016; Duan et al.| [016 Botvinick et al. 2019 Clune 2019) settings. The most relevant modular meta- learning work is that of Alet
et al. (2018), which proposes to learn modular network architecture based on MAML, however
their approach relies on pre-trained composable transformations, The goal of the current work is to
4

---

Published as a conference paper at ICLR 2021
learn the transformations (i.e decomposition of knowledge into separate modules), as well how to dynamically route information among such modules.
Recently (Bengio et al. 2019 Ke et al. 2019) used meta-learning to learn causal mechanisms or causal dependencies between a set of high- variables, that inspired the approach presented here. The 'modules' in their work are the conditional distributions for each variable in a directed causal graphical model (Schölkopf et al., 2016). The inner-loop of meta-learning also allows the modules to be adapted within an episode (corresponding to an intervention distribution), while the outer-loop of meta-learning discovers how the modules are connected (statically) to each other form the graph structure of the graphical model.
Meta-Learning to Disentangle Causal Mechanisms:
Modular Networks. The generalization capabilities of neural models may benefit from structure in the design of models, to make them resemble the kind of rules they are supposed to learn (Andreas
et al.
2016 Johnson et al. 2017/Jacoss et al. 1991, Bottou & Gallinari]
1991; Ronco et al.
1997.
Reed& & De Freitas 2015 Andreas al. 2016, Rosenbaum et al. 2017; Fernando et al.
2017
Shazeer et al. 2017 Kirsch et al. 2018 Rosenbaum al. 2019). These networks has an architecture
which is composed dynamically from several neural modules, where each module is meant to perform
a distinct function. In such architecture different modules are applied one at time, in contrast to the proposed architecture, where different modules can update their state and be used for prediction in parallel. The focus of this work is to automatically decompose knowledge into set of independent modules as well as learn input-dependent connectivity between these modules. We show that using meta-learning setup helps in learning better decomposition of knowledge such that the resulting model can generalize better out-of-distribution scenarios.
EXPERIMENTS
We evaluate the proposed Meta-RIMs networks to answer the following questions: (a) Does the proposed method improve sample efficiency? We answer this positively in section [.] (b) Does the proposed method lead to policies that generalize better to systematic changes to the training distribution? We find positive evidence for this in section 4.2 (c) Does the proposed method lead to a faster adaptation to new distributions and a better curriulumm learning regime to train agents in an incremental fashion by reusing the knowledge from previously learnt similar tasks? We evaluate this setting and find positive evidence in section 4.3 We also conduct ablation studies to understand the importance of different components of the proposed setup (attention-controll modularity and meta-Icarning) and these are summarized in section 4.4
The experiments are based a large variety of environments from the MiniGrid and BabyAl suite (Chevalier- (Chevalier-Boisvert et al. 2018) which provide an egocentric and partially ob- served view of the environment. Partial observability, sparse rewards, and procedurally generated series of environments with systematically incremental difficulty make faster learning particularly challenging for reinforcement learning agents, but make it useful to address the questions we raised above. Throughout the training, the agent is presented with sparse rewards, with a positive reward given only when the goal is reached within a certain number of timesteps. We use the metries of mean reward Rt and average success-rate S (i.e. the mean of the rewards obtained and the mean of the percentage of times the agent succeeds in achieving the goal, across multiple runs) in the results, showing the learning curves (as function of number of input frames). Please refer to Appendix A.] for additional details on the environments and hyperparameters used.
Environments:
hn
We compare the performance of Meta-RIMs Networks with the following baselines in which all the parameters of the network are learnt at the same time without any meta-learning setup: (a) A Vanilla LSTM model (b) A modular network (i.e RIMs (Goyal et al., 2019)). These are referred to as "metaRims", "vanilla" and "modular", respectively, in the following sections.
Baselines:
4.1 IMPROVED SAMPLE EFFICIENCY
Training agents in a sample efficient manner is one of the major challenges in reinforcement learning. Training different parameters of the model at different timescales in the proposed setup leads to more sample efficient training regime, evident across wide range of BabyAl and MiniGrid tasks Toocol, PickupDist, GollledBall, DoorKey, PutNear, Fetch, FourRooms GoToObj. MemoryS1 3Rnnoom), see Fig. for Mean Reward R. and Fig. 4 for Mean Success Rate S.
5

---

Published as a conference paper at ICLR 2021
Figure 3: Sample Efficiency on MiniGrid and BabyAI environments: proposed method (labeled "metaRims") with modular architecture and meta-learning (red), outperforms the "modular" RIMs (green) and "vanilla" L.STM (blue) baselines. Improvements (in mean-return vs number of frames seen) are more profound in more difficult environments.
4.2 BETTER POLICY GENERALIZATION
We evaluate the capability of the proposed model to transfer knowledge from easy environment to similar but incrementally more difficult environments, such that the environments share most aspects of the underlying structure. More specifically, the agent has to find key unlock door and reach the goal in the other end of the room. As the environments get progressively more difficult, the rooms become larger leading even longer trajectories and sparser rewards, see Fig. To Appendix. the scenario of zero-shot generalization to harder environments, the proposed "meta" method outperforms the other two baselines by substantial margin in the more difficult environment, see Table [ We found the three methods (meta, modular, and vanilla) to not degrade much in the backward direction.
Target Environment
R(MetaRims)
R( R(Vanilla)
R(Modular)
S(MetaRims)
S(Vanilla)
S(Modular)
Difficult
0.97
0.95
0.85
1.0
1.0
0.89
More Difficult
0.45
0.18
0.01
0.47
0.19
0.02
Table 1:
Zero shot Policy Transfer: The model is trained the easiest environment, and transferred a zero-shot manner to more difficult and larger environment, outperforming the baselines terms of both rewards (R) and success rates (S) the difficulty of environment increases.
4.3 EFFICIENT TRAIIIIGN AND KNOWLEDGE TRANSFER FOR CURRICULUM LEARNING
In order to evaluate how efficiently the proposed setup can leverage the previously learnt knowledge, we trained an agent on a source environment and used as pre-trained model to adapt to target environment. To allow reuse of knowledge across tasks, the source and the target environments share some underlying structure in the form of a set of core competencies, namely "Room". "Distr-Box", "Distr". As evident in Fig- a)a), we find that the proposed model adapts much better compared to the regular LSTM setup, indicating more efficient reuse of the past experience.
4.4 ABLATION AND ANALYSIS
Benefits of Meta-Learning Setup. In order to understand the impact of the proposed meta-learning setup and training different parameters at different timescales, trained vanilla LSTM baseline using a similar meta- leaanni setup, in which the parameters the LSTM and the policy head are updated more often in the inner loop, while the parameters of the value function head updated in

---

Published as a conference paper at ICLR 2021
Figure 4: Sample Efficiency: The proposed method (labeled "metaRims") (red), outperforms the baselines of "modular" RIMs (green) and "vanilla" LSTM (blue) in terms of the Mean Success Rate
PickUpLoc, Pretrained
GoToLocal
Return
3.0
03
0.7
0.5
meraflims
3.4
0.5
0.00
0.25
0.50
0.75
10
1.25
3350
1.75
2.50
Figure 5: Knowledge for Curriculum Learning and Ablation study experiments: (a) The pro- posed model and the baseline are trained on an easier environment and then fine-tuned on a more difficult environment such that the environments share subset of core competencies. The proposed method outperforms the baseline, indicating a better reuse of the previous knowledge. (b) To under- stand the importance of meta-learning, fashion, called "meta-Istm", which led to an improved the performance over vanilla setup.
trained vanilla LSTM baseline in similar meta-learning
the slow meta loop. We found an improvement in this "meta-LSTM" setup, see Fig. 5 b), showing positive evidence towards the benefits of meta-learning the proposed approach.
Sparsity in module activation and Slow-factor of Outer loop: In our experiments, we found that updating only a sparse subset of modules performs much better than updating all the modules fast learning loop, indicating that sparsity plays an important role. For the slow outer loop, we found that increasing the length of the interaction span helps in faster learning, see Fig. (a).
Value Function Visualization: We plot the value function predicted by the agent along different timesteps fixed length input sequence and plot the values across nine episodes of a meta episode, find that the value function shoots high value
Fig. 7 Teasing apart of the episodes, when the agent sees the target, and topples down as soon as the target is achieved, indicating that the agent is indeed finding good states and is acting according
properly learnt value function.
Visualizing Module Activations: We plotted module activations across two environments and found that the activations quite diverse with dead modules, indicating active participation from all

---

Published as a conference paper at ICLR 2021
Sparsity
n=5, l=3
n=5, k=2
0.6
40
60
80
100
120
0.4
Input Sequence
0.2
0.0 0.5
1.0
2.5
2.0
3.0
Figure 6: Ablations: (a) Sparsity in Module Activations and Slow-Factor for outer loop: As shown here, sparse activation and update of modules in the inner loop leads to faster learning and better performance. Also, longer spans of meta episodes in outer loop help in faster learning. (b) Module Activations: A plot of module activations (g-axis) for a fixed-length input sequence (æ-axis) for two of the environments for settings (n 5, k 3) and (n 5, k 2) shows diverse and active participation from all modules (no dead modules) to dynamically respond to the inputs received.
Value function
19.0
rightt
18.5-
Done
17.5
Episode
Frame
Frame
Frame 11
Frame
Episode
17.0
Eplsode
Episode the red box
Tramten
The predicted value function plotted across nine episodes of a meta episode. Each red marker marks a new episode and "a" is the action taken. The frames shown are the partial observations received by the agent within one such episode. We can see that the agent understands potentially good states, the right targets and the states that need to be explored further.
Figure Value function visualization:
of the modules to dynamically attend to variations in the observed input sequence, see Fig. 6 (b) and Appendix A.2 for further discussion.
Reversing the Role of Attention Parameters: In order to evaluate the choice of using attention parameters as meta parameters the proposed setup, we reversed the roles of the parameters such that the attention parameters act as quickly changing parameters in the inner loop and the parameters of the dynamic modules act as slowly changing parameters. We found that this did not perform very well, Fig. emphasizing the particular importance of the role of attention parameters acting as slowly changing meta-parameters in the proposed setup.
Importance of Fast and Slow Update Loops: We compared the proposed setup with setting in which, instead of learning attention parameters slowly as meta-parameters in the outer loop, we reduced the learning rate of the attention parameters by one-fourth of the learning rate of other parameters. This way, all the parameters are learnt at the same time, but attention parameters are updated much more slowly. We found this particular setting, model does not perform very well, and the learning slows down significantly as compared to the proposed setup, see Fig. 8
What are the Modules Learning: We visualized the module activations of trained agent for the MemoryS1 1 3Raddom environment across 20 episodes and tracked the agent's 's behavior towards the end of the episode when that (1) When the agent is about to reach the goal, module is always activated, (2) Module is activated when the agent is about to reach the junction at the end of the hallway indicating differennt roles these modules can play in different parts of the state space.
about reach the goal, see Appendix Fig. TT and Fig. [2 We found
Specific Roles of the Active Modules: To further validate the importance of the specific roles played by the active modules, we randomly turned off subset of active modules. Specifically, first randomly deactivate one module and then two of the three active modules. As expected, the performance degrades significantly, such that for 20 episodes, with 3 active modules, the agent takes 192 frames, while takes much longer as modules are gradually deactivated, with 243 frames for two active and 2179 frames with only one active, see Fig.
8

---

Published as a conference paper at ICLR 2021
MemoryS13Random
Memory513Rander
Mcan
Success Rate
matafi
metalimma
melafip
metafiip
Woan hutur
0.8
3.8
0.4
2.5
0.50
0.25
3.50
1.75
1.00
1:25
1.75
0.00
0.25
0.50
2.75
3.50
1.25
3.50
3.75
tramen
166
Trames
Figure 8: Ablation: (1) Choice of attention parameters as slow-parameters: In order to evaluate the choice of learning attention parameters as slow parameters in the proposed setup, we experimented by reversing the roles such that the attention parameters are quickly changing parameters and the module parameters act as slow parameters. As seen from the curve labeled "metaFlip", this reversed- parameter setting does not do as well in terms of both Mean Reward and Success Rate. (2) Slowing down the learning rate for Attention Parameters: We also compare the proposed setup with setting in which the learning rate of the attention parameters is reduced by one-fourth of the learning rate of other parameters and all parameters are learnt at the same time. The curve labeled "slowLK shows that just lowering the learning rate of attention parameters does not help. and in fact slows down the learning, emphasizing the importance of the two learning loops in the proposed setup. Totalo5, 3
Totalo5, Active
0.2
0.0
tramess
Figure 9: Ablation: Special role the winner modules Turning off modules that win the compe-
tition significantly degrades performance, and agent takes much longer to complete an episode.
5 CONCLUSION
This paper investigates using meta-learning approach modular architectures with sparse commu- nication (as in RIMs (Goyal et al. 2019) to capture short-term vs long-term aspects of the underlying mechanisms in the data generation process, by considering parameters of attention mechanism as meta-parameters and parameters of the recurrent modules as parameters. Meta-lcarning is directed towards sample efficiency in the face of changes in distribution and tasks. The experimental results on grounded language learning tasks in the reinforcement learning setting strongly indicate that the combination of meta-learning combined with dynamically connected modular architectures with sparse communication, leads in many ways to superior results in terms of improved sample efficiency (faster convergence, higher mean return and success rates), and an improved transfer across tasks in a curriculum, both zero-shot transfer and with adaptation. Ablation studies further confirm that using a meta-learning approach to update different parameters of the network over different timescales leads to improvements in sample efficiency compared to training all the parameters at once. Overall, these results point towards an interesting way to perform meta-learning and attention- modularization for better sample efficiency, duiti--distiibution generalization and transfer learning. In the present work, different modules are still trained end-to-end to optimize a single objective, and future area of research would investigate how to train different modules, where each module has separate objective and optimizing the module specific objective optimizes the global objective.
6 ACKNOWLEDGEMENTS
The authors would like to thank Stefan Bauer and Alex Lamb for useful discussions. This project internally at Mila was know "Meta-Attention Networks".

---

Published as a conference paper at ICLR 2021
REFERENCES
Ferran Alet, Tomás Lozano-Pérez, and Leslie P Kaelbling. Modular meta-learning. arXiv preprint arXiv: 10166, 2018.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition, pp. 39-48, 2016.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0 1409.4479 2014.
Dzmitry Bahdanau, Shikhar Murty, Michael Noukhovitch, Thien Huu Nguyen, Harm de Vries, and Aaron Courville. Systematic generalization: what is required and can it be learned? arXiv preprint arXiv:/811.12889 2018.
Yoshua Bengio. Samy Bengio, and Jocelyn Cloutier. Learning synaptic learning rule. Citeseer, 1990.
Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Rosemary Ke, Sébastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. arXiv: 2019.
Léon Bottou and Patrick Gallinari. A framework for the cooperation of learning algorithms. In Advances in neural information processing systems, pp. 781-788, 1991.
Matthew Botvinick, Sam Ritter, Jane X Wang, Zeb Kurth-Nelson, Charles Blundell, and Demis Hassabis. Reinforcement learning, fast and slow. Trends in cognitive sciences, 23(5):408 422, 2019.
Devendra Singh Chaplot, Kanthashree Mysore Sathyendra, Rama Kumar Pasumarthi, Dheeraj Rajagopal, and Ruslan Salakhutdinov. Gated-attention architectures for task-oriented language grounding. arXiv preprint arXiv: 0aciv:777000000 7706077236 2017.
Maxime Boiseeir, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: First steps towards grounded language learning with human in the loop. arXiv preprint arXiv:/810.08272, 2018.
Jeff Clune. Ai-gas: Ai-generating algorithms, an alternate paradigm for producing general artificial intelligence. arXiv preprini urXiv: 10985, 2019.
Yan Duan, John Schulman, Xi Chen, Peter Bartlett, Ilya Sutskever, and Pieter Abbeel. RI:: Fast reinforcement learning via slow reinforcement learning, arXiv preprint arXiv: 2016.
Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha, Andrei A Rusu, Alexander Pritzel, and Daan Wierstra. Pathnet: Evolution channels gradient descent in super neural networks. arXiv preprint arXiv: (eXXi:1701.000 1701.0834, 2017.
Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. arXiv preprint arXiv:2011.15091 2020.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909. 10893, 2019.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006. /6255, 2020.
Karl Moritz Hermann, Felix Hill, Simon Green, Fumin Wang, Ryan Faulkner, Hubert Soyer, David Szepesvari, Wojeiech Marian Czarnecki, Max Jaderberg, Denis Teplyashin, et al. Grounded language learning in simulated 3d world. arXiv preprint arXiv: /tt.//0666.55555 2017.
Rein Houthooft, Yuhua Chen, Phillip Isola, Bradly Stadie, Filip Wolski, OpenAI Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. In Advances in Neural Information Processing Systems, pp. 5400-5409, 2018.
10

---

Published as a conference paper at ICLR 2021
Robert A Jacobs, Michael I Jordan, Steven J Nowlan, Geoffrey E Hinton, et al. Adaptive mixtures of local experts. Neural computation, 3(1):79-88 87, 1991.
Jason Jo, Vikas Verma, and Yoshua Bengio. Modularity matters: Learning invariant relational reasoning tasks. arXiv preprint arXiv: 1806. 06765, 2018.
Justin Johnson, Bharath Hariharan, Laurens Van Der Maaten, Judy Hoffmann Li Fei-Fei, C Lawrence Zitmick, and Ross Girshick. Inferring and executing programs for visual reason- ing. In Proceedings of the IEEE International Conference on Computer Vision, pp. 2998, 2017.
Nan Rosemary Ke, Anirudh Goyal ALIAS PARTH GOYAL, Olexa Bilaniuk, Jonathan Binas, Michael C Mozer, Chris Pal, and Yoshua Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in Neural Information Processing Systems, pp. 7640-7651, 2018.
Nan Rosemary Ke, Olexa Bilaniuk, Anirudh Goyal, Stefan Bauer, Hugo Larochelle, Chris Pal, and Yoshua Bengio. Learning neural causal models from unknown interventions. arXiv preprint arXiv: 2019.
Louis Kirsch, Julius Kunze, and David Barber. Modular networks: Learning decompose neural computation. In Advances in Neural Information Processing Systems, pp. 2408-2418, 2018.
Louis Kirsch, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Improving generalization in meta reinforcement learning using learned objectives. arXiv preprint arXiv:1910.04098, 2019.
Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to-sseueene recurrent networks. In International Conference on Machine Learning, pp. 2873-2882. PMLR, 2018.
Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta- learner. arXiv preprint urXiv: 2017.
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning combine top-down and bottom-up signals in recurrent neural networks with attention over modules. arXiv preprini arXiv:2006. 16981, 2020.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint 100.00900000999, 2018
Nasim Rahaman, Anirudh Goyal, Muhammad Waleed Gondal, Manuel Wuthrich, Stefan Bauer, Yash Sharma, Yoshua Bengio, and Bernhard Schölkopf. S2rms: Spatially structured recurrent modules. arXiv preprint arXiv:2007.06533 2020.
Sachin Ravi and Hugo Larochelle. Optimization as a model for few-shot learning. 2016.
Scott Reed and Nando De Freitas. Neural programmerrinteeppetee. arXiv preprint 11X15:1511.00279, 2015.
Eric Ronco, Henrik Gollee, and Peter Gawthrop. Modular neural networks and self-deccopossttin. Technical Report CSC-96012, 1997.
Clemens Rosenbaum, Tim Klinger, and Matthew Riemer. Routing networks: Adaptive selection of non-linear functions for multi-task learning. arXiv preprint arXiv: 2017.
Clemens Rosenbaum, Ignacio Cases, Matthew Riemer, and Tim Klinger. Routing networks and the challenges of modular and compositional computation. arXiv preprint 1904://904./2770 2019.
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta- learning with memory-augmentee neural networks. In International conference on machine learning, pp. 1842-188 1850, 2016.
Adam Santoro, Ryan Faulkner, David Raposo, Jack W. Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy P. Lillicrap. Relational recurrent neural networks. CoRR, abs/ 0ts//8006001870 2018. URL URL http:htttt//tt htt/ 0122
11

---

Published as a conference paper at ICLR 2021
Jurgen Schmidhuber. Evolutionary principles in self-referential learning. On learning how to learn: The meta-meta.... hook.) Diploma thesis, Institut f. Informatik, Tech. Univ. Munich, 1:2, 1987.
B. Schölkopf, D. Janzing, and D. Lopez-Paz. Causal and statistical learning. In Oberwolfach Reports, volume 13(3), pp. 1896-1899, 2016. doi: 10.1476000W--201--33. URL https: //publications mff de httph/aannllmmoo/aaaaaoooa
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 177.v:1707.06347, 2017.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Lc, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-expenee layer. arXiv preprint arXiv: XXX1::770..065 066538, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Jane X Wang, 7eb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z. Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint ar1.v:1611.05763, 2016.
Yi Wu, Yuxin Wu, Georgia Gkioxari, and Yuandong Tian. Building generalizable agents with realistic and rich 3d environment. arXiv preprint (rXiv:/801.022 022209, 2018.
Zhongwen Xu, Hado P van Hasselt, and David Silver. Meta-gradient reinforcement learning, In Advances in neural information processing systems, pp. 2396-2407, 2018.
Haonan Yu, Haichao Zhang, and Wei Xu, Interactive grounded language acquisition and generaliza- tion in 2d world. arXiv preprint arXiv:1802.01433, 2018.
12

---

Published as a conference paper at ICLR 2021
A APPENDIX
A.1 IMPLEMENTATION DETAILS
We used variety of environments from MiniGrid and Baby Sntevalier-Boisveed et al. 2018) that provide a partial and egocentric view of the state of the environment to the agent. The reward sparse and positive reward is received only if the agent successfully reaches the goal. A penalty is awarded based on the number of steps taken to reach the goal, calculated as - 0.9n//nmas where Umaz is the maximum number of steps allowed for given environment and depends on the difficulty of the environment such that more difficult environments have a larger value of 1maax- If the agent not able to complete the task within Rmmaz steps, the episode ends and it gets zero reward. The environment have an increasing level of difficulty in an systematically incremental manner. These settings of partial observability, sparse rewards and systematic increase in the difficulty levels make the task for reinforcement learning algorithms sufficiently difficult.
The observation received by the agent consists of two parts: (1) an RGB image for the partial observation of the agent's field of view, and (2) language instruction containing the mission statement that defines the task. The image part of the observation is processed through an encoder, and an embedding is computed for the textual mission statement, which are then fed to an ensemble of recurrent modules; see Fig. for a full visualization of the architecture layout.
We used the Proximal Policy Optimization (Schulman et al. 2017) with parallelized data collection of rollouts collected by multiple parallel processes. For generalized advantage function, we used A 0.99, and discounted future rewards by factor of y 0.99. Throughout the experiments, present the rean-reward (R) and success-rate (S) of the agent, where the mean reward is the average reward across multiple runs, and the success rate represents the percentage of times the agent is able to successfully reach the goal within the Rmaa timesteps. For all of our environments, we used n = 5 total modules, with only k- = 3 of them active at any given time. Further details the specifics of each environment are provided in the Section A.4
A.2 ABLATION: VISUALIZING MODULE ACTIVATIONS
In the proposed setup, the modules are dynamically selected and activated based on their relevance to the current input. In order to visualize the diversity and frequency of module activations, and to analyze the effect of environment dynamies on the activation patterns of the modules, we tracked the module activations of a trained agent along fixed length input sequence obtained by tracking the agent navigating across the environment. We plotted these activated modules for two different environments DynamicObstacles and PutNear, see Fig. (b).
We find that for environment DynanicObstacles, in which the agent has to take more frequent decisions to figure out the best action due to quite engaging environment dynamics, such multiple distractor objects and dynamically moving obstacles, the module activations are more diversely activated. In both the environments, different modules active for different inputs, and all of them get activated at some point depicting an active engagement throughout the agent's is interaction. However for the environment, PutNear, which has relatively fewer moving parts during the interaction, some modules are consistently more active than the others. However, all modules do get activated at some point, leaving no dead or always inactive modules.
This study has a direct analogy with how humans engage with their environments: a visual input with an enriched set of objects and more frequent interactions with the environment will need higher engagement and continuous application of wider set of skill sets. On the contrary, an environment which is relatively simple and has simpler input observations would need a lower engagement, and only a few components/ / skill-sets would mostly be able to handle and decide the best course of actions.
A.3 ABLATION: VISUALIZING EPISODES AND MODULE ACTIVITY PATTERNS
In order to visualize the module activity pattern of an agent, and to understand the patterns in what the modules are learning, we plotted the module activity patterns over 20 episodes, see Fig. [add Fig. 12 On visually inspecting the behavior of the agent over these 20 episodes, we found that some
13

---

Published as a conference paper at ICLR 2021
9
Figure 10: Zero-shot transfer environments: A series of DoorKey environments with an increasing level of difficulty from left to right; as the rooms get larger, the trajectories become longer and the rewards become sparser, making the tasks progressively more difficult. At the same time however, these environments share an underlying structure, making them suitable for evaluating policy generalization and transfer. We trained our agent the leftmost (easiest) setting and evaluated it on the more difficult environments.
specific modules are always activated when a certain state space is reached: Module is always activated when the agent about to reach the matching object (the end goal), and Module is always activated when the agent is around the second to the last step and when the agent is at the end of the hallway at the split. This shows specific assignment of roles the modules undertake to achieve an optimal behavior.
A.4 ENVIRONMENTS: MINIGRID AND BABYAI
We trained our agents on several environments from MiniGrid and BabyAl AI Cheevalier-Boisverr al. 2018), such GoTolo PickupDist, Dynamic Obstacles, DoorKey, PutNear, Fetch, Four- RoomsST3, GoToObj. GoToRedBall, GoTTReeBalllGryy MultiRoom, LavaCrossingSNNi and Pick- UpLoc In each of these environments, the observation consists an RGB image of the agent's partial field of view, and an instruction textual string that contains the mission that the agent needs to solve. Fig. visualizes the model architecture and the processing the input observation. Further details on these environments, with the hyperparameters used for the recurrent modules, are provided below.
For all the environments, we used a total n modules with k active at any give time. More details each environment provided below.
Hyperparameters:
GoToLocal: The agent is presented with number of objects in single room with doors, and asked to go to one of them as specified in the mission statement.
PickupDist: A single room, with no doors, has a number of objects and the agent needs to pick the object instructed in the mission statement.
Dynamic Obstacles: This environment contains moving obstacles single room, and the agent has to reach the goal located in a comer of the room while not colliding with these dynamically moving obstacles. large penalty is subtracted if the agent collides with an obstacle. This environment particularly useful to train moving robots in dynamically moving objects setting.
DoorKey: In this environment, the agent needs to find key in the current room, pick the key and unlock a door, and then reach the goal located the other room.
The mission statement specifies an object that has to be put near another object. room of multiple objects, the agent has find the correct objects and put them in the right location as specified by the instruction statement, thus requiring both spatial and visual understanding of multiple objects.
PutNear:
14

---

Published as a conference paper at ICLR 2021
Fetch: The agent is presented with multiple objects of different colors and shapes, and it needs to find the object that is instructed to be fetched the mission instruction string. A negative reward is given if wrong object is picked.
MemoryS13 : The agent starts in a room in which it sees an object, and has to remember this object when it reaches the end of the hallway which ends split. One of the sides of this split randomly contains the same object as it saw in the beginning and it has to choose that matching object.
FourRoomsS13:
This is multi-room environment, consisting of four rooms, connected by four gaps in the walls, such that the goal is randomly placed in one the rooms. The location of the agent also random when the episode starts.
GoToObj:
This is a single room environment with single object, specified in the mission string, that the agent needs to reach
GoToRedBall:
The goal is reach red ball in a room containing distractors and other obstacles.
GoToRedBallGrey:
The room consists of multiple grey distractors and the agent needs to reach the red ball, without requiring any unblocking.
The environment consists of a series of interconnected rooms, and the agent has to unlock the door navigate the goal in the next room.
MultiRoom:
The enviromment consists of horizontal and vertical strips of lava running across the room, and the agent has to reach the goal while avoiding them. If the agent touches the lava, it dies and the episode ends with reward.
L:vaCrosinngSNNI:
PickUpLoc: A single room environment in which the goal is to pick up an object described by its location.
15

---

Published as a conference paper at ICLR 2021
Figure 11: Visualizing Activity of Modules as the Agent Navigates: Visualizing Episodes 1-12,
along with the activations of 5 modules such that are active at any give time. We find that: (1) When the agent is about reach the goal, Module is always activated, and (2) Module 2 is always activated when the agent reaches the split junction at the end of the hallway towards the last few steps of the episode, indicating a particular trend in the module activations in certain parts of the state space.
16

---

Published as a conference paper at ICLR 2021
Hi
60008
1
Figure 12: Visualizing Activity of Modules as the Agent Navigates: Visualizing Episodes 13-20,
along with the activations of 5 modules such that 3 are active at any give time. We find that: (1) When the agent is about to reach the goal, Module is always activated, and (2) Module 2 is always activated when the agent reaches the split junction the end of the hallway towards the last few steps of the episode, indicating a particular trend in the module activations in certain parts of the state space.
17


Reviews:
[{"review_id": "xQaL7nXhjFr", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nSummary:\n\nThe authors present an approach that leverages inductive biases for knowledge decomposition and relational reasoning with the aim of generalizing over out-of-distibution tasks in RL with stationary rewards.  They propose a meta-learning training framework where parameters of the network are divided into two subsets where one set is learned with fast updates and the other with slow updates capturing the slow changing dynamics of the environment. \n\nThey also propose an architecture that includes N modular components in the form of recurrent independent mechanisms (RIMS, Goyal et al. 2019) and two multi-head dot-product attention components: input attention for selecting k modules and communication attention between the selected modules.  The attention parameters are updated slowly and model the stable dynamics of decomposing task input information as meta-learning parameters while the modules are updated more frequently while the attention dynamics are held constant.  To train the policy and value function the authors chose to use Proximal Policy Optimization (PPO).  Here the value head is updated slowly with the attention meta-parameters while the policy is updated in sync with RIMs.\n\nThis approach is evaluated against GridWorld and BabyAI suite tasks where the task complexity is parameterized and involve navigating in grid worlds with partial observability conditional on input instructions.  Generalization performance is measured on more complex out-of-distibution tasks and the authors also carry out ablation study on the meta-learning setup as well as analysis on the learned value function and module communication sparsity.\n \nStrengths & Weaknesses:\n\nMeta-learning and factored representations are important areas of study if we are to build RL agents that can generalize well and learn new tasks approaching the sample efficiency or small amounts of data required, for instance, by humans.  Therefore it is key to build meta-learning into RL agents and so this work I believe is well motivated and aims to contribute to a crucial area of study.  The work here seems to be largely inspired by Bengio et al. (2019) in learning via a meta-transfer objective among causal entities where the authors use meta-learned multi-head attention mechanisms facilitating the interaction dynamics of a network where the modular components are realized as RIMs.  While this work bears some strong similarity to this work I believe that there are some key differences in the operation of the two attention mechanisms, leveraging the  composability of RIMs and learning framework that make it a novel contribution.\n\nI think there are a few things in the paper that could benefit from more detail, for instance, it would be helpful to provide the definition of recurrent dynamics D_j and how the attention, RIMs and RL all connect together quantitatively.  Also a bit more definition around the tasks in the main paper could be helpful.  I realize fitting things in can be tricky but even a small bit more description could help.\n\nIs it hard to tune the loop sizes over the meta-sequence learning?  I think the impact of this work depends on this and it could be nice to see some analysis or get a high level sense of how different loop sizes affect learning. \n\nOverall I think the results are good and the analyses are helpful in understanding the utility of the approach over the baselines in terms of sample efficiency and generalization performance over more complex tasks.  Another baseline that is able to meta-learn some parameters but without a modular architecture might have been a good point of comparison.  Although, the ablation study with the modified LSTM is a simple version of this it seems.  The analyses of the value function and module sparsity are good and I think help the overall argument.\n\nThe tasks evaluated on are GridWorld and the BabyAI.  These tasks are fairly simple toy tasks but are well suited to demonstrate the power of the model to a limit.   It would be helpful to see the approach evaluated against more complex environments where factorizing representations is more difficult to noise and complexity in the task distribution.  This might help better illustrate the horizons of this area of work.\n\nFinally, the paper is clearly written and well structured.\n\nRecommendation:\n\nI think overall this is a nice paper and would have no problem to see it accepted.  I think the work could be more ambitious but I believe that the authors have illustrated that this is a potentially useful innovation for RL.", "extra_scores": {}, "extra_text": {}}, {"review_id": "nAozgUTnyig", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Objective of paper**: Exploring methods that exhibit better systemic out-of-distribution generalization/adaptation properties. Specifically, exploring how modularity, combined with learning at different time scales, can help with this. \n\n\n**Central claim**: By expanding Recurrent Independent Mechanisms with a meta-learning layer to facilitate learning at different time scales, it is possible to achieve better within-distribution and out-of-distribution/systemic generalization and fast adaptation to new tasks.  \n\n\n**Strong points**: \n1.\tImportant research direction: I believe this research direction has a lot of value in exploring and improving the way our current deep learning systems can exhibit better systemic generalization. \n2.\tNovelty: While most of the individual components presented in the paper are not novel, the particular combination of these components and the empirical evaluation seem novel to me. (That being said, the idea of meta-learning gating mechanisms to combat catastrophic forgetting and interference (arguably the central contribution of this paper) has been explored in continual learning before – for example [1]). \n3.\tCompelling experimental results: In terms of raw performance, the proposed method clearly outperforms its Recurrent Independent Mechanisms (RIM) and LSTM counterparts in sample complexity, better systemic generalization and faster adaptation. I also think that the tasks picked for evaluation are suitable to investigate systemic generalization. \n\n\n**Weak points**: \n1.\tLack of architectural/algorithmic benchmarks in experiments: Given that systemic generalization, as well as combatting catastrophic forgetting, are active research areas, it would have been nice if the paper compared its results with alternative, competing approaches. If you think there are no other comparable approaches in the literature, then it would be nice to discuss why in the related works section. It might also be helpful to try out simpler versions of the proposed method to see what component really made the approach work. For example, to really double down on how much the meta learning component matters, could the authors try training all of the parameters at the same time (i.e. no meta-learning step) but with a small learning rate assigned for the attention parameters (maybe ¼ of the current learning rate, since you update it 4 times less than the dynamics parameters, if I’m not mistaken )? \n2.\tClaim about knowledge transfer not supported well enough: Besides the fact that the knowledge transfer results are reported on a small percentage of the environments, I’m not sure if we can conclude “knowledge transfer” is responsible for better performance when one uses models pretrained on similar tasks. Couldn’t this be due to some simpler explanation, such as better optimization landscape for pretrained models? I think the authors can address this by first pretraining the models on completely irrelevant tasks where knowledge transfer is not expected to help and showing that in this case the reported gains in performance are *not* shown. It might also help elaborate on what is meant by “knowledge”. \n3.\tHyperparameter tuning procedure not described: The hyperparameter tuning procedure was not described in the paper (I might have missed it too?). It would be particularly interesting to read more about how the authors arrived at updating the meta parameters once for every 4 inner loop updates. \n4.\tLack of “hard” environments: It seems from the results that both LSTM and RIM already perform well on most of the tasks shown in the paper (though slower and with more data). Is this a correct assessment? Are there tasks that only the proposed methods can solve? (I guess zero-shot learning results are somewhat of an example, but it would be nice to see harder environments too)\n5.\t(minor) The fact that number of modules (N) and number of active modules allowed at a time (k) are fixed: I guess this can be viewed as a weakness, thought this is shared in RIM as well, and not that relevant to what is being evaluated/investigated in the paper. \n\n\n**Decision**: This paper makes an algorithmic contribution to the systemic generalization literature, and many in the NeurIPS community who are interested in this literature would benefit from having this paper accepted to the conference. I'm in favour of acceptance. \n\n\n**Questions to authors**: \n1.\tIs there a specific reason why you reported results on knowledge transfer (i.e. section 4.3) only on a few select environments? \n2.\tAs mentioned in the “weak points” section, it would be nice if you could elaborate on \n3.\tIs it possible that the caption of Figure 4 is misplaced? That figure is referenced in Section 4.1 (Improved Sample Efficiency), but the caption suggests it has something to do with better knowledge transfer. \n4.\tIf you have the resources, I would be very interested to see how the “small learning rate for attention parameters” benchmark (described above) would compare with the proposed approach. \n5.\tIn Section 4, 1st paragraph, you write “Do the ingredients of the proposed method lead to […] a better curriculum learning regime[…]”. Could you elaborate on what you mean by this?\n\n\n[1] Beaulieu, Shawn, et al. \"Learning to continually learn.\" arXiv preprint arXiv:2002.09571 (2020).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "kkEGkU_KVj2", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of the paper:\n\nThe paper introduces a meta-learning approach to recurrent independent modules, with the goal to make RIMs adapt better to changing distributions, increasing generalization, and increasing sample efficiency. It postulates that RIMs, with their decomposition of independent information, can more easily adapt in these cases. The introduces approach wraps a \"meta-learner\" in form of two different update speeds around RIMs. The \"inner\" loop updates the RIM module parameters with large (fast) updates, while the \"outer\" loop, consisting of input and communication attention, does small (slow) updates. The authors then evaluate the approach on the reinforcement learning domains of \"MiniGrid\" and \"BabyAI\". Comparing their approach with an LSTM wrapped in a similar \"meta-learning\" layer, they show that their approach outperforms the LSTM, and is generally able to quickly adapt to changes in the environment.\n\nCommentary on the goal of the paper:\n\nLearning in non-stationary environements is an open area of research, with a lot of potential. Combining the intuitions between two different approaches - the modularity of a RIM, and the adaptivity of a \"meta-learner\" is a very appealing idea.\n\nStrengths:\n- The paper's main approach is both intuitive and simple\n- The results are good, showing clear improvements over the baseline\n- The paper is generally well-written, and well-structured\n\nWeaknesses:\n- The paper lacks detail in the description of the proposed approach. How exactly does the overall architecture work? How do the different attention mechanisms interact? From the explanations and descriptions given, I would not be able to reproduce the architecture.\n- The paper should have focused more to outline the reasoning behind the fast-slow learning design: why does the \"outer loop\" learn slowly, while the \"inner loop\" learns slowly? This is against my intuitions: I would have expected the outer loop to learn fast, and the inner slow - after all, the appeal of modular architectures is to re-use the modules, not the composition.\n- The paper overstates the novelty of their approach. Under \"modular networks\" on page 5, it claims that the approach of modular networks is \"hampered by the amount of domain knowledge required to decide or predict (Johnson et al., 2017) how the modules should be parameterized, and how they should be connected with each other\". While this is true for the early work of Andreas, this statement completely ignores later research that tries to address exactly this shortcoming (\"Routing Networks\" by Rosenbaum et al 2017; \"Modular Networks\" by Kirsch et al 2018; \"Automatically Composing Representation Transformations as a Means for Generalization\" by Chang et al 2018; \"Recursive Routing Networks\" by Cases et al, 2019). In particular \"Routing Networks and the Challenges ...\" by Rosenbaum, 2019 already mentions slowing the training of the outer loop.\n- The paper falls short on qualitative analysis: what do the modules learn, what is transfered, and how does this help with non-stationary environments?", "extra_scores": {}, "extra_text": {}}, {"review_id": "EVtLDYb6Q6O", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a meta-learning (called) method for the model based on RIMs architecture to realize fast adaptation on new unseen tasks. \n\nThe base model is the same as RIM, which is composed of multiple independent LSTM modules. An input attention is used to determine which independent modules are activated, taking output state from them. Another attention fuses information from activated modules. This paper proposes to update the parameters of attention less frequently than LSTM modules instead of updating them with the same frequency, as they believe that attention should capture a more stable distribution of the task.\n\nHowever, there are some concerns about the motivation of this paper. There is no doubt that meta-learning can achieve faster adaptation between domains. But why should we divide the parameters into two sets? How about treating all parameters as a single set and updating them under a traditional meta-learning way in which parameters are separately optimized on each distribution and the average loss among distributions is used for gradient calculation. In addition, will the distribution change much for different frames within the same task?\n\nThis paper is a bit hard to understand as there are too many long sentences with complex clauses. It is better be improved via revision.\n\nSome questions:\n\n1. concerns about motivation.\n\n2. the current strategy is similar to setting a lower learning rate for parameters of attention modules. Does the paper have any results to prove that the proposed method is better than using different learning rates?\n\n3. why setting the step length of slow update fixed as four times of fast update? Would it be better to determine it dynamically? (e.g. if the frame distribution difference has exceeded a threshold then taking outer update)\n\n4. it is unclear about the parameter updating strategy of baselines (vanilla LSTM and RIM). All parameters in them will be updated under the same frequency. But is the frequency the same as the fast update or slow update in the method.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper examines an idea that knowledge and rewards are stationary and reusable across tasks. An interesting paper that combines number of related topics (meta RL, HRL, time scale in RL, and attention), improving the speed of training.

The authors have addressed the reviewer comments, strengthening the paper. The reviewers  agree, and I concur, that the paper contributes a novel model,  valuable to the ICLR community. It is well thought-out, presented, and evaluated. 
Decision: Accept (Poster)