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: This paper proposes a transform-and-control policy to optimize the robotic agents' designs. Contributions include: - A novel perspective on agent design: rather than formulating agent design as a bi-level optimization, this paper embeds both design generation and control into a single decision-making process such that both design and control are optimized by the same RL algorithm. - In this formulation, the training experience from different designs is shared to improve sample efficiency. - Joint-specialized MLP on top of the GNN policy that further finetunes the control of individual joints.