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 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.  The idea of CPS is as follows:  * 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.  * 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.  * 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).  * The dynamic attention for other modules is similar to the implementation in RIMs using transformer-style attention.  CPS 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.).