Structure-Aware Transformer Policy for Inhomogeneous Multi-Task Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) has made remarkable successes over the last several years , achieving human-level performance in various tasks ( Silver et al. , 2016 ; Mnih et al. , 2015 ) . However , these are limited by individually training policies for each task , often requiring a large amount of interaction data . Multi-task learning , an approach to training a model jointly from diverse tasks in order to improve learning efficiency and prediction accuracy of task-specific models by exploiting commonalities and differences across tasks , has become prevalent in computer vision ( CV ) ( Dosovitskiy et al. , 2021 ) and natural language processing ( NLP ) ( Devlin et al. , 2019 ; Radford et al. , 2019 ) . In this regard , Multi-Task Reinforcement Learning ( MTRL ) is a promising approach , but training a policy for multiple tasks in the traditional manner is not straightforward in many cases , e.g. , a policy designed for a specific robot can not be reused for another one with a different embodiment . Therefore , most MTRL methods assume same state and action dimensions across tasks ( Rusu et al. , 2016 ; Parisotto et al. , 2016 ; Pinto & Gupta , 2017 ; Yang et al. , 2020 ; Kalashnikov et al. , 2021 ) and define each task by its own reward function , e.g. , grab a cup or move it with a robot arm . Instead , we are interested in a more general setting where the tasks are inhomogeneous , i.e. , of different state and action space dimensions , also known as incompatible control ( Kurin et al. , 2021 ) . One of the popular approaches to inhomogeneous MTRL is to assume a graph structure for the agent , depicted in Figure 1 as an example , where limbs and joints are represented as nodes and edges . In this modular setting , Graph Neural Networks ( GNNs ) provide a natural choice for the model of the policy , since ( 1 ) they can process graph-structured input of arbitrary sizes and connections , which allows us to obtain a single policy model that can control any agent with various morphology leading to different state and action space dimensions , and ( 2 ) they natively exploit the inductive bias that arises from the morphological structure through their message-passing ( MP ) scheme , i.e. , messages are allowed to be delivered only to nodes connected by edges . Hence , GNN-based approaches have shown to yield significantly better results in inhomogeneous MTRL , compared to the vanilla approaches that don ’ t take into account the morphological structure of the agent . Recently , Kurin et al . ( 2021 ) posed a concern that the MP scheme in GNN is prone to oversmoothing ( Li et al. , 2018 ) , where the crucial information is washed out during the multi-hop communication . They argued that the advantage of leveraging morphological information for inhomogeneous MTRL was overshadowed by this problem . They therefore advocated adopting the self-attention mechanism , which allows direct communications among nodes , while sacrificing the morphological information . Yet , it is well known that injecting the structural inductive bias into the self-attention mechanism , i.e. , positional embedding , can help significantly improve the performance ( Vaswani et al. , 2017 ) . In the same context , the states of neighbor nodes may be more important than those of non-neighbor nodes for determining actions in modular RL . In this paper , we introduce Structure-aWAre Transformer ( SWAT ) , a new modular method for inhomogeneous MTRL that effectively incorporates the agent morphology into the transformer-based policy . We propose two forms of structural embeddings of the morphology to be incorporated into the transformer model : ( 1 ) traversal-based positional embedding ( PE ) which represents absolute positions of nodes via tree-traversal algorithms , and ( 2 ) graph-based relational embedding ( RE ) which represents relative distances of node pairs reflecting their connectivity . These allow direct communication of messages among nodes while taking into account the agent morphology . The experimental results on benchmark tasks for inhomogeneous MTRL strongly support the effectiveness of our structural embedding approach , outperforming prior state-of-the-art methods in terms of sample efficiency and final performance . Through transfer learning experiments , we also demonstrate that the structural information conveyed by our approach plays an important role for downstream tasks . 2 BACKGROUND . 2.1 INHOMOGENEOUS MTRL . We assume a Markov Decision Process ( MDP ) defined by tupleM = ( S , A , T , ρ0 , r , γ ) to represent an RL task , where S is the state space , A is the action space , T ( st+1|st , at ) is the state transition probability , ρ0 is the initial state distribution , rt = r ( st , at ) ∈ R is the reward associated with state and action , and γ ∈ ( 0 , 1 ) is the discount factor . The objective of RL is to learn a probability distribution over actions conditioned on states , i.e. , policy π ( at|st ) , which maximizes the expected discounted returnR = E [ ∑ t≥0 γ trt ] . In the inhomogeneous MTRL , the agent is tasked with a set of inhomogeneous MDPs M = { M1 , M2 , · · · , MK } . We say MDPs are inhomogeneous when any pair of MDPs have different dimensionalities in their state or action spaces , i.e. , dim ( Si ) 6= dim ( Sj ) or dim ( Ai ) 6= dim ( Aj ) , whereMi , Mj ∈M and Si , Ai is the state and action space ofMi . In the context of MTRL , the goal is to find a policy that maximizes the average expected discounted return over all the environments , 1 K ∑K i=1 E [ Ri ] whereRi denotes the expected discounted return in MDP Mi . We now describe modular RL , where the agent can be represented as an undirected graph G = ( V , E ) , e.g. , the robotic agent with n limbs and l joints as in Figure 1 . Each node ( vertices ) vi ∈ V for i ∈ 1 ... n represents a limb , and an undirected edge ei , j ∈ E represents connectivity by a joint between a vi and vj . For the ease of notation , we assume an ordered set of nodes , i.e. , a node vi is i-th element in the ordered set . E can be represented as an adjacency matrix A ∈ { 0 , 1 } n×n that describes the connectivity among nodes , i.e. , if vi , vj are connected , Ai , j = 1 and∑ i , jA i , j = 2l . Thus at each time step t , the agent observes the state st , which consists of local sensory information of the limb such as limb type , coordinates , and angular velocity . Denoting the local sensory information of the limb vi at time step t as vector vit , the state st can be represented as { v1t , v2t , . . . , vnt } with morphology E . Given st , the policy outputs the action at = { a1t , a2t , . . . , ant } where ait is the torque value for the corresponding actuator controlled by the limb v i . 2.2 POLICY MODELS FOR INHOMOGENEOUS MTRL . Traditional policy model for RL , e.g. , multi-layer perceptron ( MLP ) policy , is not suitable for inhomogeneous MTRL since it can only handle the state and action spaces with the same dimension . In this regard , GNNs ( Hamilton et al. , 2017 ; Wu et al. , 2020 ) have been a natural choice for modeling policies ( Wang et al. , 2018 ; Huang et al. , 2020 ) , since they can scale to different state and action dimensions as well as incorporate the agent morphology in their message-passing ( MP ) scheme . Let us assume a graph G = ( V , E ) where each node vi ∈ V has the corresponding information , i.e. , node representation or message vi . The messages in GNNs are delivered as follows : 1 . MESSAGE AGGREGATION : mi ← σ ( { vj : j ∈ Ni } ) , ∀vi ∈ V 2 . NODE UPDATE : viupdate ← fθ ( vi , mi ) , ∀vi ∈ V , ( 1 ) whereNi is the indices of the neighborhood of vi , σ is a message aggregation function , e.g. , average or concatenation , fθ is a parameterized node update function that computes a new message given the original message vi and aggregated message mi . GNNs consist of the stack of multiple GNN-layers for multi-hop communication , where a single layer operates a single iteration . GNNs are suitable for inhomogeneous MTRL with various agent morphologies , since each node is processed in a modular manner , allowing a single policy to be learned for diverse agents . In particular , Shared Modular Policies ( SMP ) ( Huang et al. , 2020 ) assumes a tree morphology and adopts both-way MP scheme where messages are delivered back and forth between a root and leaves , showing significantly better performance than a standard monolithic policy in inhomogeneous MTRL . The self-attention mechanism , introduced in the transformer ( Vaswani et al. , 2017 ) , is another way of passing messages . Unlike GNNs , it assumes that all nodes in a graph are connected , even enabling direct communication among distant nodes . Specifically , each node feature vi ∈ Rdv is projected to three vectors , qi , ki ∈ Rdk , ṽi ∈ Rdv , where qi = viWq , ki = viWk , ṽi = viWv are a query , key , value vectors respectively , and Wq , Wk ∈ Rdv×dk , Wv ∈ Rdv×dv are learnable projection mappings . Messages are then delivered in the self-attention mechanism as : αi , j = qikjT√ dk , wi , j = exp ( αi , j ) ∑ j′ exp ( α i , j′ ) , viupdate ← ∑ j wi , jṽj , ( 2 ) where αi , j is the unnormalized attention score , and wi , j is the attention weight . When the agent morphology is a sparse graph , which is quite usual for many agents in nature , the crucial information tends to be washed out during multi-hop communication , an undesirable phenomenon known as over-smoothing . In order to avoid the issue , AMORPHEUS ( Kurin et al. , 2021 ) adopts the self-attention mechanism and computes the attention weight over all nodes in the graph , yielding better results than the previous GNN-based approaches . However , it discards the morphological information although it is potentially very useful . 2.3 POSITIONAL INFORMATION IN TRANSFORMERS . In the transformer , injecting the positional information is essential since the self-attention mechanism is permutation invariant ( Vaswani et al. , 2017 ) . In the case of NLP , a permutation of tokens can change the whole context and meaning of a sentence . Thus , the structural bias , representing the positions of tokens , is explicitly given as a form of positional embedding ( PE ) , which is found to be significantly useful at learning the contextual representations of words in different positions ( Ke et al. , 2021 ; Wang & Chen , 2020 ) . To be specific , the absolute positional information for the i-th token vi in the sequence is given by the absolute PE vector pi and it is added to the corresponding token embedding vi . On the other hand , the relative positional information between two tokens , e.g. , relative distance ( i− j ) , is given by relational embedding ( RE ) , ri , j ∈ R , which is added to the attention score as αi , j = qikjT / √ dk + r i , j ( Raffel et al. , 2020 ; Shaw et al. , 2018 ) . 3 MOTIVATION FOR STRUCTURAL BIAS AMORPHEUS ( Kurin et al. , 2021 ) proposes a plain transformer model without morphological information due to the aforementioned over-smoothing problem . However , we can still incorporate the morphological information into the transformer , i.e. , the structural embedding widely utilized in NLP to add the structural bias . In order to appreciate its effect in MTRL , we apply 2 common methods used in NLP models , the learnable embedding ( LEARN ) and the sinusoidal encoding ( SINU ) . LEARN uses a learnable vector to represent the absolute position of each token in a sequence , while SINU uses a sinusoidal function that can represent the relative position between tokens as well . First , we converted the morphology into a sequence by labeling the index of each node along with the pre-order traversal of the tree with a torso labeled as a root . We then aggregated the sinusoidal encoding or the learnable embedding with the node repre- sentations to inject the structural bias into AMORPHEUS . We run AMORPHEUS , LEARN , and SINU on Animal-Walkers , which are trained jointly for 5 morphologically different agents . The results in Figure 2 show how each of the embeddings affects the MTRL performance . LEARN improves the performance the most , because it can roughly capture the knowledge about where each limb is placed . On the other hand , SINU suffers learning instability and the performance rather worsens . The reason for the deterioration , we conjecture , is that SINU encodes morphology incorrectly . SINU , devised to capture a relative distance in sequential data , can give misguided relational information when it is applied to graph-structured data . For example , the distances from a torso to both arms are the same in a graph , but SINU regards them differently . Although this is a preliminary observation , it manifests the necessity of a pertinent embedding scheme for representing both the positional and the relational information of morphology for inhomogeneous MTRL .

Summary: This paper addresses the inhomogeneous multi-task reinforcement learning (MTRL) problem. Following existing works (Huang et al. 2020, Kurin et al. 2021), a policy is represented as a graph. Each node in the graph is an identical modular neural network. The current state-of-the-art approach,  (Kurin et al. 2021),  uses a self-attention mechanism which allows direct communication between nodes. The author argued that the sel-attention mechanism discards the morphological information, which may be critical for agent learning and better perfromance. Therefore, they proposed to encode the morphological information via traversal-based positional embedding and graph-based relational embedding. The proposed approach is simple yet outperforms the state-of-the-art approaches on the MUJOCO MTRL environments.