Transform2Act: Learning a Transform-and-Control Policy for Efficient Agent Design

1 INTRODUCTION . Automatic and efficient design of robotic agents in simulation holds great promise in complementing and guiding the traditional physical robot design process ( Ha et al. , 2017 ; 2018 ) , which can be laborious and time-consuming . In this paper , we consider a setting where we optimize both the skeletal structure and joint attributes ( e.g. , bone length , size , and motor strength ) of an agent to maximize its performance on a given task . This is a very challenging problem for two main reasons . First , the design space , i.e. , all possible skeletal structures and their joint attributes , is prohibitively vast and combinatorial , which also makes applying gradient-based continuous optimization methods difficult . Second , the problem is inherently bi-level : ( 1 ) we need to search an immensely large design space and ( 2 ) the evaluation of each candidate design entails solving a computationally expensive inner optimization to find its optimal controller . Prior work typically uses evolutionary search ( ES ) algorithms for combinatorial design optimization ( Sims , 1994 ; Wang et al. , 2019 ) . During each iteration , ES-based methods maintain a large population of agents with various designs where each agent learns to perform the task independently . When the learning ends , agents with the worst performances are eliminated while surviving agents produce child agents with random mutation to maintain the size of the population . ES-based methods have low sample efficiency since agents in the population do not share their experiences and many samples are wasted on eliminated agents . Fur- thermore , zeroth-order optimization methods such as ES are known to be sample-inefficient when the optimization search space ( i.e. , design space ) is high-dimensional ( Vemula et al. , 2019 ) . In light of the above challenges , we take a new approach to agent design optimization by incorporating the design procedure of an agent into its decision-making process . Specifically , we learn a transform-and-control policy , called Transform2Act , that first designs an agent and then controls the designed agent . In an episode , we divide the agent ’ s behavior into two consecutive stages , a transform stage and an execution stage , on which the policy is conditioned . In the transform stage , the agent applies a sequence of transform actions to modify its skeletal structure and joint attributes without interacting with the environment . In the execution stage , the agent assumes the design resulting from the transform stage and applies motor control actions to interact with the environment and receives rewards . Since the policy needs to be used across designs with a variable number of joints , we adopt graph neural networks ( GNNs ) as the policy ’ s main network architecture . Each graph node in the GNNs represents a joint and uses message passing with its neighbors to output joint-specific actions . While GNNs can improve the generalizability of learned control policies across different skeletons through weight sharing , they can also limit the specialization of each joint ’ s design since similar joints tend to output similar transform actions due to the weight sharing in GNNs . To tackle this problem , we propose to attach a joint-specialized multilayer perceptron ( JSMLP ) on top of the GNNs in the policy . The JSMLP uses different sets of weights for each joint , which allows more flexible transform actions to enable asymmetric designs with more specialized joint functions . The proposed Transform2Act policy jointly optimizes its transform and control actions via policy gradient methods , where the training batch for each iteration includes samples collected under various designs . Unlike zeroth-order optimization methods ( e.g. , ES ) that do not use a policy to change designs but instead mutate designs randomly , our approach stores information about the goodness of a design into our transform policy and uses it to select designs to be tested in the execution stage . The use of GNNs further allows the information stored in the policy to be shared across joints , enabling better experience sharing and prediction generalization for both design and control . Furthermore , in contrast to ES-based methods that do not share training samples across designs in a generation , our approach uses all the samples from all designs to train our policy , which improves sample efficiency . The main contributions of this paper are : ( 1 ) We propose a transform-and-control paradigm that formulates design optimization as learning a conditional policy , which can be solved using the rich tools of RL . ( 2 ) Our GNN-based conditional policy enables joint optimization of design and control as well as experience sharing across all designs , which improves sample efficiency substantially . ( 3 ) We further enhance the GNNs by proposing a joint-specialized MLP to balance the generalization and specialization abilities of our policy . ( 4 ) Experiments show that our approach outperforms previous methods significantly in terms of convergence speed and final performance , and is also able to discover familiar designs similar to giraffes , squids , and spiders . 2 RELATED WORK . Continuous Design Optimization . Considerable research has examined optimizing an agent ’ s continuous design parameters without changing its skeletal structure . For instance , Baykal & Alterovitz ( 2017 ) introduce a simulated annealing-based optimization framework for designing piecewise cylindrical robots . Alternatively , trajectory optimization and the implicit function theorem have been used to adapt the design of legged robots ( Ha et al. , 2017 ; 2018 ; Desai et al. , 2018 ) . Recently , deep RL has become a popular approach for design optimization . Chen et al . ( 2020 ) model robot hardware as part of the policy using computational graphs . Luck et al . ( 2020 ) learn a designconditioned value function and optimize design via CMA-ES . Ha ( 2019 ) uses a population-based policy gradient method for design optimization . Schaff et al . ( 2019 ) employs RL and evolutionary strategies to maintain a distribution over design parameters . Another line of work ( Yu et al. , 2018 ; Exarchos et al. , 2020 ; Jiang et al. , 2021 ) uses RL to find the robot parameters that best fit an incoming domain . Unlike the above works , our approach can optimize the skeletal structure of an agent in addition to its continuous design parameters . Combinatorial Design Optimization . Towards jointly optimizing the skeletal structure and design parameters , the seminal work by Sims ( 1994 ) uses evolutionary search ( ES ) to optimize the design and control of 3D blocks . Extending this method , Cheney et al . ( 2014 ; 2018 ) adopt oscillating 3D voxels as building blocks to reduce the search space . Desai et al . ( 2017 ) use human-in-the-loop tree search to optimize the design of modular robots . Wang et al . ( 2019 ) propose an evolutionary graph search method that uses GNNs to enable weight sharing between an agent and its offspring . Recently , Hejna et al . ( 2021 ) employ an information-theoretic objective to evolve task-agnostic agent designs . Due to the combinatorial nature of skeletal structure optimization , most prior works use ESbased optimization frameworks which can be sample-inefficient since agents with various designs in the population learn independently . In contrast , we learn a transform-and-control policy using samples collected from all designs , which improves sample efficiency significantly . GNN-based Control . Graph neural networks ( GNNs ) ( Scarselli et al. , 2008 ; Bruna et al. , 2013 ; Kipf & Welling , 2016 ) are a class of models that use message passing to aggregate and extract features from a graph . GNN-based control policies have been shown to greatly improve the generalizability of learned controllers across agents with different skeletons ( Wang et al. , 2018 ; 2019 ; Huang et al. , 2020 ) . Along this line , Pathak et al . ( 2019 ) use a GNN-based policy to control self-assembling modular robots to perform tasks . Recently , Kurin et al . ( 2021 ) show that GNNs can hinder learning in incompatible multitask RL and propose to use attention mechanisms instead . In this paper , we also study the lack of per-joint specialization caused by GNNs due to weight sharing , and we propose a remedy , joint-specialized MLP , to improve the specialization of GNN-based policies . 3 BACKGROUND . Reinforcement Learning . Given an agent interacting with an episodic environment , reinforcement learning ( RL ) formulates the problem as a Markov Decision Process ( MDP ) defined by a tupleM = ( S , A , T , R , γ ) of state space , action space , transition dynamics , a reward function , and a discount factor . The agent ’ s behavior is controlled by a policy π ( at|st ) , which models the probability of choosing an action at ∈ A given the current state st ∈ S . Starting from some initial state s0 , the agent iteratively samples an action at from the policy π and the environment generates the next state st+1 based on the transition dynamics T ( st+1|st , at ) and also assigns a reward rt to the agent . The goal of RL is to learn an optimal policy π∗ that maximizes the expected total discounted reward received by the agent : J ( π ) = Eπ [ ∑H t=0 γ trt ] , whereH is the variable time horizon . In this paper , we use a standard policy gradient method , PPO ( Schulman et al. , 2017 ) , to optimize our policy with both transform and control actions . PPO is particularly suitable for our approach since it has a KL regularization between current and old policies , which can prevent large changes to the transform actions and the resulting design in each optimization step , thus avoiding catastrophic failure . Design Optimization . An agent ’ s design D ∈ D plays an important role in its functionality . In our setting , the design D includes both the agent ’ s skeletal structure and joint-specific attributes ( e.g. , bone length , size , and motor strength ) . To account for changes in design , we now consider a more general transition dynamics T ( st+1|st , at , D ) conditioned on design D. The total expected reward is also now a function of design D : J ( π , D ) = Eπ , D [ ∑H t=0 γ trt ] . One main difficulty of design optimization arises from its bi-level nature , i.e. , we need to search over a large design space and solve for the optimal policy under each candidate design for evaluation . Formally , the bi-level optimization is defined as : D∗ = arg max D J ( πD , D ) ( 1 ) subject to πD = arg max π J ( π , D ) ( 2 ) The inner optimization described by Equation ( 2 ) typically requires RL , which is computationally expensive and may take up to several days depending on the task . Additionally , the design space D is extremely large and combinatorial due to a vast number of possible skeletal structures . To tackle these problems , in Section 4 we will introduce a new transform-and-control paradigm that formulates design optimization as learning a conditional policy to both design and control the agent . It uses first-order policy optimization via policy gradient methods and enables experience sharing across designs , which improves sample efficiency significantly . Graph Neural Networks . Since our goal is to learn a policy to dynamically change an agent ’ s design , we need a network architecture that can deal with variable input sizes across different skeletal structures . As skeletons can naturally be represented as graphs , graph neural networks ( GNNs ) ( Scarselli et al. , 2008 ; Bruna et al. , 2013 ; Kipf & Welling , 2016 ) are ideal for our use case . We denote a graph as G = ( V , E , A ) where u ∈ V and e ∈ E are nodes and edge respectively , and each node u also includes an input feature xu ∈ A . A GNN uses multiple GNN layers to extract and aggregate features from the graph G through message passing . For the i-th of N GNN layers , the message passing process can be written as : miu = M ( h i−1 u ) , ( 3 ) ciu = C ( { miv | ∀v ∈ N ( u ) } ) , ( 4 ) htu = U ( h i−1 u , c i u ) , ( 5 ) where a message sending module M first computes a message miu for each node u from the hidden features hi−1u ( h 0 u = xu ) of the previous layer . Each node ’ s message is then sent to neighboring nodes N ( u ) , and an aggregation module C summarizes the messages received by every node and outputs a combined message ciu . Finally , a state update module U updates the hidden state h i u of each node using ciu and previous hidden states h i−1 u . After N GNN layers , an output module P is often used to regress the final hidden features hNu to desired outputs yu = P ( h N u ) . Different designs of modules M , C , U , P lead to many variants of GNNs ( Wu et al. , 2020 ) . We use the following equation to summarize GNNs ’ operations : yu = GNN ( u , A ; V , E ) ( 6 ) where GNN ( u , · ) is used to denote the output for node u .

Summary: The paper poses the problem of morphology design for robots as RL training of one joint GNN policy. Their policy first generates the robot's morphology and then evaluates the design with a common behavior policy that conditions on it. The authors also introduce a technique called JSMPL to allow asymmetric morphologies. Experiments in the Mujoco simulator demonstrate a large improvement over evolutionary methods in terms of sample efficiency and final performance (although the latter is less clear, as neither method has clearly converged). Ablation studies show that the GNN architecture is essential, but are less clear about the impact of JSMPL.