Factored World Models for Zero-Shot Generalization in Robotic Manipulation

1 INTRODUCTION From assembly to household robots , current state-of-the-art robot learning agents can not generalize beyond a specific training task . One important aspect of generalization is the ability to understand any novel combination of known factors , a so-called combinatorial or compositional generalization . Applied to objects , combinatorial generalization ideally allows an agent to understand any arrangement of objects from only a limited number of interactions with its environment . The two key steps such agents need to perform are ( a ) decomposing a scene into individual objects and ( b ) modeling relative interactions between objects . The latter is particularly important when learning to predict the dynamics of the environment ( i.e . learning a world model ) . The purpose of our paper is to learn a world model that accurately predicts the effects of actions in the context of robotic manipulation , and generalizes to novel tasks . We assume the state of the environment is already factored into individual images for each object , and focus on learning a latent code for each object and predicting the effect of pick-and-place actions . Our model is trained using a self-supervised contrastive loss . It can predict accurate manipulation physics by using multiple graph neural network layers ( GNNs ) and action attention . GNNs consider pairwise interactions , usually between every possible pair of objects , in order to predict the state of each object one step into the future . GNNs achieve combinatorial generalization ( to the extent that pairwise interactions are expressive enough ) by equivariance to the order in which objects are presented–permutation equivariance . That is , they can not overfit to a particular ordering of objects . It is a common choice to use a single GNN layer ( in some cases called an Interaction Network ( Battaglia et al. , 2016 ) ) . But , as shown in previous work ( Kipf et al. , 2018 ; Sanchez-Gonzalez et al. , 2018 ) and confirmed by our experiments , stacking multiple GNN layers leads to much more accurate physics predictions . The ability to integrate information about actions with latent states of individual objects has been underexplored in prior work . Previous object-factored world models either assume object-action association ( a factored action space ) ( Veerapaneni et al. , 2019 ; Kipf et al. , 2020 ; Huang et al. , 2020 ) or only model sequences of states without actions ( Janner et al. , 2019 ; Bakhtin et al. , 2019 ; Qi et al. , 2021 ) . We represent actions as pick ( x , y ) and place ( x , y ) with continuous ( x , y ) coordinates ; this representation is integrated into our latent transition model through action attention and iterative refinement by a stack of residual graph neural networks . Our action attention module compares action information with the latent state of each object in order to predict the probability that an object is affected by the action . Each node in a GNN then receives a different action based on the predicted attention weights . Since each GNN in our transition model has access to the action , the responsibility for deciding which object is affected by the action and modeling its effect on said object can be distributed across the transition model . We train and evaluate our factored world model in a simulated environment involving a UR5 robotic arm manipulating blocks of various shapes ( Figure 1 ) . We instantiate two environments : Cubes includes six cubes arranged into five different structures , which take up to ten actions to build , and Shapes includes eight blocks of four types ( cube , brick , triangle , roof ) , which are arranged into sixteen different structures ( Figure 3 ) . To test zero-shot generalization — the main result of our paper — we partition the tasks into training and generalization sets . We ensure that the world model never sees the goal state of any generalization task during training . Nevertheless , our model transfers to the testing tasks in Cubes and Shapes with only a very minor drop in accuracy ( Table 1 ) . The mechanism behind successful zero-shot transfer is the ability of our model to understand the state of each object independently of others , and the ability to model sparse interactions . For example , if our model understands the interaction ” stack a triangle on top of a cube ” based on the first training task in Shapes ( Figure 3 , bottom row ) , it can use it when generalizing to four different held out tasks ( Figure 2 ) . In summary , we contribute the following : 1 . We develop a factored wold model for robotic pick-and-place tasks . This model does not require a known factored action-object association ( as was used in related models ) , owing to its use of an action-attention network . 2 . We demonstrate that the learned world model accurately predicts the outcomes of sequences of actions . Moreover , the model achieves zero-shot generalization to sequences of actions tasks and object configurations that were not seen during training , including sequences that are longer ( up to 10 actions ) than those in the training data ( up to 10 actions ) . Finally , we demonstrate transfer to a real-world dataset . 2 FACTORED WORLD MODEL FOR ROBOTIC PICK-AND-PLACE . The modeling task is to predict the effect of a sequence of actions a1 , a2 , ... , aT in an environment that is initialized to a state s0 . We will assume a setting in which the state of the world is represented as an image , which has been pre-processed into a factorized state s = 〈s1 , s2 , ... , sK〉 in which each si is an image centered on the ith object . The number of objects K can vary between episodes . This postulated factorization can be implemented using an object detection module that predicts a bounding box for each object and a tracking module that corresponds bounding boxes across time . Both problems are well-studied for natural images in computer vision ( Zhao et al. , 2019 ; Greff et al. , 2020 ) . Our goal is to learn a factorized latent representation z = 〈z1 , z2 , ... , zK〉 , along with a world model that predicts the state z′ that results from taking an action a at state z . We will show that the learned latent code is low-dimensional ; it represents object identity and position while abstracting over irrelevant features such as object color . Our world model is implemented using a graph neural network ( Gori et al. , 2005 ) that represents action-conditioned pair-wise interactions of latent factors . It is able to perform several rounds of message passing , each successively refining the prediction of the physics of picking and placing objects . We do not learn a mapping from z to s ( i.e . the world model only predicts the future within its latent code ) ; nevertheless , we show the latent transition model successfully solves prediction and action ranking tasks . Both the state encoder and the latent transition model are equivariant to the permutation of factors . That is , we can provide objects in any order as long as it is consistent within an episode . Moreover , the model is agnostic to the number of objects present in a scene — no parameters in the model are dependent on the number of objects — although the weights of the latent transition model might overfit if the number does not vary during training . We describe our world model in detail in Section 2.1 and the training procedure in Section 2.2 . 2.1 FACTORED VISUAL WORLD MODEL WITH ACTION ATTENTION . Encoder . Given a factored state s = 〈s1 , s2 , ... , sk〉 , the encoder fφ processes each factor ( image of an object ) separately . That is , z = 〈fφ ( s1 ) , fφ ( s2 ) , ... , fφ ( sK ) 〉 . The encoder uses three convolutional layers followed by average pooling and three fully-connected layers . In principle , the use of average pooling enables the encoder to process images of any size , although we use a fixed image size . As described in Section 3 , the per-object image si includes both an RGB component and a coordinate grid component , which marks the location where the crop was taken ( Figure 4 ) . Transition Model . The transition model gθ accepts a factored latent state z = 〈z1 , z2 , ... , zK〉 , zi ∈ RDz and an action a ∈ RDa . It predicts the next latent state by outputting a residual : ẑt+1 = zt + gθ ( z t , at ) . ( 1 ) We implement gθ ( zt , at ) as a stack of graph neural network layers with skip connections . Denoting the ith graph neural network layer as GNNi and the ith intermediate representation as yi , the computation with L layers is as follows : y1 = z + GNN1 ( z , a ) y2 = y1 + GNN2 ( y1 , a ) ... gθ ( z , a ) = yL−1 + GNNL ( yL−1 , a ) ( 2 ) Our implementation of individual graph neural network layers follows Kipf et al . ( 2020 ) : the network models pairwise interactions between latent state factors ( corresponding to objects in the environment ) using a fully-connected node network hn and an edge network he . Both are implemented as MLPs with one hidden layer . The edge network outputs an embedding for each directed edge ei , j = he ( zi , zj ) . The edge embeddings are then aggregated using the node network in order to update the state of each node . Finally , the graph neural network outputs a vector of updated factors z′ = 〈z′1 , z′2 , ... , z′K〉 with z′i = hn ( zi , a , ∑ j 6=i ej , i ) . ( 3 ) The summation over edges in the input of the node network assures permutation equivariance , an important property that aids generalization to novel combinations of factors . Given a permutation of factors π , we have GNN ( π ( z ) , a ) = π ( GNN ( z , a ) ) . ( 4 ) The permutation equivariance property also holds for a stack of GNN layers as well as the state encoder and action attention . Action Attention . When picking or placing objects in our experiments , the robotic arm and gripper will usually interact with only one or a small number of objects . The transition model needs to learn which latent factors change and which stay the same when an action is performed . Previous work simplified this problem by assuming that the action space is also factorized ( Sanchez-Gonzalez et al. , 2018 ; Veerapaneni et al. , 2019 ; Kipf et al. , 2020 ) , which is to say that actions are performed relative to individual objects or parts . This assumption is not realistic in typical robotics applications , where actions correspond to motor primitives that can not trivially be associated with effects on individual objects . We use pick ( x , y ) and place ( x , y ) actions with continuous ( x , y ) coordinates . We overcome this limitation by introducing an action attention layer . Instead of the action being the same for each node in the graph neural network , we create K transformations of the action that are attenuated based on how likely they are to interact with each latent factor . The mechanism follows single-head self-attention ( Vaswani et al. , 2017 ) with the following key k , query q and value v : k = 〈jk ( z1 ) , jk ( z2 ) , ... , jk ( zK ) 〉 , q = jq ( a ) , v = jv ( a ) . ( 5 ) Here , jk , jq and jv are Multi-Layer Perceptrons with one hidden layer . The factored action space a′ = 〈a′1 , a′2 , ... , a′K〉 is composed of v weighted by attention weights αi with i ∈ { 1 , ... , K } : αi = softmax ( kT1 q , k T 2 q , ... , k T Kq ) i , a′i = αiv . ( 6 ) Note the new actions are all equal up to a multiplicative factor . Figure 7 visualizes α .

Summary: This work proposes a pipeline to learn factored world models to predict the robot actions’ effects. This work proposes to leverage a graph neural network to extract state information and adopt the contrastive loss to train the encoder and latent transition model. This paper applies the proposed approach to pick-and-place tasks with simple shapes, and extensive experiments indicate its effectiveness.