Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 5: marginally below the acceptance threshold
review_content: 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.
main_review: Strengths:

s1) This work introduces traversal-based positional embedding and graph-based relational embedding to encode the morphological information.

s2) The performance of the algorithm significantly outperforms the prior state-of-the-art methods on some of the tasks.

s3) The paper is well-organized and easy to follow.

Weaknesses:

The 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_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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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. 

main_review: Strength:  
- The paper is well written and easy to read. 
- The proposed approach is straightforward. 
- The experiments show that the proposed approach outperformed competitive baselines such as SMP and AMORPHEUS.     

Weakness:   
- 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.     

- The reviewer has some questions about the technical contents. Please see the comments below:   
   - 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?   
   - 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?      

- Experiments  
   - 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.  

  - 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? 

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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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.
main_review: **Strength:**

- 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.
- 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.
- The proposed method outperforms the SOTA in both MTRL and transfer learning.
- The paper is well-written and easy to follow. The figures (e.g., Fig. 3) are nicely done and convey the main idea clearly.

**Weakness:**

- 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.
- 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?
- 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.

**Comments:**

Typo:

- "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_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.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: 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).

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).
main_review: # Strengths

- Significant performance increase from previous state-of-the-art in multi-task scenarios with diverse morphologies.
- A relevant perspective shift from SOTA in inhomogenous MTRL: an appropriate structural encoding can improve performance than capturing no structural information in the policy.
- Well-written with the relevant related work.
- Qualitative analysis of what the structural embeddings capture.

# Weakness

- 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.
- 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. 
- [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_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. 

While 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.
review_title: null

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)