T ASK -D RIVEN D ISCOVERY OF P ERCEPTUAL S CHEMAS FOR G ENERALIZATION IN R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A B S T R AC T
Deep reinforcement learning (Deep RL) has recently seen signiﬁcant 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-conﬁgurations.
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-conﬁgurations.
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.
I N T RO 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 ﬁrst 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-conﬁgurations (Figure 1, c).
We hypothesize that discovering perceptual schemas for these regularly occurring structures will facilitate zero-shot generalization to tasks deﬁned 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 beneﬁts 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:
FARM, which combines dynamic feature attention with learning multiple recurrent modules (§3).
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 ﬁlled with more distractors (§4).
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 beneﬁts that come from learning multiple recurrent modules.
Our analysis of the representations learned by FARM provide evidence that it learns perceptual schemas that are ﬂexibly distributed across combinations of recurrent modules (§4.3.1).
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.
; Chaplot et al.
In particularly, Chaplot et al.
(2018) showed that mapping language instructions to non-linear feature coefﬁcients enabled generalizing to tasks speciﬁed over unseen feature combinations in a 3D environment.
While FARM also learns non-linear feature coefﬁcients, our work has two important differences.
First, we develop a multi-head version where individual feature coefﬁcients 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 ﬁrst 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-speciﬁc 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 deﬁned by shape-color agnostic object-dynamics and 3D objects.
A R C H I T E C T U R E : F E AT U R E A T 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 ( x t , τ t ) ∈
X × T , takes an action a t ∈ A , and experiences resultant reward r t ∈ 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 π [ (cid:80) t =
γ t R t +1 | S t = a, A t = a ] .
Since the agent only observes observation-task-actionreward tuples ( x t , τ t , a t , r t ) , it needs to learn a time-series representation mapping episode histories ( x ≤ t , τ ≤ t , a <t , r <t ) to state representations s t .
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 Z t ∈ Z ∈ R m × d z where d z 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 distributed n
∈ and the subsequently experienced learn the following mapping, η (
Z × T × H ( i ) ×
A × R → H ( i ) .
For convenience, we can deﬁne 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 , a t − 1 , r t − 1 ] .
The agent integrates this context c ( i ) t with Z t to compute s t using the collection of outputs from the modules:
[Equation]
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.
Deﬁning f
( i ) update as the module update-function, we model η ( i ) as the following composition of functions:
[Equation]
Modules capture diverse perceptual structures with dynamic feature attention.
Our ﬁrst insight was that modules can attend to high-level object-dynamics information if FARM learns a recurrent observation encoder Z t = φ ( x t , Z t − 1 ) .
By doing so, information about how feature-values shift between positions (i.e. feature dynamics) can be captured in the features of Z t .
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 coefﬁcients and applies them identically to all spatial positions in Z t (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.
; Hu et al.
The operations are summarized below:
[Equation]
where (cid:12) 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 ﬂatten 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-coefﬁcients 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 deﬁne the collection of previous module-states as H t − 1 =
(cid:104) h (1) t − 1 ; . . .
; h ( n ) t − 1 ; 0 (cid:105) ∈ R ( n +1)
× d h , where 0 is a null-vector used to retrieve no information.
Each module retrieves information as follows:
[Equation]
E X P E R I M E N T S
In this section, we study the following questions:
Can FARM enable generalization of memory-retention to novel spatio-temporal compositions of object-dynamics?
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?
[Table]
Baselines.
Our ﬁrst 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 ﬁxed 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.
[Figure]
We study this with the “Ballet” 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 a distinct object-motion for 16 time-steps.
There are 15 object-motions.
The dances appear in sequence for 16 time-steps with a 48-time- delay in between. Training always consists of seeing = { always consists of seeing m = { 8 } After all ﬁnish the the
A poorly performing will obtain chance performance with
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 ﬁnd 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 a recurrent observation encoder, using multiple modules, and using feature-attention.
We conﬁrm that recurrent encoder is required.
Interestingly, we ﬁnd that either using multiple modules or using our feature-attention enables task-learning, with our feature-attention mechanism being slightly more stable.
We hypothesize that using a recurrent encoder with our attention allows for storing information about how object-motions are shifting position in a way that mitigates interference.
(Lampinen et al., 2021) 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-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. G ENERALIZING SEQUENTIAL ACTIVE PERCEPTION OF 3D OBJECTS TO LARGER ENVIRONMENTS
We study this with the 3D Unity environment from Hill et al.
(2020) shown in Figure 1 (b).
This environment tests an agent’s ability to generalize sequential composition of active perception of 3D objects to larger environments.
The agent is an embodied avatar in a room ﬁlled with task objects and distractor objects.
The agent has 46 discrete actions and has a navigation step-size of . 15 m .
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
=
A ∪ B and objects to place them on O
= C ∪ D .
During training the agent sees A × D and B × C in a 4 m × 4 m room with 4 distractors, along with A × C and B × D in a 3 m × 3 m room with 0 distractors.
We test the agent on A × C and B × D in a 4 m × 4 m room with 4 distractors.
The training curriculum also includes “Go to X” and “Lift X”.
For details see § E.3.
We present the generalization success rate in Figure 4.
All architectures and ablations had approximately 5M parameters.
We ﬁnd 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 to 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 a softmax that may saturate on individual positions
4. G ENERALIZING MEMORY - RETENTION TO LONGER TASKS
In order to study this, we create the multi-level “keybox” environment depicted in Figure 1 (c).
This environment tests an agent’s ability to learn to remember goal-information for increasingly long timehorizons.
The agent maximizes reward by completing as many levels as it can in an episode.
Each level in this environment is a hallway with a single box and a key of the same color .
The hallway for level n consists of a length- n sequence of w × w environment subsections.
Each subsection contains d 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 max where n max is the maximum level the agent can complete.
We set n max = 10 during training.
The agent has 50 n time-steps to complete a level.
The distractors pose an additional challenge of obstructing the agent’s path.
This further challenges memory since they have to be incorporated into the state-representation to predict actions but should not overwrite the goal information.
We study two generalization settings : a 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 =distractors.
In the dense setting, once the agent fails, it restarts in level n ∈ [1 , n done ] , where n done is the highest level the agent has completed.
In the sparse setting, we found we needed to always restart the agent on level 1 to enable learning.
All architectures used approximately 7.5M parameters.
We present the maximum training level reached and generalization success rates in Figure 5.
We ﬁnd that all methods achieve comparable training results.
In the dense setting, we see an LSTM quickly overﬁts in both settings.
All architectures with attention continue to improve in generalization performance as they continue training.
In the dense setting, we ﬁnd 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.
4.3. A NALYSIS 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 6 categories of regularly occurring perceptual events: pickup ball, drop ball, pickup wrong key, drop wrong 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 coefﬁcients.
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 programmatically extracting all task-relevant events for analysis is 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 Figure 6 a. While some modules are selective for different recurring events such as attending to goal information (Figure 6, 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 ﬁnd that having 4 (8.2M params) or 8 (7.5M params) modules signiﬁcantly outperforms using a single large module (7.5M params) (Figure 5, (c)).
Feature attention consistently improves performance.
D I S C U
[Equation] S I O N A N D C O N C L U S I O N
We have presented, FARM, a 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-information 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 object-motions (Figure 3) 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 two 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, §C.3).
Feature attention vs. spatial attention.
While our dynamic feature attention mechanism is ﬂexible,
FARM learns an update function that treats all positions in the attention output as unique (equation 4).
A major advantage of spatial attention is that the resultant representation is spatially-invariant , which may sometimes be more useful for a policy.
For example, we ﬁnd that RIMs has better generalization to an unseen number of distractors (§C.2).
Interestingly, using spatial attention seems to interfere with the generalization beneﬁts of learning multiple recurrent modules.
An interesting future direction might be to explore combining both forms of attention for a more ﬂexible 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, Transformers 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 observation-representations for a transformer-like architecture.
Such an architecture might enable discovering perceptual schemas over very long time-horizons.