Unsupervised Object-Based Transition Models For 3D Partially Observable Environments
Authors: Antonia Creswell
Rishabh Kabra
Christopher P. Burgess
Murray Shanahan
Venue: NeurIPS 2021 Poster (2021)
Keywords: Objects
transition model
alignment
entities
MONet

Abstract: We present a slot-wise, object-based transition model that decomposes a scene into objects, aligns them (with respect to a slot-wise object memory) to maintain a consistent order across time, and predicts how those objects evolve over successive frames. The model is trained end-to-end without supervision using transition losses at the level of the object-structured representation rather than pixels. Thanks to the introduction of our novel alignment module, the model deals properly with two issues that are not handled satisfactorily by other transition models, namely object persistence and object identity. We show that the combination of an object-level loss and correct object alignment over time enables the model to outperform a state-of-the-art baseline, and allows it to deal well with object occlusion and re-appearance in partially observable environments.

Unsupervised Object-Based Transition Models for 3D Partially Observable Environments
Antonia Creswell
Rishabh Kabra
DeepMind
DeepMind London rkabra@deepiind.com
London
tonicreeseel1@deeppinad
Christopher Burgess
Murray Shanahan
DeepMind
DeepMind
London cgburgasssddeepmind com
London
mshanahan@deepminnd com
Abstract
We present slot-wise, object-based transition model that decomposes a scene into objects, aligns them (with respect to slot-wise object memory) to maintain a consistent order across time, and predicts how those objects evolve over successive frames. The model is trained end-to-end without supervision using transition losses at the level of the object-structured representation rather than pixels. Thanks the introduction of our novel alignment module, the model deals properly with two issues that are not handled satisfactorily by other transition models, namely object persistence and object identity. We show that the combination of an object- level loss and correct object alignment over time enables model to outperform state-of of teffthe-art baseline, and allows it to deal well with object occlusion and re-appearance in partially observable environments.
1
Introduction
In spite of their well-documented ability to learn complex tasks, today's deep reinforcement learning agents are still far from matching humans at out-of-distribution generalisation or few-shot transfer
[6 21 23 Two architectural features commonly proposed to remedy this are (1) transition models
that enable the agent to internally explore paths through state space that it has never experienced (29 25 [2], and (2) compositionally structured representations that enable the agent to represent meaningful states that has never encountered [7]. These two features not exclusive; transition models that operate on compositionally structured representations are a potent combination, and these the subject of the present paper. Specifically, our focus is transition models that operate the level of objects, which are the most obvious candidates for the structural elements of representations likely to be useful for artificial agents inhabiting 3D worlds such
28
current models do not deal satisfactorily with object persistence (the concept that objects typically continue to exist when they are no longer perceptible (24) or with object identity (the concept that a token object at one time-step is the same token object at a later time-step [I]. As show, transition models that neglect object persistence tend to perform badly in complex, partially observable environments, while models that neglect object identity interactions) over time in a way that generalises to future time-steps. By proposing and incorporating novel module for aligning objects across time using slot-based memory, model handles both these concepts, and exhibits better performance as result.
While there has been progress in object-based transition models 34
35
unable to integrate information about single object (and its
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

t=
t=2
t=4
t=10
t=12
t=14
t= 16
t=
188
Imput
Target
Unroll
OAT (ours)
Unroll Unroll
Memory Neig
OP3 (Veerepaneni current
state-of-art model)
Comparing our model to baselines. Each model was trained with four input steps and to unroll for six steps, here we unroll for 15 steps. Our model, OAT, performs significantly better than OT (our model without the AlignNet) and current state-of-the-an model, OP3 [3] See Figure in the Appendix for additional OP3 roll-outs.
Figure 1:
An important feature of our transition model is that makes predictions and computes losses in representation space that is divided into objects. This contrasts with existing models that make predictions in an unstructured representation space [13 12). Making predictions and com- puting losses in an object-structured rep- resentation space facilitates learning, not only because the representation space lower dimensional than pixel space, but also because the model can exploit the fact that dynamics tend to apply objects a whole, which simplifies learning. How- compute prediction losses directly over distinct object representations (rather than first mapping predictions back to pix- els [34] 32 the objects in a predicted rep- resentation must be matched with those in the target representation, which again re- quires a proper treatment of object persis- tence and identity. The result is a model that can roll-out accurate predictions for significantly more steps than seen during training, outperforming sate-of-the-ar for comparable models.
m
z
M++
za
h
h
h+
Transition module
Transition module
LOSS
za 1-1
a
z
a
Alignment module
Alignment module
Z-1 m
m
Scene decomposition and representation module
Scene decomposition and representation module
X
X
Figure 2: Encoding steps. The scene decomposition and representation module extracts K object represen- tations, 4, E RF from images, 2t R x3 These are aligned using the aligument module to obtain of R Aligned objects are fed the slot-wise transition model with the action, at and hidden state, (, h2) E Mx to predict the object representa- tions at the next time-step, of E M well the updated memory, m E XF We also depict the transition model loss, Transition model. computed between z and 2 (shown as LOSS in the figure).
To achieve this, model, Objects- Align- Transition (OAT), combines (1) scene decomposition and representation module, MONet 2]. that transforms a raw image into slot-wise object-based representation, (2) an novel alignment module which, with the aid of a slot-wise memory, ensures that each object is represented in the same slot time, even if has temporarily disap- peared from view, and (3) novel slot- wise transition model that operates on the object representations to predict future states. All three components are differentiable, and the whole model trained end-to-end without supervision. We evaluate the model on two sequential image datasets. The first is collected from pre-trained agent moving simulated 3D environment, where we show that our model outperforms the current state-of-the the-art object-beer bosect-bbaed based transition model (OP3 32 The second dataset is collected from a real

---

robot arm interacting with various physical objects, where we demonstrate accurate roll-outs over significantly longer periods than used in training. An ablation study shows that the model's 's success depends on the combination of correct object alignment across time and the use of transition loss over object-level representations instead of over pixels.
2
Our Model: Objects-Align-Traaniiion
Our model, Objects-Aligni Transition module, in this case MONet, with our slot-wise transition module, via our alignment module, which ensures that each slot the transition model receives objects corresponding to the same token object (or identity) across time. The whole model is trained end-to-end. We now provide details of each of these modules (see Figure
Aliaaaaniinnnnnnnniiinionnn (OAT), combines scene decomposition and representation
2.1 Scene Decomposition and Representation Module
The input to OAT is sequence of RGB images, E [0, 1] with T sime- width, W, and height, H. We leave out the batch dimensions for simplicity. Each image in the sequence, Iz, is passed through a scene decomposition and representation moduld] in this case MONet, to obtain a slot-wise object representation, [21, occupying K slots, where 21, E R is an object representation vector with F features.
MONet consists of an attention module (a U-net H4, E [0, x] #xx1, for each slot and slot-wise VAF. The encoder of the slot-wise VAE is fed the predicted object segmentation masks and the input image, Zt, and outputs object representations, 2 "R The decoder reconstructs the masks, uu,,,. and each object's pixels, Tt,k. MONet's computations can be summarised by the following equations: URRt = NitentionNNttoor((t:) Zkett MONet_encoder(z;; (l.,)) and ji,.,, Tt
which predicts object segmentation masks,
(eONee__ccode((1)) (z).
may be tempting to feed the wise object representations, [= k= at time, t, directly to slot-wise transition model. However, MONet's representations [2] are not stable across time, meaning that specific object may appear in different slots different times [35] (see Figure (3). This leads to two major problems when training slot-wise transition models on slot-wise object-based representations. Firstly, difficult to compute losses between predicted and target object slots. Computing slot-wise object-level loss requires us to know how the previous object representations (and thus the predicted object representations) correspond with the target object representations. Secondly, if object representations do not occupy consistent slots it makes it harder to integrate information about a single object (and its interactions with other objects) across time (35), and makes it harder to predict the reappearance of that object (as show in Section 4.2 To address the issue of slot stability we introduce Memory AlignNet, an alignment module.
2.2 Alignment Module
Our alignment module, Memory AlignNet, must play two key roles. The first is align objects across time, enabling us to compute slot-wise object-level losses for training. The second is to learn a slot-wise memory that encodes the history of each slot across time, allowing our transition model to operate effectively in partially observable environments by keeping track of objects as they go in and out of view. This is especially important for embodied agents that take actions in 3D environments, where the agent's 's looking around frequently causes objects to move in and out of its field-of-view.
We propose Memory AlignNet, which takes slot-wise memory, m. E X with M > K slots and the (stacked) output of the scene decomposition and representation module, 2, E R R returning the aligned object representations, 2 MxF. To perform alignment, the Memory that aligns objects, 2t, with the memory, mi- This adjacency matrix allows to compute soft alignment 20.t = A, hard alignment, 2a A za,hard Hungarian( (A:) We will refer to 2 a.herd as z throughout the rest of the paper. Hungarian(-) denotes the application of the Hungarian algorithm, a non-different non-differentiable algorithm which computes permutation matrix given an adjacency matrix. The soft version of the alignment
AlignNet predicts an adjacency matrix, A₁ E R MM
We tried some variants where depended on 2 However these did not yield significantly beneficial results.
3

---

Alignment module losses (Memory AlignNet). The Memory AlignNet trained on a recon- struction loss between the softly aligned object representations, 2, 20,soft and the output of the transition model, z for the corresponding time-step. There are also regularisation losses includ- used to train the Memory AlignNet while the hard version of the alignment is passed to the slot-wise transition model. which has the same number of slots, M, the memory. Further details can be found in Section A.2 of the Appendix.
The hard-aligned objects may also be used update the memory using recurrent slot-wise model. For simplicity we use our transition model, (-).). to predict deltas for both the object representations, zf, and the memory, my. We will detail this in the next section.
2.3 Slot-Wise Transition Module
The transition model operates in both an encoding and an unrolling phase. During the encoding phase the transition model is fed aligned, observed object representations, 24, and actions, dg, to predict aligned object representations at the next time-step, =t+1 E RMMX.. During the unroll phase, the transition model is fed the predictions, 2f, and actions, de, from the current time-step predict the object representations, z[+1 at the next time-step.
More concretely in the encoding steps, our transition model, To()), takes aligned object representa- tions, of, and a hidden state, h1 E 3M** and predicts deltas for both the object representations, 24 and memory, my, The output of our transition model given by [A] Ar"), hi+1 = To at hf). The memory and the object representations are then updated as follows: M++1 = m. + A," and 2d
2141 = 2f + A, In the unroll steps, next step predictions are given by 2 4++1 = 2f + A where [A Ar",, h = (=( ag, h.). By using the transition model to predict deltas, 2 is aligned by default and does not need be re-aligned when unrolling the model.
A key feature of our transition model is that weights are shared across object slots and can be instantiated in many different ways. One simple option is independently to each slot, sharing weights between slots. An alternative, novel instantiation which found work well is first apply a transformer and Section[A.3 This allows the model to capture interactions with other objects (via the transformer) and integrate that information over time (via the SlotLSTMM, Further details are in Section A.2 of the Appendix.
a SlorLSTM; LSTM applied
KI and then the SlotLSTM (see Figure TQ
2.4 Training
When training OAT, we jointly learn the parameters of the scene decomposition and representation module, the aligmment module and the transition module. Each module has its own losses that contribute to downstream gradients and updates. For example, both the transition model and the aligmment module losses can influence the scene decomposition and representation module's 's updates. Let us define the losses for each module.
Scene decomposition and representation module losses (MONet). We train MONet using stan- dard MONet losses, |MONet. see Equation 3 in [2]. The first term in the MONet loss is scene reconstruction term. This is a spatial mixture loss parameterised by dbg and og which are the standard deviation used for each slot's component likelihood (for the first slot and remaining slots, respectively) that go into the mixture loss. The remaining terms are regularisation terms that 1) induce latent information bottleneck needed for good representation learning (scaled by B, the latent KL loss weight), and 2) ensure that predicted and target masks are similar (scaled by y, the mask KL loss weight). MONet losses computed only in the encoding phase and not used in the unroll phase. Additionally, MONet loss gradients do not affect the transition module or the alignment module weights.
ing an entropy loss, [H!()), on the adjacency matrix, A., which encourages values towards zero and one; and a loss that penalises columns that sum to more that one, avoiding the case where multiple objects are assigned to the same memory slot. The AlignNet loss, (AligaNer is defined as VeiggnNer = Ei' || - =i =g. + +H (A,) + Σド' max(0, (ΣK Ar,. 1)), where T is the total number of encoding and unroll steps. We use t/ = 0.01 for experiments presented this paper.

---

Slot-wise transition module losses. The slot-wise transition model is trained with reconstruction loss between the outputs of the transition model, 2 and the aligned observations for the same time-step,
z.. The slot-wise transition model loss, Transition model is defined as Transition model
Ez || =f 1|. Importantly, compute loss directly between object representations without the need for decoding them to obtain pixels.
OAT is trained end-to-end to minimise IMONet + AAlignNet + Çl Transition model. using = 10 for all experiments presented in this paper.
3 Related Work
Our work addresses the challenging topic of learning abject-based transition models in 3D partially observable environments without supervision. We identify two phenomena that are not dealt with adequately in current object-based transition models, namely object persistence and identity. As consequence, current models are not able to perform well in partially observable environments, and are often trained without semantically meaningful losses.
Firstly, we acknowledge previous work on temporally extended scene decomposition and representa- tion models [35 [4] (typically used for video representation learning), object-based transition models [20]7 and action conditional object-based transition models such as OP3 (32) and C-SWM (1)]. We consider Veerapaneni al.'s 's OP3 32
be the current state-of-the-art object-based transition model since C-SWM requires privileged information about exactly which object an action was applied to, while OP3 (like our model) simply requires the action taken by the agent. OP3 alternates between refinement steps and dynamics (or prediction) steps. One drawback of OP3 is that its refinement steps require access to the observation, which means refinement cannot be applied when rolling out the model to unobserved time-steps. For this reason Veerapaneni et al. [22 only demonstrate unrolls for a limited number of steps.
3.1 Current Models are not Designed for Partially Observable Environments
The world that we operate in, and which we intend our agents to operate in, is partially observable. A severe limitation of existing object-based transition models [ [9] and temporally extended object scene decomposition and representation models [T) 14 is their inability to cope with partially observable environments. A promising approach proposed by He et al. [14] uses an external memory for object tracking (but does not use transition model). Their mechanism is different to that of our proposed alignment module, AlignNet, which we use to perform slot alignment in OAT. He al. 14 train their model, TBA, for reconstruction while AlignNet incorporates dynamics and is trained using prediction, meaning that AlignNet can resolve ambiguities using dynamic while TBA cannot. Additionally, TBA can only cope with static backgrounds and so it is not applicable here. Goyal et al. (8] also develop a general model with a notion of object persistence. However they only show results in 2D environments.
3.2 Current Models are not Trained Using Semantically Meaningful Losses
Another problem with current approaches to learning object- transition models is how they are trained. In most scenes, the background accounts for most of the pixels while the objects account for only small fraction them. So while it may tempting to train transition and video--rpresenntation models using a pixel-level loss [4 K5
32
20
17
[4], we show that it is preferable to compute
losses directly between predicted and target objects (see Section and Figure T.. Furthermore,
pixel-level losses require the object representations to be decoded into an image 34 35 is often computationally expensive.
32
which
An alternative way to find object-level losses is to compute a minimum assignment loss 4 employing the Hungarian algorithm, between predicted and target object representations. However, doing so can be problematic because you need to first compute a similarity matrix on which to apply the Hungarian algorithm. We show in Section 4.2 that computing a loss using the Hungarian algorithm (using L2
For example, when two visually similar objects collide with one another, the AlignNet can use dynamics to resolve which object is which after the collision.
5

---

1
Figure 3: Outputs of our scene decomposition and representation (MONet) module across time. MONet predicts K 10 latent object representation vectors {=r,)i= each time-step. Here we visualise those vectors, in columns CI to C10, using MONet's decoder. Notice that objects switch slots across time. This makes it difficult (a) compute losses between predictions and targets and (b) integrate information about an object across time.
loss between all object pairs as the similarity matrix) leads to poor generalisation in transition models when performing longer unrolls than those seen during training.
Furthermore, at the start of training, comparing predictions with targets to compute a loss may not be meaningful, since predictions will start off essentially random. our approach, align the inputs with the targets, and condition the slot-wise prediction on the slot-wise history. This reinforces slot stability and allows us accurately predict changes, Ai, in each object representation.
Interestingly, Lowe et al. 22 use Deep Sets [36 to encode predicted and target object sets and use contrastive loss between encodings. Their approach has only been demonstrated for very simple datasets. Kipf et al. (C-SWM) [19] do compute losses their representation space. However, they avoid the object correspondence problem because they extract representations spatially, keeping the spatial ordering and using data where object movement was limited. Their approach is unlikely to scale well to partially observable environments with significant movement of objects across the field of view, especially when being trained to predict multiple steps into the future.
OP3 32 and other models 35 attempt induce weak, implicit object alignment by conditioning predicted object representations on those from the previous time-step. This does not guarantee alignment (or slots with consistent identity) over time, especially in partially observable environments. Moreover, in contrast to OAT, they do not use their implicit alignment for computing losses.
Finally, we note that while there are existing models that learn to predict future states, given actions, directly from pixels [13] 12]. we have focused primarily on related work that predicts the future states of objects, because we are particularly interested in developing models that may support future work towards object-based based agents.
Experiments and Results
In this section we (1) demonstrate OAT's performance in the 3D Playroom environment and compare the current state-of-the-ar object-based transition model, OP3 (32) (Section 4.1
[
16
investigate the benefit of alignment and object-level losses when training transition models (Section
4.2). and (3) apply OAT in real world robotics environment and show it accurately predicts both the motion of the robotic arm and its physical interaction with objects (Section
4.1 Objects-Align- Transition Results in Playroom
We train and test our model, OAT, on data collected by an agent taking actions a Playroom environment [6]. Each procedurally generated room in the dataset contains between 10 and 45 objects from 34 different classes 10 different colours and three different sizes. dataset of observation-action trajectories, (Iz, o,) is generated by agent taking actions according to learned policy in procedurally generated room. We collect 100, 000 trajectories with a 7:2:1 train:validation.test split. The top row of Figure 5/sows example trajectory.
We train OAT with four encoding steps (i.e. the model sees four observations for the first four time-steps) and six unrolling steps. MONet outputs K object representations, (e..h. We

---

slot
slot
slot 2
slot
slot 4
slot
slot
slot 7
slot
slot
slot 10 slot
Σ
Figure 4: Aligned inputs (first time-steps) and targets (last 6 time-steps). Our alignment module outputs object representation vectors which visualise here using MONet's entity decoder. Most of the objects are now in consistent slots across time, making easier to compute losses directly between object representations using simple L2 loss. Notice that while MONet outputs 10 slots, the AlignNet has 12 slots. Additional time-steps shown in Figure 11 in the Appendix.
K 10 objects, with F 32 features, and a memory with M slots. Figure visualises the MONet outputs for the first three time steps. Our model achieves good segmentation (see the Adjusted Rand Index, ARI, in Table [ segmentation metrics). Importantly, notice that the slots are not stable across time. For example, the purple object in slot C7 at t = 0 switches slots at t = 2. 0 switches slot C2 at 1 and slot C4
Similarly the yellow object in slot C5
2.
The object representations, 2, R K F), output by MONet are fed to the alignment module. The outputs of the alignment module, 2 E 3R(Mx F), are visualised in Figure 4 The outputs of the aligmment module object representations, and Figure His visualisation these representations using MONet's decoder. Our alignment module successfully keeps objects in consistent slots across time.
The output the alignment module is used to generate both inputs and targets for training the transition model. The transition model predicts latents, 2f E xF] In Figure 5 we visualise roll-outs from our model using MONet's decoder; the reconstructed scene visualisations as the mask weighted sum of the each slot's pixels, Σx (, While the model is trained to unroll for six steps, here unroll for 15 steps. The transition model only sees the first four frames. In the top example, we see that the model is able to predict the appearance of the avatar well, and in general notice that the model, given the agent's actions, able to predict the position of objects well and without the representations degrading. (For comparison to baselines see Figure T some of the examples, the targets appear to have more objects than those seen in the unroll. This because the model has only seen the first four frames and has not seen those other objects in the room, and therefore does not have enough information predict where unseen objects will appear.
generated
Ianget
Unai
target
Ingget
Figure 5: Unrolling OAT for more time-steps than seen during training. The slot-wise transition model makes predictions in latent space (one latent per object), which we visualise using MONet's decoder. The slot-wise transition model was trained with 4 inputs steps and to unroll for 6 steps. Here the model takes the first 4 steps as input and is unrolling for 15. On the top row at f = 16, it is impressive that the model able predict the appearance of the avatar's base (the red circle with the yellow stripe) when looking down. (For per-object visualisation see Figure 13 in the Appendix.)
7

---

To conclude these experiments, we compare our model to the current state-of-the-an object- object-based transition model, OP3 [32 and our ablated model, OT, without the alignment module in Table and Figure [ Qualitatively, from Figure ① we see that OAT significantly outperforms both baselines. Unrolls using OT, trained without alignment, lead to objects merging towards grey cloud in the middle of the frame. Unrolls from OP3 lead to objects fading into the background (we explore the cause of each of these pathologies in the next Section 4.2). Results were consistent for each model across multiple runs.
To quantitatively compare models we consider three metries: Encoding ARI, Unroll Pixel Error and Unroll ARI. The Encoding ARI (Adjusted Rand Index. [18 26 measures the accuracy of the object segmentation masks learned by the scene decomposition model. The Unroll Pixel Error and Unroll ARI evaluate the quality of the transition model's 's unrolls. To compute each of these we decode the latents, z predicted by unrolling the transition model, to produce images, æy and masks, Mi.. The Unroll Pixel Error is the mean-squared-emvee between the ground truth images and xy for the unroll steps only. The Unroll ART is the accuracy of the decoded masks, 1& compared with the ground truth masks. For both ARI scores we exclude background pixels from the score since accurate decomposition of the objects is the main concern here. The Unroll ARI is more meaningful evaluation of the unrolls than the Unroll Pixel Error since is not affected by the background which often accounts for most of the pixels. Note that we only use ground truth masks for evaluation purposes.
The results in Table furthhr demonstrate the crucial role that alignment (AlignNet) plays when performing unrolls. Without alignment the Unroll ART is significantly lower because we are not able to compute a semantically meaningful object-level loss for training. Additionally, we see that our model significantly outperforms OP3 on all metrics (see Figure In light of the results in Table I in the next section we will more concretely look at the role of alignment and the object-level loss (between object representations) when training transition models.
Encoding ARI
Unroll Pixel Error
Unroll ARI
OAT (ours)
0.62
0.0121
0.42
OT (ours, no AlignNet)
0.59
0.0143
0.12
OP3
0.32
0,0132
0.33
Table 1: Comparing
model OAT to baselines. We trained three OAT and OT models (three seeds each), and report the ARI score, Unroll Pixel Error and Unroll ARI score for the model with the best Unroll ARI score. Results for OP3 were obtained similarly, except that ten models were trained (ten seeds), since more variance was observed in the OP3 results (see Appendix D for details).
4.2 What Matters in Object-Centric Transition Models?
In this section demonstrate the need for both (1) alignment, which ensures that each slot in the transition model receives the same object consistently across time, and (2) object-level loss, computed between predicted and target object representations, rather than a pixel-level loss. We also compare different transition model cores and find a transformer followed by a slot-wise LSTM to be best.
In this section (Section 4.2 only, we use a MONet model that is trained using ground-truth masks instead of learning its own masks. We do this here directly analyse the benefits of aligned object representations and object-level losses for the transition model without other confounds. Using ground-truth masks allows us know each object's true identity across time and directly measure the role of alignment in transition models without confounding errors from AlignNet or MONet's segmentation quality. We use a pre-trained MONet (with fixed weights) to compute object representations to keep the model similar to the full end-to-end model described in the rest of the paper.
We train slot-wise transition models under four different conditions: feeding aligned or unaligned object representations (and target) object representations the transition mode and training the
'For the unaligned inputs. we shuffle the order of the ground-truth masks E W,,,5, that are fed to MONet along the K axis. For aligned inputs, ensure that each slot contains consistent object across time.
8

---

t=10
1=11
t=27
t=30
t=31
[=40
t=48
Targets
Aligned inputs, object-level loss
Aligned inputs, pixel-level loss
Unaligned inputs, object-level loss
Unaligned inputs, pixel-level loss
Figure 6: Comparing slot-wise transition models trained with aligned vs. unaligned input latents and pixel-level vs. object-level loss. These models were trained with 8 inputs steps and 12 unroll steps. Here they are being unrolled for 40 time-stepy The output of the transition model is slot-wise object-based representation for each time-step. Here, we visualise the object representation vectors using MONet's is decoder. Only models trained with aligned inputs were able to predict the reappearance of the chair at = 30. Only the model trained with aligned inputs and object-level loss is able to predict the appearance of the avatars "base". In models trained with pixel-level loss predictions becomes very blurred.
Reporting object-level error for models trained using:
Object-level loss
Pixel-level loss
Aligned inputs
0.0929 1 0.00469
14.4 18.4
Unaligned inputs
> 20.3 + 0.757
3312 + 4050
Table 2:
The role of alignment and object-level loss when training transition models. This table reports the std latent error between predicted and ground truth latents (for 5 seeds). For models trained using unaligned input latents we compute the object-level loss using the Hungarian algorithm which is a lower bound the actual value. Models trained with pixel-level loss become very unstable resulting in high variance between runs. Individual seeds are shown in Figure
model with pixel-level or object-level loss. Results in Table 2 clearly demonstrate the benefits of (1) using aligned object representations for training transition models and (2) training transition models with object-level loss. These results are critical for the future development of object-based transition models and, in particular, demonstrate the need for alignment. Note that for the experiments using the unaligned inputs computed the object-level loss using the Hungarian algorithm which is a lower bound estimate of the true object-level loss (since the minimum assignment in L2 may not be the correct assignment).
Figure coompprrs models trained under the four conditions listed above. We see that models trained using aligned inputs (and targets) are able to predict the re-appearance of objects while those without aligned latents are not (Figure 6 t = 30). We also see that models trained with an object-level loss and without alignment deteriorate quickly. Additionally, for models trained with pixel-level loss notice a "ghosting" effect where objects fade into the background across time, similar to the effects seen in the OP3 results (Figure I7p. Pixel-level loss leads to this "ghosting" effect because most pixels in the observation are background pixels and so the background pixels dominate the loss.
We found these results to be consistent across multiple runs (see Figure 9 in the Appendix) and different choices for slot-wise transition model architectures. For the results shown this section we used a transformer with a slot-wise LSTM. Figure 10 in the Appendix, compares object and pixel errors for different transition module cores. Training with object-level loss and aligned inputs, the transformer with slot LSTM worked best.

---

t=0
t=2
t=3
t=4
t=10
t=17
t=33
t=44
t=54
t=71
t=103
Target
Uhroll
Figure 7: Unrolling for many more steps than seen during training, OAT is trained to take four input steps and unroll for six steest Here we demonstrate the model unrolling for 100 time-steps. Impressively, we see that the model learns both to predict the behaviour of the robot arm accurately, as well as how the arm interacts with the objects. We also see that the model able to unroll for significantly more steps than seen during training. Note that our model makes predictions latent space; here we are visualising those latents using MONet's 's decoder. Additional results in Figure 15 in the Appendix.
4.3 Application to Robotics
Here we demonstrate the application of OAT to a real world dataset B of robot trajectories. These trajectories involve robot arm interacting with three objects of varying shapes and colours. The dataset is particularly challenging because our model must learn to predict not just the motion of the robot arm, given the actions, but also how the arm interacts with objects requiring some understanding of intuitive physics.
We train OAT with four input steps and to unroll for six steps. Our model achieves excellent segmentation results, shown in Figure 14 of the Appendix. Figure [ shows impressive results obtained when unrolling the OAT model for significantly more steps than seen during training. What more, in the third sample from the top. we see that the model correctly predicts the reappearance of the red object after it had been fully occluded at t = 3 for 51 time-steps, reappearing fully at t 54 in both the prediction and the target frame. Our model is able to accurately predict the reappearance of objects, even after long-term occlusion, because explicitly captures the history of each object, endowing the model with a notion of object persistence.
Conclusion
We presented Dojects-Align-- Transition (OAT), an object-centric transition model that combines scene decomposition and object representation module (MONet) with slot-wise transition module, via an alignment module. The alignment module plays two key roles. Firstly, it ensures that the slot-wise transition model receives slot-consistent object representations across time. Secondly, allows us to compute object-level loss rather than a pixel-level loss which is commonly used when training transition models.
In an ablation study, we demonstrated the essential role that alignment and object-level losses play when training transition models. Additionally, we significantly outperform existing state-of-the-an object-centriy transition models in 3D partially observable environment, and we applied our model robotics dataset, predicting many steps further into the future than seen during training.
There is room to improve our model in further work, for example by making stochastic predictions about the future, and by better modelling the uncertainty about the objects in the environment. Meanwhile, our work paves the way for future object-centric agent research, for example, enabling agents to plan over future trajectories in object representation space.
Acknowledgments and Disclosure of Funding
We would like to thank Adam Kosiorek for feedback on early drafts of the paper as well as Arunkumar Byravan, Dushyant Rao, Yusuf Aytar and Markus Wulfmeier for their help running experiments on robotics data. This work was carried out within DeepMind, who are wholly owned by Alphabet, the parent company of Google.
10

---

References
[1] R. Baillargeon, E. S. Spelke, and S. Wasserman. Object permanence five-month-old infants. Cognition, 20(3):191-208 1985.
[2] C. P. Burgess, L. Matthey, N. Watters, R. Kabra, Higgins, M. Botvinick, and A. Ler- chner. MONet: Unsupervised scene decomposition and representation. arXiv preprint arXiv: 1390, 2019.
[3] S. Cabi, S. Gómez Colmenarejo, A. Novikov, K. Konyushkova, S. Reed, R. Jeong, K. Zolna, Y. Aytar, D. Budden, M. Vecerik, et al. Scaling data-driven robotics with reward sketching and batch reinforcement learning. arXiv, pages arXiv-1909, 2019.
[4] D. Co-Reyes, R. Veerapaneni, M. Chang, M. Janner, C. Finn, J. Wu, Tenenbaum, and S. Levine. Discovering, predicting, and planning with objects. International Conference Machine Learning Workshop, 2019.
[5] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint 1rXIv: 1901:0286100 2019.
[6] M. Garnelo, K. Arulkumaran, and Shanahan. Towards deep symbolic reinforcement learning. arXiv preprint arXiv: 1609,05588, 2016.
[7] M. Garnelo and M. Shanahan. Reconeiling deep learning with symbolic artificial intelligence: representing objects and relations. Current Opinion in Behavioral Sciences, 29:17-23, 2019.
[8] A. Goyal, A. Lamb, P. Gampa, P. Beaudoin, S. Levine, C. Blundell, Y. Bengio, and M. Mozer. Object files and schemata: Factorizing declarative and procedural knowledge in dynamical systems. arXiv preprint arXiv:2006. 16225, 2020.
[9] A. Graves, M. Liwicki, S. Fernández, R. Bertolami, H. Bunke, and J. Schmidhuber. A novel connectionist system for unconstrained handwriting recognition. JEEE transactions on patter analysis and machine intelligence, 3 31(5):855-88) 868, 2008.
[10] K. Greff, R. L. Kaufmmnn, R. Kabra, N. Watters, C. Burgess, D. Zoran, L. Matthey, M. Botvinick, and A. Lerchner. Multi-object representation learning with iterative variational inference. In International Conference Machine Learning, pages 2424 2433. PMLR, 2019.
[11] K. Greff, A. Rasmus, M. Berglund, Hao, H. Valpola, and Schmidhuber. Tagger: Deep unsupervised perceptual grouping. In Advances in Neural Information Processing Systems, pages 4484 4492, 2016.
[12] D. Ha and J. Schmidhuber. World models. arXiv preprint 1013, 2018.
[13] D. Hafner, T. Lillicrap, Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: 19X.v:9912.016000 2019.
[14] Z. Hc, Li, D. Liu, He, and Barber. Tracking by animation: Unsupervised learning of multi-object attentive trackers. arXiv preprint
2018.
[15] F. Hill, A. Lampinen, R. Schneider, S. Clark, M. Botvinick, J. L.. McClelland, and A. Santoro. Environmental drivers of systematicity and generalization in a situated agent. In International Conference on Learning Representations, 2020.
[16] Hill, S. Mokra, N. Wong, and T. Harley. Human instruction-follloiing with deep reinforcement learning via transfer-learning from text. arXiv preprint arXiv:2005.09338, 2020.
[17] J.-T. Hsich, B. Liu, D.-A. Huang. L. F. Fei-Fei, and J. C. Niebles. Learning to decompose and disentangle representations for video prediction. In Advances in Neural Information Processing Systems, pages 517-526, 2018.
[18] L. Hubert and P. Arabie. Comparing partitions. Journal of classification, 2(1):193-213 1985.
11

---

[19] T. Kipf, E. van der Pol, and M. Welling. Contrastive learning of structured world models. arXiv preprint arXiv: 1911. 2019.
[20] A. Kosiorek, H. Kim, Y. W. Teh, and Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. In Advances in Neural Information Processing Systems, pages 8606-8616 2018.
[21] B. M. Lake, T. D. Ullman, J. B. Teeenbaum, and S. Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40, 2017.
[22] S. Löwe, K. Greff, R. Jonschkowski, A. Dosovitskiy, and T. Kipf. Learning object-centric video models by contrasting sets, arXiv preprint arXiv:2011.10280 2020.
[23] G. Marcus and E. Davis. Rebooting Al: Building Artificial Intelligence We Can Trust. Ballantine Books Inc., 2019.
[24] Z. W. Pylyshyn. The role of location indexes in spatial perception: A sketch of the finst spatial-index model. Cognition, 32:65-995 97, 1989.
[25] S. Racanière, T. Weber, D. Reichert, L. Buesing, A. Guez, D. Jimenez Rezende, A. Puig- domènech Badia, O. Vinyals, N. Heess, Y. Li, R. Pascanu, P. Battaglia, D. Hassabis, D. Silver, and D. Wierstra. Imagination-augmented agents for deep reinforcement learning. In Advances in Neural Information Processing Systems 30, pages 5690-5701, 2017.
[26] W. M. Rand. Objective criteria for the evaluation of clustering methods. Journal of the American Statistical association, 66(336):846 850, 1971.
[27] O. Ronneberger, Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference on Medical image computing and computer-assisted intervention, pages 234-241. Springer, 2015.
[28] M. Shanahan, B. Beyret, M. Crosby, and L. Cheke. Artificial intelligence and the common sense of animals. Trends in Cognitive Sciences, 24(11):862-8) -72, 2020.
[29] D. Silver, A. Huang, C. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of Go with deep neural networks and tree search. Nature, 529((587):484
2016.
[30] S. van Steenkiste, M. Chang, K. Greff, and J. Schmidhuber. Relational neural expecta- tion maximization: Unsupervised discovery of objects and their interactions. arXiv preprint arXiv:1802.10353, 2018
[31] A. Vaswani, N. Shazeer, N. Parmar, Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998-6009 6008, 2017.
[32] R. Veerapaneni, J. D. Co-Reyes, M. Chang, M. Janner, C. Finn, J. Wu, J. Tenenbaum, and S. Levine. Entity abstraction in visual model-based reinforcement learning. In Conference Robot Learning, pages 1439-1456. PMLR, 2020.
[33] X. Wang, R. Girshick, A. Gupta, and K. He. Non-local neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 7794-7803, 2018.
[34] N. Watters, L. Matthey, M. Bosnjak, C. P. Burgess, and Lerchner. Cobra: Data-efficient basel-based rl through unsupervised object discovery and curiosity-driven exploration, 2019.
[35] M. A. Weis, K. Chitta, Y. Sharma, W. Brendel, M. Bethge, A. Geiger, and A. S. Ecker. Unmasking the inductive biases of unsupervised object representations for video sequences, 2020.
[36] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. R. Salakhutdinov, and A. Smola. Deep sets. In Advances in neural information processing systems, pages 3391-3401, 2017.
12


Reviews:
[{"review_id": "oU_Q3zsNg94", "reviewer": "Reviewer_vmwV", "summary": "This paper faces a relevant challenge to develop intelligent agents, this is, accounting for object permanence. The key idea is to align objects across time using a slot-based memory. Specifically, the proposed model, Objects-Align-Transition (OAT), consists of 3 main steps: i) Objects: scene is decomposed into objects using an existing technique (Monet), ii) Align: relevant objects are mapped to a slot-wise memory keeping object identity and dealing with object permanence (identity and persistency), iii) Transition: an object transition model uses a \"transition prediction loss\" to predict future states, maintaining object identity across time. As a result, the method is able to track objects across time by mapping each object into a slot-wise memory and mapping subsequent views of the same object to the same slot.  As relevant features: i) Model is trained end-to-end without supervision, and ii) Transition loss operates at the level of an object-structured representation as opposed to an unstructured space, such as pixels. Results using synthetic and real constrained scenarios indicate that the model is able to handle object persistence and object identity, dealing with problem such as object occlusion and re-appearance in partially observable environments.", "questions": "", "limitations": "", "rating": 7, "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": "Overall the paper is clear and easy to follow. The main contributions are clearly stated and ablation experiments demonstrate their usefulness. Main concern is the potential lack of applicability of the proposed method to more complex scenarios. The current setup only includes simplified scenarios with few objects and limited dynamics. Real natural scenes are harder to decompose into objects (or fully represent as a set of objects). Also, more objects require more memory slots, especially if complex objects have to be further decomposed into several parts to fully model their dynamics. A related point is that the more objects there are, the more difficult it will be to align them. It's entirely possible that a model with a global representation and a \"pixel space\" loss will perform better in the suggested -hypothetical- case. However, as the experiments show, in the limited environment considered here, the proposed model shows good performances and it is able to outperform alternative approaches. In this sense, this work is part of initial efforts to develop models that can tackle problems such as object permanence. This is crucial to deploy embedded agents in natural environments, as a consequence, I recommend its acceptance to Neurips.\n\nFurther comments:\nIs the model able to capture some intuitive physics when learning the transition model?\nWhat is the distance metric in the the Adjacency Matrix used to apply the Hungarian algorithm?\nIs the method able to solve false object detections by Monet?\nCan you clarify how the model is able to improve out-of-distribution generalization or few-shot transfer?\nFigures are difficult to see, frames are really small making very hard to appreciate the results.\nThere are problems with references, several of them do not include information about publication venue.\nMissing reference: Tracking the world state with recurrent entity networks.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Not potential negative societal impacts.", "review_text": "", "extra_scores": {"ethical_concerns": "Not ethical issues", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.0 hours"}, "extra_text": {}}, {"review_id": "gtaaVCOREX4", "reviewer": "Reviewer_gBUs", "summary": "Summary: The paper proposes a slot-based memory architecture to factorize the visual world into a set of objects to preserve object persistence and object identity. Object persistence is important to make sure that the representation of objects are in the memory even though the objects are no longer perceptible, and object identity is important to make sure the same slot refers to the same object at later time-steps. \n\nThe paper conducts experiments on two sequential image datasets, and achieve better results as compared to the state of the art object centric transition models. ", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "Architecture: The proposed architecture consists of 3 components \n(a) Representation module: a module for parsing the visual world into set of slots. For this component paper proposes to use MONet. \n(b) Alignment module: the goal of this module is to align objects across slot representations at different time-steps. \n(c) Slot-wise Transition module: The paper proposes to use SlotLSTM where weights are shared across object slots and can be instantiated in many different ways. \n\nPros:\n- I appreciate and like the idea of evaluating slot-wise transition models for partially observable 3D environments. Different papers have argued before ([1, 2]) of the importance of using structured representations in 3D environments, but I'm not aware of any work which evaluates such ideas in a systematic way on such environments. \n- The paper is easy to read, and generally well written. \n- I acknowledge the importance of doing analysis in section 4.2. Its a useful result that the object level loss performs better as compared to pixel-level prediction (for concreteness object-level loss i.e., a loss which is computed directly between predicted and target object representations (as compared to a pixel-level loss). \n\nCons:\n- I'm concerned about the depth of the experiments, as well as the statements regarding the relevant literature. The authors missed the relevant literature which tries to achieve *both* object persistence and identity (though they don't evaluate on 3D partially observable environments)\n\n*Comments about the contributions*:\n\n- When I read the title of the paper, I was pretty excited because the paper seems to be one of the very few papers which tries to use factorized representations for partially observable 3D environments. Though, on reading the introduction it seems that the paper proposes a new architecture for learning slot-wise transition model, and just evaluates the proposed architecture on 3D environments, so IMO title seems to be not aligned as to what paper is trying to do (in its current form as of now)\n\n*two issues that are not handled satisfactorily by other transition models, namely object persistence and object identity*\n\nI'm slightly concerned about this statement. This statement is factually incorrect. There's some recent work which tries to achieve both object persistence and object identity in a single architecture. For ex. See SCOFF [1] and its followup Neural production systems [2]. SCOFF consists a set of slots (or object files) where each slot attends to different part of the state space (i.e.. different entity). Each object file is parameterized as a LSTM. Different slots are represented by different LSTM (with dynamic weight sharing). In SCOFF,  the state of object file persists over time (to address object identity), and different object files attend to different parts of the state space via top-down attention where the \"query\" is a function of the hidden state of the object file.  So, in essence the Slot-wise Transition module (set of LSTMs with weight sharing) used in the proposed work is a baseline in SCOFF.  In NPS, authors also claim to achieve better results as compared to OP3. That said, both in SCOFF or NPS, the authors rather evaluate such ideas on toy 2D domains or for \"easy\" downstream RL task. \n\nReferences:\n- [1] The ThreeDWorld Transport Challenge: A Visually Guided Task-and-Motion Planning Benchmark for Physically Realistic Embodied AI\nhttps://arxiv.org/abs/2103.14025\n- [2] CausalWorld, https://arxiv.org/abs/2010.04296\n- [3] Object Files and Schemata: Factorizing Declarative and Procedural Knowledge in Dynamical Systems,  \nhttps://arxiv.org/abs/2006.16225 (ICLR'21) \n- [4] Neural Production Systems (https://arxiv.org/abs/2103.01937)\n\n*Comments about the experimental results*:\n\n- The authors evaluate the proposed architecture using three downstream metrics: Encoding ARI, Unroll Pixel Error and Unroll ARI. In Table 1, paper shows the importance of proposed alignment module.  It would be useful to know to variance of both the baselines and proposed model across different seeds. I understand that such models can be a bit unstable, so it would be helpful to know the variance and standard deviation across seeds (the authors refer to Appendix for OP3, but it would also be useful for OAT). \n\n- The results mentioned in the section 4.3 are encouraging but it would also be useful to report results on some downstream metric i.e., does learning the representations using the proposed method leads to better downstream performance on some RL task ? Otherwise its hard to judge the significance of these results.\n\nQuestions for authors:\n\nI'm pretty happy to change my ratings if authors can address these points in a satisfactory way: \n\n- *Acknowledge relevant literature*: The reviewer appreciates the effort of trying to use factorized representations for 3D partially observable environments but its importance to acknowledge the relevant literature. OAT consists of 3 modules (a) MONet (b) AlignNET (c) SlotLSTM. Both MONet and SlotLSTM (like in SCOFF/NPS) are proposed and used in literature already so it would be useful to acknowledge that there exists work which has the property of both object persistence and identity and mention the limitations of previous work i.e., it does not evaluate on 3D partially observable environments. \n\n-  *Performance of the proposed method on downstream RL tasks*:vIt would be useful to report downstream metric for the robotics experiment. It would be very interesting to see, if the representations learned by the proposed model leads to more powerful policies or improved sample efficiency on the downstream RL task. \n\n-  *Performance of the proposed method on downstream RL tasks for Playroom env*:  The paper mentions that the data  is collected from a\npre-trained agent moving in a simulated 3D environment. It would again be useful to know if the learned representations lead to better generalization (or improved sample efficiency or transfer) for downstream RL task. \n\n-  *Performance of the proposed method on 2D task*: It would be useful to know how the representations learned by the proposed method are useful as compared to [1]. In order to evaluate this, authors can follow the similar experimental procedure as in [1] i.e. train OAT on sequence of frames, and use the learned representation of entities as input to Transformer (as compared to MONet used in [1]). Does the representations learned by the proposed method leads to improved downstream results ?\n\n[1] Object-based attention for spatio-temporal reasoning: Outperforming neuro-symbolic models with flexible distributed architectures, https://arxiv.org/abs/2012.08508\n\nPlease don't hesitate to ask for any clarifications, if something is not clear. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Please refer to main review for pointers for improving the presentation of the current work.  \n\n", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns. ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}, {"review_id": "NB0VWjHtJ07", "reviewer": "Reviewer_eFBe", "summary": "The goal is to learn models that can parse object entities of a scene from third person views in an unsupervised manner. A key challenge in prior work is to enforce consistent parsing over time. To this end, this paper proposes an explicit alignment module, together with a transition module on top of it. By training the whole model end-to-end, the model is claimed to able to extract consistent object-level features while aligning the entities of these features over time. Evaluation is conducted on one simulation based dataset and one real world dataset. The proposed approach is shown to outperform a prior method [31] on perceptual and image segmentation based metrics.", "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": "### Formatting Issue\nBefore proceeding with any reviewing, I want to point out that this submission has a formatting issue. The main PDF contains appendix content after the references and paper checklist (see page 12 to 21). This extra content is only allowed to show up in supplementary, but is somehow presented in the main PDF here.\n\nEven with this flaw, I'll still perform my reviewer duty and provide the review below. \n\n### Comments\n- [Fig. 1] Better explain what these images represent in the caption. These are images visualized by MONet's decoder? So the actual task is not video frame prediction (which is what it looks like in the first glance), but rather to segment the scene based on objectness (which again has non-trivial connection to these images).\n- [Fig. 1] For OP3, is a MONet decoder additionally trained to decode it's representation? Otherwise how could a pre-trained MONet decoder work out of the box?\n- [L120] \"During the encoding phase the transition model is fed aligned, observed object representations, $z^a_t$, and actions, $a_t$, to predict ...\" -> What is this \"action\"? Where does it come from? Is this an \"action\" of a policy in the sense of RL? The proposed approach is just doing unsupervised learning and thus should not have any RL stuff right?\n- [L129] What do you do with the second $\\Delta^m_t$ (predicted in the unroll step)? Is it also used to update $m_t$ following $m_{t+1}=m_t+\\Delta^{m}_t$ (L128) after the same update is done for the first $\\Delta^m_t$ from the encoding step?\n- [L146-148] \"This is a spatial mixture loss parameterised by $\\sigma_{bg}$ and $\\sigma_{fg}$, ... that go into the mixture loss.\" -> This is unclear. The authors should clearly explain the high level mechanism of the MONet loss without needing the readers to go to [2].\n- [L154-161] \"Alignment module losses ... for all experiments presented in this paper.\" -> This is also unclear. Why can minimizing the difference between $z^a$ and $z^d$ encourage alignment (or i.e. temporal consistency)? What is the intuition behind that?\n- [L228-232] Since the main baseline is [31], why not just compare on the dataset used in [31] (and potentially just use the numbers in their paper)? Why is this 3D Playroom environment dataset needed? With that you also need to re-train [31] on this new dataset right?\n- [L228] \"3D Playroom environment [14,15]\" -> Wrong citation. Should be changed to the following:\n  - Josh Abramson, Arun Ahuja, Iain Barr, Arthur Brussee, Federico Carnevale, Mary Cassin, Rachita Chhaparia, Stephen Clark, Bogdan Damoc, Andrew Dudzik, Petko Georgiev, Aurelia Guy, Tim Harley, Felix Hill, Alden Hung, Zachary Kenton, Jessica Landon, Timothy Lillicrap, Kory Mathewson, Soňa Mokrá, Alistair Muldal, Adam Santoro, Nikolay Savinov, Vikrant Varma, Greg Wayne, Duncan Williams, Nathaniel Wong, Chen Yan, Rui Zhu. Imitating Interactive Intelligence. arXiv preprint arxiv:2012.05672, 2020.\n- Overall the image quality in the paper is not at the publication level.\n  - Fig. 1 -> Images too blurry.\n  - Fig. 2 -> The text between \"Transition module\" blocks are illegible.\n  - Fig. 4 -> Images too small.\n  - Fig. 7 -> Images too blurry. Hard to tell what the content is.\n- Please add the venue of the following citations:\n  - [31] -> CoRL 2019.\n  - [33] -> arXiv preprint.\n  - [34] -> arXiv preprint.\n\n### Justification of Rating\nOn the up side, the paper is addressing an important and challenging direction. On the downside, the quality of the presentation can be improved (see the unclear details above). Also the intuition behind the design of the loss functions is unclear, and the use of a new dataset for experimental evaluation is unjustified. Given these concerns, the initial rating is kept at slightly below the borderline.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No discussions on potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "L_32YZJUK4k", "reviewer": "Reviewer_V6Rr", "summary": "This paper focuses on decomposing the scene into objects and predicting how these objects evolve in successive frames. This method is based on MONet and improves object alignment by proposing AlignNet. Compared with OP3, the quantitative results show better performance.", "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": "I thank the authors for their efforts to rewrite the paper and include the AlignNet method. I saw a few strengths in this paper: Compared with pixel-level loss, the author demonstrated better performance using object-level loss. In addition to proposing the transformation module, the authors also compared the core of different transition models and found the best one using transformer followed by LSTM. I am inclined to accept, but it seems that many of the claims made in the paper about the method's ability to process some observable scenes and capture the reproduction of objects are not fully supported by the experimental results (explained in detail in the next section). I would like to hear the authors' response to further verify my assessment.\n\nAfter reading the rebuttal, most of my concerns were resolved. I thank the authors for their response.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "It seems that many of the claims made in the paper about the method's ability to process part of the observable scene and capture the reappearance of objects are not fully supported by the experimental results.\n- The authors claim that the proposed method is better in dealing with 3D partially observable environments. In view of the fact that the image is too small to study, I think it is better to perform a quantitative detailed analysis: given the different ratios of observations (such as 10%, 50%, etc.), show the robustness of the method.\n- The authors mentioned that methods like C-SWM \"is unlikely to scale well to partially observable environments with significant movement of objects across the field of view\". But there are no quantitative or qualitative results to prove this statement.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "AXjQXfEDk2r", "reviewer": "Reviewer_sozU", "summary": "The paper presents an approach to learning implicit representations of\n(simple synthetic) 3D scenes from a few consecutive video frames.\nBased on the Monet scene decomposition and representation approach,\nthe new proposed approach adds: 1) a differentiable alignment module\n(Hungarian algorithm based) that helps maintain the same scene component in the\nsame model representation slot, and 2) a temporal transition component that\nlinks representations, 3) loss based over object-representations instead of over pixels.\nIn comparison to a state-of-art approach (OP3), a better representation is achieved and\nis better maintained many frames past training.\n", "questions": "", "limitations": "", "rating": 8, "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 approach is built on 3 components: the previous Monet component, a\nnew alignment component, and a new temporal transition component that links\nconsecutive time frames. The paper clearly explains the motivation and\ntheory behind each component, and experiments demonstrate the benefits\nfrom the components.\nGenerally, the paper is quite clear, although many of the images depend on being able to\nzoom the paper, and using higher resolution embedded images would have helped.\nThe paper does not give a precise definition of the experimental 'unroll' step, so it\nis a little hard to evaluate how much information is given as input into the unroll steps.\nWhile the proposed approach is good theoretical and experimental improvement, the\nreconstruction images show that there are still many extra and missing scene\ncomponents. This could have been discussed. The quantitative results only used the\nforeground objects (which seems reasonable given the majority of the pixels are the\nbackground and this seems to be reasonably represented) and the numerical results are\nmuch improved, but the high numbers did not relate to my visual inspection of the\nreconstruction results (fig 1 good, fig 5 top reasonable, others not so good).\nPerhaps this is a consequence of the ARI measure being based on the inferred object masks,\nand none of these seem to have been shown in the paper or supplementary materials.\nThe pixel error also seems to include the background, which is well represented and so\ndominates the error calculation.\n\nIn summary: several clearly beneficial theoretical and practical innovations\ngiving a substantial improvement to the implicit learning from video research theme,\nbut the research theme is still in an early stage.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The paper could be more critical about what was not working well, and how far the research ahs to go to be practical.", "review_text": "", "extra_scores": {"ethical_concerns": "none", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "36o47XdTAy7", "reviewer": "Reviewer_Jy59", "summary": "This paper proposes a novel object-based transition model for the decomposition of a scene into objects, followed by alignment to maintain time consistency. It then predicts how these objects will evolve over successive frames. It is trained end-to-end without supervision using transition losses at the level of object-structure representation rather than pixels. A novel alignment module enables the proper handling of two issues that other methods do not do satisfactorily: object persistence and object identity. ", "questions": "", "limitations": "", "rating": 7, "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": "The proposed OAT model (Objects-Align-Transition) has three modules, two of which are the contributions:\n\n- MONet, a scene decomposition and representation module to transform a raw image into a slot-wise object-based representation\n- A novel alignment module to ensure that each object is represented in the same slot across space, even if it has disappeared from view\n- A novel slot-wise transition model that operates on the object representations to predict future states. \n\nThe application of each module is straightforward: MONet creates an object-based representation, the slot-wise transition model uses temporal information to predict future states, and the alignment module ensures that the transition model receives information pertaining to the same object across time. Experiments show that the proposed method outperforms the current state of the art transition-based model, OP3, and that the proposed modules individually contribute to performance improvements. Particularly, the authors pinpoint drawbacks of such model, and how the proposed contributions address these limitations. I would like to have seen comparisons with other methods, especially because there have been a few proposed methods (as acknowledged by the authors) that try to address some of the mentioned issues, and it would have been useful to see to what degree they succeed relative to the proposed method. Figures and other qualitative results are also very helpful in understanding the impact of the proper modeling of object permanence and identity. \n\nThe motivation behind the contributions is well explained and related to key issues of current methods. I believe finding ways to address object persistence and object identity would have impact in several field of research, and it is interesting to see progress on that front. The related work section of this paper is well-written, and does a good job positioning the proposed method relative to other published works, showing where they overlap and where the novelty resides, particularly related to the issues of partially aligned environments and semantically meaningful losses. I am not an expert in the field, but the theory and formulations for the various losses and transition models seem solid, with no notable mistakes. Overall, the paper is well-written, clear and concise. As for improvements from previous submissions, I think the paper is self-contained and have no issues with building on top of non-peer-reviewed work for the current submission. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Some directions for future work are provided, such as making stochastic predictions about the future and better modeling of uncertainties about objects in the environment. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers agreed that this is a solid paper that deserves acceptance. Authors are highly encouraged to address the key comments reported by reviewers as well as to implement all the improvements (as indicated by authors in the rebuttal) in the final camera-ready version. 
Decision: Accept (Poster)