Unsupervised Learning of Temporal Abstractions with Slot-based Transformers
Abstract
The discovery of reusable sub-routines simpliﬁes decision-making and planning in complex reinforcement learning problems.
Previous approaches propose to learn such temporal abstractions in a purely unsupervised fashion through observing state-action trajectories gathered from executing a policy.
However, a current limitation is that they process each trajectory in an entirely sequential manner, which prevents them from revising earlier decisions about sub-routine boundary points in light of new incoming information.
In this work we propose SloTTAr, a fully parallel approach that integrates sequence processing Transformers with a Slot
Attention module for learning about sub-routines in an unsupervised fashion.
We demonstrate how SloTTAr is capable of outperforming strong baselines in terms of boundary point discovery, while being up to 30x faster on existing benchmarks.
Introduction
An intelligent goal-seeking agent situated in the real world should make decisions along various timescales to act and plan efﬁciently.
A natural approach that facilitates this behavior is via a divideand-conquer strategy, where goals are decomposed along sub-goals, and solutions to such sub-goals (i e. ‘correct’ sequences of actions) are stored as reusable ‘primitive’ sub-routines [9, 4, 35].
Viewing complex goal-directed behavior as (novel) compositions of known sub-routines simpliﬁes both decision-making and planning [36, 43], while also beneﬁting out-of-distribution generalization [32].
The options framework [42] introduced a design strategy for the hierarchical organization of behavior within the context of reinforcement learning.
Crucial to its success is the quality of each ‘option’ (or sub-routine) in terms of the level of modularity and reusability it offers.
Consider for example, the task of planning a travel itinerary to go from London to New York.
Useful sub-routines in this context may include purchasing ﬂight tickets, navigating to the airport, or boarding the ﬂight.
This factorization into sub-routines is desirable as these sub-routines capture mostly self-contained activities.
Thus they are far more likely to apply to other travel plans between two different cities, which may again involve taking ﬂights, etc.
Prior approaches propose to address this issue by learning about useful sub-routines directly from data [2, 37, 23, 27].
Of particular interest is the fully unsupervised setting, where the learner is only given access to state-action trajectories from an (expert) policy.
Two relevant works can be distinguished: CompILE [23] and OMPN [27].
In CompILE expert trajectories are modeled using an RNN-based latent variable model that infers a pre-determined number of boundary points by iteratively processing the entire trajectory multiple times to recover segments associated with reusable sub-routines.
Alternatively, OMPN equips an RNN with a multi-layer hierarchical memory, where information processing at different levels in the memory can be interpreted as belonging to separate segments.
While OMPN additionally includes a strong preference for capturing hierarchical relationships between different sub-routines, both methods are limited insofar they process the Workshop on Deep Reinforcement Learning, NeurIPS 2021. trajectory in an entirely sequential manner.
This prevents them from revising earlier decisions about boundary points in light of new information that becomes only accessible at a future stage.
Moreover, iteratively processing the sequence multiple times (as in CompILE) or interfacing with a deep hierarchical memory (as in OMPN) incurs signiﬁcant computational costs (Table 2).
In this paper we propose a novel architecture to learning sub-routines that addresses these shortcomings, which we call Slot-based Transformer for Temporal Abstraction (SloTTAr) .
Central to our approach is the similarity between the spatial grouping of pixels into visual objects and the temporal grouping of state-action pairs into self-contained sub-routines.
This motivates us to combine a parallel Transformer encoder with a Slot Attention module [26] (developed for learning object representations [16]), to group learned features at different temporal positions and recover a modular factorization into ‘slots’ (sub-routines) of the inputs.
A parallel Transformer decoder reconstructs the action sequence from each slot and outputs an unnormalized distribution over endpoints.
From these, the segment belonging to each sub-routine and a standard reconstruction objective for unsupervised training can be derived.
We demonstrate the efﬁcacy of our approach on Craft [27] and MiniGrid [6] where we typically observe signiﬁcant improvements over CompILE and OMPN in terms of recovering ‘ground-truth’ sub-routines.
In this way, we show how general principles of similarity-based grouping used to segment visual inputs into objects [40, 25, 24, 16] are also relevant for grouping other input modalities [33].
Method
Given an input sequence of actions a = [ a 1 , a 2 , ..., a L ] : a l ∈ { 1 , .., A } where A is the size of the action space and observations o =
[ o 1 , o 2 , ..., o L ] : o l ∈ R D obs of length L .
Our goal is to infer the unique constituent latent sub-routine to which a pair of ( a l , o l ) belongs and learn its associated representation ( slot_k ∈ R D slots ).
Similar to CompILE [23], we consider an unsupervised reconstruction objective for training, yet here we strive for a fully parallel approach that employs more general-purpose modules for grouping inputs based on their internal predictive structure.
Our model, SloTTAr, consists of three modules, as shown in Figure 1.
First, a Transformer encoder [44] learns suitable spatio-temporal features from the input sequences a , o .
This encoder beneﬁts from global access to the whole input sequence to learn suitable context features at each location, unlike prior purely sequential RNN-based approaches [23, 27].
Subsequently, we use a Slot
Attention module [26] to iteratively group the features at each temporal location in parallel based on their internal predictive structure and obtain K slot representations slot_k .
The slot representations are decoded in parallel by a decoder that outputs both a predicted actions ˆ a ( k ) and an unnormalized distribution over the end position of the sub-routine modelled by the respective slot ( end_logits_k ).
Inputs: inputs ∈ R L × D in , slots where slot_k ∼ N ( µ k , diag ( σ k ))
∈
R D slots ∀ k ∈ K
Params: Linear projections for attention: key, query, value ; GRU; MLP; LayerNorm x 1: for t = 0 , 1 , ..., T do 2: slots_prev = slots 3: slots = LayerNorm(slots) 4: attn =
Softmax( 1 √ D slots key (inputs) · query (slots) T , axis=‘slots’)
(cid:46) norm.
5: updates := WeightedSum(weights=attn, values= value (inputs))
(cid:46) aggregate 6: slots =
GRU(state=
slots_prev, inputs=updates) (cid:46) GRU update (per-slot) 7: slots += MLP(LayerNorm(slots))
(cid:46) residual update (per-slot) 8: end for 9: return slots
Sub-routine segment masks are generated from the unnormalized distributions to compute the aggregated action sequence logits ˆ a .
Finally, a reconstruction objective between ˆ a and a can be deﬁned.
Encoder.
We process action tokens a and observations o by Embedding and Linear layers respectively to acquire separately learned distributed representations for each.
Next, we learn a joint representation of the action and observation representations at each position using an MLP , after which a sinusoidal positional encoding p sin l ∈ R D enc is added [44].
Finally, we apply a number of standard
Transformer encoder layers to learn suitable spatio-temporal features based on the content of the entire sequence and add a learned positional encoding p learn ∈ R D enc for follow-up processing: z al , z ol
=
Embedding ( a l ) , Linear ( o l )
[Equation]
[Equation]
h =
TransformerEnc ( z
[Equation]
+ p sin )
+ p learn
Slot Attention.
We use Slot Attention [26] to group these spatio-temporal features according to their constituent sub-routines and learn associated representations given by the slots.
Slot Attention was previously only applied to the visual setting where pixels are grouped according to constituent visual objects to model images.
Here we hypothesize that sub-routines take on a similar role as modular and reusable primitives when modeling action sequences, suggesting that they similarly can be inferred using a grouping mechanism that focuses on their internal predictive structure (modularity) [16].
The iterative grouping mechanism in Slot Attention (reproduced in Algorithm 1) is implemented via key-value attention and a stateful slot update rule using a recurrent neural network ( GRU [7], see also earlier work [12]).
The input features are projected to keys and values using linear layers key ( · ) , value ( · ) , queries are computed from slots using a linear layer query ( · ) , and dot-product attention between keys and queries is used to distribute value vectors among the slots.
Initial representations for each slot ( slot_k ) are sampled from a Gaussian N ( µ k , diag ( σ k )) .
Importantly, slots compete to represent parts of the input sequence via a Softmax over slots, thereby encouraging a decomposition of the input into modular parts that can be processed separately and represented efﬁciently.
Unlike the original Slot Attention formulation, we use separate shift and scaling parameters ( µ k and σ k ) for each slot, which allows for ordering them as later needed by the decoder.
Finally, we replaced the WeightedMean (line 9 in Algorithm 3 in Appendix A.3), which we found to lead to more stable training and better performance.
Decoder.
To decode the slot representations and obtain a correspondence of actions in a to their constituent sub-routine segments, we adapt the spatial broadcast decoder architecture [45] to the sequential setting using Transformers.
Each slot representation is decoded in parallel to obtain the predicted action logits ˆa ( k ) as well as an unnormalized distribution over the end position of the sub-routine in the sequence ( end_logits_k ).
The latter is used to generate the segment masks corresponding to each sub-routine as described in Algorithm 2.
We enforce temporal contiguity of the segment modeled by each slot and use a ﬁxed ordering to compose them, which is different from the standard mixture formulation adopted by prior work on discovering visual objects [15, 14, 26].
This has a similar effect to the sequential decoding in CompILE [23] and provides a useful inductive bias for treating sub-tasks as contiguous chunks in time that can be ordered along the temporal axis.
The entire system is trained end-to-end to minimize the cross-entropy loss between the aggregated
[Table]
predicted action logits
[Table]
k mask_k · ˆa ( k ) and the ground-truth a . For additional details on to Appendix
Inputs: end_logits ∈ R K × L , masks = [], end_cdf_k = 1 , mask_upto_km1 =1: for k = 1 , ..., K do 2: end_dist_k =
Softmax(end_logits_k, axis=‘L’)
(cid:46) norm.
over sequence length 3: end_cdf_k
Additional Related Work
The Craft environment was originally introduced to evaluate the method of Andreas et al.
[2], however they rely on annotations of sub-routine sequences (‘policy sketches’) as additional supervision to the model.
TACO [37] was also proposed to address this setting, but views it as a sequence alignment and classiﬁcation problem using an LSTM [18] trained with a CTC loss [13].
In other recent work [1] the primitives are learned directly from ofﬂine data for continuous control.
However, these methods [1] learn a continuous (low-dimensional) space of primitives whereas our method represents primitives as a discrete set.
In the context of the options framework, several methods have been proposed for option and/or sub-goal discovery [3, 28, 29, 41, 8, 39, 38].
However, these methods are not easily extendable to the case with highly nonlinear function approximation.
Alternatively, other methods using nonlinear function approximation seek to maximize coverage (diversity) of learned skills by maximizing the mutual information between options and terminal states achieved by their execution [17, 11].
Many recent works have applied Transformers to modalities beyond natural language, such as images [10, 46], and other high-dimensional input domains [21, 20].
The application of Transformers to reinforcement learning has also recently received a lot of interest [31, 34, 22, 5, 19].
Experiments
We compare SLoTTAr to CompILE [23] and OMPN [27] on the partially and fully observable versions of Craft environment [27] and on a partially observable environment in MiniGrid [6].
Here our focus is on the fully unsupervised setting without the use of task sketches as an auxiliary supervision [2, 37].
We use 3 tasks in the Craft environment namely, MakeAxe, MakeBed, MakeShears consistent with prior work [27].
In MiniGrid we use the DoorKey-8x8 environment.
Episodes in this environment are procedurally generated leading to far greater variability for actions that constitute a sub-routine compared to Craft.
Further in the DoorKey-8x8 environment, two actions ( ‘PICKUP’ or ‘TOGGLE’ ) are used to determine sub-routine boundaries as opposed to in Craft where this is marked by only a single ‘USE’ action.
[Table]
Evaluation.
To quantitatively measure the quality of the action sequence decomposition, we use the F1 score (with a tolerance=1) based on the accuracy of the boundary index predictions with respect to ground-truth.
Further, we also report the alignment accuracy between the predicted and ground-truth sub-routines consistent with prior work [27].
Please refer to Appendix A.1 for further details on how ground truth sub-routine boundaries are computed.
Hyperparameter Search.
Below we report results after conducting an extensive hyperparameter search (up to 200 conﬁgurations) for each method.
We include training parameters like batch size and learning rate, capacity parameters like layer sizes, and model-speciﬁc parameters such as the number of Slot Attention iterations in case of SloTTAr.
Results are reported for the best conﬁgurations with mean and std.
dev.
computed over 5 seeds.
Details about the search parameters for each method, including details about the best conﬁgurations can be found in Appendix A.3.
Results.
Table 1 shows the results of comparing SloTTAr to both baselines on all three environments.
On the fully observable Craft tasks, it can be seen how SLoTTAr outperforms both CompILE and OMPN.
Similarly, on the partially observable settings in Craft, it can be seen how SloTTAr outperforms CompILE, although this time it performs worse compared to OMPN.
We speculate how the strong hierarchical inductive bias in OMPN gives it an edge in this setting, which closely reﬂects the hierarchical sub-task structure in Craft.
Indeed, when we only consider OMPN conﬁgurations having two levels of memory (OMPN-2) it can be observed how the difference in alignment accuracy to the best performing conﬁguration reduces 1 .
Finally, on the harder DoorKey-8x8 partially observable environment, it can observed how SloTTAr signiﬁcantly outperforms both CompILE and OMPN in terms of both metrics.
These results demonstrate the strengths of our approach and suggest that a fully parallel approach to computing sub-routine boundary points is highly beneﬁcial.
Another added advantage of using a parallel architecture is the potential gain in computational complexity.
Table 2 reports the number of tokens processed per-second during training (forward and backward passes) and at testing time (forward pass) for each model.
It can be seen how SloTTAr is about 7x faster compared to CompILE and about 30x faster compared to the OMPN model, suggesting that it also may scale better to more sophisticated environments.
Reducing the memory depth in case of OMPN only yields a marginal improvement.
Conclusion
We have proposed a novel approach to learning action segments belonging to modular sub-routines (suitable as ‘options’ [42]) in a purely unsupervised fashion.
Our approach draws insight from prior literature on learning about visual objects [26] and Transformers [44] to improve over existing purely sequential approaches.
It was shown how SloTTAr outperforms CompILE and OMPN in terms of recovering ground-truth sub-routine segments and provides a massive speed-up by leveraging parallel computation over the time axis.
To that extent we demonstrated how general principles of similarity-based grouping used to segment visual inputs are also relevant for grouping other input modalities, suggesting many additional avenues for future research.
Acknowledgements
We thank all members of our research group and the anonymous reviewers for their valuable feedback.
This research was funded by Swiss National Science Foundation grant: 200021_192356, project
NEUSYM. This work was also supported by a grant from the Swiss National Supercomputing Centre (CSCS) under project ID s1023. We also thank NVIDIA Corporation for donating a DGX-1 as part of the Pioneers of AI Research Award and to IBM for donating a Minsky machine. A Experimental Details A.1 Datasets
Craft The craft environment was initially introduced and re-implemented as a gym environment [2, 27].
Demonstration data is collected from the rollouts of the heuristic bots [27].
We use rollouts from each of the 3 tasks MakeAxe, MakeBed and MakeShears for both the fully observable and partial observable versions of these environments [27] to report the results in Table 1.
We use 10trajectories for training and another separate 1000 trajectories each for creating validation and test splits.
These are used for hyperparameter tuning and reporting evaluation metrics respectively.
MiniGrid We use the DoorKey-8x8 environment from MiniGrid [6] as an additional dataset.
This environment uses procedural generation for each episode ensuring more variability in the action sequences that make up a sub-routine in comparison to Craft.
We collect demonstration data by training an A2C agent [30] on this environment until the policy receives an episodic return of ≥ 0 .
9 , which is close to optimal.
We use 10 000 rollouts collected from this “expert” A2C agent [30] as the training set and 1000 rollouts each for the validation and test splits used to tune hyperparameters and report evaluation metrics respectively.
Ground-truth segment generation For the Craft environment, the ‘USE’ action serves as a delimiter that marks the end of sub-routines.
In DoorKey-8x8 environment, the ‘PICKUP’ and ‘TOGGLE’ actions serve as delimiters.
We extract the ground-truth boundaries of sub-routines using the indices of these action tokens in the sequences as markers.
Ground-truth sub-routine indices for every sequence are in ascending order from left-to-right where delimiting points are speciﬁed using the heuristics described above.
Table 3 shows the ground-truth sub-routines for each of the environments.
[Table]
A.2 Evaluation Metrics
Evaluation metrics used to quantitatively measure the segmentation performance of models are the alignment accuracy of sub-routine prediction and F1 score of boundary index prediction and is consistent with prior work [27, 23].
Alignment Accuracy The expression to compute alignment accuracy is shown below:
[Equation]
where L is the sequence length (possibly different for each sequence), N is the number of sequences, ˆ s l is the predicted sub-routine and s l is the ground-truth sub-routine of action a l in the n th sequence.
F1 Score
The F1 score on predicted boundary indices is computed as shown below:
[Equation]
where precision is computed as, precision = # matches of boundary predictions with ground-truth total # boundary predictions and recall is computed as, recall =
# ground-truth matches with predictions total # ground-truth boundaries
Please refer to the Appendix D.1.
in prior work [23] for further details.
Inputs: inputs ∈ R L × D in , slots where slot_k ∼ N ( µ , diag ( σ ))
∈
R D slots ∀ k ∈ K
Params: Linear projections for attention: key, query, value ; GRU; MLP; LayerNorm x 1: for t = 0 , 1 , ..., T do 2: slots_prev = slots 3: slots = LayerNorm(slots) 4: attn =
Softmax( 1 √ D slots key (inputs) · query (slots) T , axis=‘slots’)
(cid:46) norm.
5: updates := WeightedMean(weights=attn, values= value (inputs))
(cid:46) re-normalizing attn values over L 6: slots =
GRU(state=
slots_prev, inputs=updates) (cid:46) GRU update (per-slot) 7: slots += MLP(LayerNorm(slots))
(cid:46) residual update (per-slot) 8: end for 9: return slots
A.3 Architecture and Training Details
Input Processing The Embedding layer used (Section 2) maps action tokens a l to D enc dimensions (denoted by hidden size in Table 4 and Table 5).
The Linear layer used (Section 2) maps observations o l to D enc dimensions.
These D enc dimensional action and observation features are then concatenated and processed by a 1-layer MLP with D enc units and ReLU activation to learn a joint actionobservation embedding space.
Encoder The Linear layers used in self-attention have D enc units in the TransformerEnc block.
The pointwise feedforward network used in the TransformerEnc uses 2x the number of units as hidden size.
We add the standard sinusoidal positional encoding p sin [44] to the joint actionobservation features z ao .
Further, we use a Linear layer to generate the learned positional encoding p learn that is added to the outputs h from the TransformerEnc module.
The ﬁnal and sweep conﬁguration(s) for each of these hyperparameters are shown in Table 4 and Table 5 respectively.
The overall TransformerEnc block uses the same architectural template as the Transformer Encoder [44].
Slot Attention
The slot attention module uses Linear layers with D slots units to generate keys and values.
The recurrent update function is implemented using a GRU [7] (see also earlier work [12]) with D slots units (denoted by slot size in Table 5 and Table 4).
The MLP network used for the residual update is implemented using a 2-layer network with D slots units and ReLU and Linear activation for the hidden layer and output layers respectively.
Futher, the initial query vector for slot_k is sampled from a Gaussian distribution with learnable mean µ k and standard-deviation σ k .
For reference Algorithm 3 describes the original Slot Attention update [26].
Decoder The decoder architecture adapts the spatial broadcast decoder [45] in a suitable manner to decode all slots into their corresponding action segments.
We broadcast the slot representations along the sequence length L and pass the observations o l as inputs at each timestep to the Decoder.
The
Decoder module has the exact same architecture as the Transformer Encoder with the exception that an additional output Linear layer of A + [-
-]1 (where A is the size of action space) dimensions is used to decode the slot representations into predicted action logits and the end position logits.
Segment masks are computed given end position logits as shown in Algorithm 2.
Then we composite segment masks and predicted action logits to generate the full predicted action sequence.
We perform a random search of 200 hyperparameter conﬁgurations from all possible conﬁgurations in the hyperparameter sweeps shown in Table 4 on 3 seeds.
We train our model for 50 epochs on
Craft (fully-observable) and 100 epochs for both the Craft (partially-observable) and DoorKey-8x environments.
We observed that our models start overﬁtting on the training set beyond this point.
Then, we picked the best performing conﬁguration(s) for each of the 3 environments for our model based on the evaluation metrics on the validation set and run them again on 5 seeds to report ﬁnal scores (shown in Table 1).
The ﬁnal conﬁgurations used for our model are shown in Table 5.
[Table] [Table]
A.4 Baseline Models and Training Details
We developed our re-implementations of CompILE 2 and OMPN 3 by adapting the authors’ implementations available on GitHub.
We perform a random search of 200 hyperparameter conﬁgurations from all possible conﬁgurations in the hyperparameter sweeps shown in Table 7 and Table 8 for each of the 3 environments and both baseline models (CompILE and OMPN) with 3 seeds.
Then, we picked the best performing conﬁguration(s) based on their evaluation metrics on the validation set from this sweep.
We run these best performing conﬁguration(s) on 5 seeds to obtain all results shown in Table 1.
We train CompILE and OMPN models for 30 and 20 epochs respectively as we observed that these models start overﬁtting on the training set beyond this point.
We checkpoint models during the course of training based on their evaluation metric scores on the validation set.
We use the best checkpoints to report the ﬁnal scores shown in Table 1 for all the baseline models.
The ﬁnal conﬁgurations used for the baseline models on the 3 environments are shown in Table 6 and
Table 9.
[Table] [Table]
[Table] [Table]