Year: 2021

Abstract:
Modular Reinforcement Learning, where the agent is assumed to be morphologically structured as a graph, for example composed of limbs and joints, aims to learn a policy that is transferable to a structurally similar but different agent. Compared to traditional Multi-Task Reinforcement Learning, this promising approach allows us to cope with inhomogeneous tasks where the state and action space dimensions differ across tasks. Graph Neural Networks are a natural model for representing the pertinent policies, but a recent work has shown that their multihop message passing mechanism is not ideal for conveying important information to other modules and thus a transformer model without morphological information was proposed. In this work, we argue that the morphological information is still very useful and propose a transformer policy model that effectively encodes such information. Specifically, we encode the morphological information in terms of the traversal-based positional embedding and the graph-based relational embedding. We empirically show that the morphological information is crucial for modular reinforcement learning, substantially outperforming prior state-of-the-art methods on multi-task learning as well as transfer learning settings with different state and action space dimensions.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 BACKGROUND

heading: 2.1 INHOMOGENEOUS MTRL
text: 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.

heading: 2.2 POLICY MODELS FOR INHOMOGENEOUS MTRL
text: 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.

heading: 2.3 POSITIONAL INFORMATION IN TRANSFORMERS
text: 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.

heading: 4 METHOD
text: Motivated by the result in the previous section, we present the structural embedding scheme that can be effectively incorporated into the transformer-based policy. It allows our policy to leverage the morphological information while free from the over-smoothing. We introduce two components for injecting the structural bias into the transformer: traversal-based PE and graph-based RE.

heading: 4.1 TRAVERSAL-BASED POSITIONAL EMBEDDING
text: Unlike language sequences in NLP, it is not straightforward to capture the positional information of each node in a graph. Yet, if we assume tree-structured robot morphology defining a torso as the root node, as done in SMP (Huang et al., 2020), we can represent it as a combination of multiple sequences in terms of tree traversals. We can traverse any tree by several consistent orderings, e.g., pre-order or post-order traversal. Although a single traversal sequence, as in section 3, is ambiguous to reconstruct the given tree, we can uniquely identify a binary tree with in-order traversal along with another traversal (Burgdorff et al., 1987).
Inspired by this property, we introduce traversal-based PE for representing the position of each node in the given tree. Since we assume an agent morphology as a general tree with an arbitrary number of children, we apply left-child-right-sibling representation (LCRS) to represent a general tree as
a binary tree. Then, we traverse the binarized tree by pre-order, in-order, and post-order, which is sufficient to reconstruct the original tree. A node position can be represented as a tuple that consists of the position in each traversal. Concretely, pv i
pre, p vi in, p vi post ∈ N denote the indices of a node vi in each traversal ordering respectively, and we assign the PE pi ∈ Rdv to the node vi using the learnable embedding vectors pv i
pre,p vi in,p vi post as follows:
pi = AGGREGATE([pv i pre,p vi in,p vi post]),
vi ← vi + pi (3)
We simply concatenate the embedding vectors in this work in AGGREGATE. By combining multiple traversals, the PE assigned to each node can uniquely represent the node position in the whole morphology.

heading: 4.2 GRAPH-BASED RELATIONAL EMBEDDING
text: In a graph, the relation among nodes also provides structural information from another perspective than the position. Compared to the node positional information above, the relational information contains more explicit knowledge between a pair of nodes. We now introduce graph-based RE by utilizing three meaningful features extracted from the graph, which are incorporated into the transformer.
Normalized Graph Laplacian The Graph Laplacian (Merris, 1994) is a matrix that represents connectivity in terms of both adjacency and node degrees in a graph, defined as L = D−A where A is an adjacency matrix and D is a diagonal degree matrix, i.e., Di,i = |Ni|. We use the normalized graph Laplacian to bound its norm, which is defined as Rlap = D− 1 2LD− 1 2 . Then, Rlap ∈ Rn×n represents the relation among neighbors, capturing local information.
Shortest Path Distance Shortest path distance (SPD) measures the distance between two nodes in a graph, which can be easily computed by several algorithms such as breadth-first search (BFS). It is similar to the relative positional encoding approach in NLP domains, but we use real-value distance Rspd ∈ Rn×n divided by the number of nodes n to bound in [0, 1], i.e., Ri,jspd = SPD(i, j)/n where SPD(i, j) denotes the shortest path distance between vi, vj . SPD considers the direct path among all nodes, so we can say it captures the global information.
Personalized PageRank Personalized PageRank (PPR), a variation of PageRank (Page et al., 1999), represents the proximity between two nodes in a graph based on a random walk model (Park et al., 2019), which is widely utilized in the graph domain (Klicpera et al., 2019; Bojchevski et al., 2020). PPR is a node visitation probability distribution of a -discounted random walk model with a restarting node. To be concrete, the random walker travels across the Markov chain formed via the graph with its transition probability matrix P ∈ Rn×n, where Pi,j = 1/|Ni| for ∀vj ∈ Ni otherwise 0. The random walker warps to the restarting node with probability and follows the P with 1 − . Then, the node visitation probability distribution conditioned on the restarting node vi, PPR(i), can be obtained by solving a recursive form:
x← (1− )xP+ 1i, PPR(i) = 1i(I− (1− )P)−1,
where xj denotes the probability that the walker resides on vj , and 1j denotes the one-hot encoded restarting node vi. Then, the proximity between vi and vj , Ri,jppr, can be denoted by the j-th entry of PPR(i). In contrast to SPD, PPR considers every possible path between two nodes, whereby it captures the relation among nodes, without the overall graph dispensed with.
Although the aforementioned graph features all reflect the agent morphology, they are mutually complementary; for instance, while Graph Laplacian mainly focuses on neighbor information, other features provide knowledge about distant nodes, generating local and global views respectively. We thus represent the relational information among nodes in various perspectives by combining those features. In detail, we aggregate them and learn the graph-based RE R ∈ Rn×n×nhead from them:
Ri,j = gφ([R i,j lap,R i,j spd,R i,j ppr]), (4)
where gφ is a function parameterized by φ that maps the graph features to a RE vector Ri,j and nhead denotes the number of heads. Note that we represent the relational information between two nodes as a vector Ri,j ∈ Rnhead instead of a scalar to take full advantage of multi-head attention. Now we modify the attention score formula with the RE in multi-head setting as follows:
αi,j(h) = qi(h)k jT (h)√
dk/nhead +Ri,j(h), (5)
where the subscript (h) denotes the h-th head and Ri,j(h) is a h-th entry of R i,j .

heading: 4.3 STRUCTURE-AWARE TRANSFORMER POLICY
text: We now present how SWAT works in inhomogeneous control tasks, illustrated in Figure 3. SWAT is based on a 3-stage framework similar to the existing methods: 1) Encode local sensory information vi for each node with the traversal-based PE into the hidden representation by a single MLP layer shared across all nodes. We use the same notation vi for the encoded node representation for clear notation. 2) Update each node representation by MP. While SMP uses the multi-hop communication MP scheme, we use the direct communication through self-attention with graph-based RE. 3) Finally, decide the action for each non-torso node by pooling the final node representations. To learn the policy via an RL algorithm, we use TD3 (Fujimoto et al., 2018), a deterministic policy gradient algorithm in the actor-critic framework, following both SMP and AMORPHEUS.

heading: 5 EXPERIMENT

heading: 5.1 EXPERIMENT SETTING
text: We run experiments on modular MTRL benchmarks (Huang et al., 2020; Wang et al., 2018), which are created based on Gym MuJoCo locomotion tasks. 9 environments can be categorized into 2 settings, in-domain and cross-domain. For in-domain, there are 4 environments: (1) Hopper++, (2) Walker++, (3) Cheetah++, and (4) Humanoid++; they all contain both the intact morphology and its variants, e.g., Humanoid only with one leg or arm, for Humanoid++. The cross-domain environments are combinations of indomain environments: (1) Walker-Humanoid++, (2) Walker-Humanoid-Hopper++, (3) Cheetah-Walker-Humanoid++, (4) Cheetah-Walker-Humanoid-Hopper++ (CWHH++), and (5) Animal-Walkers1. See Appendix A.1 for more details of the environments.
We compare our method, SWAT, with 2 baselines, the GNN-based method SMP (Huang et al., 2020) and the morphology-free self-attention method AMORPHEUS (Kurin et al., 2021). We use
1Unlike other environments whose variants are in fact subsets of the intact morphology, Animal-Walkers consists of fundamentally different morphologies, such as a Walker, Horse, and Ostrich.
TD3 (Fujimoto et al., 2018) for training the policy over both baselines and ours for fairness. The policy is trained jointly over all morphologies in various environments, and we run every experiment for 3 random seeds to report the mean and the standard error. Details of the implementation can be found in Appendix A.2.

heading: 5.2 MTRL RESULT
text: Our results are summarized in Figure 4. As Kurin et al. (2021) pointed out, SMP, which exploits the morphological information but suffers the over-smoothing problem, shows the worst performance among all the methods. On the other hand, AMORPHEUS that discards the morphological information shows better performance than SMP owing to direct communication via self-attention mechanism.
SWAT clearly outperforms the previous state-of-the-art, AMORPHEUS, especially in the crossdomain environments, where different types of morphologies are mixed. Furthermore, the performance gap between SWAT and AMORPHEUS is notably larger in Animal-Walkers, which has the greatest morphological diversity. Another noteworthy observation is that SWAT achieves the best mean performance as well as converges faster than other baselines in both in-domain and cross-domain environments. These results consistently demonstrate that the effectiveness of our embedding method increases as more various types of robots are jointly trained. We speculate the reason for the effectiveness is that the PE and RE encourage our model to transfer commonalities among resemblant partial morphologies (e.g., the gait patterns of the lower bodies of the Walker and Humanoid), while discriminating against unrelated ones (the legs of Hopper from those of Walker). We shall take a close look at this process in subsection 5.5. By utilizing these structural embeddings across all environments, our model can learn faster and achieve higher final returns than those without the structural embeddings.
0.0 0.2 0.4 0.6 0.8 1.0 Total environment steps (1e6)
0
1000
2000 3000 Av er ag e re tu rn
Humanoid++ (Test set)
0.0 0.2 0.4 0.6 0.8 1.0 Total environment steps (1e6)
0
1000
2000
3000
4000
5000
CWHH++ (Test set)
SWAT SMP AMORPHEUS SWAT w/o transfer
Figure 5: Performance in transfer learning. 0 2 4 6 8 10 Total environment steps (1e6)
0
2000
4000
6000
Av er
ag e
re tu
rn
Animal-Walkers
SWAT SWAT\{RE} SWAT\{PE} AMORPHEUS
0 2 4 6 8 10 Total environment steps (1e6)
0
1000
2000
3000
CWHH++
Figure 6: Ablation study on PE and RE.

heading: 5.3 TRANSFER LEARNING RESULT
text: In this section, we benchmark SWAT in a transfer learning setting where the policy is trained in various tasks and then transferred to another downstream task, which is a common learning strategy in CV and NLP. We compare SWAT with SMP, AMORPHEUS, and SWAT w/o transfer, i.e., SWAT without pretraining. All models except for SWAT w/o transfer are pretrained in the train environments and transferred to the test environment with morphology unseen to them. We evaluate them on Humanoid++, which has the largest number of limbs among the in-domain, and on CWHH++, which contains all in-domain environments. Note that we only train them for 1 million time steps during transfer learning.
As shown in Figure 5, with all pretrained models outperforming SWAT w/o transfer, SWAT learns considerably faster than other baselines, showing much higher sample efficiency, and again outperforms them in average returns. Similar to multi-task learning settings, the useful knowledge obtained from the training tasks can be effectively transferred to the unseen tasks by means of the positional and relational information, enabling quick adaptation with fewer samples. Although our main focus is MTRL and transfer learning, we also conduct experiments on zero-shot setting in Appendix A.4.

heading: 5.4 ABLATION STUDY
text: In this section, we conduct additional experiments to analyze how the positional and relational information affects the performance of our model respectively. We ablate the PE and the RE, and both of them from SWAT, rendering our model identical to AMORPHEUS. We evaluate these on Animal-Walkers and CWHH++, where we expect for their morphological diversity to incite each embedding to play a more crucial role than in other environments.
As shown in Figure 6, the performance of SWAT degrades when the embeddings are removed one by one. This result supports that our proposed embeddings indeed play important roles in different manners. In other words, we need both positional and relational information to take full advantage of the given morphology. Meanwhile, as it can be seen from SWAT\{RE}, i.e., SWAT with the positional embedding only, versus SWAT\{PE}, i.e., SWAT with the relational embedding only, the performance gains from the positional information are greater than ones from relational information. We conjecture that the self-attention, originally devised to learn the implicit relations among nodes, might capture some relational information but cannot the absolute positional information. Further ablation studies about each component of PE and RE are provided in Appendix A.5

heading: 5.5 BEHAVIOR ANALYSIS
text: In this section, we investigate how the agent that is jointly trained in the various environments acts in a single environment. To examine this, we visualize trajectories of the agent trained in CWHH++. Figure 7(a) compares the mean performance of AMORPHEUS and SWAT in the single environment of intact Humanoid with the largest number of 9 limbs.
As shown in Figure 7(b), AMORPHEUS does not learn to walk but instead jumps to move forward like Hopper. This result shows the unexpected effect of multi-task learning. In multi-task learning, the knowledge from other tasks is not always beneficial to the task of interest; therefore, careless exploitation of such knowledge can hamper learning rather than help it. That is, the knowledge obtained in Hopper environments hinders learning to walk in the case of AMORPHEUS. On the
other hand, SWAT successfully learns to run since the structural embeddings allow it to transfer the appropriate knowledge exclusively where necessary. We speculate this is possible as the embeddings which have learned similar limbs across different agents are mapped closely to one another. We visualize the structural embeddings in Appendix A.3.

heading: 6 RELATED WORK
text: There have been lots of studies about MTRL with the homogeneous setting where the state and action spaces are the same across tasks but each task has a different objective (Rusu et al., 2016; Parisotto et al., 2016; Teh et al., 2017; Kalashnikov et al., 2021). The inhomogeneous setting has emerged in recent years and explored by several works. Devin et al. (2017) decompose the policy model into task-specific modules and morphology-specific ones, then reuse the combination of them for various settings. D’Eramo et al. (2020) use a shared module over all tasks that is combined with a task-specific encoder and decoder to learn a policy per task. Those approaches based on the separated module have a limitation in that they require a larger number of model parameters as the number of tasks increases. Chen et al. (2018) encodes agent morphology in a feature vector and learns the policy for different morphologies conditioned on the state and feature vector.
Another line of work that is closely related to ours is the modular approach. Wang et al. (2018) and Pathak et al. (2019) represent the agent morphology as a graph and utilize GNNs as the policy network in order to tackle the inhomogeneous setting. Both of their studies show that the GNN-based policy has large benefits over a monolithic policy in inhomogeneous MTRL. Recently, Huang et al. (2020) present SMP based on GNNs for the inhomogeneous setting. They assume each agent as a tree and propose a both-way MP scheme where the messages are propagated from leaves to root (bottom-up) and from root to leaves (top-down). Kurin et al. (2021) pose a concern about the MP in GNNs due to the over-smoothing problem, and propose AMORPHEUS, which adopts self-attention mechanism rather than GNNs for direct communication in exchange for leveraging the morphological information. In contrast, we utilize the morphological information through the structural embeddings, which can be simply incorporated into the transformer, allowing direct communication while taking full advantage of the structural inductive bias.

heading: 7 CONCLUSION
text: In this paper, we argue that the knowledge of agent morphology plays an important role in modular MTRL and present Structure-aWAre Transformer (SWAT) that leverages structural biases arising from the morphology. Instead of the MP scheme, we encode the structural information through the medium of the traversal-based PE and the graph-based RE, which can be easily incorporated into the transformer-based policy. We empirically demonstrate that our method achieves the state-of-the-art performance in inhomogeneous robot locomotion MTRL benchmarks as well as in transfer learning settings.

heading: A APPENDIX
text: A.1 ENVIRONMENT DESCRIPTION
Following Kurin et al. (2021), we conduct experiments on benchmarks as in Table 2. Huang et al. (2020) proposed Hopper++, Walker++, Cheetah++, and Humanoid++, removing some parts from the intact morphology, while every morphology is able to hop, walk, or run, i.e., they do not contain the morphology that is unable to go forward such as the robot with only arms or the one lacking its torso.
To explain the name of environment in more detail, the first part means its intact morphology and the middle part means the number of limbs belongs to its morphology, and the last part means the missing limbs. For example, humanoid 2d 7 right leg originates from Humanoid and consists of 7 limbs, lacking right legs, both right thigh and right shin. Note that we dub Walker in Wang et al. (2018) as Animal-Walkers, because it is quite confusing in that Walker++ also contains the variants of Walker. Since Animal-Walkers consists of various animal-shaped morphology, we distinguish it from Walker.
All the environments provides the morphology of the robot as a graph-structured data that is extracted from the corresponding MuJoCo XML file. The reward is given by the distance the agent move forward, penalized by the norm of action, i.e., the sum of torque values in all actuators. And it terminates when the agent slips and cannot proceed anymore or it succeeds to survive 1,000 time steps.
A.2 IMPLEMENTATION DETAIL
We implement SWAT based on AMORPHEUS which is built on TransformerEncoder from PyTorch, sharing the codebase of SMP. Additionally, we simply modify TransformerEncoder to incorporate with PE and RE, enabling relational embedding to be added per head. We use 3 independent embedding layers for each traversal to learn PE, where each embedding size has the equal size, i.e., summation of all PEs per traversal is equal to the embedding size of final PE. And we use a simple fully-connected layer for gφ in Equation 4 to learn the graph-based RE. SWAT shares hyperparameter setting used in Kurin et al. (2021) as shown in Table 1, except the replay buffer size due to the memory efficiency. Also, we use the residual connection as reported in Kurin et al. (2021).
Environment Train set Test set Hopper++
hopper 3 hopper 4 hopper 5
Walker++ walker 2 main walker 3 main walker 4 main walker 6 main walker 5 main walker 7 main
Cheetah++ cheetah 2 back cheetah 3 balanced cheetah 2 front cheetah 5 back cheetah 3 back cheetah 6 front cheetah 3 front cheetah 4 allback cheetah 4 allfront cheetah 4 back cheetah 4 front cheetah 5 balanced cheetah 5 front cheetah 6 back cheetah 7 full
Humanoid++ humanoid 2d 7 left arm humanoid 2d 7 left leg humanoid 2d 7 lower arms humanoid 2d 8 right knee humanoid 2d 7 right arm humanoid 2d 7 right leg humanoid 2d 8 left knee humanoid 2d 9 full
A.3 STRUCTURAL EMBEDDING
We visualize both PE and RE in SWAT learned in cross-domain environments, CWHH++. Figure 8 shows that the PEs of limbs placed at similar position in morphology are mapped closely. Figure 9, 10, and 11 shows SWAT learns relational information among nodes through Rlap,Rspd and Rppr, which considers not only local but global information of graph.
A.4 ZERO-SHOT EVALUATION
Table 3 and 4 show zero-shot evaluation results on in-domain setting and cross-domain setting respectively. For in-domain setting, each policy is trained on Walker++, Humanoid++, Cheetah++ and evaluated on their test set respectively. For example, SWAT trained on the train set of Walker++ is tested on walker 3 main and walker 6 main. Note that we omit the lacking part in the name of morphologies in the Table 3 and 4. For cross-domain setting, each policy is trained on CWHH++ and evaluated on its test set, i.e., the union of test sets of Walker++, Humanoid++, Cheetah++. We evaluate by the average performance and the standard error over 3 seeds, where each seed is evaluated on 100 rollouts. As in Kurin et al. (2021), the high variance due to instability in generalization can be observed in SWAT as well as in SMP and AMORPHEUS. Therefore, it would be hard to say which one is statistically advantageous to the others in zero-shot generalization task.
A.5 FURTHER ABLATION STUDY ON EMBEDDING COMPONENTS
In this section, we further conduct experiments to analyze how each component in the positional and relational embedding affects the performance of our model. We firstly verify the usefulness of components in the PE, which are pre-order (PRE), in-order (IN) and post-order (POST) traversal. As shown in Figure 12, SWAT\{POST,IN} shows the worst performance due to ambiguity, i.e., a binary tree cannot be reconstructed only with pre-order traversal. Using in-order traversal along with pre-order traversal is theoretically sufficient to recover a binary tree, so giving an additional post-order seems to be redundant. However, as it can be seen from SWAT versus SWAT\{POST}, the empirical result shows that the additional information actually helps learning more efficiently in some cases. This is because although redundant, they can provide positional information in different perspectives, making the model easier to figure out node positions. We thus decided to use all three traversals in the main experiments. Similar to the PE, we also check the effect of each component in the RE, which are PPR, SPD and LAP, shown in Figure 13. While the performance gaps among 3 variants are not large, SWAT consistently outperformed others. We thus used all three features for RE in the main experiments although they provide overlapping information.

Review summaries:
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).

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.  

This work introduces a structure-aware transformer for inhomogeneous multi-task reinforcement learning tasks. This work proposes to use traversal-based positional embedding and graph-based relational embedding to encode morphological information. The papers show that their proposed approach outperforms prior state-of-the-art methods on the module multi-task RL benchmarks and transfer learning settings.

This paper proposes investigates the role of positional and relational embeddings in transformer-based policies for multi-task reinforcement learning (MTRL) across different morphologies. Unlike prior works that completely discard the structural information of the robot in transformer-based policy, this paper proposes to leverage structural information using traversal-based positional embedding and graph-based relational embeddings. Experiments are performed for MTRL as well as transfer learning on several gym environments, which demonstrate the effectiveness of the proposed method against the baselines.