M ETA M ORPH : L EARNING U NIVERSAL C ONTROLLERS WITH T RANSFORMERS
Anonymous authors
Paper under double-blind review
A BSTRACT
Multiple domains like vision, natural language, and audio are witnessing tremendous progress by leveraging Transformers for large scale pre-training followed bytask speciﬁc ﬁne tuning.
In contrast, in robotics we primarily train a single robotfor 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 controllerover a modular robot design space.
MetaMorph is based on the insight that robotmorphology 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 robotmorphologies.
We further demonstrate that our pre-trained policy can be used forsample-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 agentswhose 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 oftabula rasa deep reinforcement learning and the challenge of designing a single robot which canperform a wide variety of tasks has led to the current dominant paradigm of ‘one robot one task’.
Instark 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, havewitnessed 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 domainsare 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 systemsprovide a natural opportunity by affording the ﬂexibility of combining a small set of general-purposebuilding blocks into a task-optimized morphology.
Indeed, modularity at the level of hardware is amotif 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 (Wanget al., 2018; Chen et al., 2018; Sanchez-Gonzalez et al., 2018) on learning policies that can generalizeacross different robot morphologies have been limited to: (1) manually constructed variations of asingle or few base morphologies, i.e. little diversity in the kinematic structure; (2) low complexityof control ( ≤ 7 degrees of freedom); (3) using Graph Neural Networks (Scarselli et al., 2008) basedon 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 universalcontroller for a modular robot design space which has the following properties: (a) generalizationto 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
[Figure]
Z e
[Figure]
o -
[Figure]
ho t G ene r a li z a t i on
MetaMorph
KinematicsVariations
Tasks
DynamicsVariations armature, material properties, gear etc.
DoFs, new morphologies, module parameters morphologies and tasks.
We instantiate the exploration of this general setting in the UNIMAL designspace introduced by Gupta et al.
We choose the UNIMAL design space as it contains achallenging ( 15 − 20 DoFs) distribution of robots that can learn locomotion and mobile manipulationin complex stochastic environments.
Learning a single universal controller for a huge variety ofrobot 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 robotdesign 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 ofdifferences in embodiment by leveraging a Transformer based architecture which takes as input asequence of tokens corresponding to the number of modules in the robot.
Each input token is createdby combining proprioceptive and morphology information at the level of constituent modules.
Thecombination of proprioceptive and embodiment modalities and large scale joint pre-training leadsto 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 thedifferences in learning speeds of different robots, we propose dynamic replay buffer balancing todynamically 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 amodular 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-shotgeneralize 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 proposedlike 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).
Wanget 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 singlemorphology but with higher control complexity.
The usage of GNNs is based on the assumptionthat the robot morphology is a good inductive bias to incorporate into neural controllers, which canbe naturally modelled by GNNs.
Recently, Kurin et al.
(2021) also proposed using Transformersfor training planar agents.
Our work differs from Kurin et al.
(2021) in the diversity and scale oftraining robots, complexity of the environments, conditioning the Transformer on morphologicalinformation, and showcasing strong generalization to unseen morphologies and tasks.
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 expressivesoft robot design spaces (Cheney et al., 2014; 2018).
In the case of rigid bodies, Ha (2019); Schaffet 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, GNNshave been leveraged to share controller parameters (Wang et al., 2019) across generations or developnovel heuristic search methods for efﬁcient exploration of the design space (Zhao et al., 2020). Incontrast to task speciﬁc morphology optimization, III et al. (2021) propose evolving morphologieswithout 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. 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 designspace, 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 rootof the tree; (2) cylinders parameterized by length, radius, and density to represent the limbs of therobot.
Two nodes of the graph can be connected via at most two motor-actuated hinge joints (i.e. Gis 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 discountfactor.
At each time step, the robot k receives an observation s kt , takes an action a kt , and is given areward r kt .
A policy π θ ( a kt | s kt ) models the conditional distribution over action a kt ∈
A given states kt ∈ S .
The goal is to ﬁnd policy parameters θ which maximize the average expected return acrossall 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 locomotionpolicies for complex high-dimensional agents from scratch, albeit with tedious hyperparamter tuning.
However, this approach is not suitable for modular design spaces containing exponentially manyrobot morphologies.
Indeed, Gupta et al.
(2021) estimates that the UNIMAL design space containsmore than 10 18 robots.
Hence, learning a separate policy for each robot is infeasible.
However, themodular nature of the design space implies that while each robot morphology is unique, it is stillconstructed from the same set of modules and potentially shares subgraphs of the kinematic treewith other morphologies.
We describe how MetaMorph exploits this structure to meet the challengeof 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 proprioceptivestates 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 toa different robot morphology, all drawn from the same modular design space.
Hence, instead oflearning a policy which is agnostic to the robot morphology, we need to learn a policy conditionedon 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 pairof deep neural networks: a policy network that produces an action distribution (Fig. 2), and a criticnetwork that predicts discounted future returns.
We use Transformers (Vaswani et al., 2017) toparametrize 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 sequenceof 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: 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 theembedding 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 thepre-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 episodelengths, whereas robot B will have longer episode lengths.
Hence, more data will be collected forrobot B in the earlier phases of training, and in turn will lead to robot B learning even faster, therebyresulting 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 theentire 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 asour performance metric. We determine P k by replacing E i with 1000 E i , where the numerator is themaximum 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 andkinematics 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 forwarddisplacement 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 duringeach 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 todesign a reward function tailored to each morphology.
In all tasks, our reward function promotesforward 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 afall (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 complex3D morphologies with variable number of joints.
Specifically, we replace how the NerveNet model receives input and produces output by our encode and decode stepsrespectively (§4.2).
In addition, the model receives thesame 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 ifthe domain-speciﬁc inductive bias of the robot kinematictree in the GNN is necessary.(2) MLP : We train all 100 robots separately with a 2-layer
MLP and report the average performance.
This baselineserves 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 fullmethod 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 MLPbaseline on both FT, and obstacle environment.
While MetaMorph performance in VT is slightlybelow the MLP baseline, we note that it has not saturated and we stopped training at 10 8 iterationsacross 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 alevel 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) sparinglyutilizes the 2D neighborhood structure of images at the beginning of the model by cutting the imageinto patches, MetaMorph uses the graph structure of the robot in the beginning by creating a 1Dsequence 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 fromscratch.
We highlight that the learned position embeddings signiﬁcantly improve the performanceof MetaMorph, just as they do in Transformer based image classiﬁcation.
Finally, without accessto the morphological information, MetaMorph-NM fails to learn a policy that can control diverserobot morphologies.
All of this substantiates our central claim that morphological state informationis necessary to learn successful control policies, although the kinematic graph need not be explicitlybaked 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 balancingis 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 modularrobot design space.
In this section, we demonstrate that MetaMorph shows favorable generalizationproperties 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, andmotor 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]
[Figure]
[Figure] and 2 − 3 × more sample efﬁciency [Figure] relevant property of all modules or joints.
See Table 2 for sampling ranges.
We then comparezero-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 GNNacross 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 notdo dynamics randomization during training, i.e., all robots in the training set have the same armatureand 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 unseenvariations in module shape parameters and changes to joint angle ranges.
In fact, changes to jointangles can signiﬁcantly alter the range of possible motion and might necessitate a different gaitfor successful locomotion.
Consequently, we ﬁnd that even though MetaMorph exhibits stronggeneralization performance compared to MetaMorph-NM and GNN in all 3 environments, thereis indeed a performance drop for the challenging setting of variations in joint angles.
However, thezero-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 werenot part of the training set, and we demonstrate that MetaMorph shows favorable sample efﬁcienttransfer 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 oftraining 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-shotperformance , 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 settingmimicking the promise of modular robots, where we are faced with a novel task and want to use anew robot morphology which may be suited for this task.
We consider the same set of 100 testrobots on two new tasks (Fig.
8): (1) Escape: The agent starts at the center of a bowl shaped terrainsurrounded 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 (thesize distribution is also different from the train task i.e. cuboid shapes).
We transfer the learnedpolicy from VT and Obstacles to Escape and Obstacles-Cylinders respectively.
Again, we ﬁnd thatthere is a strong zero-shot performance across all 3 environments and ﬁne-tuning is 2 to 3 timesmore 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), whichproposed 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 ofsets of joints (or muscles) through co-activation by a single neural signal.
Such synergies obviate theneed 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 ofthe attention matrix.
For attention matrix A l ∈ R m × m for layer l , the stable rank is deﬁned as:sr ( A l ) = (cid:107)
[Equation]
(cid:107) 2 F (cid:107)
A l (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 smalland oscillates between two values which correspond to attention maps where groups of limbs areactivated simultaneously (denoted by dark columns), a characteristic signature of motor synergies.
Hence, MetaMorph simpliﬁes the control problem by learning to activate different motor synergiesdepending 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 designspace.
To this end, we proposed MetaMorph, a Transformer approach based on the insight that robotmorphology is just another modality on which we can condition the output of a Transformer.
Weshowcased that pre-training on a large collection of diverse morphologies leads to policies which cangeneralize to unseen variations in kinematics, dynamics, new morphologies, and tasks.
We hope thatour work serves as a step towards realizing the potential of large-scale pre-training and ﬁne-tuningin the ﬁeld of robotics, a paradigm that has seen tremendous success in vision and language.
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 ).
Consider a modular robot composedof n ∈ 1 ...N k modules.
For each module s km i consists of: (1) Module Parameters : Module shapeparameters (e.g. radius, height), material information (e.g. density), and local geometric orientationand 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), andactuator 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 instantaneousstate 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 thegeneralized 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 proprioceptioninformation to the child node.
We note that the importance of proprioceptive observations has alsobeen observed in learning good locomotion (Lee et al., 2020; Kumar et al., 2021) and manipulationpolicies (Mandlekar et al., 2021).
In general, additional global sensory information ( s kg ) can be camera or depth sensor images.
Tosave 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 ofthe body increases.
All heights are expressed relative to the height of the ground immediately underthe root of the agent.
The sampling points range from 1 m behind the agent to 4 m ahead of it alongthe 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).
Fora 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 terrainand manipulation in variable terrain.
For each environment at the end of evolution, they had 100task optimized robots.
Out of these 300 , we choose a subset of 100 UNIMALS as our train set.
Weensure that no robot has the exact same kinematic tree or kinematic parameters.
Similarly, we createda test set of 100 UNIMALS.
Finally, for zero shot evaluation experiments we created the variants asdescribed in § 5.3.
For creating kinematic variations for joint angles, we randomly selected a jointrange 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.
Globalsensory 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 dynamicreplay buffer sampling with α = 0 .
1 and β = 1 . 0 .
For all our ablations and baselines, we use thesame 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 jointtraining of modular robots are shown in Algorithm 1.
We ensure that all baselines and ablationshave 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 ofseed for random number generators (Henderson et al., 2018).
Hence, we run all our baseline andablation experiments for 3 random seeds.
We found that the training curves were robust to the choiceof seeds as indicated by small standard deviation (Fig. 4, 10).
Consequently, we evaluate the zeroshot performance using pre-trained model corresponding to a single seed (Fig. 6). For our transferlearning experiments (Fig. 7, 8), we ﬁne tune the model with the random seed corresponding to theone used for pre-training (for all 3 seeds). B A DDITIONAL E XPERIMENTSB. 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 thefollowing reasons: (1) does not work with 3D morphologies with variable number of joints betweenlimbs; (2) does not incorporate exteroceptive observations; (3) Amorpheus ensures a balanced collection of experience by sequentially collecting data on each morphology and maintaining a separatereplay 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 GPUdays to train for 100 million iterations on Nvidia RTX 2080, while our method only needs 1.5 GPUdays ( ∼ 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 ingeneralized coordinates, and normalized joint angle ranges.
Note that except for joint angle ranges,
Amorpheus does not incorporate morphology information (see § A.1).
[Figure]
2) Multi-Task MLP
: We
train a 6-layer MLP
with a
hidden
dimension of 512 .
The MLP receivesthe
same
observations as MetaMorph,
and
has the same number of parameters.
Ablations : We perform additional ablation studies to understand the importance of learnt positionembeddings 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 orphologyinformation + T ask encoding): we only provide s k =
( s kp , s kg ) as inputs, i.e. the model does nothave access to information about the robot morphology.
In addition, we provide a context tokenwith binary encoding to identify the morphology; (2) MetaMorph-HPE ( H and-designed P osition
E ncoding): we provide s k = ( s km , s kp , s kg ) as inputs, with s km containing a per limb unique ID.
[Figure]
Fig.
10 shows the learning curvesfor joint training in the ﬂat terrainenvironment for 100 training morphologies.
Multi-task MLP baseline struggles to learn a good policy,which we attribute to the difﬁcultyof the task due to the rich diversityof morphologies.
We next investigate the role of learnt position embeddings and morphological information.
MetaMorph-HPE performsslightly better than MetaMorph-NPE,indicating that adding unique limbids improves performance.
However, there is still a large gap between MetaMorph and MetaMorphHPE, suggesting that learnt positionembeddings capture additional infor-
[Figure]
[Figure]
mation.
Although MetaMorph-AO performs better than MetaMorph-NM due to the addition ofunique limb ids and joint angle ranges, it is still signiﬁcantly worse than MetaMorph as it does notincorporate morphological information.
Finally, MetaMorph-NMT also performs poorly due to thelack of morphological information.
B. P OSITION E MBEDDING
To better understand the information captured by position embeddings we visualize their cosinesimilarity.
We found that across all seeds and in all environments (we only show ﬂat terrain forbrevity) a diagonal pattern emerged (Fig 11).
Although, this might suggest that learnt positionembeddings are capturing a unique position identiﬁer, we found that manually specifying a uniquelimb identiﬁer to perform much worse (see MetaMorph-HPE in Fig. 10).
Thus, indicating that theseembeddings are capturing additional useful information.
We believe that further investigation isneeded to better understand the information captured by position embeddings and defer it to futurework.
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 nodeswith 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-shottransfer of the learnt policy was not robust to an opposite ordering of nodes and the performancedropped 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 thevisitation 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 designspace.
We make a key assumption that we already had access to a large collection of robots whichwere optimized for the task of locomotion.
Hence, currently we require a two stage pipeline wherewe ﬁrst create robots and then pre-train them jointly.
An important direction of future work would beto combine these two phases in an efﬁcient manner.
Moreover, although we found that MetaMorphhas strong zero-shot generalization to unseen variations in kinematics and dynamics, there is indeeda performance drop on zero-shot generalization to new morphologies.
Hence, an important avenueof exploration is creating algorithms for sample efﬁcient transfer learning.
Finally, our current suiteof tasks focuses primarily on locomotion skills.
An important line of future work will involvedesigning general purpose controllers which could perform multiple skills.
1: Input: π θ : policy function
V φ function
R
P k : probability of sampling k 2: Initialize
Learnable parameters θ for π θ φ
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] [Table] [Table]