Original: F ACTORED W ORLD M ODELS FOR Z ERO -S HOT G ENERALIZATION IN R OBOTIC M ANIPULATION
Revised: F ACTORED W ORLD M ODELS FOR Z ERO -S HOT G ENERALIZATION IN R OBOTIC M ANIPULATION

Original: Anonymous authors
Revised: Anonymous authors

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

Original: A BSTRACT
Revised: A BSTRACT

Original: World models for environments with many objects face a combinatorial explosion of states: as the number of objects increases, the number of possible arrangements grows exponentially.
Revised: World models for environments with many objects face a combinatorial explosion of states: as the number of objects increases, the number of possible arrangements grows exponentially.

Original: In this paper, we learn to generalize over robotic pick-andplace tasks using object-factored world models, which combat the combinatorial explosion by ensuring that predictions are equivariant to permutations of objects.
Revised: In this paper, we learn to generalize over robotic pick-andplace tasks using object-factored world models, which combat the combinatorial explosion by ensuring that predictions are equivariant to permutations of objects.

Original: We build on one such model, C-SWM, which we extend to overcome the assumption that each action is associated with one object.
Revised: We build on one such model, C-SWM, which we extend to overcome the assumption that each action is associated with one object.

Original: To do so, we introduce an action attention module to determine which objects are likely to be affected by an action.
Revised: To do so, we introduce an action attention module to determine which objects are likely to be affected by an action.

Original: The attention module is used in conjunction with a residual graph neural network block that receives action information at multiple levels.
Revised: The attention module is used in conjunction with a residual graph neural network block that receives action information at multiple levels.

Original: Based on RGB images and parameterized motion primitives, our model can accurately predict the dynamics of a robot building structures from blocks of various shapes.
Revised: Based on RGB images and parameterized motion primitives, our model can accurately predict the dynamics of a robot building structures from blocks of various shapes.

Original: Our model generalizes over training structures built in different positions.
Revised: Our model generalizes over training structures built in different positions.

Original: Moreover crucially, the learned model can make predictions about tasks not represented in training data.
Revised: Moreover crucially, the learned model can make predictions about tasks not represented in training data.

Original: That is, we demonstrate successful zero-shot generalization to novel tasks. For example, we measure only 2.4% absolute decrease in  our action ranking metric in the case of a block assembly task.
Revised: That is, we demonstrate successful zero-shot generalization to novel tasks, with only a minor decrease in model performance. Furthermore, we evaluate our model on a dataset collected on a physical robot.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: From assembly to household robots, current state-of-the-art robot learning agents cannot generalize beyond a speciﬁc training task.
Revised: From assembly to household robots, current state-of-the-art robot learning agents cannot generalize beyond a speciﬁc training task.

Original: One important aspect of generalization is the ability to understand any novel combination of known factors, a so-called combinatorial or compositional generalization.
Revised: One important aspect of generalization is the ability to understand any novel combination of known factors, a so-called combinatorial or compositional generalization.

Original: Applied to objects, combinatorial generalization ideally allows an agent to understand any arrangement of objects from only a limited number of interactions with its environment.
Revised: Applied to objects, combinatorial generalization ideally allows an agent to understand any arrangement of objects from only a limited number of interactions with its environment.

Original: The two key steps such agents need to perform are (a) decomposing a scene into individual objects and (b) modeling relative interactions between objects.
Revised: The two key steps such agents need to perform are (a) decomposing a scene into individual objects and (b) modeling relative interactions between objects.

Original: The latter is particularly important when learningto predict
Revised: The latter is particularly important when learning

Original:  the dynamics of the environment (i
Revised: to predict the dynamics of the environment (i

Original: e. learning a world model).
Revised: e. learning a world model)

Original: The purpose of our paper is to learn a world model that accurately predicts the effects of actions in the context of robotic manipulation, and generalizes to novel tasks.
Revised: The purpose of our paper is to learn a world model that accurately predicts the effects of actions in the context of robotic manipulation, and generalizes to novel tasks.

Original: We assume the state of the environment is already factored into individual images for each object, and focus on learning a latent code for each object and predicting the effect of pick-and-place actions.
Revised: We assume the state of the environment is already factored into individual images for each object, and focus on learning a latent code for each object and predicting the effect of pick-and-place actions.

Original: Our model is trained using a self-supervised contrastive loss.
Revised: Our model is trained using a self-supervised contrastive loss.

Original: It can predict accurate manipulation physics by using multiple graph neural network layers (GNNs) and action attention.
Revised: It can predict accurate manipulation physics by using multiple graph neural network layers (GNNs) and action attention.

Original: GNNs consider pairwise interactions, usually between every possible pair of objects, in order to predict the state of each object one step into the future.
Revised: GNNs consider pairwise interactions, usually between every possible pair of objects, in order to predict the state of each object one step into the future.

Original: GNNs achieve combinatorial generalization (to the extent that pairwise interactions are expressive enough) by equivariance to the order in which objects are presented–permutation equivariance.
Revised: GNNs achieve combinatorial generalization (to the extent that pairwise interactions are expressive enough) by equivariance to the order in which objects are presented–permutation equivariance.

Original: That is, they cannot overﬁt to a particular ordering of objects.
Revised: That is, they cannot overﬁt to a particular ordering of objects.

Original: It is a common choice to use a single GNN layer (in some cases called an Interaction Network (Battaglia et al., 2016)).
Revised: It is a common choice to use a single GNN layer (in some cases called an Interaction Network (Battaglia et al., 2016)).

Original: But, as shown in previous work (Kipf et al., 2018; Sanchez-Gonzalez et al., 2018) and conﬁrmed by our experiments, stacking multiple GNN layers leads to much more accurate physics predictions.
Revised: But, as shown in previous work (Kipf et al., 2018; Sanchez-Gonzalez et al., 2018) and conﬁrmed by our experiments, stacking multiple GNN layers leads to much more accurate physics predictions.

Original: The ability to integrate information about actions with latent states of individual objects has been underexplored in prior work.
Revised: The ability to integrate information about actions with latent states of individual objects has been underexplored in prior work.

Original: Previous object-factored world models either assume object-action association (a factored action space)
Revised: Previous object-factored world models either assume object-action association (a factored action space)

Original: (Veerapaneni et al., 2019; Kipf et al., 2020;
Revised: (Veerapaneni et al., 2019; Kipf et al., 2020;

Original: Huang et al., 2020) or only model sequences of states without actions (Janner et al., 2019; Bakhtin et al., 2019; Qi et al., 2021).
Revised: Huang et al., 2020) or only model sequences of states without actions (Janner et al., 2019; Bakhtin et al., 2019; Qi et al., 2021).

Original: We represent actions as pick(x,y) and place(x,y) with continuous (x,y) coordinates; this representation is integrated into our latent transition model through action attention and iterative reﬁnement by a stack of residual graph neural networks.
Revised: We represent actions as pick(x,y) and place(x,y) with continuous (x,y) coordinates; this representation is integrated into our latent transition model through action attention and iterative reﬁnement by a stack of residual graph neural networks.

Original: Our action attention module compares action information with the latent state of each object in order to predict the probability that an object is affected by the action.
Revised: Our action attention module compares action information with the latent state of each object in order to predict the probability that an object is affected by the action.

Original: Each node in a GNN then receives a different action based on the predicted attention weights.
Revised: Each node in a GNN then receives a different action based on the predicted attention weights.

Original: Since each GNN in our transition model has access to the action, the responsibility for deciding which object is affected by the action and modeling its effect on said object can be distributed across the transition model.
Revised: Since each GNN in our transition model has access to the action, the responsibility for deciding which object is affected by the action and modeling its effect on said object can be distributed across the transition model.

Original: We train and evaluate our factored world model in a simulated environment involving a UR robotic arm manipulating blocks of various shapes (Figure 1).
Revised: We train and evaluate our factored world model in a simulated environment involving a UR robotic arm manipulating blocks of various shapes (Figure 1).

Original: We instantiate two environments: Cubes includes six cubes arranged into ﬁve different structures, which take up to ten actions to build, and Shapes includes eight blocks of four types (cube, brick, triangle, roof), which are arranged into sixteen different structures (Figure 3).
Revised: We instantiate two environments: Cubes includes six cubes arranged into ﬁve different structures, which take up to ten actions to build, and Shapes includes eight blocks of four types (cube, brick, triangle, roof), which are arranged into sixteen different structures (Figure 3).

Original: To test zero-shot generalization — the main result of our paper — we partition the tasks into training and generalization sets.
Revised: To test zero-shot generalization — the main result of our paper — we partition the tasks into training and generalization sets.

Original: We ensure that the world model never sees the goal state of any generalization task during training.
Revised: We ensure that the world model never sees the goal state of any generalization task during training.

Original: Nevertheless, our model transfers to the testing tasks in  Shapes with only a very minor drop in accuracyand generalizes in Cubes with a mild to moderate performance drop (Table 1).
Revised: Nevertheless, our model transfers to the testing tasks in Cubes and Shapes with only a very minor drop in accuracy (Table 1).

Original: The mechanism behind successful zero-shot transfer is the ability of our model to understand the state of each object independently of others, and the ability to model sparse interactions.
Revised: The mechanism behind successful zero-shot transfer is the ability of our model to understand the state of each object independently of others, and the ability to model sparse interactions.

Original: For example, if our model understands the interaction ”stack a triangle on top of a cube” based on the ﬁrst training task in Shapes (Figure 3, bottom row), it can use it when generalizing to four different held out tasks (Figure 2).
Revised: For example, if our model understands the interaction ”stack a triangle on top of a cube” based on the ﬁrst training task in Shapes (Figure 3, bottom row), it can use it when generalizing to four different held out tasks (Figure 2).

Original: In summary, we contribute the following:
Revised: In summary, we contribute the following:

Original: We develop a factored wold model for robotic pick-and-place tasks.
Revised: We develop a factored wold model for robotic pick-and-place tasks.

Original: This model does not require a known factored action-object association (as was used in related models), owing to its use of an action-attention network.
Revised: This model does not require a known factored action-object association (as was used in related models), owing to its use of an action-attention network.

Original: We demonstrate that the learned world model accurately predicts the outcomes of sequences of actions.
Revised: We demonstrate that the learned world model accurately predicts the outcomes of sequences of actions.

Original: Moreover, the model achieves zero-shot generalization to sequences of actions tasks and object conﬁgurations that were not seen during training, including sequences that are longer (up to 10 actions) than those in the training data (up to 10 actions). 
Revised: Moreover, the model achieves zero-shot generalization to sequences of actions tasks and object conﬁgurations that were not seen during training, including sequences that are longer (up to 10 actions) than those in the training data (up to 10 actions). Finally, we demonstrate transfer to a real-world dataset.

Original:  F ACTORED W ORLD M ODEL FOR R OBOTIC P ICK - AND -P LACE
Revised:  F ACTORED W ORLD M ODEL FOR R OBOTIC P ICK - AND -P LACE

Original: The modeling task is to predict the effect of a sequence of actions a 1 , a 2 , ..., a T in an environment that is initialized to a state s 0 .
Revised: The modeling task is to predict the effect of a sequence of actions a 1 , a 2 , ..., a T in an environment that is initialized to a state s 0 .

Original: We will assume a setting in which the state of the world is represented as an image, which has been pre-processed into a factorized state s = (cid:104) s 1 , s 2 , ..., s K (cid:105) in which each s i is an image centered on the i th object.
Revised: We will assume a setting in which the state of the world is represented as an image, which has been pre-processed into a factorized state s = (cid:104) s 1 , s 2 , ..., s K (cid:105) in which each s i is an image centered on the i th object.

Original: The number of objects K can vary between episodes.
Revised: The number of objects K can vary between episodes.

Original: This postulated factorization can be implemented using an object detection module that predicts a bounding box for each object and a tracking module that corresponds bounding boxes across time.
Revised: This postulated factorization can be implemented using an object detection module that predicts a bounding box for each object and a tracking module that corresponds bounding boxes across time.

Original: Both of these problems are well-studied for natural images in computer vision (citations).
Revised: Both of these problems are well-studied for natural images in computer vision (Zhao et al.,

Original: Our goal is to learn a factorized latent representation
Revised: Our goal is to learn a factorized latent representation

Original: z = (cid:104)
Revised: z = (cid:104)

Original: z 1 , z 2 , ..., z K (cid:105) , along with a world model that predicts the state z (cid:48) that results from taking an action a at state z .
Revised: z 1 , z 2 , ..., z K (cid:105) , along with a world model that predicts the state z (cid:48) that results from taking an action a at state z .

Original: We will show that the learned latent code is low-dimensional; it represents object identity and position while abstracting over irrelevant features such as object color.
Revised: We will show that the learned latent code is low-dimensional; it represents object identity and position while abstracting over irrelevant features such as object color.

Original: Our world model is implemented using a graph neural network (citations) that represents action-conditioned pair-wise interactions of latent factors.
Revised: Our world model is implemented using a graph neural network (Gori et al., 2005) that represents action-conditioned pair-wise interactions of latent factors.

Original: It is able to perform several rounds of message passing, each successively reﬁning the prediction of the physics of picking and placing objects.
Revised: It is able to perform several rounds of message passing, each successively reﬁning the prediction of the physics of picking and placing objects.

Original: We do not learn a mapping from z to s (i.e. the world model only predicts the future within its latent code); nevertheless, we show the latent transition model successfully solves prediction and action ranking tasks.
Revised: We do not learn a mapping from z to s (i.e. the world model only predicts the future within its latent code); nevertheless, we show the latent transition model successfully solves prediction and action ranking tasks.

Original: Both the state encoder and the latent transition model are equivariant to the permutation of factors.
Revised: Both the state encoder and the latent transition model are equivariant to the permutation of factors.

Original: That is, we can provide objects in any order as long as it is consistent within an episode.
Revised: That is, we can provide objects in any order as long as it is consistent within an episode.

Original: Moreover, the model is agnostic to the number of objects present in a scene — no parameters in the model are dependent on the number of objects — although the weights of the latent transition model might overﬁt if the number does not vary during training.
Revised: Moreover, the model is agnostic to the number of objects present in a scene — no parameters in the model are dependent on the number of objects — although the weights of the latent transition model might overﬁt if the number does not vary during training.

Original: We describe our world model in detail in Section 2.1 and the training procedure in Section 2.2.
Revised: We describe our world model in detail in Section 2.1 and the training procedure in Section 2.2.

Original: 2. F ACTORED V ISUAL W ORLD M ODEL WITH A CTION A TTENTION
Revised: 2. F ACTORED V ISUAL W ORLD M ODEL WITH A CTION A TTENTION

Original: Encoder.
Revised: Encoder.

Original: Given a factored state s = (cid:104) s 1 , s 2 , ..., s k (cid:105) , the encoder f φ processes each factor (image of an object) separately.
Revised: Given a factored state s = (cid:104) s 1 , s 2 , ..., s k (cid:105) , the encoder f φ processes each factor (image of an object) separately.

Original: That is, z = (cid:104)
Revised: That is, z = (cid:104)

Original: f φ ( s 1 ) , f φ ( s 2 ) , ..., f φ ( s K ) (cid:105) .
Revised: f φ ( s 1 ) , f φ ( s 2 ) , ..., f φ ( s K ) (cid:105) .

Original: The encoder uses three convolutional layers followed by average pooling and three fully-connected layers.
Revised: The encoder uses three convolutional layers followed by average pooling and three fully-connected layers.

Original: In principle, the use of average pooling enables the encoder to process images of any size, although we use a ﬁxed image size.
Revised: In principle, the use of average pooling enables the encoder to process images of any size, although we use a ﬁxed image size.

Original: As described in Section 3, the per-object image s i includes both an RGB component and a coordinate grid component, which marks the location where the crop was taken (Figure 6).
Revised: As described in Section 3, the per-object image s i includes both an RGB component and a coordinate grid component, which marks the location where the crop was taken (Figure 4).

Original: Transition Model.
Revised: Transition Model.

Original: The transition model g θ accepts a factored latent state z = (cid:104) z 1 , z 2 , ..., z K (cid:105) , z i ∈ R D z and an action a ∈ R D a .
Revised: The transition model g θ accepts a factored latent state z = (cid:104) z 1 , z 2 , ..., z K (cid:105) , z i ∈ R D z and an action a ∈ R D a .

Original: It predicts the next latent state by outputting a residual:
Revised: It predicts the next latent state by outputting a residual:

Original: [Equation]
Revised: [Equation]

Original: We implement g θ ( z t , a t ) as a stack of graph neural network layers with skip connections.
Revised: We implement g θ ( z t , a t ) as a stack of graph neural network layers with skip connections.

Original: Denoting the i th graph neural network layer as GNN i and the i th intermediate representation as y i , the computation with L layers is as follows:
Revised: Denoting the i th graph neural network layer as GNN i and the i th intermediate representation as y i , the computation with L layers is as follows:

Original: [Equation]
Revised: [Equation]

Original: Our implementation of individual graph neural network layers follows Kipf et al.
Revised: Our implementation of individual graph neural network layers follows Kipf et al.

Original: (2020): the network models pairwise interactions between latent state factors (corresponding to objects in the environment) using a fully-connected node network h n and an edge network h e .
Revised: (2020): the network models pairwise interactions between latent state factors (corresponding to objects in the environment) using a fully-connected node network h n and an edge network h e .

Original: Both are implemented as MLPs with one hidden layer.
Revised: Both are implemented as MLPs with one hidden layer.

Original: The edge network outputs an embedding for each directed edge e i,j = h e ( z i , z j ) .
Revised: The edge network outputs an embedding for each directed edge e i,j = h e ( z i , z j ) .

Original: The edge embeddings are then aggregated using the node network in order to update the state of each node.
Revised: The edge embeddings are then aggregated using the node network in order to update the state of each node.

Original: Finally, the graph neural network outputs a vector of updated factors z (cid:48)
Revised: Finally, the graph neural network outputs a vector of updated factors z (cid:48)

Original: = (cid:104) z (cid:48) 1 , z (cid:48) 2 , ..., z (cid:48) K (cid:105) with
Revised: = (cid:104) z (cid:48) 1 , z (cid:48) 2 , ..., z (cid:48) K (cid:105) with

Original: [Equation]
Revised: [Equation]

Original: The summation over edges in the input of the node network assures permutation equivariance, an important property that aids generalization to novel combinations of factors.
Revised: The summation over edges in the input of the node network assures permutation equivariance, an important property that aids generalization to novel combinations of factors.

Original: Given a permutation of factors π , we have
Revised: Given a permutation of factors π , we have

Original: [Equation]
Revised: [Equation]

Original: The permutation equivariance property also holds for a stack of GNN layers as well as the state encoder and action attention.
Revised: The permutation equivariance property also holds for a stack of GNN layers as well as the state encoder and action attention.

Original: Action Attention.
Revised: Action Attention.

Original: When picking or placing objects in our experiments, the robotic arm and gripper will usually interact with only one or a small number of objects.
Revised: When picking or placing objects in our experiments, the robotic arm and gripper will usually interact with only one or a small number of objects.

Original: The transition model needs to learn which latent factors change and which stay the same when an action is performed.
Revised: The transition model needs to learn which latent factors change and which stay the same when an action is performed.

Original: Previous work simpliﬁed this problem by assuming that the action space is also factorized (Sanchez-Gonzalez et al., 2018;
Revised: Previous work simpliﬁed this problem by assuming that the action space is also factorized (Sanchez-Gonzalez et al., 2018;

Original: Veerapaneni et al., 2019; Kipf et al., 2020), which is to say that actions are performed relative to individual objects or parts.
Revised: Veerapaneni et al., 2019; Kipf et al., 2020), which is to say that actions are performed relative to individual objects or parts.

Original: This assumption is not realistic in typical robotics applications, where actions correspond to motor primitives that cannot trivially be associated with effects on individual objects.
Revised: This assumption is not realistic in typical robotics applications, where actions correspond to motor primitives that cannot trivially be associated with effects on individual objects.

Original: We use pick(x,y) and place(x,y) actions with continuous (x,y) coordinates.
Revised: We use pick(x,y) and place(x,y) actions with continuous (x,y) coordinates.

Original: We overcome this limitation by introducing an action attention layer.
Revised: We overcome this limitation by introducing an action attention layer.

Original: Instead of the action being the same for each node in the graph neural network, we create K transformations of the action that are attenuated based on how likely they are to interact with each latent factor.
Revised: Instead of the action being the same for each node in the graph neural network, we create K transformations of the action that are attenuated based on how likely they are to interact with each latent factor.

Original: The mechanism follows single-head self-attention (Vaswani et al., 2017) with the following key k , query q and value v :
Revised: The mechanism follows single-head self-attention (Vaswani et al., 2017) with the following key k , query q and value v :

Original: [Equation]
Revised: [Equation]

Original: Here, j k , j q and j v are Multi-Layer Perceptrons with one hidden layer.
Revised: Here, j k , j q and j v are Multi-Layer Perceptrons with one hidden layer.

Original: The factored action space a (cid:48) =
Revised: The factored action space a (cid:48) =

Original: (cid:104) a (cid:48) 1 , a (cid:48) 2 , ..., a (cid:48) K (cid:105) is composed of v weighted by attention weights
Revised: (cid:104) a (cid:48) 1 , a (cid:48) 2 , ..., a (cid:48) K (cid:105) is composed of v weighted by attention weights

Original: α i with i ∈ { 1 , ..., K } :
Revised: α i with i ∈ { 1 , ..., K } :

Original: [Equation]
Revised: [Equation]

Original: Note the new actions are all equal up to a multiplicative factor.
Revised: Note the new actions are all equal up to a multiplicative factor.

Original: Figure 5 visualizes α .
Revised: Figure 7 visualizes α .

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

Original: We train the encoder and latent transition model using a single-step contrastive loss.
Revised: We train the encoder and latent transition model using a single-step contrastive loss.

Original: We use a contrastive loss with a single positive and negative example (Kipf et al., 2020).
Revised: We use a contrastive loss with a single positive and negative example (Kipf et al., 2020).

Original: Given a transition (cid:104) s t , a t , s t +1 (cid:105) , we encode the current and next state ( z t = f φ ( s t ) , z t +1 =
Revised: Given a transition (cid:104) s t , a t , s t +1 (cid:105) , we encode the current and next state ( z t = f φ ( s t ) , z t +1 =

Original: f φ ( s t +1 ) , ˆ z t +1 =
Revised: f φ ( s t +1 ) , ˆ z t +1 =

Original: z t + g θ ( z t , a t ) ) and make a prediction about the next latent state.
Revised: z t + g θ ( z t , a t ) ) and make a prediction about the next latent state.

Original: The contrastive loss minimizes the distance between the real and predicted next latent state, while maximizing the encoding distance between the current state and a negative state ( ¯ z = f φ (¯ s ))
Revised: The contrastive loss minimizes the distance between the real and predicted next latent state, while maximizing the encoding distance between the current state and a negative state ( ¯ z = f φ (¯ s ))

Original: up to a margin γ :
Revised: up to a margin γ :

Original: L (
Revised: L (

Original: [Equation]
Revised: [Equation]

Original: t
Revised: t

Original: [Equation]
Revised: [Equation]

Original: The negative state ¯ s is sampled by randomly permuting a training batch of current states.
Revised: The negative state ¯ s is sampled by randomly permuting a training batch of current states.

Original: It is a proxy for sampling a random state from the entire dataset (without needing to increase the batch size).
Revised: It is a proxy for sampling a random state from the entire dataset (without needing to increase the batch size).

Original: Intuitively, we want the encoder to capture the minimum information required to distinguish a randomly sampled pair of states while enabling the latent transition model to be accurate.
Revised: Intuitively, we want the encoder to capture the minimum information required to distinguish a randomly sampled pair of states while enabling the latent transition model to be accurate.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our empirical evaluation focuses on pick-and-place robotic manipulation tasks performed in simulation as well as on a real-world robotic arm (Figure 1).
Revised: Our empirical evaluation focuses on pick-and-place robotic manipulation tasks performed in simulation as well as on a real-world robotic arm (Figure 1 and Figure 8).

Original: The environments involve a robotic arm manipulating objects of four shapes (cube, brick, triangle, roof) in a 30 × 30 cm workspace.
Revised: The environments involve a robotic arm manipulating objects of four shapes (cube, brick, triangle, roof) in a 30 × 30 cm workspace.

Original: The agent controlling the arm uses pre-deﬁned PICK and PLACE motion primitives: it chooses a particular ( x, y ) location (continuous action space) in the workspace in which the arm executes a top-down pick or place action depending on if it is holding an object.
Revised: The agent controlling the arm uses pre-deﬁned pick and place motion primitives: it chooses a particular ( x, y ) location (continuous action space) in the workspace in which the arm executes a top-down pick or place action depending on if it is holding an object.

Original: The state of the environment is captured by two RGB cameras pointed at the workspace and a third
Revised: The state of the environment is captured by two RGB cameras pointed at the workspace and a third

Original: RGB camera that captures the content of the robotic gripper.
Revised: RGB camera that captures the content of the robotic gripper.

Original: We assume access to a bounding box for each object.
Revised: We assume access to a bounding box for each object.

Original: The factored state s = (cid:104) s 1 , s 2 , ..., s K (cid:105) is created by cropping the contents of each bounding box with added padding and resizing the resulting image to an 18 × 18 square.
Revised: The factored state s = (cid:104) s 1 , s 2 , ..., s K (cid:105) is created by cropping the contents of each bounding box with added padding and resizing the resulting image to an 18 × 18 square.

Original: We add bounding box coordinates into the cropped image in the form of four additional coordinate grids channels (Figure 6, Section A.1).
Revised: We add bounding box coordinates into the cropped image in the form of four additional coordinate grids channels (Figure 4, Section A.1).

Original: Hence, each view of each object results in an 18 × 18 image with 3 RGB and 4 coordinate grid channels.
Revised: Hence, each view of each object results in an 18 × 18 image with 3 RGB and 4 coordinate grid channels.

Original: Note that we only require object bounding boxes within the images taken by side-viewing cameras.
Revised: Note that we only require object bounding boxes within the images taken by side-viewing cameras.

Original: We do not need the ground-truth ( x, y, z ) object positions in order to generate the factored state.
Revised: We do not need the ground-truth ( x, y, z ) object positions in order to generate the factored state.

Original: We concatenate the two views of each object channel-wise, creating a 14-channel image for each object.
Revised: We concatenate the two views of each object channel-wise, creating a 14-channel image for each object.

Original: If an object is held by the robotic gripper, we replace the two views with two images of the gripper — one from the front and one from the side.
Revised: If an object is held by the robotic gripper, we replace the two views with two images of the gripper — one from the front and one from the side.

Original: These two images can be captured by a single camera by rotating the robotic arm in between images.
Revised: These two images can be captured by a single camera by rotating the robotic arm in between images.

Original: The coordinate grids for hand images are set to zero.
Revised: The coordinate grids for hand images are set to zero.

Original: The environment poses a difﬁcult exploration problem: a sequence of random actions is unlikely to create a large structure without knocking it over.
Revised: The environment poses a difﬁcult exploration problem: a sequence of random actions is unlikely to create a large structure without knocking it over.

Original: A model-free agent with a random exploration policy fails to learn to build structures involving more than two objects (Biza et al.
Revised: A model-free agent with a random exploration policy fails to learn to build structures involving more than two objects (Biza et al.

Original: (2021), Table 2, ”DQN RS”).
Revised: (2021), Table 2, ”DQN RS”).

Original: Therefore, we collect a dataset of demonstrations from an experiment with added randomness.
Revised: Therefore, we collect a dataset of demonstrations from an experiment with added randomness.

Original: We evaluate both the ability of our model to ﬁt tasks with expert trajectories as well as the ability to generalize to a set of held-out tasks without ﬁne-tuning (zero-shot transfer, Figure 3).
Revised: We evaluate both the ability of our model to ﬁt tasks with expert trajectories as well as the ability to generalize to a set of held-out tasks without ﬁne-tuning (zero-shot transfer, Figure 3).

Original: We aim to answer the following questions:
Revised: We aim to answer the following questions:

Original: • Can the Factored World Model ﬁt physics of robotic manipulation?
Revised: • Can the Factored World Model ﬁt physics of robotic manipulation?

Original: (Section 3.1) • Does permutation equivariance in the Factored World Model facilitate generalization to unseen tasks?
Revised: (Section 3.1) • Does permutation equivariance in the Factored World Model facilitate generalization to unseen tasks?

Original: (Section 3.2) •
Revised: (Section 3.2) •

Original: What is the contribution of the individual components of our model to its performance?
Revised: What is the contribution of the individual components of our model to its performance?

Original: Which component of the factored state is the most important?
Revised: Which component of the factored state is the most important?

Revised: (Section 3.3) •

Revised: Can we use a model trained in simulation to make prediction by trajectories generated by a physical robot?

Original: (Section 3.3)
Revised: (Section 3.4).

Original: 3. L EARNING R OBOTIC M ANIPULATION D YNAMICS
Revised: 3. L EARNING R OBOTIC M ANIPULATION D YNAMICS

Original: Setup.
Revised: Setup.

Original: To train the Factored World Model, we collect datasets from the training tasks in the Cubes and Shapes environments (Figure 3).
Revised: To train the Factored World Model, we collect datasets from the training tasks in the Cubes and Shapes environments (Figure 3).

Original: Each dataset consists of 200k transitions collect by an expert with added randomness, see Appendix A.2 for details.
Revised: Each dataset consists of 200k transitions collect by an expert with added randomness, see Appendix A.2 for details.

Original: The Factored World Model learns forepochs using contrastive learning, and is subsequently evaluated both on the training and testing tasks.
Revised: The Factored World Model learns forepochs using contrastive learning, and is subsequently evaluated both on the training and testing tasks.

Original: We use a block position prediction metric to evaluate the quality of the learned representation as well as the latent transition model.
Revised: We use a block position prediction metric to evaluate the quality of the learned representation as well as the latent transition model.

Original: We also evaluate the model in a setting where it predicts the outcomes of several action sequences.
Revised: We also evaluate the model in a setting where it predicts the outcomes of several action sequences.

Original: The action sequences are similar, but only one of them reaches the goal state of a given task.
Revised: The action sequences are similar, but only one of them reaches the goal state of a given task.

Original: By predicting the outcome of each sequence, the model guesses which one it is.
Revised: By predicting the outcome of each sequence, the model guesses which one it is.

Original: We use this setting as a proxy for planning.
Revised: We use this setting as a proxy for planning.

Original: Block position prediction (RMSE): After training, we freeze the model and train an additional decoder to predicts the ( x, y, z ) position of the center of each object from its latent representation.
Revised: Block position prediction (RMSE): After training, we freeze the model and train an additional decoder to predicts the ( x, y, z ) position of the center of each object from its latent representation.

Original: The decoder, a Multi-Layer Perceptron with two hidden layers, is trained for ﬁve epochs.
Revised: The decoder, a Multi-Layer Perceptron with two hidden layers, is trained for ﬁve epochs.

Original: It is trained using additional supervision (pairs of states and position labels) not available during training of the
Revised: It is trained using additional supervision (pairs of states and position labels) not available during training of the

Original: Factored World Model.
Revised: Factored World Model.

Original: We evaluate both the model’s ability to represent block positions in the current state ( t = 0 ) and its ability to make predictions for trajectories ( t > 0 ).
Revised: We evaluate both the model’s ability to represent block positions in the current state ( t = 0 ) and its ability to make predictions for trajectories ( t > 0 ).

Original: We report the root mean squared error in centimeters.
Revised: We report the root mean squared error in centimeters.

Original: The quantity predicted in this setting is different than the bounding box coordinates provided to the model as input.
Revised: The quantity predicted in this setting is different than the bounding box coordinates provided to the model as input.

Original: The bounding box coordinates index into a ﬂat 2D images; the model needs to account for perspective to predict the (x,y,z) positions of objects.
Revised: The bounding box coordinates index into a ﬂat 2D images; the model needs to account for perspective to predict the (x,y,z) positions of objects.

Original: [Table]
Revised: [Table]

Original: Action sequence ranking (Hits@1, MRR): We start with an optimal action sequence that achieves a goal state of a particular task.
Revised: Action sequence ranking (Hits@1, MRR): We start with an optimal action sequence that achieves a goal state of a particular task.

Original: Then, we generate ten other action sequences where each action is perturbed by noise with magnitude (cid:15) .
Revised: Then, we generate ten other action sequences where each action is perturbed by noise with magnitude (cid:15) .

Original: The noise is added to each action in the sequence by drawing (cid:15) ∼ Unif [0 , (cid:15) ] and θ ∼ Unif [0 , 2 π ] that are then combined to create a 2D vector in the direction of θ with a length of (cid:15) that is added to the (x,y) coordinate of each action.
Revised: The noise is added to each action in the sequence by drawing (cid:15) ∼ Unif [0 , (cid:15) ] and θ ∼ Unif [0 , 2 π ] that are then combined to create a 2D vector in the direction of θ with a length of (cid:15) that is added to the (x,y) coordinate of each action.

Original: Each perturbed action sequence must not achieve the given task; otherwise, we re-sample noise.
Revised: Each perturbed action sequence must not achieve the given task; otherwise, we re-sample noise.

Original: The model observes the starting state and predicts the ﬁnal state of each action sequence.
Revised: The model observes the starting state and predicts the ﬁnal state of each action sequence.

Original: We also encode the ﬁnal state of the correct action sequence and compare its distance to the predicted ﬁnal states of all action sequences.
Revised: We also encode the ﬁnal state of the correct action sequence and compare its distance to the predicted ﬁnal states of all action sequences.

Original: Hits@1 report the fraction of times the model’s prediction of the ﬁnal state of the correct action sequences was closer to the encoded ﬁnal state than the predictions for all of the incorrect action sequences.
Revised: Hits@1 report the fraction of times the model’s prediction of the ﬁnal state of the correct action sequences was closer to the encoded ﬁnal state than the predictions for all of the incorrect action sequences.

Original: Since the incorrect action sequences are chosen to be a small distance from the correct one, a model making random predictions is not expected to reach Hits@1 of 1/11.
Revised: Since the incorrect action sequences are chosen to be a small distance from the correct one, a model making random predictions is not expected to reach Hits@1 of 1/11.

Original: In fact, the model needs to be fairly competent to get a non-zero score.
Revised: In fact, the model needs to be fairly competent to get a non-zero score.

Original: Results.
Revised: Results.

Original: We report results for training tasks in Table 1 top section, Figure 7 left and Figure 4 left.
Revised: We report results for training tasks in Table 1 top section, Figure 5 left and Figure 6 left.

Original: The Factored World Model reaches low block prediction error (2 cm for Cubes and 0.9 cm for
Revised: The Factored World Model reaches low block prediction error (1.5 cm for Cubes and 0.8 cm for

Original: Shapes; the size of a cube is 3 cm for comparison) and high action ranking score (99% for Cubes and 89% for Shapes).
Revised: Shapes; the size of a cube is 3 cm for comparison) and high action ranking score (98% for Cubes and 78% for Shapes).

Original: We compare to an autoencoder baseline (FWM-AE) and the Region Proposal
Revised: We compare to an autoencoder baseline (FWM-AE) and the Region Proposal

Original: Interaction Network (RPIN, Qi et al. (2021)) with minor changes.
Revised: Interaction Network (RPIN, Qi et al. (2021)) with minor changes.

Original: For FWM-AE, we add a decoder to our model and train with an autoencoding loss (speciﬁcally the Embed to Control loss Watter et al.
Revised: For FWM-AE, we add a decoder to our model and train with an autoencoding loss (speciﬁcally the Embed to Control loss Watter et al.

Original: (2015)) instead of contrastive learning.
Revised: (2015)) instead of contrastive learning.

Original: To adapt RPIN, a factored video prediction method, we append an (x,y,pick/place) action to the input of the Prediction module in their Convolutional[Figure] [Figure] [Figure]
Revised: To adapt RPIN, a factored video prediction method, we append an (x,y,pick/place) action to the input of the Prediction module in their Convolutional

Original: [Figure]

Original: Interaction Network.
Revised: Interaction Network.

Original: We recreate our dataset in the same format as their Shape Stacking experiment ( 224 × 224 images with a bounding box for each object) and we use the same hyper-parameter.
Revised: We recreate our dataset in the same format as their Shape Stacking experiment ( 224 × 224 images with a bounding box for each object) and we use the same hyper-parameter.

Original: Compared to an autoencoding loss (FWM-AE), contrastive learning (FWM) succeeds both in Cubes and Shapes, whereas the autoencoding loss fails to ﬁt Cubes entirely.
Revised: Compared to an autoencoding loss (FWM-AE), contrastive learning (FWM) succeeds both in Cubes and Shapes, whereas the autoencoding loss fails to ﬁt Cubes.

Original: In Shapes, autoencoding outperforms contrastive loss in terms of block position prediction, but underperforms in action sequence ranking.
Revised: In Shapes, autoencoding outperforms contrastive loss in terms of block position prediction, but underperforms in action sequence ranking.

Original: We believe the difference can be explained by the focus of the two different losses: autoencoding loss directly incentivizes the model to reconstruct the state of the environment (which include bounding box coordinates in the form of a coordinate grid), whereas contrastive loss focuses on learning compact state representations predictive of the future.
Revised: We believe the difference can be explained by the focus of the two different losses: autoencoding loss directly incentivizes the model to reconstruct the state of the environment (which include bounding box coordinates in the form of a coordinate grid), whereas contrastive loss focuses on learning compact state representations predictive of the future.

Original: Our adaptation of RPIN to action-conditioned sequence prediction fails to learn a dynamics model.
Revised: Our adaptation of RPIN to action-conditioned sequence prediction fails to learn a dynamics model.

Original: The encoder of RPIN can capture some position information: it reaches around 6 cm RMSE for the task of predicting ground-truth object coordinates without forward modeling.
Revised: The encoder of RPIN can capture some position information: it reaches around 6 cm RMSE for the task of predicting ground-truth object coordinates without forward modeling.

Original: However, the forward model fails, reaching 20 cm RMSE.
Revised: However, the forward model fails, reaching 20 cm RMSE.

Original: We further study the block prediction error in Figure 4.
Revised: We further study the block prediction error in Figure 6.

Original: Figure 4 reports block prediction error for noisy trajectories generated by an expert policy with added randomness, whereas Table 1 reports errors for optimal trajectories that reach the goal state of each of the training (and generalization) tasks.
Revised: Figure 6 reports block prediction error for noisy trajectories generated by an expert policy with added randomness, whereas Table 1 reports errors for optimal trajectories that reach the goal state of each of the training (and generalization) tasks.

Original: In Shapes (Figure 4 right), we see a reversal compared to Table 1, as contrastive loss outperforms autoencoding by a small margin.
Revised: In Shapes (Figure 6 right), we see a reversal compared to Table 1, as contrastive loss outperforms autoencoding by a small margin.

Original: This result suggests that the transition model learned by autoencoding is not as accurate for blocks falling and colliding with each other.
Revised: This result suggests that the transition model learned by autoencoding is not as accurate for blocks falling and colliding with each other.

Original: Finally, we plot action ranking Hits@1 as a function of noise level (cid:15) for Shapes in Figure 7 left.
Revised: Finally, we plot action ranking Hits@1 as a function of noise level (cid:15) for Shapes in Figure 5 left.

Original: Contrastive loss outperforms autoencoding for all noise levels with the area under the curve being 84% for FWM and 64% for FWM-AE.
Revised: Contrastive loss outperforms autoencoding for all noise levels with the area under the curve being 84% for FWM and 64% for FWM-AE.

Original: 3. Z ERO -S HOT G ENERALIZING TO N OVEL T ASKS
Revised: 3. Z ERO -S HOT G ENERALIZING TO N OVEL T ASKS

Original: Setup.
Revised: Setup.

Original: We further evaluate models trained in Section 3.2 on tasks unseen during training (Figure 3).
Revised: We further evaluate models trained in Section 3.2 on tasks unseen during training (Figure 3).

Original: We speciﬁcally make sure that the training datasets do not contain a single example of a goal state of the novel tasks.
Revised: We speciﬁcally make sure that the training datasets do not contain a single example of a goal state of the novel tasks.

Original: In Cubes, the training task is building a 4-stack and the novel tasks are building a wall, stairs, a wall on ground and stairs on ground.
Revised: In Cubes, the training task is building a 4-stack and the novel tasks are building a wall, stairs, a wall on ground and stairs on ground.

Original: We found that the training data collection policy never solves the novel tasks (we use a separate policy to collect evaluation trajectories for testing tasks).
Revised: We found that the training data collection policy never solves the novel tasks (we use a separate policy to collect evaluation trajectories for testing tasks).

Original: In Shapes, we train on half of the possible structures with a height of three and roof on top, and test generalization on the other half.
Revised: In Shapes, we train on half of the possible structures with a height of three and roof on top, and test generalization on the other half.

Original: In this case, the data collection policy for the eight training tasks solves one of the testing tasks around once every 50 episodes (due to added randomness) and we delete these episodes.
Revised: In this case, the data collection policy for the eight training tasks solves one of the testing tasks around once every 50 episodes (due to added randomness) and we delete these episodes.

Original: We report block position error and action ranking Hits@1 described in
Revised: We report block position error and action ranking Hits@1 described in

Original: Section 3.1.
Revised: Section 3.1.

Original: We do not perform any ﬁne-tuning on the novel tasks in this experiment.
Revised: We do not perform any ﬁne-tuning on the novel tasks in this experiment.

Original: Result.
Revised: Result.

Original: In  Shapes, FWM transfers to the novel tasks with only a minor decrease in performance (block position error increases by  0.2 cm  and action sequence ranking Hits @1 decrease by 2%).
Revised: In both Cubes and Shapes, FWM transfers to the novel tasks with only a minor decrease in performance (block position error increases by 0.4 cm in Cubes and 0.2 cm in Shapes and action sequence ranking Hits @1 remain almost unchanged).

Original: We consider this to be an important demonstration of the generalization properties of permutation equivariant models.
Revised: We consider this to be an important demonstration of the generalization properties of permutation equivariant models.

Original: In contrast, only limited generalization and transfer properties have been shown in prior work (Li et al., 2020; Biza et al., 2021).
Revised: In contrast, only limited generalization and transfer properties have been shown in prior work (Li et al., 2020; Biza et al., 2021).

Original: FWM-AE suffers a large decrease in action sequence ranking(17% decrease) while still outperforming FWM on block position prediction.
Revised: FWM-AE suffers a large decrease in action sequence ranking while still outperforming FWM on block position prediction in Shapes.

Original: We plot action ranking error as a function of (cid:15)  in Figure 7 right; we see a pattern analogous to results for training tasks (Section 3.1).
Revised: We plot action ranking error as a function of (cid:15) for Shapes in Figure 5 right; we see a pattern analogous to results for training tasks (Section 3.1).

Original: The areas under curve are 81% for FWM and 48% for FWM-AE.
Revised: The areas under curve are 81% for FWM and 48% for FWM-AE.

Original: In Cubes, the training and testing tasks are different not only in the structure built but also in the time horizon. Building a 4-stack takes six actions whereas all novel Cubes tasks have a time horizon of ten actions. Compared to training, FWM block position error increases by 3.2 cm and the action ranking score decreases by 12%. The main difﬁculty of stacking many cubes close to each other to form a wall or stairs is modeling the possibility that the robotic gripper hits some of the blocks when attempting to place a block and topples over the structure.
Revised: 3. FWM A BLATIONS AND A TTENTION V ISUALIZATION

Original: But, we observe that FWM still retains a relatively high ability to discern successful versus unsuccessful action sequences (87% action sequence ranking Hits@1).

Original: FWM-AE did not ﬁt the training tasks well and, hence, performs poorly in the zero-shot generalization setting.

Original: We report the results of the following ablations in Table 1: • using only one graph neural network layer (1 GNN Layer),• removing action attention (No Attention), • only showing bounding box coordinates as inputs (No RGB),• only showing RGB images as inputs (No Coordinates) and• using a monolithic latent transition model (No Factorization). 
Revised: We report the results of the following ablations in Table 1: removing action attention (No Attention and 1 GNN, No Att.), using only one graph neural network layer (1 GNN Layer), only showing bounding box coordinates as inputs (No RGB), only showing RGB images as inputs (No Coordinates) and using a monolithic latent transition model (No Factorization). We did not ﬁnd a signiﬁcant difference when using FWM with and without action attention.

Original: Action attention improves block position prediction while having only a minor impact on action ranking Hits@1. In our preliminary experiments, we observed that the latent states of objects not being acted on tend todrift as the latent transition model makes predictions.
Revised: As shown in Figure 7 ﬁrst row and third row, FWM learns interpretable attention weights. During a pick action, the attention weights reﬂect the object being picked up as well as any object that is nearby. The attention module has an interesting behavior during place actions: it gives a weight of zero to the object being placed, but attends to all objects that are underneath or next to the place location.

Original: We hypothesize that action attention prevents the drift by decreasing the weights of actions applied to graph neural network nodes corresponding to static objects.
Revised: For example, in the top-right image in Figure 7, the green blocks is being placed on top of the red, blue and black blocks. The attention module only attends to these three blocks with the highest weight assigned to the red blocks, which is directly below the object being placed.

Original: The decreased drift, in turn, makes the block position decoder more accurate.
Revised: Why is there no quantitative increase in performance even though the attention module functions properly?

Original: We visualize attention weights in Figure 5. No Factorization and No Coordinates ablations show that both of these components are vital to  the 
Revised: We hypothesize that what the attention module does explicitly the four GNN layers in FWM can learn to do implicitly using the node and edge networks. To test this hypothesis, we train a model with a single GNN layer (commonly used in prior work) with and without attention.

Revised: We assume that the reduced capacity of the model will require it to use the attention module.

Revised: Here, action attention leads to large quantitative improvements.

Revised: E.g. not using attention more than doubles block prediction error in Cubes.

Revised: Furthermore, qualitative analysis shows that the attention module behaves differently with 1 GNN layer compared to 4 GNN layers.

Revised: With 1 GNN layer, the module strictly attends to the object being manipulated (Figure 7, second and fourth row), during both pick and place , while ignoring objects that are nearby.

Revised: A second interesting result in our ablation studies is the comparison of FWM with and without

Revised: RGB image.

Revised: The No RGB baseline only has access to the coordinates of the bounding boxes for each object.

Revised: We see no statistically signiﬁcant different between these two settings.

Revised: Since the

Revised: RGB baseline has access to bounding box coordinates from two different cameras, it can use this information to triangulate where an object is located in the workspace.

Revised: We believe more complex object interaction (e.g. packing a bin with household objects, building from LEGO) would put higher emphasis on images.

Revised: We leave these experiments for future work.

Revised: No Factorization and No Coordinates ablations show that both of these components are vital to the

Original: Factored World Model.
Revised: Factored World Model.

Original: Without coordinate grids, the model perceives images of objects without any information about where they are located in the environment.
Revised: Without coordinate grids, the model perceives images of objects without any information about where they are located in the environment.

Original: Without object factorization, the monolithic latent transition model is unable to generalize beyond the speciﬁc arrangements of blocks seen in the training tasks.
Revised: Without object factorization, the monolithic latent transition model is unable to generalize beyond the speciﬁc arrangements of blocks seen in the training tasks.

Original: Here, we mean in-distribution generalization to building known structures in novel locations of the workspace.
Revised: Here, we mean in-distribution generalization to building known structures in novel locations of the workspace.

Original: Naturally, the monolithic model also fails in zero-shot transfer.
Revised: Naturally, the monolithic model also fails in zero-shot transfer.

Original: Finally, we ﬁnd stacking of four residual GNNs leads to large improvements compared to a single
Revised: Finally, we ﬁnd stacking of four residual GNNs leads to large improvements compared to a single

Original: GNN layer.
Revised: GNN layer.

Original: In Shapes, the action ranking Hits@1 increase by 
Revised: In Shapes, the action ranking Hits@1 increase by 19% with four GNN layers.

Original:  R ELATED W ORK
Revised: 3. T RANSFER TO A R EAL -W ORLD D ATASET

Revised: Setup. We collect evaluation trajectories on a physical robot (Figure 8) for the ﬁve cube stacking tasks shown in Figure 3, top row. The real-world dataset includes bounding boxes (generated by segmenting cubes by color) and ground-truth block positions annotated by hand with the help of a depth camera.

Revised: We do not train our model on the real-world dataset, since we only collectepisodes for each task.

Revised: Instead, we change the background and object colors in simulation to match the real-world images.

Revised: In all cases, we train on the task of stacking four blocks on top of each other and transfer to other four tasks in the Cubes environment.

Revised: That is, we test the ability of our model to transfer from building a tower of four blocks in simulation to building stairs, walls, etc.

Revised: in the real world without any additional ﬁne-tuning.

Revised: Result.

Revised: We compare results from simulation and the real-world in Table 2.

Revised: On average, the block position prediction error increases by 0.4 cm as we transfer from simulation to the real-world dataset.

Revised: Our model appears to struggle with tasks that involve stacking blocks on top of each other in the real world compared to simulation.

Revised: Conversely, there no decrease in performance for Wall Ground (WG) and Stairs Ground (WG), where all blocks are placed on the ground.

Revised:  R ELATED W ORK

Original: World Models with Object Factorization.
Revised: World Models with Object Factorization.

Original: Action-conditioned object-factorized world models have been explored in toy object environments and Atari games (Kipf et al., 2020; Huang et al., 2020) as well as in robotic manipulation (Veerapaneni et al., 2019) and controlling billiard balls (Kossen et al., 2020).
Revised: Action-conditioned object-factorized world models have been explored in toy object environments and Atari games (Kipf et al., 2020; Huang et al., 2020) as well as in robotic manipulation (Veerapaneni et al., 2019) and controlling billiard balls (Kossen et al., 2020).

Original: We build upon C-SWM (Kipf et al., 2020), which has been further extended by (Huang et al., 2020) to handle identically looking objects (one of the assumptions of C-SWM is that each object is distinct).
Revised: We build upon C-SWM (Kipf et al., 2020), which has been further extended by (Huang et al., 2020) to handle identically looking objects (one of the assumptions of C-SWM is that each object is distinct).

Original: Veerapaneni et al.
Revised: Veerapaneni et al.

Original: (2019) used a factored transition model to learn and plan for stacking up to three cubes.
Revised: (2019) used a factored transition model to learn and plan for stacking up to three cubes.

Original: Different from these works we consider a simpliﬁed setup wherein our model does not perform symmetry breaking or object discovery from raw visual features, since we assume our environment to be already factored. 
Revised: Different from these works we consider a simpliﬁed setup wherein our model does not perform symmetry breaking or object discovery from raw visual features, since we assume our environment to be already factored. [Table] [Table] the

Revised: [Table]

Revised: four transfer tasks shown in Figure 3, top row

Original: Object-factored world models are also commonly used in the context of physics prediction in videos.
Revised: Object-factored world models are also commonly used in the context of physics prediction in videos.

Original: Janner et al.
Revised: Janner et al.

Original: ; Veerapaneni et al.
Revised: ; Veerapaneni et al.

Original: (2019) demonstrated generalization in a Tetris-like environment, where a world model predicts the outcome of dropping blocks from a height.
Revised: (2019) demonstrated generalization in a Tetris-like environment, where a world model predicts the outcome of dropping blocks from a height.

Original: Ye et al.
Revised: Ye et al.

Original: (2019) studied learning dynamics of pushing objects with a robotic arm in the real world.
Revised: (2019) studied learning dynamics of pushing objects with a robotic arm in the real world.

Original: RPIN (Qi et al., 2021) demonstrated state-of-the-art results in the tasks of 2D physics modeling (PHYRE benchmark, Bakhtin et al.
Revised: RPIN (Qi et al., 2021) demonstrated state-of-the-art results in the tasks of 2D physics modeling (PHYRE benchmark, Bakhtin et al.

Original: , modeling billiard balls and modeling videos of falling stacks of blocks.
Revised: , modeling billiard balls and modeling videos of falling stacks of blocks.

Original: As we showed in our comparison, models like RPIN cannot be naively applied to our pick-and-place tasks, both due to the problem of including action information at the right point in a transition model and due to the large changes between subsequent states in our environment (object disappearing from the environment vs. a video of an object being slowly picked up by a robotic hand).
Revised: As we showed in our comparison, models like RPIN cannot be naively applied to our pick-and-place tasks, both due to the problem of including action information at the right point in a transition model and due to the large changes between subsequent states in our environment (object disappearing from the environment vs. a video of an object being slowly picked up by a robotic hand).

Original: Object factorization is an instance of a general factored Markov Decision Process, which has been studied in the context of policy search (Guestrin et al., 2003), factor discovery (Jonsson & Barto, 2005) and MDP abstraction (Ravindran, 2004; Wolfe & Barto, 2006).
Revised: Object factorization is an instance of a general factored Markov Decision Process, which has been studied in the context of policy search (Guestrin et al., 2003), factor discovery (Jonsson & Barto, 2005) and MDP abstraction (Ravindran, 2004; Wolfe & Barto, 2006).

Original: In fact, factored MDP works often include examples where individual factors are objects (e.g. Tower of Hanoi in Ravindran (2004)).
Revised: In fact, factored MDP works often include examples where individual factors are objects (e.g. Tower of Hanoi in Ravindran (2004)).

Original: Early works speciﬁcally on object factorization include Wolfe (2006); Diuk et al.
Revised: Early works speciﬁcally on object factorization include Wolfe (2006); Diuk et al.

Original: Model-based Learning for Robotics.
Revised: Model-based Learning for Robotics.

Original: One prominent line of work focuses on optimal control or policy search in low-dimensional state space.
Revised: One prominent line of work focuses on optimal control or policy search in low-dimensional state space.

Original: If such space is directly available, PILCO ﬁts a
Revised: If such space is directly available, PILCO ﬁts a

Original: Gaussian Process to real-world robot dynamics and performs policy search through gradient descent in the model (Deisenroth & Rasmussen, 2011).
Revised: Gaussian Process to real-world robot dynamics and performs policy search through gradient descent in the model (Deisenroth & Rasmussen, 2011).

Original: Successful cart-pole swing and robotic unicycle balance is achieved within a few seconds of online learning.
Revised: Successful cart-pole swing and robotic unicycle balance is achieved within a few seconds of online learning.

Original: In the case of image states, a lowdimensional latent space can be learned by variational autoencoders (see Lesort et al.
Revised: In the case of image states, a lowdimensional latent space can be learned by variational autoencoders (see Lesort et al.

Original: (2018) for a survey).
Revised: (2018) for a survey).

Original: Watter et al.
Revised: Watter et al.

Original: (2015) learn a latent state space together with a locally linear transition model, enabling efﬁcient inference in optimal control methods.
Revised: (2015) learn a latent state space together with a locally linear transition model, enabling efﬁcient inference in optimal control methods.

Original: Watter et al.
Revised: Watter et al.

Original: (2015) demonstrate successful control of pendulum, cart-pole and a simpliﬁed robotic arm operating in 2D. Follow-up works add additional constraints to the framework in order to improve dynamics model accuracy and robustness (Karl et al., 2017; Banijamali et al., 2018; Levine et al., 2020).
Revised: (2015) demonstrate successful control of pendulum, cart-pole and a simpliﬁed robotic arm operating in 2D. Follow-up works add additional constraints to the framework in order to improve dynamics model accuracy and robustness (Karl et al., 2017; Banijamali et al., 2018; Levine et al., 2020).

Original: Other works have used models of the world to accelerate learning for robotic manipulation tasks.
Revised: Other works have used models of the world to accelerate learning for robotic manipulation tasks.

Original: In
Revised: In

Original: Nair et al.
Revised: Nair et al.

Original: (2018), a VAE, trained on image observations, is used for sampling goals during training and providing a reward signal via distance in the latent space.
Revised: (2018), a VAE, trained on image observations, is used for sampling goals during training and providing a reward signal via distance in the latent space.

Original: This method achieved similar sample complexity to a state-based method on a robotic pick-and-place task.
Revised: This method achieved similar sample complexity to a state-based method on a robotic pick-and-place task.

Original: In Nasiriany et al.
Revised: In Nasiriany et al.

Original: (2019), a challenging robotic pushing task was performed by searching for sequences of subgoals in the latent space of a VAE that could be followed by a learned policy.
Revised: (2019), a challenging robotic pushing task was performed by searching for sequences of subgoals in the latent space of a VAE that could be followed by a learned policy.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: The results in this paper demonstrate that factorized world models can successfully be adapted to manipulation tasks with continuous state and action spaces by introducing an action-attention model to isolate the effects of actions on individual objects.
Revised: The results in this paper demonstrate that factorized world models can successfully be adapted to manipulation tasks with continuous state and action spaces by introducing an action-attention model to isolate the effects of actions on individual objects.

Original: The resulting models are able to learn transition dynamics that generalize to previously unseen sequences of actions and conﬁgurations of objects at test time, and are able to predict outcomes at moderate time horizons of up to 10 actions.
Revised: The resulting models are able to learn transition dynamics that generalize to previously unseen sequences of actions and conﬁgurations of objects at test time, and are able to predict outcomes at moderate time horizons of up to 10 actions.

Original: This represents a signiﬁcant step in applying these models in practical robotics tasks.
Revised: This represents a signiﬁcant step in applying these models in practical robotics tasks.

Original: An immediate future line of work is to integrate these world models with planning algorithms, and evaluate the ability of the resulting planner to generalize to tasks that were unseen during training.
Revised: An immediate future line of work is to integrate these world models with planning algorithms, and evaluate the ability of the resulting planner to generalize to tasks that were unseen during training.

Original: R EPRODUCIBILITY S TATEMENT
Revised: R EPRODUCIBILITY S TATEMENT

Original: We will open-source our code with the camera-ready version of our paper.
Revised: We will open-source our code with the camera-ready version of our paper.

Original: We will include our simulated PyBullet environments, dataset generation scripts, Pytorch model deﬁnitions and training scripts.
Revised: We will include our simulated PyBullet environments, dataset generation scripts, Pytorch model deﬁnitions and training scripts.

Original: A A PPENDIX A. E NVIRONMENT D ETAILS
Revised: A A PPENDIX A. E NVIRONMENT D ETAILS

Original: Our PyBullet simulation consists of a UR5 robotic arm with a Robotiq gripper operating over a 30 × 30 cm workspace.
Revised: Our PyBullet simulation consists of a UR5 robotic arm with a Robotiq gripper operating over a 30 × 30 cm workspace.

Original: In comparison, the size of a cube is 3 × 3 cm.
Revised: In comparison, the size of a cube is 3 × 3 cm.

Original: The environment is captured from two sides by two cameras.
Revised: The environment is captured from two sides by two cameras.

Original: Each camera produces a 90 × 90 image.
Revised: Each camera produces a 90 × 90 image.

Original: PyBullet also provides a segmentation mask for each object, which we use to create bounding boxes.
Revised: PyBullet also provides a segmentation mask for each object, which we use to create bounding boxes.

Original: We ﬁrst draw the smallest rectangular box that captures the whole segmentation mask and add 4 px symmetric padding to it.
Revised: We ﬁrst draw the smallest rectangular box that captures the whole segmentation mask and add 4 px symmetric padding to it.

Original: If the bounding box is smaller than 18 × 18 px, then they are padded up to the minimum size.
Revised: If the bounding box is smaller than 18 × 18 px, then they are padded up to the minimum size.

Original: We crop the RGB image and four coordinate grids inside of the bounding box and resize them to an 18 × 18 image.
Revised: We crop the RGB image and four coordinate grids inside of the bounding box and resize them to an 18 × 18 image.

Original: The aspect ration of the RGB image gets corrupted in this step, but the model can still capture the size of each object based on the coordinate grids.
Revised: The aspect ration of the RGB image gets corrupted in this step, but the model can still capture the size of each object based on the coordinate grids.

Original: We have a vertical and a horizontal coordinate grid that traverses the interval of [-1, 1] either left-to-right or up-to-down.
Revised: We have a vertical and a horizontal coordinate grid that traverses the interval of [-1, 1] either left-to-right or up-to-down.

Original: We also create two additional coordinate grids that traverse the interval right-to-left and down-to-up. 
Revised: We also create two additional coordinate grids that traverse the interval right-to-left and down-to-up. Block colors.

Revised: We use a single block color (red) for all experiments in simulation.

Revised: In Figure 7, blocks are colored only for the purpose of matching blocks to attention weights.

Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Revised: actual states passed into the model.

Revised: In the real-world experiments, we assign each block a different color so that we can get bounding boxes using color segmentation (Table 2).

Revised: We use different colors both in the Sim2Real dataset (generated in simulation) and the Real dataset (generated on a physical robot),

Original: A. E XPERT P OLICIES
Revised: A. E XPERT P OLICIES

Original: Experts for training datasets.
Revised: Experts for training datasets.

Original: For Cubes, we have the following policy: with 70% probability either pick a random cube that is not covered by other objects or place a cube on top of any other cube.
Revised: For Cubes, we have the following policy: with 70% probability either pick a random cube that is not covered by other objects or place a cube on top of any other cube.

Original: Otherwise, execute a pick or place action in a random coordinate.
Revised: Otherwise, execute a pick or place action in a random coordinate.

Original: Additionally, 1 cm random noise is added to the (x,y) position of each action.
Revised: Additionally, 1 cm random noise is added to the (x,y) position of each action.

Original: The agent does not have to decide between pick and place actions: hand empty → pick, hand full → place.
Revised: The agent does not have to decide between pick and place actions: hand empty → pick, hand full → place.

Original: For Shapes, we use trained SDQfD agents provided by Biza et al.
Revised: For Shapes, we use trained SDQfD agents provided by Biza et al.

Original: With 80% probability we execute the actions predicted by the expert.
Revised: With 80% probability we execute the actions predicted by the expert.

Original: Otherwise, during pick, we pick a random object with 50% and we execute a pick at a random location otherwise; during place, we place in a random position.
Revised: Otherwise, during pick, we pick a random object with 50% and we execute a pick at a random location otherwise; during place, we place in a random position.

Original: We collect 200k training transitions for Cubes and Shapes.
Revised: We collect 200k training transitions for Cubes and Shapes.

Original: Experts for generalization datasets.
Revised: Experts for generalization datasets.

Original: We simply collect optimal trajectories for evaluation and ﬁlter out any trajectory that does not reach the goal of each task.
Revised: We simply collect optimal trajectories for evaluation and ﬁlter out any trajectory that does not reach the goal of each task.

Original: A. M ODEL D ETAILS
Revised: A. M ODEL D ETAILS

Original: Per-object Encoder CNN.
Revised: Per-object Encoder CNN.

Original: Conv2D( 5 × 5 kernel size, 32 kernels, stride 2, padding 1) → BatchNorm (Ioffe & Szegedy, 2015)
Revised: Conv2D( 5 × 5 kernel size, 32 kernels, stride 2, padding 1) → BatchNorm (Ioffe & Szegedy, 2015)

Original: → ReLU → Conv2D( 5 × 5 kernel size, 64 kernels, stride 2, padding 1) →
Revised: → ReLU → Conv2D( 5 × 5 kernel size, 64 kernels, stride 2, padding 1) →

Original: BatchNorm → ReLU → Conv2D( 5 × 5 kernel size, 64 kernels, stride 2, padding 1).
Revised: BatchNorm → ReLU → Conv2D( 5 × 5 kernel size, 64 kernels, stride 2, padding 1).

Original: GNN.
Revised: GNN.

Original: Both the node and the edge networks are MLP with one hidden layer of size 512.
Revised: Both the node and the edge networks are MLPs with one hidden layer of size 512.

Original: Each layer, except the output layer, is followed by LayerNorm (Ioffe & Szegedy, 2015) and ReLU activation.
Revised: Each layer, except the output layer, is followed by LayerNorm (Ba et al., 2016) and ReLU activation.

Original: Action Attention.
Revised: Action Attention.

Original: The key, value and query MLPs have a single hidden layer of size 512, the output size is also 512.
Revised: The key, value and query MLPs have a single hidden layer of size 512, the output size is also 512.

Original: LayerNorm and ReLU is used in the same way as in the GNN.
Revised: LayerNorm and ReLU is used in the same way as in the GNN.

Original: Decoder for FWM-AE Baseline.
Revised: Decoder for [Table] [Table] [Table]

Original: ConvTranspose2D( 5 × 5 kernel size, 64 kernels, stride 2) →
Revised: → [Table] → Conv2D( 3 × 3 kernel size, [Table]

Original: BatchNorm → ReLU → Conv2D( 3 × 3 kernel size, 32 kernels, stride 2) → BatchNorm → ReLU → Conv2D( 3 × 3 kernel size, 14 kernels, stride 1).
Revised: [Table]

Original: Training.
Revised: Training.

Original: We use the Adam (Kingma & Ba, 2015) optimizer with default parameters and a learning rate of 5 e − 5 .
Revised: We use the Adam (Kingma & Ba, 2015) optimizer with default parameters and a learning rate of 5 e − 5 .

Original: We train for 200 epochs (dataset size is 200k transitions) with a batch size of 256. 
Revised: We train for 200 epochs (dataset size is 200k transitions) with a batch size of 256. [Figure] [Figure] [Figure]

Revised: [Figure]

Revised: [Figure]