Learning an agent that interacts with objects is ubiquituous in many RL tasks. In most of them the agent's actions have sparse effects : only a small subset of objects in the visual scene will be affected by the action taken. We introduce SPECTRA, a model for learning slot-structured transitions from raw visual observations that embodies this sparsity assumption. Our model is composed of a perception module that decomposes the visual scene into a set of latent objects representations (i.e. slot-structured) and a transition module that predicts the next latent set slot-wise and in a sparse way. We show that learning a perception module jointly with a sparse slot-structured transition model not only biases the model towards more entity-centric perceptual groupings but also enables intrinsic exploration strategy that aims at maximizing the number of objects changed in the agents trajectory.
This paper introduces a model that learns a slot-based representation, along with a transition model to predict the evolution of these representations in a sparse fashion, all in a fully unsupervised way. This is done by leveraging a self-attention mechanism to decide which slots should be updated in a given transition, leaving the others untouched. The model learns to encode in a slot-wise and is trained on single step transitions.
This work tackles an important problem, and is very well motivated and presented in a very clear fashion. It reuses some known ideas and components, but combines them in a nice way. I especially liked the use of attention to select what to update, which is a good prior to have.
However, the results presented unfortunately seem to fall a bit short in this current version, and some decisions might have had too much of an effect on some of these shortcomings. Given some improvements, this work might become quite promising, but for the time being I am leaning against publication.
1.	Most modeling decisions are clear and well-motivated, however the choice to make the transition model f_trans always be applied only “slot-wise” might be too restrictive. Indeed, for a given action, this means that 2 slots have to independently learn the effect of that action (e.g. as shown in the example in Figure 2. right), and that some interactions are ~impossible to learn (e.g. in Sokoban, pushing a box requires knowing about the location of both the agent and the box). This could have been alleviated if the transition had access to “interactions outcomes” (e.g. if using a GraphNet, or in your model, if \tilde{s} was provided to the transition function f_theta). Other works (including Zambaldi et al 2018, which is cited several times), handle this appropriately. Did you try to provide \tilde{s} to the transformation operator (e.g. in Figure 7 left)?
2.	Adding a direct comparison to pure non-slotted versions of the model/baselines would have been quite useful, as currently it is unclear why certain things are failing.
3.	Similarly, finding what was the output of the CNN encoder for pixel inputs was a bit too difficult. It is explained in the Appendix that one maps into 4x4 feature maps, but that might be too large for the current environments? Indeed for Sokoban, this means that any grid is partially supported by several “slots”, which may hurt the results more than they should (especially combined with the slot-wise transition constraints expressed above).
4.	The early state of the current results are quite visible in all examples of the “Separate training” model predictions (Figure 3, 5, 8, 9, 10 and 11). None of these actually show this model performing a “correct” prediction for t+1? They only predict no changes, or nonsensical interpolations… This is not sufficient to try to make an argument about the “joint training” helping, and most discussions about “what information they contain” is strenuous at best.
5.	The paper keeps mentioning that it “implicitly imposes transitions to be sparse”, however it is never explained how that would come about? I understand that the softmax in the self-attention may tend to become “peaky” and hence only affect a few slots (and the results do seem to confirm this observation), but I was expecting an explicit loss to enforce this fact. The current emphasis seems a bit ill-funded, so I would present more evidence to it or downplay it.
6.	Some of the results shown seem hard to interpret or provide only weak evidence for the proposed model:
	a.	Figure 2. Left does seem to indicate a benefit in using the self-attention module, but it is hard to know how much of an effect the gap between the orange and red curves actually imply. This figure is overall a bit too small to interpret, and it might be better to split the 2 conditions into sub-plots. The names of the curves in the legend do not correspond to anything described in the text/caption (but I could understand them…). I was expecting more discussion of the results in Figure 2, for example at the end of Section 4.1.
	b.	As explained above, Figure 3 only shows that the Joint Training can perform a 1-step prediction, which is ok but is the bare minimum. Does it handle multi-step rollouts?
	c.	Figure 4 does not seem to provide any significant results or insights. I would interpret them as showing no significant difference between the curves, and they are too small to extract any information out of them. I would remove this figure fully.
	d.	Figure 5 is unclear about what f_k=0 really is, and once again just shows that Separate training does nothing. I expected it to be exactly reconstructing s_t (given that the others are trying to predict s_t+1)? But this is only the case for the joint training, and without knowing what x_t actually was, it is hard to trust. The fact that f_k changes what it does as it is being increased makes the whole point hard to interpret.
7.	Figure 6 was quite interesting, and I feel like this could be pushed forward in a quite interesting manner. The choice of “maximizing the number of entities selected” was fair as a first try, but I could imagine it failing to generalize to more complex environments, or to be easily exploited if one ever decides to pass gradients back into the representation from the policy. It was unfortunate that the legends do not correspond to anything expressed in the main text or caption (e.g. why do you not reuse the “valid_move”, …, ”blocked_push” names that you thoroughly introduce?). What is “random_XX”? Could you comment on why the curves seem to have a large increase in variance along the 80000-100000 updates region?
Details:
8.	Figure 2 (right) was good to explain how the model worked, but I would actually change its location and try to move it into Figure 1. Similarly, Figure 7 in the Appendix seemed rather necessary to understand the model, and belongs in the main text in my opinion.
9.	When presenting the self-attention block, it would be good to directly state that these are MLPs receiving [s_t, a_t] (as done in the Appendix).
10.	Is sigma^2 in the decoder fixed? To which value?
11.	When presenting the “sparse” and “full” settings, having access to Figure 7 and the rest of the Appendix might be beneficial, it would be good to point forward to it.