Self-supervised Visual Reinforcement Learning with Object-centric Representations
Authors: Andrii Zadaianchuk
Maximilian Seitzer
Georg Martius
Venue: ICLR 2021 Spotlight (2021)
Keywords: self-supervision
autonomous learning
object-centric representations
visual reinforcement learning

Abstract: Autonomous agents need large repertoires of skills to act reasonably on new tasks that they have not seen before. However, acquiring these skills using only a stream of high-dimensional, unstructured, and unlabeled observations is a tricky challenge for any autonomous agent. Previous methods have used variational autoencoders to encode a scene into a low-dimensional vector that can be used as a goal for an agent to discover new skills. Nevertheless, in compositional/multi-object environments it is difficult to disentangle all the factors of variation into such a fixed-length representation of the whole scene. We propose to use object-centric representations as a modular and structured observation space, which is learned with a compositional generative world model.
We show that the structure in the representations in combination with goal-conditioned attention policies helps the autonomous agent to discover and learn useful skills. These skills can be further combined to address compositional tasks like the manipulation of several different objects.

SELF-SUPERVISER VISUAL REINFORCEMENT LEARN- ING WITH OBJECT-CENTRIC REPRESENTATIONS
Andrii Zadaianchuk 1,2, Maximilian Seitzer',, Georg Martius'
1 Max Planck Institute for Intelligent Systems, Tübingen, Germany
2 Department of Computer Science, ETH Zurich andrii. zada i www.aaachhuaaaaae www.uuuwkttuiiiine...ue.e mpg.
ABSTRACT
Autonomous agents need large repertoires of skills to act reasonably on new tasks that they have not seen before. However, acquiring these skills using only a stream of high-dimensional, unstructured, and unlabeled observations is a tricky challenge for any autonomous agent. Previous methods have used variational autoencoders to encode a scene into a low-dimensional vector that can be used as goal for an agent to discover new skills. Nevertheless, in compositional/mulli object environments it is difficult to disentangle all the factors variation into such fixed-length representation of the whole scene. We propose to object-centric representations as a modular and structured observation space, which is learned with compositional generative world model. We show that the structure in the representations in combination with goal-conditioned attention policies helps the autonomous agent to discover and learn useful skills. These skills can be further combined to address compositional tasks like the manipulation of several different objects.
attps://martius-ll ww.bb.githbbb... 10. 10MOOR
INTRODUCTION
Reinforcement learning (RL) includes a promising class of algorithms that have shown capability to solve challenging tasks when those tasks are well specified by suitable reward functions. However, in the real world, people are rarely given a well-defined reward function. Indeed, humans are excellent at setting their own abstract goals and achieving them. Agents that exist persistently in the world should likewise prepare themselves to solve diverse tasks by first constructing plausible goal spaces, setting their own goals within these spaces, and then trying to achieve them. In this way, they can learn about the world around them.
In principle, the goal space for an autonomous agent could be any arbitrary function of the state space. However, when the state space is high-dimensional and unstructured, such as only images, it desirable to have goal spaces which allow efficient exploration and learning, where the factors of variation in the environment are well disentangled. Recently, unsupervised representation learning has been proposed to learn such goal spaces (Nair et al. 2018 2019, Pong et al., 2020). All existing methods in this context use variational autoencoders (VAEs) map observations into low-dimensional latent space that can later be used for sampling goals and reward shaping.
However, for complex compositional scenes consisting of multiple objects, the inductive bias of VAFs could be harmful. In contrast, representing perceptual observations in terms of entities has been shown to improve data efficiency and transfer performance on wide range of tasks (Burgess et al. 2019). Recent research has proposed a range of methods for unsupervised scene and video decomposition Greff al. [2017 Kosiorek al. 2018 Burgess et al.) 2019 Greff al. 2019 Jiang et al. 2019 Weis et al. Locatello et al. 2020). These methods learn object representations and scene decomposition jointly. The majority of them are part motivated by the fact that the learned representations are useful for downstream tasks such as image classification, object detection, or semantic segmentation. this work, we show that such learned representations are also beneficial for autonomous control and reinforcement learning.
"qqull contribution

---

Zt
Train SCALOR on data from random policy
O2
Goal- conditioned attention policy
SCALOR encoder
Train SMORL with generated goals
0g
q (z)0)
Ty (ag 2 Zgy
Zy
Evaluate SMORL new tasks from environment
Figure I: Our proposed SMORL architecture. Representations Z, are obtained from observations o, through the object-centric SCALOR encoder 4s, and processed by the goal-conditional attention policy To (arazzz kg). During training, representations of goals are sampled conditionally on the representations of the first observation ž1. At time, the agent is provided with an external goal image Og that is processed with the same SCALOR encoder to a set of potential goals {z, === After this, the goal Zg sequentially chosen from this set. This way, the agent attempts solve all the discovered sub-tasks one-by-one, not simultaneously.
We propose combine these object-centric unsupervised representation methods that represent the scene as set of potentially structured vectors with goal-conditional visual RL. In our method (illustrated in Figure []. dubbed SMORL (for self-supervised multi-object RL), a representation of raw sensory inputs learned by compositional latent variable model based on the SCALOR architecture (Jiang et al., 2019). We show that using object-centric representations simplifies the goal space learning. Autonomous agents different goals with a reward function that utilizes the structure of the learned goal space. Our main contributions follows:
those representations to learn how to achieve
We show that structured object-centric representations learned with generative world models can significantly improve the performance of the self-supervised visual RL agent.
We develop SMORL, an algorithm that uses learned representations autonomously discover and learn useful skills compositional environments with several objects using only images as inputs.
We show that with fully disentangled ground-truth representation there is large benefit from using SMORL environments with complex compositional tasks such as rearranging many objects.
RELATED WORK
Our work lies in the intersection of several actively evolving topies: visual reinforcement learning for control and robotics, and self-supervised learning. Vision-based RL for robotics is able to efficiently learn variety of behaviors such grasping, pushing and navigation (Levine al. 2016 Pathak et al. 2018 Levine et al. 20188) Kalashnikov et al.) 2018) using only images and rewards as input signals. Self-supervised learning is a form of unsupervised learning where the data provides the supervision. It was successfully used to learn powerful representations for downstream tasks in natural language processing (Devlin et al., 2018 Radford et al. 2019) and computer vision (He et al.) 2019; Chen et al. 2020). In the context of RI. self-supervision refers to the agent constructing its own reward signal and using to solve self-proposed goals (Baranes & Oudeyer 2013 Nair et al.
2018 Péré et al., 2018 Hausman et al.) 2018: Lynch al. 2019). This is especially relevant for visual RL, where reward signal is usually not naturally available. These methods can potentially acquire diverse repertoire of perepol-puppoee robotic skills that be reused and combined during test time. Such self-supervised approaches are crucial for scaling learning from narrow single-tar task learning more general agents that explore the environment their own to prepare for solving
2

---

many different tasks in the future. Next, we will cover the two most related lines of research in more detail.
Self-supervised visual RL (Nair al. 2018; 2019 Pong et al. 2020 Ghosh et al. 2019 Warde- Farley et al. 2019 Laversanne- Finot et al. 20188) tackles multi-task RL problems from images without any external reward signal. However, all previous methods assume that the environment observation can be encoded into a single vector, e.g. using VAE representations. With multiple objects being present, this assumption may result in object encodings overlapping in the representation, which is known as the binding problem (Greff et al. 2016 2020). In addition, as the reward is also constructed based on this vector, the agent is incentivized to solve tasks that are incompatible, for instance simultancously moving all objects to goal positions. In contrast, we suggest to learn object-centric representations and use them for reward shaping. This way, the agent can learn to solve tasks independently and then combine these skills during evaluation.
Learning object-centric representations in RL (Watters et al. 2019 van Steenkiste et al. 2019 Veerapaneni et al.) 2020] Kipf et al. 2020) has been suggested to approach tasks with combinatorial and compositional elements such as the manipulation of multiple objects. However, the previous work has assumed a fixed, single task and given reward signal, whereas are using the learned object-representation to construct a reward signal that helps to learn useful skills that can be used to solve multiple tasks. In addition, these methods et al. 2019) and IODINE (Greff et al. 2019). which do not explicitly contain features like position and scale. These features can be used by the agent for efficient sampling from the goal space and thus the explicit modeling of these features helps to create additional biases useful for manipulation tasks. However, we expect that other object-centric representations could also be successfully applied as suitable representations for RL tasks.
scene-mixture models such as MONET (Burgess
BACKGROUND
Our method combines goal-conditional RL with unsupervised object-oriented representation learning for multi-object environments. Before we describe each technique in detail, we briefly state some RL preliminaries. We consider Markov decision process defined by (S, A, p, r), where S and A are the continuous state and action spaces, p: representing the probability of transitioning to state s1+1 E S from state s, E S given action a E A. and r: S + R is function computing the reward for reaching state 8,+1. The agent's objective is to maximize the expected return R = Ei= Egg.op.. a. s [r(s++1] over the horizon T, where Pr (5t) is the state marginal distribution induced by the agent's s policy (ay]s::)
A + [0, oo) an unknown probability density
3.1 COAL-CONDITIONAL REINFORCEMENT LEARNING
In the standard RL setting described before, the agent only learns to solve a single task, specified by the reward function. If are interested in agent that can solve multiple tasks (each with different reward function) in an environment, we can train the agent on those tasks by telling the agent which distinct task solve at each time step. But how can we describe a task to the agent? A simple, yet not too restrictive way is to let each task correspond to an environment state the agent has to reach, denoted as the goal state g. The task is then given to the agent by conditioning its policy T (a= | st, g) on the goal g, and the agent's 's objective turns to maximize the expected goal-conditional return:
T
E
Bg-
Bay an m (r (S(()))]]
(1)
where G is some distribution over the space of goals G C the agent receives for training. The reward function can, for example, be the negative distance of the current state the goal: (s)s)) Often, we are only interested in reaching a partial state configuration, e.g. moving an object to a target position, and want to avoid using the full environment state the goal. In this case, we have to provide a mapping m: S + 9 of states to the desired goal space; the mapping is then used to compute the reward function, i.e. (s) [mmms)) gll-
||s gil-
3

---

As the reward is computed within the goal space, it is clear that the choice of goal space plays crucial role in determining the difficulty of the learning task. If the goal space is low-dimensional and structured, e.g. in terms of ground truth positions of objects, rewards provide meaningful signal towards reaching goals. However, if we only have access to high-dimensional, unstructured observations, e.g. camera images, and naively choose this space as the goal space, optimization becomes hard as there is little correspondence between the reward and the distance of the underlying world states (Nair al. 2018).
One option to deal with such difficult observation spaces is to learn goal space in which the RL task becomes easier. For instance, we can try to find a low-dimensional latent space Z and use both as the input space to our policy and the space in which we specify goals. If the environment is composed of independent parts that we intend to control separately, intuitively, learning to control is easiest the latent space is also structured in terms of those independent components. Previous research (Nair et al. 2018 Pong et al.) 2020) relied on the disentangling properties of representation learning models such as the 3-VAE (Higgins et al. 2017) for this purpose. However, these models become insufficient when faced with multi-object scenarios due to the increasing combinatorial complexity of the scene, model explicitly geared towards inferring object-structured representations, which we introduce in the next section.
show in Sec. 5.2 and in App. A.2, Instead,
3.2 STRUCTURED REPRESENTATION LEARNING WITH SCALOR
SCALOR (Jiang et al. 2019) is a probabilistic generative world model for learning object-oriente representations of a video or stream of high-dimensional environment observations. SCALOR assumes that the environment observation o, at step generated by the background latent variable z; and the foreground latent variable z, The foreground further factorized into a set of object representations ft [Zt,n ( InEo., where O is the set of recognised object indices. To combine the information from previous time steps, a appagatiiondiscoverr In SCALOR, an object is represented by 2t,n (=i,m pres where s.the)) The scalar 2|,m defines if the
model is used (Kosiorek et al., 2018).
object is present in the scene, whereas the vector Zi," encodes object appearance. The component dopth With
zieeer is further decomposed into the object's center position z scale zqiile and depth 2t, this, the generative process of SCALOR can be written as:
p(oLT, (1T) p(u, (z he) II p(o, p(2 2 p(2 x p()
222))
rendering
background transition
discrveer
propagation
where z; = (z, z contains latent variables of objects discovered in the present step, and z? contains latent variables of objects propagated from the previous step. Due to the intractability of the true posterior distribution p(Zı:T 01:T), SCALOR trained using variational inference with the following posterior approximation:
bg
T
q(Z1. (IT))
9(z ZZc 2c1,
II 9
|2,, o.) 2. (z))((z) (zz, (3)
t=1
t=1
by maximizing the following evidence lower bound [(B,
Σ E (os( (esltzlzee [og 2x)] - DKL. [go(z] nee,,
pz( zewi]] (4)
where DKL denotes the Bullback-Leible divergence, As we are using SCALOR in an active setting, we additionally condition the next step posterior predictions on the actions a, taken by the agent. For more details and hyperparameters used to train SCALOR, we refer App. D33 In the next section, we describe how the structured representations learned by SCALOR can be used in downstream RL tasks such as goal- cond-conditional visual RL.
4

---

4
MULTI-OBJECT REINFORCEMENT LEARNING
Learning from flexible representations obtained from unsupervised scene decomposition methods such as SCALOR creates several challenges for RL agents. In particular, these representations consist of sets of vectors, whereas standard policy architectures assume fixed-length state vectors as input. We propose to use goal-conditioned allention policy that can handle sets as inputs and flexibly learns to attend to those parts of the representation needed achieve the goal at hand.
In the setting we consider, the agent is not given any reward signal or goals from the environment at the training stage. Thus, discover useful skills that can be used during evaluation tasks, the agent needs to rely on self-supervision in the form of an internally constructed reward signal and self-proposed goals. Previous VAE-based methods used latent distances the goal state as the reward signal. However, for compositional goals, this means that the agent needs to master the simultaneous manipulation of all objects. In our experiments Sec. [.I] we show that even with fully disentangled. ground-truth representations of the scene, this is challenging setting for state-of-the-ar model-free RI. agents. Instead, we propose to use the discovered structure of the learned goal and state spaces twofold: the structure within each representation, namely object position and appearance, to construct a reward signal, and the set-based structure between representations to construct sub-goals that correspond to manipulating individual objects.
4.1 POLICY WITH GOAL-CONDITIONED ATTENTION
We use the multi-heade attention mechanism Vaswani et al. 2017) as the first stage of our policy ma to deal with the challenge of set-based input representations. As the policy needs to flexibly vary its behavior based on the goal at hand, it appears sensible to steer the attention using a goal-dependent query (2,y) 2g W Each object is allowed to match with the query via object-dependent key K(z,)= 2, W and contribute the attention's 's output through the value V(z,) 2, W which is weighted by the similarity between ((Zg) and K(z)). As inputs, we concatenate the representations 2.berre, for object vectors Zt [zithat, z 22, J. and similarly the goal representation to 2g
where, depth The attention head A₁ is computed as (z, 22 2g
Zy (ZWWk)) ZeW",
Ax softmax
where Ze is packed matrix of all Ze, 's WA, Wh W constitute learned linear transformations and de is the common key, value and query dimensionality. The final attention output A is concatenation of all the attention heads A = [Ai;-..; A Ax). In general, we expect it to be beneficial for the policy to not only attend to entities conditional on the goal; we thus some heads attend based on a set of input independent, learned queries, which are not conditioned on the goal. We go into more details about the attention mechanism in App. D. and ablate the impact of different choices in App. B
The second stage of our policy is fully connecco neural network f that takes as inputs A and the goal representation 2g and outputs an action a,. The full policy Ty can thus be described by
To (fze, InEo, 2g) = f(A, zg).
(6)
4.2 SELF-SUPERVISED TRAINING
In principle, our policy can be trained with any coal-conditional model-free RL algorithm. For our experiments, we picked soft-actor critic (SAC) (Haarnoja et al., (018b) as a state-of state-of-theeart method for continuous action spaces, using hindsight experience replay (HER) (Andrychowicz et al. 2017) standard way to improve sample-efficiency in the goal-conditional setting.
The training algorithm is summarized in Alg. | We first train SCALOR on data collected from a random policy and fit a distribution wbere) to representations z "here of collected data. Each rollout, we generate new goal for the agent by picking a random z what from the initial observation Z and sampling a new where from the fitted distribution phereer)) The policy is then rolled out using this goal. During off-policy training, we are relabeling goals with HER, and, similar to RIG (Nair et al. 2018), also with "imagined goals" produced in the same way as the rollout goals.
5

---

Algorithm 1 SMORL: Self-Supervised Multi-Object RL (Training)
Require: SCALOR encoder ge. goal-conditional policy ne, goal-conditional SAC trainer, number of training
episodes K.
1: Train SCALOR sequences uniformly sampled from D using loss described in Eq. 4
Fit prior p(z "here 22hal the latent encodings of observations.
3: for n== I.... K episodes do
4: Sample goal 3, = (2 x)
5: Collect episode data with policy (at Ze, (4,7g) and SCALOOR representations of observations (Z:
Z-c,; Ogg))-
Store transitions (ze, au; 214-1 Z,,) into replay buffer R.
7:
Sample transitions from replay buffer (z, a, z xzz)) ~ R.
8: Relabel z whher goal components
combination of future and m(zeee z 2
9: Compute matching reward signal R = r(z'; 7,))
10: Update policy (ak Z,, kny) using I. with SAC trainer.
11: end for
We also refer to Alg. in App. D.2 for more detailed description of the algorithm.
A challenge with compositional representations is how to measure the progress of the agent towards achieving the chosen goal. As the goal always corresponds a single object, we have to extract the state of this object in the current observation in order to compute reward. One way is to rely on the tracking objects. as was shown possible e.g. by SCALOR (Jiang et al., 2019). However, as the agent learns, we noticed that it would discover some flaws of the tracking and exploit them to get maximal reward that is not connected with environment changes, but rather with internal vision and tracking flaws (details in App. ED.
We follow an alternative approach, namely to use the z zhaa component of discovered objects and match them with the current goal representation zwhat As the z what space encodes the appearance of objects, two detections corresponding to the same object should close this space (we verify that this hypothesis holds in App. A. Thus, it is easy to find the object corresponding to the current goal what
what ||. case of failure to discover a close representation, object using the distance min& 1 z
iee. when all z have distance larger than some threshold a to the goal representation 2 " we use fixed negative reward Nno goo_ to incentivise the agent avoid this situation.
Our reward signal is thus
1-112 where if ming | What what < a otherwise,
r(z, 2n,)
(7)
where & = arg ming
|| whar
4.3 COMPOSING INDEPENDENT SUB SUB-GOALS DURING EVALUATION
At evaluation time, the agent receives goal image from the environment showing the state to achieve. The goal image is processed by SCALOR to yield a set of goal vectors. For our experiments, we assume that these sub-goals achieve them by cycling through them until all of them are solved. The evaluation algorithm is summarized in Alg. with more details added in App. D.2
independent of each other and that the agent thus sequentially

---

5 EXPERIMENTS
We have done computational experiments to address the following questions:
How well does our method scale to challenging tasks with large number of objects in case when ground-trun representations are provided?
How does our method perform com- pared to prior visual goal-conditioned RL methods on image-basee multi- object continuous control tasks?
(a) View from top
(b) Agent observation
Figure 2: Multi-Object Visual Push and Rearrange envi- ronments with 2 objects and a Sawyer robotic arm.
How suitable are the representations learned by the compositional genera- tive world model for discovering and solving RL tasks?
To answer these questions, we constructed the Multi Aulfi-Object Visual Push and Multi Multi-Object Visual Rearrange environments. Both environments based on MuJoCo (Todorov al.) (0122) and the Multiworld package for image-based continuous control tasks introduced by Nair et al. (2018). and contain 7-dof Sawyer of small picks on a table. In the first environment, the objects are located on fixed positions in front of the robot arm that the arm must push random target positions. We included this environment as it largely corresponds to the Visual Pusher environments of Nair et al. (2018). In the second environment, the task is to rearrange the objects from random starting positions to random target positions. This task is more challenging for RL algorithms due to the randomness of initial object positions. For both environments, we measure the performance distance of all pucks to their goal positions on the last step of the episode. Our code, as well as the multi-objects environments will be made public after the paper publication.
where the agent needs be controlled manipulate variable number
the algorithms as the average
5.1 SMORL WITH GROUND-TRUTH (GT) STATE REPRESENTATION
We first compared SMORL with ground-truth representation with Soft Actor-Critic (SAC) (Haarnoja et al. 2018a) with Hindsight Experience Replay (HER) relabeling (Andrychowicz et al., 2017) that takes an unstructured vector of all objects coordinates input. We are using a one-hot encoding for object identities z what and object and arm coordinates as z where components. With such a representa- tion, the matching task becomes trivial so our main focus in this experiment on the benefits of the goal-conditioned attention policy and the sequential solving of independent sub-tasks. We show the results in Fig. 3 While for 2 objects, SAC+HER is performing similarly, for and 4 objects, SAC+HER fails to rearrange any of the objects. In contrast, SMORL equipped with ground-truth representation is still able to rearrange 3 and 4 objects, and it can solve the more simple sub-tasks of moving each object independently. This shows that provided with good representations, SMORL can use them for constructing useful sub-tasks and learn how to solve them.
5.2 VISUAL RL METHODS COMPARISON
We compare the performance of our algorithm with two other self-supervised, multi-task visual RL algorithms on our two environments, with one and two objects. The first one, RIG (Nair et al. 2019), uses the VAE latent space to sample goals and to estimate the reward signal. The second one, Skew-Fit (Pong et al. 2020), also uses the VAE latent space, however, is additionally biased on rare observations that were not modeled well by the VAE on previously collected data. In terms of computational complexity. both method and RIG need to train generative model before RL training. We note that training SCALOR is more costly than training RIG's VAE due to the sequence processing utilized by SCALOR. However, compared to RIG's VAE during RL training, and compared to Skew-Fit, our method still faster to train Skew-Fit needs continuously retrain its VAF.
trained, SCALOR only adds little overhead
7

---

2 objects
objects
4 objects
Avg. object dist.
0.300
0.2
0.0
0.33
0.1
0.18
0350
0.16-
3.16
0.44
0.14
012
0.12
0.10
Timesteps
Timesteps
Timesteps
SMORL+GT
SAC+GT
Passive policy
Figure 3: Average distance of objects goal positions, comparing SMORL using ground truth representations to SAC with ground truth representations in the Rearrange environment with different number of objects. SAC struggles improve performance when the combinatorial complexity of the scene rises. The dotted line indicates the performance of passive policy that performs no movements. Results averaged over random seeds, shaded regions indicate one standard deviation.
Visual Push
object
objects
0.150
Avg. object dist.
0.10
0.100
0.08
0.075
0.06
0.04
0.025
0.0
0.5
1.0
2.0
3.0
0.0
0.5
1.0
1.5
2.0
Timesteps
x10
Timesteps
x10
Visual Rearrange
object
2 objects
0.200
0.175
Avg. object dist.
0.150
0.125
0.100
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Timesteps
x10
Timesteps
x10
SMORL
RIG
Skew-Fit
SAC+GT
Passive policy
Figure 4: Average distance of objects to goal positions, comparing SMORL to Visual RL Baselines. In addition the baselines, show SAC's performance with ground truth representations. Results averaged over 5 random seeds, shaded regions indicate one standard deviation.
We show the results in Fig. 4 For the simpler Multi-Ohject Visual Push environment, the performance of SMORL comparable to the best performing baseline, while for the more challenging Multi- Object Visual Rearrange environment, SMORL is significantly better then both RIG and Skew-Fit. This shows that learning of object-oriented representations brings benefits for goal sampling and self-supervised learning of useful skills. However, method is still significantly worse than SAC with ground-truth representations. We hypothesize that one reason for this could be that SCALOR right now does not properly deal with occluded objects, which makes the environment
8

---

partially observable from the point of view of the agent. On top of this, suspect noise in the representations, misdetections and an imperfect matching signal to slow down training and ultimately hurt performance. Thus, we expect that adding recurrence to the policy improving SCALOR itself could help close the gap to an agent with perfect information.
5.3 OUT-OF-DISTRIBUTTE
FOR DIFFERENT NUMBER OF OBJECTS
One important advantage of structured policies is that they could potentially still be applicable for observa- tions that are from different, but related distributions. Standard visual RL algorithms were shown to be sensitive to small changes unrelated to the current task (Higgins al. 2018). To see how algorithm can generalize to a changing environment, we tested SMORL agent trained on observations of the Rearrange environment with 2 objects on the environ-
object environment
0.18
Avg. object dist.
SMORL (2 objects)
0.16
SMORL (1 object)
0.14
0.12
ment with object. As be seen from Fig. 5 the
performance of such an agent increases during train- SMORL
0.0
0.6
Training examples
X10*
ing up performance comparable agent that was trained on the I object environment.
Figure 5: Out-of-distribution generalization of SMORL agent training on Visual Rear- range with two objects and being tested with one object. Green line shows final perfor- mance when training with one object.
6 CONCLUSION AND FUTURE WORK
In this work, we have shown that discovering structure in the observations of the environment with a compositional generative world models and using it for controlling different parts of the environment is crucial for solving tasks in compositional environments. Learning to manipulate different parts of object-centric representations is powerful way to acquire useful skills such object manipulation. Our SMORL agent learns how to control different entities in the environment and can then combine the learned skills achieve more complex compositional goals such as rearranging several objects using only the final image of the arrangement.
Given the results presented so far, there are a number of interesting directions to take this work. First, one can combine leamed sub-tasks with planning algorithm to achieve a particular goal. Currently, the agent is simply sequentially cycling through all discovered sub-tasks, so we expect that a more complex planning algorithm as e.g. described by Nasiriany et al. (2019) could allow solving more challenging tasks and improve the overall performance of the policy. To this end, considering interactions between objects in the manner of Fetaya et al. (2018) or Kipf et al. (2020) could help to lift the assumption of independence of sub-tasks. Second, prioritizing certain sub-tasks during learning, similar to Blaes al. (2019), could accelerate the training of the agent. Finally, an active training of SCALOR to combine the object-oriented bias of SCALOR with a bias towards independently controllable objects (Thomas al. 2018) is interesting direction for future research.
ACKNOWLEDGEMENTS
The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Maximilian Seitzer. Andrii Zadaianchuk is supported by the Max Planck ETH Center for Learning Systems. We acknowledge the support from the German Federal Ministry Education and Research (BMBF) through the Tübingen Al Center (FKZ: 011S (833BB...

---

REFERENCES
Marcin Andrychowicz. Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. In I. Guyon, U. V. Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. 5048- 5058. Curran Associates, Inc., 2017. 2017
Adrien Baranes and Pierre- Yves Oudeyer. Active learning of inverse models with intrinsically motivated goal exploration in robots. Robotics Auton. Syst. 61:49-73, 2013.
Sebastian Blaes, Marin Vlastelica Pogančić, Jiajie Zhu, and Georg Martius. Control what you can: Intrinsically motivated task-planning agent. In H. Wallach, H. Larochelle, A. Beygelzimer, F. Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 12541-12552. Curran Associates, Inc.. 2019. URL httttsts:
neurips 232.2///00/112277776677733333320
Christopher P. Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation, 2019.
Ricky TQ Chen, Xuechen Li, Roger Grosse, and David Duvenaud. Isolating sources of disentanglement in variational autoencoders. In Advances in Neural Information Processing Systems, pp. 2610-2620, 2018.
Ting Chen, Simon Korbllith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05709. 2020.
Jacob Devlin, Ming-We Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: wxXiv:8110.04805 2018.
T. Fetaya, E. Wang. K.-C. Welling, M. Zemel, Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, M. Welling, and R. Zemel. Neural relational inference for interacting systems. In ICML, 2018.
D. Ghosh, A. Gupta, and S. Levine. Learning actionable representations with goal-conditioned policies. ArXiv, abs/1811 1.07199 2019.
Klaus Greff, Rupesh Kumar Srivastava, and Jürgen Schmidhuber. Binding via reconstruction clustering, 2016.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. In Advances in Neural Information Processing Systems, pp. 6691-6701, 2017.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object representation learning with iterative variational inference. Proceedings of the 36nd International Conference on Machine Learning, 2019.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks, 2020.
Haarnoja, Aurick Zhou, Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In ICML, 2018a.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with stochastic actor. arXiv preprint 101.0:1801.01290. 2018b.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. International Conference on Learning Representations, 2018. URL httttpp:e//pennreeeieeeeeee/e/ee
10

---

Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. arXiv preprint arXiv:1911.05722 2019.
1. Higgins, Loic Matthey, A. Pal, C. Burgess, Xavier Glorot, M. Botvinick, S. Mohamed, and Alexander Lerchner. bet:-vae: Learning basic visual concepts with a constrained variational framework. In ICLR, 2017.
Irina Higgins, Arka Pal, Andrei A. Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer reinforcement learning, 2018.
Jindong Jiang, Sepehr Janghorbani, Gerard de Melo, and Sungjin Ahn. Scalable object-oriented sequential generative models. arXiv preprint arXiv:1910.02386 2019.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for rision-basee robotic manipulation. arXiv preprint arXiv:1806.1029: 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. CoRR, Jbbs/1412.6900, 2015.
Thomas Kipf, Elise van der Pol, and Max Welling, Contrastive learning of structured world models. In International Conference on Learning Representations, 2020. URL htttps://opeeeeeae//
net fret/forie Eeroumiii---.................... Hlgx66V
Adam Roman Kosiorek, Hyunjik Kim, Ingmar Posner, and Yee Whye Teh. Sequential attend, infer, repeat: Generative modelling of moving objects. In Advances in Neural Information Processing Systems, 2018.
hawp/////////1...../// 11../8806.000 ab..110..01199999 01794
Adrien Laversanne-Finn Finot, Alexandre Pere, and Pierre- Yves Oudeyer. Curiosity driven explo- ration of learned disentangled goal spaces. Aude Billard, Anca Dragan, Jan Peters, and Jun Morimoto (eds.), Proceedings of The 2nd Conference Robot Learning, volume 87 of Proceedings of Machine Learning Research, pp. 487-504. PMLR, 29-31 Oct 2018. URL html
http:://rroceeding: www..rr............. press/ v8 www.///aaaaaaeeeaaaeaeaa
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):1334-1373 2016.
Sergey Levine, Peter Pastor, Alex Krizhevsky, Julian Ibarz, and Deirdre Quillen. Learning hand- eye coordination for robotic grasping with deep learning and large-scale data collection. The International Journal of Robotics Research, 37(4-5):421 436, 2018
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning dis- entangled representations. In International Conference Machine Learning, pp. 41 14 4124, 2019.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention, 2020.
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. Conference Robot Learning (CoRL), 2019. URL htt.:/tppss//rci...rrrr/ss//////
01119993..111331 01973
Ashvin Nair, Shikhar Bahl, Alexander Khazatsky, Vitchyr H. Pong, G. Berseth, and S. Levine. Contextual imagined goals for self-supervised robotic learning. In CoRL, 2019.
Ashvin V Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. In Advances Neural Information Processing Systems, pp. 9191-9200, 2018.
11

---

Soroush Nasiriany, Vitchyr Pong. Steven Lin, and Sergey Levine. Planning with goal-conditioned policies. In Advances in Neural Information Processing Systems, pp. 14843-144885 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, pereomerfo deep learning library. In H. Wallach, H. Larochelle. A. Beygelzimer, F. d Alché-Bus E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 8026- 8037. Curran Associates, Inc., 2019. WtLLihtts:///poootttt
wwwrrnnurs.... cc/paper//p/pppcpp pdf
Deepak Pathak, Parsa Mahmoudich, Guanghao Luo, Pulkit Agrawal, Dian Chen, Yide Shentu, Evan Shelhamer, Jitendra Malik, Alexei A Efros, and Trevor Darrell. Zero-shot visual imitation. In Proceedings the IEEE Conference on Computer Vision and Pattern Recognition Workshops, pp- 2050-2053. 2018.
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew-fit: State-covering self-supervised reinforcement learning. In Proceedings of the 37nd International Conference on Machine Learning, volume 42 of JMLR Workshop and Conference Proceedings. JMLR, 2020.
Alexandre Péré, Sébastien Forestier, Olivier Sigaud, and Pierre Yves Oudeyer. Unsupervised learning of goal spaces for intrinsically motivated goal exploration. International Conference on Learning Representations, 2018.
wwt. fort:mfoummm wtt.rrummmddSS1PPPP1111
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea directions (by accident). In Proceedings IEEE Conf. on Computer Vision and Pattern Recognition thecvf. com/ 2eweeent______1999_
(CVPR), June 2019. [ Pursue PCA Directions. papers/Ro Rtci:/oolnnnn inek Variational Autoencoders by Acc.deiinntt... CVPR 2019 paper. pdf
Valentin Thomas, Emmanuel Bengio, William Fedus, Jules Pondard, Philippe Beaudoin, Hugo Larochelle, Joelle Pineau, Doina Precup, and Yoshua Bengio. Disentangling the independently controllable factors of variation by interacting with the world. arXiv preprint arXiv: 018880000444. 2018.
E. Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. 2012 JEEE/RSJ International Conference Intelligent Robots and Systems, pp. 5026-5033, 2012.
Sjoerd van Steenkiste, Klaus Greff, and Jürgen Schmidhuber. perspective on objects and systematic generalization in model-based rl. arXiv preprini arXiv:1906.0// 1906.01035,, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Rishi Veerapaneni, John D. Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua B. Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning, 2020.
David Warde-Farley Tom Van de Wiele, Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih. Unsupervised control through non-parametric discriminative rewards. In International Conference on Learning Representations, 2019. URL http:tts///pperreiiiieee
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P. Burgess, and Alexander Lerchner. Cobra: eita-efficient model-baser rl through unsupervised object discovery and curiosity-driven exploration, 2019.
12

---

Marissa A. Weis, Kashyap Chitta, Yash Sharma, Wieland Brendel, Matthias Bethge, Andreas Geiger, and Alexander S. Ecker. Unmasking the inductive biases of unsupervised object representations for video sequences, 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ Salakhut- dinov, and Alexander J Smola. Deep sets. In I. Guyon, U. V. Luxburg, S. Ben- gio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gamett (eds.). Advances in Neural Information Processing Systems, volume 30, pp. 3391-3401. Curran Asso- ciates, Inc., 2017. URL 77te44/47...2277e333335660000722
neurips neuri /tteaarpppscppppppr///////////
pdf
13

---

APPENDIX
A ANALYSIS OF REPRESENTATIONS LEARNED SCALOR
A.1 CLUSTERING z WHAT COMPONENTS
In this section, analyze the representations learned by SCALOR. First, looked how well different detections of the same object cluster together in the z what space SCALOR learns. This is important in order find whether objects which necessary to compute rewards for the agent (see Sec. 4.2) well separated z what space also indicates the usefulness SCALOR's representations for other potential downstream tasks such as classification. In Fig. 6 what space, and color each point according the mean pixel value of the foreground object in the crop detected by SCALOR. As one can see, the three objects (green, blue, and red points) and the robotic arm (darker red points) are quite well separated, with relatively low intra cluster variance. The robotic arm cluster shows larger variance it is observed in more different poses than the objects. There is also small cluster of misdetections (center top), with the gray color of the table. Overall, this shows that the z what space is well suited for the purpose of matching.
distances in z whaa space match corresponding
we plot the first and second principal component of points in
Analysis of clustering in SCALOR's what space
2nd PCA dimension of z what
1st PCA dimension of z what
Figure 6: First and second PCA dimension of 2 what space of SCALOR trained Visual Rearrange with objects. The plot shows 3000 random what points collected from a random policy. Each point colored the mean of the foreground pixels the crop detected by SCALOR. For each cluster, the highlighted point shows an example crop. Dashed lines indicate the Voronoi partitions according to cluster centers found by running k-means clustering. Figure is best viewed screen.
A.2 DISENTANGLEMENT ANALYSIS REPRESENTATION LEARNED BY SCALOR AND VAE
After seeing that SCALOR representation can be successfully used for object classification, further examined the quality the object location information learned by SCALOR by evaluating how disentangled they are. For this, we computed Mutual Information Gap (MIG) (Chen al. 2018) scores for SCALOR and VAE components. As SCALOR representations are unordered sets of vectors, used the clusters obtained from the cluster analysis (see App. A. to produce vector where that has consistent dimension ordering by matching z what components to clusters. In the case Zyec
of object not being recognized in image, imputed zeros values to its part in the vector 2
14

---

We estimated MIG by adapting the di disentanglement 1ib (Locatello et al. 2019), with an additional discretization of the continuous ground truth factors in the same way the continuous latent space is discretized.
The results in Fig. 7 show that SCALOR's s z where components are more disentangled and thus are better suited for the construction of independent RL sub-tasks. In addition, it can be that the VAE disentanglement score is quite low, potentially because different factors variation (object coordinates) have the same variance and thus could be more difficult to disentangle (Rolinek et al. 2019).
2.5
0.3
Whooe SCALOR features
0
2.0
1.5
MIG
0.2
1.0
0.1
0.5
-0.0
6 7
0.0
SCALOR
VAB
GT coordinates
(a) MIG score (higher is hetter).
(b) Mutual information matrix for SCALOR representations.
Figure 7: Comparison of VAE and SCALOR representations. (a) shows MIG scores of VAE and SCALOR representations on data obtained from running a random policy in the Visual Rearrange environment with 3 objects (with whisker showing the standard deviation over runs), (b) shows the mutual information matrix for SCALOR representations on the same data.
A.3 SCALOR TRAJECTORY TRAVERSALS
One of the ways to evaluate the quality of learned representations to show how reconstructs the scene. To this end, Fig. Big.Bshoo shows some example environment traversals and how SCALOR processes them. SCALOR is not only able to reconstruct the final image, but in addition is also able to locate objects and produce accurate segmentation masks for each object.
B ABLATION ANALYSIS OF GOAL-CONDITIONED ATTENTION POLICY
To understand how important the contribution of the goal-conditioned attention policy is to the performance and the generalization properties of our architecture, we have compared it with several other options for processing the set of SCALOR representations. In particular, test more variants of our attention mechanism: one where we use only goal-conditional attention heads, and one where we use only sanundunooiiio heads with learned, input-independent queries. We hypothesize that using only goal-conditional heads reduces the ability of the policy to easily concentrate on parts of the environment that are globally relevant for all tasks. Using only goal-unconditional heads with learned queries should still allow the policy to learn to order the input representations and produce consistent fixed length vector; however, removes the ability flexibly select parts of the inputs based on the task at hand. Finally, we also implemented the DeepSet method (Zaheer et al. (017) as an alternative approach to process inputs of sets of vector representation. In our case, we instantiate DeepSets by transforming each component embedding with activation to feature vectors of dimensionality 128 and then summing up these vectors.
hidden layer MLP with ReLU
The results in Fig. 9 show that both types of attention heads are necessary achieve the best results, with the goal-conditional heads having a larger impact on the final performance. Without the goal- conditional heads, the SMORL algorithm performs significantly worse. In addition, observe that SMORL with DeepSets can also perform competitively on the two objects tasks, however, is significantly worse the -fitt--diitiibutiinn task with one object.
15

---

Visual Rearrange environment with 2 objects
a)
b)
d)
Visual Rearrange environment with 3 objects
d)
Figure 8: Reconstructions of scene observations using learned SCALOR representation and decoder. Rows are a) original images (green boxes for recognized objects, red boxes for pooppropdtteed objects), b) full reconstructions, e) bounding boxes of recognized objects produced using where d) foreground object reconstructions, e) segmentation masks of objects generated by SCALOR.
C LONGER TRAINING FOR VISUAL REARRANGE WITH TWo OBJECTS
For the challenging Visual Rearrange environment with 2 objects, we trained a SMORL agent for twice as long as in the main plot in Fig. 4 to better understand the final convergence performance (see Fig. 10 Whereas the RIG baseline still shows no signs of progress after one million timesteps, SMORL agent is continuing to improve performance. This result hints at that with more training steps, SMORL might eventually reach the performance of a SAC agent that has privileged information of the ground-truth state.
16

---

2 objects
objects policy in object environment
0.18
Avg. object dist.
Avg. object dist.
0.17
0.16
0.15
0.16
0.14
0.13
0.15
0.12
Training examples
x10
Training examples
x100
SMORL
SMORL with DeepSet
SMORL only unconditional head
SMORL only goal-conditonal heads
Figure 9: Ablation study of goal-conditioned attention policy on Visual Rearrange with two objects (left) and out-of-distribtionn testing Visual Rearrange with object (right). We compare variants of the attention policy with only goal-conditional and only goal-unconditional attention heads, plus an alternative approach aggregate of vector representations in the form of DeepSets (Zaheer et al. 2017). Our results demonstrate that both types of attention heads are necessary to achieve the best results.
2 objects
0.0
0.2
0.4
0.6
0.8
Timesteps
x10
SMORL
RIG
SAC+GT
Passive policy
Figure 10: Performance of a SMORL agent trained for 10t timesteps on Visual Rearrange with objects.
D IMPLEMENTATION DETAILS
D.1 DETAILS OF ATTENTION MECHANISM
As discussed in Sec. 1. the policy should be able to select from the set of input representations based on the goal it needs to solve. We implement this running attention with goal-dependent query Q(z,) 2, W However, there might be some parts of the input state that are always relevant to the policy, regardless of the current goal. For example, in our experiments, we expect the state of the robotic arm to always be important, as it needed to manipulate objects. To simplify the extraction of this information for the policy, we optionally add M learned, input-independent, goal-unconditional queries Q(P9) to the goal-dependent query. Pa Md is simply matrix of parameters that is trained via backpropagation; we initialize Pa by sampling from (0, 0.02). Furthermore, we use two separate sets of attention heads to process the goal-conditional and -unconditional queries, i.e. each set of attention head has its own set of projections W W", W The output of both sets of attention heads is simply concatenated before feeding it to the next stage of the policy.
We use Pytorch's 's (Paszke et al. 2019) torch. nn Mult t whehhhhd waw.detttttt.t tent module to implement the attention mechanism. In practice, two separate instantiations of these modules to implement
17

---

goal-conditional and goal- cocl-ddncnii heads. In accordance to the original transformer attention formulation Vassaani et al. 2017), this module also includes a linear transformation that mixes the outputs of the different heads together. As think this transformation is not strictly necessary, have omitted for notational clarity. Moreover, note that we also linearly embed the policy inputs z and Re," 's into common space of dimensionality de before processing them further, which we have found to slightly improve performance.
D.2 FULL SMORL TRAINING AND EVALUATION ALGORITHMS
We display fully detailed version of the training algorithm in Alg 2 In addition, Alg. shows how we apply SMORL during evaluation. For evaluation, the agent receives a goal image to achieve from the environment. After processing this image into latent representations with SCALOR, the agent picks one of the recognized objects as its sub-goal and attempts to achieve for a fixed number of time steps. Following this, the agent sequentially moves that is not solved and repeats this process until either all goals are solved or the agent runs out of evaluation time steps. For our purpose, we define goal solved when the z whhere component of the best matching object from the observation is closer to the z whore component of the sub-goal than some threshold.
the next object in the goal image
Algorithm SMORL: Self-Supervised Multi-object RL (Training with Details)
Require: SCALOR encoder qo goal-conditioned policy me. goal-conditioned value function Qw, number of data points from random policy N, number training episodes K, number time steps in the episode If
1: Collect D = {o.) i= using random initial policy.
Train SCALOR on sequences data uniformly sampled from using loss described Eq.
Fit prior p(z where hhat to the latent encodings of observations { {(2 where
,2,
what
i=1
obtained using
(a( Zet;, 0c1)-
4: for ..., K episades do
for = I..... steps do 6: if i-==1 then
7:
Generate goal Mn (2gghr 2n (haa)) using SCALOR and initial observation (pick random detected object & and substitute z where by sampled from prior swhere ~ p where z )).
8:
end if
Encode z, using (Z Zet, Ogt)-
10:
action au ~ Iee (a: Zr, 2y).
11:
Execute a, and get next state observation o:41 from environment.
12:
Encode 2 using (s(2+41 igg,, 0g
13:
Store (zt, at, zt+1, 2g) into replay buffer R.
14:
With probability 0.5, replace 2e with a sample p(z where y wht)).
D. Sample "imagined" goals
15:
Sample transition (z, z², .zg) R.
16:
Compute matching reward signal It r(x',(x, zg) using Eq. 7
Minimize Bellman Error using (z, a, Zy, R). end for
18:
19:
for I t,.., H steps do
20:
Sample future state Oh; that has matching component in observation representation set Za, to the original goal zg. < h2 H
Sample HER "future" goals Store (zi, at, Zi+1, k) into (for such that z is matching the original goal Zg).
21:
22:
end for
23:
end for
D.3 SCALOR
We are using the SCALOR implementation from the original authors The parameters that were modified from the default settings can be found in Table In particular, equal to 8 for object and equal to for two objects. We observed that using smaller dimensionalities for what makes the training more stable, if it is possible train SCALOR with it. As for our purpose, the background model is not important and our environments have stable background, for this work, are modeling the background with small Zig
using z what dimension
https://gi httpg//thhh..oo//oooooooooooaaaa ang/ www.aa/ggaaaaaac SCALORR
18

---

Algorithm 3 SMORL (Evaluation)
Require:
Trained SMORL agent To, goal image og. SCALOR encoder qa. evaluation episode length L, sub-goal episode length /
1: Get goal representation 2, (zm ) m== (o(),, where N is the number of recognized objects.
2: Get the number of attempts K
3: Initialize goal index = 1.
4: Initialize evaluation step f 1. for I.,, steps Obtain initial observation to and pick sub-goal 2m-
7: for = 1,-., steps do
8: Encode using ( Z,,,, Z,t, 0gt)-
Get action a, [a, 2i, Zm..
Execute a, and get next observation 0:411 from environment. 1.
11:
12:
end for
13:
if all sub-goals 2m solved then
14:
Stop evaluation.
end if
16:
Set m = (m + 1) mod N.
17:
while is solved do
18:
Set (m + 1) mod N.
end while
20:
end for
During RL training, we process the first observation 0₁ of each episode 5 times with SCALOR which we found to stabilize the inferred representations, During evaluation, we do the same with the goal images Og given from the environment.
Hyper-parameter Optimizer
Value
Adam Kingma & Ba
2015
with default settings
Number of iterations
5000
Learning rate Batch size
0.0001
Explained Ratio Threshold
0.1
Number of training points
10000
Number of cells
4
Size bias
0.22
Size variance
0.12
Ratio bias
1.0
Ratio variance
0.3
Table SCALOR hyper--prameene
D.4 SMORL
We refer Table 2 for general papam-parameters of SMORL and to Table 3 for environment specific hyper-parameters of SMORL.
D.5
PRIOR WORK
For the baselines, iee. SAC, RIG and Skew-Fit, we started from standard settings and made environment-specifee tweaks tune them for best performance. In particular, significant hyper- parameter search effort (>500 runs) was spent on finding the best SAC parameters for Objjet Visual Rearrange 2, 3, and 4 objects.
19

---

Hyper-parameter
Value
Optimizer
Adam with default settings None (SAC policy is stochastic) 2048
Exploration Noise
RL Batch Size
Reward Scaling
Automatic SAC entropy tuning
yes 0.05
SAC Soft Update Rate
# Training Batches per Time Step
1
Hidden Activation
ReLU
Network Initialization
Xavier uniform
Separate Attention for Policy & Q-Function
yes 100000
Replay Buffer Size
Relabeling Fractions Rollout/Future/Imaginee Goals Number Initial Random Samples
0.1 0.4// 0.5 10000
Table 2: General
used by SMORI. for visual environments.
E PROBLEMS WITH SCALOR TRACKING DURING RL TRAINING
During our experimentation with the reward specification, we first considered SCALOR's internal tracking of objects. SCALOR assigns each discovered object ID, and these IDs principle propagated over time steps. By matching IDs, one can easily compute distances to the goal 2g in the space of the 2 where component (because we pick the episode goal from the objects discovered in the first observation during RL training). However, with such reward specification, the agent was easily finding ways to exploit the biases towards position in the propagation of the representation to the next time step.
Hlpeer-paraeeter
Push, Obj.
Push, 2 Obj. 15
Training Path Length
15
Evaluation Path Length
Learning Rate
0.001
0.0007
Discount Factor
0.925
0.95
Matching Threshold
1.2
1.3
No Match Reward Tno goal what Dim
0.75
1.0
Embedding Dim de Number of Cond./Uncond. Heads
48
32
3/0
Number of Input- iaput-Independen Queries
0
Policy Hidden Sizes
[128, 128)
[128, 128, 128)
Q-Function Hidden Sizes
[256, 256, 256]
128, 128, 128
parammerrr
Rearrange, Obj.
Rearrange, 2 Obj.
Training Path Length
20
20
Evaluation Path Length
60
100
Learning Rate
0.001
0.0005
Discount Factor
0.95
0.925
Matching Threshold a
1.2
1.3
No Match Reward Tho poal z what Dim
0.75
1.5
Embedding Dim de Number of Uncon/U Heads Number of Input- -nuut-Inddeedennt Queries
48
32
3/0
1/1
Policy Hidden Sizes
[64, 64]
1128, 128, 128) 128, 128, 128
Q-Function Hidden Sizes
[128, 128, 128]
Table 3: Environment specific hyper-parameters used by SMORL for visual environments.
20

---

In particular, one underlying assumption of SCALOR is that "two objects cannot coexist in the same position" " (Jiang et al. 2019). However, due to projjecct the 3D objects and possible occlusions, this assumption is not always fulfilled and the RL agent was able to exploit this during training. For example, the agent learned to position the robotic arm exactly above the object, and due to the positional propagation, this object's component was then propagated to the arm. After this, the agent was able to "manipulate" this component just by positioning his arm to the object's ''s goal. This shows the importance of evaluating learned representations in downstream tasks.
21


Reviews:
[{"review_id": "PgXrL96qlba", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\nThe paper proposes to use object-centric representations for RL, which can efficiently handle multiple objects in the scene. To learn a policy that can take a variable number of object observations, the paper proposes the goal-conditioned attention policy, which can focus on objects of interests to achieve each sub-goal, and thus reduce the combinatorial complexity of multiple objects. The goal-conditioned attention policy can be efficiently trained with hindsight experience replay on the object-centric goal representations. The experiments demonstrate the superior performance of the goal-conditioned attention policy on dealing with multiple objects.\n\n### Strengths\n- The idea of learning composable object-centric visual representations and goal-conditioned attention policy is an intuitive and plausible way to tackle combinatorial challenges in multi-object manipulation tasks.\n- The experiments with ground truth states show that the proposed goal-conditioned attention policy can effectively handle multiple object manipulation tasks. \n\n### Weaknesses\n\n- Based on Figure 4, any of the methods without ground truth states solve the tasks. Although the proposed method shows better learning performance in the Visual Rearranging task, the improvement is marginal to claim the proposed method can solve the tasks.\n- Why only up to 2 objects are considered in Figure 4? The proposed method has the advantage of dealing with multiple objects but did not show the benefit. It would be more convincing if the proposed method can reasonably deal with more than 2 objects.\n- The baseline could include recent visual policy learning methods using data augmentation [1,2].\n- The paper claims \"Self-supervised RL\" but it is not clear which part of the method is trained with self-supervised learning. The proposed method seems to consist of unsupervised representation learning and reinforcement learning.\n- The proposed method assumes that the sub-goals are independent of each other but it is not true in many cases, e.g., collisions between objects. \n- One of the claims in the paper is that the proposed representations and policy can work with a variable number of objects but the experiments do not cover this setup.\n\n### Questions and additional feedback\n- It would be better to include the architecture of the visual encoders for baselines and the proposed method.\n- It could be good to show the quality of learned representations, such as object pose prediction and classification.\n- The training time might be too short. The proposed method can be trained more (e.g., 1e6 environment steps).\n- The website link is provided but nothing is there.\n- How does the policy decide when to switch to the next sub-goal? Including how to rollout an episode toward sequential sub-goals with the proposed model would be helpful.\n\n### Overall assessment\nThe proposed method is intuitive and tackles an important problem of multi-object manipulation. However, the experiments and results are not yet convincing to claim the advantage in dealing with multiple objects. Overall, the reviewer thinks the paper requires more thorough experiments and is not ready to be published.\n\n\n[1] Laskin et al., Reinforcement Learning with Augmented Data\n\n[2] Kostrikov et al., Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels", "extra_scores": {}, "extra_text": {}}, {"review_id": "LmH31ZH8-OD", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work proposes to use object-centric unsupervised representation learning for self-supervised goal-conditioned RL, as opposed to prior work that assumes no particular structure on the learned representations (eg. VAEs). The proposed method, self-supervised multi-object RL (SMORL), uses the SCALOR architecture from prior work, then modifies the policy representation with single-object attention and also the reward function in RL with imagined goals (RIG). The results show that the method can learn simulated pushing and rearranging tasks in a self-supervised way with up to 4 objects in the scene, and outperforms RIG and Skew-Fit on pushing tasks. The proposed method is sufficiently novel, explores an important direction for self-supervised learning, and the results are quite strong.\n\nThe main motivation argued is that in more complex environments, it is difficult for fully unstructured reconstruction-based representation learning methods such as VAEs to recover a disentangled representation. This then causes difficulties running self-supervised RL algorithms such as RIG, which use the representation to compress the input, to set meaningful exploration goals, and for evaluating the reward. Using object-centric representations makes the representation more disentangled and improves RL, as demonstrated by the results. The paper would be strengthened by direct analysis of this hypothesis: for instance, correlations or measures of disentanglement between ground truth state and learned representation for SCALOR vs. VAE in the environments tested, particularly as number of objects increase.\n\nThe key novel contributions lie in how SCALOR is integrated with self-supervised learning. First, after learning the SCALOR representation from data, the proposed policy uses an attention mechanism to pay attention to reaching the goal for a single object at a time. One detail I did not understand was how the policy operates at test time. At training time, exactly one z^where is changed and the policy attempts to match that object. At test time, potentially many z^where could be different, so do you cycle between all the objects? Also, the paper would benefit from an ablation where you explore the choice of attention architecture; you could imagine that simply learning the object-centric representation and treating it as a flat representation like a VAE could also provide gains, so it is important to disentangle that effect from the novel policy architecture. The policy contribution is evaluated in Figure 3, where the success of SMORL+GT shows that the architecture makes manipulating a large number of objects possible.\n\nThe other differences to RIG have to do with goals and rewards. During self-supervised training, goals are sampled by sampling a new z^where for a single object, encouraging manipulation of exactly one object. This proposal seems logical, although in the long run it could potentially be an assumption that would not scale beyond object repositioning tasks. The reward is also modified to use the SCALOR latent, to penalize distance to the closest z^what object as the current goal with a threshold alpha for detecting the matching object. Again, the proposed reward function would be better evaluated with an ablation where you use the original reward from prior work -||z - z_g||.\n\nThe experiments show that the proposed method SMORL outperforms RIG and Skew-Fit on visual pushing tasks with many objects (and “rearranging”, which is pushing with random initial positions of objects - having both sets of experiments potentially seems a bit redundant since it seems rearranging is strictly more difficult). SMORL is worse than an oracle (SAC+GT) which uses ground truth state information, but it seems to tend towards the oracle performance on even the more difficult tasks (seeing longer versions of the learning curve that show asymptotic performance would help understand this better).\n\nGenerally, the results on multi-object manipulation and self-supervised learning are strong. Further experiments as mentioned above would better allow the contributions to understood independently.\n\nMinor comments\n\nThe related work is a bit thin and narrow - it addresses the nearest-neighbor works well but does not address self-supervised methods and robotics methods more broadly; it would be best to use the related work to make the paper more understandable to the broader community who are not embedded in goal-conditioned RL (and potentially put it in at section 2 instead of at the end).\n\nPage 6: “Out code” -> “our code”\n\n“In general, we expect that it is beneficial for the policy to not always attend to entities conditional on the goal; we thus allow some heads to only attend to additional learned parametric queries (left out above for notational clarity).” - did not understand this, it would be good to explain further what type of information this would include, and perhaps describe formally in the appendix.", "extra_scores": {}, "extra_text": {}}, {"review_id": "BdHSLGJnp3-", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\n\nThe paper combines an existing generative world model (SCALOR, Jiang et al. 2019) with goal-conditioned attention policy. The method is evaluated on object manipulation environments based on MuJoCo (Todorov et al., 2012), Multiworld (Nair et al. 2018) and a Sawyer arm. The paper is clearly written; the authors discuss challenges and motivate their design choices well throughout the paper.\n\nScore justification:\n\nThe paper is mostly incremental but it provides enough contributions for acceptance. SMORL (Algorithm 1; the proposed method) is well-defined and motivated. The method outperforms a strong baseline (Soft Actor-Critic with Hindsight Experience Replay) in an object manipulation task with available ground-truth state representations. In the visual rearranging task, the proposed method performs better than existing self-supervised RL algorithms: RIG (Nair et al., 2019) and Skew-Fit (Pong et al., 2020) when using 1 and 2 objects. Experiments with a larger set of objects which demonstrate the compositional benefit of SMORL would strengthen the paper. I would also be interested in a visualization of the latent object representations learned by SMORL.\n\nPros:\n\nThe paper contributes to improving scene decomposition and object representation learning in model-free RL which has practical applications in robotics and object-oriented RL.\n\nThere is a discussion of existing limitations and challenges (limitations of VAEs in visual RL, defining reward functions in goal-conditioned RL), and how SMORL is meant to address them (goal-conditioned attention policy to handle set inputs, incorporating goal and object representations in the reward). \n\nExperiments in the multi-object environments showing that SMORL might be promising with and without ground-truth representations.\n\nCons:\n\nThere is no discussion of the computation cost of SMORL in comparison to the baselines (SAC with HER, RIG, Skew-Fit).\n\nThe “compositional” aspect is unclear in the Experiments section. How does the “compositional generative world model” translate to productivity, substitivity or other forms of compositional generalization with respect to the objects in the image?\n\nExperiments with a larger set of objects would help in highlighting the advantage of SMORL in a general multi-object visual RL setting.\n\nQuestions during rebuttal period: \n\nPlease address and clarify the cons above  \n\nTypos and structure:\n\nThe link in the abstract leads to an empty page\n\nEquation 1: G (distribution over the set of goals) is once bolded, once in italic\n\nIt seems there are two different citation styles used in the paper: (Jiang et al., 2019) and Jiang et al. (2019)\n\nAlgorithm 1, line 1: typo “sequences data”\n\nSection header Experiments can be moved to the next page\n\nTypos: “our method scale challenging tasks”, “out code”, “objects identities”, “2 object”, “3 object”, “4 object”", "extra_scores": {}, "extra_text": {}}, {"review_id": "1KyMSB9iEek", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 9, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper combines object-centric representations and self-supervised HER goal-conditioned policy learning to learn efficient RL policies for a robot manipulation task.\n\nThey use SCALOR as an object-based state representation, and use it to propose semantically meaningful goals for a SAC policy to achieve. They can then leverage this on new tasks to solve them efficiently..\n\nOverall, I found this paper very interesting, clearly written, well executed (with very sensible decisions throughout) and presenting several good ideas, especially in how to leverage the additional object structure effectively. It demonstrates good early results in the novel field of Object-oriented RL.\n\nI have a few comments/questions:\n\n1. The explanation of how the goal-conditioned policies were trained was very clear, and I especially like how you use z_what and z_where to construct novel meaningful goals (which will tend to just force the policy to move objects around, but that is a good prior for your environment!). However how the “evaluation on a novel task” is done wasn’t as clear (i.e. when we try to implement a given goal or sequence of goals). More precisely, it is said in several places that the goal z_g is decomposed in sub-tasks where only one of the slots is used as the target. \n   1. Could you provide more details on exactly how that is done? Do you learn p(z^where) on task later?\n   2. What happens when some of the “objects” aren’t achievable or controllable? E.g. I’d expect that you see a slot which represents the robot arm, is this treated differently?\n2. How good are SCALOR representations in your environment? \n   1. It would be very helpful to show samples / traversals in the Appendix.\n   2. Similarly, comparing to the GT information you provide would be interesting (e.g. try to decode it? I understand you’d have to match the slots up unfortunately)\n   3. Did you try continuing training SCALOR in the RL phase? It would make the results stronger and less reliant on a good random exploration strategy.\n3. Did the hard matching cause issues while learning? I’d guess that the argmin is not too problematic because it is used in the reward computation only, but if you’d consider extending this setting to learning the goal proposal function z_g, this seems like a limitation?\n   1. The explanation about the issue of using tracking as part of the model directly wasn’t especially clear to me. It might deserve a bit more expansion, especially in the Appendix?\n4. How complex are the observations of the environment?\n   1. Could you add more samples of the environment’s observations?\n   2. It seems like the environment chosen is extremely similar to the standard Gym Fetch environment, did you try using it instead? https://gym.openai.com/envs/FetchPickAndPlace-v0/ \n5. It is not entirely true that MONet/IODINE “do not contain disentangled and interpretable features like position and scale”. \nIt is true that they are not explicitly enforced (like done in SCALOR), but they do arise quite easily purely unsupervised. \nEspecially, in my experience with both of these models, obtaining (and identifying) “position” latents is rather easy. See for example Figure 5 in [1] and this animation [2].\n\nSo in summary, I believe this is a strong paper in a budding field, which deserves publication at ICLR and may interest many people there.\n\n* [1] https://arxiv.org/abs/1901.11390\n* [2] https://twitter.com/cpburgess_/status/1091220207941701632 \n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a self supervised learning algorithm to compute object-centric representations for efficient RL in the context of robot manipulation tasks.

The key idea is to learn an object-centric representation (using prior work on SCALOR) and use this to intrinsically generate goals for a SAC policy to achieve. The policy is a goal-conditioned attention policy. The evaluation metric is a set of tasks to manipulate objects for a visual rearrangement task. 

${\bf Pros}: $
1. The baselines are reasonable and consist of other unsupervised RL algorithms in recent literature. 

2. Object-oriented RL is a growing area of interest and this paper proposes a reasonably novel and validated set of ideas in this domain. I believe it will be of significant interest and potentially make an impact on research in robotics and deep reinforcement learning.

3. The goal-conditioned attention policy can handle realistic scenarios, namely -- multi-object manipulation tasks

4. The attention mechanism also provides a reasonable solution to mitigate combinatorial hardness in multi-object environments

${\bf Cons}$: 

1. Some of the reviewers felt that the experimental results from pixel inputs could have been pushed further. However, since the setup and algorithm is relatively novel, there are already many moving parts and this paper seems like a step in that direction

2. Experiments with larger set of objects would have been interesting to investigate and report. 

   
Decision: Accept (Spotlight)