Year: 2021

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-andplace 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, with only a minor decrease in model performance. Furthermore, we evaluate our model on a dataset collected on a physical robot.

text: 1 INTRODUCTION
From assembly to household robots, current state-of-the-art robot learning agents cannot generalize beyond a specific training task. One important aspect of generalization is the ability to understand any novel combination of known factors, a so-called combinatorial or compositional generalization. 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. The two key steps such agents need to perform are (a) decomposing a scene into individual objects and (b) modeling relative interactions between objects. The latter is particularly important when learning to predict the dynamics of the environment (i.e. learning a world model).
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 a 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 to predict the state of each object one step into the future. 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. That is, they cannot overfit to a particular ordering of objects. It is a 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 et al., 2018; Sanchez-Gonzalez 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 et al., 2019; Kipf et al., 2020; Huang et al., 2020) or only model sequences of states without actions (Janner et al., 2019; Bakhtin et al., 2019; Qi et 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 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.
We train and evaluate our factored world model in a simulated environment involving a UR5 robotic arm manipulating blocks of various shapes (Figure 1). We instantiate two environments: Cubes includes six cubes arranged into five 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). 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 a very minor drop in accuracy (Table 1). 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. 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 it when generalizing to four different held out tasks (Figure 2).
In summary, we contribute the following:
1. We develop a factored wold model for robotic pick-and-place tasks. 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.
2. We demonstrate that the learned world model accurately predicts the outcomes of sequences of actions. Moreover, the model achieves zero-shot generalization to sequences of actions tasks and object configurations 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.

heading: 2 FACTORED WORLD MODEL FOR ROBOTIC PICK-AND-PLACE
text: The modeling task is to predict the effect of a sequence of actions a1, a2, ..., aT in an environment that is initialized to a state s0. 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 = 〈s1, s2, ..., sK〉 in which each si 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 a 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).
Our goal is to learn a factorized latent representation z = 〈z1, z2, ..., zK〉, along with a world model that predicts the state z′ that results from taking an action a at state z. 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. Our world model is implemented using a graph neural network (Gori et al., 2005) that represents action-conditioned pair-wise interactions of latent factors. It 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 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.
Both the state encoder and the latent transition model are equivariant to the permutation of factors. That is, we can provide objects in any order as long as it 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.

heading: 2.1 FACTORED VISUAL WORLD MODEL WITH ACTION ATTENTION
text: Encoder. Given a factored state s = 〈s1, s2, ..., sk〉, the encoder fφ processes each factor (image of an object) separately. That is, z = 〈fφ(s1), fφ(s2), ..., fφ(sK)〉. The encoder uses three convolutional 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 a fixed image size. As described in Section 3, the per-object image si includes both an RGB component and a coordinate grid component, which marks the location where the crop was taken (Figure 4).
Transition Model. The transition model gθ accepts a factored latent state z = 〈z1, z2, ..., zK〉, zi ∈ RDz and an action a ∈ RDa . It predicts the next latent state by outputting a residual:
ẑt+1 = zt + gθ(z t, at). (1)
We implement gθ(zt, at) as a stack of graph neural network layers with skip connections. Denoting the ith graph neural network layer as GNNi and the ith intermediate representation as yi, the computation with L layers is as follows:
y1 = z + GNN1(z, a) y2 = y1 + GNN2(y1, a) ...
gθ(z, a) = yL−1 + GNNL(yL−1, a) (2)
Our implementation of individual graph neural network layers follows Kipf et al. (2020): the network models pairwise interactions between latent state factors (corresponding to objects in the environment) using a fully-connected node network hn and an edge network he. Both are implemented as MLPs with one hidden layer. The edge network outputs an embedding for each directed edge ei,j = he(zi, zj). The edge embeddings are then aggregated using the node network in order to update the state of each node. Finally, the graph neural network outputs a vector of updated factors z′ = 〈z′1, z′2, ..., z′K〉 with
z′i = hn ( zi, a, ∑ j 6=i ej,i ) . (3)
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 π, we have
GNN(π(z), a) = π(GNN(z, a)). (4)
The permutation equivariance property also holds for a stack of GNN layers as well as the state encoder and action attention.
Action Attention. 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. 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 (Sanchez-Gonzalez et al., 2018; Veerapaneni et al., 2019; Kipf et al., 2020), 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.
We overcome this limitation by introducing an action attention layer. 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. The mechanism follows single-head self-attention (Vaswani et al., 2017) with the following key k, query q and value v:
k = 〈jk(z1), jk(z2), ..., jk(zK)〉, q = jq(a),
v = jv(a). (5)
Here, jk, jq and jv are Multi-Layer Perceptrons with one hidden layer. The factored action space a′ = 〈a′1, a′2, ..., a′K〉 is composed of v weighted by attention weights αi with i ∈ {1, ...,K}:
αi = softmax(kT1 q, k T 2 q, ..., k T Kq)i, a′i = αiv. (6)
Note the new actions are all equal up to a multiplicative factor. Figure 7 visualizes α.

heading: 2.2 LEARNING BY CONTRASTIVE LOSS
text: We train the encoder and latent transition model using a single-step contrastive loss. We use a contrastive loss with a single positive and negative example (Kipf et al., 2020). Given a transition 〈st, at, st+1〉, we encode the current and next state (zt = fφ(st), zt+1 = fφ(st+1), ẑt+1 = zt + gθ(z
t, at)) and make a 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 a negative state (z̄ = fφ(s̄)) up to a margin γ:
L(zt, zt+1, ẑt+1, z̄) = 1
2Kσ2 K∑ i=1 ∥∥∥zt+1i − ẑit+1∥∥∥2 2 + max 0, γ − 12Kσ2 K∑ i=1 ∥∥zti − z̄i∥∥22  . (7)
The negative state s̄ is sampled by randomly permuting a training batch of current states. It is a proxy for sampling a random state from the entire dataset (without needing to increase the batch size). 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.

heading: 3 EXPERIMENTS
text: 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). The environments involve a robotic arm manipulating objects of four shapes (cube, brick, triangle, roof) in a 30×30 cm workspace. The agent controlling the arm uses pre-defined 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.
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 = 〈s1, s2, ..., sK〉 is created by cropping the contents of each bounding box with added padding and resizing the resulting image to an 18×18 square. We add bounding box coordinates into the cropped image in the form of four additional coordinate grids channels (Figure 4, Section A.1). Hence, each view of each object results in an 18×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 (x, y, z) object positions in order to generate the factored state. We concatenate the two views of each object channel-wise, creating a 14-channel image for each object. 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. These two images can be captured by a single camera by rotating the robotic arm in between images. The coordinate grids for hand images are set to zero.
The environment poses a difficult exploration problem: a sequence of random actions is unlikely to create a large structure without knocking it over. A model-free agent with a random exploration policy fails to learn to build structures involving more than two objects (Biza et al. (2021), Table 2, ”DQN RS”). Therefore, we collect a dataset of demonstrations from an experiment with added randomness. We evaluate both the ability of our model to fit tasks with expert trajectories as well as the ability to generalize to a set of held-out tasks without fine-tuning (zero-shot transfer, Figure 3).
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 in simulation to make prediction by trajectories generated by a
physical robot? (Section 3.4).

heading: 3.1 LEARNING ROBOTIC MANIPULATION DYNAMICS
text: Setup. To train the Factored World Model, we collect datasets from the training tasks in 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 as the latent transition model. We also evaluate the model in a 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 decoder to predicts the (x, y, z) 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 current 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 is 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 to predict the (x,y,z) positions of objects.
Action sequence ranking (Hits@1, MRR): We start with an optimal action sequence that achieves a goal state of a particular task. Then, we generate ten other action sequences where each action is perturbed by noise with magnitude . The noise is added to each action in the sequence by drawing ∼ Unif[0, ] and θ ∼ Unif[0, 2π] that are then combined to create a 2D vector in the direction of θ with a length of that is added to the (x,y) coordinate of each action. Each perturbed action sequence must not 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 to the predicted final states of all action sequences. Hits@1 report the fraction of times the model’s prediction of the final state of the correct action sequences was closer to 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, a model making random predictions is not expected to reach Hits@1 of 1/11. In fact, the model needs to be fairly competent to get a non-zero score.
Results. We report results for training tasks in Table 1 top section, Figure 5 left and Figure 6 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 is 3 cm for comparison) and high action ranking score (98% for Cubes and 78% for Shapes). We compare to an 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. (2015)) instead of contrastive learning. 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
Interaction Network. 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. Compared to an 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-conditioned 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 cm RMSE.
We further study the block prediction error in Figure 6. 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. In Shapes (Figure 6 right), we see a reversal compared to Table 1, as contrastive loss outperforms 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.
Finally, we plot action ranking Hits@1 as a function of noise level for Shapes in Figure 5 left. Contrastive loss outperforms autoencoding for all noise levels with the area under the curve being 84% for FWM and 64% for FWM-AE.

heading: 3.2 ZERO-SHOT GENERALIZING TO NOVEL TASKS
text: 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 a goal state of the novel tasks. 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. 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). 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. 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. We report block position error and action ranking Hits@1 described in Section 3.1. We do not perform any fine-tuning on the novel tasks in this experiment.
Result. 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). 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; Biza et al., 2021). FWM-AE suffers a 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 for Shapes in Figure 5 right; we see a pattern analogous to results for training tasks (Section 3.1). The areas under curve are 81% for FWM and 48% for FWM-AE.

heading: 3.3 FWM ABLATIONS AND ATTENTION VISUALIZATION
text: 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 find a significant difference when using FWM with and without action attention. As shown in Figure 7 first row and third row, FWM learns interpretable attention weights. During a pick action, the attention weights reflect 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. 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.
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 1 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 in 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 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. We believe more complex object interaction (e.g. packing a 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 to 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-distribution 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 to large improvements compared to a single GNN layer. In Shapes, the action ranking Hits@1 increase by 19% with four GNN layers.

heading: 3.4 TRANSFER TO A REAL-WORLD DATASET
text: 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.
Result. We compare results from simulation and the real-world in Table 2. 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 to simulation. Conversely, there no decrease in performance for Wall Ground (WG) and Stairs Ground (WG), where all blocks are placed on the ground.

heading: 4 RELATED WORK
text: World Models with Object Factorization. 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). 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). Veerapaneni et al. (2019) used a factored transition model to learn and plan for stacking up to three cubes. Different from these works we consider a simplified 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.
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 environment, 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-art results in the tasks of 2D physics modeling (PHYRE benchmark, Bakhtin et al. (2019)), modeling billiard balls and modeling videos of falling stacks of blocks. 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).
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); Diuk 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, 2011). Successful cart-pole swing and robotic unicycle balance is achieved within a few seconds of online learning. In the case of image states, a lowdimensional latent space can be learned by variational autoencoders (see Lesort et al. (2018) for a survey). Watter et al. (2015) learn a latent state space together with a locally linear transition model, enabling efficient inference in optimal control methods. Watter et al. (2015) demonstrate successful control of pendulum, cart-pole and a 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 al., 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, is 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 a 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 a learned policy.

heading: 5 CONCLUSION
text: 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 a significant step in applying these models in practical robotics tasks. 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.
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.

heading: A EXPERIMENT DETAILS
text: A.1 ENVIRONMENT DETAILS
Our PyBullet simulation consists of a UR5 robotic arm with a Robotiq gripper operating over a 30×30 cm workspace. In comparison, the size of a cube is 3×3 cm. The environment is captured from two sides by two cameras. Each camera produces a 90×90 image. PyBullet also provides a 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 4 px symmetric padding to it. If the bounding box is smaller than 18×18 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 an 18×18 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 on the coordinate grids. We have a vertical and a 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 a single block color (red) for all experiments in simulation. In Figure 7, blocks are colored only for the purpose of matching blocks to attention weights. All blocks are red in the
actual states passed into the model. In the real-world experiments, we assign each block a different color so that we can get bounding boxes using color segmentation (Table 2). We use different colors both in the Sim2Real dataset (generated in simulation) and the Real dataset (generated on a 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 a cube on top of any other cube. Otherwise, execute a pick or place action in a random coordinate. Additionally, 1 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, we use trained SDQfD agents provided by Biza et al. (2021). With 80% probability we execute the actions predicted by the expert. 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.
We collect 200k training transitions for Cubes and Shapes.
Experts for generalization datasets. We simply collect optimal trajectories for evaluation and filter out any trajectory that does not reach the goal of each task.
A.3 MODEL DETAILS
Per-object Encoder CNN. Conv2D(5×5 kernel size, 32 kernels, stride 2, padding 1)→ BatchNorm (Ioffe & Szegedy, 2015)→ ReLU→ Conv2D(5×5 kernel size, 64 kernels, stride 2, padding 1)→ BatchNorm→ ReLU→ Conv2D(5×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 a 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. ConvTranspose2D(5×5 kernel size, 64 kernels, stride 2) → BatchNorm → ReLU → Conv2D(3×3 kernel size, 32 kernels, stride 2) → BatchNorm → ReLU → Conv2D(3×3 kernel size, 14 kernels, stride 1). Training. We use the Adam (Kingma & Ba, 2015) optimizer with default parameters and a learning rate of 5e− 5. We train for 200 epochs (dataset size is 200k transitions) with a batch size of 256.

heading: B ADDITIONAL EXPERIMENTS
text: B.1 C-SWM WITH HEURISTIC ACTION FACTORIZATION
We train the C-SWM model (Kipf et 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 to our environment to help C-SWM factor it: (1) we give each object a distinct color (Figure 9) so that the model can potentially learn color-specific filters and (2) we create a heuristically factored action space. The factored action space only provides action ati to the ith node in the C-SWM graph neural network if the ith object changed between state st and st+1. Otherwise, the ith node receives a null action.
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 a custom encoder with the following architecture: Conv2D(5×5 kernel size, 64 kernels, stride 2) → BatchNorm → LeakyReLU → Conv2D(5×5 kernel size, 64 kernels, stride 2) → BatchNorm → LeakyReLU → Conv2D(5×5 kernel size, 6 kernels, stride 1)→ BatchNorm→ ReLU. The output of the encoder is a 16×16 feature map for each object. Across a range of learning rate, C-SWM does not learn to factor the state space of Cube stacking. By our metrics, C-SWM is on-par with an unfactored model (Table 1). Figure 9 visualizes the learned feature maps for each object slot: the maps follow an ABAB pattern, where the model appears to only distinguish between the robot holding or not holding an object. This pattern holds across episodes.

Review summaries:
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.

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. 

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.

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 position as well action ranking (how well the final state of the action sequence matches the final latent state).