Original: L EARNING U NSUPERVISED F ORWARD M ODELS FROM O BJECT K EYPOINTS
Revised: L EARNING U NSUPERVISED F ORWARD M ODELS FROM O BJECT K EYPOINTS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Object-centric representation is an essential abstraction for forward prediction.
Revised: Object-centric representation is an essential abstraction for forward prediction.

Original: Most existing forward models learn this representation through extensive supervision (e.g., object class and bounding box) although such ground-truth information is not readily accessible in reality.
Revised: Most existing forward models learn this representation through extensive supervision (e.g., object class and bounding box) although such ground-truth information is not readily accessible in reality.

Original: To address this, we introduce KINet (Keypoint
Revised: To address this, we introduce KINet (Keypoint

Original: Interaction Network)—an end-to-end unsupervised framework to reason about object interactions based on a keypoint representation.
Revised: Interaction Network)—an end-to-end unsupervised framework to reason about object interactions based on a keypoint representation.

Original: Using visual observations, our model learns to associate objects with keypoint coordinates and discovers a graph representation of the system as a set of keypoint embeddings and their relations.
Revised: Using visual observations, our model learns to associate objects with keypoint coordinates and discovers a graph representation of the system as a set of keypoint embeddings and their relations.

Original: It then learns an action-conditioned forward model using contrastive estimation to predict future keypoint states.
Revised: It then learns an action-conditioned forward model using contrastive estimation to predict future keypoint states.

Original: By learning to perform physical reasoning in the keypoint space, our model automatically generalizes to scenarios with a different number of objects, novel backgrounds, and unseen object geometries.
Revised: By learning to perform physical reasoning in the keypoint space, our model automatically generalizes to scenarios with a different number of objects, novel backgrounds, and unseen object geometries.

Original: Experiments demonstrate the effectiveness of our model in accurately performing forward prediction and learning plannable object-centric representations which can also be used in downstream robotic manipulation tasks.
Revised: Experiments demonstrate the effectiveness of our model in accurately performing forward prediction and learning plannable object-centric representations which can also be used in downstream robotic manipulation tasks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Discovering a structured representation of the world allows humans to perform a wide repertoire of motor tasks such as interacting with objects.
Revised: Discovering a structured representation of the world allows humans to perform a wide repertoire of motor tasks such as interacting with objects.

Original: At the core of this process is the ability to predict the response of the environment to applying an action (Miall & Wolpert, 1996; Wolpert & Kawato, 1998).
Revised: The core of this process is learning to predict the response of the environment to applying an action (Miall & Wolpert, 1996; Wolpert & Kawato, 1998).

Original: Such an internal model, often referred to as the forward model, aims to come up with an estimation of the future states of the world given its current state and an action.
Revised: The internal models, often referred to as the forward models, come up with an estimation of the future states of the world given its current state and the action.

Original: 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).
Revised: 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).

Original: Recently,various deep learning architectures have been proposed to perform forward modeling using an object-centric representation of the system (Ye et al., 2020;
Revised: Recently, deep learning architectures have been proposed to perform forward prediction using an object-centric representation of the system (Ye et al., 2020;

Original: Chen et al., 2021b; Li et al., 2020; Qi et al., 2020).
Revised: Chen et al., 2021b; Li et al., 2020a; Qi et al., 2020).

Original: This 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).
Revised: This 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).

Original: We identified two major limitations in the recent work: First, they 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).
Revised: We identified two major limitations in the existing work: First, they either assume access to the ground-truth object states (Battaglia et al., 2016; Li et al., 2020a) or predict them using idealized techniques such as pre-trained object detection or instance segmentation models (Ye et al., 2020; Qi et al., 2020).

Original: However, obtaining ground truth object states is not always feasible in practice.
Revised: However, obtaining ground truth object states is not always feasible in practice.

Original: 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 which is often infeasible in real-world settings.
Revised: 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 which is often infeasible in real-world settings.

Original: Second, factorizing the scene into a fixed number of object instances limits the generalization of forward models to scenarios with a different number of objects.
Revised: Second, factorizing the scene into a fixed number of object instances limits the generalization of the model models to scenarios with a different number of objects.

Original: In this paper, we address both of these limitations by proposing to learn forward models using a keypoint-based object-centric representation.
Revised: In this paper, we address both of these limitations by proposing to learn forward models using a keypoint representation.

Original: Keypoints represent a set of salient locations of moving entities.
Revised: Keypoints represent a set of salient locations of moving entities.

Original: Our model KINet (Keypoint Interaction Network) learns an unsupervised forward model in threemajor steps: (1) A keypoint extractor factorizes the scene into keypoint coordinates with no supervision other than raw visual observations.
Revised: Our model KINet (Keypoint Interaction Network) learns an unsupervised forward model in three steps: (1) A keypoint extractor factorizes the scene into keypoints with no supervision other than raw visual observations.

Original: A probabilistic graph representation of the system is learned where each node corresponds to a keypoint and edges are keypoints relations.
Revised: A probabilistic graph representation of the system is learned where each node corresponds to a keypoint and edges are keypoints relations.

Original: Node features carry implicit object-centric representation as well as explicit keypoint state information.
Revised: Node features carry implicit object-centric features as well as explicit keypoint state information.

Original: (3) With probabilistic message passing, our model learns an action-conditional forward model to predict  future locations of the keypoints. Using this prediction, it also reconstructs the future appearances of the scene.
Revised: (3) With probabilistic message passing, our model learns an action-conditional forward model to predict the future location of keypoints and reconstruct the future appearances of the scene.

Original: We evaluate KINet’s forward prediction accuracy and demonstrate that, by learning the physical reasoning in a keypoint coordinate, our modelcan effectively re-purpose this knowledge and generalize it to complex unseen circumstances.
Revised: We evaluate KINet’s forward prediction accuracy and demonstrate that, by learning forward prediction in a keypoint coordinate, our model effectively re-purposes this knowledge and generalizes it to complex unseen circumstances.

Original: 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.
Revised: Our key contributions are: (1) We introduce KINet, an end-to-end method for learning unsupervised action-conditional forward models from visual observations (2)

Original: (2) We introduce Probabilistic Interaction Networks for efficient message-passing to aggregate relevant information.
Revised: We introduce probabilistic Interaction Networks for efficient message-passing to aggregate relevant information.

Original: (3) We introduce the GraphMPCalgorithm for accurate action planning based on graph similarity.
Revised: (3) We introduce the GraphMPC for accurate action planning using graph similarity.

Original: (4) We demonstratethat learning forward models in keypoint coordinates enables zero-shot generalization to complex unseen scenarios.
Revised: (4) We demonstrate learning forward models in keypoint coordinates enables zero-shot generalization to complex unseen scenarios.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Unsupervised keypoint extraction methods have been usedsuccessfully 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; Manuelli et al., 2020).
Revised: Unsupervised keypoint extraction. Keypoints have been used in computer vision areas 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; Manuelli et al., 2020).

Original: Recent work explored keypoint extraction for control tasks in reinforcement learningenvironments to project the visual observationspace to a lower-dimensional keypoint space (Kulkarni et al., 2019;
Revised: Recent work explored keypoints for control tasks in reinforcement learning to project the visual observation to a lower-dimensional keypoint space (Kulkarni et al., 2019;

Original: Chen et al., 2021a; Jakab et al.,2018; Suwajanakorn et al., 2018).
Revised: Chen et al., 2021a; Jakab et al., 2018).

Original: We extend these methods for the object manipulation setting.

Original: Forward prediction models estimate the future state of a system given its current state. Neural networks can be used to estimate such models for various applications. Example of earlier approaches are Nguyen & Widrow (1990) and Grzeszczuk et al. (1998) for modeling state transition in dynamic systems.
Revised: Forward models.

Original: The most fundamentally relevant work to ours isthe Interaction Network (IN) (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018) and follow-up workon using graph neural network in forward modeling (Pfaff et al., 2020;
Revised: The most fundamentally relevant work to ours is Interaction Networks (IN) (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018) and follow-up work using graph neural network for forward simulation (Pfaff et al., 2020;

Original: Li et al., 2019; Kipf et al., 2018; Mrowca et al., 2018).
Revised: Li et al., 2019; Kipf et al., 2018; Mrowca et al., 2018).

Original: Although these approaches demonstrated powerful forward modeling, they rely onbuilding explicit object representations based on the ground-truth state information of objects.
Revised: These methods rely on the ground-truth states of objects to build explicit object representations.

Original: In an attempt to bridge this gap several methods have been proposed to use IN-based models based on visual featuresextracted from image observations (Watters et al., 2017; Qi et al., 2020; Ye et al., 2020).
Revised: Several approaches extended IN by combining explicit object positional information with implicit visual features from images (Watters et al., 2017; Qi et al., 2020; Ye et al., 2020).

Original: However, two main concerns remain unaddressed in these methods.
Revised: However, two main concerns remain unaddressed.

Original: First,the visual features are often extracted from object bounding boxes using object detection or segmentation model (Janner et al., 2018;
Revised: First, visual features are often extracted from object bounding boxes using object detection or segmentation models (Janner et al., 2018;

Original: Qi et al., 2020; Kipf et al., 2018) that is either pretrained on the environment (Qi et al., 2020) or has the ground-truth object position (Ye et al., 2020).
Revised: Qi et al., 2020; Kipf et al., 2018) that are either pretrained on the setup (Qi et al., 2020) or use the ground-truth object position (Ye et al., 2020).

Original: Second, these approaches lack generalization as they are formulated on a fixed number of objects.
Revised: Second, these approaches lack generalization as they are formulated on a fixed number of objects.

Original: Learning action-conditional forward models is an active challenge.
Revised: Minderer et al. (2019) used keypoints for video prediction given a history of previous frames. However, their dynamics model is not formulated on external action and cannot be used for action planning applications. Action-conditional forward models.

Original: The original IN augments the actionvector as an external effect  to the node embeddings. 
Revised: Battaglia et al. (2016) augments the action as an external effect augmented to the node embeddings. Ye et al. (2020) included action as an additional node in a fully connected graph with other nodes representing objects.

Original: For probabilistic forward models Henaff et al.
Revised: For probabilistic forward models Henaff et al.

Original: (2019)interestingly suggests using a latent variable dropout mechanism to properly condition theforward model on the action (Gal & Ghahramani, 2016).
Revised: (2019) suggests using a latent variable dropout mechanism to properly condition the model on the action (Gal & Ghahramani, 2016).

Original: In a more relevant application to ours, Yan et al.
Revised: In a more relevant application to ours, Yan et al.

Original: (2020) highlighted the effectiveness of contrastive estimation (Oord et al., 2018) to learn proper actionable object-centric representations.
Revised: (2020) highlighted the effectiveness of contrastive estimation (Oord et al., 2018) to learn proper actionable object-centric representations.

Original: Graph representation inference methods have been introduced in variational graph autoencoders (Kipf & Welling, 2016). To formulate the graph structure inference we were also inspired by Simonovsky & Komodakis (2018) where they use a variational graph autoencoder to generate small graphs of molecular structures.
Revised: Unsupervised Forwad Models.

Original: More relevant to our work is Kipf et al.
Revised: Kipf et al.

Original: (2019) where an object-level contrastive losshas been used to learn object-centric abstractions in multibody systems.
Revised: (2019) uses an object-level contrastive loss to learn object-centric abstractions in multibody environments with deterministic structures and minimal visual features such as 2d shapes.

Original: However, this work is only based on fixed environments with oversimplified objects such as two-dimensional geometries (Watters et al., 2019; Stani´c et al., 2020).
Revised: In our work, we randomize the properties of the system and examine realistic objects. Veerapaneni et al. (2020) infers a set of entities in the image based on their depth and predicts future entity states.

Original: In our work, we experiment with more complex three-dimensional objects with randomized properties to replicate real-world settings.
Revised: Entities are mixed using weight parameters of their distance from the camera. Our work, on the other hand, does not make any assumption on inferring depth and relies on keypoints.

Original: K EYPOINT I NTERACTION N ETWORKS (KIN ET )
Revised: Kossen et al. (2019) uses images to infer a set of explicit states (e.g., position and velocity) for a fixed number of objects in a dynamic system to predict the future state of each object using graph networks. Although learning unsupervised state representation, this method is formulated on a fixed number of objects and only tested on environments with simple 2d geometries. Li et al.

Revised: (2020b) infers the causal structure and makes future predictions on a fixed dynamic system using a pretrained keypoint extractor on topview images.

Revised: In our work, we experiment with other camera angles and 3d objects with random properties such as geometry and texture.

Revised:  K EYPOINT I NTERACTION N ETWORKS (KIN ET )

Original: We assume access to observational data that consists of RGB image, action vector, and the resulting image after applying the action:
Revised: We assume access to observational data that consists of RGB image, action vector, and the resulting image after applying the action:

Original: D = { ( I t , u t , I t +1 ) } .
Revised: D = { ( I t , u t , I t +1 ) } .

Original: Our goal is to learn a forward model that predicts  future states of the system with no supervision above the observational data.
Revised: Our goal is to learn a forward model that predicts the future states of the system with no supervision above the observational data.

Original: We describe our approach in two main steps (see Figure 1): learning to encode visual observations into keypoints andthen learning an action-conditioned forward model in the keypoint space.
Revised: We describe our approach in two main steps (see Figure 1): learning to encode visual observations into keypoints and learning an action-conditioned forward model in the keypoint space.

Original: 3. U NSUPERVISED K EYPOINT D ETECTION
Revised: 3. U NSUPERVISED K EYPOINT D ETECTION

Original: The keypoint detector ( f kp , Fig.1) is a mapping from visual observations to a lower-dimensional set of K keypoint coordinates { x kt } k
Revised: The keypoint detector ( f kp , Fig.1) is a mapping from visual observations to a lower-dimensional set of K keypoint coordinates { x kt } k

Original: =1 ...K
Revised: =1 ...K

Original: = f kp ( I t ) .
Revised: = f kp ( I t ) .

Original: The keypoint coordinates are learned through capturing the spatial appearance of the objects in an unsupervised manner.
Revised: The keypoint coordinates are learned by capturing the spatial appearance of the objects in an unsupervised manner.

Original: Specifically, the keypoint detector receives a pair of initial and current images ( I 0 , I t ) and uses a convolutional encoder to compute a K -dimensional feature map for each image Φ( I 0 ) , Φ( I t ) ∈
Revised: Specifically, the keypoint detector receives a pair of initial and current images ( I 0 , I t ) and uses a convolutional encoder to compute a K -dimensional feature map for each image Φ( I 0 ) , Φ( I t ) ∈

Original: R H ′
Revised: R H ′

Original: × W ′ × K .
Revised: × W ′ × K .

Original: The expected number of keypoints is set by the dimension K .
Revised: The expected number of keypoints is set by the dimension K .

Original: Next, the feature maps are marginalized into a 2D keypoint coordinate { x k 0 , x kt } k =1 ...K ∈ R 2 .
Revised: Next, the feature maps are marginalized into a 2D keypoint coordinate { x k 0 , x kt } k =1 ...K ∈ R 2 .

Original: We use a convolutional image reconstruction model ( f rec , Fig.1) with skip connections to inpaint the current image frame using the initial image and the predicted keypoint coordinates ˆ I t
Revised: We use a convolutional image reconstruction model ( f rec , Fig.1) with skip connections to inpaint the current image frame using the initial image and the predicted keypoint coordinates ˆ I t

Original: = f rec ( I 0 , { x k 0 , x kt } k
Revised: = f rec ( I 0 , { x k 0 , x kt } k

Original: =1 ...K ) .
Revised: =1 ...K ) .

Original: With this formulation, f kp and f rec create a bottleneck to encode the visual observation in a temporally consistent lower-dimensional set of keypoint coordinates (Kulkarni et al., 2019).
Revised: With this formulation, f kp and f rec create a bottleneck to encode the visual observation in a temporally consistent lower-dimensional set of keypoint coordinates (Kulkarni et al., 2019).

Original: 3. G RAPH R EPRESENTATION OF S YSTEM
Revised: 3. G RAPH R EPRESENTATION OF S YSTEM

Original: After factorizing the system into K keypoints, we build a graph G t =
Revised: After factorizing the system into K keypoints, we build a graph G t =

Original: ( V t , E t , Z ) (undirected, no self-loop) where keypoints and their pairwise relations become the graph nodes and edges.
Revised: ( V t , E t , Z ) (undirected, no self-loop) where keypoints and their pairwise relations are the graph nodes and edges.

Original: Keypoint positional and visual information are encoded intothe embedding of nodes { n kt } k =1 ...K ∈ V t and edges { e ijt } ∈ E t .
Revised: Keypoint positional and visual information are encoded into embedding of nodes { n kt } k =1 ...K ∈ V t and edges { e ijt } ∈ E t .

Original: We also define an adjacency matrix to specify thegraph connectivity as Z ∈ R K × K where [ Z ] ij ∈ [0 , 1] specifies the probability of the edge { e ij } ∈ E .
Revised: We also use an adjacency matrix to specify the connectivity as Z ∈ R K × K where [ Z ] ij ∈ [0 , 1] specifies the probability of the edge { e ij } ∈ E .

Original: Ateach timestep t , node embeddings encode keypoint visualfeatures and positional information { n kt } = (cid:2)
Revised: At timestep t , node embeddings encode keypoint visual and positional information { n kt } = (cid:2)

Original: Φ kt , x kt (cid:3) .
Revised: Φ kt , x kt (cid:3) .

Original: Edge embeddingsalso contain relative positional information of each node pair { e ijt } = (cid:2) x it − x jt , ∥ x it − x jt ∥ 22 (cid:3) .
Revised: Edge embeddings contain relative positional information of each node pair { e ijt } = (cid:2) x it − x jt , ∥ x it − x jt ∥ 22 (cid:3) .

Original: 3. P ROBABILISTIC I NTERACTION N ETWORKS
Revised: 3. P ROBABILISTIC I NTERACTION N ETWORKS

Original: To build a forward model, we extend the recent approaches and propose a probabilistic variation of the Interaction Networks (IN) (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018).
Revised: To build a forward model, we extend the recent approaches and propose a probabilistic variation of the Interaction Networks (IN) (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2018).

Original: The core of probabilistic IN is generating edge-level latent variables z
Revised: The core of probabilistic IN is generating edge-level latent variables z

Original: ij ∈ R d
Revised: ij ∈ R d

Original: that represents the edge probability [ Z ] ij = z ij .
Revised: that represents the edge probability [ Z ] ij = z ij .

Original: A posterior network p θ infers the elements of the adjacency matrix given the graph representation of the scene.
Revised: A posterior network p θ infers the elements of the adjacency matrix given the graph representation of the scene.

Original: In particular, we model the posterior as p θ ( z ij |G t )
Revised: In particular, we model the posterior as p θ ( z ij |G t )

Original: =
Revised: =

Original: σ ( f enc ([ n it , n jt )]) where σ ( . ) is the sigmoid function.
Revised: σ ( f enc ([ n it , n jt ])) where σ ( . ) is the sigmoid function.

Original: The probabilistic IN forward model ˆ G t +1 =
Revised: The probabilistic IN forward model ˆ G t +1 =

Original: f fwd ( G t , u t , Z ) predicts the graph representation at the next timestep by taking as input the current graph representation and action ( f fwd , Fig.1).
Revised: f fwd ( G t , u t , Z ) predicts the graph representation at the next timestep by taking as input the current graph representation and action ( f fwd , Fig.1).

Original: The message-passing operation in the forward model can be described as, { ˆ e ij } ← f e ( n i , n j , e ij ) , { ˆ n k } ← f n ( n k , (cid:80) i ∈ N ( k ) Z ik ˆ e ik , u t ) where the edge-specific function f
Revised: The message-passing operation in the forward model can be described as, { ˆ e ij } ← f e ( n i , n j , e ij ) , { ˆ n k } ← f n ( n k , (cid:80) i ∈ N ( k ) Z ik ˆ e ik , u t ) where the edge-specific function f

Original: e first updates edge embeddings, then the node-specific function f n updates each node embedding ˆ n k by probabilistically aggregating its neighboring nodes N ( k ) information (i.e., the edge probabilities from the inferred adjacency matrix Z is used as weights in the neighbor aggregation).
Revised: e first updates edge embeddings, then the node-specific function f n updates each node embedding ˆ n k by probabilistically aggregating its neighboring nodes N ( k ) information (i.e., the edge probabilities from the inferred adjacency matrix Z is used as weights in the neighbor aggregation).

Original: The action vector u t is also an input to the neighbor aggregation step.
Revised: The action vector u t is also an input to the neighbor aggregation step.

Original: Note that the functions f enc , f e , and f n are MLPs.
Revised: Note that the functions f enc , f e , and f n are MLPs.

Original: Recent models for forward prediction rely on fully connected graphs for message passing (Qi et al., 2020;
Revised: Recent models for forward prediction rely on fully connected graphs for message passing (Qi et al., 2020;

Original: Ye et al., 2020; Li et al., 2018).
Revised: Ye et al., 2020; Li et al., 2018).

Original: Our model, however, learns to probabilistically sample the neighbor information.
Revised: Our model, however, learns to probabilistically sample the neighbor information.

Original: Intuitively, this adaptive sampling allows the network to efficiently aggregate the most relevant neighboring information.
Revised: Intuitively, this adaptive sampling allows the network to efficiently aggregate the most relevant neighboring information.

Original: This is specifically essential in our model since keypoints could provide redundant information if they are in very close proximity.
Revised: This is specifically essential in our model since keypoints could provide redundant information if they are in very close proximity.

Original: 3. F ORWARD P REDICTION
Revised: 3. F ORWARD P REDICTION

Original: The state decoder ( f state , Fig.1) transforms the predicted node embeddings of the updated graph ˆ
Revised: The state decoder ( f state , Fig.1) transforms the predicted node embeddings of the updated graph ˆ

Original: G t +1 to a first-order difference which is integrated once to predict the position of the keypoints in the next timestep { ˆ x kt +1 }
Revised: G t +1 to a first-order difference which is integrated once to predict the position of the keypoints in the next timestep { ˆ x kt +1 }

Original: = { x kt } + f state
Revised: = { x kt } + f state

Original: ( { ˆn kt +1 } ) .
Revised: ( { ˆn kt +1 } ) .

Original: To reconstruct the image at the next timestep, we borrow the reconstruction model f rec from the keypoint detector
Revised: To reconstruct the image at the next timestep, we borrow the reconstruction model f rec from the keypoint detector

Original: ˆ I t +1 =
Revised: ˆ I t +1 =

Original: f rec ( I 0 , { x k 0 , ˆ x kt +1 } ) .
Revised: f rec ( I 0 , { x k 0 , ˆ x kt +1 } ) .

Original: 3. L EARNING KIN ET
Revised: 3. L EARNING KIN ET

Original: Reconstruction loss.
Revised: Reconstruction loss.

Original: The keypoint detector is trained using the distance between the ground truth image and the reconstructed image at each timestep L rec =
Revised: The keypoint detector is trained using the distance between the ground truth image and the reconstructed image at each timestep L rec =

Original: ∥ ˆ I t − I t ∥ 22 .
Revised: ∥ ˆ I t − I t ∥ 22 .

Original: As suggested by Minderer et al.
Revised: As suggested by Minderer et al.

Original: (2019), errors from the keypoint detector were not backpropagated to other modules of the model.
Revised: (2019), errors from the keypoint detector were not backpropagated to other modules of the model.

Original: This is a necessary step to ensure the model does not conflate errors from image modules and reasoning modules.
Revised: This is a necessary step to ensure the model does not conflate errors from image modules and reasoning modules.

Original: Forward loss.
Revised: Forward loss.

Original: The model is also optimized to predict the next state of the keypoints.
Revised: The model is also optimized to predict the next state of the keypoints.

Original: A forward loss penalizes the distance between the estimated future keypoint locations using first-order state decoder and the keypoint extractor predictions: L fwd =
Revised: A forward loss penalizes the distance between the estimated future keypoint locations using first-order state decoder and the keypoint extractor predictions: L fwd =

Original: (cid:80) K
Revised: (cid:80) K

Original: ∥ ˆ x kt +1 − f kp ( I t +1 ) k
Revised: ∥ ˆ x kt +1 − f kp ( I t +1 ) k

Original: ∥ 22 .
Revised: ∥ 22 .

Original: Inference loss.
Revised: Inference loss.

Original: Our model is also trained to minimize the KL-divergence between the posterior and prior distributions: L infer =
Revised: Our model is also trained to minimize the KL-divergence between the posterior and prior distributions: L infer =

Original: D KL (cid:0) p ϕ ( Z |G ) (cid:13)(cid:13) p ( Z ) (cid:1) .
Revised: D KL (cid:0) p ϕ ( Z |G ) (cid:13)(cid:13) p ( Z ) (cid:1) .

Original: We use independent Gaussian prior p ( Z ) =
Revised: We use independent Gaussian prior p ( Z ) =

Original: (cid:81)
Revised: (cid:81)

Original: i N ( z i ) and use reparameterization trick for training (Kingma & Welling, 2013).
Revised: i N ( z i ) and use reparameterization trick for training (Kingma & Welling, 2013).

Original: Contrastive loss.
Revised: Contrastive loss.

Original: To further encourage the model to learn actionable representations, we make useof the contrastive estimation method.
Revised: We use the contrastive estimation method to further enhance the learned graph representations.

Original: We add a contrastive loss (Oord et al., 2018; Yan et al., 2020) and reframe it for graph embeddings as: L ctr = − E D [log( S (ˆ G t +1 , G + t +1 )  (cid:80) S
Revised: We add a contrastive loss (Oord et al., 2018; Yan et al., 2020) and reframe it for graph embeddings as: L ctr = − E D [log( S ( ˆ G t +1 , G + t +1 ) / (cid:80) S

Original: (ˆ G t +1 , G − t +1 ) )] .
Revised: ( ˆ G t +1 , G − t +1 ))] .

Original: The predicted graph representations ˆ G t +1 are maximally similar to their corresponding positive sample pair G + t +1 :
Revised: The predicted graph representations ˆ G t +1 are maximally similar to their corresponding positive sample pair G + t +1 :

Original: =
Revised: =

Original: G t +1 and maximally distant from the negative sample pairs G − t
Revised: G t +1 and maximally distant from the negative sample pairs G − t

Original: +

Revised: +

Original: = G τ ∀ τ ̸ = t + 1 .
Revised: = G τ ∀ τ ̸ = t + 1 .

Original: We use a simple node embedding similarity as the graph matching algorithm S ( G 1 , G 2 ) =
Revised: We use a simple node embedding similarity as the graph matching algorithm S ( G 1 , G 2 ) =

Original: (cid:80) K { n k 1 } .
Revised: (cid:80) K { n k 1 } .

Original: { n k 2 } .
Revised: { n k 2 } .

Revised: The motivation behind adding a contrastive loss is aligning the graph representation of similar object configurations while pushing apart those of dissimilar configurations in the embedding space to enhance the learned graphs.

Original: Finally, the combined loss is: L = λ rec L rec + λ fwd L fwd + λ infer L infer + λ ctr L ctr .
Revised: Finally, the combined loss is: L = λ rec L rec + λ fwd L fwd + λ infer L infer + λ ctr L ctr .

Original: 3. G RAPH MPC P LANNING WITH KIN ET
Revised: 3. G RAPH MPC P LANNING WITH KIN ET

Original: We use a learned KINet model and select actions based on Model Predictive Control (MPC) algorithm
Revised: We use a learned KINet model and select actions based on Model Predictive Control (MPC) algorithm

Original: (Finn & Levine, 2017) in the graph embedding space (GraphMPC).
Revised: (Finn & Levine, 2017) in the graph embedding space (GraphMPC).

Original: With KINet forward prediction, we compute the predicted graph representation of the next timestep for multiple sampled actions.
Revised: With KINet forward prediction, we compute the predicted graph representation of the next timestep for multiple sampled actions.

Original: We then select the optimal action that produces the most similar graph representation to the goal graph representation G goal .
Revised: We then select the optimal action that produces the most similar graph representation to the goal graph representation G goal .

Original: We describe our GraphMPC algorithm with a time horizon of T as: u ∗ t = arg max {S (cid:0) G goal , f fwd ( G t , { u t : T } ) (cid:1) } ; t ∈ [0 , T ] .
Revised: We describe our GraphMPC algorithm with a time horizon of T as: u ∗ t = arg max {S (cid:0) G goal , f fwd ( G t , { u t : T } ) (cid:1) } ; t ∈ [0 , T ] .

Original: Unlike performing conventional MPC only with respect to positional states, GraphMPC allows for accurately bringing the system to a goal state both explicitly (i.e, position) and implicitly (i.e, pose, orientation, and visual appearance).
Revised: Unlike performing conventional MPC only with respect to positional states, GraphMPC allows for accurately bringing the system to a goal state both explicitly (i.e, position) and implicitly (i.e, pose, orientation, and visual appearance).

Original:  E XPERIMENTAL S ETUPS
Revised:  E XPERIMENTAL S ETUPS

Original: Our experiments seek to address the following: (1) How accurate is our forward model?
Revised: Our experiments seek to address the following: (1) How accurate is our forward model?

Original: (2) Can we use the model for action planning?
Revised: (2) Can we use the model for action planning?

Original: and (3) Does the model generalize to unseen circumstances?
Revised: and (3) Does the model generalize to unseen circumstances?

Original: 4. D ATASET
Revised: 4. D ATASET

Original: We apply our approach to learn a forward model for multi-object manipulation tasks.
Revised: We apply our approach to learn a forward model for multi-object manipulation tasks.

Original: The task involves rearranging multiple objects in the scene and bringing them to the desired goal state usingpoint-to-point pushing actions.
Revised: The task involves rearranging multiple objects in the scene and bringing them to the desired goal state using pushing actions.

Original: We use MuJoCo 2.0 (Todorov et al., 2012) to generate simulated training and testing scenarios.
Revised: In the RealBlocks dataset, we use the Sawyer robot pushing dataset in Ye et al. (2020) to exemplify how our proposed framework applies to real settings.

Original:  Wealso use a real Sawyer robot pushing dataset from Ye et al.
Revised: RealBlocks data includes 5K random pushing actions on 7 blocks. For each action, an RGB image pair with the action vector in the image coordinates is captured. We use MuJoCo 2.0 (Todorov et al., 2012) to generate two sets of simulated scenarios.

Original: We simulate 10K episodes of random pushing on multiple objects (1-5 objects) where a simple robot end-effector applies randomized pushing for 90 timesteps per episode.
Revised: In BlockObjs dataset, we simulate 10K episodes of random pushing on multiple objects (1-5 objects) where a simple robot end-effector applies randomized pushing for 90 timesteps per episode (Figure A.1.2).

Original: Each object has randomized geometry and color.
Revised: Each object has randomized geometry and color.

Original: We sample the object geometry from a predefined continuous range ( geom train for training and geom gen for generalization to unseen geometries).
Revised: We sample the object geometry from a predefined continuous range denoted as geom train for training and geom gen for generalization to unseen geometries.

Original: Unseen geometries are designed to have elongated shapes to create complex out-of-distribution cases (see Figure 5).
Revised: Unseen geometries are designed to have elongated shapes to create complex out-ofdistribution cases (Figure A.1.1).

Revised: In YCBObjs dataset, we simulate 10K episodes (60 timesteps) of randomly pushing a subset of YCB objects placed on a wooden table (1-5 objects) (Calli et al., 2015) that includes objects of daily life with diverse properties such as shape, size, and texture (Figure A.1.3).

Original: Weonly collect the 4-dimensional action vector (pushing start and end location) and RGB imagesof the scene before and after the action is applied.
Revised: We collect the 4-dimensional action vector (pushing start and end location) and RGB images before and after each action is applied.

Original: Images are obtained using an overhead ( Top View ) and an angled camera ( Angled View )
Revised: Images are obtained using an overhead ( Top View ) and an angled camera ( Angled View )

Original: (see Appendix A.1 for more details).
Revised: (see Appendix A.1 for more details).

Original: For the real-robot data, we use the Sawyer robot pushing dataset in Ye et al.

Original: (2020) to exemplify how our proposed framework applies to the real settings.

Original: The real-robot data includes 5K random point-to-point pushing actions on 7 simple blocks.

Original: For each pushing action, an RGB image pair is provided along with the action vector in the image coordinates.

Original: 4. B ASELINES
Revised: 4. B ASELINES

Original: We compare our approach with existing methods for learning object-centric forward models:
Revised: We compare our approach with existing methods for learning object-centric forward models:

Original: Forward Model (Forw) :
Revised: Forward Model (Forw) :

Original: We train a convolutional encoder to extract visual features of the scene image ( Img ) and learn a forward model in the feature space.
Revised: We train a convolutional encoder to extract visual features of the scene image ( Img ) and learn a forward model in the feature space.

Original: Forward-Inverse Model (ForwInv) :
Revised: Forward-Inverse Model (ForwInv) :

Original: We train a convolutional encoder to extract visual features of the scene image ( Img ) and jointly learn forward and inverse modelsin the feature space as in Agrawal et al.
Revised: We train a convolutional encoder to extract visual features of the scene image ( Img ) and jointly learn forward and inverse models Agrawal et al.

Original: Each object has a separate feature channel.

Original: Interaction Network (IN) :
Revised: Interaction Network (IN) :

Original: We follow Battaglia et al.
Revised: We follow Battaglia et al.

Original: (2016); Sanchez-Gonzalez et al.
Revised: (2016); Sanchez-Gonzalez et al.

Original: (2018) to build an Interaction Network based on the ground truth location of the objects.
Revised: (2018) to build an Interaction Network based on the ground truth location of the objects.

Original: Each object representation contains the ground-truth position and velocity of the objects ( GT state ).
Revised: Each object representation contains the ground-truth position and velocity of the objects ( GT state ).

Original: This approach is only applicable to scenarios where the number of objects in the scene is known and fixed.
Revised: This approach is only applicable to scenarios where the number of objects in the scene is known and fixed.

Original: Visual Interaction Network (VisIN) :
Revised: Visual Interaction Network (VisIN) :

Original: To replicate Ye et al.
Revised: To replicate Ye et al.

Original: ; Watters et al.
Revised: ; Watters et al.

Original: (2017), we train a convolutional encoder to extract visual features of fixed-size bounding boxes centered on groundtruth object locations ( GT state + Img ).
Revised: (2017), we train a convolutional encoder to extract visual features of fixed-size bounding boxes centered on groundtruth object locations ( GT state + Img ).

Original: We use the extracted visual features as object representations in the Interaction Network.
Revised: We use the extracted visual features as object representations in the Interaction Network.

Original: This approachalso requires prior knowledge of the number of objects in the scene.
Revised: This approach requires prior knowledge of the number of objects.

Revised: Causal Discovery from Videos (V-CDN) :

Revised: Li et al.

Revised: (2020b) infers the causal structure of a fixed physical system from visual observations and makes future predictions for that system.

Revised: A pretrained perception module extracts keypoints that are used in an inference module to predict a graph representation for the visual observation which is then used in an IN-based dynamics module to predict the future location of the keypoints.

Original: 4. T RAINING AND E VALUATION S ETTING .
Revised: 4. T RAINING AND E VALUATION S ETTING .

Original: All models are trained on a subset of the simulated data  with 3 objects (8K episodes: 80% training, 10% validation, and 10% testing sets).
Revised: All models are trained on a subset of the simulated data ( BlockObjs and YCBObjs ) with 3 objects (8K episodes: 80% training, 10% validation, and 10% testing sets).

Original: To evaluate generalization to a different number of objects, we use other subsets of data with 1, 2, 4, and 5 objects ( ∼ 400 episodes for each case).
Revised: To evaluate generalization to a different number of objects, we use other subsets of data with 1, 2, 4, and 5 objects ( ∼ 400 episodes for each case).

Original: We also measure generalization for unseen geometries ( geom gen ).

Original: We train our model separately on images obtained from the overhead camera ( Top View ) and the angled camera ( Angled View ).
Revised: We train our model separately on images obtained from the overhead camera ( Top View ) and the angled camera ( Angled View ).

Original: For the realrobot data, we only provide qualitative results as the ground-truth location of the objects is unknown.
Revised: For RealBlocks data, we only provide qualitative results as the ground-truth location of the objects is unknown.

Original: We set the expected number of keypoints K = 6 for the simulation and K = 14 for the real robot.
Revised: We set the expected number of keypoints K = 6 for BlockObjs , K = 9 for YCBObjs , and K = 14 for RealBlocks .

Original:  R ESULTS
Revised:  R ESULTS

Original: This section is organized to answer a series of questions to thoroughly evaluate our model and justify the choices we made in formulating our approach.
Revised: This section is organized to answer a series of questions to thoroughly evaluate our model and justify the choices we made in formulating our approach.

Original: 5. D OES THE MODEL ACCURATELY LEARN A FORWARD MODEL ?
Revised: 5. D OES THE MODEL ACCURATELY LEARN A FORWARD MODEL ?

Original: First [Table] we evaluate the forward prediction accuracy.
Revised: First, we evaluate the forward prediction accuracy.

Original: Figure 2 showcases qualitative results of forward predictionsbased on real-robot data.
Revised: Figure 2 showcases qualitative results of forward predictions on RealBlocks .

Original: Our model factorizes the observation into a keypoint representation and accurately estimates the future appearance of the scene conditioned on external action.
Revised: Our model factorizes the observation into a keypoint representation and accurately estimates the future appearance of the scene conditioned on external action.

Original: These qualitative resultsalso highlight that our model learnsto predict the effect of the action on objects as well as object-to-object interactions (see Appendix A.3 for more examples).
Revised: The qualitative results highlight that our model learns the effect of the action on objects as well as objectto-object interactions (see Appendix A.4 for more examples).

Original: [Table]
Revised: [Table]

Original: We quantify the effectiveness of our model in comparison with Forw, ForwInv, IN, and VisIN baseline performances (Table 1) on the simulated data.
Revised: We quantify the effectiveness of our model in comparison with Forw, ForwInv, IN, and VisIN baselines (Table 1) on BlockObjs data.

Original: We separately train and examine each model on 
Revised: We separately train and examine each model on Top View and

Original: Top View and Angled View images.
Revised: Angled View images.

Original: The prediction error is computed as the average distance between the predicted and ground-truth positional states.
Revised: The prediction error is computed as the average distance between the predicted and ground-truth positional states.

Original: VisIN performs  best among baseline models as it builds object representations with explicit ground-truth object positions and their visual features.
Revised: VisIN performs the best among baseline models as it builds object representations with explicit ground-truth object positions and their visual features.

Original: Our model, on the other hand, achieves a comparable performance to VisIN while it does not rely on any supervision beyond the scene images.
Revised: Our model, on the other hand, achieves a comparable performance to VisIN while it does not rely on any supervision beyond the scene images.

Original: Forw and ForwInv baselines have similar supervision to ours but are significantly less accurate.
Revised: Forw and ForwInv baselines have similar supervision to ours but are significantly less accurate.

Original: This emphasizes the capability of our approach to learn a rich graph representation of the scene and an accurate forward model while relaxing prevailing assumptions of the prior work on the structure of the environment and availability of ground-truth state information (see Appendix A.2 formore insight into the learned graph representations).
Revised: This emphasizes the capability of our approach in learning a rich graph representation of the scene and an accurate forward model while relaxing prevailing assumptions of the prior work on the structure of the environment and availability of ground-truth state information (see Appendix A.2 for the learned graph representations).

Original: 5. C AN WE USE THE MODEL IN CONTROL TASKS ?
Revised: 5. C AN WE USE THE MODEL IN CONTROL TASKS ?

Original: To further compare our approach with baseline methods, we examine a robotic manipulation task of bringing a set of objects to the desired goal state using MPC.
Revised: We design a robotic manipulation task of rearranging a set of objects to the desired goal state using MPC with pushing actions.

Original: For all models, we run 1K trials where at each trial object geometries and initial pose are randomized and a random goal configuration of objects is also generated.
Revised: For all models, we run 1K episodes with randomized object geometries and initial poses and a random goal configuration of objects.

Original: The planning horizon is set to T = 40 timesteps.
Revised: The planning horizon is set to T =timesteps in each episode.

Original: For our model, we perform GraphMPC based on graph embedding similarity as described in Section 3.6.
Revised: For our model, we perform GraphMPC based on graph embedding similarity as described in Section 3.6.

Original: For all baseline models, we perform MPC directlybased on the distance to the goal.
Revised: For all baseline models, we perform MPC directly on the distance to the goal.

Original: Figure 3 shows MPC results over the planning horizon based on Top View observations.
Revised: Figure 5.2-a shows MPC results of BlockObjs based on Top View observations.

Original: Our approach is consistently faster than the baseline models in reaching the goal configuration.
Revised: Our approach is consistently faster than the baseline models in reaching the goal configuration.

Original: Additionally, we compute the pixel-wise distance to the goal image to evaluate whether the planned actions reach the goal configuration in terms of object appearances (i.e., orientation).

Original: 5. D OES THE MODEL GENERALIZE TO UNSEEN CIRCUMSTANCES ?
Revised: 5. D OES THE MODEL GENERALIZE TO UNSEEN CIRCUMSTANCES ?

Original: One of our main motivations to learn a forward model in keypoint space is to eliminate the dependency of the model formulation on the number of objects in the system.
Revised: One of our main motivations to learn a forward model in keypoint space is to eliminate the dependency of the model formulation on the number of objects in the system which allows for generalization to an unseen number of objects, object geometries, background textures, etc. BlockObjs .

Original: We  test for zero-shot generalization to a different number of objects (1, 2, 4, and 5 objects) and unseen object geometries ( geom gen ).
Revised: We train KINet on 3 randomized blocks with ( geom train ) and test for zero-shot generalization to an unseen number of objects (1, 2, 4, 5), unseen object geometries ( geom gen ), and unseen background texture.

Original: Figure 5 shows qualitative generalization results (see Appendix A.4 for more details).
Revised: Figure 4 shows qualitative generalization results.

Original: We separately train and examine for generalization on Top View and Angled View observations.
Revised: We separately train and examine for generalization on Top View and Angled View observations with a planning horizon of T = 80 .

Original: Note that baseline models, with the exception of Forw , are not able to generalize to a different number of objects because this number is hard-coded in their formulation (e.g., in the number of feature channels or the number of graph nodes).

Original: We set the planning horizon to T = 80 .

Original: Since our model learns to perform forward modeling in the keypoint space,with zero-shot generalization, it reassigns the keypoints to unseen objects andthen makes forward predictions.
Revised: Since our model learns to perform forward modeling in the keypoint space, it reassigns the keypoints to unseen objects and makes forward predictions.

Original: We also quantitatively evaluate thequality of generalization scenarios (Table 2).
Revised: Table 2 summarizes the generalization performance.

Original: As expected, by increasing the number of objects the average distance to the goal position increases.
Revised: As expected, by increasing the number of objects the average distance to the goal position increases.

Original: Also, objects with  geometriesthat were sampled out of the training dimension range have more distance to the goal position.
Revised: Also, objects with out-of-distribution geometries have more distance to the goal position.

Original: Our model’s generalization performance is significantly superior to Forw baseline which is a simple action-conditioned video prediction model that  uses the same supervision as our model.[Figure]
Revised: Our model’s generalization performance is significantly superior to Forw baseline which is a simple action-conditioned video prediction model that does not depend on the number of objects and uses the same supervision as our model.

Original: To further assess the generalizability, we testfor the performance of the model on unseen background textures (i.e., the table texture).
Revised: Further, we test the performance of the model on unseen background textures (i.e., the table texture).

Original: Our model is trained on 3 objects placed on a white-colored table ( domain train ).

Original: We test for zero-shot generalization to randomized table textures ( domain gen ). Figure 6) shows qualitative results of the control task performed in the unseen background textures. Since the keypoint extractionstep relies on visual features of salient objects, our model successfully performs the control tasks by ignoring the background andappropriately assigning keypoints to the  objects. 
Revised: Since the keypoint extraction relies on visual features of salient objects, our model is able to perform the control tasks by ignoring the background and assigning keypoints to the moving objects. Figure 5.2-b compares the MPC results for the KINet trained on a fixed white background ( fixed train ), zero-shot generalization of KINet trained on the fixed background to randomized backgrounds ( rand gen ), and KINet trained directly on randomized backgrounds ( rand train ).

Original: In extreme out-of-distribution cases such as the dirt texture (Fig 6, last row), a few keypoints are assigned tospecific parts of the background. We speculate that this is because the unseen texture, unlike the training set, has rich visual features (see Fig.
Revised: As expected, although the MPC converges, the final distance to the goal configuration is larger for rand gen . This final error is statistically at the same level of accuracy for fixed train and rand train . For qualitative examples, see Figure 4 (third row) and Figure A.5.2.

Original: A.5. for failed cases).
Revised: [Table] YCBObjs .

Original: Figure 4 compares the MPC results for the unseen randomized texture with the simple background. As expected, although the MPC converges, the final distance to thegoal configuration is larger for generalization to unseen textures.
Revised: We evaluate the generalization of our model on a set of realistic YCB objects (Calli et al., 2015) with challenging diverse properties such as color, texture, mass, and geometry.

Original: 5. A NALYSIS AND A BLATION
Revised: We train our model on a random subset of 3 YCB objects and test for generalization to an unseen number of objects (1,2,4,5) with a planning horizon of T = 40 . As shown in Figure 5, our method generalizes well to an unseen number of objects and performs the control task accurately. Importantly, assigning multiple keypoints to each object allows our framework to implicitly capture the orientation of each object, as well as their position without any supervision on the object pose (e.g., compare the power drill pose in Fig, 5). We compare the performance of our framework with V-CDN (Li et al., 2020b) baseline which is also a keypoint-based model to learn the structure of physical systems and perform future predictions and potentially generalize to an unseen number of objects.

Original: We justify the major choices we made to formulate the model through ablation studies.
Revised: Although V-CDN is formulated to extract the causal structure of a fixed system through visual observations, we pretrain its perception module on our randomized YCBObjs dataset for a direct comparison to our model. Also, to perform a control task, we condition the V-CDN on the action by adding an encoding of the action vector to each keypoint embedding. Figure 6 compares the MPC results (normalized to the number of objects) of our method and V-CDN both trained on a subset of 3 YCB objects. The MPC performance of both models is comparable for rearranging 3 objects (green lines).

Original: Specifically, we examine two elements in our approach: the probabilistic graph representation, and the contrastive [Figure]
Revised: However, V-CDN is most accurate for the number of objects it has been trained on and significantly less accurate when generalizing to an unseen number of objects. We attribute this to two reasons: (1) Although using keypoints,

Original: [Table] [Figure]
Revised: V-CDN is formulated to infer explicit causal structure and attributes for the graph representation of a fixed system setup which does not necessarily carry over to unseen circumstances.

Original: [Table]
Revised: (2) Unlike our method, V-CDN does not take into account the visual features in the model and only uses the keypoint positions.

Original: [Figure] [Table]
Revised: 5. A NALYSIS AND A BLATION

Original: [Figure] [Table]
Revised: We justify the major choices we made to formulate the model with ablation studies.

Original:  loss. We train two variants of KINet: (1) KINet - deterministic in which the graph representation is fully connectedand not probabilistic and there’s no inference on the graph adjacency matrix. KINet - no ctr losswhere we train the model without the contrastive loss.
Revised: We examine two elements in our approach: the probabilistic graph representation, and the contrastive loss. We train two variants of KINet: (1) KINet - deterministic with a fully connected graph instead of probabilistic (2) KINet - no ctr loss without the contrastive loss.

Original: The best forward predictionperformance for both Top View and Angled View observations is achieved when the model is probabilistic and trained with a contrastive loss (Table 1).
Revised: The best forward prediction for both Top View and Angled View images is achieved when the model is probabilistic and trained with a contrastive loss (Table 1).

Original: The contrastive loss is an essential element in our approach to ensure the learned forward model is accurately action conditional.
Revised: The contrastive loss is an essential element in our approach to ensure the learned forward model is accurately action conditional.

Original: Also, a probabilistic graph representation allows for more efficient information aggregation of neighboring nodes.

Original: With a probabilistic graph representation, our model achieves higher generalizationperformance compared to its deterministic variant.
Revised: Also, with a probabilistic graph representation, our model achieves better generalization compared to the deterministic variant.

Original: This performance gap is more evident when generalizing to unseen geometries (Table 2).
Revised: This performance gap is more evident when generalizing to unseen geometries (Table 2).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Summary.

Original: In this paper, we proposed a method for learning action-conditioned forward models based only on image observations.
Revised: In this paper, we proposed a method for learning action-conditioned forward models based only on image observations.

Original: We showed that our approach effectively makes forward predictions through keypoint factorization.
Revised: We showed that our approach effectively makes forward predictions through keypoint factorization.

Original: Also, we demonstrated that a keypoint-based forward model, unlike prior work, does not make assumptions about the number of objects which allows for automatic generalization to a variety of unseen circumstances.
Revised: Also, we demonstrated that a keypoint-based forward model, unlike prior work, does not make assumptions about the number of objects which allows for automatic generalization to a variety of unseen circumstances.

Original: Importantly, our model learns a forward model without explicit supervision on ground-truth object state information.
Revised: Importantly, our model learns a forward model without explicit supervision on ground-truth object state information.

Original: Limitations and future work.

Original: One limitation in our formulation is fixing the number of expected keypoints; however, we showed this gives more generalizability compared to fixing the number of objects.
Revised: One limitation in our formulation is fixing the number of expected keypoints (see Appendix A.3). However, we showed this gives more generalizability compared to fixing the number of objects.

Original: We also observed inconsistency in the predicted keypoints for real-robot data scenarios where most of the objects were pushed out of the image frame (see Appendix A.5).
Revised: We also observed inconsistency in the predicted keypoints for real-robot data scenarios where most of the objects were pushed out of the image frame (see Appendix A.6).

Original: An interesting future direction is to focus on the keypoint extraction module to further enhance forward models for real settings.
Revised: An interesting future direction is to focus on the keypoint extraction module to further enhance forward models for real settings.

Original: Finally, we hope our general approach inspires future research on physical reasoning in settings where ground-truth information is hard to obtain.
Revised: Finally, we hope our general approach inspires future research on physical reasoning in settings where ground-truth information is hard to obtain.