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: The paper proposes SWAT to incorporate the agent's morphology into transformer-based policy in multi-task RL. Assuming that the agent's limbs are nodes and joints are edges of a graph, two embeddings capture the morphological structure: first, positional embedding (PE) to represent the absolute position of limb nodes by tree-traversal (pre-order, in-order, post-order), and second, relational embedding (RE) to represent the relative distances of limb nodes in terms of their graph connectivity (normalized Laplacian, shortest path distance, personalized page rank).  The hypothesis is SWAT's traversal-based PE and graph-based RE enable effective multi-task policy learning in diverse morphological tasks. In its support, the paper demonstrates that SWAT achieves a higher average return than the baselines GNN-based method (SMP) and morphology-free transformer (AMORPHEUS) in Gym Mujoco locomotion tasks. While SWAT and AMORPHEUS are close in performance on four in-domain locomotion tasks, SWAT shows a significantly higher average return than all baselines in five cross-domain tasks with diverse morphologies. Two ablation experiments analyze the performance impact on excluding PE and RE. The paper also qualitatively examines the behavior patterns for SWAT and AMORPHEUS policies trained in cross-domain (CWHH++) and tested in a single task (Humanoid).