MetaMorph: Learning Universal Controllers with Transformers

Multiple domains like vision , natural language , and audio are witnessing tremen-1 dous progress by leveraging Transformers for large scale pre-training followed by2 task specific fine tuning . In contrast , in robotics we primarily train a single robot3 for a single task . However , modular robot systems now allow for the flexible com-4 bination of general-purpose building blocks into task optimized morphologies.5 However , given the exponentially large number of possible robot morphologies,6 training a controller for each new design is impractical . In this work , we pro-7 pose MetaMorph , a Transformer based approach to learn a universal controller8 over a modular robot design space . MetaMorph is based on the insight that robot9 morphology is just another modality on which we can condition the output of a10 Transformer . Through extensive experiments we demonstrate that large scale pre-11 training on a variety of robot morphologies results in policies with combinato-12 rial generalization capabilities , including zero shot generalization to unseen robot13 morphologies . We further demonstrate that our pre-trained policy can be used for14 sample-efficient transfer to completely new robot morphologies and tasks.15 1 INTRODUCTION16 The field of embodied intelligence posits that intelligent behaviours can be rapidly learned by agents17 whose morphologies are well adapted to their environment ( Brooks , 1991 ; Pfeifer & Scheier , 2001 ; 18 Bongard , 2014 ; Gupta et al. , 2021 ) . Based on this insight , a robot designer is faced with a predica-19 ment : should the robot design be task specific or general ? However , the sample inefficiency of20 tabula rasa deep reinforcement learning and the challenge of designing a single robot which can21 perform a wide variety of tasks has led to the current dominant paradigm of ‘ one robot one task ’ . In22 stark contrast , domains like vision ( Girshick et al. , 2014 ; He et al. , 2020 ) and language ( Dai & Le,23 2015 ; Radford et al. , 2018 ) , which are not plagued by the challenges of physical embodiment , have24 witnessed tremendous progress especially by leveraging large scale pre-training followed by trans-25 fer learning to many tasks through limited task-specific fine-tuning . Moreover , multiple domains26 are witnessing a confluence , with domain specific architectures being replaced by Transformers27 ( Vaswani et al. , 2017 ) , a general-purpose architecture with no domain-specific inductive biases.28 How can we bring to bear the advances in large-scale pre-training , transfer learning and general-29 purpose Transformer architectures , to the field of robotics ? We believe that modular robot systems30 provide a natural opportunity by affording the flexibility of combining a small set of general-purpose31 building blocks into a task-optimized morphology . Indeed , modularity at the level of hardware is a32 motif which is extensively utilized by evolution in biological systems ( Hartwell et al. , 1999 ; Kashtan33 & Alon , 2005 ) and by humans in many modern engineered systems . However , prior works ( Wang34 et al. , 2018 ; Chen et al. , 2018 ; Sanchez-Gonzalez et al. , 2018 ) on learning policies that can generalize35 across different robot morphologies have been limited to : ( 1 ) manually constructed variations of a36 single or few base morphologies , i.e . little diversity in the kinematic structure ; ( 2 ) low complexity37 of control ( ≤ 7 degrees of freedom ) ; ( 3 ) using Graph Neural Networks ( Scarselli et al. , 2008 ) based38 on the assumption that kinematic structure of the robot is the correct inductive bias.39 In this work , we take a step towards a more challenging setting ( Fig . 1 ) of learning a universal40 controller for a modular robot design space which has the following properties : ( a ) generalization41 to unseen variations in dynamics ( e.g . joint damping , armature , module mass ) and kinematics ( e.g.42 degree of freedom , morphology , module shape parameters ) and ( b ) sample-efficient transfer to new43 morphologies and tasks . We instantiate the exploration of this general setting in the UNIMAL design44 space introduced by Gupta et al . ( 2021 ) . We choose the UNIMAL design space as it contains a45 challenging ( 15−20 DoFs ) distribution of robots that can learn locomotion and mobile manipulation46 in complex stochastic environments . Learning a single universal controller for a huge variety of47 robot morphologies is difficult due to : ( 1 ) differences in action space , sensory input , morphology,48 dynamics , etc . ( 2 ) given a modular design space , not all robots are equally adept at learning a task,49 e.g . some robots might inherently be less sample-efficient ( Gupta et al. , 2021 ) .50 To this end , we propose MetaMorph , a method to learn a universal controller for a modular robot51 design space . MetaMorph is based on the insight that robot morphology is just another modal-52 ity on which we can condition the output of a Transformer . MetaMorph tackles the challenge of53 differences in embodiment by leveraging a Transformer based architecture which takes as input a54 sequence of tokens corresponding to the number of modules in the robot . Each input token is created55 by combining proprioceptive and morphology information at the level of constituent modules . The56 combination of proprioceptive and embodiment modalities and large scale joint pre-training leads57 to policies which exhibit zero-shot generalization to unseen variations in dynamics and kinemat-58 ics parameters and sample-efficient transfer to new morphologies and tasks . Finally , to tackle the59 differences in learning speeds of different robots , we propose dynamic replay buffer balancing to60 dynamically balance the amount of experience collection for a robot based on its performance.61 In sum , our key contributions are : ( 1 ) we introduce MetaMorph to learn a universal controller for a62 modular design space consisting of robots with high control complexity for challenging 3D locomo-63 tion tasks in stochastic environments ; ( 2 ) we showcase that our learned policy is able to zero-shot64 generalize to unseen variations in dynamics , kinematics , new morphologies and tasks , which is par-65 ticularly useful in real-world settings where controllers need to be robust to hardware failures ; ( 3 ) 66 we analyze the learned attention mask and discover the emergence of motor synergies ( Bernstein,67 1966 ) , which partially explains how MetaMorph is able to control a large number of robots.68 2 RELATED WORK69 Prior works on learning control policies which can generalize across robot morphologies have pri-70 marily focused on parametric variations of a single ( Chen et al. , 2018 ) or few ( 2 − 3 ) robot types71 ( Wang et al. , 2018 ; Sanchez-Gonzalez et al. , 2018 ; Huang et al. , 2020 ; Kurin et al. , 2021 ) . For gener-72 alizing across parametric variations of a single morphology , various approaches have been proposed73 like using a learned hardware embedding ( Chen et al. , 2018 ) , meta-learning for policy adaptation74 ( Al-Shedivat et al. , 2017 ; Ghadirzadeh et al. , 2021 ) , kinematics randomization ( Exarchos et al.,75 2020 ) , and dynamics randomization ( Peng et al. , 2018 ) . In case of multiple different morpholo-76 gies , one approach to tackle the challenge of differences in action and state spaces is to leverage77 Graph Neural Networks ( Scarselli et al. , 2008 ; Kipf & Welling , 2017 ; Battaglia et al. , 2018 ) . Wang78 et al . ( 2018 ) ; Huang et al . ( 2020 ) use GNNs to learn joint controllers for planar agents ( ≤ 7 DoFs ) .79 Blake et al . ( 2021 ) propose freezing selected parts of networks to enable training GNNs for a single80 morphology but with higher control complexity . The usage of GNNs is based on the assumption81 that the robot morphology is a good inductive bias to incorporate into neural controllers , which can82 be naturally modelled by GNNs . Recently , Kurin et al . ( 2021 ) also proposed using Transformers83 for training planar agents . Our work differs from Kurin et al . ( 2021 ) in the diversity and scale of84 training robots , complexity of the environments , conditioning the Transformer on morphological85 information , and showcasing strong generalization to unseen morphologies and tasks ( see § B.1 ) .86 Another closely related line of work is the design of modular robot design spaces and developing al-87 gorithms for co-optimizing morphology and control ( Sims , 1994 ) within a design space to find task-88 optimized combinations of controller and robot morphology . When the control complexity is low,89 evolutionary strategies have been successfully applied to find diverse morphologies in expressive90 soft robot design spaces ( Cheney et al. , 2014 ; 2018 ) . In the case of rigid bodies , Ha ( 2019 ) ; Schaff91 et al . ( 2019 ) ; Liao et al . ( 2019 ) have proposed using RL for finding optimal module parameters of92 fixed hand-designed morphology for rigid body robots . For more expressive design spaces , GNNs93 have been leveraged to share controller parameters ( Wang et al. , 2019 ) across generations or develop94 novel heuristic search methods for efficient exploration of the design space ( Zhao et al. , 2020 ) . In95 contrast to task specific morphology optimization , III et al . ( 2021 ) propose evolving morphologies96 without any task or reward specification . Finally , for self reconfigurable modular robots ( Fukuda & 97 Nakagawa , 1988 ; Yim et al. , 2007 ) , modular control has been utilized in both real ( Rubenstein et al.,98 2014 ; Mathews et al. , 2017 ) and simulated ( Pathak et al. , 2019 ) systems.99 3 LEARNING A UNIVERSAL CONTROLLER100 We begin by reviewing the UNIMAL design space and formulating the problem of learning a uni-101 versal controller for a modular robot design space as a multi-task reinforcement learning problem.102 3.1 THE UNIMAL DESIGN SPACE103 An agent morphology can be naturally represented as a kinematic tree , or a directed acyclic graph,104 corresponding to a hierarchy of articulated 3D rigid parts connected via motor actuated hinge joints.105 The graph G : = ( V , E ) consists of vertices V = { v1 , ... , vn } corresponding to modules of the design106 space , and edges eij ∈ E corresponding to joints between vi and vj . Concretely , in the UNIMAL107 ( Gupta et al. , 2021 ) design space , each node represents a component which can be one of two types:108 ( 1 ) a sphere parameterized by radius and density to represent the head of the agent and form the root109 of the tree ; ( 2 ) cylinders parameterized by length , radius , and density to represent the limbs of the110 robot . Two nodes of the graph can be connected via at most two motor-actuated hinge joints ( i.e . G111 is a multi-graph ) , parameterized by joint axis , joint limits and a motor gear ratio.112 3.2 JOINT POLICY OPTIMIZATION113 The problem of learning a universal controller for a set of K robots drawn from a modular robot de-114 sign space is a multi-task RL problem . Specifically , the control problem for each robot is an infinite-115 horizon discounted Markov decision process ( MDP ) represented by a tuple ( S , A , T , R , H , γ ) ,116 where S represents the set of states , A represents the set of available actions , T ( st+1|st , at ) rep-117 resents the transition dynamics , R ( s , a ) is a reward function , H is the horizon and γ is the discount118 factor . At each time step , the robot k receives an observation skt , takes an action a k t , and is given a119 reward rkt . A policy πθ ( a k t |skt ) models the conditional distribution over action akt ∈ A given state120 skt ∈ S. The goal is to find policy parameters θ which maximize the average expected return across121 all tasks : R = 1K ∑K k=0 ∑H t=0 γ trkt . We use Proximal Policy Optimization ( PPO ) ( Schulman et al.,122 2017 ) , a popular policy gradient ( Williams , 1992 ) method for optimizing this objective.123 4 METAMORPH124 Progress in model-free reinforcement learning algorithms has made it possible to train locomotion125 policies for complex high-dimensional agents from scratch , albeit with tedious hyperparamter tun-126 ing . However , this approach is not suitable for modular design spaces containing exponentially many127 robot morphologies . Indeed , Gupta et al . ( 2021 ) estimates that the UNIMAL design space contains128 more than 1018 robots . Hence , learning a separate policy for each robot is infeasible . However , the129 modular nature of the design space implies that while each robot morphology is unique , it is still130 constructed from the same set of modules and potentially shares subgraphs of the kinematic tree131 with other morphologies . We describe how MetaMorph exploits this structure to meet the challenge132 of learning a universal controller for different morphologies.133 4.1 FUSING PROPRIOCEPTIVE STATES AND MORPHOLOGY REPRESENTATIONS134 To learn policies that can generalize across morphologies , we must encode not only proprioceptive135 states essential for controlling a single robot , but also morphological information . From a multi-task136 RL perspective , this information can be viewed as a task identifier , where each task corresponds to137 a different robot morphology , all drawn from the same modular design space . Hence , instead of138 learning a policy which is agnostic to the robot morphology , we need to learn a policy conditioned139 on the robot morphology . Consequently , at each time step t ( we drop the time subscript for brevity ) ,140 the robot receives an observation sk = ( skm , s k p , s k g ) which is composed of the morphology repre-141 sentation ( skm ) , the proprioceptive states ( s k p ) , and additional global sensory information ( s k g ) . See142 § A.1 for a detailed description of each observation type.143 4.2 MORPHOLOGY AWARE TRANSFORMER144 The robot chooses its action via a stochastic policy πθ ( akt |skt ) where θ are the parameters of a pair145 of deep neural networks : a policy network that produces an action distribution ( Fig . 2 ) , and a critic146 network that predicts discounted future returns . We use Transformers ( Vaswani et al. , 2017 ) to147 parametrize both policy and critic networks as described in detail below.148 Encode . We make a distinction between how we process local and global state information . Concretely , let sk = ( skl , s k g ) where s k l = ( s k m , s k p ) . Since the number of joints between two modules can vary , we zero pad skli to ensure that input observation vectors are of the same size , i.e . s k l ∈ RN×M . In order to provide an arbitrary robot morphology as input to the Transformer , we first create a 1D sequence of local observation vectors by traversing the kinematic tree in depth first order starting at the root node ( torso in case of the UNIMAL design space ) . We then apply a single layer MLP independently to skli to create a D dimensional module embedding . We also add learnable 1D position embeddings to the module embeddings to automatically learn positional information : m0 = [ φ ( s k l1 ; We ) ; · · · ; φ ( s k lN ; We ) ] + Wpos , ( 1 ) where φ ( · ) is the embedding function , We ∈ RM×D are the embedding weights , and Wpos ∈149 RN×D are the learned positional embeddings . Note that in practice we zero pad the input sequence150 of local observation vectors for efficient batch processing of multiple morphologies.151 Process . From the module embeddings described above , we obtain the output feature vectors as : m′ ` = MSA ( LN ( m ` −1 ) ) + m ` −1 , ` = 1 . . . L ( 2 ) m ` = MLP ( LN ( m ′ ` ) ) + m ′ ` , ` = 1 . . . L ( 3 ) where MSA is multi-head attention ( Vaswani et al. , 2017 ) and LN is Layernorm ( Lei Ba et al. , 2016 ) .152 Decode . We integrate the global state information skg consisting of high-dimensional sensory input from camera or depth sensors . Naively concatenating skg and s k l in the encoder has two downsides : ( 1 ) it dilutes the importance of low-dimensional local sensory and morphological information ; ( 2 ) 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 : g = γ ( skg ; Wg ) , µ ( s k ) = φ ( mli , g ; Wd ) , πθ ( a k|sk ) = N ( { µ ( ski ) } Ni , Σ ) , ( 4 ) where γ ( · ) is a 2-layer MLP with parameters Wg , φ ( · ) is an embedding function with Wd as the153 embedding weights . The action distribution is modeled as a Gaussian distribution with a state-154 dependent mean µ ( ski ) and a fixed diagonal covariance matrix Σ . Similarly , for the critic network,155 we estimate value for the whole morphology by averaging the value per limb.156 4.3 DYNAMIC REPLAY BUFFER BALANCING157 Joint training of diverse morphologies is challenging as different morphologies are adept at per-158 forming different tasks . Consequently , some morphologies might be inherently better suited for the159 pre-training task . Let us consider two robots : ( A ) Robot A locomotes in a falling forward manner,160 i.e. , robot A is not passively stable ; ( B ) Robot B is passively stable . Especially with early termina-161 tion , robot A will keep falling during the initial phases of training , which results in shorter episode162 lengths , whereas robot B will have longer episode lengths . Hence , more data will be collected for163 robot B in the earlier phases of training , and in turn will lead to robot B learning even faster , thereby164 resulting in a ‘ rich gets richer ’ training dynamic . However , our goal is to ensure that all morpholo-165 gies have a similar level of performance at the end of training as we want to generalize across the166 entire distribution of morphologies.167 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 Ek 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 kth robot with probability Pk , given by : Pk = Eβk∑ Eβi Eτk = αEτk + ( 1− α ) E ( τ−1 ) k , ( 5 ) where α ∈ [ 0 , 1 ] is the discount factor and the exponent β determines the degree of dynamic pri-168 oritization , with β = 0 corresponding to the uniform case . In practice , we use episode length as169 our performance metric . We determine Pk by replacing Ei with 1000Ei , where the numerator is the170 maximum episode length.171 5 EXPERIMENTS172 In this section , we evaluate our method MetaMorph in different environments , perform extensive ab-173 lation studies of different design choices , test zero-shot generalization to variations in dynamics and174 kinematics parameters , and demonstrate sample efficient transfer to new morphologies and tasks.175 For qualitative results , please refer to the video on our project website 1.176 1https : //metamorph-iclr.github.io/site/ 5.1 EXPERIMENTAL SETUP177 We create a training set of 100 robots from the UNIMAL design space ( Gupta et al. , 2021 ) ( see178 § A.2 ) . We evaluate MetaMorph on three different environments ( Fig . 3 ) using the MuJoCo sim-179 ulator ( Todorov et al. , 2012 ) . In all 3 environments the goal of the agent is to maximize forward180 displacement over the course of an episode which lasts 1000 timesteps . The 3 environments are : ( 1 ) 181 Flat terrain ( FT ) ; ( 2 ) Variable terrain ( VT ) : VT is an extremely challenging environment as during182 each episode a new terrain is generated by randomly sampling a sequence of terrain types and inter-183 leaving them with flat terrain . We consider 3 types of terrains : hills , steps , and rubble ; ( 3 ) Obstacles:184 cuboid shaped obstacles of varying sizes on flat terrain.185 We use a dense morphology independent reward function for all our tasks as it is not feasible to186 design a reward function tailored to each morphology . In all tasks , our reward function promotes187 forward movement using small joint torques ( the latter obtained via a small energy usage penalty ) .188 In addition , as described in §4.3 , we use early termination across all environments when we detect a189 fall ( i.e . if the torso height drops below 50 % ( FT , Obstacles ) or 30 % ( VT ) of its initial height ) .190 5.2 BASELINES AND ABLATIONS191 Baselines : We compare against the following baselines:192 ( 1 ) GNN : We modify the NerveNet model proposed by193 Wang et al . ( 2018 ) to learn control policies for complex194 3D morphologies with variable number of joints . Specif-195 ically , we replace how the NerveNet model receives in-196 put and produces output by our encode and decode steps197 respectively ( §4.2 ) . In addition , the model receives the198 same observation as MetaMorph i.e . sk = ( skm , s k p , s k g ) 199 and is trained with dynamic replay buffer sampling . Thus,200 the only difference is in the process step . This helps test if201 the domain-specific inductive bias of the robot kinematic202 tree in the GNN is necessary.203 ( 2 ) MLP : We train all 100 robots separately with a 2-layer204 MLP and report the average performance . This baseline205 serves as an approximate upper bound for our method.206 Ablations : We also do an ablation study of different de-207 sign choices involved in our method . We refer our full208 method as MetaMorph and consider the following abla-209 tions : ( 1 ) MetaMorph-NPE : no learned position embeddings ; ( 2 ) MetaMorph-NM : we only provide210 sk = ( skp , s k g ) as inputs , i.e . the model does not have access to information about the robot morphol-211 ogy.212 Fig . 4 shows learning curves across 3 environments for training 100 morphologies . In all environ-213 ments MetaMorph can successfully match the average reward achieved via the per morphology MLP214 baseline on both FT , and obstacle environment . While MetaMorph performance in VT is slightly215 below the MLP baseline , we note that it has not saturated and we stopped training at 108 iterations216 across all three environments . Moreover , MetaMorph is significantly more sample-efficient ( 5× ) 217 than training independent MLPs ( 5 × 106 iterations per robot ) . The GNN baseline saturates at a218 level 2 to 3 times below MetaMorph . In GNN based models , locality and neighborhood connectiv-219 ity is explicitly baked into the model . Interestingly , just like ViT ( Dosovitskiy et al. , 2021 ) sparingly220 utilizes the 2D neighborhood structure of images at the beginning of the model by cutting the image221 into patches , MetaMorph uses the graph structure of the robot in the beginning by creating a 1D222 sequence corresponding to the kinematic tree by traversing the graph in depth first order . More-223 over , the position embeddings carry no information about the graph structure and are learned from224 scratch . We highlight that the learned position embeddings significantly improve the performance225 of MetaMorph , just as they do in Transformer based image classification . Finally , without access226 to the morphological information , MetaMorph-NM fails to learn a policy that can control diverse227 robot morphologies . All of this substantiates our central claim that morphological state information228 is necessary to learn successful control policies , although the kinematic graph need not be explicitly229 baked into neural architectures to learn policies capable of controlling diverse robot morphologies.230 Finally , we test the importance of dynamic replay buffer balancing in Fig . 5 , and find that balancing231 is necessary to learn a good control policy in 10− 15 % of robots across all 3 environments.232 5.3 ZERO-SHOT GENERALIZATION233 Our focus in this work is to learn policies that can generalize to unseen robots drawn from a modular234 robot design space . In this section , we demonstrate that MetaMorph shows favorable generalization235 properties across many different kinematic and dynamic variations.236 Experimental Setup . For each of the 100 training robots , we create a dedicated test set to test zero-237 shot transfer performance across two types of variations : dynamics ( armature , density , damping , and238 motor gear ) and kinematics ( module shape parameters like radius , and length of cylindrical limbs,239 and joint angles ) . For each training robot , we randomly create 4 different variants for each property,240 i.e . 400 robots with armature variations , and so on . While creating a new variant , we change the241 0 2 4 6 50 100 150 5e w ar d ( sFaSe 0 1 2 3 4 0 500 1000 1500 2000 2EstaFOes ( FyOinders ) ) inetune 6FratFh 0 . 0.2 0.4 0.6 0.8 1.0 IteratiRns ( ×107 ) 0.00 0.25 0.50 0.75 1.0 Obstacles ( cylinders ) Escape 2x3x 0 2 4 6 50 100 150 5e w ar d ( sFaSe 0 1 2 3 4 0 500 1000 1500 2000 2EstaFOes ( FyOinders ) ) inetune 6FratFh 0 . 0.2 0.4 0.6 0.8 1.0 IteratiRns ( ×107 ) 0.00 0.25 0.50 0.75 1.0 0 2 4 6 50 100 150 5e w ar d ( sFaSe 0 1 2 3 4 0 500 1000 1500 2000 2EstaFOes ( FyOinders ) ) inetune 6FratFh 0 . 0.2 0.4 0.6 0.8 1.0 IteratiRns ( ×107 ) 0.00 0.25 0.50 0.75 1.0 Figure 8 : Fine tuning : New robot morphologies and tasks . Left : Test environments . Right : Comparison of reward progression of 100 test robots averaged over 3 runs for pre-trained MetaMorph ( VT→ Escape , Obstacles→ Obstacles ( cylinders ) ) vs from scratch . Shaded regions denotes standard deviation . Across all environments pre-training leads to strong zero-shot performance and 2− 3× savings in training iterations to achieve the same level of average reward . relevant property of all modules or joints . See Table 2 for sampling ranges . We then compare242 zero-shot performance averaged over 10 trials.243 Generalization : Dynamics . First , we consider generalization to different dynamics ( Fig . 6 ) . We244 find consistently that MetaMorph performs significantly better than MetaMorph-NM and GNN245 across all types of dynamic variations and all environments . In fact , the difference is more pro-246 nounced for harder tasks like VT and Obstacles . We note that this result is surprising as we do not247 do dynamics randomization during training , i.e. , all robots in the training set have the same armature248 and damping parameters . Despite this , we see strong generalization performance.249 Generalization : Kinematics . Next , we consider generalization to different kinematics parameters250 ( Fig . 6 ) . This is a significantly more challenging setting as the model has to generalize to unseen251 variations in module shape parameters and changes to joint angle ranges . In fact , changes to joint252 angles can significantly alter the range of possible motion and might necessitate a different gait253 for successful locomotion . Consequently , we find that even though MetaMorph exhibits strong254 generalization performance compared to MetaMorph-NM and GNN in all 3 environments , there255 is indeed a performance drop for the challenging setting of variations in joint angles . However , the256 zero-shot performance is encouraging and motivates our next set of experiments on transfer learning.257 258 5.4 SAMPLE EFFICIENT TRANSFER LEARNING259 Experimental Setup . Here we create 100 robots from the UNIMAL design space which were260 not part of the training set , and we demonstrate that MetaMorph shows favorable sample efficient261 transfer to unseen morphologies , and even unseen morphologies performing novel tasks.262 Different morphologies . We first consider sample efficient learning of controllers for new mor-263 phologies on the same task by taking a pre-trained MetaMorph model on an environment , and then264 fine tuning it to learn to control morphologies in the test set . In Fig . 7 we compare the number of265 training iterations required to achieve a particular performance level when we fine tune MetaMorph266 vs training from scratch on the test set . Across all 3 environments we not only find strong zero-shot267 performance , but also 2 to 3 times higher sample efficiency compared to training from scratch.268 Different morphologies and novel tasks . Finally , we consider the most realistic and general setting269 mimicking the promise of modular robots , where we are faced with a novel task and want to use a270 new robot morphology which may be suited for this task . We consider the same set of 100 test271 robots on two new tasks ( Fig . 8 ) : ( 1 ) Escape : The agent starts at the center of a bowl shaped terrain272 surrounded by small hills ( bumps ) , and has to maximize the geodesic distance from the start location273 ( escape the hilly region ) . ( 2 ) Obstacles ( cylinders ) : cylinder shaped obstacles of varying sizes ( the274 size distribution is also different from the train task i.e . cuboid shapes ) . We transfer the learned275 policy from VT and Obstacles to Escape and Obstacles-Cylinders respectively . Again , we find that276 there is a strong zero-shot performance across all 3 environments and fine-tuning is 2 to 3 times277 more sample efficient than training from scratch.278 5.5 EMERGENT MOTOR SYNERGIES279 We next search for a potential explanation for how MetaMorph can coordinate the large number of280 DoFs ( ∼ 16 × 100 ) across several agents . Our hypothesis is inspired by Bernstein ( 1966 ) , which281 proposed the existence of muscle synergies as a neural strategy to simplify the control of multiple282 DoFs by the central nervous system . A muscle synergy corresponds to the constrained movements of283 sets of joints ( or muscles ) through co-activation by a single neural signal . Such synergies obviate the284 need to control all joints independently , by coupling sets of joints into adaptive functional groups.285 Although the definition of synergies ( Bruton & O ’ Dwyer , 2018 ) vary in the literature , dimension-286 ality reduction is generally accepted as a signature of synergistic control ( Todorov & Ghahramani,287 2004 ; Todorov , 2004 ) . Consequently , to test this hypothesis , in Fig . 9 we plot the stable rank of288 the attention matrix . For attention matrix Al ∈ Rm×m for layer l , the stable rank is defined as:289 sr ( Al ) = ‖Al‖2F ‖Al‖2 = ∑ σ2i σ2max , where σi are the singular values of Al . We note that sr ( Al ) is small290 and oscillates between two values which correspond to attention maps where groups of limbs are291 activated simultaneously ( denoted by dark columns ) , a characteristic signature of motor synergies.292 Hence , MetaMorph simplifies the control problem by learning to activate different motor synergies293 depending on both skm and s k p .294 6 CONCLUSION295 In this work , we explored how we can learn a universal controller for a large modular robot design296 space . To this end , we proposed MetaMorph , a Transformer approach based on the insight that robot297 morphology is just another modality on which we can condition the output of a Transformer . We298 showcased that pre-training on a large collection of diverse morphologies leads to policies which can299 generalize to unseen variations in kinematics , dynamics , new morphologies , and tasks . We hope that300 our work serves as a step towards realizing the potential of large-scale pre-training and fine-tuning301 in the field of robotics , a paradigm that has seen tremendous success in vision and language.302 REFERENCES303 Maruan Al-Shedivat , Trapit Bansal , Yuri Burda , Ilya Sutskever , Igor Mordatch , and Pieter Abbeel.304 Continuous adaptation via meta-learning in nonstationary and competitive environments . arXiv305 preprint arXiv:1710.03641 , 2017.306 Peter W Battaglia , Jessica B Hamrick , Victor Bapst , Alvaro Sanchez-Gonzalez , Vinicius Zambaldi,307 Mateusz Malinowski , Andrea Tacchetti , David Raposo , Adam Santoro , Ryan Faulkner , et al.308 Relational inductive biases , deep learning , and graph networks . arXiv preprint arXiv:1806.01261,309 2018.310 Nikolai Bernstein . The co-ordination and regulation of movements . The co-ordination and regula-311 tion of movements , 1966.312 Charlie Blake , Vitaly Kurin , Maximilian Igl , and Shimon Whiteson . Snowflake : Scaling gnns to313 high-dimensional continuous control via parameter freezing . arXiv preprint arXiv:2103.01009,314 2021.315 Josh Bongard . Why morphology matters . The horizons of evolutionary robotics , 6:125–152 , 2014.316 Rodney A Brooks . New approaches to robotics . Science , 253 ( 5025 ) :1227–1232 , 1991.317 Michaela Bruton and Nicholas O ’ Dwyer . Synergies in coordination : A comprehensive overview318 of neural , computational , and behavioral approaches . Journal of Neurophysiology , 120 ( 6 ) :2761–319 2774 , 2018.320 Tao Chen , Adithyavairavan Murali , and Abhinav Gupta . Hardware conditioned policies for multi-321 robot transfer learning . In NIPS , 2018.322 Nick Cheney , Robert MacCurdy , Jeff Clune , and Hod Lipson . Unshackling evolution : Evolving323 soft robots with multiple materials and a powerful generative encoding . SIGEVOlution , 7 ( 1 ) :324 11–23 , August 2014. doi : 10.1145/2661735.2661737 . URL https : //doi.org/10.1145/325 2661735.2661737.326 Nick Cheney , Josh Bongard , Vytas SunSpiral , and Hod Lipson . Scalable co-optimization of mor-327 phology and control in embodied machines . Journal of The Royal Society Interface , 15 ( 143 ) :328 20170937 , 2018.329 Andrew M Dai and Quoc V Le . Semi-supervised sequence learning . Advances in neural information330 processing systems , 28:3079–3087 , 2015.331 Alexey Dosovitskiy , Lucas Beyer , Alexander Kolesnikov , Dirk Weissenborn , Xiaohua Zhai , Thomas332 Unterthiner , Mostafa Dehghani , Matthias Minderer , Georg Heigold , Sylvain Gelly , Jakob Uszko-333 reit , and Neil Houlsby . An image is worth 16x16 words : Transformers for image recognition at334 scale . In ICLR , 2021.335 Ioannis Exarchos , Yifeng Jiang , Wenhao Yu , and C Karen Liu . Policy transfer via kinematic domain336 randomization and adaptation . arXiv preprint arXiv:2011.01891 , 2020.337 Toshio Fukuda and Seiya Nakagawa . Dynamically reconfigurable robotic system . In ICRA , pp.338 1581–1586 . IEEE , 1988.339 Ali Ghadirzadeh , Xi Chen , Petra Poklukar , Chelsea Finn , Mårten Björkman , and Danica Kragic.340 Bayesian meta-learning for few-shot policy adaptation across robotic platforms . arXiv preprint341 arXiv:2103.03697 , 2021.342 Ross Girshick , Jeff Donahue , Trevor Darrell , and Jitendra Malik . Rich feature hierarchies for ac-343 curate object detection and semantic segmentation . In Proceedings of the IEEE conference on344 computer vision and pattern recognition , pp . 580–587 , 2014.345 Agrim Gupta , Silvio Savarese , Surya Ganguli , and Li Fei-Fei . Embodied intelligence via learning346 and evolution . Nature communications , 12 ( 1 ) :5721 , 2021.347 David Ha . Reinforcement learning for improving agent design . Artificial life , 25 ( 4 ) :352–365 , 2019.348 Leland H Hartwell , John J Hopfield , Stanislas Leibler , and Andrew W Murray . From molecular to349 modular cell biology . Nature , 402 ( 6761 ) : C47–C52 , 1999.350 Kaiming He , Haoqi Fan , Yuxin Wu , Saining Xie , and Ross Girshick . Momentum contrast for351 unsupervised visual representation learning . In CVPR , pp . 9729–9738 , 2020.352 Peter Henderson , Riashat Islam , Philip Bachman , Joelle Pineau , Doina Precup , and David Meger.353 Deep reinforcement learning that matters . In Thirty-Second AAAI Conference on Artificial Intel-354 ligence , 2018.355 Wenlong Huang , Igor Mordatch , and Deepak Pathak . One policy to control them all : Shared modular356 policies for agent-agnostic control . In ICML , pp . 4455–4464 . PMLR , 2020.357 Donald Joseph Hejna III , Pieter Abbeel , and Lerrel Pinto . Task-agnostic morphology evolution . In358 International Conference on Learning Representations , 2021 . URL https : //openreview.359 net/forum ? id=CGQ6ENUMX6.360 Nadav Kashtan and Uri Alon . Spontaneous evolution of modularity and network motifs . Proceed-361 ings of the National Academy of Sciences , 102 ( 39 ) :13773–13778 , 2005.362 Thomas Kipf and Max Welling . Semi-supervised classification with graph convolutional networks.363 ArXiv , abs/1609.02907 , 2017.364 Ashish Kumar , Zipeng Fu , Deepak Pathak , and Jitendra Malik . Rma : Rapid motor adaptation for365 legged robots . arXiv preprint arXiv:2107.04034 , 2021.366 Vitaly Kurin , Maximilian Igl , Tim Rocktäschel , Wendelin Boehmer , and Shimon Whiteson . My367 body is a cage : the role of morphology in graph-based incompatible control . In ICLR , 2021.368 Joonho Lee , Jemin Hwangbo , Lorenz Wellhausen , Vladlen Koltun , and Marco Hutter . Learning369 quadrupedal locomotion over challenging terrain . Science robotics , 5 ( 47 ) , 2020.370 Jimmy Lei Ba , Jamie Ryan Kiros , and Geoffrey E. Hinton . Layer Normalization . arXiv e-prints,371 art . arXiv:1607.06450 , July 2016.372 T. Liao , G. Wang , B. Yang , R. Lee , K. Pister , S. Levine , and R. Calandra . Data-efficient learning of373 morphology and controller for a microrobot . In 2019 International Conference on Robotics and374 Automation ( ICRA ) , pp . 2488–2494 , 2019. doi : 10.1109/ICRA.2019.8793802.375 Ajay Mandlekar , Danfei Xu , Josiah Wong , Soroush Nasiriany , Chen Wang , Rohun Kulkarni , Li Fei-376 Fei , Silvio Savarese , Yuke Zhu , and Roberto Martı́n-Martı́n . What matters in learning from offline377 human demonstrations for robot manipulation . In 5th Annual Conference on Robot Learning,378 2021 . URL https : //openreview.net/forum ? id=JrsfBJtDFdI.379 Nithin Mathews , Anders Lyhne Christensen , Rehan O ’ Grady , Francesco Mondada , and Marco380 Dorigo . Mergeable nervous systems for robots . Nature communications , 8 ( 1 ) :1–7 , 2017.381 Deepak Pathak , Christopher Lu , Trevor Darrell , Phillip Isola , and Alexei A Efros . Learning to382 control self-assembling morphologies : A study of generalization via modularity . In H. Wallach,383 H. Larochelle , A. Beygelzimer , F. d'Alché-Buc , E. Fox , and R. Garnett ( eds . ) , Advances in Neural384 Information Processing Systems , volume 32 . Curran Associates , Inc. , 2019.385 Xue Bin Peng , Marcin Andrychowicz , Wojciech Zaremba , and Pieter Abbeel . Sim-to-real transfer of386 robotic control with dynamics randomization . In 2018 IEEE international conference on robotics387 and automation ( ICRA ) , pp . 3803–3810 . IEEE , 2018.388 Rolf Pfeifer and Christian Scheier . Understanding intelligence . MIT press , 2001.389 Alec Radford , Karthik Narasimhan , Tim Salimans , and Ilya Sutskever . Improving language under-390 standing by generative pre-training ( 2018 ) , 2018.391 Michael Rubenstein , Alejandro Cornejo , and Radhika Nagpal . Programmable self-assembly in a392 thousand-robot swarm . Science , 345 ( 6198 ) :795–799 , 2014.393 Alvaro Sanchez-Gonzalez , Nicolas Heess , Jost Tobias Springenberg , Josh Merel , Martin Riedmiller,394 Raia Hadsell , and Peter Battaglia . Graph networks as learnable physics engines for inference and395 control . In ICML , pp . 4470–4479 . PMLR , 2018.396 Franco Scarselli , Marco Gori , Ah Chung Tsoi , Markus Hagenbuchner , and Gabriele Monfardini.397 The graph neural network model . IEEE transactions on neural networks , 20 ( 1 ) :61–80 , 2008.398 Charles Schaff , David Yunis , Ayan Chakrabarti , and Matthew R Walter . Jointly learning to construct399 and control agents using deep reinforcement learning . In ICRA , pp . 9798–9805 . IEEE , 2019.400 John Schulman , Filip Wolski , Prafulla Dhariwal , Alec Radford , and Oleg Klimov . Proximal Policy401 Optimization Algorithms . arXiv e-prints , art . arXiv:1707.06347 , July 2017.402 Karl Sims . Evolving 3d morphology and behavior by competition . Artificial life , 1 ( 4 ) :353–372,403 1994.404 Emanuel Todorov . Optimality principles in sensorimotor control . Nature neuroscience , 7 ( 9 ) :907–405 915 , 2004.406 Emanuel Todorov and Zoubin Ghahramani . Analysis of the synergies underlying complex hand407 manipulation . In The 26th Annual International Conference of the IEEE Engineering in Medicine408 and Biology Society , volume 2 , pp . 4637–4640 . IEEE , 2004.409 Emanuel Todorov , Tom Erez , and Yuval Tassa . Mujoco : A physics engine for model-based control.410 In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems , pp . 5026–5033.411 IEEE , 2012.412 Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,413 Łukasz Kaiser , and Illia Polosukhin . Attention is all you need . In NIPS , 2017.414 Tingwu Wang , Renjie Liao , Jimmy Ba , and Sanja Fidler . Nervenet : Learning structured policy with415 graph neural networks . In ICLR , 2018.416 Tingwu Wang , Yuhao Zhou , Sanja Fidler , and Jimmy Ba . Neural graph evolution : Automatic robot417 design . In ICLR , 2019.418 Ronald J Williams . Simple statistical gradient-following algorithms for connectionist reinforcement419 learning . Machine learning , 8 ( 3-4 ) :229–256 , 1992.420 Mark Yim , Wei-Min Shen , Behnam Salemi , Daniela Rus , Mark Moll , Hod Lipson , Eric Klavins , and421 Gregory S Chirikjian . Modular self-reconfigurable robot systems [ grand challenges of robotics ] .422 IEEE Robotics & Automation Magazine , 14 ( 1 ) :43–52 , 2007.423 Allan Zhao , Jie Xu , Mina Konaković-Luković , Josephine Hughes , Andrew Spielberg , Daniela Rus,424 and Wojciech Matusik . Robogrammar : graph grammar for terrain-optimized robot design . ACM425 Transactions on Graphics ( TOG ) , 39 ( 6 ) :1–16 , 2020.426 A IMPLEMENTATION DETAILS427 In this section , we provide additional implementation details for our method.428 A.1 INPUT OBSERVATIONS429 In our setup , at each time step t ( we drop the time subscript for brevity ) the robot receives an430 observation sk = ( skm , s k p , s k g ) which is composed of the morphology representation ( s k m ) , the pro-431 prioceptive states ( skp ) and additional global sensory information ( s k g ) . We note that prior work432 ( Kurin et al. , 2021 ; Huang et al. , 2020 ) has defined morphology as the connectivity of the limbs.433 However , connectivity is only one aspect of morphology , and in general a substantial amount of434 additional information may be required to adequately describe the morphology . Consider a modular435 robot composed of n ∈ 1 ... Nk modules . For each module skmi consists of : ( 1 ) Module Parame-436 ters : Module shape parameters ( e.g . radius , height ) , material information ( e.g . density ) , and local437 geometric orientation and position of the child module with respect to the parent module . ( 2 ) Joint438 Parameters : This consists of information about joint type ( e.g . hinge ) and its properties ( e.g . joint439 range and axis ) , and actuator type ( e.g . motor ) and its properties ( e.g . gear ) . All this information can440 be found for example , in the Universal Robot Description Format ( URDF ) or in simulator specific441 kinematic trees ( e.g . MuJoCo XML ( Todorov et al. , 2012 ) ) .442 Similarly for each module skpi consists of the instantaneous state of the system : ( 1 ) Module Pro-443 prioception : 3D Cartesian position , 4D quaternion orientation , 3D linear velocity , and 3D angular444 velocity . ( 2 ) Joint Proprioception : Position and velocity in the generalized coordinates . Except the445 root node , each module will be connected to its parent module via a set of joints . We adopt the con-446 vention of providing joint parameters and proprioception information to the child node . We note that447 the importance of proprioceptive observations has also been observed in learning good locomotion448 ( Lee et al. , 2020 ; Kumar et al. , 2021 ) and manipulation policies ( Mandlekar et al. , 2021 ) .449 In general , additional global sensory information ( skg ) can be camera or depth sensor images . To450 save computation , we provide the information about the terrain as 2D heightmap sampled on a non-451 uniform grid . The grid is created by decreasing the sampling density as the distance from the root of452 the body increases . All heights are expressed relative to the height of the ground immediately under453 the root of the agent . The sampling points range from 1m behind the agent to 4m ahead of it along454 the direction of motion , as well as 4m to the left and right.455 A.2 ENVIRONMENTS456 All our training and testing environments are implemented in MuJoCo ( Todorov et al. , 2012 ) . For457 a detailed description about the environments and reward functions please refer Gupta et al . ( 2021 ) .458 Gupta et al . ( 2021 ) evolved UNIMALS in 3 different environments : flat terrain , variable terrain459 and manipulation in variable terrain . For each environment at the end of evolution , they had 100460 task optimized robots . Out of these 300 , we choose a subset of 100 UNIMALS as our train set . We461 ensure that no robot has the exact same kinematic tree or kinematic parameters . Similarly , we created462 a test set of 100 UNIMALS . Finally , for zero shot evaluation experiments we created the variants as463 described in § 5.3 . For creating kinematic variations for joint angles , we randomly selected a joint464 range for each joint from Table 2 which had alteast 50 % overlap with the original joint angle range.465 This helped in preventing robot variants which could not be trained.466 A.3 TRAINING HYPERPARAMETERS467 We use Transformers ( Vaswani et al. , 2017 ) to parametrize both policy and critic networks . Global468 sensory information is encoded using a 2-layer MLP with hidden dimensions [ 64 , 64 ] . We use469 Proximal Policy Optimization ( Schulman et al. , 2017 ) for joint training of agents in all environments.470 All hyperparameters for Transformer and PPO are listed in Table 1 . In addition , we use dynamic471 replay buffer sampling with α = 0.1 and β = 1.0 . For all our ablations and baselines , we use the472 same hyperparameters . For MLP baseline we use a 2-layer MLP with hidden dimensions [ 64 , 64 ] 473 and for GNN baseline we use a 5-layer MLP with hidden dimension 512 . The details of the joint474 training of modular robots are shown in Algorithm 1 . We ensure that all baselines and ablations475 have approximately the same number of parameters ( ∼ 3.3 Million ) .476 A.4 EVALUATION METHODOLOGY477 In general , the performance of RL algorithms is known to be strongly dependent on the choice of478 seed for random number generators ( Henderson et al. , 2018 ) . Hence , we run all our baseline and479 ablation experiments for 3 random seeds . We found that the training curves were robust to the choice480 of seeds as indicated by small standard deviation ( Fig . 4 , 10 ) . Consequently , we evaluate the zero481 shot performance using pre-trained model corresponding to a single seed ( Fig . 6 ) . For our transfer482 learning experiments ( Fig . 7 , 8 ) , we fine tune the model with the random seed corresponding to the483 one used for pre-training ( for all 3 seeds ) .484 B ADDITIONAL EXPERIMENTS485 B.1 BASELINES AND ABLATIONS486 Baselines : We compare against the following baselines:487 ( 1 ) MetaMorph-AO : Direct comparison to Amorpheus ( Kurin et al. , 2021 ) is not feasible due to the488 following reasons : ( 1 ) does not work with 3D morphologies with variable number of joints between489 limbs ; ( 2 ) does not incorporate exteroceptive observations ; ( 3 ) Amorpheus ensures a balanced col-490 lection of experience by sequentially collecting data on each morphology and maintaining a separate491 replay buffer per morphology . Further , updates to policy parameters are also performed sequentially.492 This sequential nature of the algorithm is not amenable to scaling , and would require ∼ 30 GPU493 days to train for 100 million iterations on Nvidia RTX 2080 , while our method only needs 1.5 GPU494 days ( ∼ 20x training speedup ) .495 Hence , we compare with MetaMorph-AO ( Amorpheus Observation ) as the closest variant to Amor-496 pheus , where we provide the same input observations as described in Kurin et al . ( 2021 ) . Specifi-497 cally , we provide the following input observations : one hot encoding of unique limb ID , 3D Carte-498 sian position , 4D quaternion orientation , 3D linear velocity , and 3D angular velocity , position in499 generalized coordinates , and normalized joint angle ranges . Note that although both MetaMorph500 and Amorpheus don ’ t explicitly incorporate the graph structure as input to the policy , our input ob-501 servations include additional morphology information ( see § B.1 ) . In contrast , except for joint angle502 ranges , Amorpheus does not incorporate morphology information.503 ( 2 ) Multi-Task MLP : We train a 6-504 layer MLP with a hidden dimension505 of 512 . The MLP receives the same506 observations as MetaMorph , and has507 the same number of parameters.508 Ablations : We perform additional509 ablation studies to understand the im-510 portance of learnt position embed-511 dings and morphology information512 in input observation . We refer our513 full method as MetaMorph and con-514 sider the following additional ab-515 lations : ( 1 ) MetaMorph-NMT ( No516 Morphology information + Task en-517 coding ) : we only provide sk =518 ( skp , s k g ) as inputs , i.e . the model does519 not have access to information about520 the robot morphology . In addition , we provide a context token with binary encoding to iden-521 tify the morphology ; ( 2 ) MetaMorph-HPE ( Hand-designed Position Encoding ) : we provide sk =522 ( skm , s k p , s k g ) as inputs , with s k m containing a per limb unique ID.523 Fig . 10 shows the learning curves for joint training in the flat terrain environment for 100 training524 morphologies . Multi-task MLP baseline struggles to learn a good policy , which we attribute to the525 difficulty of the task due to the rich diversity of morphologies . We next investigate the role of learnt526 position embeddings and morphological information . MetaMorph-HPE performs slightly better than527 MetaMorph-NPE , indicating that adding unique limb ids improves performance . However , there is528 still a large gap between MetaMorph and MetaMorph-HPE , suggesting that learnt position embed-529 dings capture additional information . Although MetaMorph-AO performs better than MetaMorph-530 NM due to the addition of unique limb ids and joint angle ranges , it is still significantly worse than531 MetaMorph as it does not incorporate morphological information . Finally , MetaMorph-NMT also532 performs poorly due to the lack of morphological information.533 B.2 POSITION EMBEDDING534 To better understand the information captured by position embeddings we visualize their cosine535 similarity . We found that across all seeds and in all environments ( we only show flat terrain for536 brevity ) a diagonal pattern emerged ( Fig 11 ) . Although , this might suggest that learnt position537 embeddings are capturing a unique position identifier , we found that manually specifying a unique538 limb identifier to perform much worse ( see MetaMorph-HPE in Fig . 10 ) . Thus , indicating that these539 embeddings are capturing additional useful information . We believe that further investigation is540 needed to better understand the information captured by position embeddings and defer it to future541 work.542 B.3 KINEMATIC TREE TRAVERSAL ORDER543 We first process an arbitrary robot by creating a 1D sequence of tokens corresponding to the depth544 first traversal of its kinematic tree . We note that the DFS ordering is not unique , and for nodes545 with similar depth , we adopted the convention of visiting the nodes which come first in the MuJoCo546 XML . This convention is also adopted by MuJoCo when parsing the XML . We found that zero-shot547 transfer of the learnt policy was not robust to an opposite ordering of nodes and the performance548 dropped by ∼ 75 % . However , we found that the policy could be made robust to variations in the549 DFS ordering by training the model with a simple data augmentation strategy of randomizing the550 visitation order of nodes at the same depth.551 C LIMITATIONS AND FUTURE WORK552 In this work , we explored how we can learn a universal controller for a large modular robot design553 space . We make a key assumption that we already had access to a large collection of robots which554 were optimized for the task of locomotion . Hence , currently we require a two stage pipeline where555 we first create robots and then pre-train them jointly . An important direction of future work would be556 to combine these two phases in an efficient manner . Moreover , although we found that MetaMorph557 has strong zero-shot generalization to unseen variations in kinematics and dynamics , there is indeed558 a performance drop on zero-shot generalization to new morphologies . Hence , an important avenue559 of exploration is creating algorithms for sample efficient transfer learning . Finally , our current suite560 of tasks focuses primarily on locomotion skills . An important line of future work will involve561 designing general purpose controllers which could perform multiple skills.562 Algorithm 1 MetaMorph : Joint Training of Modular Robots 1 : Input : πθ : policy function Vφ : value function R : replay buffer K : training pool of robots Pk : probability of sampling robot k 2 : Initialize : Learnable parameters θ for πθ , φ for Vφ , P uniform distribution 3 : for i=1,2 , .. Niter do 4 : # Collect robot experience in parallel 5 : for j=1,2 , .. Nworkers do 6 : Sample a robot k ∈ K , according to Equation 5 7 : # Collect one episode of experience 8 : τj ≡ { st , at , st+1 , rt+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 , .. Nepochs 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 > = Nwarmup then 20 : P ← samplingProbUpdate ( R ) , according to Equation 5 21 : end if 22 : end for

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