Unsupervised Object Keypoint Learning using Local Spatial Predictability

1 INTRODUCTION . An intelligent agent situated in the visual world critically depends on a suitable representation of its incoming sensory information . For example , a representation that captures only information about relevant aspects of the world makes it easier to learn downstream tasks efficiently ( Barlow , 1989 ; Bengio et al. , 2013 ) . Similarly , when explicitly distinguishing abstract concepts , such as objects , at a representational level , it is easier to generalize ( systematically ) to novel scenes that are composed of these same abstract building blocks ( Lake et al. , 2017 ; van Steenkiste et al. , 2019 ; Greff et al. , 2020 ) . In recent work , several methods have been proposed to learn unsupervised representations of images that aim to facilitate agents in this way ( Veerapaneni et al. , 2019 ; Janner et al. , 2019 ) . Of particular interest are methods based on learned object keypoints that correspond to highly informative ( salient ) regions in the image as indicated by the presence of object parts ( Zhang et al. , 2018 ; Jakab et al. , 2018 ; Kulkarni et al. , 2019 ; Minderer et al. , 2019 ) . Many real world tasks primarily revolve around ( physical ) interactions between objects and agents . Therefore it is expected that a representation based on a set of task-agnostic object keypoints can be re-purposed to facilitate downstream learning ( and generalization ) on many different tasks ( Lake et al. , 2017 ) . One of the main challenges for learning representations based on object keypoints is to discover salient regions belonging to objects in an image without supervision . Recent methods take an information bottleneck approach , where a neural network is trained to allocate a fixed number of keypoints ( and learn corresponding representations ) in a way that helps making predictions about an image that has undergone some transformation ( Jakab et al. , 2018 ; Minderer et al. , 2019 ; Kulkarni et al. , 2019 ) . However , keypoints that are discovered in this way strongly depend on the specific transformation that is considered and therefore lack generality . For example , as we will confirm in our experiments , the recent Transporter ( Kulkarni et al. , 2019 ) learns to prioritize image regions that change over time , even when they are otherwise uninformative . Indeed , when relying on extrinsic object properties ( i.e . that are not unique to objects ) one becomes highly susceptible to distractors as we will demonstrate . In this work , we propose PermaKey a novel representation learning approach based on object keypoints that does not overly bias keypoints in this way . The key idea underlying our approach is to view objects as local regions in the image that have high internal predictive structure ( self-information ) . We argue that local predictability is an intrinsic property of an object and therefore more reliably captures objectness in images ( Alexe et al. , 2010 ) . This allows us to formulate a local spatial prediction problem to infer which of the image regions contain object parts . We perform this prediction task in the learned feature space of a convolutional neural network ( CNN ) to assess predictability based on a rich collection of learned low-level features . Using PointNet ( Jakab et al. , 2018 ) we can then convert these predictability maps to highly informative object keypoints . We extensively evaluate our approach on a number of Atari environments and compare to Transporter ( Kulkarni et al. , 2019 ) . We demonstrate how our method is able to discover keypoints that focus on image regions that are unpredictable and which often correspond to salient object parts . By leveraging local predictability to learn about objects , our method profits from a simpler yet better generalizable definition of an object . Indeed , we demonstrate how it learns keypoints that do not solely focus on temporal motion ( or any other extrinsic object property ) and is more robust to uninformative ( but predictable ) distractors in the environment , such as moving background . On Atari games , agents equipped with our keypoints outperform those using Transporter keypoints . Our method shows good performance even on challenging environments such as Battlezone involving shifting viewpoints where the Transporters ’ explicit motion bias fails to capture any task-relevant objects . As a final contribution , we investigate the use of graph neural networks ( Battaglia et al. , 2018 ) for processing keypoints , which potentially better accommodates their discrete nature when reasoning about their interactions , and provide an ablation study . 2 METHOD . To learn representations based on task-agnostic object keypoints we require a suitable definition of an object that can be applied in an unsupervised manner . At a high level , we define objects as abstract patterns in the visual input that can serve as modular building blocks ( i.e . they are self-contained and reusable independent of context ) for solving a particular task , in the sense that they can be separately intervened upon or reasoned with ( Greff et al. , 2020 ) . This lets us treat objects as local regions in input space that have high internal predictive structure based on statistical co-occurrence of features such as color , shape , etc . across a large number of samples . Hence , our focus is on their local predictability , which can be viewed as an “ intrinsic ” object property according to this definition . For example , Bruce & Tsotsos ( 2005 ) have previously shown that local regions with high self-information typically correspond to salient objects . More generally , self-information approximated via a set of cues involving center-surround feature differences has been used to quantify objectness ( Alexe et al. , 2010 ) . In this paper we introduce Prediction ERror MAp based KEYpoints ( PermaKey ) , which leverages this definition to learn about object keypoints and corresponding representations . The main component of PermaKey is a local spatial prediction network ( LSPN ) , which is trained to solve a local spatial prediction problem in feature space ( light-blue trapezoid in Figure 1 ) . It involves predicting the value of a feature from its surrounding neighbours , which can only be solved accurately when they belong to the same object . Hence , the error map ( predictability map ) that is obtained by evaluating the LSPN at different locations carves the feature space up into regions that have high internal predictive structure ( see rows 4 & 5 in Figure 2 ( a ) ) . In what follows , we delve into each of the 3 modules that constitute our PermaKey system : 1 ) the VAE to learn spatial features 2 ) the LSPN to solve the local spatial prediction task and 3 ) PointNet to group error maps into keypoints ( see Figure 1 for an overview ) . Spatial Feature Embedding To solve the local spatial prediction task we require an informative set of features at each image location . While prior approaches focus on low-level features of image patches ( eg . RGB values Isola et al . ( 2014 ) , or ICA features Bruce & Tsotsos ( 2005 ) ) , we propose to learn features using a Variational Auto-Encoder ( VAE ; Kingma & Welling ( 2014 ) ) . It consists of an Encoder that parametrizes the approximate posterior qφ ( z|x ) and a Decoder that parametrizes the generative model pθ ( xi|z ) , which are trained to model the observed data via the ELBO objective : Le ( θ , φ ) = Eqφ ( z|x ) log [ pθ ( x|z ) ] −DKL [ qφ ( z|x ) ||p ( z ) ] . ( 1 ) The encoder is based on several layers of convolutions that offer progressively more abstract image features by trading-off spatial resolution with depth . Hence , which layer ( s ) we choose as our feature embedding will have an effect on the outcome of the local spatial prediction problem . While more abstract high-level features are expected to better capture the internal predictive structure of an object , it will be more difficult to attribute the error of the prediction network to the exact image location . On the other hand , while more low-level features can be localized more accurately , they may lack the expressiveness to capture high-level properties of objects . Nonetheless , in practice we find that a spatial feature embedding based on earlier layers of the encoder works well ( see also Section 5.3 for an ablation ) . Local Spatial Prediction Task Using the learned spatial feature embedding we seek out salient regions of the input image that correspond to object parts . Our approach is based on the idea that objects correspond to local regions in feature space that have high internal predictive structure , which allows us to formulate the following local spatial prediction ( LSP ) task . For each location in the learned spatial feature embedding , we seek to predict the value of the features ( across the feature maps ) from its neighbouring feature values . When neighbouring areas correspond to the same object- ( part ) , i.e . they regularly appear together , we expect that this prediction problem is easy ( green arrow in Figure 3 ) . In contrast , this is much harder when insufficient context is available ( red arrow Figure 3 : Two different prediction scenarios encountered when attempting to solve the LSP task . In the first case , too little context information is available to predict the center patch from its surroundings , likely yielding a prediction error ( assuming the character is not always present ) . Alternatively , when sufficient context is available as in the second case , prediction becomes easy . in Figure 3 ) , or when features rarely co-occur . Similarly , when parts of an image remain fixed across a large number observations it is easily predictable regardless of how much context is available . We implement the local spatial prediction task by training a neural network ( LSPN in Figure 1 ) to predict a center patch of features from the 8 first-order neighbouring patches using the learned spatial feature embedding at layer l of the VAE encoder . The weights of this network are trained to solve this prediction problem across images and across feature locations using the following objective : Lp ( ψl ) = 1 NHW ∑ x H , W∑ i , j ( fLSPNψl ( A [ ne ( i , j ) ] ) −A [ i , j ] ) 2 , ( 2 ) where A ∈ RH×W×C is the set of feature maps at layer l of the encoder qφ ( z|x ) , and ne ( i , j ) are the coordinates of the first-order neighbours surrounding i , j and H , W are the height and width of the spatial feature maps respectively . The features of the neighbouring locations are concatenated before being fed to the network . We train a separate LSPN for each of the chosen feature embedding specified by the encoder layer ( s ) and record their error at each location to obtain predictability maps that are expected to indicate the presence of objects ( rows 4 & 5 in Figure 2 ( a ) ) . Extracting Keypoints from Error Maps The predictability maps can be thought of as being a superposition of the local predictability errors due to each of the individual objects in the scene . Hence , what remains is to extract the k most salient regions ( belonging to object parts ) as our object keypoints from these predictability maps . We do so by training a PointNet ( Jakab et al. , 2018 ) to reconstruct the error maps through a bottleneck consisting of k fixed-width 2D Gaussian windows with learnable centers µ that are the keypoints ( Lg loss in Figure 1 ) as shown below : Lg ( ω ) = 1 MHW M∑ m H , W∑ i , j ( fPointNetω ( H [ m , i , j ] ) −H [ m , i , j ] ) 2 . ( 3 ) HereH ∈ RM×H×W is the concatenation of layer-wise LSP error maps for various encoder layer ( s ) l ( M in total ) , and H , W are the height and width of the error maps ( potentially up-sampled ) . Since the bottleneck layer has only limited capacity it is forced to consider locations that are most helpful in minimizing reconstruction error . This also requires the PointNet to consider statistical regularities across the predictability maps that can efficiently be described . In that sense this procedure can be thought of as a way of ( spatially ) clustering the local errors belonging to the same object by placing keypoints .

Summary: This paper works on unsupervised discovering keypoints in an Atari game frame to help improving Atari game performance. The keypoint discovery is based on predicting "predictable" local structure. I.e., the authors consider points that can not be predicted from its neighbor as good. Experiments show the learned keypoints performs better on 3 Atari games (Table. 1) than a counterpart keypoint discovery method, Transporter.