Unsupervised Object Keypoint Learning using Local Spatial Predictability
Authors: Anand Gopalakrishnan
Sjoerd van Steenkiste
Jürgen Schmidhuber
Venue: ICLR 2021 Spotlight (2021)
Keywords: unsupervised representation learning
object-keypoint representations
visual saliency

Abstract: We propose PermaKey, a novel approach to representation learning based on object keypoints. It leverages the predictability of local image regions from spatial neighborhoods to identify salient regions that correspond to object parts, which are then converted to keypoints. Unlike prior approaches, it utilizes predictability to discover object keypoints, an intrinsic property of objects. This ensures that it does not overly bias keypoints to focus on characteristics that are not unique to objects, such as movement, shape, colour etc.  We demonstrate the efficacy of PermaKey on Atari where it learns keypoints corresponding to the most salient object parts and is robust to certain visual distractors. Further, on downstream RL tasks in the Atari domain we demonstrate how agents equipped with our keypoints outperform those using competing alternatives, even on challenging environments with moving backgrounds or distractor objects.


Published as a conference paper at ICLR 2021
UNSUPERVISED OBJECT KEYPOINT LEARNING USING LOCAL SPATIAL PREDICTABILITY
Anand Gopalakrishnan, Sjoerd van Steenkiste, Jürgen Schmidhuber The Swiss Lab IDSIA, USI, SUPSI (anand, soorrd, juergen) dsiid.iaa....
ABSTRACT
We propose PermaKey, novel approach to representation learning based on object keypoints. It leverages the predictability of local image regions from spatial neighborhoods to identify salient regions that correspond to object parts, which are then converted to keypoints. Unlike prior approaches, it utilizes predictability discover object keypoints, an intrinsic property of objects. This ensures that it does not overly bias keypoints to focus on characteristics that are not unique to objects, such as movement, shape, colour etc. We demonstrate the efficacy of Permakey on Atari where it learns keypoints corresponding to the most salient object parts and is robust certain visual distractors. Further, downstream RL tasks in the Atari domain we demonstrate how agents equipped with our keypoints outperform those using competing alternatives, even challenging environments with moving backgrounds or distractor objects.
INTRODUCTION
An intelligent agent situated in the visual world critically depends on 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 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. (219). Of particular interest are methods based on learned object keypoints that correspond to highly informative (salient) regions in the image 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 representation based on a set of task-agnostic object keypoints can be re-purposed to facilitate downstream learning (and generalization) many different tasks (Lake 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 neural network is trained to allocate fixed number of keypoints (and learn corresponding representations) in a way that helps making predictions about 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, 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 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 local regions in the image that have high internal predictive structure infomminoomaa..))... We argue that local predictability is an intrinsic property an object and therefore more reliably captures objectness in images (Alexe et al. 2010). This allows to formulate local spatial prediction

---

Published as a conference paper at ICLR 2021
Enooder
Decoder
+3E22
original image
Reconstructed Image
contect featuree
O
LSPN
min.
LBPN layar
layer
original patch
Layera
Maps
Masks
Figure 1: PermaKey consists of three modules (encapsulated by the dotted lines): learning a suitable spatial feature embedding (1), solving a local spatial prediction task (2), and converting error maps to keypoints (3). Objective functions used to learn each of the 3 modules shown within dotted blocks.
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 PointNer (Jakab et al. 2018) these predictability maps to highly informative object keypoints.
then convert
We extensively evaluate approach number of Atari environments and compare to Trans- porter (Kulkarni et al. 2019). We demonstrate how our method is able to discover keypoints that focus image regions that unpredictable and which often correspond to salient object parts. By leveraging local predictability to learn about objects, our method profits from 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 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 tesk-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.
METHOD
To learn representations based on task-agnostic object keypoints we require a suitable definition an object that can be applied in unsupervised manner. At high level, define objects abstract patterns the visual input that can serve as modular building blocks (i.c. they are self-contained and reusable independent of context) for solving particular task, in the intervened upon or reasoned with (Greeff et al.) 2020). This lets us treat objects as local regions in input space that have high internal predictive structure based 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 example, BBruce & Tsotsos (2055) have previously shown that local regions with high self informfootttoon typically correspond to salient objects. More generally, self self-information approximated via of cues involving center-surround feature differences has been used to quantify objectness (Alexe et al. 2010).
that they can be separately
"intrinsic" object property according this definition. For
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 trained to solve a local spatial prediction problem in feature space (light-blue trapezoid in Figure 1). It involves predicting the value of 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, delve into each of the 3 modules that constitute
2

---

Published as a conference paper at ICLR 2021
Vahwaman
Erier Pouge Laver
(a) Atari
(b) Noisy Atari
Figure 2: Image (row 1), Transporter keypoints overlaid on image (row 2), PermaKey (ours) keypoints
overlaid on image (row 3), predictability maps from feature layers 0 and (rows 4 & 5 respectively).
our PermaKey system: 1) the VAE learn spatial features 2) the LSPN solve the local spatial prediction task and 3) PointNet to group error maps into keypoints (see Figure I for an overview).
Spatial Feature Embedding To solve the local spatial prediction task require informative set of features at each image location. While prior approaches focus on low-level features of image patches (eg. RGB values [sola al. (2014), or ICA features Bruce Tsotsos (2055) propose to learn features using Variational Auto-Encoder (VAE; Kingma & Welling (2014)). It consists of Encoder that parametrizes the approximate posterior (2 x)) and a Decoder that parametrizes the generative model pe 22).).. which are trained to model the observed data via the ELBO objective:
L. (6, の) [=)(+(=)) logips [z|2)] (xL[go
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 our feature embedding will have an effect on the outcome of the local spatial prediction problem. While more ab- stract 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 expres- siveness to capture high-level properties of objects. Nonetheless, in practice we find that spatial fea- ture embedding based on earlier layers of the encoder works well (see also Section 5.3 for 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.c. 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
3

---

Published as a conference paper at ICLR 2021
Figure 3: Two different prediction sce- narios encountered when attempting to solve the LSP task. In the first case, too little context information is available to predict the center patch from sur- roundings, likely yielding prediction error (assuming the character is not al- ways present). Alternatively, when suffi- cient context is available as the second case, prediction becomes easy.
Insufficiett context Predicttoo
Redundant Information Prediction easy
in Figure 3). or when features rarely co-occur. Similarly, when parts of an image remain fixed across a large number observations is easily predictable regardless of how much context is available.
We implement the local spatial prediction task by training neural network (LSPN in Figure to predict center patch of features from the 8 first-order neighhouring patches using the learned spatial feature embedding at layer I 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:
H.W Σ
Lp((h) NHW
Σ E
(L (A(ne((i))))) Ali,,)))
where A E HxXxC is the set of feature maps layer 1 of the encoder 4s (2(2z, and ne(i, the coordinates of the first-order neighbours surrounding i. and H, W are the height and width of the spatial feature maps respectively. The features of the neighbouring locations before being fed to the network. We train 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)
concatenated
Extracting Keypoints from Error Maps The predictability maps can be thought of as being superposition of the local predictability errors due to each of the individual objects in the scene. Hence, what remains is to extract the most salient regions (belonging to object parts) as our object keypoints from these predictability maps. We do so by training PointNet (Jakab et al. (0188) to reconstruct the error maps through a bottleneck consisting fixed-width 2D Gaussian windows with learnable centers μ that are the keypoints (Ly loss in Figure as shown below:
Lg(w) MHW Σ
Σ
FL PointNet (H[m, (, H.m,i,)))..
i.)
(3)
Here H E MxXxx is the concatenation of laser-wise LSP error maps for various encoder layer(s) / (M in total), and H, W are the height and width of the error maps (potentially up- up-sampled). Since the bottleneck layer has only limited capacity it is forced to consider locations that are most helpful in min- imizing 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.
USING KEYPOINT REPRESENTATIONS FOR DOWNSTREAM PROCESSING
Given a set of keypoint locations and their associated features learned an purely unsupervised framework we examine how to incorporate them into state representation suitable to solve set of downstream RL tasks. Here we will consider Atari games as an example, which essentially consist of entities such player, enemies, reward objects like health potions, food etc. Player(s) obtain positive reward through collecting these reward objects or killing enemies and in turn risk getting attacked by these enemies. Successful game-play in such environments would therefore require understanding of these interaction effects between player, enemies and reward objects in order to estimate the best set of actions to take maximize the game score.
4

---

Published as a conference paper at ICLR 2021
Prior work like the Transporter (Kularni et al. 2019) uses CNN to encode features of object keypoints for downstream RL tasks on the Atari domain. However, a CNN permits an implicit encoding of relational factors between keypoints only in limited sense since it measures interactions between spatially close keypoints through the weight-sharing scheme of the convolutional kernel. An alternative is to use graph neural network (GNN) (Scarselli et al.. [209 BBattaglia et al.] 2018) (see also Pollack| (1990); Küchler& & Goller (196)) for earlier approaches) which have been applied to model relational inductive biases in the state representations of deep RL agents (Zambaldi et al. 2019). In a similar spirit, we consider treating keypoints as an unordered set and using GNN to explicitly encode relational factors between pairs of such entities. In this case, we spatially average the convolutional features within each keypoint mask to obtain keypoint feature. These are appended with learned positional embedding and initialized as node values of the underlying graphs. Underlying graphs are fully-connected and edge factors initialized with zeros. We use the Sicode-Prroees Deecese design strategy (Battaglia et al. 2018) for the graph keypoint encoder in our PKey+GNN model for Atari in (Section 5.1] The Encode and Decode blocks independently process nodes and edges of underlying graphs. We use an Interaction Network (Battaglia et al.] (2016) for the Process block with single message- saassggg step to compute updated edges e' and nodes v' as shown below:
(4) Here Gij is the edge connecting nodes o, and U/, f is the edge update function, pr is the edge aggregation function, Ej represents all incoming edges to node v; and S is the node update function. After one round of pessaag-passsing by the Interaction Network the resultant nodes are decoded as outputs. These output node features of all nodes are then concatenated and passed through an MLP before being input to the agent. For further implementation details we refer to Appppddix A.2.3
c f eiy),
=
4 RELATED WORK
Inferring object keypoints offers a flexible approach to segmenting a scene into a set of informative parts useful for downstream processing. Recent approaches (Jakab et al., 2018 Zhang al. 2018 Kulkarni et al. 2019 Minderer et al. 2019) apply a transformation (i.e. rotation, deformation, view- point shift or temporal- -shift) on source image to generate target image. An autoencoder network with a keypoint-bottleneck layer is trained to reconstruct the target image given the source image and thereby learns to infer the most salient image regions that differ between source and target image pairs (Jakab et al.) 2018). However, since the initial transformation does not focus on intrinsic object properties, they are more susceptible to 'false- cosst-postives' and 'false-negatives.. In extreme scenarios, such approaches may even fail to infer any useful object keypoints at all (e.g. as shown in Figure 2(h)).
bottleneck
The more general framework of discovering salient image regions (Itti et al. 1998) is closely related to inferring object keypoints. Early work explored the connections between information theoretic mea- sures of information gain or local entropy (Kadir & Brady 200T Jagersand| 1995) and its role in driv- ing fixation( attention) patterns in human vision. For example, [Bruce& & Tsotsos (2005) use Shannon's self information metric as a measure of saliency, while Zhmoginov et al. (2019) leverage mutual infor- mation to guide saliency masks. An interesting distinction is between local and global saliency, where the former is concerned with (dis) (diiisiiil between local image regions and the latter focuses on how (in)frequently information content occurs across an entire dataset (Borji & Ittil 2012).. Our prediction objective accounts for both: the LSP account for local saliency, while more global saliency is achieved having the prediction network solve this local prediction problem across an entire dataset of images. In this way, it can also reason about statistical regularities that only occur at this more global level.
An alternative to inferring object keypoints (and corresponding representations) is to learn entire object representations. Recent approaches can be broadly categorized into spatial mixture models (Greff et al., 2017 2019 Burgess et al., 2019), sequential attention models (Eslami et al., 2016) Kosiorek et al.) 2018) or hybrid models that are combinations of both (Lin et al.) 2020). While these methods have shown promising results, they have yet to scale to more complex visual settings. In contrast, object keypoints focus on more low-level information (i.e. object parts) that can be inferred more easily, even in complex visual settings (Minderer et al., 2019) Recent work on object detection (Duan et al. 2019 Zhou et al., 2019) have explored the use of center point (i.e. keypoint) and a bounding box as an object abstraction. However, they use ground ground-truth bounding boxes of objects infer keypoint locations whereas our method is fully unsupervised. Alternatively Isola et al.
5

---

Published as a conference paper at ICLR 2021
(2015) propose self-supervised framework for instance segmentation where local image patches classified as belonging together based on their spatial or temporal distance.
Graph neural networks (GNNs) have been previously used for intuitive physics models (Battaglia et al., 2016 Chang et al. [017]] van Steenkiste et al. 2018 Janner et al., 2019, Veerapaneni et al. 2219]]111] Kipl al. 2020 Stanié al., 2021). Other work instantiate the policy (Wang et al. 2018 value function (Bapst et al. 2019 as GNNs for RL tasks on continuous control or physical construction domains. However, the latter assume access underlying ground-truth object(-part) information. In contrast, we learn our keypoint representations in fully unsupervised manner from raw images. Further, we do not instantiate the policy or value function as a GNN, but rather use it to model relational factors between learned keypoint for downstream RL tasks in similar spirit to Zambaldi et al. (2019).
EXPERIMENTS
We empirically evaluate the efficacy of our approach on subset of Atari games (Bellemare et al. 2013) that ensure a good variability in terms of object sizes, colour, types, counts, and move- ment. As baseline, compare against the recently proposed Transporter (Kulkarni et al. 2019). which was previously evaluated on Atari. To evaluate the efficacy of object keypoints to facilitate sample-efficient learning on RL tasks we train agents the low-data regime of 100,000 interac- tions following the same experimental protocol as in Kulkarni et al.
(2019): Kaiser et al. 2020)
We use a deep recurrent Q-learning variant (Hausknecht & Stone 2015 using
double-Q learning (Hasselt et al. 2016). target networks (Mnih et al. (015
Frostbite
and a 3-step return as our RL algorithm for all agents with a window size of 8 and batch size of 16. To implement the Q-function we use 128 unit LSTM
network (Hochreiter & Schmidhuber 1997 Gers 1999). Full experimental
Image
details can be found in Appendix Alan and additional results in Appendix B
5.1 ATARI
Keypoints (ours)
We train Permakey on Atari and find that it is frequently able to recover key- points corresponding to object parts. For example, in Figure 2(a) 3) it can be seen how our method discovers keypoints correponding to the road and traf- fie in Enduro (column 3) and player and enemy tank parts (eg. wheels, cannon) in Battlezone (column 4). On Frostbite we find that it learns about the borders of the ice platforms, the player, scoreboard and birds, while on MsPaeman we observe that it tracks the various characters like the blue ghosts and Pac-Man.
Keyooint Regions
We emphasize that while the precise location of the keypoints often seems to focus on the borders of predictability, their associated window of attention (i.e.
in PointNer) is able to capture most of the object part. This is better seen in Figure 4] (and in Figure 10 in Appendiix B). where the corresponding windows of attention (available to an agent) are shown. From the error maps (rows 4 and in Figure 2(a) it can be seen how the LSPN implicitly learns to discover salient parts of the image, and how the choice of feature embedding shifts the focus from local edges to entire objects. In Figure 5] (igures 8 and 9 it is shown how PermaKey is stable across multiple different runs.
Figure 4: Keypoint regions highlighted.
Compared to Transporter (Figure 2(a) row 2) we find that our approach often produces qualitatively better keypoints. For example, on Enduro and Batrlezone the Transporter fails to place keypoints on salient object parts belonging the road and traffic or parts of the player and enemy tank (eg. its canon, wheels etc.). This is intuitive, since in these environments moving image regions do not always correspond to objects. Indeed, on Frosthite and MsPacman, where this is not the case, the keypoints produced by Transporter and our method are of comparable quality.
In Table we provide quantitative results of training agents that use these keypoints on 4 Atari environments. We observe that agents using PermuKey keypoints (PKey-CNN in Table 1) outperform corresponding variants using Transporter keypoints (original from Kulkarni et al. (2019 and our
"Code reproduce all experiments is available at
42/ 42/pernakey

---

Published as a conference paper at ICLR 2021
OSSALL
DOLOA
Figure 5: Observed Enduro frames (odd columns) and with overlaid keypoints from 5 random seeds having a unique color (even columns). The inferred keypoints stable across different runs and consistently focus on salient object parts. Additional frames can be seen in Figure 8 in Appendix
Game
Rainbow
PPO
SimPLE
Transporter (orig.)
Transporter (re-imp.)
Transporter -GNN
PKcy-CNN
PKcy-GNN
(ours)
(ours)
3363.5
5300.0
5184.4
N/A
N/A
N/A
10266.7
12566.7
Battlezone
(523.8)
(3655.1)
(1347.5)
(3272.1)
(3297.9)
364.3
496.0
762.8
999.4
983.0
838.3
1038.5
906.3
MsPacman
(20.4)
(379.8)
(331.5)
(145.4)
(806.3)
(537.3)
(417.1)
(382.2)
206.3
370.0
370.9
236.7
455.3
296.0
520.0
375.3
Seaquest
(17.1)
(103.3)
(128.2)
(22.2)
(160.6)
(82.5)
(159.9)
(75.4)
Frostbite
140.1
174.0
254.7
388.3
310.7
657.3
(2.7)
(40.7)
(4.9)
(142.1)
(14.9)
(442.1)
(150.1)
(556.8)
Table 1: Atari Results: mean scores (and std. dev.) obtained for each agent. For PKey and Transporter (re-imp.) we report the mean for 3 different policies (using the median evaluation run over 10 environment seeds). Results for Transporter (orig.) are taken from Kulkarni et al. (2019) and for SimPLE, PPO, and Rainbow from (Kaiser et al. 20200,. where a similar experimental set-up used. "N/A" indicates that policy did not make any learning progress w.r.t a random policy.
re-implementation) on 3 environments Notice how on Battlezone we were unable to train a policy that is better than random when using Transporter keypoints. This is intuitive, since in Figure 2(a) we observed how Transporter learns keypoints that are not concerned with salient objects on this environment. For comparison,
also include scores from Kaiser al. (2020) in Table I where similar experimental setup is used. It can be seen that SimPLE (Kaiser et al. 2020). Rainbow (Hessel
al. 2018) and PPO (Schulman al., 2017) also achieve worse scores than our PKey agents.
When using GNN for downstream processing (PKey+GNN) we are able to obtain additional im- provements on the Battlezone and Frostbite environments. Meanwhile on MsPacman and Seaquest performance drops, indicating that the stronger spatial inductive bias (i.e. nearby keypoints are processed similarly) implemented by CNNs can be advantageous as well. Note that the relative improvements of (Transporter+GNNY compared to (Transporter+CNN)) largely mirrors the relative changes shown by analogous PermaKey variants (CNN vs GNN keypoint encoder) across all envi- ronments. Further, can be seen how PermaKey consistently outperforms Transporter, using GNN keypoint encoders for both. On MsPacman we observed a larger standard- for Transporter, compared the results reported by Kulkarni et al. (2019), although did tune the hyper-parameters for our re-implementation separately. A possible reason could be due to us using lower batch size of 16 (compared 32) and lower window size of (compared 16). We were unable to investigate this further due to constraints on the available amount of compute. On Frostbite.
when
observe large variances for both (PKey+GNN) and (Transporter+ agents as some policies complete the first level (resulting in very high scores of about 2000 points), while others learn sub-optimal strategies yielding episodic returns around 220.
5.2 ROBUSTNESS TO DISTRACTORS
In order to demonstrate the limitations of the explicit motion bias in Transporter we synthesized a modified version of Atari ("noisy" Atari) that contains coloured strip (either horizontal, vertical or both) at random locations in the environment (see row 1 in Figure 2(b)), thereby creating the illusion
Unfortunately, Kulkarni et al. 2019 have not open-sourced their code (needed to reproduce RL results and apply Transporter other environments) and their self-generated dataset containing 'ground-truth' keypoints for evaluation Atari. We able to obtain certain hyper-parameters through personal communication.
7

---

Published as a conference paper at ICLR 2021
Game
Transp. (re-imp.)
PKey + CNN 562.0 (114.78)
Table 2: "Noisy" Atari Results: mean (and std dev.) for different policies (using the median over environment seeds as before).
Seaquest
357.3 (99.10)
MsPacman
470.0 (212.90)
574.7 (290.07)
MsPacman
5 keypoints
7 keypoints
10 keypoints
Table 3: Keypoint ablation on the MsPacman env.
Transp. (re-imp.)
(433..9))
983.0 (806.3)
907.0 (317.41)
PKey CNN
1004.3 (319.15)
1038.5 (417.1)
1003.3 (313.07)
of motion. Figure 2(b) (row 3) demonstrates how our method is more resistant against this distractor by focusing on predictability and able to recover keypoints belonging to object parts as before (eg. player, skull, ladder positions in Montezuma's 's Revenge). Indeed, notice how the distractor does not show up in the error maps (rows 4 and 5 in Figure 2(b) as it corresponds to a predictable pattern that was frequently encountered across the entire datase In contrast, find that Transporter (row 2) dedicates many of the available keypoints to the distractor as opposed to the static salient object parts.
As quantitative experiment, consider the downstream performance of agents that receive either Transporter keypoints or PermaKey keypoints on "Noisy" versions of Seaquest and MsPacman. In Table it can be seen how the performance of the PermaKey agent (PKey + CNN) remains largely unaffected on Seaquest (compared to TTble TT. while the sensitivity of Transporter to such moving distractors causes a significant drop in performance. On "Noisy" MsPacman, both agents perform worse compared to before, but the PermaKey only to a lesser extent. We hypothesize that the general performance drop is caused by the distractor bars occluding the fruits (collecting fruits yields rewards) located along the horizontal and vertical columns at every timestep. In general, we observe that Permakey consistently outperforms Transporter on these tasks.
5.3 ABLATION STUDY
Number of keypoints We gain additional insight our approach by modulating the number of keypoints and observe its qualitative effects. On the Frostbite environment use 15, 20, 25 keypoints (shown in Figure 6). while keeping all the other hyperparameters same as described in [Appendix A When using additional keypoints we observe that Permakey prioritizes regions in the error map that contribute most to the reconstruction loss. On the other hand, when the number of keypoints is too large, redundant keypoints are placed seemingly random locations e.g. the border of the image.
We also quantitatively measure the effects of varying the number of keypoints (5, 7, 10} produced by PermaKey and Transporter downstream performance in MsPacman (Table 3). using the same evalu- ation protocol as before. The agent using Permakey outperforms the one using Transporter keypoints across the entire keypoint range. Since the Transporter keypoints are biased towards image regions that change across frames, additional keypoints only seem to track the different characters that move around the maze. With PermaKey focusing on predictability, additional keypoints tend to capture other salient parts of the maze, although this did not yield improved downstream performance in this case.
Spatial resolution of the feature embedding We experiment with the choice of feature layer(s) used for the local spatial prediction task and observe its qualitative effects on the keypoints discovered. On the Space Invaders environment we use the following sets of feature layer(s) {(0). (o, 1). (2. 3). (0, 1, 3.)} and retain the same values for all the other hyperparameters as described in Appendix A] Results are shown in Figure 7in Appendix B When only the lower layers are used, we find that the corresponding error maps tend to focus more on highly local image features. In contrast, when only higher layers are used, we find that the loss of spatial resolution leads to highly coarse keypoint localization. In practice, we found that using feature layers 0 and offers the right balance between spatial locality and expressiveness of features.
The ice floats in Frostbite, which are captured as keypoints (Figure 2(a)), serve as a useful example that illustrates the importance of integrating multiple spatial resolutions (i.e. that trade off local and global information as also noted in Itti et al. [998); Jagersand| (1995)). Ice floats seen at the default image resolution (layer-0 with 'cony" receptive field=4) might seem very predictable in it's highly local context of other ice float parts. However, casting the same local prediction task of predicting
"This observation also confirms that the LSPN does not collapse to naive color thresholder.
8

---

Published as a conference paper at ICLR 2021
the ice float features at layer-1 (image resolution 0.5x due to stride=2 in conv-encoder) given only blue background context would lead to some error spikes. This is because at many locations in the image, it seems reasonable to expect blue background patches to surround other blue background patches. Only at certain locations white ice floats are interspersed with blue background patches thereby driving some amount of error (surprise) in local predictability (In Figure 2(a] column 1, compare rows 4&5). Hence, varying the spatial resolution of the feature embedding and integrating the corresponding error maps lets us accommodate objects of varying shapes, sizes and geometry etc.
6 DISCUSSION
We have seen how PermaKey, by focusing on local predictability, is able to reliably discover salient object keypoints while being more robust certain distractors compared Transporter (Kulkarni et al. 2019). However, it should be noted that Permakey can only account for a "bottom-up" notion of saliency (ltti 2007) based on the observed inputs, and does not consider dask-dependent information. Therefore a potential limitation of this approach is that it might capture objects that are not relevant for solving the task at hand (i.e. corresponding to "unpredictable distractors"). On the other hand, a purely unsupervised approach to extracting keypoints may allow for greater re-use and generalizability. The same (overcomplete) set of keypoints can now be used to faciliate a number of tasks, for example by querying (and attending to) only the relevant keypoints using top-down attention (Mott al. 2019] Mittal et al. 2020). In this way it may readily facilitate other tasks, such as a "new object" acting as the key in Montezuma's 's Revenge, without having to re-train the bottom-up vision module.
Other kinds of "umpredictable distractors" could include certain types of image noise having highly unpredictable local structure, such salt-and-pepper noise. It is unclear how this affects the local spatial prediction problem, since the learned spatial embedding of the VAE may readily suppress such information (due it being trained for reconstruction). In that case, since Permakey acts on these learned spatial embeddings, it could leave the quality of error map(s) relatively unchanged. Further, we note that such distractors are typically the result of artificial noise and uncommon in the real world, unlike distractors due to motion, which are therefore more important to handle effectively.
In this work, we have trained Permakey in a purely unsupervised fashion by pre- training once on a dataset of collected episodes of game- game-play. On games where new objects are introduced only at much later levels would be beneficial to continue refining the keypoint module using the unsupervised loss (Kulkarni et al.) 2019). Further, it may be interesting to consider whether there are benefits to learning both the keypoint module and policy network in an end-to-end fashion. While this is expected to reduce the re-use and generalizability of the discovered keypoints, it may make it casier to learn relevant keypoints due to directly considering task-specific information (e.g. via policy gradients).
CONCLUSION
We proposed PermaKey, novel approach to learning object keypoints (and representations) that leverages predictability to identify salient regions that correspond to object parts. Through extensive experiments Atari it was shown how our method is able to learn accurate object keypoints that are more robust to distractors, unlike the Transporter baseline approach. RL agents that employ Permakey keypoints as input representations show sample-efficient learning on several Atari environments and outperform several other baselines. Further, it was shown how additional improvements can sometimes be obtained when processing keypoints using graph neural networks. Future work includes scaling PermaKey to more complex 3D visual worlds involving changing viewpoints, occlusions, egocentric vision etc. Similarly, designing quantitative evaluation metrics to directly measure the quality of keypoints extracted by different methods without the need for computationally expensive downstream tasks would allow for efficient performance analysis and model design as we begin to scale-up these ideas to more complex visual domains.
ACKNOWLEDGMENTS
We thank Aleksandar Stanié and Aditya Ramesh for valuable discussions. This research was partially funded by ERC Advanced grant no: 742870 and by Swiss National Science Foundation grants: 200021 165675/1 & 200021 192356. We also thank NVIDIA Corporation for donating a DGX-1 as part of the Pioneers of AI Research Award and IBM for donating a Minsky machine.

---

Published as a conference paper at ICLR 2021
REFERENCES
B. Alexe, Thomas Deselaers, and V. Ferrari. What is an object? Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 73-80, 2010.
V. Bapst, Alvaro
Doersch, Kimberly L. Stachenfeld, Pushmeet Kohli, P. Battaglia, and Jessica B. Hamrick. Structured agents for physical construction. In ICML, 2019.
Horace Barlow. Unsupervised learning. Neural computation, 1(3):295-311, 1989.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. In Advances in neural information processing 2016.
systems, pp.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonnzlez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint 186.:0800001211 2018.
Mare Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. J. Artif. Int. Res., 47(1):253-279, May 2013. ISSN 1076-9757.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions pattern analysis and machine intelligence, 35(8):1798-1822 2013.
Ali Borji and Laurent Itti. Exploiting local and global patch rarities for saliency detection. 2012 IEEE Conference Computer Vision and Pattern Recognition, pp. 478 178-485,2 2012.
Neil D. B. Bruce and John K. Tsotsos. Saliency based information maximization. In Proceedings of the 18th International Conference on Neural Information Processing Systems, NIPS' 05, pp- MIT Press.
155-162, Cambridge, MA, USA, 2005.
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. arXiv preprint arXiv: 11300 2019.
Michael Chang, Tomer Ullman, Antonio Torralba, and Joshua B. Tenenbaum. compositional object-based approach to learning physical dynamics. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings, 2017.
Kaiwen Duan, Song Bai, Lingxi Xie, Honggang Qi, Qingming Huang, and Qi Tian. Centernet: Keypoint triplets for object detection. Proceedings of the IEEE International Conference on Computer Vision, pp. 6569 607866578827 2019.
SM Ali Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E. Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. In Advances in Neural Information Processing Systems, pp. 3225-3233, 2016.
F.A. Gers, Learning to forget: continual prediction with Istm. IET Conference Proceedings, pp. 850-855(5), January 1999.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. In Advances in Neural Information Processing Systems, pp. 6691-6701, 2017.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2424-22222 2433, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
10

---

Published as a conference paper at ICLR 2021
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05208, 2020.
H. V. Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning, In AAAI, 2016.
Matthew J. Hausknecht and P. Stone. Deep recurrent q-learning for partially observable mdps. In AAAI Fall Symposia, 2015.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. Sheila A. Mcllraith and Kilian Q. Weinberger AAAI Conference Artificial Intelligence, (AAAI-18),
(eds.), Proceedings of the New Orleans, Louisiana, USA, February 2-7, pp. 3215-3222. AAAI Press, 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1997.
Phillip Isola, Daniel Zoran, Dilip Krishnan, and Edward Adelson. Crisp boundary detection using pointwise mutual information. In European Conference on Computer Vision, pp. 199-814. Springer, 2014.
Phillip Isola, Daniel Zoran, Dilip Krishnan, and Edward H Adelson. Learning visual groups from co-occurrences in space and time. arXiv preprint arXiv:151.00611,, 2015.
L. Itti. Visual salience. Scholarpedia, 2(9):3327, 2007. doi: 10.42499scholarpedia.3333 revision #72776.
Laurent Itti, Christof Koch, and Ernst Niebur. model of saliency-based visual attention for rapid scene analysis. IEEE Transactions pattern analysis and machine intelligence, 20(11): 1259-1259, 1998.
Martin Jagersand. Saliency maps and attention selection in scale and spatial coordinates: An information theoretic approach. In Proceedings of IEEE International Conference on Computer Vision, pp. 195-202. IEEE, 1995.
Tomas Jakab, Ankush Gupta, Hakan Bilen, and Andrea Vedaldi. Unsupervised learning of object landmarks through conditional image generation. In Advances in Neural Information Processing Systems, pp. 4016 2018.
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-centric models. In International Conference on Learning Representations, 2019.
Timor Kadir and Michael Brady. Saliency, scale and image description. Int. J. Comput. Vision, 45(2): 83-105 November 2001. ISSN 0920-5691. doi: 10. 10.1023/A.10122444113555
Lukasz Kaiser, Mohammad Babaeizadeh. Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (cds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA. USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun (eds.). 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14- 16, 2014, Conference Track Proceedings, 2014.
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In International Conference on Learning Representations, 2020.

---

Published as a conference paper at ICLR 2021
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In Advances in Neural Information Processing Systems, pp. 8606-8616 2018.
Andreas Küccher and Christoph Goller. Inductive learning in symbolic domains using structure- recurrent neural networks. In Anmual Conference Artificial Intelligence, pp. 183-197. Springer, 1996.
Tejas D Kulkarni, Ankush Gupta, Catalin lonescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. In Advances in Neural Information Processing Systems, pp- 10723-10733, 2019.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
Zhixuan Lin, Yi-Fu Wu, Skand Vishwanath Peri, Weihao Sun, Gautam Singh, Fei Deng, Jindong Jiang, and Sungjin Ahn. SPACE: unsupervised object-oriented scene representation via spatial attention and decomposition. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020.
Matthias Minderer, Chen Sun, Ruben Villegas, Forrester Cole, Kevin P Murphy, and Honglak Lec. Unsupervised learning of object structure and dynamics from videos. In Advances in Neural Information Processing Systems, pp. 92-102, 2019
Sarthak Mittal, Alex Lamb, Anirudh Goyal, Vikram Voleti, Murray Shanahan, Guillaume Lajoie, Michael Mozer, and Yoshua Bengio. Learning to combine top-down and bottom-up signals in recurrent neural networks with attention over modules. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 6972-6986. PMLR, 13-18 Jul 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540):5295333 2015.
Alexander Mott, Daniel Zoran, Mike Chrzanowski, Daan Wierstra, and Danilo Jimenez Rezende. Towards interpretable reinforcement learning using attention augmented agents. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 32, pp. 12350-12359, Curran Associates, Inc., 2019.
Jordan B. Pollack. Recursive distributed representatiions Artificial Intelligence, 46(1):77 105, 1990.
ISSN 0004-3702. doi:
Scarselli, M. Gori, A. C. Tsoi, M. Hagenbuchner,, and G. Monfardini. The graph neural network model. IEEE Transactions on Neural Networks, 20( 1):61-00, 2009.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Aleksandar Stanić, Sjoerd van Steenkiste, and Jürgen Schmidhuber. Hierarchical relational inference. In Proceedings of the Thirty-Fifih AAAI Conference on Artificial Intelligence, 2021.
Felipe Petroski Such, Vashisht Madhavan, Rosanne Liu, Rui Wang. Pablo Samuel Castro, Yulun Li, Ludwig Schubert, Marc G. Bellemare, Jeff Clune, and Joel Lehman. An atari model zoo for analyzing, visualizing, and comparing deep reinforcement learning agents. In IJCAI, 2019.
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery of objects and their interactions. In International Conference on Learning Representations, 2018.
Sjoerd van Steenkiste, Klaus Greff, and Jürgen Schmidhuber. A perspective objects and systematic generalization in model- rl. ICML Workshop on Generative Modeling and Model-Based Reasoning for Robotics and AI, 2019.
12

---

Published as a conference paper at ICLR 2021
Rishi Veerapaneni, John D Co- CeyRe, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua B Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Conference on Robot Learning, 2019.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. Nervenet: Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018.
Vinicius Zambaldi, David Raposo, Adam Santoro, Victor Bapst, Yujia Li, Igor Babuschkin, Karl Tuyls, David Reichert, Timothy Lillicrap, Edward Lockhart, Murray Shanahan, Victoria Langston, Razvan Pascanu, Matthew Botvinick, Oriol Vinyals, and Peter Battaglia. Deep reinforcement learning with relational inductive biases. In International Conference on Learning Representations, 2019.
Yuting Zhang, Yijie Guo, Yixin Jin, Yijun Luo, Zhiyuan He, and Honglak Lee. Unsupervised discovery of object landmarks as structural representations. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 2694-2703, 2018.
Andrey Zhmoginov, lan Fischer, and Mark Sandler. Information-bottlena approach salient region discovery. ICML Workshop on Self-Supervised Learning, 2019.
Xingyi Zhou, Dequan Wang, and Philipp Krähenbühl. Objects points. In arXiv preprint 1rXiv:1904.07850, 2019.
13

---

Published as a conference paper at ICLR 2021
A EXPERIMENT DETAILS
A. DATASETS
To obtain Atari game frames, we use rollouts of various pre-trained agents in the Atari Model Zoo (Such et al., 2019). We split the aggregated set of game frames for each the chosen environments into separate train, validation and test sets of 85,000, 5000 and 5000 samples respectively.
For "noisy" Atari, we start with the regular Atari dataset and superimpose colored bars (either horizontal, vertical or both) centered at random x-y co-ordinates. This is done on-the-fly during traiinngevaluuation to generate the required Noisy Atari samples.
A.2 ARCHITECTURE AND TRAINING DETAILS
A.2.1 PERMAKEY
We train our method using the Adam optimizer (Kingma & Ba 2015) with initial learning rate of 0.0002 and decay rate of 0.85 every 10000 steps. We use batch size of 32 and in all cases train for 100 epochs, using use early stopping of 10 epochs on the validation set to prevent overfitting. All modules (i.e. the VAE, the prediction network, and PointNet) are trained in parallel using their respective losses. We do not backpropagate gradients between modules to improve stability. The same hyperparameters are used across all environments except for the number of keypoints.
Variational Autoencoder We use convolutional neural network with Conv-BatchNorm-RRe ReLU layers for the encoder network of the VAE. The encoder network uses kernel sizes [4, 3, 3, 3]. filters [32, 64, 64, 128] and strides [1, 2, 2, 1]. The architecture of the decoder the transpose of that of the encoder with 2x bi-linear upsampling used to undo the striding.
Preliminary experiments were used to determine the kernel size of the convolutional layers in the encoder. There it was observed that a kernel size of 4 on the first layer for 84 84 images achieves the best results and that larger kernel sizes and strides (2 2) reduce the spatial resolution of the feature maps and lead to coarser keypoint localization.
Prediction Network We use separate 3-layer MLP with hidden layer sizes [8 p x p C, 151,,26,, p X p x C1 where denotes height and width of activation patch and C' the num- ber of channels of the activation map of encoder CNN. We use linear output activations for the prediction network and use a separate network for each of the selected layers of the VAE encoder that perform separate local spatial prediction task. We use encoder layers (0, 1] and p = for the local spatial prediction task.
PointNet For the PointNet network we use the same encoder architecture for the VAE, but add a final x 1 regressor to K feature-maps corresponding to K keypoints (Jakab et al.) 2018) with standard deviation of 0.1 for the Gaussian masks. We resize predictability maps to 84 84 and concatenate them channel-wise before feeding it to the PointNet.
A.2.2 TRANSPORTER RE-IMPLEMENTATION
We re-implemented the Transporter model (Klkkrni et al. 2019) for our baseline method. We used an encoder network of Hoan-BBatcNNomm-RRe
layers. The feature extractor network + uses kernel sizes [3, 3, 3). filters [16, 16, 32, 32] and strides [1, 1, 1]. The PointNet V uses the same architecture but includes a final X 1 regressor to K feature-maps corresponding to K keypoints. computed from these K maps described in Jakab et al. (2018). The RefineNer
2D co-ordinates uses the transpose of t with 2x bi-linear upsampling to undo striding. We used the Adam optimizer with learning rate of 0.0002 with decay rate of 0.9 every 30000 steps and batch size of 64. We trained for 100 epochs with an early-stopping parameter of 5 to prevent over-fitting.
The same hyperparameters are used across all environments except for the number of keypoints.
14

---

Published as a conference paper at ICLR 2021
A.2.3 RL AGENT AND TRAINING
The convolutional keypoint encoder in the Transporter and PKey+CNN models consists of Conv- BatchNorm-ReLU layers with kernel sizes [3, 3, Strdds
strides [1, 1, 2, 1] and filters [128, 128, 128, 128].
The final convolution layer activation is flattened and passed through a Dense layer with 128 units and ReLU activation. The convolutional keypoint encoder receives the keypoint features obtained in the following manner: let (xre) ) H. W, CJ] tensor the convolutional features and H(s,) [H. W. K] tensor be the keypoint masks for an input image Keypoint features are obtained by multiplying keypoint masks with convolutional features and superposing masked features of all the K keypoints.
The MLP architecture used as the building block throughout the graph neural network consists of a 2-layers, each having 64 ReLU units. We use similar MLP (but having linear output units) as the positional encoder network to produce learned positional embeddings given keypoint centers. Output node values from the Decode block of the GNN concatenated and fed into a 2-layer MLP, with 128 ReI.U units each before being input to the agent.
For training the RL agents, update the target Q-network. We use Adam optimizer (Kingma & Ba 2013) with a learning rate of 0.0002 and clip gradients to maximum norm of 10. For e-greedy exploration we linearly anneal from to 0.1 over the entire 100,000 steps (or 400,000 frames).
Polyak weight averaging scheme at every step (constant=0.005) to
During training we checkpoint policies by running 10 episodes on separate validation environment initialized with a seed different from the one used for training. We checkpoint policies based on their mean scores on this validation environment. For final evaluation (reporting test scores), we take the best checkpointed models from 3 different runs and run 10 evaluation episodes for each of the policies on 10 previously unseen test environment seeds. We report the mean (and std. dev.) for the different policies using the median evaluation run over the 10 test seeds. We emphasize that different environment seeds are used for training, validation (checkpointing), and testing. Further we do not add any noise in our policy used for evaluation episodes.
B ADDITIONAL RESULTS
B. ADDITIONAL VISUALIZATIONS
B.1.1 NUMBER OF KEYPOINTS
On the Frostbite environment we modulate the number of keypoints produced by the PointNet while keeping all other hyperparameters same as described above Amppendix A (visualization shown in Figure 6).
B.1.2 EFFECT OF LAYER CHOICE
On the Space Invaders' environment we modulate the feature layer(s) chosen = {(0), (0, 1), (2, 3), (0, 1, 2, 3)) for the local spatial prediction task while keeping all other hyperparameters same as described above in Appendix A
B.1.3 SEEDS
We evaluate the stability of our keypoint discovery method over random seeds on Enduro and Battlezone environments with all hyperparameter settings same described in Appendix Aland the qualitative results are shown in Figure 8 and Figure 9
B.1.4 KEYPOINT MASKED REGIONS
Figure 10]shows the keypoint mean (center) as well as Gaussian mask around it. We can see clearly from Figure 10) that although the keypoint centers might be slightly focused towards the borders of predictability their corresponding attention windows ensure good coverage of salient object parts in the scene.
15

---

Published as a conference paper at ICLR 2021
Num
F
0
Num
Figure 6: Image (row 1), keypoints (15, 20 and 25) overlaid on image (rows 2, and 4). With 15 keypoints (row 2)
that our model prioritizes the keypoints to capture only the most salient object parts in the scene i.e. ice floats, player, scoreboard etc. As we increase the number of keypoints available it places the excess keypoints the right border, after having captured all the highly salient object parts i.c. ice floats, player, birds etc.
16

---

Published as a conference paper at ICLR 2021
1240
1310
D0OD
1200
0180
DOOO
2
W
Image
3O00
0000
3000
L3yen 0
0
Layers 0-3
12
B
13 E3
S9P0 23
Layers 0-1-2-3
Figure 7: Image (row 1), keypoints using feature layer(s) = {0), (0, 1), (2, 3), (0, 1, 2,2,3), 3)} (rows 2, that using feature layers 0 and 1 achieves the best keypoint quality. Including higher feature layers (rows 4 and 5) results in more uneven grouping of rows of enemy aliens keypoints (columns 2 and 6 row and failure track orange shields (rows 4 and 5).
4 and 5 respectively) for local spatial predictability task. We
17

---

Published as a conference paper at ICLR 2021
00.11
Figure 8: image (row 1), keypoints colour coded from 5 random seeds overlaid simultancously on image (row 2). we can see the learned keypoints remain stable across different runs and capture the salient object parts i.e. player, edges of the road, oncoming traffic, scoreboard etc. each time.
2019
35013
23013
9909
39013
Figure 9: image (row 1), keypoints colour coded from 5 random seeds overlaid simultaneously on image (row 2). As we can see the learned keypoints remain stable across different runs and capture the salient object parts i.e. player tank parts, enemy tank, orientation dial at the top etc. each time.
18

---

Published as a conference paper at ICLR 2021
Frostbite
MsPacman
Enduro
Battlezone
Image
Keypoints (ours)
Keypoint Regions
Figure 10: image (row 1), keypoint centers (row 2) and keypoint windows/masks, (row 3).
19


Reviews:
[{"review_id": "u6hVfZBqDrL", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors tackle the problem of self-supervised representation learning, and validate their approach on downstream Reinforcement Learning tasks. Building on the insight that predictability in local patches is a good inductive bias for salient regions that characterize objects, the authors propose a well-reasoned, well-engineered and thoroughly validated pipeline for learning object keypoints without supervision. The authors present a wide range of ablative studies to validate their design choices, and demonstrate the superiority of their method both illustratively as well as quantitatively on a number of standard Atari benchmarks. \n\nThe paper is very well written and clearly explained, with the assumptions clearly stated, and design choices thoroughly validated through ablative studies. Overall the authors make a very compelling argument for using local predictability as an intrinsic property of objects, leading me to recommend accepting this paper for publication.\n\nPros:\n+ The intro motivates the problem well, contrasting the proposed method with a number of key recent methods. The implementation details are well recorded in the supplementary, with the added mention of releasing the source code \n+ The keypoint detection pipeline is well reasoned and well explained: using the error map obtained through the spatial prediction task to recover keypoints via a bottleneck with limited capacity is a neat idea. The authors ablate a number of design choices (number of keypoints, which encoder layers to use); Figure 1 and 2 are great at showing the high-level components of the method as well as (intermediate) outputs\n+ The comparison against Transporter is thorough and well analyzed. Fig2.b. provides a very clear insight into the limitations of Transporter, showing that the method proposed by the authors is able to achieve some robustness to visual distractors. PKey-CNN uses a similar method as Transporter for encoding keypoint features for downstream tasks, and thus serves to show that the keypoints identified are indeed superior. PKey-GNN further increases performance on a number of Atari games. \n+ Very good ablative analysis and qualitative examples.\n\nSome questions:\n+ Do the authors have any further insights regarding why PKey-GNN would perform worse than PKey-CNN? While the authors’ reasoning makes sense, in my understanding a GNN based approach should be able to model any kind of interaction. \n+ The authors demonstrate impressive results on a number of Atari games. I am wondering how this method would perform on a slightly more complex environment, i.e. CarRacing in OpenAI’s gym environment, or maybe even going as far as CARLA?\n+ As I understand, PermaKey is first trained on Atari game rollouts, with the policy trained afterwards. Would it be possible to optimize both the keypoints and the policy together, end-to-end?\n\nPost Rebuttal:\nI thank the authors for their detailed and thorough response. All my questions and concerns were addressed and I appreciate the discussion on end-to-end learning as well as the “Transporter + GNN experiment”. I am happy to maintain my original rating and recommend acceptance.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "hCRTWCKTmof", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "#### Summary\n\nThis 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. \n\n#### Strength\n\n- The key idea of finding non-locally-predictable points as a representation of the game state is interesting, and specifically suitable for Games, where the backgrounds are mostly static and predictable. \n\n- The technical implementation of the framework (Fig. 2) is clear and makes sense to me.\n\n- Experimental results in Table. 1 is healthy. It shows the proposed method decently outperforms the Transporter counterparts.\n\n#### Weaknesses\n\n- The ablation studies are not exciting. These (number of keypoints/ spatial resolution of the embedding) are mostly design choice experiments and are better to be in the supplement materials. A more interesting ablation would be to quantitatively evaluate the quality of the points. Currently the paper only qualitatively shows the keypoint discovering results in Fig. 2 and claims an advantage to Transporter. This is not clear to the reviewer. The reviewer understands that there is no existing metric to evaluate keypoint discovery quality. However some proxy evaluation would also be helpful. For example, are the learned points temporally stable?\n\n- As the key contribution is the keypoint discovery, it would be more convincing to compare with other unsupervised keypoint discovery methods besides Transporter if applicable. E.g., PointNet (Jakab et al 2018) that considers keypoints as a pixel-level reconstruction bottle-neck. \n\n\n#### Summary\n\nThe paper proposed an interesting idea with reasonable results (better than a recent counterpart, Transporter). However, the reviewer does not have backgrounds in the specific experimental settings (Atari games), and can not assess the significance of the improvements. Comparisons with more keypoint discovery methods would make the results more convincing. My current rating is 6, but might change based on other reviews. \n\n#### Post rebuttal\n\nThank you for providing the rebuttal. The rebuttal addressed my concern on comparing to other baselines. And it's fine to keep the design choice experiments in the main paper. However a proxy evaluation of key point evaluation is still missing and it will further strengthen the paper (I don't have a clear idea for the evaluation either). I keep my original rating of 6. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "_z-yMq4pNdP", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose a novel approach to unsupervised key point detection based on predictability. The demonstrate their model on Atari tasks comparing to other key point detectors.\n\n Quality:\n\nThe authors compare their work both qualitatively and quantitatively to Transporter. \n\nThe authors show that their model picks out important key points that Transporter does not. Figure 5 is a great! It would also be good to show the distribution of predicted key points over multiple runs for other levels.\n\nThe authors train agents on Atari and compare their model to suitable baselines. It’s interesting that the GNN does not always outperform the CNN. This paper could be improved by also comparing to Transporter + GNN. \n\nWhen constructing the error map, is this approach very sensitive to the receptive field and the number/ location of the neighbours? How would this approach handle larger / more complex objects?\n\nSection 5.2 is interesting. What about if you have new objects that are not predictable, but are distractors? Would the model not create key points here? For example, adding in some smaller distractor shapes? Some randomly coloured dots? Or some missing pixels?\n\nThe ablation study is great!\n\nThis results section of this paper is very thorough and addressed a lot of the questions that came to mind when reading the introduction and methods section of this paper. \n\nClarity:\n\nThe authors argue that local predictability is an intrinsic property of an object with out giving more evidence for this. Perhaps the authors are hypothesising that this interpretation of objects will be more useful for downstream tasks? It’s not clear to me otherwise why this is an intrinsic property of an object?\n\nThe authors could improve their paper by being very clear about distinguishing focus on image regions that are unpredictable” from “local regions in input space that have high internal predictive structure” when describing objects and key points.\n\nWith the exception to the above, the introduction is well written and the methods section is easy to follow.\n\nThe model is designed to pick out points that are harder to predict, which is useful for ignoring background and finding the agent as nicely demonstrated in Figure 2 of Montezuma Revenge, but it’s not clear that this is a good definition of an object? For example the platforms in Frostbite may be very easy to predict, but you would need to know where they are in order to successfully navigate the environment. Also, it seems that predictability of a feature may depend heavily on the environment and any new object in an environment would be immediately picked out even if they are irrelevant to the task (i.e. new colours etc).\n\nCould you explain better why you thing that the platforms in frostbite are assigned key points as they are?\n\nFigure 3 is a really clear and nice example.\n\nOriginality and Significance:\n\nThis approach is novel and interesting and offers a new perspective on what an object can be and what definitions of an object may be useful for training agents.\n\nPros:\n- Well written (with minor exception)\n- Thorough results section.\n- A novel approach to thinking about what an object is.\n- Improvement over Transporter on Atari tasks.\n\nCons:\n- Some confusion in the introduction about their definition of an object.\n- There may be some limitation of the types (size and shape) of the objects that this model can assign key points to.\n- There are some examples of objects being detected where it is not clear why those points are being detected according to the definition given in the paper. Explaining this more clearly would improve the paper.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers all agreed that this submission has an interesting new idea for learning object/keypoint representations: parts of a visual scene that are not easily predictable from their neighborhoods are good object candidates. Experimental gains on various Atari games are convincing. The main drawback at this point is that the evaluation is limited to visually rather simple settings, and it is unclear how the approach will scale to more realistic scenes. 
Decision: Accept (Spotlight)