Structure-Aware Transformer Policy for Inhomogeneous Multi-Task Reinforcement Learning
Authors: Sunghoon Hong
Deunsol Yoon
Kee-Eung Kim
Venue: ICLR 2022 Poster (2022)
Keywords: Multitask Reinforcement Learning
Modular Reinforcement Learning
Transfer Learning
Transformer
Structural Embedding

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 multi-hop 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.

Published as a conference paper at ICLR 2022
STRUCTURE-AWARE TRANSFORMER POLICY FOR INHOMOGENEOUS MULTI-TASK REINFORCEMENT LEARNING
Sunghoon Hongl.3, Deunsol Yoon 1.3 Kee-Eung Kim
"iim Jaechul Graduate School AI, KAIST, Daejeon, Republic of Korea Schooo of Computing, KAIST, Daejeon, Republic of Korea LG Al Research, Seoul, Republic of Korea (shhong01, solpino, kekim] ac..iistta...... ac.kr
ABSTRACT
Modular Reinforcement Learning, where the agent is assumed to be morpholog- ically structured graph, for example composed of limbs and joints, aims learn a policy that transferable to a structurally similar but different agent. Com- pared 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 recent work has shown that their multi- hop message passing mechanism is not ideal for conveying important information to other modules and thus a transformer model without morphological informa- tion was proposed. In this work, we argue that the morphological information is still very useful and propose transformer policy model that effectively encodes such information. Specifically, we encode the morphological information terms of the traversal- based positional embedding and the graph- based relational em- bedding. 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.
INTRODUCTION
Deep reinforcement learning (RL) has made remarkable successes over the last several years, achiev- ing human-leww level performance in various tasks (Silver et al. 2016 Mnih et al. 2013). However, these are limited by individually training policies for each task, often requiring large amount of interac- tion 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 com- monalities and differences across tasks, has become prevalent in computer vision (CV) Dosovitskiy al. 202I) and natural language processing (NLP) (Devlin al. 2019 Radford et al. 2019).
In this regard, Multi-Task Reinforcement Learning (MTRL) is a promising approach, but training policy for multiple tasks in the traditional manner is not straightforward in many cases, e.g., 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 al.
2016 Parisotto et al. 2016 Pinto & Gupta 2017 Yang et al. 2020 Kalashnikov et al. 2021 I) and define each task by its own reward function, e.g., grab a cup or move with a robot arm. Instead, are interested in more general setting where the tasks inhomogeneous, i.c. of different state and action space dimensions, also known as incompatible control (Kurin et al. (022].
One of the popular approaches to inhomogeneous MTRL is to assume graph structure for the agent, depicted in [Figure I] as an example, where limbs and joints are represented as nodes and edges. In this modular setting, Graph Neural Networks (GNNs) provide natural choice for the model of the policy, since (1) they can process graph-structured input of arbitrary sizes and connections, which allows us obtain single policy model that can control any agent with various morphology leading

---

Published as a conference paper at ICLR 2022
Inhomogeneous MDPs M
M1
Shared Policy
Environment
M2
Torque
Node
Limb
Actustor
Timestep
Timestep +++
Figure I: Illustration of modular robot locomotion task. Each node is corresponded to limb, and each non-torso limb has an actuator execute action, which controls the torque of the joint connecting two limbs.
to different state and action space dimensions, and (2) they natively exploit the inductive bias that arises from the morphological structure through their message- message-passing (MP) scheme, i.e., messages 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 GNN prone to over- smoothing (Li et al.j 2018). where the crucial information is washed out during the multi-hop communication. They argued that the advantage of leveraging morphological information for in- homogeneous 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, is well known that injecting the structural inductive bias into the self-attention mechanism, i.e., positional embedding, can help significantly improve the per- formance (Vaswani et al. 2017). the same context, the states of neighbor nodes may be more important than those of nen-neigh nodes for determining actions in modular RI.
this paper, we introduce Structure-aWAre Transformer (SWAT), new modular method for in- homogeneous MTRL that effectively incorporates the agent morphology into the transformer-based policy. 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 po- sitions 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 commu- nication of messages among nodes while taking into account the agent morphology.
The experimental results benchmark tasks for inhomogeneous MTRL strongly support the ef- fectiveness of our structural embedding approach, outperforming prior state-of-the-art methods in terms of sample efficiency and final performance. Through transfer learning experiments, 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 tuple (8, A, T, To.. r, to represent an RL task, where S is the state space, A is the action space, T(3 [st, az) is the state transition probability, Pu is the initial state distribution, rt = r(st, az) E is the reward associated with state and action, and y E (0, 1) is the discount factor. The objective of RL is to learn probability distribution over actions conditioned on states, i.e., policy π (s) which maximizes the expected discounted return R ElE,20 7/re)-
2

---

Published as a conference paper at ICLR 2022
In the inhomogencous MTRL, the agent is tasked with set of inhomogeneous MDPs M {M1, ... ,MK' ). We say MDPs are inhomogeneous when any pair of MDPs have different dimensionalities in their state or action spaces, i.c., (Sm(S,) # dim(S,) or dim((A,) # dim( where M,, Mj E M and S, A, is the state and action space of M,. In the context of MTRI, the goal is to find policy that maximizes the average expected discounted return over all the environments, K EK B[R,] where R, denotes the expected discounted return in MDP M,.
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 1 joints as in Figure Each node (vertices) v² E V for E 1...n represents a limb, and an undirected edge eth E & represents connectivity by a joint between a v² and v². For the ease of notation, we assume an ordered set nodes, i.c., node v' is i-th element in the ordered set. & can be represented {0. 1) 1)nxn that describes the connectivity among nodes, i.e., if o², v² are connected, A = 1 and Eig A 21. Thus at each time step t, the agent observes the state St, which consists of local sensory information of the limb such limb type, coordinates, and angular velocity. Denoting the local sensory information of the limb v' at time step as vector vi. the state 8; can be represented as (v), vi :....... vi } with morphology E. Given 8t, the policy outputs the action a, {at 2 ..... af where a is the torque value for the corresponding actuator controlled by the limb v².
adjacency matrix A E
2.2 POLICY MODELS FOR INHOMOGENEOUS MTRL.
Traditional policy model for RL, c.g., multi-layer perceptron (MLP) policy, not suitable for inho- mogeneous MTRL since it can only handle the state and action spaces with the same dimension. In this regard, GNNs (Hamilton et al. 2017 Wu 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 -aaggggggainnggggg (MP) scheme.
Let us assume graph G = (V, E) where each node v² E V has the corresponding information, i.e.,
node representation or message v². The messages GNNs are delivered follows:
1. MESSAGE AGGREGATION, m +- a(v' : Ni)), Vo' EV
(1)
2. NODE UPDATE: Vipdaate fo(v', m²), Vo EV,
where Nf is the indices of the neighborhood of v², is message aggregation function, e.g., average or concatenation, fo is parameterized node update function that computes a new message given the original message vf and aggregated message m'. 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 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 root and leaves, showing significantly better performance than a standard monolithic policy 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 v² E dRd is projected to three vectors, q' k' E Rdx V E Rd", where q' = v v'Wa q, k' v'W v = ''W, are a query, key, value vectors respectively, and Wq. Wk E Rd W, E Rd. xd, are learnable projection mappings. Messages are then delivered in the self-attention mechanism as:
9'k
exp(a*>)
w
v updatte + Σ Ewisss
ViL
ES. exp(a"""))
where ad is the unnormalized attention score, and wis 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.
3

---

Published as a conference paper at ICLR 2022
2.3 POSITIONAL INFORMATION IN TRANSFORMERS
In the transformer, injecting the positional information is essential since the self-attention mecha- nism is permutation invariant (Vaswani et al.] 2017). In the case NLP, a permutation of tokens change the whole context and meaning of 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 2/ in the sequence is given by the absolute PE vector p and it is added to the corresponding token embedding v'. On the other hand, the relative positional information between two tokens, e.g., relative distance (i (1-)), j). given by relational embedding (RE), r R. which is added the attention score as a q'k /V +r (Raffel et al. 2020 Shaw et al. 2018).
MOTIVATION FOR STRUCTURAL BIAS
AMORPHEUS (Kurin et al., 2021 proposes a plain transformer
model without morphological information due to the aforementioned over-smog problem. However, we can still incorporate the mor- phological information into the transformer, i.e., the structural embed- ding widely utilized in NLP to add the structural bias.
Animal-Walkers
5000
Average return
4000
30000
In order to appreciate its effect in MTRL, we apply 2 common methods used in NLP models, the learnable embedding (LEARN) and the sinu- soidal encoding (SINU). LEARN uses a learnable vector to represent the absolute position of each token in 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 traver- sal of the tree with a torso labeled as 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 An animal. -Wa 1 kers, which are trained jointly for 5 morphologically different agents.
2000
AMORPHEUS
1000
LEARN
SINU
2 Total environment steps (1a6)
1
Figure 2: Training curve An.imm-aaaalee..
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 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 graph, but SINU regards them differently. Although this is a preliminary observation, manifests the necessity of pertinent embedding scheme for representing both the positional and the relational information of morphology for inhomogeneous MTRL.
4 METHOD
Motivated by the result in the previous section, we present the structural embedding scheme that can be effectively incorporated into the transformer-basee 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 RF.
4.1 TRAVERSAL-BASEE POSITIONAL EMBEDDING
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 stresttrut robot morphology defining a torso as the root node, as done in SMP (Huang al. 2020), we can represent it as combination of multiple sequences in terms of tree traversals. We can traverse any tree by several consistent orderings, e.g., pre-order post-order traversal. Although single traversal sequence, as in section 3] is ambiguous
4

---

Published as a conference paper at ICLR 2022
to reconstruct the given tree, uniquely identify binary tree with in-order traversal along with another traversal (Burgdorfi et al. 1987).
Inspired by this property, we introduce traversal- 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, apply left-child-right-sillnng representation (LCRS) represent general tree as a binary tree. Then, we traverse the binarized tree by pre-order, in-order, and post-order, which is sufficient reconstruct the original tree. A node position can be represented as tuple that consists of the position in each traversal. Concretely, Peres Pins Proet E denote the indices of node of in each traversal ordering respectively, and we assign the PE p² E Rdv to the node of using the learnable embedding vectors P P Pppat as follows:
p' = AGGREGATE Prre Pins Ppoat))
+'++'p'
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.
4.2 TRAPH-BASED RELATIONAL EMBEDDING
In 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 introduce graph-based RE by utiliz- ing 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 graph, defined == D A where A is an adjacency matrix and D is a diagonal degree matrix, i.e., D graph Laplacian to bound its norm, which is defined as Ryap - LLD Then, Rtop RR represents the relation among neighbors, capturing local information.
[i.]. We use the normalized
Shortest Path Distance Shortest path distance (SPD) measures the distance between two nodes in graph, which can be easily computed by several algorithms such breadth- first search (BFS). It is similar to the relative positional encoding approach in NLP domains, but we use real-value distance Rapd E Roxn divided by the number of nodes n to bound in [o, 1]. i.e., R.j = SPD(i, j)/n where SPD(i,,) denotes the shortest path distance between of, vú. 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 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 node visitation probability distribution of a e-discounted random walk model with restarting node. To be concrete, the random walker travels across the Markov chain formed via the graph with its transition probability matrix P E R=*. where p 1/|NV| for Vv E Ag otherwise 0. The random walker warps to the restarting node with probability and follows the P with e. Then, the node visitation probability distribution conditioned on the restarting node vf. PPR(i), can be obtained by solving recursive form:
(1 (1-6)x x) + el,
PPR())= -1,-))-c
where x' denotes the probability that the walker resides on v2, and 1, denotes the one-hot encoded restarting node v²_ Then, the proximity between v' and vن, Rig ppr. can be denoted by the j-th entry of Pi)( contrast to SPD, PPR considers every possible path between two nodes, whereby captures the relation among nodes, without the overall graph dispensed with.
5

---

Published as a conference paper at ICLR 2022
1) Encode
2) Message Passing
3) Decode
Node
+ Ineer
MatMMll
Linear
Linear
Softmas
Linser
Score
Mer Morphologr
Positonnl Embedding
Embedding
Figure 3: The overview of SWAT.
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, aggregate them and learn the graph-based RE R BaxnnRReed from them:
R = g ( R. Lap spd-
R
i. ppr
(4)
where Jo is function parameterized by Φ that maps the graph features RE vector R and Whead denotes the number of heads. Note that we represent the relational information between two nodes as vector R E R instead of 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:
Ins dr. Theaad
1.j
K. R
(5)
(h)
(A)
where the subscript (h) denotes the h-th head and R is a h-th entry of R
(h)
4.3 STRUCTURE-AWARE TRANSFORMER POLICY
We now present how SWAT works in inhomogeneous control tasks, illustrated in Figure 3 SWAT is based on 3-stage framework similar to the existing methods: 1) Encode local sensory information y for each node with the traversal-based PE into the hidden representation by single MLP layer shared across all nodes.
use the same notation v for the encoded node representation for clear notation. 2) Update each node representation by MP. While SMP uses the multi-hop hopti-hop communica- tion 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 RI. algorithm, we use TD3 (Fujimoto et al. 2018). a deterministic policy gradient algorithm in the actor-critic framework, following both SMP and AMORPHEUS.
EXPERIMENT
5.1 EXPERIMENT SETTING
We run experiments on modular MTRL benchmarks (Huang et al.- 2020, Wang et al.) 2018), which are created based on Gym MuloCo locomotion tasks. 9 environments can be categorized into 2 settings, in-domain and cross-domain. For in-domain, there are 4 en- vironments: (1) Hopper+t. (2) Walkert+, (3) Cheetah-t, 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 in- domain environments: (1) HalkkerHumanoidtt. (2) Wa Walker-Humance -eannedd-Hopee. d-

---

Published as a conference paper at ICLR 2022
Hop Happeree
Walkeres
Choetahe-
5000
3000
6000
25:0
2000
1000
SWAT
1600
SMP
1000
AMGRPHEUS
steas (1
20
0
Teol crwforrmet sieps
3
20
$2.00
Humanold+
Valler Humanald
50.00
4000
30:00
AMORPHEUS
2000
1000
muttmnmen deps (1)0)
ThiM enuforment Sips (1e5)
Th andirrment Steps (led)
Chestah-WolerrHurrrmet-
Animal-Wakkers
3000
3020
2500
25:00
5030
2000
6010
1530
1000
500
SMP
530
AMGRPHEUS
T3aa enviteeeet stese (1et)
10
Trral envirrmmnt (2e6)
10
Tabal enxtromment cteps (1ed]
10
Figure 4: Training curves on 9 environments. We evaluate on 5 different seeds and plot the mean of average returns over all morphologies. The shaded area represents the standard error.
(3)
(4) Molerra--d----e-----eeeoooddddd (CWHH++), and (5) malimmlllallaaa See Appendix A.1 for more details.
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. (021). We use TD3 (Fujimoto al. 2018) for training the policy over both baselines and ours for fairness. The policy is trained jointly over all morphologies various environments, and we run every experiment for 5 random seeds to report the mean and the standard error.
5.2 MTRL RESULT
Our results are summarized in Figure 4 the morphological information but suffers the svor-shog problem, shows the worst perfor- mance 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.
As Kurin et al. (2021) pointed out, SMP, which exploits
SWAT clearly outperforms the previous state-of-the-art, AMORPHEUS, especially in the cross- domain environments, where different types of morphologies are mixed. Furthermore, the perfor- mance gap between SWAT and AMORPHEUS is notably larger in Animal Wa kers, 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 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 commonalitics
Unlike other environments whose variants in fact subsets of the intact morphology, Animal -wal kers consists of fundamentally different morphologies, such Walker, Horse, and Ostrich.

---

Published as a conference paper at ICLR 2022
Hu m + fTest set)
CWHH++ (Test set)
Animal- Walkers
CWHi+++
5000
3000
3000
1000
Averge retum
5000
2000
2000
3000
8000
21000
2000
SONTI/RE)
1080
1000
AMORPHELS
500TT/P))
Gi0 we moneter
0
ARIORMEUUS
0.2
1:0
52 0.4
3.6
088
Tata envronment
enviromment
Figure 5: Performance in transfer learning.
Figure 6: Ablation study PE and RE
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). By utilizing these structural embeddings across all environments, our model can learn faster and achieve higher final returns than those without the structural embeddings.
5.3 TRANSFER LEARNING RESULT
In this section, we benchmark SWAT in transfer learning setting where the policy is trained in var- ious 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 the train environ- ments and transferred to the test environment with morphology unseen to them. We evaluate them on Humanoidi +++ 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 outper- forms them in average returns. Similar 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
5.4 ABLATION STUDY
In this section, conduct additional experiments to analyze how the positional and relational in- formation 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-Walen kers 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, 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 provided in Appendix A.5
it can be seen from SWAT (RE), i.e., SWAT with the
5.5 BEHAVIOR ANALYSIS
In this section, we investigate how the agent that is jointly trained in the various environments acts in single environment. To examine this, visualize trajectories of the agent trained in C (() the mean performance of AMORPHEUS and SWAT in the single environment of intact Humanoid with the largest number of 9 limbs.
8

---

Published as a conference paper at ICLR 2022
Humanoid limbs)
AMORPHEUS
3000
4000
3000
2000
SWAT
SWAT AMORPHEUS
environment 1166) 10
(a)
(b)
Figure 7: (a) AMORPHEUS and SWAT performance. (b) We visualize trajectories from AMOR PHEUS and SWAT on Humanoid with 9 limbs in CWHH++.
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
6 RELATED WORK
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 Tch 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-specife ones, then reuse the combination of them for various settings. D' D'Eramo et al. (2020) use a shared module over all tasks that is combined with task-specific encoder and decoder learn policy per task. Those approaches based the separated module have limitation in that they require a larger number of model parameters as the number of tasks increases. [Cen 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 closely related to ours is the modular approach. Wang al. (0018) Pathak et al. (2019) represent the agent morphology network in order to tackle the inhomogeneous setting. Both of their studies show that the GNN- based policy has large benefits over monolithic policy in inhomogeneous MTRL. Recently, Huang jet al. (2020) present SMP based on GNNs for the inhomogeneous setting. They assume each agent as a tree and propose 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. (202I) 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 the morphological information. In contrast, utilize the morphological information through the structural embeddings, allowing direct communication while taking advantage of the structural bias.
graph and utilize GNNs as the policy
exchange for leveraging
7 CONCLUSION
In this paper, we argue that the knowledge of agent morphology plays an important role modular MTRL and present 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 em- pirically demonstrate that our method achieves the tfate-of-the-- art performance in inhomogeneous robot locomotion MTRL benchmarks as well in transfer learning settings.

---

Published as a conference paper at ICLR 2022
ACKNOWLEDGMENTS
The authors thank Hyolim Kang for contributing the improvement of SWAT. This work was supported by the National Research Foundation (NRF) of Korea TTTF=0019RAAC1087660 NRF- 202 102IMBIAAON 1009933))) and Institute of Information communications Technology Planning Eval- uation (ITPP) grant funded by the Korea government (MSIT) (No.2019-0-000)) No.2020-0-00400 No. 20.1 -0002068)
10

---

Published as a conference paper at ICLR 2022
REFERENCES
Aleksandar Bojchevski, Johannes Klicpera, Bryan Perozzi, Amol Kapoor, Martin Blais, Benedek Rózemberczki, Michal Lukasik, and Stephan Günnemann. Scaling graph neural networks with approximate pagerank. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2464 2473, 2020.
HA Burgdorff, Sushil Jajodia, Frederick N. Springsteel, and Yechezkel Zalestein. Alternative meth- ods for the reconstruction of trees from their traversals. BIT Numerical Mathematics, 27(2):
133-140, 1987.
Tao Chen, Adithyavairavan Murali, and Abhinav Gupta. Hardware conditioned policies for multi- robot transfer learning. In Advances Neural Information Processing Systems 31: Annual Con- ference Neural Information Processing Systems 2018, NeurlPS 2018, pp. 9355-9366, 2018.
Carlo D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Sharing knowl- edge in multi-task deep reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020. OpenReview.net, 2020.
Coline Devin, Abhishek Gupta, Trevor Darrell, Picter Abbeel, and Sergey Levine. Learning mod- ular neural network policies for multi-task and multi-robot transfer. In 2017 IEEE Interna- tional Conference on Robotics and Automation, ICRA 2017, pp. 2169-2176. IEEE, 2017. doi: N0.1109//CRA2221.7799222
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-IILT 2019, pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021. OpenReview.net, 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International Conference Machine Learning, pp. 1587-1596. PMLR, 2018.
William I. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Proceedings of the 31st International Conference on Neural Information Processing Systems, pp. 1025-1035, 2017.
Wenlong Huang, Igor Mordatch, and Deepak Pathak. One policy to control them all: Shared modular policies for agent-agnostic control. In International Conference on Machine Learning, pp. 4455- 4464. PMLR, 2020.
Dmitry Kalashnikov, Jacob Varley, Yevgen Chebotar, Benjamin Swanson, Rico Jonschkowski, Chelsea Finn, Sergey Levine, and Karol Hausman. Mt-opt: Continuous multi-task robotic re- inforcement learning scale. CoRR, abs/2 2021.
Guolin Ke, Di He. and Tie- Yan Liu. Rethinking positional encoding language pre-training. In 9th International Conference on Learning Representations, ICLR 2021, 2021.
Johannes Klicpera, Aleksandar Bojchevski, and Stephan Günnemann. Combining neural networks with personalized pagerank for classification on graphs. In 7th International Conference on Learming Representations, ICLR 2019, 2019.
Vitaly Kurin, Maximilian Igl, Tim Rocktäschel, Wendelin Bochmer, and Shimon Whiteson. My body is cage: the role of morphology in graph-based incompatible control. In 9th International Conference on Learning Representations, ICLR 2021, 2021.
Qimai Li, Zhichao Han, and Xino-Ming Wu. Deeper insights into graph convolutional networks for semi-supervised learning. In Thirfy-Second AAAI conference artificial intelligence, 2018.

---

Published as a conference paper at ICLR 2022
Russell Merris. Laplacian matrices of graphs: a survey. Linear Algebra and its Applications, 197- 198:143-176, 1994. ISSN 0024-3795, doi: httpss//oo.oo///00000060044777/1
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wier- stra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529 533, 2015.
Lawrence Page, Sergey Brin, Rajeev Motwani, and Terry Winograd. The pagerank citation ranking: Bringing order to the web. Technical report, Stanford InfoLab, 1999.
Emilio Parisotto, Lei Jimmy Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. In 4th International Conference on Learning Representations,
ICLR 2016,, 2016.
Sungchan Park. Wonseok Lee, Byeongseo Choe, and Sang-Goo Lee. A survey on personal- ized pagerank computation algorithms. IEEE Access, 7:163049-163062, 2019. doi: 10.1109/
Deepak Pathak, Christopher Lu, Trevor Darrell, Phillip Isola, and Alexei A. Efros. Learning to control self-assembling morphologies: A study of generalization via modularity. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Process- ing Systems 2019, NeurlPS 2019, pp. 2292-2302, 2019.
Lerrel Pinto and Abhinav Gupta. Learning push by grasping: Using multiple tasks for effective learning. In 2017 IEEE International Conference on Robotics and Automation, ICRA 2017, pp- 2161-2168. IEEE, 2017. doi: 10. 1109/CCAA.201..9992222
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21:1-67, 2020.
Andrei A. Rusu, Sergio Gomez Colmenarejo, Çaglar Gülçehre, Guillaume Desjardins, James Kirk- patrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distil- lation. In 4th International Conference on Learning Representations, ICLR 2016, 2016.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representa- tions. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers), pp-
464-468, 2018
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484-489, 2016.
Yee Whye Teh, Victor Bapst, Wojeiech M. Czamecki, John Quan, James Kirkpatrick. Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Ad- vances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, pp. 4496-4506, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances Neural Infor- mation Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, pp. 5998-6008, 2017.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In 6th International Conference on Learning Representations, ICLR 2018, 2018.
12

---

Published as a conference paper at ICLR 2022
Yu-An Wang and Yun-Nung Chen. What do position embeddings learn? empirical study of pre- trained language model positional encoding. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 6840-6844 2020.
Zonghan Wu, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. JEEE transactions on neural networks and learning systems, 32(1):4 4-14):4422 24, 2020.
Ruihan Yang. Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, 2020.
13

---

Published as a conference paper at ICLR 2022
A APPENDIX
A.1 ENVIRONMENT DESCRIPTION
Following Kurin et al. 2021
conduct experiments on benchmarks as Table Huang et al. 2020) proposed Hopper + +, Wa Walkertt, Cheetah++, and Humanoid++ d+++ removing some parts from the intact morphology, while every morphology is able 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 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, numanoid.2a..7riin
originates from Humanoid and consists of 7 limbs, lacking right legs, both right thigh and right shin. Note that we dub Walker Maamal-aaakeee
in Wang et al. (2018) also contains the variants of Walker. Since Animal -alaaaaaakeesssee consists of various animal-shaped morphology, distinguish it from Walker.
because it is quite confusing in that Walker++
All the environments provides the morphology of the robot as graph-structured data that is ex- tracted from the corresponding MuJoCo XML. file. The reward is given by the distance the agent move forward, penalized by the norm of action, i.c., the sum of torque values in all actuators. And terminates when the agent slips and cannot proceed anymore or it succeeds survive 1,000 time steps.
A.2 IMPLEMENTATION DETAIL.
We implement SWAT based on AMORPHEUS which built on TransformerEncoder from Py- Torch, sharing the codebase of SMP. Additionally, simply modify TransformerEncoderr to incor- porate with PE and RE, enabling relational embedding to added per head. We use 3 independent embedding layers for each traversal 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 Os in Equation to learn the graph-based RE. SWAT shares hyperparam- eter setting used in Kurin et al. (2021) as shown in Table I| except the replay buffer size due to the memory efficiency. Also, we use the residual connection reported in Kurin et al. (2021).
Hyperparemeter
Value
Learning rate
0.0001
Gradient clipping
0.1
Normalization
LayerNorm
Attention layers
3
Attention heads
2
Attention hidden size
256
Encoder output size
128
Mini-batch size
100
Replay buffer size
500K
Embedding size
128
Table Hyperparameter setting in SWAT
14

---

Published as a conference paper at ICLR 2022
Environment
Train set
Test set
#opper++
hopper_3
hopper_4
hopper.5
Wa Walker+t
walker.2.mall
walker-3.ma main walker_6.main
walker_4
Balker_4_maake
walker.5_mas
walker.7.main
Cheetaht+
cheetah 2 back cheetah f.front cheetah.3.back cheetah.3.front cheetah.4.allback
cheetah 3 balanced cheetah back cheetah_6_frons
cheetah.4.allfront
cheetah_4.back
cheetah 4 front
cheetah 5 ba balanced cheetah.5.front cheetah.6.back cheetah.7.full
Humanoidt+
humanoid wd_.d_222___. a.d.7.eef.... arm humanoid d [dw.rlower arms
humancid 2d_7_left leg
humanoid 2d..8.... 1.rrrrrgttttrr. knee
humanoi w_t__ooddddd__riie___ ght_amm
Numanoidd22..7.iighhe.
humamccid.2.._eeft..eee
bumanoid.2d.9.ful]
(WH++)
Union of Walker-+ and Humanoidt+.
Nalker-Burranodd-Hoppe-- (WHH++)
Haleaah--WWlleeee-uaaaaeiaoo H and Hoppert+.
Union Wa HumanoidH,,
(CWH++)
Union of H,teeehhht, Walkertt, and Humanoiát+.
Cheetah-Wet karalkee-Huaaaae -
(CWHR++)
Union of Cheetah++, Walker++ Humanoidt+, and Hopper+
Ani Animal-Walken Weers
FullCheetah (Wolf)
Hopper HalfCheetah (Horse)
HalfHumanoid
Ostrich
Table 2: Full list of environments.
15

---

Published as a conference paper at ICLR 2022
A.3 STRUCTURAL EMBEDDING
We visualize both PE and RE in SWAT learned in cross-domain environments, +W Figure shows that the PEs of limbs placed similar position morphology are mapped closely. Figure 9
To and IT shows SWAT learns relational information among nodes through Rap: R. Repd and R ppr which considers not only local but global information of graph.
Humanoid
+ Walker
+
(alker/ffoot
X Cheetah
X
(footah\fool)
Hopper
torso
thigh1
shin1
foot1
thigh2
shin2
foot2
legs
arms
Figure 8: t-SNE visualization of positional embeddings learned from multiple robots with various morphology. The points are corresponded to the reduced positional embedding of each node in different morphologies. The shape denotes the morphology it belongs to and the color denotes its limb type. Walker(Cheetah)), (foot) are the variants lacking of feet.
16

---

Published as a conference paper at ICLR 2022
-02
lorso
-04
thigh
lower leg
0.4
leg
leg
0.8
thigh
lower
leg
tonsa
foo!
Kirsa
thigh
leg
loser leg
foot
tirso
thih
leg
lower lerg
foot
Figure 9: Illustration of the relational embeddings from two heads on Hopper (5 limbs)
torso
zhin2
-0.3
thigh
thigih
thigh2
-0.4
-0.6
00.6
shin2
-1.0
thigh1
-0.8
foot
foot
-1.0
torso
tright
shint
fooz1
thign2
thinz
terso
thight
think
foot)
thigh
thin2
foor2
Figure 10: Illustration of the relational embeddings from two heads on Walker limbs)
-0.3
upper
Upper
upper erm2
-044
lower armi
lower
Tower
upper armi
-05
forso
shin2
-066
thigh2
thigh
thigh
shinl
-0.0
thigh1
-0.8
shin
shin
torso
-1.
torao biiant
shin
upan
lower arm
lower arma upper arim2
tonso
tight
tuight
shin
opaer armi
upser armn lover bmnt
lower armz
arm!
Figure 11: Illustration of the relational embeddings from two heads on Humanoid (9 limbs)
A.4 ZERO-SHOT EVALUATION
Table 3 and show zero-shot evaluation results on in-domain setting and cross-domain set- ting respectively. For in-domain setting, each policy trained on Walker+t, Humanoid+4, Cheetah++ and evaluated on their test set respectively. For example, SWAT trained on the train set of Walkert+ tested on wa 1err ma in and wa ker imain. Note that we omit the lacking part in the name of morphologies in the Table 3 [nd For cross-domain setting, each pol- icy trained on + and evaluated on its test set, i.c., the union test sets of Wa 1 kertt. Humanoid i++ Ch 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. (021). the high variance due to instability in generalization can be observed in SWAT as well as in SMP and AMORPHEUS. Therefore, would be hard to say which one statistically advantageous to the others in zero-shot generalization task.
17

---

Published as a conference paper at ICLR 2022
SWAT
AMORPHEUS
SMP
walker_3
220.49 66.62
264.04 = 37.39
218.07 49.27
walker_6
396.10 + 136.98
745.28 + 635.61
900.64 172.95
humanoid.2d.7
2428.38 = 1252.61
3042.23 = 1154.78
bumanoid.2d.8
1L17.417111
3624.95 + 1035.22
3798.49 = 137.23
cheetah_3
10.94
- 224.54
171.07
230.04
139.09 268.49
cheetah.5
1011.95 = 945.82
3027.78 735.40
2768.24 + 229.80
cheetah. ah
2610.97 + 747.48
4989.3 + 257.68
5861.86 + 307.39
Table 3: Comparison zero-shot evaluation on in-domain test set.
SWAT
AMORPHEUS
SMP
walker_3
256.62 + 133.57
216.73 40.24
52.30 = 24.39
walker_6
1564.45 = 1229.78
3209.67 + 1425.86
852.41 = 363.50
humanoid_2d_7
1916.45 + 1358.3
2590.07 + 905.40
532.02 + 314.94
humanoid.2d..
4040.86 = 296.85
1201.67 + 565.64
1205.34 + 683.62
cheetah_3
-89.33 41.44
50.13 + 199.91
133.56 68.88
cheetah.5
855.30 + 243.15
1842.01 + 40.70
290.19 = 52.70
cheetah_6
2625.68 + 291.62
3113.09 + 262.48
2778.24 + 391.44
Table Comparison zero-shot evaluation on cross-domain test set.
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 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 Figure 12 SWATT(POST.IN
shows the worst performance due ambiguity, i.e., binary tree cannot be reconstructed only with pre-order traversal. Using in-order traversal along with pre-order traversal theoretically sufficient to recover 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 perspectives, making the model casier to figure out node positions. We thus decided to use three traversals in the main experiments. Similar to the PE, we also check the effect of cach component in the RE, which are PPR. SPD and LAP, shown in Figure 13] While the performance gaps among 3 variants not large, SWAT consistently outperformed others. We thus used all three features for RE in the main experiments although they provide overlapping information.
provide positional information in different
18

---

Published as a conference paper at ICLR 2022
WH+++
WH++++
CWH++
5000
4000
3000
Average return
4000
3000
3000
2000
1000
SWART
SWATXPOST)
1000
(WAT(PPOST.IN)
10
2
Total anvironment steps (1e6)
(1o6)
Total anvironment stops (1e6)
Total environment stops (1of)
Animal-Walkers
3000
6000
Averagge return
1000
2000
SWATYPPOSTS
SWATT(POSTJIN)
8 Total environment steps (1e6)
10
Total environment steps (1e6)
Figure 12: Ablation study on craversa--based positional embedding.
WH++
WHH++
CWH+
5000
4000
3000
average raturn
4000
3000
3000
2000
1000
SWATLELAP)
1000
SWPP)
10
2
10
Total environment steps (1e6)
(1e6)
Total environment steps (1e6)
Total environment-steep (1e8)
CWHH
Animal-Walkers
Average return
4000
1000
SWAT)(LAP)
SWATI(LAP. SPO)
0:
8
10
Total environment steps (1e6)
Total environment steps (1e6)
(leps(1e6)
Figure 13: Ablation study on graph-based relational embedding.
19


Reviews:
[{"review_id": "ujCx-5Lhf_f", "reviewer": "Reviewer_4eRD", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\ns1) This work introduces traversal-based positional embedding and graph-based relational embedding to encode the morphological information.\n\ns2) The performance of the algorithm significantly outperforms the prior state-of-the-art methods on some of the tasks.\n\ns3) The paper is well-organized and easy to follow.\n\nWeaknesses:\n\nThe major weakness of this paper is the experimental evaluation. For figure 4, only three different seeds are used. I would recommend 10 seeds. ", "summary_of_the_paper": "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.", "summary_of_the_review": "This work proposes to use traversal-based positional embedding and graph-based relational embedding to encode morphological information for inhomogeneous multi-task reinforcement learning tasks. The proposed approach seems reasonable. However, the technical novelty is limited. Experiments with more seeds are needed to strengthen the paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oCU0SRGbqlY", "reviewer": "Reviewer_xDqR", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:  \n- The paper is well written and easy to read. \n- The proposed approach is straightforward. \n- The experiments show that the proposed approach outperformed competitive baselines such as SMP and AMORPHEUS.     \n\nWeakness:   \n- The reviewer has some concerns about the novelty of the paper. Positional encoding for transformers has been widely used and recognized in NLP and computer vision literature. The paper’s main contribution is showing that transformers with positional encoding outperforms the ones without the positional encoding in MTRL tasks. The reviewer thinks it is reasonable but the contribution seems somewhat incremental.     \n\n- The reviewer has some questions about the technical contents. Please see the comments below:   \n   - To reconstruct a binary tree, inorder traversal along with preorder or postorder traversal is sufficient. However, in Eq (3) all the three traversal results are aggregated together. What is the purpose of adding redundant information?   \n   - In inhomogeneous MTRL, each task has different state dimensions. However, in Section 4.3, it says a single MLP layer is shared across all nodes. Is some padding mechanism used here to make a single MLP layer applicable to all nodes?      \n\n- Experiments  \n   - For relational encoding, the authors combine three features: (1) normalized graph Laplacian (2) shortest path distance, (3) page rank. It would be great to check the effectiveness of each of the features vai an ablation study.  \n\n  - In Section A2, the authors mentioned that they used a different replay buffer size from the baseline. For a fair comparison, why not using the same size? \n", "summary_of_the_paper": "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. \n", "summary_of_the_review": "The reviewer thinks this paper is well written and the proposed positional encoding is simple yet seems effective. However, the reviewer has some concerns about the novelty of the paper, since positional encoding is widely used along with transformers. \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HHQ2azm2l8O", "reviewer": "Reviewer_BRC4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strengths\n\n- Significant performance increase from previous state-of-the-art in multi-task scenarios with diverse morphologies.\n- A relevant perspective shift from SOTA in inhomogenous MTRL: an appropriate structural encoding can improve performance than capturing no structural information in the policy.\n- Well-written with the relevant related work.\n- Qualitative analysis of what the structural embeddings capture.\n\n# Weakness\n\n- While the paper shows ablation experiments by excluding PE and RE from SWAT, it does not discuss the effect of different tree-traversals in PE and graph connectivity representations in RE. It seems that there is redundant information in the proposed PE and RE representations. For example, is there any extra information provided by Normalized Laplacian and SPD that PPR cannot capture? What if only PPR is used to represent the graph-based RE? Further ablation analysis would be required to answer if  embedding that include different and overlapping information are necessary for generalization and transfer.\n- Is SWAT is computationally more expensive than the baseline approaches? SWAT requires computing additional PE and RE embeddings compared to AMORPHEUS and uses of quadratic attention compared to SMP. What would be the wall-clock time for taking an \"environment step\" in each of the three models? What is the overhead computation required for ε-discounted random walk for PPR embedding? A discussion in this regard would be helpful for reproducibility and future research directions. \n- [Clarification needed] In Transfer learning, the train set of environments is a bit unclear for Humanoid++ and CWHH++ test tasks. Is it consisting of all the 9 train environments except the test environment as shown in Appendix A.2 Table 2? Has the policy for test setting of CWHH++ seen the morphologies of Cheetah/Walker/Humanoid/Hopper in the in-domain tasks? Is only difference in train and test setting is an unseen state and action space size?", "summary_of_the_paper": "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).\n\nThe 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.\nTwo 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).", "summary_of_the_review": "I enjoyed the proposed simple additions to existing state-of-the-art, showing how certain body morphology encodings can significantly improve the performance. While this work again shows that structural information can be helpful in MTRL, it brings up the question of the appropriate ways to encode the structure for inhomogeneous MTRL for successful policy learning. \n\nWhile there seems no proven \"optimal\" way for encoding structural information, the proposed PE and RE features in transformer-based policy seem an excellent first step.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "7L-8z9IUZEF", "reviewer": "Reviewer_bVEF", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strength:**\n\n- The paper is well-motivated in that structural and morphological information about the robot could help improve the performance in MTRL and transfer learning, as this information can make the transformer-based policy more specialized for each morphology.\n- The proposed traversal-based embedding seems new and achieves good performance as shown in the ablation studies. The overall method seems quite easy to implement and has the potential to improve other GNN-based tasks.\n- The proposed method outperforms the SOTA in both MTRL and transfer learning.\n- The paper is well-written and easy to follow. The figures (e.g., Fig. 3) are nicely done and convey the main idea clearly.\n\n**Weakness:**\n\n- Base on its design, I feel the proposed traversal-based positional embedding is not very suitable for generalization. The index of each joint is quite global and not actually transferrable among different morphologies. For example, a foot joint with index i in robot A may have index j in robot B, but index j may represent a completely different joint (e.g., head) in robot A, so the learned embedding p_j (learned for the head joint) may not be suitable for the foot joint in robot B. I think the proposed method can still achieve good performance in MTRL because the training in MTRL allows the policy to overfit to the learnable positional embeddings. But the proposed embedding has limited generalization ability so it is not suitable for tasks like zero-shot policy transfer which has been demonstrated in [Huang et al., 2020]. I think some experiments on zero-shot policy transfer are required to shed some light on this, or the paper should discuss the limitation in generalization.\n- The paper spent a fair amount of space on relational embedding (RE), however, the performance gain from it seems a bit small. Most performance gain seems to come from the positional embedding (PE). As discussed in my first point, is it because the PEs already allows the network to overfit to them in MTRL?\n- Ablation studies on all the environments should be provided, especially because the proposed RE seems not very effective in Fig. 6. Additionally, both the PE and RE have multiple components. It would be nice to do some ablation studies in those components to understand their usefulness.\n\n**Comments:**\n\nTypo:\n\n- \"the performance gains from the positional information(SWAT\\RE) are greater than ones from relational information (SWAT\\PE)\" —-  switch RE and PE?", "summary_of_the_paper": "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.", "summary_of_the_review": "Overall, this is a well-motivated paper with good performance. The proposed traversal-based PE seems quite useful for MTRL task. However, the proposed RE seems less useful and the ablation studies are not quite sufficient. I also have doubts in terms of the PE's generalization ability, which could be addressed through zero-shot policy transfer experiments. I vote for a weak accept for now, but may change my score based on the authors' response.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the role of positional and relational embedding s for multi-task reinforcement learning with transformer-based policies, The paper is well-motivated, the experiment shows its effectiveness against other competitive methods. In the rebuttal period, the authors solved most of the reviews’ questions such as novelty and ablation studies. There are still some concerns about the generalizability of this approach for other tasks and more experiments are needed.
Decision: Accept (Poster)