Task-driven Discovery of Perceptual Schemas for Generalization in Reinforcement Learning

TASK-DRIVEN DISCOVERY OF PERCEPTUAL SCHEMAS FOR GENERALIZATION IN REINFORCEMENT LEARNING Anonymous authors Paper under double-blind review A B S T R A C T 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 remembering goal-information over sequences of regularly occurring object-configurations . These diverse task structures all share an 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 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 reccurring segments of an agent ’ s experience , “ perceptual schemas ” . We propose Feature Attending Recurrent Modules ( FARM ) , which learns a 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 to generalize across the diverse object-centric domains we study . 1 I N T R O D U C T I O N Cognitive scientists theorize that humans generalize broadly with “ schemas ” they discover for regularly occurring structures within their experience ( Minsky , 1979 ; Rumelhart , 1980 ; Rumelhart et 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 car motions ( such as accelerating or stopping ) , and common car arrangements ( such as a row of cars ) . Importantly , schemas are composable representations over portions of our observations . This allows us to 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 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 to one type of generalization . For example , one algorithm will generalize to novel compositions of familiar shapes and colors ( Higgins et al. , 2017 ; Chaplot et al. , 2018 ) , whereas another algorithm will generalize to longer sequences of observed subtasks ( Sohn et al. , 2018 ; 2021 ; Brooks et al. , 2021 ) . In this work , we hypothesize that we can develop a single deep RL architecture that can exhibit multiple types of generalization if it can learn schema-like representations for regularly occurring structures within its experience . As a 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 1 ) . Across these environments , test tasks are novel compositions 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 1 , a ) ; generalization for the “ Place X on Y task ” involves generalizing active perception of and navigation towards regularly appearing 3D objects ( Figure 1 , b ) ; and generalization for the “ Keybox task ” involves generalizing memory-retention of goal-information to larger environments composed of sequences of regularly occurring object-configurations ( Figure 1 , 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 ) , a state representation learning architecture 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 information than a 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 a module can represent , each module employs a novel dynamic feature attention mechanism to dynamically attends to important features in the agent ’ s observation . When combined with spatio-temporal features , our results suggest that the perceptual schemas FARM discovers capture diverse structures including object motions , 3D objects , and spatial-relationships 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 a 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 ( §3 ) . 2 . We show that FARM ’ s components synergestically enable generalizing ( a ) recall to novel compositions 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 ( §4 ) . 3 . We show that spatial attention—which reduces observation features to a weighted average over spatial positions—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.1 ) . 2 R E L AT E D W O R K O N G E N E R A L I Z AT I O N I N D E E P R L While a large body of work has focused on studying generalization in deep RL ( Kansky et al. , 2017 ; Witty et al. , 2018 ; Farebrother et al. , 2018 ; Zhang et al. , 2018b ; Cobbe et al. , 2019 ; Raileanu & Fergus , 2021 ) , 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. , 2018a ; Lampinen et al. , 2021 ) or executing novel subtask structures ( Sohn et al. , 2018 ; 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. , 2017 ; Chaplot et al. , 2018 ; Lee et al. , 2020 ; Hill et al. , 2020 ) , novel backgrounds ( Zhang et al. , 2021 ; Agarwal et al. , 2021 ) , and novel distractors ( Mott et al. , 2019 ; Goyal et al. , 2020b ) . We attempt to bridge these prior strands of research by developing a 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 . Additionally , our procedurally-generated KeyBox task follows previous research on using procedural level generalization for faster learning and generalization ( Justesen et 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 . ( 2018 ) . In particularly , Chaplot et al . ( 2018 ) showed that mapping language instructions to non-linear feature coefficients enabled generalizing 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 are produced by their own recurrent modules . This enables FARM to leverage this form of attention in settings where language instructions don ’ t indicate what to 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 ) . 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 et al. , 2020a ; Didolkar et al. , 2021 ) and improved information sharing among modules ( Mittal et al. , 2020 ; Goyal et al. , 2021 ) , there hasn ’ t been an emphasis on which components enable generalization in 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 learning of tasks defined by shape-color agnostic object-dynamics and 3D objects . 3 A R C H I T E C T U R E : F E AT U R E AT T E N D I N G R E C U R R E N T M O D U L E S We study an agent that experiences a partial observation and task description ( xt , τt ) ∈ X × T , takes an action at ∈ A , and experiences resultant reward rt ∈ R. The agent needs a representation for state that enables learning a policy π : S → A that maximizes the expected reward from taking an action at a state Eπ [ ∑ t=0 γ tRt+1|St = a , At = a ] . Since the agent only observes observation-task-actionreward tuples ( xt , τt , at , rt ) , it needs to learn a time-series representation mapping episode histories ( x≤t , τ≤t , a < t , r < t ) to state representations st . The agent learns these representations with Feature Attending Recurrent Modules ( FARM ) . Exploiting structured observation features for generalization . We assume that FARM can generalize if it represents task-relevant regularly occurring perceptual structures , i.e . perceptual schemas . To discovering perceptual schemas , we assume an observation encoder φ ( · ) that produces observation features Zt ∈ Z ∈ Rm×dz where dz features are shared across m rows representing different portions of the observation . For an 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 { η ( i ) } . Module i uses a distinct initial module-state h ( i ) 0 ∈ H ( i ) and the subsequently experienced observation-task-action-reward tuple ( xt , τt , at , rt ) to learn the following mapping , η ( i ) : Z ×T ×H ( i ) ×A×R → H ( i ) . For convenience , we can define a module context c ( i ) t as the concatenation of the previous module-state , the action chosen with it , the resultant reward , and the current task-description as c ( i ) t = [ τt , h ( i ) t−1 , at−1 , rt−1 ] . The agent integrates this context c ( i ) t with Zt to compute st using the collection of outputs from the modules : st = [ η ( 1 ) ( Zt , c ( 1 ) t ) , . . . , η ( n ) ( Zt , c ( n ) t ) ] ( 1 ) = [ h ( 1 ) t , . . . , h ( n ) t ] . ( 2 ) Module update rule . Each module will use its context c ( i ) t to ( a ) attend to some aspect of the observation and ( b ) to retrieve information from other modules . This manifests with two functions , ( a ) f ( i ) att and ( b ) f ( i ) share , respectively . Defining f ( i ) update as the module update-function , we model η ( i ) as the following composition of functions : η ( i ) ( Zt , c ( i ) t ) = f ( i ) update ( c ( i ) t , f ( i ) att ( Zt , c ( i ) t ) , ︸ ︷︷ ︸ attend to observation f ( i ) share ( c ( i ) t , { h ( j ) t−1 } n1 ) ︸ ︷︷ ︸ share information ) ( 3 ) 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 a recurrent observation encoder Zt = φ ( xt , Zt−1 ) . By doing so , information about how feature-values shift between positions ( i.e . feature dynamics ) can be captured in the features of Zt . Our second insight was that feature attention is an 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 Zt ( Perez et al. , 2018 ; Chaplot et 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 are summarized below : f ( i ) att ( Zt , c ( i ) t ) = ( ZtW1 σ ( W iattc ( i ) t ) ) W2 ( 4 ) where denotes an element-wise product over the feature dimension and σ is a sigmoid nonlinearity . 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 4 and give this as input to an RNN . Flattening leads all spatial positions to be treated uniquely and enables a module to represent aspects of the observation that span multiple positions , such as 3D objects ( § 4.2 ) and spatial arrangements of objects ( § 4.3 ) . Since the feature-coefficients for the next time-step are produced with observation features from the current time-step , modules can dynamically shift their attention when task-relevant events occur ( see Figure 6 , c ) . 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 transformerstyle attention ( Vaswani et al. , 2017 ) . We define the collection of previous module-states as Ht−1 = [ h ( 1 ) t−1 ; . . . ; h ( n ) t−1 ; 0 ] ∈ R ( n+1 ) ×dh , where 0 is a null-vector used to retrieve no information . Each module retrieves information as follows : f ( i ) share ( c ( i ) t , { h ( j ) t−1 } n1 ) = softmax ( c ( i ) t W q i ) ( Ht−1W k i ) > √ dh Ht−1W vi ( 5 ) 4 E X P E R I M E N T S In this section , we study the following questions : 1 . Can FARM enable generalization of memory-retention to novel spatio-temporal compositions of object-dynamics ? 2 . Can FARM generalize sequential active perception of 3D objects to larger environments ? 3 . Can FARM generalize memory-retention of goal-information to longer tasks composed of longer sequences of observed obstacles ? Baselines . Our first baseline is the canonical choice for learning state-representations , a Long Short-term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) . Our other two baselines — the Attention Augmented Agent ( AAA ) ( Mott et al. , 2019 ) and Recurrent Independent Mechanisms ( RIMs ) ( Goyal et al. , 2020b ) — also employ top-down attention over observation features . However , they both employ transformer-style attention ( Locatello et al. , 2020 ; Vaswani et al. , 2017 ) to dynamically attend to spatial positions in the 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 a Convolutional LSTM ( ConvLSTM ) ( Shi et al. , 2015 ) . We implement the update function of each module with an LSTM . We used multihead-attention ( Vaswani et al. , 2017 ) for f ( i ) share . We trained the architecture with the IMPALA algorithm ( Espeholt et al. , 2018 ) and an Adam optimizer ( Kingma & Ba , 2015 ) . We tune hyperparameters for all architectures with the “ Place X next to Y ” task from the BabyAI environment ( Chevalier-Boisvert 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 on hyperparameters , see §B .

Summary: 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: 1.  a recurrent encoder produces a set of $m$ feature vectors of dimension $d_z$, one vector for a different spatial location 2.  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)}$. 3. The update rule of each schema module produces the new hidden state $h_{t}^{(i)}$ from three inputs:     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.     2. $c_t^{(i)}$     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]$  The 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.  The authors apply their architecture to three generalization scenarios: 1. " 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. 2. " 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. 3. "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.