Original: C ONTRASTIVE L EARNING OF S TRUCTURED W ORLD M ODELS
Revised: C ONTRASTIVE L EARNING OF S TRUCTURED W ORLD M ODELS

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A structured understanding of our world in terms of objects, relations, and hierarchies is an important component of human cognition.
Revised: A structured understanding of our world in terms of objects, relations and hierarchies is an important component of human cognition.

Original: Learning such a structured world model from raw sensory data remains a challenge.
Revised: Learning such a structured world model from raw sensory data remains a challenge.

Original: As a step towards this goal we introduce Contrastively-trained Structured World Models C SWMs).
Revised: As a step towards this goal, we introduce Contrastively-trained Structured World Models (C-SWMs).

Original: C SWMs utilize a contrastive approach for representation learning in environments with compositional structure.
Revised: CSWMs utilize a contrastive approach for representation learning in environments with compositional structure.

Original: We structure each state embedding as a set of ob ject representations and their relations , modeled by a graph neural network.
Revised: We structure each state embedding as a set of object representations and their relations, modeled by a graph neural network.

Original: This allows objects to be discovered from raw pixel observations without direct super- vision as part of the learning process.
Revised: This allows objects to be discovered from raw pixel observations without direct supervision as part of the learning process.

Original: We evaluate C SWMs on compositional environments involving multiple interacting objects that can be manipulated independently by an agent simple Atari games, and a multi object physics simulation
Revised: We evaluate C-SWMs on compositional environments involving multiple interacting objects that can be manipulated independently by an agent, simple Atari games and a multi-object physics simulation.

Original: Our experiments demonstrate that C-SWMs can overcome limitations of models based on pixel reconstruction and outperform typical representatives of this model class in highly structured environments, while learning interpretable object-based representations.
Revised: Our experiments demonstrate that C-SWMs can overcome limitations of models based on pixel reconstruction and outperform typical representatives of this model class in highly structured environments, while learning interpretable object-based representations

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Compositional reasoning in terms of objects, relations, and actions is a central ability in human cognition (Spelke & Kinzler, 2007).
Revised: Compositional reasoning in terms of objects, relations, and actions is a central ability in human cognition (Spelke & Kinzler, 2007).

Original: This ability serves as a core motivation behind a range of recent works that aim at enriching machine learning models with the ability to disentangle scenes into objects, their properties, and relations between them (Chang et al., 2016; Battaglia et al., 2016;
Revised: This ability serves as a core motivation behind a range of recent works that aim at enriching machine learning models with the ability to disentangle scenes into objects, their properties, and relations between them (Chang et al., 2016; Battaglia et al., 2016;

Original: Watters et al., 2017; van Steenkiste et al., 2018; Kipf et al., 2018; Sun et al., 2018; 2019b;
Revised: Watters et al., 2017; van Steenkiste et al., 2018; Kipf et al., 2018; Sun et al., 2018; 2019b;

Original: Xu et al., 2019).
Revised: Xu et al., 2019).

Original: These structured neural models greatly facilitate predicting physical dynamics and the consequences of actions, and provide a strong inductive bias for generalization to novel environment situations, allowing models to answer counterfactual questions such as “ What would happen if I pushed this block instead of pulling it? ”.
Revised: These structured neural models greatly facilitate predicting physical dynamics and the consequences of actions, and provide a strong inductive bias for generalization to novel environment situations, allowing models to answer counterfactual questions such as “ What would happen if I pushed this block instead of pulling it? ”.

Original: Arriving at a structured description of the world in terms of objects and relations in the ﬁrst place, however, is a challenging problem.
Revised: Arriving at a structured description of the world in terms of objects and relations in the ﬁrst place, however, is a challenging problem.

Original: While most methods in this area require some form of human annotation for the extraction of objects or relations, several recent works study the problem of object discovery from visual data in a completely unsupervised or self-supervised manner (Greff et al., 2017; Nash et al., 2017; van Steenkiste et al., 2018;  Janner et al., 2019; Xu et al., 2019; Burgess et al., 2019; Greff et al., 2019; Engelcke et al., 2019).
Revised: While most methods in this area require some form of human annotation for the extraction of objects or relations, several recent works study the problem of object discovery from visual data in a completely unsupervised or self-supervised manner (Eslami et al., 2016; Greff et al., 2017; Nash et al., 2017; van Steenkiste et al., 2018; Kosiorek et al., 2018; Janner et al., 2019; Xu et al., 2019; Burgess et al., 2019; Greff et al., 2019; Engelcke et al., 2019).

Original: These methods follow a generative approach, i.e., they learn to discover object-based representations by performing visual predictions or reconstruction and by optimizing an objective in pixel space.
Revised: These methods follow a generative approach, i.e., they learn to discover object-based representations by performing visual predictions or reconstruction and by optimizing an objective in pixel space.

Original: Placing a loss in pixel space requires carefully trading off structural constraints on latent variables vs. accuracy of pixel-based reconstruction.
Revised: Placing a loss in pixel space requires carefully trading off structural constraints on latent variables vs. accuracy of pixel-based reconstruction.

Original: Typical failure modes include ignoring visually small, but relevant features for predicting the future, such as a bullet in an Atari game (Kaiser et al., 2019), or wasting model capacity on visually rich, but otherwise potentially irrelevant features, such as static backgrounds.
Revised: Typical failure modes include ignoring visually small, but relevant features for predicting the future, such as a bullet in an Atari game (Kaiser et al., 2019), or wasting model capacity on visually rich, but otherwise potentially irrelevant features, such as static backgrounds.

Original: To avoid such failure modes, we propose to adopt a discriminative approach using contrastive learning, which scores real against fake experiences in the form of state-action-state triples from an experience buffer (Lin, 1992), in a similar fashion as typical graph embedding approaches score true facts in the form of entity-relation-entity triples against corrupted triples or fake facts.
Revised: To avoid such failure modes, we propose to adopt a discriminative approach using contrastive learning, which scores real against fake experiences in the form of state-action-state triples from an experience buffer (Lin, 1992), in a similar fashion as typical graph embedding approaches score true facts in the form of entity-relation-entity triples against corrupted triples or fake facts.

Original: We introduce Contrastively-trained Structured World Models (C-SWMs), a class of models for learning abstract state representations from observations in an environment.
Revised: We introduce Contrastively-trained Structured World Models (C-SWMs), a class of models for learning abstract state representations from observations in an environment.

Original: C-SWMs learn a set of abstract state variables, one for each object in a particular observation.
Revised: C-SWMs learn a set of abstract state variables, one for each object in a particular observation.

Original: Environment transitions are modeled using a graph neural network (Scarselli et al., 2009;
Revised: Environment transitions are modeled using a graph neural network (Scarselli et al., 2009;

Original: Li et al., 2015; Kipf & Welling, 2016;
Revised: Li et al., 2015; Kipf & Welling, 2016;

Original: Gilmer et al., 2017; Battaglia et al., 2018) that operates on latent abstract representations.
Revised: Gilmer et al., 2017; Battaglia et al., 2018) that operates on latent abstract representations.

Original: This paper further introduces a novel object-level contrastive loss for unsupervised learning of object-based representations.
Revised: This paper further introduces a novel object-level contrastive loss for unsupervised learning of object-based representations.

Original: We arrive at this formulation by adapting methods for learning translational graph embeddings (Bordes et al., 2013; Wang et al., 2014) to our use case.
Revised: We arrive at this formulation by adapting methods for learning translational graph embeddings (Bordes et al., 2013; Wang et al., 2014) to our use case.

Original: By establishing a connection between contrastive learning of state abstractions (Franc¸ois-Lavet et al., 2018; Thomas et al., 2018) and relational graph embeddings (Nickel et al., 2016a), we hope to provide inspiration and guidance for future model improvements in both ﬁelds.
Revised: By establishing a connection between contrastive learning of state abstractions (Franc¸ois-Lavet et al., 2018; Thomas et al., 2018) and relational graph embeddings (Nickel et al., 2016a), we hope to provide inspiration and guidance for future model improvements in both ﬁelds.

Original: In a set of experiments, where we use a novel ranking-based evaluation strategy, we demonstrate that
Revised: In a set of experiments, where we use a novel ranking-based evaluation strategy, we demonstrate that

Original: C-SWMs learn interpretable object-level state abstractions, accurately learn to predict state transitions many steps into the future, demonstrate combinatorial generalization to novel environment conﬁgurations and learn to identify objects from scenes without supervision.
Revised: C-SWMs learn interpretable object-level state abstractions, accurately learn to predict state transitions many steps into the future, demonstrate combinatorial generalization to novel environment conﬁgurations and learn to identify objects from scenes without supervision.

Original:  S TRUCTURED W ORLD M ODELS
Revised:  S TRUCTURED W ORLD M ODELS

Original: Our goal is to learn an object-oriented abstraction of a particular observation or environment state.
Revised: Our goal is to learn an object-oriented abstraction of a particular observation or environment state.

Original: In addition, we would like to learn an action-conditioned transition model of the environment that takes object representations and their relations and interactions into account.
Revised: In addition, we would like to learn an action-conditioned transition model of the environment that takes object representations and their relations and interactions into account.

Original: We start by introducing the general framework for contrastive learning of state abstractions and transition models without object factorization in Sections 2.1–2.2, and in the following describe a variant that utilizes object-factorized state representations, which we term a Structured World Model.
Revised: We start by introducing the general framework for contrastive learning of state abstractions and transition models without object factorization in Sections 2.1–2.2, and in the following describe a variant that utilizes object-factorized state representations, which we term a Structured World Model.

Original: 2. S TATE A BSTRACTION
Revised: 2. S TATE A BSTRACTION

Original: We consider an off-policy setting, where we operate solely on a buffer of ofﬂine experience, e.g., obtained from an exploration policy.
Revised: We consider an off-policy setting, where we operate solely on a buffer of ofﬂine experience, e.g., obtained from an exploration policy.

Original: Formally, this experience buffer B = { ( s t , a t , s t +1 ) }
Revised: Formally, this experience buffer B = { ( s t , a t , s t +1 ) }

Original: Tt =1 contains T tuples of states s t ∈ S , actions a t ∈ A , and follow-up states s t +1 ∈ S , which are reached after taking action a t .
Revised: Tt =1 contains T tuples of states s t ∈ S , actions a t ∈ A , and follow-up states s t +1 ∈ S , which are reached after taking action a t .

Original: We do not consider rewards as part of our framework for simplicity.
Revised: We do not consider rewards as part of our framework for simplicity.

Original: Our goal is to learn abstract or latent representations z t ∈ Z of environment states s t ∈ S that discard any information which is not necessary to predict the abstract representation of
Revised: Our goal is to learn abstract or latent representations z t ∈ Z of environment states s t ∈ S that discard any information which is not necessary to predict the abstract representation of

Original: the followup state z t +1 ∈
Revised: the followup state z t +1 ∈

Original: Z after taking action a t .
Revised: Z after taking action a t .

Original: Formally, we have an encoder E : S → Z which maps observed states to abstract state representations and a transition model T : Z ×
Revised: Formally, we have an encoder E : S → Z which maps observed states to abstract state representations and a transition model T : Z ×

Original: A
Revised: A

Original: →
Revised: →

Original: Z operating solely on abstract state representations.
Revised: Z operating solely on abstract state representations.

Original: 2. C ONTRASTIVE L EARNING
Revised: 2. C ONTRASTIVE L EARNING

Original: Our starting point is the graph embedding method TransE (Bordes et al., 2013):
Revised: Our starting point is the graph embedding method TransE (Bordes et al., 2013):

Original: TransE embeds facts from a knowledge base K = { ( e t , r t , o t ) }
Revised: TransE embeds facts from a knowledge base K = { ( e t , r t , o t ) }

Original: Tt =1 , which consists of entity-relation-entity triples ( e t , r t , o t ) , where e t is the subject entity (analogous to the source state s t in our case), r t is the relation (analogous to the action a t in our experience buffer), and o t is the object entity (analogous to the target state s t +1 ).
Revised: Tt =1 , which consists of entity-relation-entity triples ( e t , r t , o t ) , where e t is the subject entity (analogous to the source state s t in our case), r t is the relation (analogous to the action a t in our experience buffer), and o t is the object entity (analogous to the target state s t +1 ).

Original: TransE deﬁnes the energy of a triple ( e t , r t , o t ) as H =
Revised: TransE deﬁnes the energy of a triple ( e t , r t , o t ) as H =

Original: d ( F ( e t )
Revised: d ( F ( e t )

Original: + G ( r t )
Revised: + G ( r t )

Original: , F ( o t )) , where F (and
Revised: , F ( o t )) , where F (and

Original: G ) are embedding functions that map discrete entities (and relations) to R D , where D is the dimensionality of the embedding space, and d ( · , · ) denotes the squared Euclidean distance.
Revised: G ) are embedding functions that map discrete entities (and relations) to R D , where D is the dimensionality of the embedding space, and d ( · , · ) denotes the squared Euclidean distance.

Original: Training is carried out with an energy-based hinge loss (LeCun et al., 2006), with negative samples obtained by replacing the entities in a fact with random entities from the knowledge base.
Revised: Training is carried out with an energy-based hinge loss (LeCun et al., 2006), with negative samples obtained by replacing the entities in a fact with random entities from the knowledge base.

Original: We can port TransE to our setting with only minor modiﬁcations.
Revised: We can port TransE to our setting with only minor modiﬁcations.

Original: As the effect of an action is in general not independent of the source state, we replace G ( r t ) with T ( z t , a t ) , i.e., with the transition function, conditioned on both the action and the (embedded) source state via z t =
Revised: As the effect of an action is in general not independent of the source state, we replace G ( r t ) with T ( z t , a t ) , i.e., with the transition function, conditioned on both the action and the (embedded) source state via z t =

Original: E ( s t ) .
Revised: E ( s t ) .

Original: The overall energy of a state-action-state triple then can be deﬁned as follows:
Revised: The overall energy of a state-action-state triple then can be deﬁned as follows:

Original: H =
Revised: H =

Original: d ( z t + T ( z t , a t ) , z t +1 ) .
Revised: d ( z t + T ( z t , a t ) , z t +1 ) .

Original: This additive form of the transition model provides a strong inductive bias for modeling effects of actions in the environment as translations in the abstract state space.
Revised: This additive form of the transition model provides a strong inductive bias for modeling effects of actions in the environment as translations in the abstract state space.

Original: Alternatively, one could model effects as linear transformations or rotations in the abstract state space, which motivates the use of a graph embedding method such as RESCAL (Nickel et al., 2011), CompleX (Trouillon et al., 2016), or HolE (Nickel et al., 2016b).
Revised: Alternatively, one could model effects as linear transformations or rotations in the abstract state space, which motivates the use of a graph embedding method such as RESCAL (Nickel et al., 2011), CompleX (Trouillon et al., 2016), or HolE (Nickel et al., 2016b).

Original: With the aforementioned modiﬁcations, we arrive at the following energy-based hinge loss:
Revised: With the aforementioned modiﬁcations, we arrive at the following energy-based hinge loss:

Original: [Equation]
Revised: [Equation]

Original: deﬁned for a single ( s t , a t , s t +1 ) with a corrupted abstract state ˜
Revised: deﬁned for a single ( s t , a t , s t +1 ) with a corrupted abstract state ˜

Original: z t = E (˜ s t ) .
Revised: z t = E (˜ s t ) .

Original: ˜ s t is sampled at random from the experience buffer.
Revised: ˜ s t is sampled at random from the experience buffer.

Original: The margin γ is a hyperparameter for which we found γ = 1 to be a good choice.
Revised: The margin γ is a hyperparameter for which we found γ = 1 to be a good choice.

Original: Unlike Bordes et al.
Revised: Unlike Bordes et al.

Original: (2013), we place the hinge only on the negative term instead of on the full loss and we do not constrain the norm of the abstract states z t , which we found to work better in our context (see Appendix A.3).
Revised: (2013), we place the hinge only on the negative term instead of on the full loss and we do not constrain the norm of the abstract states z t , which we found to work better in our context (see Appendix A.3).

Original: The overall loss is to be understood as an expectation of the above over samples from the experience buffer B .
Revised: The overall loss is to be understood as an expectation of the above over samples from the experience buffer B .

Original: 2. O BJECT - ORIENTED S TATE F ACTORIZATION
Revised: 2. O BJECT - ORIENTED S TATE F ACTORIZATION

Original: Our goal is to take into account the compositional nature of visual scenes, and hence we would like to learn a relational and object-oriented model of the environment that operates on a factored abstract state space Z
Revised: Our goal is to take into account the compositional nature of visual scenes, and hence we would like to learn a relational and object-oriented model of the environment that operates on a factored abstract state space Z

Original: = Z 1 × . . .
Revised: = Z 1 × . . .

Original: ×Z K , where K is the number of available object slots.
Revised: ×Z K , where K is the number of available object slots.

Original: We further assume an object-factorized action space
Revised: We further assume an object-factorized action space

Original: A =
Revised: A =

Original: A 1 × . . .
Revised: A 1 × . . .

Original: × A K .
Revised: × A K .

Original: This factorization ensures that each object is independently represented and it allows for efﬁcient sharing of model parameters across objects in the transition model.
Revised: This factorization ensures that each object is independently represented and it allows for efﬁcient sharing of model parameters across objects in the transition model.

Original: This serves as a strong inductive bias for better generalization to novel scenes and facilitates learning and object discovery.
Revised: This serves as a strong inductive bias for better generalization to novel scenes and facilitates learning and object discovery.

Original: The overall C-SWM model architecture using object-factorized representations is shown in Figure 1.
Revised: The overall C-SWM model architecture using object-factorized representations is shown in Figure 1.

Original: [Figure]
Revised: [Figure]

Original: model is composed of the following components: 1) a CNN-based object extractor, 2) an MLP-based object encoder, 3) a GNN-based relational transition model, and 4) an object-factorized contrastive loss.

Original: Colored blocks denote abstract states for a particular object.

Original: Encoder and Object Extractor
Revised: Encoder and Object Extractor

Original: We split the encoder into two separate modules: 1) a CNN-based object extractor E ext , and 2) an MLP-based object encoder E enc .
Revised: We split the encoder into two separate modules: 1) a CNN-based object extractor E ext , and 2) an MLP-based object encoder E enc .

Original: The object extractor module is a
Revised: The object extractor module is a

Original: CNN operating directly on image-based observations from the environment with K feature maps in its last layer.
Revised: CNN operating directly on image-based observations from the environment with K feature maps in its last layer.

Original: Each feature map m k t =
Revised: Each feature map m kt =

Original: [ E ext ( s t )]
Revised: [ E ext ( s t )]

Original: k can be interpreted as an object mask corresponding to one particular object slot, where [ . . . ]
Revised: k can be interpreted as an object mask corresponding to one particular object slot, where [ . . . ]

Original: k denotes selection of the k -th feature map.
Revised: k denotes selection of the k -th feature map.

Original: For simplicity, we only assign a single feature map per object slot which sufﬁced for the experiments considered in this work (see Appendix A.4).
Revised: For simplicity, we only assign a single feature map per object slot which sufﬁced for the experiments considered in this work (see Appendix A.4).

Original: To allow for encoding of more complex object features (other than, e.g., position/velocity), the object extractor can be adapted to produce multiple feature maps per object slot.
Revised: To allow for encoding of more complex object features (other than, e.g., position/velocity), the object extractor can be adapted to produce multiple feature maps per object slot.

Original: After the object extractor module, we ﬂatten each feature map m kt (object mask) and feed it into the object encoder E enc .
Revised: After the object extractor module, we ﬂatten each feature map m kt (object mask) and feed it into the object encoder E enc .

Original: The object encoder shares weights across objects andreturns an abstract state representation: z kt =
Revised: object encoder shares weights across objects and abstract state representation: z k t =

Original: E enc ( m kt ) with z kt ∈ Z k .
Revised: E enc ( m kt ) with z k t ∈ Z k .

Original: We set Z k
Revised: We set Z k

Original: = R D in the following, where D is a hyperparameter.
Revised: = R D in the following, where D is a hyperparameter.

Original: Relational Transition Model We implement the transition model as a graph neural network (Scarselli et al., 2009;
Revised: Relational Transition Model We implement the transition model as a graph neural network (Scarselli et al., 2009;

Original: Li et al., 2015; Kipf & Welling, 2016; Battaglia et al., 2016;
Revised: Li et al., 2015; Kipf & Welling, 2016; Battaglia et al., 2016;

Original: Gilmer et al., 2017; Battaglia et al., 2018), which allows us to model pairwise interactions between object states while being invariant to the order in which objects are represented.
Revised: Gilmer et al., 2017; Battaglia et al., 2018), which allows us to model pairwise interactions between object states while being invariant to the order in which objects are represented.

Original: After the encoder stage, we have an abstract state description z kt ∈ Z k and an action a kt ∈
Revised: After the encoder stage, we have an abstract state description z kt ∈ Z k and an action a kt ∈

Original: A k for every object in the scene.
Revised: A k for every object in the scene.

Original: We represent actions as one-hot vectors (or a vector of zeros if no action is applied to a particular object), but note that other choices are possible, e.g., for continuous action spaces.
Revised: We represent actions as one-hot vectors (or a vector of zeros if no action is applied to a particular object), but note that other choices are possible, e.g., for continuous action spaces.

Original: The transition function then takes as input the tuple of object representations z t
Revised: The transition function then takes as input the tuple of object representations z t

Original: = ( z 1 t , . . .
Revised: = ( z 1 t , . . .

Original: , z Kt ) and actions a t = ( a 1 t , . . .
Revised: , z Kt ) and actions a t = ( a 1 t , . . .

Original: , a Kt ) at a particular time step:
Revised: , a Kt ) at a particular time step:

Original: [Equation]
Revised: [Equation]

Original: T ( z t , a t ) is implemented as a graph neural network (GNN) that takes z kt as input node features.
Revised: T ( z t , a t ) is implemented as a graph neural network (GNN) that takes z kt as input node features.

Original: The model predicts updates ∆ z t = (∆ z 1 t , . . .
Revised: The model predicts updates ∆ z t = (∆ z 1 t , . . .

Original: , ∆ z Kt ) .
Revised: , ∆ z Kt ) .

Original: The object representations for the next time step are obtained via z t +1 =
Revised: The object representations for the next time step are obtained via z t +1 =

Original: ( z 1 t + ∆ z 1 t , . . .
Revised: ( z 1 t + ∆ z 1 t , . . .

Original: , z Kt + ∆ z Kt ) .
Revised: , z Kt + ∆ z Kt ) .

Original: The GNN consists of node update functions f node and edge update functions f edge with shared parameters across all nodes and edges.
Revised: The GNN consists of node update functions f node and edge update functions f edge with shared parameters across all nodes and edges.

Original: These functions are implemented as MLPs and we choose the following form of message passing updates:
Revised: These functions are implemented as MLPs and we choose the following form of message passing updates:

Original: [Equation]
Revised: [Equation]

Original: where e ( i,j ) t is an intermediate representation of the edge or interaction between nodes i and j .
Revised: where e ( i,j ) t is an intermediate representation of the edge or interaction between nodes i and j .

Original: This corresponds to a single round of node-to-edge and edge-to-node message passing.
Revised: This corresponds to a single round of node-to-edge and edge-to-node message passing.

Original: Alternatively, one could apply multiple rounds of message passing, but we did not ﬁnd this to be necessary for the experiments considered in this work.
Revised: Alternatively, one could apply multiple rounds of message passing, but we did not ﬁnd this to be necessary for the experiments considered in this work.

Original: Note that this update rule corresponds to message passing on a fully-connected scene graph, which is O ( K 2 ) .
Revised: Note that this update rule corresponds to message passing on a fully-connected scene graph, which is O ( K 2 ) .

Original: This can be reduced to linear complexity by reducing connectivity to nearest neighbors in the abstract state space, which we leave for future work.
Revised: This can be reduced to linear complexity by reducing connectivity to nearest neighbors in the abstract state space, which we leave for future work.

Original: We denote the output of the transition function for the k -th object as ∆ z kt =
Revised: We denote the output of the transition function for the k -th object as ∆ z kt =

Original: T k ( z t , a t ) in the following.
Revised: T k ( z t , a t ) in the following.

Original: Multi-object Contrastive Loss We only need to change the energy function to take the factorization of the abstract state space into account, which yields the following energy H for positive triples and ˜ H for negative samples:
Revised: Multi-object Contrastive Loss We only need to change the energy function to take the factorization of the abstract state space into account, which yields the following energy H for positive triples and ˜ H for negative samples:

Original: [Equation]
Revised: [Equation]

Original: where ˜ z kt is the k -th object representation of the negative state sample
Revised: where ˜ z kt is the k -th object representation of the negative state sample

Original: ˜ z t = E
Revised: ˜ z t = E

Original: ˜ s t ) .
Revised: ˜ s t ) .

Original: The overall contrastive loss for a single state-action-state sample from the experience buffer then takes the form:
Revised: The overall contrastive loss for a single state-action-state sample from the experience buffer then takes the form:

Original: [Equation]
Revised: [Equation]

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

Original: For coverage of related work in the area of object discovery with autoencoder-based models, we refer the reader to the Introduction section.
Revised: For coverage of related work in the area of object discovery with autoencoder-based models, we refer the reader to the Introduction section.

Original: We further discuss related work on relational graph embeddings in Section 2.2.
Revised: We further discuss related work on relational graph embeddings in Section 2.2.

Original: Structured Models of Environments
Revised: Structured Models of Environments

Original: Recent work on modeling structured environments such as interacting multi-object or multi-agent systems has made great strides in improving predictive accuracy by explicitly taking into account the structured nature of such systems (Sukhbaatar et al., 2016;
Revised: Recent work on modeling structured environments such as interacting multi-object or multi-agent systems has made great strides in improving predictive accuracy by explicitly taking into account the structured nature of such systems (Sukhbaatar et al., 2016;

Original: Chang et al., 2016; Battaglia et al., 2016; Watters et al., 2017; Hoshen, 2017; Wang et al., 2018; van Steenkiste et al., 2018; Kipf et al., 2018; Sanchez-Gonzalez et al., 2018; Xu et al., 2019).
Revised: Chang et al., 2016; Battaglia et al., 2016; Watters et al., 2017; Hoshen, 2017; Wang et al., 2018; van Steenkiste et al., 2018; Kipf et al., 2018; Sanchez-Gonzalez et al., 2018; Xu et al., 2019).

Original: These methods generally make use of some form of graph neural network, where node update functions model the dynamics of individual objects, parts or agents and edge update functions model their interactions and relations.
Revised: These methods generally make use of some form of graph neural network, where node update functions model the dynamics of individual objects, parts or agents and edge update functions model their interactions and relations.

Original: Several recent works succeed in learning such structured models directly from pixels (Watters et al., 2017; van Steenkiste et al., 2018; Xu et al., 2019; Watters et al., 2019), but in contrast to our work rely on pixel-based loss functions.
Revised: Several recent works succeed in learning such structured models directly from pixels (Watters et al., 2017; van Steenkiste et al., 2018; Xu et al., 2019; Watters et al., 2019), but in contrast to our work rely on pixel-based loss functions.

Original: The latest example in this line of research is the COBRA model (Watters et al., 2019), which learns an action-conditioned transition policy on object representations obtained from an unsupervised object discovery model (Burgess et al., 2019).
Revised: The latest example in this line of research is the COBRA model (Watters et al., 2019), which learns an action-conditioned transition policy on object representations obtained from an unsupervised object discovery model (Burgess et al., 2019).

Original: Unlike C-SWM, COBRA does not model interactions between object slots and relies on a pixel-based loss for training.
Revised: Unlike C-SWM, COBRA does not model interactions between object slots and relies on a pixel-based loss for training.

Original: Our object encoder, however, is more limited and utilizing an iterative object encoding process such as in MONet (Burgess et al., 2019) would be interesting for future work.
Revised: Our object encoder, however, is more limited and utilizing an iterative object encoding process such as in MONet (Burgess et al., 2019) would be interesting for future work.

Original: Contrastive Learning Contrastive learning methods are widely used in the ﬁeld of graph representation learning (Bordes et al., 2013; Perozzi et al., 2014; Grover & Leskovec, 2016;
Revised: Contrastive Learning Contrastive learning methods are widely used in the ﬁeld of graph representation learning (Bordes et al., 2013; Perozzi et al., 2014; Grover & Leskovec, 2016;

Original: Bordes et al., 2013; Schlichtkrull et al., 2018; Veliˇckovi´c et al., 2018), and for learning word representations (Mnih & Teh, 2012; Mikolov et al., 2013).
Revised: Bordes et al., 2013; Schlichtkrull et al., 2018; Veliˇckovi´c et al., 2018), and for learning word representations (Mnih & Teh, 2012; Mikolov et al., 2013).

Original: The main idea is to construct pairs of related data examples (positive examples, e.g., connected by an edge in a graph or co-occuring words in a sentence) and pairs of unrelated or corrupted data examples (negative examples), and use a loss function that scores positive and negative pairs in a different way.
Revised: The main idea is to construct pairs of related data examples (positive examples, e.g., connected by an edge in a graph or co-occuring words in a sentence) and pairs of unrelated or corrupted data examples (negative examples), and use a loss function that scores positive and negative pairs in a different way.

Original: Most energy-based losses (LeCun et al., 2006) are suitable for this task.
Revised: Most energy-based losses (LeCun et al., 2006) are suitable for this task.

Original: Recent works (Oord et al., 2018; Hjelm et al., 2018; H´enaff et al., 2019;
Revised: Recent works (Oord et al., 2018; Hjelm et al., 2018; H´enaff et al., 2019;

Original: Sun et al., 2019a; Anand et al., 2019) connect objectives of this kind to the principle of learning representations by maximizing mutual information between data and learned representations, and successfully apply these methods to image, speech, and video data.
Revised: Sun et al., 2019a; Anand et al., 2019) connect objectives of this kind to the principle of learning representations by maximizing mutual information between data and learned representations, and successfully apply these methods to image, speech, and video data.

Original: State Representation Learning State representation learning in environments similar to ours is often approached by models based on autoencoders (Corneil et al., 2018;
Revised: State Representation Learning State representation learning in environments similar to ours is often approached by models based on autoencoders (Corneil et al., 2018;

Original: Watter et al., 2015; Ha &
Revised: Watter et al., 2015; Ha &

Original: Schmidhuber, 2018; Hafner et al., 2019; Laversanne-Finot et al., 2018) or via adversarial learning (Kurutach et al., 2018; Wang et al., 2019).
Revised: Schmidhuber, 2018; Hafner et al., 2019; Laversanne-Finot et al., 2018) or via adversarial learning (Kurutach et al., 2018; Wang et al., 2019).

Original: Some recent methods learn state representations without requiring a decoder back into pixel space.
Revised: Some recent methods learn state representations without requiring a decoder back into pixel space.

Original: Examples include the selectivity objective in Thomas et al.
Revised: Examples include the selectivity objective in Thomas et al.

Original: (2018), the contrastive objective in Franc¸ois-Lavet et al.
Revised: (2018), the contrastive objective in Franc¸ois-Lavet et al.

Revised: (2018), the mutual information objective in

Revised: Anand et al.

Original: (2018), the distribution matching objective in Gelada et al.
Revised: (2019), the distribution matching objective in Gelada et al.

Original: or using causality-based losses and physical priors in latent space (Jonschkowski & Brock, 2015;
Revised: or using causalitybased losses and physical priors in latent space (Jonschkowski & Brock, 2015;

Original: Ehrhardt et al., 2018).
Revised: Ehrhardt et al., 2018).

Original: Most notably, Ehrhardt et al.
Revised: Most notably, Ehrhardt et al.

Original: (2018) propose a method to learn an object detection module and a physics module jointly from raw video data without pixel-based losses.
Revised: (2018) propose a method to learn an object detection module and a physics module jointly from raw video data without pixel-based losses.

Original: This approach, however, can only track a single object at a time and requires careful balancing of multiple loss functions.
Revised: This approach, however, can only track a single object at a time and requires careful balancing of multiple loss functions.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our goal of this experimental section is to verify whether C-SWMs can 1) learn to discover object representations from environment interactions without supervision, 2) learn an accurate transition model in latent space, and 3) generalize to novel, unseen scenes.
Revised: Our goal of this experimental section is to verify whether C-SWMs can 1) learn to discover object representations from environment interactions without supervision, 2) learn an accurate transition model in latent space, and 3) generalize to novel, unseen scenes.

Revised: Our implementation is available under https://github.com/tkipf/c-swm .

Original: 4. E NVIRONMENTS
Revised: 4. E NVIRONMENTS

Original: We evaluate C-SWMs on two novel grid world environments (2D shapes and 3D blocks) involving multiple interacting objects that can be manipulated independently by an agent, two Atarigames (Atari Pong and Space Invaders), and a multi-object physics simulation (3-body physics).
Revised: We evaluate C-SWMs on two novel grid world environments (2D shapes and 3D blocks) involving multiple interacting objects that can be manipulated independently by an agent, two Atarigames (Atari Pong and Space Invaders), and a multi-object physics simulation (3-body physics).

Original: See Figure 2 for example observations.
Revised: See Figure 2 for example observations.

Original: For all environments, we use a random policy to collect experience for both training and evaluation.
Revised: For all environments, we use a random policy to collect experience for both training and evaluation.

Original: Observations are provided as 50 × 50 × 3 color images for the grid world environments and as 50 × 50 × 6 tensors (two concatenated consecutive frames) for the Atari and 3-body physics environments.
Revised: Observations are provided as 50 × 50 × 3 color images for the grid world environments and as 50 × 50 × 6 tensors (two concatenated consecutive frames) for the Atari and 3-body physics environments.

Original: Additional details on environments and dataset creation can be found in Appendix B.
Revised: Additional details on environments and dataset creation can be found in Appendix B.

Original: [Figure]
Revised: [Figure]

Original: 4. E VALUATION M ETRICS
Revised: 4. E VALUATION M ETRICS

Original: In order to evaluate model performance directly in latent space, we make use of ranking metrics, which are commonly used for the evaluation of link prediction models, as in, e.g., Bordes et al.
Revised: In order to evaluate model performance directly in latent space, we make use of ranking metrics, which are commonly used for the evaluation of link prediction models, as in, e.g., Bordes et al.

Original: This allows us to assess the quality of learned representations directly without relying on auxiliary metrics such as pixel-based reconstruction losses, or performance in downstream tasks such as planning.
Revised: This allows us to assess the quality of learned representations directly without relying on auxiliary metrics such as pixel-based reconstruction losses, or performance in downstream tasks such as planning.

Original: Given an observation encoded by the model and an action, we use the model to predict the representation of the next state, reached after taking the action in the environment.
Revised: Given an observation encoded by the model and an action, we use the model to predict the representation of the next state, reached after taking the action in the environment.

Original: This predicted state representation is then compared to the encoded true observation after taking the action in the environment and a set of reference states (observations encoded by the model) obtained from the experience buffer.
Revised: This predicted state representation is then compared to the encoded true observation after taking the action in the environment and a set of reference states (observations encoded by the model) obtained from the experience buffer.

Original: We measure and report both Hits at Rank 1 (H@1) and Mean Reciprocal Rank (MRR).
Revised: We measure and report both Hits at Rank 1 (H@1) and Mean Reciprocal Rank (MRR).

Original: Additional details on these evaluation metrics can be found in Appendix C.
Revised: Additional details on these evaluation metrics can be found in Appendix C.

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

Original: Autoencoder-based World Models The predominant method for state representation learning is based on autoencoders, and often on the VAE (Kingma & Welling, 2013; Rezende et al., 2014) model in particular.
Revised: Autoencoder-based World Models The predominant method for state representation learning is based on autoencoders, and often on the VAE (Kingma & Welling, 2013; Rezende et al., 2014) model in particular.

Original: This World Model baseline is inspired by Ha & Schmidhuber (2018) and uses either a deterministic autoencoder (AE) or a VAE to learn state representations.
Revised: This World Model baseline is inspired by Ha & Schmidhuber (2018) and uses either a deterministic autoencoder (AE) or a VAE to learn state representations.

Original: Finally, an MLP is used to predict the next state after taking an action.
Revised: Finally, an MLP is used to predict the next state after taking an action.

Original: Physics As Inverse Graphics (PAIG)
Revised: Physics As Inverse Graphics (PAIG)

Original: This model by Jaques et al.
Revised: This model by Jaques et al.

Original: (2019) is based on an encoderdecoder architecture and trained with pixel-based reconstruction losses, but uses a differentiable physics engine in the latent space that operates on explicit position and velocity representations for each object.
Revised: (2019) is based on an encoderdecoder architecture and trained with pixel-based reconstruction losses, but uses a differentiable physics engine in the latent space that operates on explicit position and velocity representations for each object.

Original: Thus, this model is only applicable to the 3-body physics environment.
Revised: Thus, this model is only applicable to the 3-body physics environment.

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

Original: We train C-SWMs on an experience buffer obtained by running a random policy on the respective environment.
Revised: We train C-SWMs on an experience buffer obtained by running a random policy on the respective environment.

Original: We choose 1000 episodes with 100 environment steps each for the grid world environments, 1000 episodes with 10 steps each for the Atari environments and 5000 episodes withsteps each for the 3-body physics environment.
Revised: We choose 1000 episodes with 100 environment steps each for the grid world environments, 1000 episodes with 10 steps each for the Atari environments and 5000 episodes withsteps each for the 3-body physics environment.

Original: For evaluation, we populate a separate experience buffer with 10 environment steps per episode and a total of 10.000 episodes for the grid world environments, 100 episodes for the Atari environments and 1000 episodes for the physics environment.
Revised: For evaluation, we populate a separate experience buffer with 10 environment steps per episode and a total of 10.000 episodes for the grid world environments, 100 episodes for the Atari environments and 1000 episodes for the physics environment.

Original: For the Atari environments, we minimize train/test overlap by ‘warm-starting’ experience collection in these environments with random actions before we start populating the experience buffer (see Appendix B), and we ensure that not a single full test set episode coincides exactly with an episode from the training set.
Revised: For the Atari environments, we minimize train/test overlap by ‘warm-starting’ experience collection in these environments with random actions before we start populating the experience buffer (see Appendix B), and we ensure that not a single full test set episode coincides exactly with an episode from the training set.

Original: The state spaces of the grid world environments are large (approx.
Revised: The state spaces of the grid world environments are large (approx.

Original: 6.4M unique states) and hence train/test coincidence of a full 10-step episode is unlikely.
Revised: 6.4M unique states) and hence train/test coincidence of a full 10-step episode is unlikely.

Original: Overlap is similarly unlikely for the physics environment which has a continuous state space.
Revised: Overlap is similarly unlikely for the physics environment which has a continuous state space.

Original: Hence, performing well on these tasks will require some form of generalization to new environment conﬁgurations or an unseen sequence of states and actions.
Revised: Hence, performing well on these tasks will require some form of generalization to new environment conﬁgurations or an unseen sequence of states and actions.

Original: All models are trained for 100 epochs (200 for Atari games) using the Adam (Kingma & Ba, 2014) optimizer with a learning rate of 5 · 10 − 4 and a batch size of 1024 (512 for baselines with decoders due to higher memory demands, and 100 for PAIG as suggested by the authors).
Revised: All models are trained for 100 epochs (200 for Atari games) using the Adam (Kingma & Ba, 2014) optimizer with a learning rate of 5 · 10 − 4 and a batch size of 1024 (512 for baselines with decoders due to higher memory demands, and 100 for PAIG as suggested by the authors).

Original: Model architecture details are provided in Appendix D.
Revised: Model architecture details are provided in Appendix D.

Original: 4. Q UALITATIVE R ESULTS
Revised: 4. Q UALITATIVE R ESULTS

Original: We present qualitative results for the grid world environments in Figure 3 and for the 3-body physics environment in Figure 4.
Revised: We present qualitative results for the grid world environments in Figure 3 and for the 3-body physics environment in Figure 4.

Original: All results are obtained on hold-out test data.
Revised: All results are obtained on hold-out test data.

Original: Further qualitative results (incl.
Revised: Further qualitative results (incl.

Original: on Atari games) can be found in Appendix A.
Revised: on Atari games) can be found in Appendix A.

Original: In the grid world environments, we can observe that C-SWM reliably discovers object-speciﬁc ﬁlters for a particular scene, without direct supervision.
Revised: In the grid world environments, we can observe that C-SWM reliably discovers object-speciﬁc ﬁlters for a particular scene, without direct supervision.

Original: Further, each object is represented by two coordinates which correspond (up to a random linear transformation) to the true object position in
Revised: Further, each object is represented by two coordinates which correspond (up to a random linear transformation) to the true object position in

Original: [Figure]
Revised: [Figure]

Original: [Figure][Figure] [Figure]
Revised: [Figure]

Original: [Figure][Figure] [Figure] [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: the scene.
Revised: the scene.

Original: Although we choose a two-dimensional latent representation per object for easier visualization, we ﬁnd that results remain unchanged if we increase the dimensionality of the latent representation.
Revised: Although we choose a two-dimensional latent representation per object for easier visualization, we ﬁnd that results remain unchanged if we increase the dimensionality of the latent representation.

Original: The edges in this learned abstract transition graph correspond to the effect of a particular action applied to the object.
Revised: The edges in this learned abstract transition graph correspond to the effect of a particular action applied to the object.

Original: The structure of the learned latent representation accurately captures the underlying grid structure of the environment.
Revised: The structure of the learned latent representation accurately captures the underlying grid structure of the environment.

Original: We further ﬁnd that the transition model, which only has access to latent representations, correctly captures whether an action has an effect or not, e.g., if a neighboring position is blocked by another object.
Revised: We further ﬁnd that the transition model, which only has access to latent representations, correctly captures whether an action has an effect or not, e.g., if a neighboring position is blocked by another object.

Original: Similarly, we ﬁnd that the model can learn object-speciﬁc encoders in the 3-body physics environment and can learn object-speciﬁc latent representations that track location and velocity of a particular object, while learning an accurate latent transition model that generalizes well to unseen environment instances.
Revised: Similarly, we ﬁnd that the model can learn object-speciﬁc encoders in the 3-body physics environment and can learn object-speciﬁc latent representations that track location and velocity of a particular object, while learning an accurate latent transition model that generalizes well to unseen environment instances.

Original: 4. Q UANTITATIVE R ESULTS
Revised: 4. Q UANTITATIVE R ESULTS

Original: We set up quantitative experiments for evaluating the quality of both object discovery and the quality of the learned transition model.
Revised: We set up quantitative experiments for evaluating the quality of both object discovery and the quality of the learned transition model.

Original: We compare against autoencoder baselines and model variants that do not represent the environment in an object-factorized manner, do not use a GNN, or do not make use of contrastive learning.
Revised: We compare against autoencoder baselines and model variants that do not represent the environment in an object-factorized manner, do not use a GNN, or do not make use of contrastive learning.

Original: Performing well under this evaluation setting requires some degree of (combinatorial) generalization to unseen environment instances.
Revised: Performing well under this evaluation setting requires some degree of (combinatorial) generalization to unseen environment instances.

Original: We report ranking scores (in %) in latent space, after encoding source and target observations, and taking steps in the latent space using the learned model.
Revised: We report ranking scores (in %) in latent space, after encoding source and target observations, and taking steps in the latent space using the learned model.

Original: Reported results are mean and standard error of scores over 4 runs on hold-out environment instances.
Revised: Reported results are mean and standard error of scores over 4 runs on hold-out environment instances.

Original: Results are summarized in Table 1.[Table] R [Table]
Revised: Results are summarized in Table 1.

Original: P

Original: ONG

Original: [Table]
Revised: [Table]

Original: We ﬁnd that baselines that make use of reconstruction losses in pixel space (incl.
Revised: We ﬁnd that baselines that make use of reconstruction losses in pixel space (incl.

Original: the C-SWM model variant without contrastive loss) typically generalize less well to unseen scenes and learn a latent space conﬁguration that makes it difﬁcult for the transition model to learn the correct transition function.
Revised: the C-SWM model variant without contrastive loss) typically generalize less well to unseen scenes and learn a latent space conﬁguration that makes it difﬁcult for the transition model to learn the correct transition function.

Original: See Appendix A.1 for a visual analysis of such latent state transition graphs.
Revised: See Appendix A.1 for a visual analysis of such latent state transition graphs.

Original: This effect appears to be even stronger when using a VAE-based World Model, where the prior puts further constraints on the latent representations.
Revised: This effect appears to be even stronger when using a VAE-based World Model, where the prior puts further constraints on the latent representations.

Original: C-SWM recovers this structure well, see Figure 3.
Revised: C-SWM recovers this structure well, see Figure 3.

Original: On the grid-world environments (2D shapes and 3D blocks), C-SWM models latent transitions almost perfectly, which requires taking interactions between latent representations of objects into account.
Revised: On the grid-world environments (2D shapes and 3D blocks), C-SWM models latent transitions almost perfectly, which requires taking interactions between latent representations of objects into account.

Original: Removing the interaction component, i.e., replacing the latent GNN with an object-wise
Revised: Removing the interaction component, i.e., replacing the latent GNN with an object-wise

Original: MLP, makes the model insensitive to pairwise interactions and hence the ability to predict future states deteriorates.
Revised: MLP, makes the model insensitive to pairwise interactions and hence the ability to predict future states deteriorates.

Original: Similarly, if we remove the state factorization, the model has difﬁculties generalizing to unseen environment conﬁgurations.
Revised: Similarly, if we remove the state factorization, the model has difﬁculties generalizing to unseen environment conﬁgurations.

Original: We further explore variants of the grid-world environments in Appendix A.5, where 1) certain actions have no effect, and 2) one object moves randomly and independent of agent actions, to test the robustness of our approach.
Revised: We further explore variants of the grid-world environments in Appendix A.5, where 1) certain actions have no effect, and 2) one object moves randomly and independent of agent actions, to test the robustness of our approach.

Original: For the Atari 2600 experiments, we ﬁnd that results can have a high variance, and that the task is more difﬁcult, as both the World Model baseline and C-SWM struggle to make perfect longterm predictions.
Revised: For the Atari 2600 experiments, we ﬁnd that results can have a high variance, and that the task is more difﬁcult, as both the World Model baseline and C-SWM struggle to make perfect longterm predictions.

Original: While for Space Invaders, a large number of object slots ( K = 5 ) appears to be beneﬁcial, C-SWM achieves best results with only a single object slot in Atari Pong.
Revised: While for Space Invaders, a large number of object slots ( K = 5 ) appears to be beneﬁcial, C-SWM achieves best results with only a single object slot in Atari Pong.

Original: This suggests that one should determine the optimal value of K based on performance on a validation set if it is not known a-priori.
Revised: This suggests that one should determine the optimal value of K based on performance on a validation set if it is not known a-priori.

Original: Using an iterative object encoding mechanism, such as in MONet (Burgess et al., 2019), would enable the model to assign ‘empty’ slots which could improve robustness w.r.t.
Revised: Using an iterative object encoding mechanism, such as in MONet (Burgess et al., 2019), would enable the model to assign ‘empty’ slots which could improve robustness w.r.t.

Original: the choice of K , which we leave for future work.
Revised: the choice of K , which we leave for future work.

Original: We ﬁnd that both C-SWMs and the autoencoder-based World Model baseline excel at short-term predictions in the 3-body physics environment, with C-SWM having a slight edge in the 10 step prediction setting.
Revised: We ﬁnd that both C-SWMs and the autoencoder-based World Model baseline excel at short-term predictions in the 3-body physics environment, with C-SWM having a slight edge in the 10 step prediction setting.

Original: Under our evaluation setting, the PAIG baseline (Jaques et al., 2019) underperforms using the hyperparameter setting recommended by the authors.
Revised: Under our evaluation setting, the PAIG baseline (Jaques et al., 2019) underperforms using the hyperparameter setting recommended by the authors.

Original: Note that we do not tune hyperparameters of C-SWM separately for this task and use the same settings as in other environments.
Revised: Note that we do not tune hyperparameters of C-SWM separately for this task and use the same settings as in other environments.

Original: 4. L IMITATIONS
Revised: 4. L IMITATIONS

Original: Instance Disambiguation In our experiments, we chose a simple feed-forward CNN architecture for the object extractor module.
Revised: Instance Disambiguation In our experiments, we chose a simple feed-forward CNN architecture for the object extractor module.

Original: This type of architecture cannot disambiguate multiple instances of the same object present in one scene and relies on distinct visual features or labels (e.g., the green square ) for object extraction.
Revised: This type of architecture cannot disambiguate multiple instances of the same object present in one scene and relies on distinct visual features or labels (e.g., the green square ) for object extraction.

Original: To better handle scenes which contain potentially multiple copies of the same object (e.g., in the Atari Space Invaders game), one would require some form of iterative disambiguation procedure to break symmetries and dynamically bind individual objects to slots or object ﬁles (Kahneman & Treisman, 1984; Kahneman et al., 1992), such as in the style of dynamic routing (Sabour et al., 2017), iterative inference (Greff et al., 2019; Engelcke et al., 2019) or sequential masking (Burgess et al., 2019; Kipf et al., 2019).
Revised: To better handle scenes which contain potentially multiple copies of the same object (e.g., in the Atari Space Invaders game), one would require some form of iterative disambiguation procedure to break symmetries and dynamically bind individual objects to slots or object ﬁles (Kahneman & Treisman, 1984; Kahneman et al., 1992), such as in the style of dynamic routing (Sabour et al., 2017), iterative inference (Greff et al., 2019; Engelcke et al., 2019) or sequential masking (Burgess et al., 2019; Kipf et al., 2019).

Original: Stochasticity & Markov Assumption Our formulation of C-SWMs does not take into account stochasticity in environment transitions or observations, and hence is limited to fully deterministic worlds.
Revised: Stochasticity & Markov Assumption Our formulation of C-SWMs does not take into account stochasticity in environment transitions or observations, and hence is limited to fully deterministic worlds.

Original: A probabilistic extension of C-SWMs is an interesting avenue for future work.
Revised: A probabilistic extension of C-SWMs is an interesting avenue for future work.

Original: For simplicity, we make the Markov assumption: state and action contain all the information necessary to predict the next state.
Revised: For simplicity, we make the Markov assumption: state and action contain all the information necessary to predict the next state.

Original: This allows us to look at single state-action-state triples in isolation.
Revised: This allows us to look at single state-action-state triples in isolation.

Original: To go beyond this limitation, one would require some form of memory mechanism, such as an RNN as part of the model architecture, which we leave for future work.
Revised: To go beyond this limitation, one would require some form of memory mechanism, such as an RNN as part of the model architecture, which we leave for future work.

Original:  C ONCLUSIONS
Revised:  C ONCLUSIONS

Original: Structured world models offer compelling advantages over pure connectionist methods, by enabling stronger inductive biases for generalization, without necessarily constraining the generality of the model: for example, the contrastively trained model on the 3-body physics environment is free to store identical representations in each object slot and ignore pairwise interactions, i.e., an unstructured world model still exists as a special case.
Revised: Structured world models offer compelling advantages over pure connectionist methods, by enabling stronger inductive biases for generalization, without necessarily constraining the generality of the model: for example, the contrastively trained model on the 3-body physics environment is free to store identical representations in each object slot and ignore pairwise interactions, i e., an unstructured world model still exists as a special case.

Original: Experimentally, we ﬁnd that C-SWMs make effective use of this additional structure, likely because it allows for a transition model of signiﬁcantly lower complexity, and learn object-oriented models that generalize better to unseen situations.
Revised: Experimentally, we ﬁnd that C-SWMs make effective use of this additional structure, likely because it allows for a transition model of signiﬁcantly lower complexity, and learn object-oriented models that generalize better to unseen situations.

Original: We are excited about the prospect of using C-SWMs for model-based planning and reinforcement learning in future work, where object-oriented representations will likely allow for more accurate counterfactual reasoning about effects of actions and novel interactions in the environment.
Revised: We are excited about the prospect of using C-SWMs for model-based planning and reinforcement learning in future work, where object-oriented representations will likely allow for more accurate counterfactual reasoning about effects of actions and novel interactions in the environment.

Original: We further hope to inspire future work to think beyond autoencoder-based approaches for object-based, structured representation learning, and to address some of the limitations outlined in this paper
Revised: We further hope to inspire future work to think beyond autoencoder-based approaches for object-based, structured representation learning and to address some of the limitations outlined in this paper

Revised: A CKNOWLEDGEMENTS

Revised: We would like to thank Marco Federici and Adam Kosiorek for helpful discussions.

Revised: We would further like to thank the anonymous reviewers for valuable feedback.

Revised: T

Revised: K. acknowledges funding by

Original:  A A DDITIONAL R ESULTS AND D ISCUSSION A. O BJECT - SPECIFIC R EPRESENTATIONS
Revised: SAP SE. A A DDITIONAL R ESULTS AND D ISCUSSION A. O BJECT - SPECIFIC R EPRESENTATIONS

Original: We visualize abstract state transition graphs separated by object slot for the 3D cubes environment in Figure 5.
Revised: We visualize abstract state transition graphs separated by object slot for the 3D cubes environment in Figure 5.

Original: Discovered object representations in the 2D shapes dataset (not shown) are qualitatively very similar.
Revised: Discovered object representations in the 2D shapes dataset (not shown) are qualitatively very similar.

Original: We apply the same visualization technique to the model variant without contrastive loss, which is instead trained with a decoder model and a loss in pixel space.
Revised: We apply the same visualization technique to the model variant without contrastive loss, which is instead trained with a decoder model and a loss in pixel space.

Original: See Figure 6 for this baseline and note that the regular structure in the latent space is lost, which makes it difﬁcult for the transition model to learn transitions which generalize to unseen environment instances.
Revised: See Figure 6 for this baseline and note that the regular structure in the latent space is lost, which makes it difﬁcult for the transition model to learn transitions which generalize to unseen environment instances.

Revised: Qualitative results for the 3-body physics dataset are summarized in Figures 7 and 8 for two different random seeds.

Original: [Figure]
Revised: [Figure]

Original: nearly identical for each object, which illustrates that the model successfully represents objects in the same manner despite their visual differences.
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: Qualitative results for the 3-body physics dataset are summarized in Figures 7 and 8 for two different random seeds.

Original: For the Atari 2600 environments, we generally found latent object representations to be less interpretable.
Revised: For the Atari 2600 environments, we generally found latent object representations to be less interpretable.

Original: We attribute this to the fact that a) objects have different roles and are in general not exchangeable (in contrast to the block pushing grid world environments and the 3-body physics environment), b) actions affect only one object directly, but many other objects indirectly in two consecutive frames, and c) due to multiple objects in one scene sharing the same visual features.[Figure] [Figure]
Revised: We attribute this to the fact that a) objects have different roles and are in general not exchangeable (in contrast to the block pushing grid world environments and the 3-body physics environment), b) actions affect only one object directly, but many other objects indirectly in two consecutive frames, and c) due to multiple objects in one scene sharing the same visual features.

Original: See Figure 9 for an example of learned representations in Atari Pong and Figure 10 for an example in Space Invaders.
Revised: See Figure 9 for an example of learned representations in Atari Pong and Figure 10 for an example in Space Invaders.

Original: [Figure]
Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Revised: multiple prediction steps into the future.

Original: A. M ODEL C OMPARISON IN P IXEL S PACE
Revised: A. M ODEL C OMPARISON IN P IXEL S PACE

Original: To supplement our model comparison in latent space using ranking metrics, we here show a direct comparison in pixel space at the example of the 2D shapes environment.
Revised: To supplement our model comparison in latent space using ranking metrics, we here show a direct comparison in pixel space at the example of the 2D shapes environment.

Original: This requires training a separate decoder model for C-SWM.
Revised: This requires training a separate decoder model for C-SWM.

Original: For fairness of this comparison, we use the same protocol to train a separate decoder for the World Model (AE) baseline (discarding the one obtained from the original end-to-end auto-encoding training procedure).
Revised: For fairness of this comparison, we use the same protocol to train a separate decoder for the World Model (AE) baseline (discarding the one obtained from the original end-to-end auto-encoding training procedure).

Original: This decoder has the same architecture as in the other baseline models and is trained for 100 epochs.
Revised: This decoder has the same architecture as in the other baseline models and is trained for 100 epochs.

Original: [Figure]
Revised: [Figure]

Original: [Figure]

Original: For a qualitative comparison, see Figure 11.
Revised: For a qualitative comparison, see Figure 11.

Original: The
Revised: The

Original: C-SWM model, as expected from our ranking analysis in latent space, performs almost perfectly at this task.
Revised: C-SWM model, as expected from our ranking analysis in latent space, performs almost perfectly at this task.

Original: Although the World Model (AE) baseline makes clear mistakes which compound over time, it nonetheless often gets several object positions correct after many time steps.
Revised: Although the World Model (AE) baseline makes clear mistakes which compound over time, it nonetheless often gets several object positions correct after many time steps.

Original: The ranking loss in latent space captures this behaviour well, and, for example, assigns an almost perfect score for 1-step prediction to the World Model (AE) baseline.
Revised: The ranking loss in latent space captures this behaviour well, and, for example, assigns an almost perfect score for 1-step prediction to the World Model (AE) baseline.

Original: The typically used mean-squared error (MSE) in pixel space (see Figure 12), however, differs by several orders of magnitude between the two models and does not capture any of the nuanced differences in qualitative predictive behavior.
Revised: The typically used mean-squared error (MSE) in pixel space (see Figure 12), however, differs by several orders of magnitude between the two models and does not capture any of the nuanced differences in qualitative predictive behavior.

Original: We hence strongly encourage researchers in this area to consider ranking-based evaluations directly in latent space in addition to, or as a replacement for comparisons in pixel space.
Revised: We hence strongly encourage researchers in this area to consider ranking-based evaluations directly in latent space in addition to, or as a replacement for comparisons in pixel space.

Original: A. H INGE L OSS
Revised: A. H INGE L OSS

Original: In Table 2, we summarize results of a comparison between our multi-object contrastive loss from
Revised: In Table 2, we summarize results of a comparison between our multi-object contrastive loss from

Original: Eq. 6, denoted by C-SWM (original) , and a full-hinge (or triplet ) loss that places the hinge over both the positive and the negative energy term,
Revised: Eq. 6, denoted by C-SWM (original) , and a full-hinge (or triplet ) loss that places the hinge over both the positive and the negative energy term,

Original: i
Revised: i

Original: [Table]
Revised: [Table]

Original: e., L = max(0 , γ + H − ˜ H ) .
Revised: e., L = max(0 , γ + H − ˜ H ) .

Original: The latter loss is similar to the loss employed in Bordes et al.
Revised: The latter loss is similar to the loss employed in Bordes et al.

Original: (2013), but without imposing a norm constraint on the embeddings z t .
Revised: (2013), but without imposing a norm constraint on the embeddings z t .

Original: Reported results aremean andstandard error over 4 runs on hold-out environment instances.
Revised: Reported results are and on hold- environment instances.

Original: [Table]
Revised: [Table]

Original: full hinge ( γ = 1 )
Revised: full hinge ( γ = 1 )

Original: [Table]
Revised: [Table]

Original: with full hinge ( γ = 5 )
Revised: full hinge ( γ = 5 )

Original: [Table]
Revised: [Table]

Original: with full hinge ( γ = 10 )
Revised: with full hinge ( γ = 10 )

Original: [Table]
Revised: [Table]

Original: We ﬁnd that placing the hinge only on the negative energy term signiﬁcantly improves results, likely because the full hinge loss can be minimized by simply increasing the norms of the embeddings z t , as outlined in Bordes et al.
Revised: We ﬁnd that placing the hinge only on the negative energy term signiﬁcantly improves results, likely because the full hinge loss can be minimized by simply increasing the norms of the embeddings z t , as outlined in Bordes et al.

Original: This could be addressed by introducing additional constraints, such as by forcing the embeddings z t to lie on the unit-hypersphere, as done in TransE (Bordes et al., 2013), which we do not consider in this work as most of the environment dynamics in our experiments obey a ﬂat Euclidean geometry (e.g., translation on a grid).
Revised: This could be addressed by introducing additional constraints, such as by forcing the embeddings z t to lie on the unit-hypersphere, as done in TransE (Bordes et al., 2013), which we do not consider in this work as most of the environment dynamics in our experiments obey a ﬂat Euclidean geometry (e.g., translation on a grid).

Original: A. M ULTIPLE F EATURE M APS PER O BJECT S LOT
Revised: A. M ULTIPLE F EATURE M APS PER O BJECT S LOT

Original: We experimented with assigning multiple feature maps per object slot to investigate whether this can improve performance on the Space Invaders environment (as an example of one of the more complex environments considered in this work).
Revised: We experimented with assigning multiple feature maps per object slot to investigate whether this can improve performance on the Space Invaders environment (as an example of one of the more complex environments considered in this work).

Original: All other model settings are left unchanged.
Revised: All other model settings are left unchanged.

Original: Results are summarized in Table 3.
Revised: Results are summarized in Table 3.

Original: The original model variant with one feature map per object slot (using a total of K = 5 slots) is denoted by C-SWM (original) .
Revised: The original model variant with one feature map per object slot (using a total of K = 5 slots) is denoted by C-SWM (original) .

Original: Reported results are mean and standard error over 4 runs on hold-out environment instances.
Revised: Reported results are mean and standard error over 4 runs on hold-out environment instances.

Original: [Table]
Revised: [Table]

Original: with 2 feature maps
Revised: with 2 feature maps

Original: [Table]
Revised: [Table]

Original: with 5 feature 
Revised: with 5 feature maps

Original: [Table]
Revised: [Table]

Revised: with

Revised: [Table]

Original: 10 feature
Revised: feature

Original: [Table]
Revised: [Table]

Original: We ﬁnd that there is no clear advantage in using multiple feature maps per object slot for the Space
Revised: We ﬁnd that there is no clear advantage in using multiple feature maps per object slot for the Space

Original: Invaders environment, but this might nonetheless prove useful for environments of more complex nature.
Revised: Invaders environment [Table] but this might nonetheless prove useful for environments of more complex nature.

Original: A. G RID -W ORLD E NVIRONMENT V ARIANTS
Revised: A. G RID -W ORLD E NVIRONMENT V ARIANTS

Original: To further evaluate the robustness of our approach
Revised: To further evaluate the robustness of our approach

Original: [Table]
Revised: [Table]

Original: we experimented with two variants of the 2D shapes grid-world environment with ﬁve objects.
Revised: we experimented with two variants of the 2D shapes grid-world environment with ﬁve objects.

Original: In the ﬁrst variant, termed no-op action , we include an additional ‘no-op’ action for each object that has no effect.
Revised: In the ﬁrst variant, termed no-op action , we include an additional ‘no-op’ action for each object that has no effect.

Original: This mirrors the setting present in many Atari 2600 benchmark environments that have an explicit no-op action.
Revised: This mirrors the setting present in many Atari 2600 benchmark environments that have an explicit no-op action.

Original: In the second variant, termed random object , we assign a special property to one out of the ﬁve objects: a random action is applied to this object at every turn, i [Table] e., the agent only has control over the other four objects and in every environment step, a total of two objects receive an action (one by the agent, which only affects one out of the ﬁrst four objects, and one random action on the last object).
Revised: In the second variant, termed random object , we assign a special property to one out of the ﬁve objects: a random action is applied to this object at every turn, i.e., the agent only has control over the other four objects and in every environment step, a total of two objects receive an action (one by the agent, which only affects one out of the ﬁrst four objects, and one random action on the last object).

Original: This is to test how the model behaves in the presence of other agents that act randomly.
Revised: This is to test how the model behaves in the presence of other agents that act randomly.

Original: As our model is fully deterministic, it is to be expected that adding a source of randomness in the environment will pose a challenge.
Revised: As our model is fully deterministic, it is to be expected that adding a source of randomness in the environment will pose a challenge.

Original: We report results in Table 4.
Revised: We report results in Table 4.

Original: Reported numbers are mean and standard error over 4 runs on hold-out environment instances.
Revised: Reported numbers are mean and standard error over 4 runs on hold-out environment instances.

Original: We ﬁnd that adding a no-op action has little effect on our model, but adding a randomly moving object reduces predictive performance.
Revised: We ﬁnd that adding a no-op action has little effect on our model, but adding a randomly moving object reduces predictive performance.

Original: Performance in the latter case could potentially be improved by explicitly modeling stochasticity in the environment.
Revised: Performance in the latter case could potentially be improved by explicitly modeling stochasticity in the environment.

Original: [Table]
Revised: [Table]

Original: 99.9 ± 0. + random object98.4 ± 0. 99.2 ± 0.
Revised: + random object

Original: 95.9 ± 0. 97.6 ± 0. 90.4 ± 0. 93.7 ± 0.
Revised: [Table]

Original: A. S TABILITY
Revised: A. S TABILITY

Original: The discovered object representations and identiﬁcations can vary between runs.
Revised: The discovered object representations and identiﬁcations can vary between runs.

Original: While we found that this process is very stable for the simple grid world environments where actions only affect a single object, we found that results can be initialization-dependent on the Atari environments, where actions can have effects across all objects, and the 3-body physics simulation (see Figures 7 and 8), which does not have any actions.
Revised: While we found that this process is very stable for the simple grid world environments where actions only affect a single object, we found that results can be initialization-dependent on the Atari environments, where actions can have effects across all objects, and the 3-body physics simulation (see Figures 7 and 8), which does not have any actions.

Original: In some cases, the discovered representation can be less suitable for forward prediction in time or for generalization to novel scenes, which explains the variance in some of our results on these datasets.
Revised: In some cases, the discovered representation can be less suitable for forward prediction in time or for generalization to novel scenes, which explains the variance in some of our results on these datasets.

Original: A. T RAINING T IME
Revised: A. T RAINING T IME

Original: We found that the overall training time of the C-SWM model was comparable to that of the World
Revised: We found that the overall training time of the C-SWM model was comparable to that of the World

Original: Model baseline.
Revised: Model baseline.

Original: Both C-SWM and the World Model baseline trained for approx.
Revised: Both C-SWM and the World Model baseline trained for approx.

Original: 1 hour wall-clock time on the 2D shapes dataset, approx.
Revised: 1 hour wall-clock time on the 2D shapes dataset, approx.

Original: 2 hours on the 3D cubes dataset, and approx.
Revised: 2 hours on the 3D cubes dataset, and approx.

Original: 30min on the 3-body physics environment using a single NVIDIA GTX1080Ti GPU.
Revised: 30min on the 3-body physics environment using a single NVIDIA GTX1080Ti GPU.

Original: The Atari Pong and Space
Revised: The Atari Pong and Space

Original: Invaders models trained for typically less than 20 minutes.
Revised: Invaders models trained for typically less than 20 minutes.

Original: A notable exception is the PAIG baseline model (Jaques et al., 2019) which trained for approx.
Revised: A notable exception is the PAIG baseline model (Jaques et al., 2019) which trained for approx.

Original: 6 hours on a NVIDIA TitanX Pascal GPU using the recommended settings by the authors of the paper.
Revised: 6 hours on a NVIDIA TitanX Pascal GPU using the recommended settings by the authors of the paper.

Original: B D ATASETS B. G RID W ORLDS
Revised: B D ATASETS B. G RID W ORLDS

Original: To generate an experience buffer for training, we initialize the environment with random object placements and uniformly sample an object and an object-speciﬁc action at every time step.
Revised: To generate an experience buffer for training, we initialize the environment with random object placements and uniformly sample an object and an object-speciﬁc action at every time step.

Original: We provide state observations as 50 × 50 × 3 tensors with RGB color channels, normalized to [0 , 1] .
Revised: We provide state observations as 50 × 50 × 3 tensors with RGB color channels, normalized to [0 , 1] .

Original: Actions are provided as a 4-dim one-hot vector (if an action is applied) or a vector of zeros per object slot in the environment.
Revised: Actions are provided as a 4-dim one-hot vector (if an action is applied) or a vector of zeros per object slot in the environment.

Original: The action one-hot vector encodes the directional movement action applied to a particular object, or is represented as a vector of zeros if no action is applied to a particular object.
Revised: The action one-hot vector encodes the directional movement action applied to a particular object, or is represented as a vector of zeros if no action is applied to a particular object.

Original: Note that only a single object receives an action per time step.
Revised: Note that only a single object receives an action per time step.

Original: For the Atari environments, we provide a copy of the one-hot encoded action vector to every object slot, and for the 3-body physics environment, which has no actions, we do not provide an action vector.
Revised: For the Atari environments, we provide a copy of the one-hot encoded action vector to every object slot, and for the 3-body physics environment, which has no actions, we do not provide an action vector.

Original: 2D Shapes
Revised: 2D Shapes

Original: This environment is a 5 × 5 grid world with 5 different objects placed at random positions.
Revised: This environment is a 5 × 5 grid world with 5 different objects placed at random positions.

Original: Each location can only be occupied by at maximum one object.
Revised: Each location can only be occupied by at maximum one object.

Original: Each object is represented by a unique shape/color combination, occupying 10 × 10 pixels on the overall 50 × 50 pixel grid.
Revised: Each object is represented by a unique shape/color combination, occupying 10 × 10 pixels on the overall 50 × 50 pixel grid.

Original: At each time step, one object can be selected and moved by one position along the four cardinal directions.
Revised: At each time step, one object can be selected and moved by one position along the four cardinal directions.

Original: See Figure 2a for an example.
Revised: See Figure 2a for an example.

Original: The action has no effect if the target location in a particular direction is occupied by another object or outside of the 5 × 5 grid.
Revised: The action has no effect if the target location in a particular direction is occupied by another object or outside of the 5 × 5 grid.

Original: Thus, a learned transition model needs to take pairwise interactions between object properties (i.e., their locations) into account, as it would otherwise be unable to predict the effect of an action correctly.
Revised: Thus, a learned transition model needs to take pairwise interactions between object properties (i.e., their locations) into account, as it would otherwise be unable to predict the effect of an action correctly.

Original: 3D Blocks To investigate to what degree our model is robust to partial occlusion and perspective changes, we implement a simple block pushing environment using Matplotlib (Hunter, 2007) as a rendering tool.
Revised: 3D Blocks To investigate to what degree our model is robust to partial occlusion and perspective changes, we implement a simple block pushing environment using Matplotlib (Hunter, 2007) as a rendering tool.

Original: The underlying environment dynamics are the same as in the 2D Shapes dataset, and we only change the rendering component to make for a visually more challenging task that involves a different perspective and partial occlusions.
Revised: The underlying environment dynamics are the same as in the 2D Shapes dataset, and we only change the rendering component to make for a visually more challenging task that involves a different perspective and partial occlusions.

Original: See Figure 2b for an example.
Revised: See Figure 2b for an example.

Original: B. A TARI 2600 G AMES
Revised: B. A TARI 2600 G AMES

Original: Atari Pong We make use of the Arcade Learning Environment (Bellemare et al., 2013) to create a small environment based on the Atari 2600 game Pong which is restricted to the ﬁrst interaction between the ball and the player-controlled paddle, i.e., we discard the ﬁrst couple of frames from the experience buffer where the opponent behaves completely independent of any player action.
Revised: Atari Pong We make use of the Arcade Learning Environment (Bellemare et al., 2013) to create a small environment based on the Atari 2600 game Pong which is restricted to the ﬁrst interaction between the ball and the player-controlled paddle, i.e., we discard the ﬁrst couple of frames from the experience buffer where the opponent behaves completely independent of any player action.

Original: Speciﬁcally, we discard the ﬁrst 58 (random) environment interactions.
Revised: Speciﬁcally, we discard the ﬁrst 58 (random) environment interactions.

Original: We use the P ONG D ETERMINISTICV 4 variant of the environment in OpenAI Gym (Brockman et al., 2016).
Revised: We use the P ONG D ETERMINISTICV 4 variant of the environment in OpenAI Gym (Brockman et al., 2016).

Original: We use a random policy and populate an experience buffer with 10 environment interactions per episode, i.e., T = 10 .
Revised: We use a random policy and populate an experience buffer with 10 environment interactions per episode, i.e., T = 10 .

Original: An observation consists of two consecutive frames, cropped (to exclude the score) and resized to 50 ×pixels each.
Revised: An observation consists of two consecutive frames, cropped (to exclude the score) and resized to 50 ×pixels each.

Original: See Figure 2c for an example.
Revised: See Figure 2c for an example.

Original: Space Invaders
Revised: Space Invaders

Original: This environment is based on the Atari 2600 game Space Invaders, using the
Revised: This environment is based on the Atari 2600 game Space Invaders, using the

Original: S PACE I NVADERS D ETERMINISTIC - V 4 variant of the environment in OpenAI Gym (Brockman et al., 2016), and processed / restricted in a similar manner as the Pong environment.
Revised: S PACE I NVADERS D ETERMINISTIC - V 4 variant of the environment in OpenAI Gym (Brockman et al., 2016), and processed / restricted in a similar manner as the Pong environment.

Original: We discard the ﬁrst 50 (random) environment interactions for each episode and only begin populating the experience buffer thereafter.
Revised: We discard the ﬁrst 50 (random) environment interactions for each episode and only begin populating the experience buffer thereafter.

Original: See Figure 2d for an example observation.
Revised: See Figure 2d for an example observation.

Original: B. 3-B ODY P HYSICS
Revised: B. 3-B ODY P HYSICS

Original: The 3-body physics simulation environment is an interacting system that evolves according to classical gravitational dynamics.
Revised: The 3-body physics simulation environment is an interacting system that evolves according to classical gravitational dynamics.

Original: Different from the other environments considered here, there are no actions.
Revised: Different from the other environments considered here, there are no actions.

Original: This environment is adapted from Jaques et al.
Revised: This environment is adapted from Jaques et al.

Original: (2019) using their publicly available implementation 1 , where we set the step size ( dt ) to 2 . 

Revised: (2019) using their publicly available implementation 1 , where we set the step size ( dt ) to 2 .

Original: and the initial maximum x and y velocities to 0 .
Revised: and the initial maximum x and y velocities to 0 .

Original: We concatenate two consecutive frames of 50 × 50 pixels each to provide the model with (implicit) velocity information.
Revised: We concatenate two consecutive frames of 50 × 50 pixels each to provide the model with (implicit) velocity information.

Original: See Figure 4 for an example observation sequence.
Revised: See Figure 4 for an example observation sequence.

Original: C E VALUATION M ETRICS C. H ITS AT R ANK K (H@ K )
Revised: C E VALUATION M ETRICS C. H ITS AT R ANK K (H@ K )

Original: This score is 1 for a particular example if the predicted state representation is in the k-nearest neighbor set around the encoded true observation, where we deﬁne the neighborhood of a node to include the node itself.
Revised: This score is 1 for a particular example if the predicted state representation is in the k-nearest neighbor set around the encoded true observation, where we deﬁne the neighborhood of a node to include the node itself.

Original: Otherwise this score is 0.
Revised: Otherwise this score is 0.

Original: In other words, this score measures whether the rank of the predicted state representation is smaller than or equal to k, when ranking all reference state representations by distance to the true state representation.
Revised: In other words, this score measures whether the rank of the predicted state representation is smaller than or equal to k, when ranking all reference state representations by distance to the true state representation.

Original: We report the average of this score over a particular evaluation dataset.
Revised: We report the average of this score over a particular evaluation dataset.

Original: C.M EAN R ECIPROCAL R ANK (MRR) 
Revised: C.

Revised: [Equation]

Original: This score is deﬁned as the average inverse rank, i.e., MRR = 1 N (cid:80) Nn = 1 rank n , where rank n is the rank of the n -th sample.
Revised: This score is deﬁned as the average inverse rank, i.e., MRR = 1 N (cid:80) Nn = 1 rank n , where rank n is the rank of the n -th sample.

Original: D A RCHITECTURE AND H YPERPARAMETERS D. O BJECT E XTRACTOR
Revised: D A RCHITECTURE AND H YPERPARAMETERS D. O BJECT E XTRACTOR

Original: For the 3D cubes environment, the object extractor is a 4-layer CNN with 3 × 3 ﬁlters, zero-padding, and 16 feature maps per layer, with the exception of the last layer, which has K = 5 feature maps, i.e., one per object slot.
Revised: For the 3D cubes environment, the object extractor is a 4-layer CNN with 3 × 3 ﬁlters, zero-padding, and 16 feature maps per layer, with the exception of the last layer, which has K = 5 feature maps, i.e., one per object slot.

Original: After each layer, we apply BatchNorm (Ioffe & Szegedy, 2015)  a ReLU( x ) = max(0 , x ) activation function.
Revised: After each layer, we apply BatchNorm (Ioffe & Szegedy, 2015) and a ReLU( x ) = max(0 , x ) activation function.

Original: For the 2D shapes environment, we choose a simpler
Revised: For the 2D shapes environment, we choose a simpler

Original: CNN architecture with only a single convolutional layer with 10 × 10 ﬁlters and a stride of 10 , followed by BatchNorm and a ReLU activation (LeakyReLU (Xu et al., 2015) for the Atariand physics environments).
Revised: CNN architecture with only a single convolutional layer with 10 × 10 ﬁlters and a stride of 10 , followed by BatchNorm and a ReLU activation or LeakyReLU (Xu et al., 2015) for the Atariand physics environments.

Original: This layer has 16 feature maps and is followed by a channel-wise linear transformation (i.e., a 1 × 1 convolution), with 5 feature maps as output.
Revised: This layer has 16 feature maps and is followed by a channel-wise linear transformation (i.e., a 1 × 1 convolution), with 5 feature maps as output.

Original: For both models, we choose a sigmoid activation function after the last layer to obtain object masks with values in (0 , 1) .
Revised: For both models, we choose a sigmoid activation function after the last layer to obtain object masks with values in (0 , 1) .

Original: We use the same two-layer architecture for the Atari 2600 environments and the 3-body physics environment, but with 9 × 9 ﬁlters (and zero-padding) in the ﬁrst layer, and 5 × 5 ﬁlters with a stride of 5 in the second layer.
Revised: We use the same two-layer architecture for the Atari 2600 environments and the 3-body physics environment, but with 9 × 9 ﬁlters (and zero-padding) in the ﬁrst layer, and 5 × 5 ﬁlters with a stride of 5 in the second layer.

Original: D. O BJECT E NCODER
Revised: D. O BJECT E NCODER

Original: After reshaping/ﬂattening the output of the object extractor, we obtain a vector representation per object (2500-dim for the 3D cubes environment, 25-dim for the 2D shapes environment, and 1000- dim for Atari 2600 and physics environments).
Revised: After reshaping/ﬂattening the output of the object extractor, we obtain a vector representation per object (2500-dim for the 3D cubes environment, 25-dim for the 2D shapes environment, and 1000dim for Atari 2600 and physics environments).

Original: The object encoder is an MLP with two hidden layers of 512 units and each, followed by ReLU activation functions.
Revised: The object encoder is an MLP with two hidden layers of 512 units and each, followed by ReLU activation functions.

Original: We further use LayerNorm (Ba et al., 2016) before the activation function of the second hidden layer.
Revised: We further use LayerNorm (Ba et al., 2016) before the activation function of the second hidden layer.

Original: The output of the ﬁnal output layer is 2-dimensional (4-dimensional for Atari 2600 and physics environments), reﬂecting the ground truth object state, i.e., the object coordinates in 2D (although this is not provided to the model), and velocity (if applicable).
Revised: The output of the ﬁnal output layer is 2-dimensional (4-dimensional for Atari 2600 and physics environments), reﬂecting the ground truth object state, i.e., the object coordinates in 2D (although this is not provided to the model), and velocity (if applicable).

Original: D. T RANSITION M ODEL
Revised: D. T RANSITION M ODEL

Original: Both the node and the edge model in the GNN-based transition model are MLPs with the same architecture / number of hidden units as the object encoder model, i.e., two hidden layers ofunits each, LayerNorm and ReLU activations.
Revised: Both the node and the edge model in the GNN-based transition model are MLPs with the same architecture / number of hidden units as the object encoder model, i.e., two hidden layers ofunits each, LayerNorm and ReLU activations.

Original: D. L OSS FUNCTION
Revised: D. L OSS FUNCTION

Original: The margin in the hinge loss is chosen as γ = 1 .
Revised: The margin in the hinge loss is chosen as γ = 1 .

Original: We further multiply the squared Euclidean distance d ( x, y ) in the loss function with a factor of 0 .
Revised: We further multiply the squared Euclidean distance d ( x, y ) in the loss function with a factor of 0 .

Original: 5 /σ 2 with σ = 0 .
Revised: 5 /σ 2 with σ = 0 .

Original: 5 to control the spread of the embeddings.
Revised: 5 to control the spread of the embeddings.

Original: We use the same setting in all experiments.
Revised: We use the same setting in all experiments.

Original: D. B ASELINES
Revised: D. B ASELINES

Original: World Model Baseline The World Model baseline is trained in two stages:
Revised: World Model Baseline The World Model baseline is trained in two stages:

Original: First, we train an autoencoder or a VAE with a 32-dim latent space, where the encoder is a CNN with the same architecture as the object extractor used in the C-SWM model, followed by an MLP with the same architecture as the object encoder module in C-SWM on the ﬂattened representation of the output of the encoder
Revised: First, we train an autoencoder or a VAE with a 32-dim latent space, where the encoder is a CNN with the same architecture as the object extractor used in the C-SWM model, followed by an MLP with the same architecture as the object encoder module in C-SWM on the ﬂattened representation of the output of the encoder

Original: CNN.
Revised: CNN.

Original: The decoder exactly mirrors this architecture where we replace convolutional layers with deconvolutional layers.
Revised: The decoder exactly mirrors this architecture where we replace convolutional layers with deconvolutional layers.

Original: We veriﬁed that this architecture can successfully build representations of single frames.
Revised: We veriﬁed that this architecture can successfully build representations of single frames.

Original: [Figure]
Revised: [Figure]

Original: Example reconstructions from the latent code are shown in Figure 13.
Revised: Example reconstructions from the latent code are shown in Figure 13.

Original: We experimented both with mean squared error and binary cross entropy (using the continuous channel values in [0, 1] as targets) as reconstruction loss in the (V)AE models, both of which are typical choices in most practical implementations.
Revised: We experimented both with mean squared error and binary cross entropy (using the continuous channel values in [0, 1] as targets) as reconstruction loss in the (V)AE models, both of which are typical choices in most practical implementations.

Original: We generally found binary cross entropy to be more stable to optimize and to produce better results, which is why we opted for this loss in all the baselines using decoders considered in the experimental section.
Revised: We generally found binary cross entropy to be more stable to optimize and to produce better results, which is why we opted for this loss in all the baselines using decoders considered in the experimental section.

Original: In the second stage, we freeze the model parameters of the auto-encoder and train a transition model with mean-squared error on the latent representations.
Revised: In the second stage, we freeze the model parameters of the auto-encoder and train a transition model with mean-squared error on the latent representations.

Original: For the VAE model, we use the predicted mean values of the latent representations.
Revised: For the VAE model, we use the predicted mean values of the latent representations.

Original: This transition model takes the form of an MLP with the same architecture and number of hidden units as the node model in C-SWM.
Revised: This transition model takes the form of an MLP with the same architecture and number of hidden units as the node model in C-SWM.

Original: We experimented both with a translational transition model (i.e., the transition model only predicts the latent state difference, instead of the full next state) and direct prediction of the next state.
Revised: We experimented both with a translational transition model (i.e., the transition model only predicts the latent state difference, instead of the full next state) and direct prediction of the next state.

Original: We generally found that the translational transition model performed better and used it throughout all the reported results.
Revised: We generally found that the translational transition model performed better and used it throughout all the reported results.

Original: The World Model baselines are trained with a smaller batch size of 512, which slightly improved performance and simpliﬁed memory management.
Revised: The World Model baselines are trained with a smaller batch size of 512, which slightly improved performance and simpliﬁed memory management.

Original: Ablations We perform the following ablations: 1) we replace the latent GNN with an MLP (per object, i.e., we remove the edge update function) to investigate whether a structured transition model is necessary, 2) we remove the state factorization and embed the full scene into a single latent variable of higher dimensionality (original dimensionality × number of original object slots) with an
Revised: Ablations We perform the following ablations: 1) we replace the latent GNN with an MLP (per object, i.e., we remove the edge update function) to investigate whether a structured transition model is necessary, 2) we remove the state factorization and embed the full scene into a single latent variable of higher dimensionality (original dimensionality × number of original object slots) with an

Original: MLP as transition model, and 3) we replace the contrastive loss with a pixel-based reconstruction loss on both the current state and the predicted next state (we add a decoder that mirrors the architecture of the encoder).
Revised: MLP as transition model, and 3) we replace the contrastive loss with a pixel-based reconstruction loss on both the current state and the predicted next state (we add a decoder that mirrors the architecture of the encoder).

Original: Physics-as-Inverse-Graphics (PAIG)
Revised: Physics-as-Inverse-Graphics (PAIG)

Original: For this baseline, we train the PAIG model from Jaques et al.
Revised: For this baseline, we train the PAIG model from Jaques et al.

Original: (2019) with the code provided by the authors 2 on our dataset with the standard settings recommended by the authors for this particular task, namely: model=
Revised: (2019) with the code provided by the authors 2 on our dataset with the standard settings recommended by the authors for this particular task, namely: model=

Original: PhysicsNet, epochs=500, batch size=100, base lr=1e-3, autoencoder loss=5.0, anneal lr=true, color=true, and cell type=gravity ode cell.
Revised: PhysicsNet, epochs=500, batch size=100, base lr=1e-3, autoencoder loss=5.0, anneal lr=true, color=true, and cell type=gravity ode cell.

Original: We use the same input size as in C-SWM, i.e., frames are of shape  50 × 50 × 3 .
Revised: We use the same input size as in C-SWM, i.e., frames are of shape 2 https:/ 50 × 50 × 3 .

Original: We further set input steps=2, pred steps=10 and extrap steps=0 to match our setting of predicting for a total of 10 frames while conditioning on a pair of 2 initial frames to obtain initial position and velocity information.
Revised: We further set input steps=2, pred steps=10 and extrap steps=0 to match our setting of predicting for a total of 10 frames while conditioning on a pair of 2 initial frames to obtain initial position and velocity information.

Original: We train this model with four different random seeds.
Revised: We train this model with four different random seeds.

Original: For evaluation, we extract learned position representations for all frames of the test set and further run the latent physics prediction module (conditioned on the ﬁrst two initial frames) to obtain model predictions.
Revised: For evaluation, we extract learned position representations for all frames of the test set and further run the latent physics prediction module (conditioned on the ﬁrst two initial frames) to obtain model predictions.

Original: We further augment position representations with velocity information by taking the difference between two consecutive position representations and concatenating this representation with the 2-dim position representation, which we found to slightly improve results.
Revised: We further augment position representations with velocity information by taking the difference between two consecutive position representations and concatenating this representation with the 2-dim position representation, which we found to slightly improve results.

Original: In total, we obtain a 12-dim (4-dim × 3 objects) representation for each time step, i.e., the same as C-SWM.
Revised: In total, we obtain a 12-dim (4-dim × 3 objects) representation for each time step, i.e., the same as C-SWM.

Original: From this, we obtain ranking metrics.
Revised: From this, we obtain ranking metrics.

Original: We found that one of the training runs (seed=2) collapsed all latent representations to a single point.
Revised: We found that one of the training runs (seed=2) collapsed all latent representations to a single point.

Original: We exclude this run in the results reported in Table 1, i.e., we report average and standard error over the three other runs only.
Revised: We exclude this run in the results reported in Table 1, i.e., we report average and standard error over the three other runs only.