KINet: Keypoint Interaction Networks for Unsupervised Forward Modeling

1 INTRODUCTION . Discovering a structured causal representation of the world allows humans to perform a wide repertoire of motor tasks such as interacting with objects . At the core of this process lies the ability to predict the response of the environment to applying an action ( Miall & Wolpert , 1996 ; Wolpert & Kawato , 1998 ) . Such an internal model , often referred to as the forward model , aims to come up with an estimation of future states of the world given its current state and an action . By cascading the predictions of a forward model it is also possible to plan a sequence of actions that would bring the world from an initial state to a desired goal state ( Wolpert et al. , 1998 ; 1995 ) . Recently , various deep learning architectures have been proposed to perform forward modeling using an object-centric representation of the system ( Ye et al. , 2020 ; Chen et al. , 2021b ; Li et al. , 2020 ; Qi et al. , 2020 ) . This object-centric representation is learned from the visual observation by factorizing the scene into the underlying object instances using ground-truth object states ( e.g. , object class , position , and bounding box ) . We identified two major limitations in the recent work : First , existing approaches either assume access to the ground-truth object states ( Battaglia et al. , 2016 ; Li et al. , 2020 ) or predict them using idealized techniques such as pre-trained object detection or instance segmentation models ( Ye et al. , 2020 ; Qi et al. , 2020 ) . However , obtaining ground truth objects state information is not feasible in practice . Relying on object detection and segmentation tools , on the other hand , makes the forward model fragile and dependent on the flawless performance of these tools . More often than not , pretrained object detection or segmentation models suffer from poor generalization to unseen objects . Second , factorizing the scene into object instances limits the generalization of forward models to scenarios with a different number of objects . In this paper , we address both of these limitations by proposing to learn forward models using a keypoint-based object-centric representation . Keypoints represent a set of salient locations of moving entities . Our model KINet ( Keypoint Interaction Network ) learns unsupervised forward modeling in three major steps : ( 1 ) A keypoint extractor factorizes the scene into keypoint coordinates with no supervision other than raw visual observations . ( 2 ) A probabilistic graph representation of the system is inferred where each node corresponds to a keypoint and edges are keypoints relations . Therefore , each object is represented in this graph by at least one distinct node . Node features carry implicit object-centric representation as well as explicit keypoint state information . ( 3 ) With probabilistic message passing , our model performs physical reasoning and learns an action-conditional forward model to predict future locations of the keypoints . Using this prediction , it also reconstructs the future appearances of the system . We evaluate KINet ’ s forward prediction accuracy and demonstrate that , by learning the physical reasoning in a keypoint coordinate , our model can effectively re-purpose this knowledge and generalize it to complex unseen circumstances . Our key contributions are as follows : ( 1 ) We introduce KINet , an end-to-end method for learning unsupervised action-conditional forward models from visual observations using keypoint coordinates . ( 2 ) We introduce Probabilistic Interaction Networks for efficient message-passing in fully connected graphs by learning to adaptively aggregate a subset of relevant neighborhood information . ( 3 ) We introduce the GraphMPC algorithm for accurate action planning based on graph similarity . ( 4 ) We demonstrate that learning forward models in keypoint coordinates enables zero-shot generalization to complex unseen scenarios . 2 RELATED WORK . Unsupervised keypoint extraction methods have been used successfully in computer vision tasks such as pose tracking ( Zhang et al. , 2018 ; Yao et al. , 2019 ) and video prediction ( Minderer et al. , 2019 ; Zhang et al. , 2018 ; Xue et al. , 2016 ) . Recent work explored keypoint extraction for control tasks in reinforcement learning environments to project the visual observation space to a lowerdimensional keypoint space ( Kulkarni et al. , 2019 ; Chen et al. , 2021a ; Jakab et al. , 2018 ) . We extend these methods for the object manipulation setting . Forward prediction models estimate the future state of a system given its current state . There is a large line of work on using neural networks to build such models for various applications . These approaches date back to Nguyen & Widrow ( 1990 ) ; Grzeszczuk et al . ( 1998 ) for modeling state transition in dynamic systems . The most fundamentally relevant approach to our model is the Interaction Network ( IN ) ( Battaglia et al. , 2016 ; Sanchez-Gonzalez et al. , 2018 ) and other follow-up work on using graph neural network in forward modeling ( Pfaff et al. , 2020 ; Li et al. , 2019 ; Kipf et al. , 2018 ; Mrowca et al. , 2018 ) . Although these approaches demonstrated powerful forward modeling capabilities , they often rely on building explicit object representations based on ground-truth state information of the objects in the system . In an attempt to bridge this gap several methods have been proposed to use IN-based models using visual features extracted from image observations . ( Watters et al. , 2017 ; Qi et al. , 2020 ; Ye et al. , 2020 ) . However , two main concerns remain unaddressed in these methods . First , the visual features of objects are often extracted from object bounding boxes using pretrained object detection or segmentation model ( Janner et al. , 2018 ; Qi et al. , 2020 ; Kipf et al. , 2018 ) that is either pretrained on the environment ( Qi et al. , 2020 ) or assume prior knowledge of the object position ( Ye et al. , 2020 ) . Second , implementing these approaches entails training the model on a fixed number of objects . Learning action-conditional forward models is an active challenge although the original IN incorporates the action as an external effect by aggregating the action vector to node and edge embeddings . For the case of probabilistic forward models Henaff et al . ( 2019 ) interestingly suggests using a latent variable dropout to ensure proper conditioning of the forward model on the action information following Gal & Ghahramani ( 2016 ) . In a more relevant application to our model , Yan et al . ( 2020 ) highlighted the effectiveness of contrastive estimation ( Oord et al. , 2018 ) to learn proper plannable object-centric representations . Graph representation inference methods has been originally introduced in variational graph autoencoders by Kipf & Welling ( 2016 ) . To formulate the graph structure inference module we were also inspired by Simonovsky & Komodakis ( 2018 ) where they use a variational graph autoencoder to generate small graphs of molecular structures . More relevant to our work is Kipf et al . ( 2019 ) in which an object-level contrastive loss has been used to learn object-centric abstractions in multibody systems . However , this work is only restricted to fixed environments with oversimplified objects such as two-dimensional geometries ( Watters et al. , 2019 ; Stanić et al. , 2020 ) . In our work , we consider experimenting with more complex three-dimensional objects where the objects are randomized and replicate real-world object manipulation . 3 KEYPOINT INTERACTION NETWORKS ( KINET ) . We assume access to observational data that consists of images , actions , and resulting images after applying actions : D = { ( It , ut , It+1 ) } where It , It+1 ∈ RH×W×C are images before and after applying action and ut ∈ R4 is the action vector . Our goal is to learn a forward model that predicts future states of the system with no supervision above the observational data . In this section , we describe our approach which is composed of two main steps : learning to encode visual observations into keypoint coordinates and then learning an action-conditioned forward model in the keypoint space to reason about the system and predict its future state . 3.1 UNSUPERVISED KEYPOINT DETECTION . The keypoint detector fkp is a mapping from visual observations to a lower-dimensional set of K keypoint coordinates { xkt } k=1 ... K = fkp ( It ) . The keypoint coordinates are learned through capturing the spatial appearance of the objects in the system in an unsupervised manner . The detector design is following the idea of Kulkarni et al . ( 2019 ) ; Jakab et al . ( 2018 ) , which we extend for object manipulation setting that necessitates consideration of external actions . Specifically , The keypoint detector receives a pair of initial and current image frames ( I0 , It ) and uses a convolutional encoder to compute a K-dimensional feature map for each image Φ ( I0 ) , Φ ( It ) ∈ RH ′×W ′×K . The expected number of keypoints in the system is captured by the dimension K. Next , each keypoint feature map is marginalized into a 2D keypoint coordinate { xk0 , xkt } k=1 ... K ∈ R2 . We use a convolutional image reconstruction model frec with skip connections to inpaint the current image frame using the initial image and the predicted keypoint coordinates Ît = frec ( I0 , { xk0 , xkt } ) . With this formulation , fkp and frec create a bottleneck to encode the visual observation in a temporally consistent lower-dimensional keypoint coordinate representation which is distributed across the visual observation . 3.2 GRAPH REPRESENTATION OF SYSTEM . The primary idea of our approach is representing the system as a graph . In particular , after factorizing the system into K keypoints , we build a directed graph Gt = ( Vt , Et , At ) where keypoints become the graph nodes and their pairwise relations become the graph edges . Keypoint positional and visual information are encoded into the feature embeddings of nodes { nkt } k=1 ... K ∈ Vt and edges { eijt } ∈ Et . We also use an adjacency matrix to further specify the graph connectivity as At ∈ R|V|×|V| where { eij } ∈ E if [ At ] ij = 1 . At each timestep t the observational data is encoded into the graph . Node embeddings are an encoding of the keypoint visual features and positional information extracted in the keypoint detector { nkt } = [ xkt , Φ k t ] . Edge embeddings contain a relative positional information of the sender and receiver nodes { eijt } = [ xit − x j t , ‖xit − x j t‖22 ] as suggested by Pfaff et al . ( 2020 ) . We assume that there is no self-loop in the graph nodes ( i.e , diagonal elements in the adjacency matrix is set to 0 ) . Note that the graph representation is built on keypoints coordinates and , hence , we do not impose any prior assumption on the number of objects in the system . 3.3 PROBABILISTIC INTERACTION NETWORKS . To build a forward model we extend the recent approaches based on Graph Neural Networks ( Battaglia et al. , 2016 ; Sanchez-Gonzalez et al. , 2018 ; Pfaff et al. , 2020 ) and propose a probabilistic variation of the Interaction Networks ( IN ) . The core of the probabilistic IN is generating node-level latent variables Zt ∈ R|V|×d ( Kipf & Welling , 2016 ; Simonovsky & Komodakis , 2018 ) . The node-level prior latent variables are independently sampled from a fixed isotropic Gaussian prior distribution p ( Zt ) . During training , the posterior distribution of the latent variable is obtained from the next timestep graph , qφ ( Zt| { nkt+1 } k=1 ... K ) = fenc ( Gt+1 ) ( 1 ) We use a probablistic decoder to derive the posterior probability of the adjacency matrix given the latent variable . To build this decoder , the existance of edges in the graph is modeled as Bernoulli variables with probabilities defined as , pθ ( [ At ] ij = 1|zit , z j t ) = σ ( fdec ( Gt , Zt ) ) ( 2 ) where σ ( · ) is the sigmoid function . We propose to use the posterior adjacency matrix distribution along with the latent variables to infer a probabilistic graph representation of the system G̃t = ( Ṽt , Ẽt , Ãt ) . In particular , the inferred adjacency matrix Ãt ∼ pθ ( At|Zt ) includes the edge probabilities . The node embeddings are also appropriately aggregated with the node-level latent variables { ñkt : [ nkt , zkt ] } k=1 ... K ∈ Ṽt . The probabilistic IN forward model Ĝt+1 = ffwd ( G̃t , ut ; Zt , Ãt ) predicts the graph representation at the next timestep by taking as input the current probabilistic graph and action . The message-passing operation in the forward model can be described as , { êij } ← fe ( ñi , ñj , ẽij ) , { n̂k } ← fn ( ñk , ∑ i∈N ( k ) êik , ut ) ( 3 ) where the edge-specific function fe first updates edge embeddings , then the node-specific function fn updates node embeddings using neighboring edge information . Note that the message-passing operation takes into account the node-level latent variables and the neighborhood aggregation N ( k ) is also performed based on the probabilistic adjacency matrix . Recent models for forward prediction often rely on fully connected graphs for message passing ( Qi et al. , 2020 ; Ye et al. , 2020 ; Li et al. , 2018 ) . Our model , however , learns to dynamically sample the neighborhood of each node at each timestep conditioned on the latent variable . Intuitively , this adaptive sampling allows the network to efficiently aggregate long-range context information only by selecting the most relevant neighboring nodes . This is specifically essential in our model since keypoints could provide redundant information if they are in very close proximity .

Summary: The paper proposes a keypoint-based forward dynamics model for generalization to unseen number of objects. The model employs unsupervised keypoint extraction methods from the literature, infers probabilistic graph representations over the keypoints, and makes forward predictions by message passing. The model is evaluated on multi-object manipulation tasks. It is shown that the model prediction error is close to that of a graph net with access to groundtruth object positions. The model is also shown to converge to the goal configuration faster than the baselines.