M ETA M ORPH : L EARNING U NIVERSAL C ONTROLLERS WITH T RANSFORMERS
A BSTRACT
Multiple domains like vision, natural language, and audio are witnessing tremendous progress by leveraging Transformers for large scale pre-training followed by task speciﬁc ﬁne tuning.
In contrast, in robotics we primarily train a single robot for a single task.
However, modular robot systems now allow for the ﬂexible 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 pretraining 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-efﬁcient transfer to completely new robot morphologies and tasks.
I NTRODUCTION
The ﬁeld of embodied intelligence posits that intelligent behaviours can be rapidly learned by agents whose morphologies are well adapted to their environment (Brooks, 1991; Pfeifer & Scheier, 2001;
Bongard, 2014; Gupta et al., 2021).
Based on this insight, a robot designer is faced with a predicament: should the robot design be task speciﬁc or general?
However, the sample inefﬁciency of tabula rasa deep reinforcement learning and the challenge of designing a single robot which can perform a wide variety of tasks has led to the current dominant paradigm of ‘one robot one task’.
In stark contrast, domains like vision (Girshick et al., 2014;
He et al., 2020) and language (Dai & Le, 2015; Radford et al., 2018), which are not plagued by the challenges of physical embodiment, have witnessed tremendous progress especially by leveraging large scale pre-training followed by transfer learning to many tasks through limited task-speciﬁc ﬁne-tuning.
Moreover, multiple domains are witnessing a conﬂuence, with domain speciﬁc architectures being replaced by Transformers (Vaswani et al., 2017), a general-purpose architecture with no domain-speciﬁc inductive biases.
How can we bring to bear the advances in large-scale pre-training, transfer learning and generalpurpose Transformer architectures, to the ﬁeld of robotics?
We believe that modular robot systems provide a natural opportunity by affording the ﬂexibility of combining a small set of general-purpose building blocks into a task-optimized morphology.
Indeed, modularity at the level of hardware is a motif which is extensively utilized by evolution in biological systems (Hartwell et al., 1999; Kashtan & Alon, 2005) and by humans in many modern engineered systems.
However, prior works (Wang et al., 2018; Chen et al., 2018; Sanchez-Gonzalez et al., 2018) on learning policies that can generalize across different robot morphologies have been limited to: (1) manually constructed variations of a single or few base morphologies, i.e. little diversity in the kinematic structure; (2) low complexity of control ( ≤ 7 degrees of freedom); (3) using Graph Neural Networks (Scarselli et al., 2008) based on the assumption that kinematic structure of the robot is the correct inductive bias.
In this work, we take a step towards a more challenging setting (Fig. 1) of learning a universal controller for a modular robot design space which has the following properties: (a) generalization #$%& '()%(*#$#$+ ,"-./*(&0"1"%&2)3#+$&45*6) 7 ) ( " 8 4" %& : )$ ) ( * /# ; * % # "$ ,)%*,"( <#$)=
*%#63>*(#*%#"$ ?
*3@ 2A$*=#63>*(#*%#"$ !"
#!$%"&'(#!$&")!
*(+",+&"$)&-'( .&!"(&$/ 1,2-'(3&4(#,"+5,*,.)&-'(#,6%*&( +!"!
#&$&"- to unseen variations in dynamics (e.g. joint damping, armature, module mass) and kinematics (e.g. degree of freedom, morphology, module shape parameters) and (b) sample-efﬁcient transfer to new morphologies and tasks.
We instantiate the exploration of this general setting in the UNIMAL design space introduced by Gupta et al.
We choose the UNIMAL design space as it contains a challenging ( 15 − 20 DoFs) distribution of robots that can learn locomotion and mobile manipulation in complex stochastic environments.
Learning a single universal controller for a huge variety of robot morphologies is difﬁcult due to: (1) differences in action space, sensory input, morphology, dynamics, etc.
(2) given a modular design space, not all robots are equally adept at learning a task, e.g. some robots might inherently be less sample-efﬁcient (Gupta et al., 2021).
To this end, we propose MetaMorph, a method to learn a universal controller for 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.
MetaMorph tackles the challenge of differences in embodiment by leveraging a Transformer based architecture which takes as input a sequence of tokens corresponding to the number of modules in the robot.
Each input token is created by combining proprioceptive and morphology information at the level of constituent modules.
The combination of proprioceptive and embodiment modalities and large scale joint pre-training leads to policies which exhibit zero-shot generalization to unseen variations in dynamics and kinematics parameters and sample-efﬁcient transfer to new morphologies and tasks.
Finally, to tackle the differences in learning speeds of different robots, we propose dynamic replay buffer balancing to dynamically balance the amount of experience collection for a robot based on its performance.
In sum, our key contributions are: (1) we introduce MetaMorph to learn a universal controller for a modular design space consisting of robots with high control complexity for challenging 3D locomotion tasks in stochastic environments; (2) we showcase that our learned policy is able to zero-shot generalize to unseen variations in dynamics, kinematics, new morphologies and tasks, which is particularly useful in real-world settings where controllers need to be robust to hardware failures; (3) we analyze the learned attention mask and discover the emergence of motor synergies (Bernstein, 1966), which partially explains how MetaMorph is able to control a large number of robots.
R ELATED W ORK
Prior works on learning control policies which can generalize across robot morphologies have primarily focused on parametric variations of a single (Chen et al., 2018) or few ( 2 − 3 ) robot types Wang et al., 2018; Sanchez-Gonzalez et al., 2018; Huang et al., 2020; Kurin et al., 2021).
For generalizing across parametric variations of a single morphology, various approaches have been proposed like using a learned hardware embedding (Chen et al., 2018), meta-learning for policy adaptation (Al-Shedivat et al., 2017; Ghadirzadeh et al., 2021), kinematics randomization (Exarchos et al., 2020), and dynamics randomization (Peng et al., 2018).
In case of multiple different morphologies, one approach to tackle the challenge of differences in action and state spaces is to leverage
Graph Neural Networks (Scarselli et al., 2008; Kipf & Welling, 2017; Battaglia et al., 2018).
Wang et al.
; Huang et al.
(2020) use GNNs to learn joint controllers for planar agents ( ≤ 7 DoFs).
Blake et al.
propose freezing selected parts of networks to enable training GNNs for a single morphology but with higher control complexity.
The usage of GNNs is based on the assumption that the robot morphology is a good inductive bias to incorporate into neural controllers, which can be naturally modelled by GNNs.
Recently, Kurin et al.
(2021) also proposed using Transformers for training planar agents.
Our work differs from Kurin et al.
(2021) in the diversity and scale of training robots, complexity of the environments, conditioning the Transformer on morphological information, and showcasing strong generalization to unseen morphologies and tasks (see § B.1).
Another closely related line of work is the design of modular robot design spaces and developing algorithms for co-optimizing morphology and control (Sims, 1994) within a design space to ﬁnd taskoptimized combinations of controller and robot morphology.
When the control complexity is low, evolutionary strategies have been successfully applied to ﬁnd diverse morphologies in expressive soft robot design spaces (Cheney et al., 2014; 2018).
In the case of rigid bodies, Ha (2019); Schaff et al.
; Liao et al.
(2019) have proposed using RL for ﬁnding optimal module parameters of ﬁxed hand-designed morphology for rigid body robots.
For more expressive design spaces, GNNs have been leveraged to share controller parameters (Wang et al., 2019) across generations or develop novel heuristic search methods for efﬁcient exploration of the design space (Zhao et al., 2020).
In contrast to task speciﬁc morphology optimization, III et al. (2021) propose evolving morphologies without any task or reward speciﬁcation.
Finally, for self reconﬁgurable modular robots (Fukuda &
Nakagawa, 1988; Yim et al., 2007), modular control has been utilized in both real
(Rubenstein et al., 2014; Mathews et al., 2017) and simulated (Pathak et al., 2019) systems.
L EARNING A U NIVERSAL C ONTROLLER
We begin by reviewing the UNIMAL design space and formulating the problem of learning a universal controller for a modular robot design space as a multi-task reinforcement learning problem.
3. T HE UNIMAL D ESIGN S PACE
An agent morphology can be naturally represented as a kinematic tree, or a directed acyclic graph, corresponding to a hierarchy of articulated 3D rigid parts connected via motor actuated hinge joints.
The graph G := ( V , E ) consists of vertices V = { v 1 , ..., v n } corresponding to modules of the design space, and edges e ij ∈ E corresponding to joints between v i and v j .
Concretely, in the UNIMAL (Gupta et al., 2021) design space, each node represents a component which can be one of two types: (1) a sphere parameterized by radius and density to represent the head of the agent and form the root of the tree; (2) cylinders parameterized by length, radius, and density to represent the limbs of the robot.
Two nodes of the graph can be connected via at most two motor-actuated hinge joints (i.e. G is a multi-graph), parameterized by joint axis, joint limits and a motor gear ratio.
3. J OINT P OLICY O PTIMIZATION
The problem of learning a universal controller for a set of K robots drawn from a modular robot design space is a multi-task RL problem.
Speciﬁcally, the control problem for each robot is an inﬁnitehorizon discounted Markov decision process (MDP) represented by a tuple ( S, A, T, R, H, γ ) , where S represents the set of states, A represents the set of available actions, T ( s t +1 | s t , a t ) represents the transition dynamics, R ( s, a ) is a reward function, H is the horizon and γ is the discount factor.
At each time step, the robot k receives an observation s k t , takes an action a k t , and is given a reward r kt .
A policy π θ ( a kt | s kt ) models the conditional distribution over action a kt ∈
A given state s kt ∈ S .
The goal is to ﬁnd policy parameters θ which maximize the average expected return across all tasks: R = 1 K (cid:80)
Kk =0 (cid:80)
Ht
=
γ t r kt .
We use Proximal Policy Optimization (PPO)
(Schulman et al., 2017), a popular policy gradient (Williams, 1992) method for optimizing this objective.
M ETA M ORPH
Progress in model-free reinforcement learning algorithms has made it possible to train locomotion policies for complex high-dimensional agents from scratch, albeit with tedious hyperparamter tuning.
However, this approach is not suitable for modular design spaces containing exponentially many robot morphologies.
Indeed, Gupta et al.
(2021) estimates that the UNIMAL design space contains more than 10 18 robots.
Hence, learning a separate policy for each robot is infeasible.
However, the modular nature of the design space implies that while each robot morphology is unique, it is still constructed from the same set of modules and potentially shares subgraphs of the kinematic tree with other morphologies.
We describe how MetaMorph exploits this structure to meet the challenge of learning a universal controller for different morphologies.
P e r j o i n t c on t r o ll e r ou t pu t
L i nea r P r
o j e c t i on + + + + + + + + +
M o r pho l og y
A w a r e
T r an s f o r m e r
D e c ode r
MLP + Norm
Norm
Multi-HeadAttention + L ×
ModuleEmbeddings
TransformerEncoder
P e r m odu l e t
o k
en
Learnt Position Embeddings
MLP
Exteroceptive Obs.
4. F USING P ROPRIOCEPTIVE S TATES AND M ORPHOLOGY R EPRESENTATIONS
To learn policies that can generalize across morphologies, we must encode not only proprioceptive states essential for controlling a single robot, but also morphological information.
From a multi-task
RL perspective, this information can be viewed as a task identiﬁer, where each task corresponds to a different robot morphology, all drawn from the same modular design space.
Hence, instead of learning a policy which is agnostic to the robot morphology, we need to learn a policy conditioned on the robot morphology.
Consequently, at each time step t (we drop the time subscript for brevity), the robot receives an observation s k = ( s km , s kp , s kg ) which is composed of the morphology representation ( s km ), the proprioceptive states ( s kp ), and additional global sensory information ( s kg ).
See § A.1 for a detailed description of each observation type.
4. M ORPHOLOGY A WARE T RANSFORMER
The robot chooses its action via a stochastic policy π θ ( a kt | s kt ) where θ are the parameters of a pair of deep neural networks: a policy network that produces an action distribution (Fig. 2), and a critic network that predicts discounted future returns.
We use Transformers (Vaswani et al., 2017) to parametrize both policy and critic networks as described in detail below.
Encode.
We make a distinction between how we process local and global state information.
Concretely, let s k =
( s kl , s kg ) where s kl =
( s km , s kp ) .
Since the number of joints between two modules can vary, we zero pad s kl i to ensure that input observation vectors are of the same size, i.e. s kl ∈ R N × M .
In order to provide an arbitrary robot morphology as input to the Transformer, we ﬁrst create a 1D sequence of local observation vectors by traversing the kinematic tree in depth ﬁrst order starting at the root node (torso in case of the UNIMAL design space).
We then apply a single layer MLP independently to s kl i to create a D dimensional module embedding.
We also add learnable 1D position embeddings to the module embeddings to automatically learn positional information:
[Equation]
where φ ( · ) is the embedding function, W e ∈ R M × D are the embedding weights, and W pos ∈
R N × D are the learned positional embeddings.
Note that in practice we zero pad the input sequence of local observation vectors for efﬁcient batch processing of multiple morphologies.
Process.
From the module embeddings described above, we obtain the output feature vectors as:
[Equation]
where MSA is multi-head attention (Vaswani et al., 2017) and LN is Layernorm (Lei Ba et al., 2016).
Decode.
We integrate the global state information s kg consisting of high-dimensional sensory input from camera or depth sensors.
Naively concatenating s kg and s kl in the encoder has two downsides: (1) it dilutes the importance of low-dimensional local sensory and morphological information; (2)
[Figure]
it increases the number of Transformer parameters due to an increase in the dimensionality of the input embedding ( D ).
Instead, we obtain the outputs of policy network as follows:
[Equation]
where γ ( · ) is a 2-layer MLP with parameters W g , φ ( · ) is an embedding function with W d as the embedding weights.
The action distribution is modeled as a Gaussian distribution with a statedependent mean µ ( s ki ) and a ﬁxed diagonal covariance matrix Σ .
Similarly, for the critic network, we estimate value for the whole morphology by averaging the value per limb.
4. D YNAMIC R EPLAY B UFFER B ALANCING
Joint training of diverse morphologies is challenging as different morphologies are adept at performing different tasks.
Consequently, some morphologies might be inherently better suited for the pre-training task.
Let us consider two robots: (A) Robot A locomotes in a falling forward manner, i.e., robot A is not passively stable; (B) Robot B is passively stable.
Especially with early termination, robot A will keep falling during the initial phases of training, which results in shorter episode lengths, whereas robot B will have longer episode lengths.
Hence, more data will be collected for robot B in the earlier phases of training, and in turn will lead to robot B learning even faster, thereby resulting in a ‘ rich gets richer ’ training dynamic.
However, our goal is to ensure that all morphologies have a similar level of performance at the end of training as we want to generalize across the entire distribution of morphologies.
To address this issue, we propose a simple dynamic replay buffer balancing scheme.
On-policy algorithms like PPO (Schulman et al., 2017) proceed in iterations that alternate between experience collection and policy parameter update.
Let E k be any performance metric of choice, e.g. normalized reward, episode length, success ratio, etc.
Let τ be the training iteration number.
At each iteration, we sample the k th robot with probability P k , given by:
[Equation]
where α ∈ [0 , 1] is the discount factor and the exponent β determines the degree of dynamic prioritization, with β = 0 corresponding to the uniform case.
In practice, we use episode length as our performance metric.
We determine P k by replacing E i with 1000 E i , where the numerator is the maximum episode length.
E XPERIMENTS
In this section, we evaluate our method MetaMorph in different environments, perform extensive ablation studies of different design choices, test zero-shot generalization to variations in dynamics and kinematics parameters, and demonstrate sample efﬁcient transfer to new morphologies and tasks.
For qualitative results, please refer to the video on our project website 1 .
[Figure]
5. E XPERIMENTAL S ETUP
We create a training set of 100 robots from the UNIMAL design space (Gupta et al., 2021)
(see § A.2).
We evaluate MetaMorph on three different environments (Fig. 3) using the MuJoCo simulator (Todorov et al., 2012).
In all 3 environments the goal of the agent is to maximize forward displacement over the course of an episode which lasts 1000 timesteps.
The 3 environments are: (1)
Flat terrain (FT);
(2) Variable terrain (VT): VT is an extremely challenging environment as during each episode a new terrain is generated by randomly sampling a sequence of terrain types and interleaving them with ﬂat terrain.
We consider 3 types of terrains: hills, steps, and rubble; (3) Obstacles: cuboid shaped obstacles of varying sizes on ﬂat terrain.
We use a dense morphology independent reward function for all our tasks as it is not feasible to design a reward function tailored to each morphology.
In all tasks, our reward function promotes forward movement using small joint torques (the latter obtained via a small energy usage penalty).
In addition, as described in §4.3, we use early termination across all environments when we detect a fall (i.e. if the torso height drops below 50% (FT, Obstacles) or 30% (VT) of its initial height).
5. B ASELINES AND A BLATIONS
[Figure]
We compare the percentage of robots with ﬁnal performance greater than 75% of the MLP baseline performance when trained jointly.
Across all 3 environments, on average for 10 − 15% robots,
MetaMorph w/o balancing is unable to learn a good control policy.
Baselines : We compare against the following baselines: (1) GNN : We modify the NerveNet model proposed by
Wang et al.
(2018) to learn control policies for complex 3D morphologies with variable number of joints.
Specifically, we replace how the NerveNet model receives input and produces output by our encode and decode steps respectively (§4.2).
In addition, the model receives the same observation as MetaMorph i.e. s k = ( s km , s kp , s kg ) and is trained with dynamic replay buffer sampling.
Thus, the only difference is in the process step.
This helps test if the domain-speciﬁc inductive bias of the robot kinematic tree in the GNN is necessary. (2) MLP : We train all 100 robots separately with a 2-layer
MLP and report the average performance.
This baseline serves as an approximate upper bound for our method.
Ablations : We also do an ablation study of different design choices involved in our method.
We refer our full method as MetaMorph and consider the following ablations: (1) MetaMorph-NPE: no learned position embeddings; (2) MetaMorph-NM: we only provide
[Figure]
s k
= ( s kp , s kg ) as inputs, i.e. the model does not have access to information about the robot morphology.
Fig.
4 shows learning curves across 3 environments for training 100 morphologies.
In all environments MetaMorph can successfully match the average reward achieved via the per morphology MLP baseline on both FT, and obstacle environment.
While MetaMorph performance in VT is slightly below the MLP baseline, we note that it has not saturated and we stopped training at 10 8 iterations across all three environments.
Moreover, MetaMorph is signiﬁcantly more sample-efﬁcient ( 5 × ) than training independent MLPs ( 5 × 10 6 iterations per robot).
The GNN baseline saturates at a level 2 to 3 times below MetaMorph.
In GNN based models, locality and neighborhood connectivity is explicitly baked into the model.
Interestingly, just like ViT (Dosovitskiy et al., 2021) sparingly utilizes the 2D neighborhood structure of images at the beginning of the model by cutting the image into patches, MetaMorph uses the graph structure of the robot in the beginning by creating a 1D sequence corresponding to the kinematic tree by traversing the graph in depth ﬁrst order.
Moreover, the position embeddings carry no information about the graph structure and are learned from scratch.
We highlight that the learned position embeddings signiﬁcantly improve the performance of MetaMorph, just as they do in Transformer based image classiﬁcation.
Finally, without access to the morphological information, MetaMorph-NM fails to learn a policy that can control diverse robot morphologies.
All of this substantiates our central claim that morphological state information is necessary to learn successful control policies, although the kinematic graph need not be explicitly baked into neural architectures to learn policies capable of controlling diverse robot morphologies.
Finally, we test the importance of dynamic replay buffer balancing in Fig.
5, and ﬁnd that balancing is necessary to learn a good control policy in 10 − 15% of robots across all 3 environments.
5. Z ERO -S HOT G ENERALIZATION
Our focus in this work is to learn policies that can generalize to unseen robots drawn from a modular robot design space.
In this section, we demonstrate that MetaMorph shows favorable generalization properties across many different kinematic and dynamic variations.
Experimental Setup.
For each of the 100 training robots, we create a dedicated test set to test zeroshot transfer performance across two types of variations: dynamics (armature, density, damping, and motor gear) and kinematics (module shape parameters like radius, and length of cylindrical limbs, and joint angles).
For each training robot, we randomly create 4 different variants for each property, i.e. 400 robots with armature variations, and so on.
While creating a new variant, we change the
[Figure]
[Figure]
2 − 3 × more sample efﬁciency
[Figure]
relevant property of all modules or joints.
See Table 2 for sampling ranges.
We then compare zero-shot performance averaged over 10 trials
[Figure]
Generalization: Dynamics.
First, we consider generalization to different dynamics (Fig. 6).
We ﬁnd consistently that MetaMorph performs signiﬁcantly better than MetaMorph-NM and GNN across all types of dynamic variations and all environments.
In fact, the difference is more pronounced for harder tasks like VT and Obstacles.
We note that this result is surprising as we do not do dynamics randomization during training, i.e., all robots in the training set have the same armature and damping parameters.
Despite this, we see strong generalization performance.
Generalization: Kinematics.
Next, we consider generalization to different kinematics parameters (Fig. 6).
This is a signiﬁcantly more challenging setting as the model has to generalize to unseen variations in module shape parameters and changes to joint angle ranges.
In fact, changes to joint angles can signiﬁcantly alter the range of possible motion and might necessitate a different gait for successful locomotion.
Consequently, we ﬁnd that even though MetaMorph exhibits strong generalization performance compared to MetaMorph-NM and GNN in all 3 environments, there is indeed a performance drop for the challenging setting of variations in joint angles.
However, the zero-shot performance is encouraging and motivates our next set of experiments on transfer learning.
5. S AMPLE E FFICIENT T RANSFER L EARNING
Experimental Setup.
Here we create 100 robots from the UNIMAL design space which were not part of the training set, and we demonstrate that MetaMorph shows favorable sample efﬁcient transfer to unseen morphologies, and even unseen morphologies performing novel tasks.
Different morphologies.
We ﬁrst consider sample efﬁcient learning of controllers for new morphologies on the same task by taking a pre-trained MetaMorph model on an environment, and then ﬁne tuning it to learn to control morphologies in the test set.
In Fig.
7 we compare the number of training iterations required to achieve a particular performance level when we ﬁne tune MetaMorph
[Figure]
vs training from scratch on the test set.
Across all 3 environments we not only ﬁnd strong zero-shot performance , but also 2 to 3 times higher sample efﬁciency compared to training from scratch.
Different morphologies and novel tasks.
Finally, we consider the most realistic and general setting mimicking the promise of modular robots, where we are faced with a novel task and want to use a new robot morphology which may be suited for this task.
We consider the same set of 100 test robots on two new tasks (Fig.
8): (1) Escape: The agent starts at the center of a bowl shaped terrain surrounded by small hills (bumps), and has to maximize the geodesic distance from the start location (escape the hilly region).
(2) Obstacles (cylinders): cylinder shaped obstacles of varying sizes (the size distribution is also different from the train task i.e. cuboid shapes).
We transfer the learned policy from VT and Obstacles to Escape and Obstacles-Cylinders respectively.
Again, we ﬁnd that there is a strong zero-shot performance across all 3 environments and ﬁne-tuning is 2 to 3 times more sample efﬁcient than training from scratch.
5. E MERGENT M OTOR S YNERGIES
We next search for a potential explanation for how MetaMorph can coordinate the large number of
DoFs ( ∼ 16 × 100 ) across several agents.
Our hypothesis is inspired by Bernstein (1966), which proposed the existence of muscle synergies as a neural strategy to simplify the control of multiple
DoFs by the central nervous system.
A muscle synergy corresponds to the constrained movements of sets of joints (or muscles) through co-activation by a single neural signal.
Such synergies obviate the need to control all joints independently, by coupling sets of joints into adaptive functional groups.
Although the deﬁnition of synergies (Bruton & O’Dwyer, 2018) vary in the literature, dimensionality reduction is generally accepted as a signature of synergistic control (Todorov & Ghahramani, 2004; Todorov, 2004).
Consequently, to test this hypothesis, in Fig.
9 we plot the stable rank of the attention matrix.
For attention matrix A l ∈ R m × m for layer l , the stable rank is deﬁned as: sr ( A l ) =
[Equation] A l [Equation] F
[Equation]
A
[Equation] (cid:107) 2 = (cid:80) σ 2 i σ 2 max , where σ i are the singular values of A l .
We note that 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 simpliﬁes the control problem by learning to activate different motor synergies depending on both s km and s kp .
C ONCLUSION
In this work, we explored how we can learn a universal controller for a large modular robot design space.
To this end, we proposed MetaMorph, a Transformer approach based on the insight that robot morphology is just another modality on which we can condition the output of a Transformer.
We showcased that pre-training on a large collection of diverse morphologies leads to policies which can generalize to unseen variations in kinematics, dynamics, new morphologies, and tasks.
We hope that our work serves as a step towards realizing the potential of large-scale pre-training and ﬁne-tuning in the ﬁeld of robotics, a paradigm that has seen tremendous success in vision and language.  A CKNOWLEDGEMENT
We gratefully acknowledge the support by Department of Navy awards (N00014-16-1-2127,
N00014-19-1-2477) issued by the Ofﬁce of Naval Research.
R EPRODUCIBILITY S TATEMENT
We have released a PyTorch (Paszke et al., 2019) implementation of MetaMorph on GitHub ( https://github.com/agrimgupta92/metamorph ).
In addition, the repository also includes all the robots and environments used for benchmarking
A I MPLEMENTATION D ETAILS
In this section, we provide additional implementation details for our method.
A. I NPUT O BSERVATIONS
In our setup, at each time step t (we drop the time subscript for brevity) the robot receives an observation s k = ( s km , s kp , s kg ) which is composed of the morphology representation ( s km ), the proprioceptive states ( s kp ) and additional global sensory information ( s kg ).
We note that prior work (Kurin et al., 2021; Huang et al., 2020) has deﬁned morphology as the connectivity of the limbs.
However, connectivity is only one aspect of morphology, and in general a substantial amount of additional information may be required to adequately describe the morphology.
Consider a modular robot composed of n ∈ 1 ...N k modules.
For each module s km 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).
All this information can be found for example, in the Universal Robot Description Format (URDF) or in simulator speciﬁc kinematic trees (e.g. MuJoCo XML (Todorov et al., 2012)).
Similarly for each module s kp 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.
Except the root node, each module will be connected to its parent module via a set of joints.
We adopt the convention of providing joint parameters and proprioception information to the child node.
We note that the importance of proprioceptive observations has also been observed in learning good locomotion (Lee et al., 2020; Kumar et al., 2021) and manipulation policies (Mandlekar et al., 2021).
In general, additional global sensory information ( s kg ) can be camera or depth sensor images.
To save computation, we provide the information about the terrain as 2D heightmap sampled on a nonuniform grid.
The grid is created by decreasing the sampling density as the distance from the root of the body increases.
All heights are expressed relative to the height of the ground immediately under the root of the agent.
The sampling points range from 1 m behind the agent to 4 m ahead of it along the direction of motion, as well as 4 m to the left and right.
A. E NVIRONMENTS
All our training and testing environments are implemented in MuJoCo (Todorov et al., 2012).
For a detailed description about the environments and reward functions please refer Gupta et al.
Gupta et al.
(2021) evolved UNIMALS in 3 different environments: ﬂat terrain, variable terrain and manipulation in variable terrain.
For each environment at the end of evolution, they hadtask optimized robots.
Out of these 300 , we choose a subset of 100 UNIMALS as our train set.
We ensure that no robot has the exact same kinematic tree or kinematic parameters.
Similarly, we created a test set of 100 UNIMALS.
Finally, for zero shot evaluation experiments we created the variants as described in § 5.3.
For creating kinematic variations for joint angles, we randomly selected a joint range for each joint from Table 2 which had alteast 50% overlap with the original joint angle range.
This helped in preventing robot variants which could not be trained.
A. T RAINING H YPERPARAMETERS
We use Transformers (Vaswani et al., 2017) to parametrize both policy and critic networks.
Global sensory information is encoded using a 2-layer MLP with hidden dimensions [64 , 64] .
We use
Proximal Policy Optimization (Schulman et al., 2017) for joint training of agents in all environments.
All hyperparameters for Transformer and PPO are listed in Table 1.
In addition, we use dynamic replay buffer sampling with α = 0 .
1 and β = 1 . 0 .
For all our ablations and baselines, we use the same hyperparameters.
For MLP baseline we use a 2-layer MLP with hidden dimensions [64 , 64] and for GNN baseline we use a 5-layer MLP with hidden dimension 512 .
The details of the joint training of modular robots are shown in Algorithm 1.
We ensure that all baselines and ablations have approximately the same number of parameters ( ∼ 3 . 3 Million).
A. E VALUATION M ETHODOLOGY
In general, the performance of RL algorithms is known to be strongly dependent on the choice of seed for random number generators (Henderson et al., 2018).
Hence, we run all our baseline and ablation experiments for 3 random seeds.
We found that the training curves were robust to the choice of seeds as indicated by small standard deviation (Fig. 4, 10).
Consequently, we evaluate the zero shot performance using pre-trained model corresponding to a single seed (Fig. 6).
For our transfer learning experiments (Fig. 7, 8), we ﬁne tune the model with the random seed corresponding to the one used for pre-training (for all 3 seeds).
B A DDITIONAL E XPERIMENTS B. B ASELINES AND A BLATIONS
Baselines : We compare against the following baselines: (1) MetaMorph-AO : Direct comparison to Amorpheus (Kurin et al., 2021) 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 a 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 ( ∼ 20x 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.
Speciﬁcally, we provide the following input observations: one hot encoding of unique limb ID, 3D Cartesian position, 4D quaternion orientation, 3D linear velocity, and 3D angular velocity, position in generalized coordinates, and normalized joint angle ranges.
Note that although both MetaMorph and Amorpheus don’t explicitly incorporate the graph structure as input to the policy, our input observations include additional morphology information (see § B.1).
In contrast, except for joint angle ranges, Amorpheus does not incorporate morphology information.
[Figure]
(2) Multi-Task MLP : We train a 6layer MLP with a hidden dimension of 512 .
The MLP receives the same observations as MetaMorph, and has the same number of parameters
Ablations : We perform additional ablation studies to understand the importance of learnt position embeddings and morphology information in input observation.
We refer our full method as MetaMorph and consider the following additional ablations: (1) MetaMorph-NMT ( N o
M orphology information + T ask encoding): we only provide s k =
( s kp , s kg ) as inputs, i.e. the model does not have access to information about the robot morphology.
In addition, we provide a context token with binary encoding to identify the morphology; (2) MetaMorph-HPE ( H and-designed P osition E ncoding): we provide s k = ( s k m , s k p , s k g ) as inputs, with s k m containing a per limb unique ID.
Fig.
10 shows the learning curves for joint training in the ﬂat terrain environment for 100 training morphologies.
Multi-task MLP baseline struggles to learn a good policy, which we attribute to the difﬁculty of the task due to the rich diversity of morphologies.
We next investigate the role of learnt position embeddings and morphological information.
MetaMorph-HPE performs slightly better than
MetaMorph-NPE, indicating that adding unique limb ids improves performance.
However, there is
[Figure]
still a large gap between MetaMorph and MetaMorph-HPE, suggesting that learnt position embeddings capture additional information.
Although MetaMorph-AO performs better than MetaMorphNM due to the addition of unique limb ids and joint angle ranges, it is still signiﬁcantly worse than
MetaMorph as it does not incorporate morphological information.
Finally, MetaMorph-NMT also performs poorly due to the lack of morphological information.
B. P OSITION E MBEDDING
To better understand the information captured by position embeddings we visualize their cosine similarity.
We found that across all seeds and in all environments (we only show ﬂat terrain for brevity) a diagonal pattern emerged (Fig 11).
Although, this might suggest that learnt position embeddings are capturing a unique position identiﬁer, we found that manually specifying a unique limb identiﬁer to perform much worse (see MetaMorph-HPE in Fig. 10).
Thus, indicating that these embeddings are capturing additional useful information.
We believe that further investigation is needed to better understand the information captured by position embeddings and defer it to future work.
B. K INEMATIC TREE TRAVERSAL ORDER
We ﬁrst process an arbitrary robot by creating a 1D sequence of tokens corresponding to the depth ﬁrst traversal of its kinematic tree.
We note that the DFS ordering is not unique, and for nodes with similar depth, we adopted the convention of visiting the nodes which come ﬁrst in the MuJoCo
XML.
This convention is also adopted by MuJoCo when parsing the XML.
We found that zero-shot transfer of the learnt policy was not robust to an opposite ordering of nodes and the performance dropped by ∼ 75% .
However, we found that the policy could be made robust to variations in the
DFS ordering by training the model with a simple data augmentation strategy of randomizing the visitation order of nodes at the same depth.
C L IMITATIONS A ND F UTURE W ORK
In this work, we explored how we can learn a universal controller for a large modular robot design space.
We make a key assumption that we already had access to a large collection of robots which were optimized for the task of locomotion.
Hence, currently we require a two stage pipeline where we ﬁrst create robots and then pre-train them jointly.
An important direction of future work would be to combine these two phases in an efﬁcient manner.
Moreover, although we found that MetaMorph has strong zero-shot generalization to unseen variations in kinematics and dynamics, there is indeed a performance drop on zero-shot generalization to new morphologies.
Hence, an important avenue of exploration is creating algorithms for sample efﬁcient transfer learning.
Finally, our current suite of tasks focuses primarily on locomotion skills.
An important line of future work will involve designing general purpose controllers which could perform multiple skills.
1: Input: π θ : policy function
V φ function
R
P k probability k 2: Initialize θ π φ
V φ
P distribution 3: for i
=1,2,..
N iter do 4:
# Collect robot experience in parallel 5: for j =1,2,..
N workers do 6:
Sample a robot k ∈ K , according to Equation7:
# Collect one episode of experience 8: τ j ≡ { s t , a t , s t +1 , r t +1 } j ∼ π θ 9: # Add data to the buffer 10:
R ← R ∪ τ j 11: end for 12: # Update policy and value functions 13: for j =1,2,..
N epochs do 14:
Sample a minibatch of data r from R 15: π θ ← PPOUpdate( π θ , r ) 16:
V φ ← PPOUpdate( V φ , r ) 17: end for 18: # Sample all robots uniformly at random initially 19: if i > = N warmup then 20:
P ← samplingProbUpdate(R), according to Equation21: end if 22: end for
[Table] [Table]
−
[Table]
[Table]