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 main idea of the paper**  This paper combines two powerful ideas, 1. unsupervised keypoint extraction for objects and 2. object-centric representation for forward modeling. KINET first detects objects' key points in the video in an unsupervised manner and extract object-centric representation for future prediction. Experiments are conducted in a new dataset built on MuJoCo and the performance in table 1 shows the proposed framework achieves better performance than the baseline methods.