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: In this paper, the authors expand on previous works on contrastive learning with factored world models (specifically the C-SWM model by Kipf et al. - 2020) and apply their model for the task of predicting the dynamics of object manipulation with a continuous action space. Such previous works have shown that by leveraging residual graph neural networks (GNNs) to contextualize the learned latent states with possible actions, the model learns representations that are equivariant to permutation of objects in the scene, thus allowing scaling to a high number of objects regardless of the combinatorial explosion of the state space. Unlike previous works, where new objects are discovered by the model from raw pixels, the authors here use the simulation environment to segment objects appearing in the scene and factor them as part of the state representation. Another difference is that instead of factorizing the action space, they use two action primitives for picking and placing objects from given continuous coordinates and propose an action attention module that predicts how each action will affect each object in the next state. They claim that by departing from the standard action-object dependence assumption made by previous works, the current model can effectively model the effect of future actions for structured tasks unseen during training. They finally propose to stack multiple GNN layers instead of a single layer as mostly practiced by earlier works. The authors evaluate their methodology on two pick-and-place tasks regarding building structures of abstract object shapes or cubes. The objective is to correctly predict the effect of certain action sequences in other blocks in order to build a certain structure (e.g. pile). They sample noise from the correct action sequence in order to generate negative examples for contrastive learning. They train in some sequence configurations and test for zero-shot performance in unseen tasks (new structures of objects). They report results for both mean square error of the block position as well action ranking (how well the final state of the action sequence matches the final latent state).