MetaMorph: Learning Universal Controllers with Transformers

Abstract: Multiple domains like vision, natural language, and audio are witnessing tremendous progress by leveraging Transformers for large scale pre-training followed by task specific fine tuning. In contrast, in robotics we primarily train a single robot for a single task. However, modular robot systems now allow for the flexible combination of general-purpose building blocks into task optimized morphologies. However, given the exponentially large number of possible robot morphologies, training a controller for each new design is impractical. In this work, we propose MetaMorph, a Transformer based approach to learn a universal controller over a modular robot design space. MetaMorph is based on the insight that robot morphology is just another modality on which we can condition the output of a Transformer. Through extensive experiments we demonstrate that large scale pre-training on a variety of robot morphologies results in policies with combinatorial generalization capabilities, including zero shot generalization to unseen robot morphologies. We further demonstrate that our pre-trained policy can be used for sample-efficient transfer to completely new robot morphologies and tasks.

Reviews and author responses:

 Thanks for clarifying my doubts and addressing my concerns. I will therefore keep my (already very good!) score.

The paper presents MetaMorph (MM), a transformer-based architecture for generalization over robot morphologies. Extensive experimentation demonstrates that policies learned by MM transfer well to new robot morphologies. This is beautiful work! The approach is well-motivated, and the experiments are thorough, covering a wide diversity of robots and environments, and demonstrating remarkable generalization. On the whole, the paper is exceptionally well-written and informative.

The main weakness I see is a failure to adequately distinguish the proposed approach (MM) from the very similar AMORPHEUS approach (AM) by Kurin et al., 2021. The current work claims that MM conditions the transformer on morphological information, while AM does not. However, it seems to me that MM and AM condition the transformer on essentially the same information. The confusion seems to arise because the two papers use the term *morphology* in two separate ways. This requires careful explanation.

Each robot consists of one body and some number of limbs branching from it. (This paper calls the body and limbs *modules*, but I’ll call them all limbs for ease of visualization.) The information about the robot’s limbs can be divided into 3 categories:  limb-to-limb connectivity, static limb state (like limb type and its joint’s range of motion), and dynamic limb state (like instantaneous angular position and velocity). For processing by the network, standard practice is to combine the state of each joint (both static and dynamic) with the state of the limb that it swings.

AM (Kurin et al., 2021) is a modification of SMP (Huang et al., 2020). In all three architectures (MM, AM & SMP), observations include both static and dynamic state information of joints and limbs, but not explicit limb-to-limb connectivity information. SMP uses limb connectivity to define the computational graph over which messages are passed, from limb-to-limb. By contrast, MM and AM use a transformer encoder which passes messages between every pair of limbs, whether directly connected to each other in the robot or not. Kurin et al., 2021, demonstrate that AM’s all-to-all message passing performs better than SMP’s restricted message passing.

Now to the terminology. Kurin et al., 2021, uses the term *morphology* to refer to the *connectivity* of the limbs. But *this* paper uses the term *morphology* to refer to the *static state* of limbs (or their joints). This seems to explain why the paper states that MM differs from AM in “conditioning the Transformer on morphological information”, as if AM didn’t. This is incorrect, because both AM and MM condition the transformer on both static and dynamic state, and neither one conditions the transformer on connectivity information. If AM leaves out some important piece of static information, or if there is some other essential difference between AM and MM, the paper in its current form does not make this difference clear.

I agree that the current work differs from that of Kurin et al., 2021, “in the diversity and scale of training robots, complexity of the environments, … and showcasing strong generalization to unseen morphologies and tasks”. But I don’t see an important architectural difference between the models. (As the paper states, “the position embeddings carry no information about the graph structure”, so limb-to-limb connectivity is not carried by MM’s position embeddings.) AM is clearly the most meaningful baseline model to compare MM against, but this was not done, despite the availability of the AM code.

The experimental results show that MM’s position embeddings boost performance, but I find no explanation for why they do. Could it just be that they provide unique IDs to help the transformer differentiate the modules from each other, similar to a scheme employed by SMP?
 The paper presents compelling experimental results, but does not adequately differentiate MetaMorph from AMORPHEUS by Kurin et al., 2021, and does not include AMORPHEUS as a baseline in the experiments.

--- POST-DISCUSSION UPDATE ---

The authors have clarified the distinction between their work and that of Kurin et al., 2021, and have performed new experiments comparing MetaMorph to an ablated version (MetaMorph-AO) that approximates the more restricted morphological observations used by Amorpheus. 

Although MetaMorph and Amorpheus use differing sets of input features, they are still fundamentally the same architecture. (The "morphology aware transformer" used by MetaMorph is just a standard transformer encoder, like that of Amorpheus.) So in my view, there is little novelty in the Amorpheus architecture, and the work's value lies in the scope of its experiments. I have raised my evaluation accordingly.

 Thank you, this latest draft adequately clarifies the difference in how the two works use the term.

 We thank the reviewer for the suggestions. We now explicitly mention terminology differences (see L432-435). In addition, based on our terminology we also explicitly mention input observations of Amorpheus (L496-500) and highlight that both methods don’t use explicit connectivity information (see L500-503). 

 I acknowledge the authors' response and tend to increase the score. However, I would like to have a discussion with other reviewers first.

 I appreciate the authors' responses and their revisions to the paper. I have revised my scores accordingly. 

This paper presents MetaMorph, a Transformer based universal controller to learn behaviors across different robot morphologies. Their Transformer module takes a sequence of tokens as input, corresponding to the number of modules in the robot. Each input token comprises the proprioceptive and morphology information for each constituent robot’s module. MetaMorph is trained using a standard model-free Proximal Policy Optimization (PPO) method. However, to ensure all robots get to explore given environments and gather a similar amount of experience, the paper introduces a dynamic replay buffer which prioritizes robots’ selection for experience collection based on their previous iteration's episode lengths. The results are compared against MLP based experts trained for all individual robots, Graph Neural Network-based universal controller, and ablated modules of MetaMorph to highlight design choices. Moreover, results also demonstrate the generalization capacity of MetaMorph to different robots through zero-shot transfer and fine-tuning. 
 Strengths:

Overall the paper is well written. It proposes an interesting way to use Transformers for learning a universal controller for different robot morphologies. The idea of a dynamic replay buffer is simple but exhibits improvement in learned behaviors. The results are thorough and cover most aspects of their approach and exhibit improved performance over prior methods. The ablation studies are also well-suited to describe the design choices. 

Weakness:
The contributions on algorithmic and network designs are marginal. Their approach is a direct application of Transformers and Reinforcement Learning approaches. 

Clarification: The description of section 4.3 can be improved, it is slightly confusing to follow. It seems Eq 5 will give a higher probability to robots with longer episode lengths which is in contrast to their description. Please elaborate or clarify if I missed something.  
 It proposes an interesting way to use Transformers for learning a universal controller for different robot morphologies and presents an extensive set of experiments. However, the overall contributions to algorithms and network designs are limited.


 Thank you for confirming that. It’s one more commonality between MetaMorph and Amorpheus that should be pointed out for readers. 

Your recent explanations here regarding terminology would provide clarity that is currently missing from the paper. Except that it doesn’t seem necessary (and comes across as pedantic) to argue that your definition of morphology is somehow *better* than that of Kurin et al. In my opinion, it’s enough to simply explain the difference in how the two papers use the term.

 No, MetaMorph's observations does not include any explicit limb-connectivity information. 

 ` "connectivity is only one aspect of morphology" `

Do MetaMorph's observations include any *explicit* limb-connectivity information?

 We are glad that the reviewer found our additional experiments helpful. We briefly clarify additional points raised by the reviewer below:
>*Refer Appendix B.1 in related work*

Done. See L86.

>*morphology terminology*

We thank the reviewer for the follow-up discussion. Amorpheus and SMP define the morphology as only the connectivity of the limbs (note that NerveNet does not use the term morphology). However, connectivity is only one aspect of morphology (see [1]), and in general a substantial amount of additional information may be required to adequately describe the morphology. Hence, our design incorporates local relative poses, information about joints types, motor types, gear ratios, density of material used etc. to describe the morphology of an agent holistically. 

[1] Bongard, Josh. "Why morphology matters." The horizons of evolutionary robotics 6 (2014): 125-152.

>*Clarification on contributions*

As we stated in our related work section, both Amorpheus and MetaMorph use Transformer architecture for policy learning. However, the main contribution of Amorpheus is showcasing that message passing in GNN constrained on morphology connectivity might hurt policy learning. In contrast, our core contribution is to showcase at scale that pre-training on large distribution of morphologies using Transformers (conditioned on morphology information) leads to favourable generalization properties in challenging environments for complex 3D morphologies. 


 The new experiments in Appendix B.1 are very helpful, especially the MetaMorph-AO baseline, since MetaMorph is so closely related to Amorpheus. In discussing Kurin et al., the Related Work section should reference this appendix. And the shift in terminology needs to be clearly explained. In Kurin et al., **morphology** refers to “**connectivity** of those limbs, i.e. **morphology**”, which was used by prior works to constrain message passing. But in your paper, *morphology* simply refers to a particular type of observation feature. This difference in definitions obscures the fact that MetaMorph is best seen as an extension of Amorpheus, given that the key contributions of Amorpheus are fully shared by MetaMorph. 

 We are humbled by the reviewers’ positive comments. We are glad that the reviewers found: 

(1) our problem formulation novel and persuasive (reviewer rhAB)

(2) our analysis/experiments to be thorough and convincing (all reviewers)

(3) our idea of dynamic replay buffer to be simple, effective, and possibly applicable to other multi task settings (reviewer M7dm, Hxo2, XpGx)

(4) our generalization results to be remarkable, impressive and at scale (all reviewers) 

(5) our idea of conditioning the transformer output on morphology information sound, interesting and well motivated (reviewer XpGx, rhAB). 

The main aim of this rebuttal is to make this paper stronger by incorporating reviewer suggestions and comments. Specifically, we have implemented **all** additional baselines and ablations as requested by reviewers (please see the newly added sections in the appendix i.e. A.4, B, and C).


 We are glad and humbled by the reviewers' positive feedback. Here we address reviewer suggestions: 
1. We agree that robustness to morphologies in the same training distribution is indeed very valuable. We have added this in the introduction (see L64-66).

2. Done. See L444.

3. We apologize for the confusion. We determine $P_k$ by replacing $\mathcal{E}_i$ with $\frac{1000}{\mathcal{E}_i}$, where the numerator is the maximum episode length. Hence, there is an inverse relationship i.e. longer episode lengths will lead to lower probability of selection. We have fixed this in the manuscript (see L169-171)

4. We speculate that a potential reason for this could be that Obstacle is a challenging environment, and 5 million iterations is perhaps not enough for MLP to learn the optimal policy. However, our goal with MLP baseline was simply to suggest an approximate upper bound, and hence we keep the number of training iterations for the MLP baseline same across all 3 environments. Our result might suggest that joint pre-training can be especially helpful for harder tasks/environments. 

5. We have added the baseline. Please see Appendix B.1 and Figure 10. 

6. We completely agree that further analysis would be interesting to understand why certain limbs are highly correlated. However, our preliminary investigations did not provide us any additional insight, and hence we defer this investigation to future work.

7. We note that we have not claimed in the text that MetaMorph performs few-shot generalization to robots which are out of distribution. We describe this setting as sample efficient transfer i.e. we do see 2-3x reductions in training iterations, and strong zero-shot performance. Nevertheless, we agree that there is further room for improvement and have added a section about limitations and future work (see Appendix C).


 We are glad that the reviewer found our problem formulation novel and persuasive. We address the reviewer concern below:
>*In the 5 Experiments, the model sizes of the other design choices for comparison are smaller than the proposed model, which may lead to unfair comparison.*

For all our ablations and experiments we have ensured that the model sizes are the same to ensure fair comparison. We now explicitly mention this in L473-473.


 We are glad that the reviewer found our work well-motivated, and demonstration of generalization impressive. We will discuss the three main points raised by the reviewer below. Please let us know in case we missed any concern, as we work towards improving our current submission. 
>*Lack of comparison with Amorpheus*

In our revised manuscript, we discuss in Appendix B.1 why direct comparison with Amorpheus is not feasible. For your convenience, we reiterate the same here. 
Direct comparison to Amorpheus is not feasible due to the following reasons: (1) does not work with 3D morphologies with variable number of joints between limbs; (2) does not incorporate exteroceptive observations; (3) Amorpheus ensures balanced collection of experience by sequentially collecting data on each morphology and maintaining a separate replay buffer per morphology. Further, updates to policy parameters are also performed sequentially. This sequential nature of the algorithm is not amenable to scaling, and would require ~30 GPU days to train for 100 million iterations on Nvidia RTX 2080, while our method only needs 1.5 GPU days **($\sim20$x training speedup)**. Hence, we compare with **MetaMorph-AO (**A**morpheus **O**bservation) as the closest variant to Amorpheus, where we provide the same input observations as described in Kurin et al, ICLR 2021.** As shown in Figure 10, MetaMorph outperforms MetaMorph-AO by a substantial margin.

>*Differences in input observations between Amorpheus and MetaMorph*

In our setup, at each time step $t$ (we drop the time subscript for brevity) the robot receives an observation $s^k = (s^k_m, s^k_p, s^k_g)$ which is composed of the morphology representation ($s^k_m$), the proprioceptive states ($s^k_p$) and additional global sensory information ($s^k_g$). For each module $s^k_{p_i}$ consists of the instantaneous state of the system: (1) Module Proprioception: 3D Cartesian position, 4D quaternion orientation, 3D linear velocity, and 3D angular velocity. (2) Joint Proprioception: Position and velocity in the generalized coordinates. **According to this terminology, Amorpheus conditions on a subset of $s^k_{p}$, normalized joint angle ranges, and per limb unique id.**
Finally, for each module $s^k_{m_i}$ consists of: (1) Module Parameters: Module shape parameters (e.g. radius, height), material information (e.g. density), and local geometric orientation and position of the child module with respect to the parent module. (2) Joint Parameters: This consists of information about joint type (e.g. hinge) and its properties (e.g. joint range and axis), and actuator type (e.g. motor) and its properties (e.g. gear). 
**We emphasize that the only morphology information which Amorpheus conditions on is joint angles ranges.** Please see Appendix A.1 where we describe the input observations in detail. 

>*Role of position embedding*

Both Amorpheus and SMP provide a unique ID to differentiate between modules. We conduct an additional ablation experiment where instead of learnt position embeddings, we provide similar unique ids which we call MetaMorph-HPE (**H**and-designed **P**osition **E**ncoding; please see Appendix B.1 and Figure 10). We note that MetaMorph-HPE performs considerably worse than MetaMorph, indicating that position embeddings contain additional useful information than simple IDs. We acknowledge that further investigation is required to understand the nature of the extra information learned by the  position embeddings. Please also see section B.2.


 >*How do you learn positional embeddings? W_{pos} should take a k-th row for the dimensions to make sense, not adding the full matrix.*

We believe that our notation might have led the reviewer to believe that we learn a separate position embedding for each morphology. We use the same position embeddings for all morphologies, and have now fixed the notation in section 4.2. Specifically, instead of using $N_k$ we now refer to all matrix dimensions as $N$ where $N$ is the maximum number of modules for morphologies in the training set. The position embeddings are randomly initialized learnable paprateters and trained from scratch like other network parameters. 

>*How do you obtain a positional embedding matrix for your zero-shot transfer experiments?*

We use the same positional embedding matrix learnt during training for our zero-shot transfer experiments. 

>*How do the position embeddings look like if you plot them as a matrix?*

We found that across all seeds and in all environments (we only show flat terrain for brevity) a diagonal pattern emerged. Please see Figure 11 and discussion in section B.2.

>*Does NerveNet benefit from balancing? When you compare with GNN in Figure 4, does it include the balancing?*

Yes, the GNN baseline does include dynamic replay buffer balancing. In our preliminary experiments, we found benefits of dynamic replay buffer balancing to be agnostic to the architecture i.e. Transformer or GNN. 

>*Amorpheus (Kurin et al, ICLR 2021) shows that zero-shot generalisation exhibits quite a large variance across the seeds, however, your approach does not suffer from this issue, why do you think this is the case?*

We speculate that the large variance in the case of Amorpheus was due to a small number of training examples. We believe that large scale training alleviates this issue.

>*When you study zero-shot/transfer behaviour, … original training runs?*

We found that the training curves were robust to the choice of seeds as indicated by small standard deviation (see Figure 4). Consequently, we evaluate the zero shot performance using a pre-trained model corresponding to a single seed which is the same for all the baselines and ablations. 

For our transfer learning experiments, we fine tune the model with the random seed corresponding to the one used for pre-training (for all 3 seeds). We have also added this discussion in Appendix A.4.

>*DFS order*

We thank the reviewer for suggesting this important experiment and data augmentation strategy. We re-trained MetaMorph, by randomizing the order of tree traversal and found the learnt policy to be robust to changes in the order. Without this data augmentation strategy we found the zero shot performance to drop to ~75% when tested on morphologies with opposite ordering. We have added the experiment, and the discussion in Appendix B.3. We reiterate here for your convenience:

We first process an arbitrary robot by creating a 1D sequence of tokens corresponding to depth first traversal of its kinematic tree. As noted by the reviewer the DFS ordering is not unique, and for nodes at similar depth, we adopted the convention of visiting the nodes which come first in the MuJoCo XML. This convention is also adopted by MuJoCo when parsing the XML. We found that zero-shot transfer of learnt policy was not robust to an opposite ordering of nodes. However, we re-trained the model with the following data augmentation strategy: randomize the order of visiting the nodes during training. We found this simple augmentation strategy made the policy robust to changes in DFS ordering.

>*Are Figure 9 results consistent across different seeds?*

We found them to be consistent across different seeds, however we have not exhaustively validated this across all morphologies. 

>*Missing citation to Snowflake*

Done



 We are glad that the reviewer found our empirical results impressive, appreciated the scale of our effort and acknowledged the generality of dynamic replay buffer balancing. We have addressed all the comments, and we discuss some of the main questions/concerns below:
>*Comparison with Amorpheus. What ablation in Figure 4 is the closest to Amorpheus?*

In our revised manuscript, we discuss in Appendix B.1 why direct comparison with Amorpheus is not feasible. For your convenience, we reiterate the same here.
Direct comparison to Amorpheus is not feasible due to the following reasons: (1) does not work with 3D morphologies with variable number of joints between limbs; (2) does not incorporate exteroceptive observations; (3) Amorpheus ensures balanced collection of experience by sequentially collecting data on each morphology and maintaining a separate replay buffer per morphology. Further, updates to policy parameters are also performed sequentially. This sequential nature of the algorithm is not amenable to scaling, and would require  approximately 30 GPU days to train for 100 million iterations on Nvidia RTX 2080, while our method only needs 1.5 GPU days **($\sim 20$x training speedup)**. Hence, we compare with **MetaMorph-AO (Amorpheus observation) as the closest variant to Amorpheus**, where we provide the same input observations as described in Kurin et al, ICLR 2021. As shown in Figure 10, MetaMorph outperforms MetaMorph-AO by a substantial margin.

>*Dynamic replay balancing will only work if the environments are somewhat similar given the used metric (e.g. similar number of episode steps).*

We agree. However, we note that if the environments are extremely different, joint training might not be feasible at all. In general, when training robots jointly one would ideally train for general purpose skills and avoid training in extremely diverse environments. After general purpose pre-training, we would then fine-tune on a much narrower selection of morphologies for a specific and potentially very different environment. 

>*The relationship to muscle synergies sounds cool and sells well, but I don't think it makes a lot of sense to me in this setting or it requires more explanation.*

We believe that one of the core advantages of Transformer architecture is that visualizing attention weights helps us better interpret how the model is making decisions. In that spirit, there has been increased interest in both NLP [1] and vision [2] communities in showcasing emergent properties by analysing the attention weights. We view our work in the same spirit. 

Muscle synergies are neural strategies to simplify the problem of controlling multiple degrees of freedom. In general, muscle synergy is defined as a pattern of co-activation of muscles by a single neural command. The analogy in transformer-based methods is to demonstrate that certain joints get activated together. This in turn would imply that the rank of the attention matrix would be low. However, in deep learning systems fluctuations due to stochastic gradient descent and numerical imprecision generally cause the matrix to be full-rank. Hence, we use the stable rank of the matrix. For attention matrix $A_l \in \mathbb{R}^{m \times m}$ for layer $l$, the stable rank is defined as: ${\text{sr}}(A_l) = \frac{\lVert A_l \rVert^2_F}{\lVert A_l \rVert_2} = \frac{\sum \sigma_i^2}{\sigma_{max}^2}$, where $\sigma_i$ are the singular values of $A_l$. We note that ${\text{sr}}(A_l)$ is small and oscillates between two values which correspond to attention maps where groups of limbs are activated simultaneously (denoted by dark columns), a characteristic signature of motor synergies. Hence, MetaMorph learns to simplify the control problem by activating different motor synergies depending on both $s^k_m$ and $s^k_p$. Note that this behavior emerges without any hard-coding.

[1] Clark, Kevin, et al. "What does bert look at? an analysis of bert's attention." arXiv preprint arXiv:1906.04341 (2019).
[2] Caron, Mathilde, et al. "Emerging properties in self-supervised vision transformers." arXiv preprint arXiv:2104.14294 (2021).

>*A fairer ablation would be providing a one-hot encoding of a task instead of some meaningful morphological information. Can you run this ablation during the rebuttal period?*

We have now added this ablation as MetaMorph-NMT (**N**o **M**orphology information + **T**ask encoding). Please see Appendix B.1 and Figure 10.

>*Some of the very important details are omitted, i.e. what does $s^k_m$?*

We note that we had provided a detailed description of $s^k_m$ in Appendix A.1.

**See part 2**



 We thank the reviewer for finding our approach interesting and appreciating our exhaustive set of experiments. We address the reviewer concerns below and please let us know if there are further questions:
>*The contributions on algorithmic and network designs are marginal. Their approach is a direct application of Transformers and Reinforcement Learning approaches.*

Our goal in this work was to bring to bear the advances in large-scale pre-training, transfer learning and general purpose Transformer architectures, to the field of robotics. This is a non-trivial endeavor as scaling and training transformers is especially challenging in reinforcement learning settings (see [1]). Hence, in addition to this demonstration, we believe we make the following scientific and empirical contributions: (1) conditioning the transformer on morphology information; (2) balanced replay buffer sampling; (3) showcasing scalability and generalization of pre-trained transformers in the challenging environments with diverse morphologies. 

[1] Parisotto, Emilio, et al. "Stabilizing transformers for reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

>*It seems Eq 5 will give a higher probability to robots with longer episode lengths which is in contrast to their description*

We apologize for the confusion. We determine $P_k$ by replacing $\mathcal{E}_i$ with $\frac{1000}{\mathcal{E}_i}$, where the numerator is the maximum episode length. Hence, there is an inverse relationship i.e. longer episode lengths will lead to lower probability of selection. We have fixed this in the manuscript (see L169-171)


The paper uses transformers for incompatible multitask continuous control and shows impressive empirical results on a wide range of tasks in multitask training, zero-shot generalisation and transfer. The authors propose morphology-based conditioning and dynamic replay buffer balancing to increase the performance. ### Pros

- Strong empirical results for multitask training/zero-shot generalisation and transfer.
- Dynamic replay balancing is an idea that might be useful for multitask RL outside this particular setting.
- The paper is mostly clearly written and has nice graphics.
- Well-written related work section.

### Cons

- I find it confusing that the method name has 'meta' in it?
- Lack of clear comparison with Amorpheus (Kurin et al, ICLR 2021) that also uses transformers for incompatible multitask control.
- Dynamic replay balancing will only work if the environments are somewhat similar given the used metric (e.g. similar number of episode steps).
- Section 5.5 is confusingly written. The relationship to muscle synergies sounds cool and sells well, but I don't think it makes a lot of sense to me in this setting or it requires more explanation.
- I don't think that MetaMorph-NM is a fair ablation that tells how morphological information is useful to training. A fairer ablation would be providing a one-hot encoding of a task instead of some meaningful morphological information. Can you run this ablation during the rebuttal period?
- Some of the very important details are omitted, i.e. what does $s^k_m$ contain? How do you learn positional embeddings? See more in Questions/Comments section.

### Questions/Comments

- Does NerveNet benefit from balancing? When you compare with GNN in Figure 4, does it include the balancing?
- Amorpheus (Kurin et al, ICLR 2021) shows that zero-shot generalisation exhibits quite a large variance across the seeds, however, your approach does not suffer from this issue, why do you think this is the case?
- When you study zero-shot/transfer behaviour, what is the experimental protocol? You run multiple seeds → select a model → run zero-shot evaluation/finetune for transfer. How do you select a model? Do you select a single model or one per each of the original training runs?
- You don't compare to Amorpheus, why? What ablation in Figure 4 is the closest to Amorpheus?
- W_{pos} should take a k-th row for the dimensions to make sense, not adding the full matrix.
- How do you obtain a positional embedding matrix for your zero-shot transfer experiments?
- DFS order depends on the internal representation of your tree, i.e. A→B←C might go to A or C after B depending on the order you store them in. What happens to your method if you do zero-shot generalisation transfer on the mirrored agent (so that DFS returns a different ordering of the same nodes)? Can you provide the results of this experiment?
- How do you learn the positional encodings? How do they look like if you plot them as a matrix?
- What does the morphological information contain, i.e. what is $s_m^k$?
- typo on page 2 'planer' → 'planar'
- Are Figure 9 results consistent across different seeds?
- The related work section is missing a recent paper: "*Snowflake: Scaling GNNs to high-dimensional continuous control via parameter freezing", Blake et al, NeurIPS 2021* This is a good paper with impressive empirical results. While transformers have been applied for incompatible multitask continuous control before, none of the works scaled the setting that much. Dynamic replay balancing proposed in the paper might be useful in multitask RL in general. I would give this paper a 7, but the system does not allow it. I am putting 6 for now and am willing to increase the score if the authors address my questions/comments during the discussion period.

In robotics, researchers primarily train a single robot for a single task, while modular robot systems suggest the flexible combination of general-purpose building blocks into task optimized morphologies. The authors proposed a Transformer approach to learn a universal controller over a modular design space leveraging the “large-scale pre-training and fine-tuning” scheme. Strengths
- The problem formulation (especially 4.2 Morphology Aware Transformer) is novel and quite persuasive.
- The proposed claim is appropriate in that robot morphology can be understood as a modality on which we can condition the output of a Transformer.
- Implementation details are shared so as easy to reproduce.

Weaknesses
- In the 5 Experiments, the model sizes of the other design choices for comparison are smaller than the proposed model, which may lead to unfair comparison. I generally agree with the argument of the paper. As the authors mentioned in the Abstract, Deep learning research fields like vision, natural language, and audio actively leverage Transformers for large-scale pre-training followed by task-specific fine-tuning. The idea makes sense that it deals with robot morphology as a modality on which we can condition the output of a Transformer in the modular robot design setting.

This paper proposes a Transformer-based universal policy to control modular robots. Being the space of modular robots combinatorial, it is (almost) impossible to train specific policy for each sample of the "robot distribution". Therefore, the paper proposes a transformer-based architecture design with explicit information about the robot morphology that, trained on a large number of possible morphologies. This universal controller generalizes zero-shot to unseen robot designs belonging to the training distribution and can be efficiently fine-tuned on new robots and tasks. The approach is thoroughly evaluated on different tasks. Overall, I think this paper is very interesting. In my opinion, its main strengths are:
1. The design changes to the standard transformer architecture are sound and interesting for the problem setup. The main idea, i.e. conditioning the architecture to the robot morphology, is sound and interesting. Also, the dynamic replay buffer (basically storing samples in function of performance) is a simple but very nice idea. None of those are really breakthroughs but they are sound technical contributions that could influence future robotics research.
2. The experimental results are solid (even though some things could be improved, more on this later).
3. Writing is good and easy to follow, with figures helping to understand the main concepts.

The following is a set of weaknesses that I think could improve the contribution. None are major weaknesses or flaws, but just recommendations to improve the current submission:

1. The motivation to study the problem of a universal controller for robot morphologies could be improved. It looks like the main motivation is to "bring the advantages of pre-training into robotics", which sounds more like a methodology than a research problem. I think a more solid connection to real-world problems could be the situation when some parts of the robot fail (a propeller for a drone, a joint for a manipulator, etc.), and the control policy should be robust to these changes. The robustness to morphologies in the same training distribution could be very valuable there I think.
2. Several prior works on legged locomotion ( like R1 or R2) condition the control policy on some environment and robot characteristics. This is in a way similar to the proposed approach of conditioning the policy on morphology, so I think it deserves mention.
3. I did not fully get Eq. 5, which explains the dynamic replay buffering. Assuming $\beta>0$ (which seems okay according to the text),  the longer the episode length, the higher the probability of being sampled. This is in contrast with the motivation of dynamic replay buffering balancing, as explained in the example in section 4.3, where one wants to sample robots with short episode length (thus difficult to train).
4. I find it interesting that in the presence of obstacles, the meta-policy works much better than the per-robot MLP. Why is this the case? Why don't we observe the same in the flat and variable terrain?
5. As baseline, it would be nice to have a naive approach that substitutes the transformer architecture with a simple MLP. This MLP would have the same observations of the transformer and the same number of parameters to make the comparison as fair as possible. This naive baseline could give more intuitions on the advantages of the transformers for the task.
6. In Section 5.5, it would be nice to analysize a specific robot morphology and try to decode the resulting motor synergies. For example, would be nice to answer the question: why is limb N highly correlated to all other parts of the robot?
7. Finally, it would be nice to have a section discussion limitations of the approach. The biggest now seems that the generalization to robot out of distribution (in particular with different kinematics) is limited, and 2x10^7 iterations or more might be required to fine-tune on a new robot (which is not really few-shot). Overall, it would be nice to be open about the limitations to facilitate future work.
 
[R1] RMA: Rapid Motor Adaptation for Legged Robots, Kumar et al.

[R2] Learning Quadrupedal Locomotion over Challenging Terrain, Lee et al. The paper tackles an interesting problem in a sound way introducing some key technical contributions. Addressing some weaknesses in the presentation and in the experimental part, as well as being more open about limitations, could improve the quality of the submission.

Meta-review: The paper aims to improve generalization and sample efficiency in robotic control. In this regard authors note that modular robot systems, which provide building blocks for a task-specific morphology, can be considered just another domain where transformers can be used. The authors propose to learn a universal controller over this modular design space and leverage successful “large-scale pre-training and fine-tuning” scheme for transformer. We thank the reviewers and authors for engaging in an active discussion. Reviewers found that methodological novelty is limited ("The contributions on algorithmic and network designs are marginal. Their approach is a direct application of Transformers and Reinforcement Learning approaches.") and there are very similar approaches in the literature (e.g., AMORPHEUS approach (AM) by Kurin et al., 2021.) Despite these shortcomings, overall the reviewers found the experimental results to be strong and analysis to be thorough, which will be valuable to the community.

Decision: Accept (Poster)