Factored World Models for Zero-Shot Generalization in Robotic Manipulation
Authors: Ondrej Biza
Thomas Kipf
David Klee
Robert Platt
Jan-Willem van de Meent
Lawson L.S. Wong
Venue: ICLR 2022 Submitted (2022)
Keywords: reinforcement learning
world models
robotic manipulation
zero-shot transfer

Abstract: 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. In this paper, we learn to generalize over robotic pick-and-place tasks using object-factored world models, which combat the combinatorial explosion by ensuring that predictions are equivariant to permutations of objects. We build on one such model, C-SWM, which we extend to overcome the assumption that each action is associated with one object. To do so, we introduce an action attention module to determine which objects are likely to be affected by an action. The attention module is used in conjunction with a residual graph neural network block that receives action information at multiple levels. 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. Our model generalizes over training structures built in different positions. Moreover crucially, the learned model can make predictions about tasks not represented in training data. 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.

Under review as a conference paper at ICLR 2022
FACTORED WORLD MODELS FOR ZERO-SHOT GENERALIZATION IN ROBOTIC MANIPULATION
Anonymous authors Paper under double-blind review
ABSTRACT
World models for environments with many objects face combinatorial explosion of states: as the number of objects increases, the number of possible arrangements grows exponentially.
this paper, we learn generalize over robotic pick-and- place tasks using object-factored world models, which combat the combinatorial explosion by ensuring that predictions are equivariant to permutations of objects. We build on one such model, C-SWM, which we extend to overcome the assump- tion that each action is associated with one object. To do so, we introduce an action attention module to determine which objects are likely to be affected by action. The attention module is used in conjunction with a residual graph neural network block that receives action information multiple levels. Based on RGB images and parameterized motion primitives, our model can accurately predict the dynamics of robot building structures from blocks of various shapes. Our model generalizes over training structures built in different positions. Moreover crucially, the learned model can make predictions about tasks not represented training data. That is, we demonstrate successful zero-shot generalization to novel tasks, with only a minor decrease in model performance. Furthermore, we evalu- ate our model on dataset collected on physical robot.
INTRODUCTION
From assembly to household robots, current state-of-the-art robot learning agents cannot generalize beyond specific training task. One important aspect of generalization is the ability to understand any novel combination of known factors, a so-called combinatorial or composi- tional generalization. Applied to objects, com- binatorial generalization ideally allows an agent to understand any arrangement of objects from only a limited number of interactions with its environment. The two key steps such agents need to perform are (a) decomposing a scene into individual objects and (b) modeling rela- tive interactions between objects. The latter is particularly important when learning to predict the dynamics of the environment (i.e. learning a world model).
//
Figure PyBullet environment with a URS arm.
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. 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. Our model is trained using self-supervised contrastive loss. It can predict accurate manipulation physics by using multiple graph neural network layers (GNNs) and action attention. GNNs consider pairwise interactions, usually between every possible pair of objects. in order predict the state of each object one step into the future. GNNs achieve combinatorial generalization (to the extent that pairwise interactions expressive enough) by equivariance to the order in which objects presented-permutatiin

---

Under review as a conference paper at ICLR 2022
equivariance. That is, they cannot overfit particular ordering of objects. It is common choice to use a single GNN layer (in some cases called an Interaction Network (Battaglia et al., 2016))
But, as shown in previous work (Kipf al. 2018 (aanchez- Goatheezz-nzzzzzzze et al. 2018) and confirmed by our experiments, stacking multiple GNN layers leads to much more accurate physics predictions.
The ability to integrate information about actions with latent states of individual objects has been underexplored in prior work. Previous object-factored world models either assume object-action
association (a factored action space) (Veerapaneni al., 2019 Kipf et al., 2020; Huang al. 2020) or only model sequences of states without actions (Janner et al. 2019 Bakhtin et al.) 2019
al., 2021]. 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 refinement by a stack of residual graph neural networks. 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. Each node in a GNN then receives a different action based on the predicted attention weights. Since each GNN in transition model has access to the action, the responsibility for deciding which object affected by the action and modeling its effect on said object can be distributed across the transition model.
We train and evaluate our factored world model in simulated environment involving a URS robotic arm manipulating blocks of various shapes (Figure D We instantiate two environ- ments: Cubes includes six cubes arranged into five different structures, which take up to ac- tions to build, and Shapes includes eight blocks of four types (cube, brick, triangle, roof), which are arranged into sixteen different structures (Figure 3). To test zero-shot generalization
the main result of our paper we partition the tasks into training and generalization sets. We ensure that the world model never sees the goal state of any generalization task during training. Nevertheless, our model transfers to the testing tasks in Cubes and Shapes with only very minor drop in accuracy (Table T). The mechanism behind successful zero-shot transfer is the ability of model to understand the state of each object independently of others, and the ability to model sparse interactions. For example, if our model understands the interaction "stack a triangle on top of a cube" based on the first training task in Shapes (Figure 3. bottom row), it can use when generalizing to four different held out tasks 2).
Figure 2: Example of generalization.
In summary, we contribute the following:
1. We develop a factored wold model for robotic pick-and- place tasks. This model does not require known factored action-object association (as was used in related models), owing to its use of an action-attention network.
2. We demonstrate that the learned world model accurately predicts the outcomes of quences of actions. Moreover, the model achieves zero-shot generalization to sequences of actions tasks and object configurations that were not seen during training, including se- quences that are longer (up to 10 actions) than those in the training data (up to 10 actions). Finally, we demonstrate transfer real-world dataset.
FACTORED WORLD MODEL FOR ROBOTIC PICK-AND-PLACE
The modeling task to predict the effect of a sequence of actions al, a² ..... aT in an environment that is initialized to state s". We will assume setting in which the state of the world is represented as an image, which has been pre-processed into a factorized state = (s1, 82,,.. sK) in which each is an image centered on the ith object. The number of objects K can vary between episodes. This postulated factorization can be implemented using an object detection module that predicts a bounding box for each object and tracking module that corresponds bounding boxes across time. Both problems are well-studied for natural images in computer vision (Zhao et al. 2019 Greff et al. 2020
2

---

Under review as a conference paper at ICLR 2022
Training Zero-Shot Transfer
Figure 3: Breakdown of tasks in Cubes (top row) and Shapes (bottom row) environments. In Cubes, we train on the task of making stack of four blocks and test on two wall stacking and two stair stacking tasks. In shapes, we instantiate 16 tasks with structures of height of three and roof and top, and split them 50/50 training/testing. We provide expert demonstrations for training tasks and test zero-shot generalization in prediction on testing tasks.
Our goal is to learn factorized latent representation (21, 22, =k), along with world model that predicts the state that results from taking an action a at state z. We will show that learned latent code is low-dimensional; it represents object identity and position while abstracting over irrelevant features such as object color. Our world model implemented using graph neural network (Gori al. 2005) that represents setion-oondiiioned pair-wise interactions of latent factors. is able to perform several rounds of message passing, each successively refining the prediction of the physics of picking and placing objects. We do not learn mapping from (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.
Both the state encoder and the latent transition model are equivariant to the permutation of factors. That is, we can provide objects any order as long as is consistent within an episode. 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 overfit if the number does not vary during training. We describe our world model in detail in Section 2.1 and the training procedure in Section 2.2
2.1 FACTORED VISUAL WORLD MODEL WITH ACTION ATTENTION
Encoder. Given factored state = (81, sk), the encoder fo processes each factor (image of (sx)). The encoder uses three convolu-
an object) separately, That is, z tional layers followed by average pooling and three fully-connected layers. In principle, the use of average pooling enables the encoder to process images of any size, although we use fixed image size. As described in Section 3 the per-object image s, includes both RGB component and coordinate grid component, which marks the location where the crop was taken (Figure
Transition Model. The transition model go accepts factored latent state (21, 22, -..-... zk), 2 E D: and an action E TR: It predicts the next latent state by outputting a residual:
gt+l = (2 a²).
(1)
We implement (=", ing the ith graph neural network layer as GNN, and the ith intermediate representation yi, the computation with L layers is as follows:
(aa,,,d a²) as a stack of graph neural network layers with skip connections. Denot-
/1 = GNN, (z, a)
1/2 31 GNN:()1.
go (0(2,0) a) WL-1 GNN
3

---

Our implementation of individual graph neural network layers follows Kipf al. (2020): the net- work models pairwise interactions between latent state factors (corresponding to objects in the envi- ronment) using fully-connected node network h, and an edge network h.- Both are implemented as MLPs with one hidden layer. The edge network outputs an embedding for each directed edge Under review as a conference paper at ICLR 2022
eij (z 2z). The edge embeddings then aggregated using the node network in order to update the state of each node. Finally, the graph neural network outputs vector of updated factors (z1, 22 w.. ²k) with
2 = 1mn 01
Σ f
E
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. Given a permutation of factors n, we have
GNN( = (GNNN22,))
(4)
The permutation equivariance property also holds for a stack of GNN layers as well as the state encoder and action attention.
When picking or placing objects in our experiments, the robotic arm and gripper will usually interact with only one small number of objects. The transition model needs to learn which latent factors change and which stay the same when an action is performed. Previous work simplified this problem by assuming that the action space is also factorized 2018 Veerapaneni et al. 2019; Kipl et al. (020). which is to say that actions are performed relative to individual objects or parts. 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. We use pick (x, y) and place (x, y) actions with continuous (x, y) coordinates.
Action Attention.
Saunhez-GGnzzale
al.,
We overcome this limitation by introducing action attention layer. Instead of the action being the same for each node the graph neural network, we create transformations of the action that are attenuated based how likely they are to interact with each latent factor. The mechanism follows single-head self-attention (Vaswani et al. 2017) with the following key k. query q and value v:
k = (ja (21), (x(x..)))
q ia(a),
V = jo(a).
(5)
Here, ju. ig and jo are Multi-Layer Perceptrons with one hidden layer. The factored action space a' (a/, 022 -. - axx) is composed of weighted by attention weights og with E {1..... K}:
0;- = softmax (aimaaxXiffor q, hig k
a = a.V.
(6)
Note the new actions are all equal up to multiplicative factor. Figure visualizes
2.2 LEARNING BY CONTRASTIVE LOss
We train the encoder and latent transition model using single-step contrastive loss. We use contrastive loss with single positive and negative example (Kipf et al. 2020). Given transition (s', a", 1 we encode the current and next state (z* = fo( t (s*t), gt, = + go(2", a²)) and make prediction about the next latent state. The contrastive loss minimizes the distance between the real and predicted next latent state, while maximizing the encoding distance between the current state and negative state (a = fo()) up to margin y::
K
K
L( 211,, 2111, 2,, 2) 2Ko
2, 21
+ max 0, 2K
2
&
2
(7)
The negative state ø is sampled by randomly permuting a training batch of current states. is a proxy for sampling random state from the entire dataset (without needing increase the batch size). Intuitively, we want the encoder to capture the minimum information required to distinguish randomly sampled pair of states while enabling the latent transition model to be accurate.
4

---

Under review as a conference paper at ICLR 2022
3 EXPERIMENTS
Our empirical evaluation focuses on pick-and- -aaaaadddeeceeeee robotic manipulation tasks performed in simu- lation as well a robotic arm manipulating objects of four shapes (cube, brick, triangle, roof) in a 30x30 cm workspace. The agent controlling the uses pre-defined pick and place motion primitives: chooses a particular (z, y) location (continuous action space) in the workspace in which the arm executes a top-down pick or place action depending if it is holding an object.
a real-world robotic arm (Figure [ and Figure 8) The environments involve
The state of the environment is captured by two RGB cameras pointed at the workspace and a third RGB camera that captures the content of the robotic gripper. We assume access to a bounding box for each object. The factored state s = ($1, $2,..., 3K) ) is created by cropping the contents of each bounding box with added padding and resizing the resulting image to an 18x square. We add bounding box coordinates into the cropped image in the form of four additional coordinate grids channels (Figure Section[A.I] Hence, each view of each object results in an 18x 18 image with 3 RGB and 4 coordinate grid channels. Note that we only require object bounding boxes within the images taken by side-viewing cameras. We do not need the ground-truth (ar,) y. object positions in order to generate the factored state. We concatenate the two views of each object channel- creating 14-channel image for each object. 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. These two images can be captured by a single camera by rotating the robotic arm in between images. The coordinate grids for hand images set to zero.
The environment poses a difficult exploration problem: a sequence of random actions is unlikely create a large structure without knocking it over. A model-free agent with random exploration policy fails to learn to build structures involving more than two objects (Biza et al. (2021), Table 2, "DQN RS"). Therefore, randomness. We evaluate both the ability of our model to fit tasks with expert trajectories as well as the ability generalize to set of held-out tasks without fine-tuning (zero-shot transfer, Figure 3)
collect dataset of demonstrations from an experiment with added
We aim to answer the following questions:
Can the Factored World Model fit physics of robotic manipulation? (Section 3.1
Does permutation equivariance in the Factored World Model facilitate generalization to unseen tasks? (Section 3.2)
What is the contribution of the individual components of our model to its performance? Which component of the factored state is the most important? (Section 3.3
Can we use a model trained simulation to make prediction by trajectories generated by a physical robot? (Section 3.4)
3.1 LEARNING ROBOTIC MANIPULATION DYNAMICS
Setup. To train the Factored World Model, we collect datasets from the training tasks the Cubes and Shapes environments (Figure 3. Each dataset consists of 200k transitions collect by an expert with added randomness, see Appendix [A.2 for details. The Factored World Model learns for 200 epochs using contrastive learning, and is subsequently evaluated both on the training and testing tasks. We use a block position prediction metric to evaluate the quality of the learned representation as well the latent transition model. We also evaluate the model in setting where it predicts the outcomes of several action sequences. The action sequences are similar, but only one of them reaches the goal state of a given task. By predicting the outcome of each sequence, the model guesses which one it is. We use this setting as a proxy for planning.
Block position prediction (RMSE): After training, we freeze the model and train an additional de- coder to predicts the (xr, y, 2) position of the center of each object from its latent representation. The decoder, a Multi-Layer Perceptron with two hidden layers, is trained for five epochs. It is trained using additional supervision (pairs of states and position labels) not available during training of the Factored World Model. We evaluate both the model's ability to represent block positions in the cur- rent state (t = 0) and its ability to make predictions for trajectories (t > 0). We report the root mean squared error in centimeters. The quantity predicted in this setting different than the bounding box coordinates provided to the model as input. The bounding box coordinates index into a flat 2D images; the model needs to account for perspective predict the (x, y z) positions of objects.
5

---

Under review as a conference paper at ICLR 2022
Method
Cubes-Test Set
Shapes-Test Set
RMSE (cm)
Hits @ 1 (%)
RMSE (cm)
Hits @ I (%)
RPIN (Qi et al.
2021
20.0
C-SWM (Kipi et al. 2020
11.39
Ozea
FWM-AE
2.01
60.7
0.67 1:032
65. 1:022
FWM (our)
1,46
98.3
0.82::0.14
78.3
FWM - GNN Layer
1.78
83.5
0.81 1:008
57.4
FWM GNN, No Att.
3.60
43.1 1143
1.34
46.2
FWM No Attention
1.63
97.7
0.89 20008
76. 16611:20 1:24
FWM No RGB
1. 1 +00007
97.3 #0.5
0.61 +-0.06
78.0
FWM No Coordinates
11.64....0
0.0.0
9.56 +0.97
0.:000
FWM No Factorization
1
6. 19 20.81
O
Method
Cubes-Zero- Zero-Shot
Shapes Zero-Shot
RMSE (cm)
@its@1 1 (%)
RMSE (cm)
Hits 1 (%)
RPIN (Qi et al. 202
21.7
0
C-SWM (Kipt et al. (020
9.15
O=0.
FWM-AE
2.77
37.9
54.0#265
FWM (our)
1.88
98.4
1.02 #0016
77.4
FWM GNN Layer
2.05
95,0
1.02 d:0.0s
56.5
FWM GNN, No Att.
5.89 #0.71
33.2
1.60 +..0s
51.0#5.2
FWM No Attention
2.22
97.7
1 1.09 20:08
75.7 #300
FWM No RGB
1.39 =0.38
98.2
0. 0.76
79.8 24.5
FWM No Coordinates
9.82
0=0.
9.81 10:0
O.00
FWM- No Factorization
8.39 :.39::111
0.0.0
6.79 00082
0..00
Table 1: Comparison between Factored World Models, baselines and ablations in Cubes and Shapes environments. We evaluate the models both on the tasks they were trained on (top section) and on un- seen tasks without additional fine-tuning (bottom section). While the bottom section reports results for generalization to unseen tasks, reaching high scores in the top section requires in-distribution generalization to known structures built in new positions in the workspace. Hence the terms "Test Set" and "Zero-Shot", We report block position error (RMSE, the lower the better, unbounded) and action sequence ranking score (Hits@1 1, the higher the better, bounded 0 - 100) (Section 3.1 Each model was run with 8 random seeds and we report means and 95% confidence intervals.
Action sequence ranking (Hits@1 1, MRR): We start with an optimal action sequence that achieves a goal state of particular task. Then, we generate ten other action sequences where each action is perturbed by noise with magnitude e. The noise is added to each action the sequence by drawing O0, c] and 0 ~ Unif(0, 2n] that are then combined to create 2D vector in the direction of with a length of e that is added to the (x, coordinate of each action. Each perturbed action sequence must
achieve the given task; otherwise, we re-sample noise. The model observes the starting state and predicts the final state of each action sequence. We also encode the final state of the correct action sequence and compare its distance the predicted final states of all action sequences. Hits@ @ report the fraction of times the model's prediction of the final state of the correct action sequences was closer the encoded final state than the predictions for all of the incorrect action sequences. Since the incorrect action sequences are chosen to be a small distance from the correct one, model making random predictions is not expected reach Hits@ I of 1/11. fact, the model needs to be fairly competent to get a non-zero score.
Results. We report results for training tasks in Table T top section, Figure 5 left and Figure left. The Factored World Model reaches low block prediction error (1.5 cm for Cubes and 0.8 cm for Shapes; the size of a cube cm for comparison) and high action ranking score (98% for Cubes and 78% for Shapes). We compare to autoencoder baseline (FWM-AE) and the Region Proposal Interaction Network (RPIN, Qi et al. (2021) with minor changes. For FWM-AE, we add a decoder to our model and train with an autoencoding loss (specifically the Embed to Control loss Watter et al. (215)) instead of contrastive learning. To adapt RPIN, a factored video prediction method, we ap- pend an (x, y, pi ck / /p aaee) action to the input of the Prediction module in their Convolutional

---

Under review as a conference paper at ICLR 2022
Interaction Network. We recreate our dataset in the same format their Shape Stacking experiment (224x224 images with a bounding box for each object) and we use the same
Compared to autoencoding loss (FWM-AE), contrastive learning (FWM) succeeds both in Cubes and Shapes, whereas the autoencoding loss fails to fit Cubes. In Shapes, autoencoding outperforms contrastive loss in terms of block position prediction, but underperforms in action sequence ranking. 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.
Our adaptation of RPIN to action conditoondi sequence prediction fails to learn a dynamics model. 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. However, the forward model fails, reaching 20 em RMSE.
6 reports block prediction error for noisy trajectories generated by an expert policy with added randomness, whereas Table reports errors for optimal trajectories that reach the goal state of each of the training (and generalization) tasks. In Shapes (Figure 6 right), we see reversal compared to Table D as contrastive loss out- performs autoencoding by a small margin. This result suggests that the transition model learned by autoencoding is not as accurate for blocks falling and colliding with each other.
We further study the block prediction error in Figure 6
Finally, we plot action ranking Hits @1 function of noise level for Shapes in Figure left. Contrastive loss outperforms autoencoding for all noise levels with the area under the curve being 84% for FWM and 64% for FWM-AE.
3.2 ZERO-SHOT GENERALIZING TO NOVEL TASKS
Setup. We further evaluate models trained in Section 3.2 on tasks unseen during training (Figure 3 We specifically make sure that the training datasets do not contain a single example of goal state of the novel tasks. In Cubes, the training task is building 4-stack and the novel tasks are building wall, stairs, a wall on ground and stairs on ground. We found that the training data collection policy never solves the novel tasks (we use a separate policy collect evaluation trajectories for testing tasks) In Shapes, we train on half of the possible structures with a height of three and roof on top, and test generalization tasks solves one of the testing tasks around once every 50 episodes (due to added randomness) and we delete these episodes. We report block position error and action ranking Hits@/ described in Section 3.1 We do not perform any fine-tuning on the novel tasks in this experiment.
the other half. In this case, the data collection policy for the eight training
Result. In both Cubes and Shapes, FWM transfers to the novel tasks with only a minor decrease in performance (block position eтror increases by 0.4 cm in Cubes and 0.2 cm in Shapes and ac- tion sequence ranking Hits remain almost unchanged). We consider this to be an important demonstration of the generalization properties of permutation equivariant models. In contrast, only limited generalization and transfer properties have been shown in prior work (Li et al.. 2020 [Bzaa et al. 2021). FWM-AE suffers large decrease in action sequence ranking while still outperforming FWM on block position prediction in Shapes. We plot action ranking error as a function of e for Shapes in Figure 5 right; see a pattern analogous results for training tasks (Section 3.1 The areas under curve are 81% % for FWM and 48% for FWM-AE.
3.3 FWM ABLATIONS AND ATTENTION VISUALIZATION
We report the results of the following ablations in Table I removing action attention (No Attention and 1 GNN, No Att.), using only one graph neural network layer (I GNN Layer), only showing bounding box coordinates as inputs (No RGB), only showing RGB images as inputs (No Coordi- nates) and using a monolithic latent transition model (No Factorization).
We did not find a significant difference when using FWM with and without action attention. As shown in Figure / first row and third row, FWM learns interpretable attention weights During pick action, the attention weights reflect the object being picked up as well as any object that is nearby. The attention module has interesting behavior during place actions: it gives weight of
7

---

Under review as a conference paper at ICLR 2022
zero the object being placed, but attends to all objects that underneath or next to the place location. 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 these three blocks with the highest weight assigned to the red blocks, which is directly below the object being placed.
Why is there no quantitative increase in performance even though the attention module functions properly? 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. We assume that the reduced capacity of the model will require it to use the attention module. Here, action attention leads to large quantitative improvements. E.g. not using attention more than doubles block prediction error in Cubes. Furthermore, qualitative analysis shows that the attention module behaves differently with GNN layer compared to 4 GNN layers. 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.
A second interesting result our ablation studies is the comparison of FWM with and without RGB image. The No RGB baseline only has access to the coordinates of the bounding boxes for each object. We see no statistically significant different between these two settings. Since the RGB bascline 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. We believe more complex object interaction (e.g- packing bin with household objects, building from LEGO) would put higher emphasis on images. We leave these experiments for future work.
No Factorization and No Coordinates ablations show that both of these components are vital the Factored World Model. Without coordinate grids, the model perceives images of objects without any information about where they are located in the environment. Without object factorization, the monolithic latent transition model is unable to generalize beyond the specific arrangements of blocks seen in the training tasks. Here, we mean in midisttiibttinn generalization to building known structures in novel locations of the workspace. Naturally, the monolithic model also fails in zero-shot transfer. Finally, we find stacking of four residual GNNs leads large improvements compared to a single GNN layer. In Shapes, the action ranking Hits@ increase by 19% with four GNN layers.
3.4 TRANSFER TO WORLD DATASET
Setup. We collect evaluation trajectories on a physical robot (Figure 8 for the five 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. We do not train our model on the real-world dataset, since we only collect 20 episodes for each task. Instead, we change the background and object colors in simulation to match the real-world images. 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. 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. in the real world without any additional fine-tuning.
We compare results from simulation and the real-world in Table On average, the block position prediction error increases by 0.4 cm as we transfer from simulation to the real-world dataset. Our model appears to struggle with tasks that involve stacking blocks on top of each other in the real world compared simulation. Conversely, there no decrease in performance for Wall Ground (WG) and Stairs Ground (WG), where all blocks are placed on the ground.
Result.
4 RELATED WORK
World Models with Object Factorization. Action-conditionel ofjecc-factorized world models have been explored in toy object environments and Atari games (Kipf et al.) 2020; Huang et al.) 2020) well as in robotic manipulation /Veerapaneni et al. 2019) and controlling billiard balls Kossen et al. 2020). We build upon C-SWM (Kipf et al. 22020). 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). Veerapaneni et al. (2019) used a factored transition model to learn and plan for stacking up to three cubes. Different from these works consider simplified setup wherein
8

---

Under review as a conference paper at ICLR 2022
Training
Evaluation
RMSE (cm)
4-Stack
Wall
Stairs
WG
SG
Avg
Sim 4-Stack Sim2Real 4-Stack
Sim All
1.46
1.76
2.14
2. 2.
1.90
1.88
Real All
2.33
2.97
2.40
1.72
1.81+15
2.25 +0.
Table 2: Block position predictions results for real-world evaluation dataset. The first row is breakdown of results previously reported in Table I (FWM). The second row reports results for model trained in simulation and evaluated on trajectories captured on physical robot. 4-Stack, Wall, Stairs, Wall Ground (WG) and Stairs Ground (SG) correspond to the tasks shown in Figure
our model does not perform symmetry breaking object discovery from raw visual features, since we assume our environment be already factored.
Object-factored world models are also commonly used in the context of physics prediction in videos.
Janner et al. (2019); Veerapaneni et al. (2019) demonstrated generalization in a Tetris-like environ- ment, where a world model predicts the outcome of dropping blocks from a height. Ye et al. (2019) studied learning dynamics of pushing objects with a robotic arm in the real world. RPIN (Qi et al,
2021) demonstrated state-of-the-arr results in the tasks of 2D physics modeling (PHYRE bench- mark, (akktin et al. (2019)) modeling billiard balls and modeling videos of falling stacks of blocks. As showed in our comparison, models like RPIN cannot be naively applied our pick-and-place tasks, both due to the problem of including action information at the right point in a transition model and due 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 robotic hand).
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). In fact, factored MDP works
often include examples where individual factors are objects (e.g. Tower of Hanoi in Ravindran (2004) ) Early works specifically on object factorization include Wolfe (2006): Diik et al. (2008).
Model-based Learning for Robotics. One prominent line of work focuses on optimal control or policy search in low-dimensional state space. If such space is directly available, PILCO fits a Gaussian Process to real-world robot dynamics and performs policy search through gradient descent in the model (Deisenroth & Rasmussen 201T). Successful cart-pole swing and robotic unicycle balance is achieved within few seconds of online learning. In the case of image states, low- dimensional latent space can be learned by variational autoencoders (see Lesort et al. (2018) for a survey). Watter et al. (2015) learn latent state space together with locally linear transition model, enabling efficient inference in optimal control methods. Watter et al. (2015) demonstrate successful control of pendulum, cart-pole and simplified 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
2017 Banijamali et al. 2018 Levine et al. 2020).
Other works have used models of the world to accelerate learning for robotic manipulation tasks. In Nair et al. (2018), a VAE, trained on image observations, used for sampling goals during training and providing a reward signal via distance in the latent space. This method achieved similar sample complexity to a state-based method on robotic pick-and-place task. In Nasiriany et al.] (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 learned policy.
5 CONCLUSION
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. The resulting models are able to learn transition dynamics that generalize to previously unseen sequences of actions and configurations of objects at test time, and are able to predict outcomes at moderate time horizons of up to 10 actions. This represents significant step in applying these models in practical robotics tasks. An immediate future line work is to integrate these world models with planning algorithms, and evaluate the ability of the resulting planner generalize to tasks that were unseen during training.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
We will open-source our code with the camera-ready version of our paper. We will include our simulated PyBullet environments, dataset generation scripts, Pytorch model definitions and training scripts.
REFERENCES
Lei Jimmy Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Abs/1607.064555 2016.
Layer normalization. CoRR,
Anton Bakhtin, Laurens van der Maaten, Justin Johnson, Laura Gustafson, and Ross B. Girshick. PHYRE: A new benchmark for physical reasoning. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neu- ral Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 5083-5094, 2019.
Ershad Banijamali, Rui Shu, Mohammad Ghavamzadeh, Hung Hai Bui, and Ali Ghodsi. Robust locally-lineas controllable embedding. In Amos Storkey and Fernando Pérez-Cruz (eds.), In- ternational Conference on Artificial Intelligence and Statisties, AISTATS 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pp. 1751-1759. PMLR, 2018.
Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.), Ad- vances Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, pp. 4502-4510, 2016.
Ondrej Biza, Dian Wang, Robert Platt Jr., Jan- Willem van de Meent, and Lawson L. Wong. Action priors for large action spaces in robotics. In Frank Dignum, Alessio Lomuscio and/ Ulle Endriss, and Ann Nowé (eds.), AAMAS '21: 20th International Conference on Autonomous Agents and Multiagem Systems, Virtual Event, United Kingdom, May 3-7, 2021, pp. 205-213. ACM, 2021.
Mare Peter Deisenroth and Carl Edward Rasmussen. PILCO: A model-based and data-efficient approach to policy search. In Lise Getoor and Tobias Scheffer (eds.), Proceedings of the 28th International Conference on Machine Learning, ICML 2011, Bellevue, Washington, USA, June 28 July 2, 2011, pp. 465-472. Omnipress, 2011.
Carlos Diuk, Andre Cohen, and Michael L. Littman. An object-oriented representation for efficient reinforcement learning. In William W. Cohen, Andrew McCallum, and Sam T. Roweis (eds.), Ma- chine Learning, Proceedings of the Twenty-Fijh FFfth International Conference (ICMI. 2008), Helsinki, Finland, June 5-9, 2008, volume 307 of ACM International Conference Proceeding Series, pp- 240-247. ACM, 2008.
M. Gori, G. Monfardini, and F. Scarselli. new model for learning in graph domains. In Pro- ceedings. 2005 IEEE International Joint Conference Neural Networks, 2005., volume 2, pp- 729-734 vol. 2, 2005.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. CoRR, abs/2012.05208, 2020.
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algo- rithms for factored mdps. J. Artif. Intell. Res., 19:399-468, 468, 2003.
Qian Huang, Horace He, Abhay Singh, Yan Zhang, Ser-Nam Lim, and Austin R. Benson. Better representations for relational reasoning, In Advances in Neural Information Processing Sys- tems 33: Anmual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, viriual, 2020.
10

---

Under review as a conference paper at ICLR 2022
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis R. Bach and David M. Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning. ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 148-456. JMI..org. 2015.
Michael Janner, Sergey Levine, William T. Freeman, Joshua B. Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. In 7th International Conference Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
Anders Jonsson and Andrew G. Barto. A causal approach to hierarchical decomposition of factored mdps. In Lue De Raedt and Stefan Wrobel (eds.), Machine Learning, Proceedings of the Twenty- Second International Conference (ICML 2005), Bonn, Germany, August 7-11, 2005, volume 119 of ACM International Conference Proceeding Series, pp. 401-408. ACM, 2005.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Con- ference Track Proceedings. OpenReview.nel, 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pp. 2688- 2697. PMLR, 2018.
Thomas N. Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Jannik Kossen, Karl Stelzner, Marcel Hussing, Claas Voelcker, and Kristian Kersting, Structured object-aware physics prediction for video modeling and planning. In 8th International Con- ference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
Timothée Lesort, Natalia Díaz Rodríguez, Jean-François Goudou, and David Filliat. State represen- tation learning for control: An overview. Neural Networks, 108:379-392, 2018.
Nir Levine, Yinlam Chow, Rui Shu, Ang Li, Mohammad Ghavamzadeh, and Hung Bui. Prediction, Representation learning for locally-lineal control. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020.
consistency, curvature:
Richard Li, Allan Jabri, Trevor Darrell, and Pulkit Agrawal. Towards practical multi-object ma- nipulation using relational reinforcement learning. In 2020 IEEE International Conference on Robotics and Automation, ICRA 2020, Paris, France, May 31 August 31, 2020, pp. 4051-4058. IEEE, 2020.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot tention. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Con- ference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12 2020, virtual, 2020.
Ashvin Nair, Vitchyr Pong. Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. arXiv preprint arXiv:1807.0417 04740442.,4 2018.
Soroush Nasiriany, Vitchyr H Pong. Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. arXiv preprint arXiv: /9/ 108453, 2019.

---

Under review as a conference paper at ICLR 2022
Haozhi Qi, Xiaolong Wang. Deepak Pathak, Yi Ma, and Jitendra Malik. Learning long-tern visual dynamics with region proposal interaction networks. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
Balaraman Ravindran. An Algebraic Approach Abstraction Reinforcement Learning. PhD thesis, 2004. AA13118325.
Alvaro Sanzhez- Nicolas Heess, Jost Tobias Springenberg. Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In International Conference on Machine Learning, pp. 4470-4479, PMLR, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5998-60988 2017.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua B. Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforce- ment learning. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura (eds.). 3rd Annual Conference on Robot Learning. CoRL 2019, Osaka, Japan, October 30 November 1, 2019, Pro- ceedings, volume 100 of Proceedings of Machine Learning Research, pp. 1439-1456 PMLR, 2019.
Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin A. Riedmiller, Embed to control: A locally linear latent dynamics model for control from raw images. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Ganett (eds.). Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Process- ing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pp. 2746-2726 2754, 2015.
Alicia P. Wolfe. Defining object types and options using mdp homomorphisms. 2006.
Alicia P. Wolfe and Andrew G. Barto. Decision tree methods for finding reusable MDP homo- morphisms. In Proceedings, The Twenty-Fing National Conference on Artificial Intelligence and the Eighteenth Innovative Applications of Artificial Intelligence Conference, July 16-20, 2006, Boston, Massachuselts, USA, pp. 530-535. AAAI Press, 2006.
Yufei Ye, Dhiraj Gandhi, Abhinav Gupta, and Shubham Tulsiani. Object-centrie forward modeling for model predictive control. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura (eds.), 3rd Annual Conference Robot Learning, CoRL 2019, Osaka, Japan, October 30 November 1, 2019, Proceedings, volume 100 of Proceedings of Machine Learning Research, pp- 100-109. PMLR, 2019.
Zhong- Qiu Zhao, Peng Zheng, Shou-tao Xu, and Xindong Wu. Object detection with deep learning: A review. IEEE Trans. Neural Networks Learn. Syst., 30(1 30(11):3212-3232 2019.
12

---

Under review as a conference paper at ICLR 2022
RGB Component
Vertical Grid
Horizontal Grid
Cropped RGB:
Cropped Grids:
Concatenated Frame + Bounding Box
Figure 4: Schema of our pipeline for creating cropped object images. We concatenate an RGB image with horizontal and vertical coordinate grids. Using a bounding box with added padding, we crop an 18x 18 image from both the RGB component and the coordinate grid components. By observing the coordinate grids, our agent known where in the image the object was cropped. Note that the coordinate grids are derived from object bounding boxes, the environment. Hence, we do not need to know the ground-truth object positions order to generate factored states. We add additional coordinate grids by mirroring the vertical and horizontal grids (similar to positional encodings in Locatello et al. (020)).
the actual (x, 2) object positions
A EXPERIMENT DETAILS
A. ENVIRONMENT DETAILS
Our PyBullet simulation consists of a URS robotic arm with a Robotiq gripper operating over 30x30 em workspace. In comparison, the size of cube is 3x3 em. The environment is captured from two sides by two cameras. Each camera produces 90x90 image. PyBullet also provides segmentation mask for each object, which we use to create bounding boxes. We first draw the smallest rectangular box that captures the whole segmentation mask and add symmetric padding it. If the bounding box smaller than 18x px, then they are padded up to the minimum size. We crop the RGB image and four coordinate grids inside of the bounding box and resize them to 18x image. The aspect ration of the RGB image gets corrupted in this step, but the model can still capture the size of each object based the coordinate grids. We have vertical and horizontal coordinate grid that traverses the interval of [-1, 1] either left-to-right or up-to-down. We also create two additional coordinate grids that traverse the interval right-to-left and down-to-up.
Block colors. We use single block color (red) for all experiments in simulation. In Figure 7 blocks colored only for the purpose of matching blocks to attention weights. All blocks are red in the
13

---

Under review as a conference paper at ICLR 2022
60
80
Hits@3 (%)
60
Hits@3
40
40
20
Method
Method
FWM
FWM
FWM-AE
FWM-AE
0-
2.0
3.0
3.5
4.0
4.5
5.s
6.0
2.0
3.5
4,0 Epsilon (cm)
4.5
5.0
5.5
Epsilon (cm)
Figure 5: Action sequence ranking Hits@1 conditioned on action noise e. The noise level controls the degree to which negative action sequences differ from the position one (Section (.1). Left: results for training tasks; right: results for zero-shot generalization. We compare Factored World Model against baselines from Table [ Means and 95% confidence intervals over 5 random seeds are reported.
actual states passed into the model. In the real-world experiments, we assign each block a different color that can get bounding boxes using color segmentation (Table We use different colors both in the Sim2Real dataset (generated in simulation) and the Real dataset (generated on physical robot),
A.2 EXPERT POLICIES
Experts for training datasets. For Cubes, we have the following policy: with 70% probability either pick a random cube that is not covered by other objects or place cube top of any other cube. Otherwise, execute pick or place action a random coordinate. Additionally, cm random noise is added to the (x, y) position of each action. The agent does not have to decide between pick and place actions: hand empty + pick, hand full + place.
For Shapes,
trained SDQfD agents provided by Biza al. (001). With 80% probability we execute the actions predicted by the expert. Otherwise, during pick, we pick a random object with 50% and we execute pick random location otherwise; during place, we place in a random position.
We collect 200k training transitions for Cubes and Shapes.
Experts for generalization datasets. We simply collect optimal trajectories for evaluation and filter any trajectory that does not reach the goal of each task.
A.3 MODEL. DETAILS
Per-object Encoder CNN. Conv2D(5x5 kernel size, 32 kernels, stride 2, padding 1) - BatchNorm
(loffe & Szegedy 2015) ReLU Conv2D(5xx x5 kernel size, 64 kernels, stride 2, padding 1) + BatchNorm + ReLU Conv2D(5 kernel size, 64 kernels, stride 2, padding 1).
GNN. Both the node and the edge networks are MLPs with one hidden layer of size 512. Each layer, except the output layer, is followed by LayerNorm (Ba et al. 2016) and ReLU activation.
Action Attention. The key, value and query MLPs have single hidden layer of size 512. the output size is also 512. LayerNorm and ReLU is used in the same way as in the GNN.
Decoder for FWM-AE Baseline. Conv Con/Traaspose22((xx kernel size, 64 kernels, stride 2) BatchNorm - ReLU > Conv2D(3x3 kernel size, 32 kernels, stride 2) > BatchNorm - ReLU + Conv2D(3x3 3 kernel size, 14 kernels, stride 1).
Training. We use the Adam (Kingma & Ba 2015) optimizer with default parameters and learning rate of 5e 5. We train for 200 epochs (dataset size is 200k transitions) with batch size of 256.
14

---

Under review as a conference paper at ICLR 2022
FVN
3.5
FWN
FWN AE
PWN AE
3.0
3.0-
RMSE (tm)
2.5
2.5
RMSE
2.0
1.5
1.-
10-
0.5
0.5
Figure 6: Block position prediction error function of the number of predicted time steps (Left: Cubes, right: Shapes). The validation set used in this figure contains noisy trajectories for training tasks, whereas the errors reported in Table are for optimal trajectories that reach the goal of each task. We compare Factored World Model against baselines from Table [ Means and 95% confidence intervals over random seeds reported.
B ADDITIONAL EXPERIMENTS
B.1 C-SWM WITH HEURISTIC ACTION FACTORIZATION
We train the C-SWM model Kippf al. 2020 in the Cubes environment. Unlike FWM, C-SWM does not receive a factored state space; instead, it can choose what information is captured in each object slot. We make two changes object a distinct color (Figure 9 so that the model can potentially learn color-specific filters and (2) create heuristically factored action space. The factored action space only provides action of to the ith node in the C-SWM graph neural network if the ith object changed between state and s1+1 Otherwise, the ith node receives a null action.
environment help C-SWM factor it: (1) give each
The model receives four images concatenated channel-wise two images of the workspace and two images of the robot hand (which indicate if the robot is holding an object). We use custom encoder with the following architecture: Conv2D(5x) x5 kernel size, 64 kernels, stride 2) - BatchNorm + Lea -> Conx2D(5x5 kernel size, 64 kernels, stride 2) -> BatchNorm LeakyReLU Conv2D(5x5 kernel size, 6 kernels, stride 1) + BatchNorm + ReLU. The output of the encoder is a 16x 16 feature map for each object.
Across range of learning rate, C-SWM does not learn to factor the state space of Cube stacking. By our metrics, C-SWM on-par with an unfactored model (Table [ Figure visualizes the learned feature maps for each object slot: the maps follow ABAB pattern, where the model appears to only distinguish between the robot holding or not holding an object. This pattern holds across episodes.
15

---

Under review as a conference paper at ICLR 2022
Attention weights:
State: t=00
[=1
t=2
[=3
t=4
(=5
[=6
t=0
F1
F2
t=3
5
[=6
Stacking four blocks. First row: FWM with four GNN layers; second
FWM with GNN layer.
Zee
Sem
28
28m
t=2
(=3
[=4
[=6
(=7
1=8
(=9
Zem
Dmm
2Bm
:BC
t=
F-1
t=2
F
F=6
F-
t=8
F-
10
(b) Building wall. First row; FWM with four GNN layers; second
FWM with GNN layer.
Figure 7: Visualizing action attention weights for a sequence of building stack of four blocks and
wall. Each bar in the histogram associated with particular object by color. (a) the agent executes
pick action at f = 1, = 3, and place action at f 2, = 4, t 6. (b) pi ck action at
1, 5, 7, and place action at 2, 4,
8, 10.
16

---

Under review as a conference paper at ICLR 2022
Camera
URS Arm
Camera 2
Workspace
Camera 3
Figure 8: Left: our real-world experimental setup includes a UR5 arm with a Robotiq gripper, two RGB cameras facing the workspace (Camera and 2) and one RGB camera to take image of an object the robot holding (Camera 3). The gripper moves inside of the box the bottom-left corner after every successful pick action. Top-right: an in-hand image taken by Camera 3. Bottom-right. the four images represent a single factored state of the environment with a tower of four blocks. Each image is centered one block starting from the top block going to the bottom. State factorization explained in Figure
17

---

Under review as a conference paper at ICLR 2022
View 1
View
Sett
Site
Slct
5004
Skst
Siet
Slat
View
View
S
Slot
Siet
Sloc
Skst
Sloe
View
View
Skd
Sict
Slot
View
Skot
Siet
Sixc
Slot
Sixt
Figure 9: Visualization feature maps learned by C-SWM with heuristic action factorization. The first two columns show the two views of the environment provided to the model. The next six columns show the 18x 18 feature maps for each object slot given by the C-SWM object extractor. The colormap is scaled between 0 and 0.4.
18


Reviews:
[{"review_id": "lbTGix6l8Eq", "reviewer": "Reviewer_toZ5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\ns1)  The paper proposes to leverage a graph neural network to extract the state feature and learn a latent dynamic model. The contrastive loss is applied to update the neural network weights.\n\ns2) The paper is cleary-written and easy to follow. \n\nWeaknesses:\n\nw1) In Sec 2, \"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 = $ <s_1 , s_2 , ..., s_K >$ in which each si is an image centered on the ith object\". The assumption may not hold in many applications where some objects are fully occluded and can not be perceived.\n\nw2) In the experiments, the evaluation setting is too simple. It would strengthen the paper if more complicated manipulation tasks (including more complex shapes and manipulation tasks pushing, pulling, throwing) are included.", "summary_of_the_paper": "This work proposes a pipeline to learn factored world models to predict the robot actions’ effects. This work proposes to leverage a graph neural network to extract state information and adopt the contrastive loss to train the encoder and latent transition model. This paper applies the proposed approach to pick-and-place tasks with simple shapes, and extensive experiments indicate its effectiveness.", "summary_of_the_review": "This work proposes a pipeline to learn factored world models to predict the robot actions' effects leveraging the graph neural network to encode the latent space and the contrastive loss to update the neural network's weight. However, the technical novelty is limited, and the current setting is too simple. It's unclear whether the proposed pipeline could work with complicated manipulation tasks.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "hWdS--8flt4", "reviewer": "Reviewer_WeZs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors present an interesting approach with several novel ideas. Using an attention mechanism to address the complexity of environments containing larger numbers of objects to focus only those objects that are going to be affected by an action is interesting and novel. Based on the ablation study, however, that this mechanism does seem overly important. Would that change if the number of objects in the environment increases dramatically? The approach performs well in the scenario where the task was shown during training. More importantly, this is also the case for the scenario were the task that was not shown during training and shows the generalization abilities of the approach. The numbers for the two metrics (RMSE and Hits @ 1) are reasonable. The baselines, however, don't seem to be sufficient since (Qi et al. 2021) is either not failing to provide results or scores sufficiently worse compared to the presented approach, and the second baseline is a modified version of the presented approach. Veerapaneni et al. (2019) and (Kipf et al. 2020) use a similar test environment and (Huang et al., 2020) compare their results to (Kipf et al. 2020). The authors neither use those methods in their  evaluation nor do the discuss their reasons for excluding those. It would be great if the authors could provide more details on this part.\nThe submission is well written, the authors are able to communicate their ideas well, but placing the related work section at the end of the submission is confusing. ", "summary_of_the_paper": "The presented work addresses the task of robotic manipulation tailored to environments containing many objects. Task planning in such environments entails large number of possible combinations of actions given the number of objects. To address this challenge, the authors propose to use an attention module in combination with a graph neural network to predict the effects of the manipulation action, limiting the objects that need to be included into the execution planning. Experiments are conducted in simulated environment using a robotic arm interacting with two sets of shapes. The model performs well when for both test scenarios, one where the model operates on a task that is was trained on, and two on a task that was not presented during training showing the generalization capabilities of the presented approach.", "summary_of_the_review": "The presented approach introduces interesting and novel ideas and shows that the model is able to generalize to tasks that it was not trained on. Only using one exterior baseline method that does not perform well seems not sufficient given that three related publications are evaluated using a similar setup. While the experimental section needs to be improved, I would rate the positive aspects the approach, generalization abilities and the novelty higher and would, therefore, argue for the submission to be accepted.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Jq3hvZBbGcm", "reviewer": "Reviewer_iwuh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1) The paper is well written and is easy to understand in a go. \n2) The action-attention module seems to be a novel idea especially if I think about a multi-agent scenario where it'd be useful to have such a module to direct the policy (or an accurate world model learning) and be sample-efficient (as opposed to just concatenating the action's latent to every object's latent state).\n\n3) I appreciate the author(s) adding Figure 5 as it provides a  good visualization of the action-attention module.\n\n\nWeaknesses:\n\n\n4) Action Attention Module: Although I like the idea of having such a module, I am not convinced by the experiments performed to evaluate this. The task of pick and place is simple enough to implicitly figure out which object does an action correspond to -- this can be seen from Table 1, FWM and FWM - No attention. There is barely any improvement in the Hits@1 (and even the block's position if you consider the std. deviation along with the mean). This implies that not having action attention suffices for the network to figure out which object is the action intended on.\n\nThis is not to say that the action attention is useless. It is to say that for this specific task of pick and place, it is not very useful.\n\n5) Claim about combinatorial generalization: The combinatorial generalization is due to the use of GNNs which is a core module of C-SWM [1]. Hence I do not agree with the claim that the zero-shot generalization shown in the paper is a contribution by FWM rather the author(s) just extends the claim made by C-SWM by showcasing it in a different environment (the robotic manipulation environment).\n\n6) Zero-Shot generalization: Based on my points (5) and (6), it is important to make the distinction that the action attention does not contribute to the generalization (at least from the results shown on pick and place task in Table 1 on both Cubes and Shapes environment), rather it is the GNN that contributes to the zero-shot generalization. These two are often worded one after the other which makes the reader think that it is the action-attention that is major contributing factor to the generalization. Kindly make this clear in the text.\n\n7) \"No RGB\" Experiment: There is no discussion about this experiment but in the environments considered by the author(s) all the objects are in red and hence RGB doesn't really play a role in the task. This is also shown clearly in Table 1, where the performance is almost similar (a slight increment/decrease on the RMSE/Hits@1) to FWM.\n\n8) Real-world experiments? - The first line of the Experiments section says \"Our empirical evaluation focuses on pick-and-place robotic manipulation tasks performed in simulation as well as on a real-world robotic arm\", but I do not see any real-world robotic arm experiments. Could the author(s) clarify this?\n\n9) Why only Robotic manipulation? - The action-attention module and (other components from C-SWM) are very generic, so why did the author(s) restrict themselves to only robotic manipulation environments? What about Atari games or other object-based RL environments? It would be interesting to see if action-attention plays an important role there (especially in getting to know the agent in Atari game a.k.a the most controllable agent).\n\n10) In Fig 5, are the objects colored only for visualization purposes or is this the RGB input?\n\n[Minor thing which I have not considered for rating the paper]\n\n11) There are some places where the citations have not been mentioned (rather the literal word \"citations\" appear twice in the text). For example: pg 2(last line) and pg 3 (first paragraph, 5th line). This is being disregarded by me for evaluating the paper since it seems to be a genuine mistake caused by overlooking. I'd request to add the citations in the future version of the manuscript.\n\n---\nReferences\n\n[1] Contrastive Learning of Structured World Models, Thomas Kipf et. al, ICLR 2020", "summary_of_the_paper": "The paper proposes Object-Factored world models, a world model learned using contrastive loss, that tends to generalize over the number of objects in the scene. It also alleviates the assumption that previous world models hold about the object-action association and the author(s) propose an \"action-attention module\" that estimates the probability of an action affecting a particular object (using self-attention). Their experiments mainly focus on robotic manipulation tasks and they show generalization on unseen-test tasks with little drop in performance as compared to the training set.\n", "summary_of_the_review": "Overall I feel that the paper lacks a concrete quantitative evaluation of the action attention module (which is the core contribution of the paper). The experiments don't prove that this module is helpful for robotic manipulation tasks (though it can be useful for some other tasks which the author(s) have to investigate further). Hence, I would like to vote for a weak rejection of the paper at this stage. I will be making a final decision after engaging with the author(s) post rebuttal.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "9r-clhErPyE", "reviewer": "Reviewer_4oxK", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n - The authors demonstrate that factored world models can indeed propagate multi-object structured relation modeling in robotic manipulation scenarios while enabling zero-shotting in unseen object/action configurations. \n - The addition of the attention module as well as stacking multiple GNN layers evidently increases  (although marginally and not always) the performance of the model across both tasks and evaluation metrics\n - The proposed factorization of the state representation is relatively straightforward to implement in the context  of robotic manipulation applications, as actuators are most commonly trained in simulation environments that  come packed with object segmentation tools (e.g. segmentation masks in Pybullet) \n\nWeaknesses:\n - Even though this work seems to be an extension of the factorized world model proposed by Kipf et al. (2020) there are no direct comparisons given with their model in the experiments section.\n - The authors claim they wish to explore the potential of factored world models in generalizing to unseen tasks, but in the concerned environments, new \"tasks\" simply correspond to novel sequences of picking and placing actions (i.e. novel configurations), something previous works that impose permutation equivariance already demonstrate that they can achieve. This work's novelty focuses on the application of an attention mechanism and stacking multiple GNN layers. Both additions provide minor performance boosts (and not robustly across all experiments, as seen in rows 3, 4, and 5 of the bottom part of Table 1). \n - A few grammatical mistakes and forgotten to actually cite works in the first two paragraphs of Section 2.\n- Inconsistent presentation of figures and tables throughout the text.", "summary_of_the_paper": "In this paper, the authors expand on previous works on contrastive learning with factored world models (specifically the C-SWM model by Kipf et al. - 2020) and apply their model for the task of predicting the dynamics of object manipulation with a continuous action space. Such previous works have shown that by leveraging residual graph neural networks (GNNs) to contextualize the learned latent states with possible actions, the model learns representations that are equivariant to permutation of objects in the scene, thus allowing scaling to a high number of objects regardless of the combinatorial explosion of the state space. Unlike previous works, where new objects are discovered by the model from raw pixels, the authors here use the simulation environment to segment objects appearing in the scene and factor them as part of the state representation. Another difference is that instead of factorizing the action space, they use two action primitives for picking and placing objects from given continuous coordinates and propose an action attention module that predicts how each action will affect each object in the next state. They claim that by departing from the standard action-object dependence assumption made by previous works, the current model can effectively model the effect of future actions for structured tasks unseen during training. They finally propose to stack multiple GNN layers instead of a single layer as mostly practiced by earlier works. The authors evaluate their methodology on two pick-and-place tasks regarding building structures of abstract object shapes or cubes. The objective is to correctly predict the effect of certain action sequences in other blocks in order to build a certain structure (e.g. pile). They sample noise from the correct action sequence in order to generate negative examples for contrastive learning. They train in some sequence configurations and test for zero-shot performance in unseen tasks (new structures of objects). They report results for both mean square error of the block\nposition as well action ranking (how well the final state of the action sequence matches the final latent state). ", "summary_of_the_review": "This work presents an adaptation of an existing factored world model specifically tailored for robotic manipulation, using motion primitives (picking/placing objects) as the action space and pre-segmenting objects in order to factorize them into the state representation. The paper's novelty comprises the addition of an attention operation to contextualize future actions with latent object states. Even though experimental results suggest positively towards the potential of this method in task generalization, the lack of extensive comparison with previous factored\n(structured) world models do not convince me 100% of the merits of the proposed methodology, especially when considering the capability of the standard C-SWM model to explore new object categories which are lost in the proposed model in favour of a \"black-box\" object detection system that factorizes the scene as input to the model.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a GNN-based attention mechanism and tests it on a robotic stacking task.

While all the reviewers agree that this work is novel and interesting, they also are unanimous (even after the rebuttal) in pointing to the insufficient experimental evaluation of the proposed method.

I encourage the authors to incorporate the feedback of all the reviewers.
Decision: Reject