Task-driven Discovery of Perceptual Schemas for Generalization in Reinforcement Learning
Authors: Wilka Torrico Carvalho
Andrew Kyle Lampinen
Kyriacos Nikiforou
Felix Hill
Murray Shanahan
Venue: ICLR 2022 Submitted (2022)
Keywords: deep reinforcement learning
reinforcement learning
deep learning
compositional generalization
generalization
recurrent architecture

Abstract: Deep reinforcement learning (Deep RL) has recently seen significant progress in developing algorithms for generalization. However, most algorithms target a single type of generalization setting. In this work, we study generalization across three disparate task structures: (a) tasks composed of spatial and temporal compositions of regularly occurring object motions; (b) tasks composed of active perception of and navigation towards regularly occurring 3D objects; and (c) tasks composed of navigating through sequences of regularly occurring object-configurations. These diverse task structures all share an underlying idea of compositionality: task completion always involves combining reoccurring segments of task-oriented perception and behavior. We hypothesize that an agent can generalize within a task structure if it can discover representations that capture these reoccurring task-segments. For our tasks, this corresponds to representations for recognizing individual object motions, for navigation towards 3D objects, and for navigating through object-configurations. Taking inspiration from cognitive science, we term representations for reoccurring segments of an agent's experience, "perceptual schemas". We propose Composable Perceptual Schemas (CPS), which learns a composable state representation where perceptual schemas are distributed across multiple, relatively small recurrent "subschema" modules. Our main technical novelty is an expressive attention function that enables subschemas to dynamically attend to features shared across all positions in the agent's observation. Our experiments indicate our feature-attention mechanism enables CPS to generalize better than recurrent architectures that attend to observations with spatial attention.

Under review as a conference paper at ICLR 2022
TASK-DRIVEN DISCOVERY OF PERCEPTUAL SCHEMAS FOR GENERALIZATION IN REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
Deep reinforcement learning (Deep RL) has recently seen significant progress in developing algorithms for generalization. However, most algorithms target a single type of generalization setting. In this work, we study generalization across three disparate task structures: (a) tasks composed of spatial and temporal com- positions of regularly occurring object motions; (b) tasks composed of active perception of and navigation towards regularly occurring 3D objects; and (c) tasks composed of remembering goal- goal-information over sequences of regularly occur- ring object-configuratoons. These diverse task structures all share underlying idea of compositionality: task completion always involves combining reccurring segments of task-oriented perception and behavior. We hypothesize that an agent can generalize within a task structure if it can discover representations that capture these reccurring sesk-segmntss For our tasks, this corresponds to representations for recognizing individual object motions, for navigation towards 3D objects, and for navigating through jnjeec-configurations. Taking inspiration from cognitive science, we term representations for reccurring segments of an agent's 's experience, "perceptual schemas". We propose Feature Attending Recurrent Modules (FARM), which learns state representation where perceptual schemas are distributed across multiple, relatively small recurrent modules. We compare FARM to recurrent architectures that leverage spatial attention, which reduces observation features to a weighted average over spatial positions. Our experiments indicate that our feature-attention mechanism better enables FARM generalize across the diverse object-centric domains we study.
INTRODUCTION
Cognitive scientists theorize that humans generalize broadly with "schemas" they discover for regularly occurring structures within their experience (Minsky 1979 Rumelhart 1980 Rumelhart al. 1986). Schemas are representations that capture common features over diverse aspects of the environment. For example, when we learn to drive, we learn schemas for common car types (such as sedans), common motions (such accelerating stopping), and common car arrangements (such as row of cars). Importantly, schemas are composable representations over portions of our observations. This allows us recombine them in novel ways. For example, once we learn schemas for sedans, car motions, and rows of cars, we can recognize many rows of sedans moving in opposite directions even if we've 've never seen this before.
While substantial progress has been made on developing deep reinforcement learning (deep RL) algorithms which can generalize, algorithms are typically limited one type of generalization. For example, one algorithm will generalize to novel compositions of familiar shapes and colors (Higgins) al. 2017, Chaplot et al.] (2018), whereas another algorithm will generalize to longer sequences of observed subtasks (Sohn al. 2018 2021 Brooks et al. (021). In this work, hypothesize that we can develop a single deep RL architecture that can exhibit multiple types of generalization if it can leam schema-like representations for regularly occurring structures within its experience. As first step, we study learning schemas for perception (i.e. perceptual schemas) that support generalization within a diverse set of tasks and environments.
We study generalization across three diverse environments and task structures, each with their own regularly occurring structures (Figure TT. Across these environments, test tasks novel compositions

---

Under review as a conference paper at ICLR 2022
(a) Ballet task: observe and remember n distinct object-mations. Language given finish.
[b]30 Place task
Agent is never given (Pickup X, Place on Y) task decomposition.
Regularly occurring structures:
Regularly occurrig structures:
Abject-motions
Perception of 3D objects
Lessetad
Agant
Objects to pick uD
Objects to place
Train: agent recall 2 motions
Test: agent must recall motions
Train: smaller room
Test: larger room
Languaae instruction: "Go to object that moved clockwise"
Language instruction: "Go to object that updown"
Language instruction: "Place toothbrush table"
Language instruction:
"Place toothbrush on table"
(c) KeyBox task: observe the goal box and bring back the key matching box color from across the hallway.
Regularly occurring structures: object configurations
Task:
Train
1. Begin on left side
2. Box determines color
Navigate acroes hallway
3. Oblain goal key H
4. Bring goal key goel box
Test
Figure 1: We study three diverse compositional task structures in both 2D and 3D environments: (a) composing memory of object-motions, (b) composing 3D objects with larger environment, and (c) composing memory of goal information longer tasks made of recurring object configurations. We test an agent's 's ability to generalize within a task structure by testing it on novel compositions of the regularly occurring structures the agent experienced during training. We hypothesize that an agent can successfully generalize at test time if it can discover perceptual schemas that capture these recurring structures. Videos our agent performing these tasks: hntpss//bit.ly///CkAAAg
of the regularly occurring structures the agent experiences during training. Generalization for the "Ballet task" involves recalling novel spatial and temporal compositions of regularly occurring object motions (Figure [] a); generalization for the "Place X on Y task" involves generalizing active perception of and navigation towards regularly appearing 3D objects (Figure [ b); and generalization for the "Keybox task" involves generalizing memory--reennionn of noal-information larger environments composed of sequences of regularly occurring object-configuratoon
(Figure c). We hypothesize that discovering perceptual schemas for these regularly occurring structures will facilitate zero-shot generalization to tasks defined over novel compositions of these structures.
We propose Feature Attending Recurrent Modules (FARM), state representation learning archi- tecture for discovering task-relevant perceptual schemas. FARM learns perceptual schemas that are distributed across multiple, smaller recurrent modules. To consider why this might be helpful, consider the benefits of using word embeddings. A word embedding can represent more informa- tion than one-hot encoding of the same dimension because it can represent different patterns of word usage with the same dimension. Analogously, learning multiple modules enables FARM to represent different patterns of an agent's experience-i.e. different perceptual schemas-with the same module. To maximize the expressivity of the patterns module represent, each module employs a novel dynamic franmic feature attention mechanism to dynamically attends to important features in the agent's 's observation. When combined with spatio-temporal features, results suggest that the perceptual schemas FARM discovers capture diverse structures including object motions, 3D objects, and rpatiall-elationssiiss between objects. To have the modules coordinate what they attend to, they share information using transformer-style attention (Vaswani et al. 2017).
Recent work indicates that spatial attention is simple inductive bias for strong performance on object-centric vision tasks (Greff et al.) 2020 Locatello et al. 2020): Goyal et al. 2020b a). We compare FARM to recurrent architectures that employ spatial attention and contribute the following: 1. FARM, which combines dynamic feature attention with learning multiple recurrent modules
2

---

Under review as a conference paper at ICLR 2022
2. We show that FARM's components synergestically enable generalizing (a) recall to novel com- positions of object motions; (b) active perception of 3D objects to larger environments; and (c) generalizing memory of goal-information to longer tasks filled with more distractors {
3. We show that spatial attention-which reduces observation features weighted average over spatial positions-Ca can be detrimental to reinforcement learning of our diverse object-centric tasks and interfere with the benefits that come from learning multiple recurrent modules.
4. Our analysis of the representations learned by FARM provide evidence that it learns perceptual schemas that are flexibly distributed across combinations of recurrent modules 4.3.
2 RELATED WORK ON GENERALIZATION IN DEEP R
While large body of work has focused on studying generalization in deep RL. (Kansky et al. 2017 Witty et al. 2018 Farebrother et al.] 2018 8 Zhang al. 2018b; Cobbe et al. 2019 Raileanu & Fergus 2021). J. there has been less emphasis on studying generalization within diverse environments
and task structures. One research direction has focused on generalizing to longer tasks, e.g. executing longer sequences (Oh et al.] 2017 Zhang et al. [0118|] Lampinen et al. 2021) or executing novel subtask structures (Sohn et al. 201 8 2021 Brooks et al. 2021). Another direction has focused on generalizing to tasks with novel features, e.g., novel shape-color combinations (Higgins et al.
Chaploo al. 2018 Lee et al. 2020, Hill et al. 2020), novel backgrounds (Zhang et al. [202T] Agarwal et al. 2022). and novel distractors (Mott et al. 22199 Goyal et al., (020b). We attempt to bridge these prior strands of research by developing single architecture for (a) generalizing recall to
novel compositions of object motions; (b) generalizing sequential perception of 3D objects to larger environments; and (c) generalizing memory of goal information to longer task lengths.
Task-driven generalization. Recent work which has shown that a diverse training curriculum can promote generalization (Tobin et al. 2017: Packer et al., 2018 Hill et al. 2020). This research inspired our task-driven approach to discovering generalizable "schema-like" representations. Addi- tionally, our procedurally-generaae KeyBox task follows previous research on using procedural level generalization for faster learning and generalization (Justesen al. 2019 Jiang et al.) 2021).
Generalizing with feature attention. Most similar to our feature-attention mechanism are the attention mechanisms by Perez et al. (2018); Chaplot et al. (0188) In particularly, Chaplot et al. (2018) showed that mapping language instructions to non-linear feature coefficients enabled general- izing to tasks specified over unseen feature combinations in a 3D environment. While FARM also learns non-linear feature coefficients, our work has two important differences. First, we develop a multi-head version where individual feature coefficients This enables FARM to leverage this form of attention in settings where language instructions don't t indicate what attend to (this is true in 2/3 of our tasks). Second, we are the first to show that feature attention facilitates generalizing recall of object dynamics (Figure 4.1 and generalizing memory-retention to larger environments (Figure 4.3).
produced by their own recurrent modules.
Generalizing with top-down spatial attention. Most similar to FARM are the Attention Augmented Agent (AAA) (Mott et al., 2019) and Recurrent Independent Mechanisms (RIMs) (Goyal et al.. 2020b). Both are recurrent architectures that produce top-down spatial attention over observations. Additionally, RIMs also uses a modular state representation produced by a set of recurrent modules. AAA showed generalization to unseen distractors and RIMs to more distractors than trained on. While follow-up work on RIMs has addressed problems such as learning type-specific update rules (Goyal 2020a Dido et al.) 2021) and improved information sharing among modules Mittal et al. Goyal et al. 2021). there hasn't 't been an emphasis on which components enable generalization reinforcement learning across diverse domains. The major difference between AAA, RIMs, and FARM is that FARM attends to an observation with feature attention as opposed to spatial attention. Our experiments indicate that spatial attention can be an impediment to reinforcement leaming of tasks defined by shape-color agnostic object-dynamic and 3D objects.
et al.
2020
3 ARCHITECTURE FEATURE ATTENDING RECURRENT MODULES
We study agent that experiences a partial observation and task description (2t, T2) E X' x T, takes an action a E A. and experiences resultant reward r, E R. The agent needs a representation for state that enables learning policy æ S + A that maximizes the expected reward from taking an action
3

---

Under review as a conference paper at ICLR 2022
state E RS Retts a, -,.,,.: a]. Since the agent only observes observation petervationttss-aaoii reward tuples (Te, Te: ae. rt), needs to learn time-series representation mapping episode histories (z Tst, dct.Te)) to state representations St- The agent learns these representations with Feature Attending Recurrent Modules (FARM).
Agent's State Representation
Module y Computations
Composition of Module Spaces
St
90 x Hg
(m)
Perceptull schemo for gaing around obstacie
Cencel
Share info ((3
Phrcaptual acchems picking up goal key
Modulo
Moculen 37 (1)
r7(4)
rgddato
Legend
Contest
Recument connection
( att
Doeft.
Obsenvaion Enceder
Apulll moullo-eempore Fasturos
Set Coordineting Modules
2,
Previcus
spetial
action
posiltions
Go around obstacle
Pickup gool key
Figure 2: FARM learns representations for perceptual schemas that are distributed across nI recurrent modules { To discover perceptual schemas, FARM exploits structured observation features = Zz-1) E Rmxd, that share d, spatio-temporal features across spatial positions. To capture diverse structures, each module uses feature attention to compute coefficients for important features: fatt- This allows a module to update with important features present across all spatial positions. The modules share information to coordinate what they represent with Jahare.
Exploiting structured observation features for generalization. We assume that FARM generalize if it represents task-relevant regularly occurring perceptual structures, i.e. perceptual schemas. To discovering perceptual schemas, assume observation encoder <b(-) that produces observation features Z, E Z E Rmxd, where d, features are shared across rows representing different portions of the observation. For image, different rows might correspond to different spatial positions; for audio, to different frequency bands; or for robotic proprioception, to spatial information about different body parts. In this work, we focus on observations in the form of images.
Feature Attending Recurrent Modules. FARM learns representations for perceptual schemas that are distributed over n recurrent modules {)) Module i uses distinct initial modulc-state tuple (xi, Tr, 4ty rx) to R H For convenience, we can define a module context of as the concatenation of the previous module-state, the action chosen with it, the resultant reward, and the current task task-description c! [Te, R4 01-1; The agent integrates this context ct with Z, to compute s; using the collection of outputs from the modules:
( (()) and the subsequently experienced
learn the following mapping, (7 x T H00)
9; (n) (2)))22) (Ze, o ()) ..... (1) (Z,, (() (1)
(Z)
[h h [m)]
(2)
Module update rule. Each module will use its context c, to (a) attend to some aspect of the observation and (b) retrieve information from other modules. This manifests with two functions, (a) /( and (b) ((0) share: respectively, Defining
(()) Supdate as the module update- function,, we model ,,()
as the following composition of functions:
(Zr,
s(i) Supdate (c)
f(()) (Ze, (Ze, (2), p (2..
attend to observation
shure information
Modules capture diverse perceptual structures with dynamic feature attention. Our first insight was that modules can attend to high-level object-dynamics information if FARM learns recurrent observation encoder Z, = ó(x,, 1). By doing so, information about how feature-values shift between positions (i.e. feature dynamics) be captured in the features of Z,. Our second insight
4

---

Under review as a conference paper at ICLR 2022
was that feature attention is expressive attention mechanism for selecting observation information to update with. Each module predicts its own feature attention coefficients and applies them identically to all spatial positions in Z, (Perez al. 2018 Chaplot al. 2018). We found it useful to linearly project the features before and after using shared parameters as in Andreas et al. (2016); Hu et al. (2018). The operations summarized below:
face Ze. (0) = (Z,Wi O att ( ()))
where @ denotes an element-wise product over the feature dimension and a is a sigmoid non- linearity. Equation 4 computes the degree to which important features are expressed across all positions. Since our features capture dynamics information, this allows a module to attend to dynamics ($ 4.1. When updating with equation 3 we flatten the output of equation f and give this input an RNN. Flattening leads all spatial positions to be treated uniquely and enables module to represent aspects the observation that span multiple positions, such as 3D objects ($ 4.2) and spatial arrangements of objects ($ 4.3) Since the fcature-coefficients for the next time-step
produced with observation features from the current time-step, modules can dynamically shift their
attention when task-relevant events occur (see Figure
Modules share information to coordinate what they represent. Similar to RIMs Goyal et al. (2020b), before updating, each module retrieves information from other modules using transformer- style attention /Vaswani et al. 2017). We define the collection of previous module-states as H,-1 [ M ; ())xx,d,,, where a null-vector used to retrieve no information. Each
module retrieves information as follows:
(dOW?) (H, (H!+) T Vdh
P (c) .t ( Fi
softmax
H. W
(5)
4 EXPERIMENTS
In this section, we study the following questions: Can FARM enable generalization of memory-retention novel spatio-temporal compositions of object-dynamien
Can FARM generalize sequential active perception of 3D objects to larger environments?
3. Can FARM generalize memory-retention of ooal-information to longer tasks composed of longer sequences of observed obstacles?
Baselines. Our first baseline is the canonical choice for learning state-representations,, Long Short-term Memory (LSTM) (Hochreiter & Schmidhuber 1997). Our other two baselines the Attention Augmented Agent (AAA) (Mot)
Observation Atlention
Medular
Method
State
LSTM
X
et al. 2019) and Recurrent Independent Mechanisms (RIMs) (Goyal et al. 2020b) also employ top-down at-
AAA
Spatial
RIMs
Spetial
FARM (Ours)
Feature
tention over observation features. However, they both employ transformer-style attention (Locatello et al., 2020 Vaswani et al. (017) to dynamically attend to spatial posilions in the
Table 1: Comparison of baselines.
observation; whereas, we dynamically attend to features shared across all spatial positions. RIMs, like FARM, composes state with a set of recurrent modules.
Implementation details. We implement our recurrent observation encoder, Φ. as a ResNet (He et al. 2016) followed by Convolutional LSTM (ConvLSTMM) (Shi et al.) 2015). We implement the update function of each module with an LSTM. We used multihead-attentionn ( (Vaswani et al. 2017) for f(]] We trained the architecture with the IMPALA algorithm (Espeholt et al.) (0188) and an Adam optimizer (Kingma & Ba, 2015). We tune hyperparameters for all architectures with the "Place X next to Y" task from the Baby AI environment (Chevalier-Bae Chevaller-Boisvertt et al.) 2019) ($ C.2 We keep most hyperparameters fixed across our tasks. Our main change is to make the RNNs employed by each architecture larger for the KeyBox task. For details hyperparameters,
5

---

Under review as a conference paper at ICLR 2022
4.1 GENERALIZING MEMORY-RETENTE TO NOVEL COMPOSITIONS OF OBJECT-DYNAMICS
(a) Main Results
(b) Ablation
2 Sarrars [trin)
Derrarrs (trair)
Carearss Hhed?
ddnncers
Arrodin No Atannce Secureis
NEMuItI Recamett
130 110
Figure 3: FARM enables generalizing recall to novel spatio-temporal compositions of object motions. We present the success rate means and standard errors computed using 5 seeds. (a) Only FARM is able go above chance performance for each setting.
(b) Given recurrent features, we find that learning perceptual schemas with either multiple modules or feature attention enables learning recall of spatio-temporal compositions. These results suggest that spatial attention removes the benefits of learning multiple modules when learning to recall object-dynamics.
We study this with the "Baller" grid-world (Lampinen et al. 2021 shown in Figure 1 (a). The agent is a white square in the middle of the grid. There are m other "ballet-dancer" objects that move with distinct object-motion for 16 time-steps. There are 15 object-motions. The dances appear in sequence for 16 time-steps with a 48 48-time-step delay in between. Training always consists of seeing m {2, 4] dancers; testing always consists of seeing m {8} dancers. After all dancers finish, the agent is given a language instruction and it must go to the correct ballet-dancer. The agent gets reward of if it goes to the correct dancer, and 0 otherwise. All shapes and colors are randomized making the dynamics the only feature that indicates the goal-object. We study the success rate of an agent. A poorly performing agent will obtain chance performance, 1/m. All agents learn with a sample budget of 2 billion frames and use approximately 7M parameters.
We present the training and generalization success rates in Figure 3 We learned recurrent observation features with RIMs and AAA for a fair comparison. We find that only FARM is able to obtain above chance performance for training and testing. In order to understand the source of our performance, we ablate using recurrent observation encoder, using multiple modules, and using feature-attentinn confirm that recurrent encoder is required. Interestingly, we find that either using multiple modules or using our feature- enables task-learning, with more stable. We hypothesize that using a recurrent encoder with our attention allows for storing information about how object-motions are shifting position way that mitigates interference.
We
aaturen mechanism being slightly
Lampinen et al., 1 showed that a hierarchical transformer architecture was able to learn and generalize recall of sequential object-motions in the Ballet task. However, their architecture uses 6 distinct 8 8-head hierarchical attention operations and 8 MLP layers, making it much deeper and more computationally intensive than ours. Their architecture had 13.3M parameters whereas ours has 7.1M. We show that when using recurrent observation features, simply learning top-down feature attention or using multiple RNNs can enable learning this task.
4.2 GENERALIZING SEQUENTIAL ACTIVE PERCEPTION OF 3D OBJECTS TO LARGER ENVIRONMENTS
We study this with the 3D Unity environment from Hill et al. (0200) shown in Figure ! (b). This environment tests an agent's 's ability to generalize sequential composition of active perception of 3D objects to larger environments. The agent is an embodied avatar room filled with task objects and distractor objects. The agent has 46 discrete actions and has navigation step-size of .15m. The agent receives a language instruction of the form "X on Y" -e.g., "toothbrush on bed" and receives a reward of 1 if it completes the task and 0 otherwise. We partition objects into two sets as follows: pickup-able objects O = AU B and objects to place them on O₂ = C D. During training the agent sees A x D and B C' in a 4m x 4m room with 4 distractors, along with A x C and B x D in a 3m x 3m room with 0 distractors. We test the agent on A x C' and B x D in a 4m X 4m room with 4 distractors. The training curriculum also includes "Go X" and "Lift X". For details $

---

Under review as a conference paper at ICLR 2022
Gencraliaation Place
larger
Setting Ablation
Settting
Setting
Setting
Succes Eate
'1u
Pharmss (kdimms
169
Premees 21 Borns
Prames Diliien
1e9
Prommes (2.0cc)
305
Peccpnal Schames
LST
Agunt
attrnton
LSTN
Figure 4: FARM enables generalizing sequential active perception of 3D objects to a larger environment. We present the success rate means and standard errors computed using 3 seeds. (a) FARM generalizes best. (b) Our performance benefits come mainly from learning multiple modules, though feature attention slightly improves performance and lowers variance. These results suggest that spatial attention interferes with reinforcement learning of 3D objects.
We present the generalization success rate in Figure All architectures and ablations had approxi- mately 5M parameters. We find that baselines which used spatial attention learn more slowly than an LSTM or FARM. Additionally both models using spatial attention have poor performance until the end of training where AAA begins improve. FARM achieves relatively good performance, achieving a success rate of 60% and 80% on the two test settings, respectively. We hypothesize that our feature attention mechanism can more easily attend to 3D objects spanning multiple positions as opposed to spatial attention which uses softmax that may saturate on individual positions.
4.3 GENERALIZINE MEMORY-RETENTE
TO LONGER TASKS
In order to study this, we create the multi-level "keybox" environment depicted in Figure (c). This environment tests an agent's ability to learn remember ool-infoormation for increasingly long time- horizons. The agent maximizes reward by completing as many levels as it can in an episode. Each level in this environment a hallway with single box and a key of the same color. The hallway for level n consists of a length-n sequence of w X environment subsections. Each subsection contains distractor objects that can either be a ball of any color or keys of non-goal colors. The agent and the box always starts in the left-most subsection and the goal key always starts in the right-most subsection. Each time the agent succeeds, it is teleported to the next level and gets a reward of n/ n/Bnax where Thax is the maximum level the agent can complete. We set Tlaax = 10 during training. The agent has 50n time-steps to complete a level. The distractors pose an additional challenge of obstructing the agent's 's path. This further challenges memory since they have to be incorporated into the state-representaation to predict actions but should not overwrite the goal information. We study two generalization settings: densely populated setting that has subsections of width w = 3 with d = 2 distractors and sparsely populated setting that has subsections of width w = 5 with d = 4 distractors. In the dense setting, once the agent fails, it restarts in level n [1, Naone], where Aoone is the highest level the agent has completed. In the sparse setting, we found we needed to always restart the agent on level enable learning. All architectures used approximately 7.5M parameters.
We present the maximum training level reached and generalization success rates in Figure 5 find that all methods achieve comparable training results. In the dense setting, we see an LSTM quickly overfits in both settings. All architectures with attention continue to improve in generalization performance as they continue training. In the dense setting, we find that FARM tends to generalize better (by about 20% for AAA and about 30% for RIMs). In the sparse setting, both RIMs and an LSTM fail to generalize above 30%. FARMbetter generalizes than the AAA for level 20 but gets comparable performance for level 30. In some ways, this is our most surprising result since it is not obvious that uniformly attending to features across all spatial positions should help with this task. In the next section we study possible sources of our generalization performance.
We
4.3.1 1 ANALYSIS OF REPRESENTATIONS FOR REGULARLY OCCURRING EVENTS
We study the representations FARM learns for categories of regularly occurring perceptual events. We collect 2000 generalization episodes in level 20. We segment these episodes into categories of regularly occurring perceptual events: pickup ball, drop ball, pickup wrong key, drop wrong
7

---

Under review as a conference paper at ICLR 2022
(a) Deneely populated setting
(c) Ablation ueing Sparse setting
Taining (easler) Max Level=10, 10, Length 60
Eval: Laval-20
Evai: Lavel-30
Evali Lovol=20
Length s 50. Distractorsn40 40
Length & 1120, Detractorene0
Longth - 120. Ditracters-80
1.0
05
LSTN
0.2
haselle babonns
380
00
Foare Jn sikedd
(b) Sparsely populated setting
Foal: eovim4n
(ergth 240, Distactors-100 180
Twining (harder)
Cuel: Level=20
Evel: 30
-Atertor Seanchrre
Level-10, 10, Length
Length Distactors--00
Length- 130, Distractors-120
Pereecuad Schenoe
ictentor + Seaichara
arternie eagronded Agyer
Figure 5: FARM enable generalizing memory-retention of goal-information to longer horizons of obstacles than trained on. On the left-most panel of (a.b), we present the maximum training level reached by each agent. All others panels in (a-c) show the generalization success rate. For all quantities, we present the mean and standard eπтor computed using 10 seeds. (a) All architectures achieve comparable training performance. FARM better generalizes to longer hallways with more dis- tractors in the densely populated setting. (b) FARM and AAA get similar generalization performance in the sparsely populated setting. (c) Using multiple modules and feature attention both improve generalization. These results suggest that spatial attention interferes with generalization benefits of learning multiple modules while feature attention acts synergistically.
key, pickup correct key, and drop correct key. In order to understand how FARM represents these categories, we study time-series of the L2 norm of each module LSTM-state and of their attention coefficients. For reference, we also show the L2 norm for the entire episode.
We present a subset of results in Figure 6 For all results, see $D It is worth noting that program- matically extracting all task-relevant events for analysis challenging. For example, we were able to observe salient state activity from module 0 when the agent moved around obstacles but found no way to extract this signal. We show an example in GigureBaa. While some modules are selective for different recurring events such as attending to goal information (Figure a-c), it seems that the perceptual schema used for recurring events are generally distributed across multiple modules. We hypothesize that this enables FARM to leverage some modules to store goal-information and other modules to represent and navigate around obstacles in increasingly long environments. This is further supported by our ablation where we find that having 4 (8.2M params) or 8 (7.5M params) modules significantly outperforms using a single large module (7.5M params) (Figure (c)). Feature attention consistently improves performance.
DISCUSSION AND CONCLUSION
We have presented, FARM, novel recurrent state representation learning architecture. Our results show that a single architecture can support the following three disparate types of generalization: (a) generalizing recall to novel spatio-temporal compositions of object-motions; (b) generalizing active perception of 3D objects to larger environments; and (c) generalizing memory of goal- inal-mnfoomttion to longer sequences of obstacles. Our results suggest that learning feature attention and learning multiple recurrent modules can act synergistically, with different components being more helpful in different settings. For example, feature attention enables more stable learning of mojecctmmtions (Figure
8

---

Under review as a conference paper at ICLR 2022
Goal Box
Full mapp during generalization on level 20
Goal
Partial Observations
t=85
t=86
1=88
1=112
t=117
t=118
turn left
go forward
turn right
Observe goal
Pickup goal key
Take back goal
Module States
12 Norm of
Module
Active
200
Modulo
Module
Channel dimensions
Modile
1.5
Modulo
Module
Timestep
Temporal window
(a) L2 norm of produle-states as agent moves around obstacles
(b) Attention coefficients for module the agent picks goal key
(c) Average L2 norm of module-states
Average pair-ise correlation norm of module-states
Tutl ep aode
modile
modie
Commerinn webtt NAS clest
Mocider and
Moduless and
(ha))
pious kay
pabco cornect bey
Correiation
Anti-correlatoon
Figure 6: While (a-c) show that individual modules have salient activity for different events, (d) shows that different combinations of modules coordinate their activity for different events. This indicates that event representations are distributed flexibly across the modules.
(a) Module commonly exhibits salient activity when the agent moves around an obstacle. (b) Module shows selective activity for representing goal information. (c) Module 6 activates its attention coefficients as the agent picks up the goal key. (d) Some modules activate together (i.e. correlate) to represent some recurring events but activate in opposite directions (i.e. anti-correlate) to represent other recurring events. For example, modules and 6 correlate for picking up the correct key but anti-correlate for dropping the wrong object. This is similar to when neurons in word embeddings correlate for some words (c.g. man and king), but anti-correlate for other words (c.g. man and woman). These results suggest that FARM is distributing representations for perceptual schemas across the modules in complicated ways rather than representing individual schema with separate modules. Videos the state-activity and attention coefficients: https://bit.ly/3qxxaa
while multiple modules better enables generalization with 3D objects (Figure 4 Importantly, neither component hinders the other and an architecture with both succeeds on a broader set of tasks. Together, the enable the discovery of representations for shape-color agnostic motion ($4.1) for active perception of 3D objects ($4.2). and for spatial relationships between objects
4..3
Feature attention vs. spatial attention. While our dynamic feature attention mechanism flexible, FARM learns an update function that treats all positions in the attention output as unique (equation 4). major advantage spatial attention that the resultant representation is spatially-invarian. which may sometimes be more useful for policy. For example, we find that RIMs has better generalization an unseen number of distractors C.2 Interestingly, using spatial attention seems interfere with the generalization benefits of learning multiple recurrent modules. An interesting future direction might be to explore combining both forms of attention for more flexible policy.
RNNs vs. Transformers. Another interesting future direction involves discovering perceptual schemas with transformers (Vaswani et al., 2017). When used as memory for an RL agent, Transform- ers have shown strong performance for representing state over longer sequences but show weaker performance for reactive behavior (Parisotto et al. 2020). An interesting next-step might be to use a hierarchical architecture where FARM provides representations for relatively short sequences. These could be used as primitive ceteevaatooorppeeeetttttt for a transformer architecture. Such an architecture might enable discovering perceptual schemas very long tmee-hoizzons.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. ICLR, 2021.
Jacob Andreas, Marcus Rohrbach, Trevor Darrell, and Dan Klein. Neural module networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 39-48, 2016.
James Bradbury. Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs. 2018. URL
http: /tte://gittti /gggttttthu.gggggggg/ggggg.ghh wtt.com///ooo/ www./ooogoe///ooooooooooooooooeo jax
Ethan A Brooks, Janarthanan Rajendran, Richard Lewis, and Satinder Singh. Reinforcement learning of implicit and explicit control flow in instructions. ICML, 2021.
Devendra Singh Chaplot, Kanthashree Mysore Sathyendra, Rama Kumar Pasumarthi, Dheeraj Rajagopal, and Ruslan Salakhutdinov. Gated-attention architectures for task-oriented language grounding. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
Maxime Chevalier- Boiiveer,,, 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. In International Conference on Learning Representations, 2019. URL https:://ppenreiii
net. wowrmmidd:reeeoooo
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. ICML, pp. 1282-1289. PMLR, 2019.
Aniket Didolkar, Anirudh Goyal, Rosemary Ke, Charles Blundell, Philippe Beaudoin, Nicolas Manfred Otto Heess, Michael Mozer, and Yoshua Bengio. Neural production systems. In NeurlPS,
2021.
Lasse Espcholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In ICML, pp. 1407-1416. PMLR, 2018.
Jesse Farebrother, Marlos C Machado, and Michael Bowling. Generalization and regularization in dqn. arXiv, 2018.
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, 2020a.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. ICLR, 2020b.
Anirudh Goyal, Aniket Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Mozer, and Yoshua Bengio. Coordination among neural modules through a shared global workspace. arXiv preprint arXiv:2103.0/ /97, 2021.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv, 2020.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. CVPR, pp. 770-778, 2016.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In ICML, pp. 1480-1490. PMLR, 2017.
Felix Hill. Andrew Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James L McClelland, and Adam Santoro. Environmental drivers systematicity and generalization in situated agent. ICLR, 2020.
10

---

Under review as a conference paper at ICLR 2022
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735- 1780, 1997.
Jie Hu, Li Shen, and Gang Sun. gqueeze-and-exciiee networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 7132-7141,, 2018.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint orXiv1611.00399 2016.
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In ICML, pp. 4940 4950. PMLR, 2021.
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illuminating generalization in deep reinforcement learning through procedural level generation. AAAI, 2019.
Ken Kansky, Tom Silver, David A Mély, Mohamed Eldawy, Miguel Lázaro-Gredilla Xinghua Lou, Nimrod Dorfman, Szymon Sidor, Scott Phoenix, and Dileep George. Schema networks: Zero-shot transfer with a generative causal model of intuitive physics. ICML, pp. 1809-1818. PMLR, 2017.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Andrew Kyle Lampinen, Stephanie CY Chan, Andrea Banino, and Felix Hill. Towards mental time travel: hierarchical memory for reinforcement learning agents. arXiv, 2021.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In ICLR, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. arXiv preprint arXiv:2006. 5505,, 2020.
Marvin Minsky, A framework for representing knowledge. De Gruyter, 1979.
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning to combine top-down and bottom-up signals in recurrent neural networks with attention over modules. International Conference on Machine Learning, pp. 6972-6986. PMLR, 2020.
Alex Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo J Rezende. Towards interpretable reinforcement learning using attention augmented agents. NeurlPS, 2019.
Junhyuk Oh, Satinder Singh, Honglak Lee, and Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. ICML, abell700.000000 706.05064, 2017.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization deep reinforcement learning. arXiv. 2018.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In ICML,, pp- 7487-7498. PMLR, 2020.
Ethan Perez, Florian Strub, Harm Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI Conference Artificial Intelligence, 2018.
Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement learning. ICLR, 2021.
David E Rumelhart. Schemata: The building blocks of cognition. Routledge, 1980.
11

---

Under review as a conference paper at ICLR 2022
David E Rumelhart, Paul Smolensky, James L McClelland, and G Hinton. Sequential thought processes in pdp models. Parallel distributed processing: explorations in the microstructures of cognition, 2:3-57, 1986.
Xingjian Shi, Zhourong Chen, Hao Wang. Dit- Yan Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional Istm network: A machine learning approach for precipitation nowcasting. Advances in neural information processing systems, 28, 2015.
Sungryull Sohn, Junhyuk Oh, and Honglak Lee. Hierarchical reinforcement learning for zero-shot generalization with subtask dependencies. NeurIPS, 2018.
Sungryull Sohn, Hyunjae Woo, Jongwook Choi, and Honglak Lee. Meta reinforcement learning with autonomous inference of subtask dependencies. ICLR, 2021.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation the real world. In IROS, pp. 23-30. IEEE, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. arXiv, 2017.
Sam Witty, Jun Ki Lee, Emma Tosch, Akanksha Atrey, Michael Littman, and David Jensen. Measur- ing and characterizing generalization in deep reinforcement learning. arXiv,, 2018.
Amy Zhang, Sainbayar Sukhbaatar, Adam Lerer, Arthur Szlam, and Rob Fergus. Composable planning with attributes. In International Conference on Machine Learning, pp. 5842-5851.
PMLR, 2018a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. ICLR, 2021.
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study overfitting in deep reinforcement learning. arXiv, 2018b.
12


Reviews:
[{"review_id": "s-2A3pDP_DM", "reviewer": "Reviewer_jU4K", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "The focus of this paper is on addressing three particular types of generalization in RL as listed above. A commonality between them is that test tasks are novel compositions of regularly occurring \"structures\" the agent experiences during training. In the first case, these are the object-motions, and in the second and third case the visual objects. The difference in generalization behavior in the second and third cases is less clear to me since both can be viewed as requiring generalization to a larger environment and novel object compositions.\nRegardless, it is interesting to consider a representation learning architecture that benefits each of these settings.\n\nThe main hypothesis, thus, is that a modular representation is desirable since individual modules can learn to capture these separate structures. This is achieved in CPS by having multiple different RNNs process the input separately and exchange information with one another to form a state representation, similar to how this is done in RIMs. Unlike in RIMs, however, here the input attention is over the features produced by an observation encoder, and attention is to all positions, while RIMs attend to particular spatial positions.\nThe motivation for this is not quite clear to me, other than that it is now possible to capture information about an object that covers multiple spatial positions. A potential downside is that it is no longer possible to attend to a particular location.\n\nThe technical contribution of CPS and technical novelty compared to prior work is therefore highly limited, especially since similar forms to the multiplicative input attention across positions explored here have also been explored previously, eg. in FiLM (as is also correctly pointed out in the related work section). Regarding the motivation of CPS based on \"perceptual schemas\", the connection is not quite clear to me either. In essence, the modularity in CPS and the reason why this may be beneficial for generalization is identical to how this is motivated in prior work such as RIMs. Therefore, it is not clear why modules or RIMs are now rebranded as perceptual subschema's, and one should view the learned representation as a perceptual schema. What exactly is new here that warrants a connection to \"schemas\" and a rebranding of existing terminology from previous papers?\n\nThe empirical contribution of CPS does appear substantial since in all the considered domains CPS is shown to outperform the baselines, which include RIMs. However, there are several experimental results regarding CPS and baselines that I would like to see clarified:\n\n* On the Ballet task it is found that either feature attention or modules are sufficient for generalization. At the same time, RIMs and AAA are shown to perform very poorly. The only difference (minus minor architectural differences) of CPS with no attention to RIMs is the lack of a recurrent encoder in case of RIMs. Is the hypothesis thus that the lack of recurrent features is mainly responsible for these baselines performing poorly? If this is the case then it seems that this task is quite adversarial to allow for a fair comparison to RIMs. Would it be possible to endow RIMs with similar capabilities (eg. a recurrent input encoder) as a more representative baseline? Similarly, although AAA does have access to the recurrent input features it also performs poorly compared to the LSTM. What is the hypothesis for this behavior? Can it be tested?\n\n* Similarly on the 3D Unity environment, the supposedly stronger baselines of AAA and RIMs struggle greatly compared to CPS and the LSTM. In this case it is not clear that recurrent input features are needed, and so the previous comparison to CPS with 1 slot + attention and to CPS with multiple slots and no attention would be particularly revealing here. I would like to see this comparison here or some other experiment to better explain the performance gap to CPS and AAA. \n\nThe multi-level keybox experiments are more revealing since it is now demonstrated how there is a clear benefit to having attention and modules, which was not the case for Ballet and couldn't be observed on 3D Unity. At the same time, it can now also be observed how AAA performs quite well and the performance gap to this baseline is small. In general, although there is often a substantial gap to the baselines it is not quite clear to me yet which components are contributing the most, and therefore how significant the empirical contribution in reality is. The analysis of representations for regularly occurring events is interesting and indeed suggests some specialization among the modules as was also previously found in RIMs.\n\nFinally, regarding the overall clarity and presentation of the work, I find that it could be improved substantially. In particular, the choice of different types of generalization seems rather arbitrary and it is unclear why these particular types are significant. Regarding the second and third generalization regime it is also not quite clear what the dimensions are that are being composed over. To me it seems that they should be the objects, yet the text talks about composable segments as well. The correspondence of perceptual schemas to parts of the learned representation is also a bit confusing. At first, it is suggested that perceptual schemas take on things like car motions, sedans, etc. yet these are the structures that are captured by the modules, which are referred to as subschemas. My question thus is, which are the schemas, the individual modules or their composition? Regarding the method details, I would strongly encourage the authors to include some pseudo-code for this algorithm, and better highlight what are the novel contributions. It also isn't clear how the policy network is structured to leverage the underlying modularity of the representation, which seems like an important choice to discuss. The text below eq. 4 about the potential capabilities of using attention in this way in combination with eq 3. was not clear to me, and it would be good to provide concrete examples of how the things outlined there could be achieved. \n\nSome minor comments:\n\n* There is a mention of f_message, which is mentioned nowhere else. I assume this is an artifact of a previous version.\n\n* I encountered several typo's, which a spell checker should reveal for you as well.\n", "summary_of_the_paper": "This paper proposes Composable Perceptual Schemas (CPS), a modular state representation learning architecture for reinforcement learning that combines modular RNNs as in RIMs (Goyal et al., 2020) with a dynamic feature attention mechanism. It is hypothesized that this lets CPS exhibit multiple different kinds of generalization encountered previously in the literature: (1) generalization to environments given by novel compositions of known object motions, (2) generalizing active perception of 3D objects to larger environments, and (3) generalizing goal-directed behavior and memory retention to larger environments.\n\nThe idea of CPS is as follows:\n\n* A recurrent observation encoder produces observation features for different parses of the input observation (here spatial locations in the image) to yield a feature matrix. Here each row corresponds to a different parse.\n\n* A number of recurrent modules (here called \"subschemas\"), each process the observation features + context (action, reward, task encoding) and the previous state of all other modules using dynamic attention. From this, an updated state is computed using an LSTM.\n\n* The dynamic attention for observations + context proceeds by scaling the observation features at all positions with a vector of learned attention coefficient. In this way, feature selection can take place (but shared across positions).\n\n* The dynamic attention for other modules is similar to the implementation in RIMs using transformer-style attention.\n\nCPS is evaluated on three environments and compared to several baselines. In general, it can be seen to exhibit the desired generalization behavior and outperform the baselines. Finally, some evidence is presented that subschemas are selective for semantically meaningful aspects of the input (eg. goal information, distractor information, etc.).", "summary_of_the_review": "This paper presents a method for learning modular representations for RL that allow for generalization in three regimes. The technical novelty is highly limited, but the empirical contribution is promising. At the same time, I would like to see a number of the presented results clarified and some additional analysis done to understand why the baselines perform so poorly and which parts of the proposed method contribute the most in different regimes.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "nFbV9eFkyg-", "reviewer": "Reviewer_4R1W", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Strengths:\n\nInteresting set of diagnostic tasks/environments have been developed to study generalization of deepRL agents along various axes such as object types, object numbers, sequence length of task solution etc. \n\nExperiments showing qualitative analysis of the subschema representations learned by the CPS model are interesting. More such qualitative studies/ablations or visualizations (various attention coefficients during an episode, subschema activity during an episode? etc.), could further help dissect and gain insight into the various components of the proposed CPS model and their roles in the generalization performance. \n\n\nWeaknesses:\n\nProposed model and architectural novelty is quite similar in nature to several of the follow-up works on RIMs [1, 2, 3, 4]. Discussion and comparison with highly relevant follow-up work on RIMs [1, 2, 3, 4] which also propose solutions to the same problems of i) how to share information among modules [1, 3] ii) learning type-specific update rules [2, 4] has not been addressed or evaluated. How does the proposed novel feature attention mechanism and the general CPS model architecture benefit over other solutions proposed by these works [1, 2, 3, 4] to the similar problems? \n\nAnother issue is with the lack of relevant baseline models being used to judge the benefits of the proposed CPS model architecture and its novel feature attention mechanism. The RIMs baseline model struggles to learn on several tasks (Figure 3 (a) (b), Figure 4), sometimes shows even worse performance than the vanilla LSTM. RMC [5] which models relational information among memory slots and uses this information to update the memory slots using key-value attention would make for a good and (relatively) simple baseline model to compare against.  \n\n\nResults:\n\nIt’s rather interesting that the vanilla LSTM is able to generalize better than RIMs on a few tasks such as Dancers (Fig. 3.a) and 3D Place X on Y task (Fig. 4). This is surprising since RIMs have inductive biases that explicitly favour modularity and compositionality in representation over the vanilla LSTM, which are beneficial to generalize on these tasks. Could the authors offer \nsome explanation for this behavior?\n\nAre the total number of learnable parameters always kept the same across CPS model variants with 1/2/4/8 subschemas in all the experiments? \n\nThe paper would improve if the authors could add/move some of the visualizations from the Appendix to the section 4.3.1 to expand this discussion and qualitative analysis on the extent of specialization and modularity in representations/behavior captured by various subschemas. This would further strengthen the author’s hypothesis of “learning n subschemas, each with 1/n parameters of a single monolithic module….. encourages the discovery of regularly occurring structures.” \n\nI didn’t quite understand how to interpret the visualizations on the pairwise correlation between subschemas in Fig. 6.b. While they show that subschemas interact with one another during an event, shouldn’t this interaction be rather sparse? If the subschemas are to specialize to unique recurring input patterns in an episode, shouldn’t just one subschema ideally activate only for the recurring events (ex: pickup key, drop key etc.) it has specialized to represent? If so, then how can there be a high correlation between subschema activations over an event? Does that indicate that multiple subschemas are modelling the same recurring pattern/event amongst them? Could the authors provide some clarification on this?\n\nWriting/Presentation: \n\nHow are the attention coefficients (alpha^i) in the f_att function computed? \n\n$f_{att}^{(i)} (x_t, e_{t-1}^{(i))} = (Z_t … )$. What is $e_{t-1}^{(i)}$? Do you mean $c_{t-1}^{(i)}$?\n\n“We used multihead-attention for $f_{message}$”. What is $f_{message}$? It was not defined earlier.\n\nPlease include highly relevant information about certain key agent/training/environment hyperparameters in the main text that is needed to judge the validity of empirical results/ablations. The reading experience can greatly improve if one doesn’t have to constantly jump from main text to Appendix to gather information needed to judge the experimental settings and validity of results.  \n\n“... note that “Place A on C” and “Place B on D” are both in our training curriculum …” wording is confusing. Does this mean that the agent has seen the same task structure in training but not with the same objects to pick and place? \n\nMinor typos: \n\n“Reoccuring” -> “recurring”?\n\n“Pickupable” -> “pickable”\n\n“Policy-state” -> “state”\n\n\n\n[1] Mittal et. al, “Learning to Combine Top-Down and Bottom-Up Signals in Recurrent Neural Networks with Attention over Modules”, ICML 2020.\n\n[2] Goyal et. al, “Factorizing Declarative and Procedural Knowledge in Structured, Dynamical Environments”, ICLR 2021.\n\n[3] Goyal et. al, “Coordination Among Neural Modules Through a Shared Global Workspace”,  \n\n[4] Goyal et. al, “Neural Production Systems”, https://arxiv.org/abs/2103.01937\n\n[5] Santoro, Adam, et al. \"Relational recurrent neural networks.\", NeurIPS 2018.\n", "summary_of_the_paper": "The paper proposes a modular state representation learning architecture called Composable Perceptual Schemas to discover regularly recurring patterns in its observations into “schemas” and process them in a modular and factorized manner. The proposed model is tested for generalization in the deepRL setting along various axes such as varying object types, object numbers, distractors, task solution lengths etc. and shows promising improvements over other recurrent baseline models.\n", "summary_of_the_review": "The paper develops an interesting suite of diagnostic tasks/environments to study generalization in deepRL agents. While the technical contributions and results are interesting, it is difficult to fully judge the merits of the proposed model and the main contribution of the feature attention mechanism if it is not referenced, discussed or evaluated in the context of related prior works [1, 2, 3, 4] that study the similar problems and propose similar solution methods. Further, the paper would greatly improve with more qualitative analysis into the CPS model’s components and their role in its generalization performance. For these reasons I’m giving a score of 5. I’m willing to increase the score if my concerns are addressed.     \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lDR6JDpj8pO", "reviewer": "Reviewer_Y5Kp", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "It is nice that this paper is motivated by representing schemas from cognitive science, and to my knowledge this paper provides a novel architecture attempting to do so. From what I understand, the architecture aims for each module to specialize to represent a particular spatial configuration of a scene, and different modules might represent different spatial configurations.\n\nI have the following concerns:\n\n1. It was difficult for me to take away from this paper a set of general principles that can be applied to future research. I would expect that the paper either (1) goes the scientific route: takes the best theory of schemas from cognitive science and show how their architecture is entailed by that theory, in which case the experiments would help us refute alternative theories of schemas proposed in cognitive science and give us insight about the nature of schemas from the abstract perspective, or (2) goes the engineering route: proposes an architecture that gives much better performance on the evaluation tasks, in which case the experiments would help us understand what engineering principles are useful to incorporate into future methods. This paper does not do (1), so I would expect it to do (2). However, the experimental results suggest that CPS is often similar to the LSTM baseline (section 4.1, 4.2), and gets the same performance as the AAA baseline (section 4.3). Furthermore, for the Ballet task, the paper does not compare against the method (HCAM) in the paper that proposed the Ballet task (https://arxiv.org/pdf/2105.14039.pdf), so it is unclear whether CPS offers insight on the best way to solve such a task.\n2. It was not sure what to conclude based on the analysis in section 4.3.1 and Appendix C. For example, could it be possible that the plots of Figure 6b, Figure 8, and Figure 9 could have just as well have been generated by just looking at the random initial weights of the schemas? Given that the authors are motivated by modeling schemas that model independent aspects of the task, what measurement could the authors make to evaluate whether the schemas have specialized to learning something meaningful about different aspects of the task, in a way that would be distinguishable from the measurement take on randomly initialized schema modules?\n3. I am unsure what to conclude from Figure 3c. It seems that having the ablations of having 4 subschema modules without attention and having only 1 subschema module with attention performs similarly to the full CPS. If this is the case, how would I evaluate the CPS as a contribution, if ablations that remove the core parts of the method (>1 schema, or feature attention) still perform similarly well? If the authors hypothesize that \"We hypothesize that learning n subschemas, each with 1/n parameters of a single monolithic module, reduces overfitting and encourages the discovery of regularly occurring structures,\" does Figure 3c actually refute their hypothesis?\n4. In various occasions the paper makes claims about the abilities of the CPS module that are not sufficiently tested. For example, the authors state in 4.1 that \" Learning this task tests a recurrent architecture’s ability to recognize and maintain separate,\nindependent dynamics in an agent’s state representation\" --> can the authors please conduct the experiment that tests whether the schema modules actually learn independent dynamics with an appropriate metric of independence? As another example, in section 3: \" Since the feature-coefficients for the next time-step are produced with observation features from the current time-step, subschemas can dynamically shift their attention when task-relevant events occur\" --> can the authors please conduct the experiment that evaluates whether this shift in attention is actually true? If there is a way to visualize how each schema module's attention to different features changes over time, and match that attention activation to semantically meaningful events in the episode, that would give much needed insight on how the method works. Another example is, in section 3: \"To encourage specialization, each subschema has $\\approx \\frac{1}{n}$ the parameters a single module would have.\" --> can the authors please conduct an experiment that measures the degree of specialization with an appropriate evaluation metric, for various different numbers of parameters of the subschema? It is unclear to me that simply restricting the parameter count is enough to induce specialization.", "summary_of_the_paper": "This paper proposes the Composable Perceptual Schemas recurrent architecture. This recurrent architecture is decomposed into $n$ LSTMs, each of which the authors call a _schema module_. At each time-step:\n1.  a recurrent encoder produces a set of $m$ feature vectors of dimension $d_z$, one vector for a different spatial location\n2.  the task information $\\tau_t$, the previous reward $r_{t-1}$, and previous action $a_{t-1}$ gets concatenate with the previous hidden state $h_{t-1}^{(i)}$ of module $i$. The authors call this concatenation the _context_ $c_t^{(i)}$.\n3. The update rule of each schema module produces the new hidden state $h_{t}^{(i)}$ from three inputs:\n    1. The output of $f_{att}$: A single soft mask of length $d_z$ is computed from $c_t^{(i)}$ and is applied to each of the $m$ feature vectors.\n    2. $c_t^{(i)}$\n    3. The output of a multi-head attention layer, where the queries are the $c_t^{(i)}$, and the keys and values are the previous hidden states of the schemas, along with the null vector: $[h_{t-1}^{(1)}, ... h_{t-1}^{(n)}, 0]$\n\nThe main difference between CPS and the baseline recurrent architectures (AAA and RIMs) is in step 1 of the update, how $f_{att}$ is computed: instead of attending over the different spatial locations CPS attends over the same dimensions of all the spatial features.\n\nThe authors apply their architecture to three generalization scenarios:\n1. \" Can CPS enable generalization of memory-retention to novel spatial and temporal (spatiotemporal) compositions of object-dynamics?\" --> tested on the Ballet task from https://arxiv.org/pdf/2105.14039.pdf. CPS performs slightly better than the LSTM when generalizing to different spatial compositions, and much better than all baselines when generalizing to different spatial and temporal compositions.\n2. \" Can CPS generalize sequential active perception of 3D objects to larger environments?\" --> tested on the \"Place X on Y\" task from https://arxiv.org/pdf/1910.00571.pdf. CPS learns faster than the baselines.\n3. \"Can CPS generalize goal-oriented behavior and memory-retention to environments composed of longer sequences of observed object-configurations?\" --> tested on a custom KeyBox environment written in the BabyAI framework. CPS performs similarly to the AAA baseline.\n\n\n\n\n", "summary_of_the_review": "Overall, while this paper is tackling an important question, its execution could be substantially improved. It was not clear what the contribution of the paper is from a scientific or engineering perspective. The analysis could be more informative about what the modules are actually learning. The central claim of the necessity of feature attention and multiple schema modules seems to be refuted from Figure 3c. The paper needs to either temper its claims or provide empirical evaluation for its claims. More details about these concerns are in the Main Review. It could be that I may have misunderstood something, in which case I would happily engage in discussion with the authors. \n\n== AFTER REBUTTAL ==\nI thank the authors for their response and additional experiments. I have raised my score to a 5, and I especially appreciated the clarifications and the qualitative analysis. After reading the rebuttal and the revised paper, one concern I still have is that it is difficult to glean a single coherent empirical message from the paper: on some tasks FARM performs the baselines (Figure 3, 4), on other tasks FARM performs similarly (e.g. Figure 5b, 7, 8). Perhaps this could be something that could be fixed through the writing, but overall the case for the following questions could be made stronger: (1) Why these three types of generalization rather than some other types of generalization? What is the underlying property that all these three types of generalization have in common, and why is this property important for AI research? (2) What is the explanation for why the proposed method must be the method for tackling tasks with this underlying property?\n\nRegarding (1), it would strengthen the paper's overall message to identify a single unifying characteristic of the problems considered, because at the moment the paper is written in a way that highlights the disparate nature of these generalization tasks, but without an underlying theme connecting these generalization tasks, they seem to be chosen in an ad hoc way. Regarding (2), this could be addressed with stronger empirical results.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper develops a mechanism for learning modular state representations in RL that organize recurring patterns into composable schemas. The approach combines modular RNNs as in RIMs (Goyal et al., 2020) with a dynamic feature attention mechanism. There were a variety of concerns in the initial reviews that were addressed by the authors through a set of clarifications and improved empirical analysis, substantially improving the paper. However, there still remain some issues in clarity of presentation and inconsistent empirical results, especially in the form of clear take-aways from the empirical analysis and broader insights from the paper, as detailed in the individual reviews. The authors are encouraged to take these aspects into consideration in revising their manuscript.
Decision: Reject