SELF-SUPERVISED VISUAL REINFORCEMENT LEARN-
Year: 2021

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/multiobject 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. https://martius-lab.github.io/SMORL

heading: 1 INTRODUCTION
text: 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 is 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) to map observations into a 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 VAEs could be harmful. In contrast, representing perceptual observations in terms of entities has been shown to improve data efficiency and transfer performance on a wide range of tasks (Burgess et al., 2019). Recent research has proposed a range of methods for unsupervised scene and video decomposition (Greff et al., 2017; Kosiorek et al., 2018; Burgess et al., 2019; Greff et al., 2019; Jiang et al., 2019; Weis et al., 2020; Locatello et al., 2020). These methods learn object representations and scene decomposition jointly. The majority of them are in part motivated by the fact that the learned representations are useful for downstream tasks such as image classification, object detection, or semantic segmentation. In this work, we show that such learned representations are also beneficial for autonomous control and reinforcement learning. ∗equal contribution
We propose to combine these object-centric unsupervised representation methods that represent the scene as a set of potentially structured vectors with goal-conditional visual RL. In our method (illustrated in Figure 1), dubbed SMORL (for self-supervised multi-object RL), a representation of raw sensory inputs is learned by a 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 can use those representations to learn how to achieve different goals with a reward function that utilizes the structure of the learned goal space. Our main contributions are as follows:
• 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 to autonomously
discover and learn useful skills in compositional environments with several objects using only images as inputs. • We show that even with fully disentangled ground-truth representation there is a large benefit
from using SMORL in environments with complex compositional tasks such as rearranging many objects.

heading: 2 RELATED WORK
text: Our work lies in the intersection of several actively evolving topics: visual reinforcement learning for control and robotics, and self-supervised learning. Vision-based RL for robotics is able to efficiently learn a variety of behaviors such as grasping, pushing and navigation (Levine et al., 2016; Pathak et al., 2018; Levine et al., 2018; 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 RL, self-supervision refers to the agent constructing its own reward signal and using it to solve self-proposed goals (Baranes & Oudeyer, 2013; Nair et al., 2018; Péré et al., 2018; Hausman et al., 2018; Lynch et al., 2019). This is especially relevant for visual RL, where a reward signal is usually not naturally available. These methods can potentially acquire a diverse repertoire of general-purpose robotic skills that can be reused and combined during test time. Such self-supervised approaches are crucial for scaling learning from narrow single-task learning to more general agents that explore the environment on their own to prepare for solving
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 et al., 2018; 2019; Pong et al., 2020; Ghosh et al., 2019; WardeFarley et al., 2019; Laversanne-Finot et al., 2018) 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 simultaneously 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 a given reward signal, whereas we are using the learned object-representations to construct a reward signal that helps to learn useful skills that can be used to solve multiple tasks. In addition, these methods use scene-mixture models such as MONET (Burgess 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 more 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.

heading: 3 BACKGROUND
text: 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 a Markov decision process defined by (S,A, p, r), where S and A are the continuous state and action spaces, p : S × S ×A 7→ [0,∞) is an unknown probability density representing the probability of transitioning to state st+1 ∈ S from state st ∈ S given action at ∈ A, and r : S 7→ R is a function computing the reward for reaching state st+1. The agent’s objective is to maximize the expected return R = ∑T t=1 Est∼ρπ,at∼π,st+1∼p [r(st+1)] over the horizon T , where ρπ(st) is the state marginal distribution induced by the agent’s policy π(at|st).

heading: 3.1 GOAL-CONDITIONAL REINFORCEMENT LEARNING
text: In the standard RL setting described before, the agent only learns to solve a single task, specified by the reward function. If we are interested in an agent that can solve multiple tasks (each with a different reward function) in an environment, we can train the agent on those tasks by telling the agent which distinct task to 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 π(at | st, g) on the goal g, and the agent’s objective turns to maximize the expected goal-conditional return:
Eg∼G [ T∑ t=1 Est∼ρπ,at∼π,st+1∼p [rg(st+1)] ] (1)
whereG is some distribution over the space of goals G ⊆ S the agent receives for training. The reward function can, for example, be the negative distance of the current state to the goal: rg(s) = −‖s−g‖. 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 as the goal. In this case, we have to provide a mapping m : S 7→ G of states to the desired goal space; the mapping is then used to compute the reward function, i.e. rg(s) = −‖m(s)− g‖.
As the reward is computed within the goal space, it is clear that the choice of goal space plays a 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 a meaningful signal towards reaching goals. However, if we only have access to high-dimensional, unstructured observations, e.g. camera images, and we 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 et al., 2018).
One option to deal with such difficult observation spaces is to learn a goal space in which the RL task becomes easier. For instance, we can try to find a low-dimensional latent space Z and use it 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 if 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 β-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, as we show in Sec. 5.2 and in App. A.2. Instead, we use a model explicitly geared towards inferring object-structured representations, which we introduce in the next section.

heading: 3.2 STRUCTURED REPRESENTATION LEARNING WITH SCALOR
text: SCALOR (Jiang et al., 2019) is a probabilistic generative world model for learning object-oriented representations of a video or stream of high-dimensional environment observations. SCALOR assumes that the environment observation ot at step t is generated by the background latent variable zbgt and the foreground latent variable z fg t . The foreground is further factorized into a set of object representations zfgt = {zt,n}n∈Ot , where Ot is the set of recognised object indices. To combine the information from previous time steps, a propagation-discovery model is used (Kosiorek et al., 2018). In SCALOR, an object is represented by zt,n = ( zprest,n , z where t,n , z what t,n ) . The scalar zprest,n defines if the object is present in the scene, whereas the vector zwhatt,n encodes object appearance. The component zwheret,n is further decomposed into the object’s center position z pos t,n, scale z scale t,n , and depth z depth t,n . With this, the generative process of SCALOR can be written as:
p(o1:T , z1:T ) = p(z D 1 )p(z bg 1 ) T∏ t=2 p(ot | zt)︸ ︷︷ ︸ rendering p(zbgt | zbg<t, zfgt )︸ ︷︷ ︸ background transition p(zDt | zPt )︸ ︷︷ ︸ discovery p(zPt | z<t)︸ ︷︷ ︸ propagation , (2)
where zt = (z bg t , z fg t ), z D t contains latent variables of objects discovered in the present step, and z P t contains latent variables of objects propagated from the previous step. Due to the intractability of the true posterior distribution p(z1:T |o1:T ), SCALOR is trained using variational inference with the following posterior approximation:
q(z1:T | o1:T ) = T∏ t=1 q(zt | z<t,o≤t) = T∏ t=1 q(zbgt | zfgt ,ot) q(zDt | zPt ,o≤t) q(zPt | z<t,o≤t), (3)
by maximizing the following evidence lower bound L(θ, φ) = T∑ t=1 Eqφ(z<t|o<t) [ Eqφ(zt|z<t,o≤t) [ log pθ(ot | zt) ] −DKL [ qφ(zt | z<t,o≤t) ‖ pθ(zt | z<t) ]] , (4)
where DKL denotes the Kullback-Leibler divergence. As we are using SCALOR in an active setting, we additionally condition the next step posterior predictions on the actions at taken by the agent. For more details and hyperparameters used to train SCALOR, we refer to App. D.3. In the next section, we describe how the structured representations learned by SCALOR can be used in downstream RL tasks such as goal-conditional visual RL.

heading: 4 SELF-SUPERVISED MULTI-OBJECT REINFORCEMENT LEARNING
text: 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 a goal-conditioned attention policy that can handle sets as inputs and flexibly learns to attend to those parts of the representation needed to 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, to 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 to 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 in Sec. 5.1, we show that even with fully disentangled, ground-truth representations of the scene, this is a challenging setting for state-of-the-art model-free RL 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.

heading: 4.1 POLICY WITH GOAL-CONDITIONED ATTENTION
text: We use the multi-head attention mechanism (Vaswani et al., 2017) as the first stage of our policy πθ 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 Q(zg) = zgW q. Each object is allowed to match with the query via an object-dependent key K(zt) = ztW
k and contribute to the attention’s output through the value V (zt) = ztW v, which is weighted by the similarity between Q(zg) and K(zt). As inputs, we concatenate the representations for object n to vectors zt,n = [zwhatt,n ; z where t,n ; z depth t,n ], and similarly the goal representation to zg = [zwhatg ; z where g ; z depth g ]. The attention head Ak is computed as
Ak = softmax ( zgW q(ZtW k)T√
de
) ZtW v, (5)
where Zt is a packed matrix of all zt,n’s, W q , W k, W v constitute learned linear transformations and de is the common key, value and query dimensionality. The final attention output A is a concatenation of all the attention heads A = [A1; . . . ;AK ]. In general, we expect it to be beneficial for the policy to not only attend to entities conditional on the goal; we thus let 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.1 and ablate the impact of different choices in App. B.
The second stage of our policy is a fully-connected neural network f that takes as inputs A and the goal representation zg and outputs an action at. The full policy πθ can thus be described by
πθ ({zt,n}n∈Ot , zg) = f(A, zg). (6)

heading: 4.2 SELF-SUPERVISED TRAINING
text: In principle, our policy can be trained with any goal-conditional model-free RL algorithm. For our experiments, we picked soft-actor critic (SAC) (Haarnoja et al., 2018b) as a state-of-the-art method for continuous action spaces, using hindsight experience replay (HER) (Andrychowicz et al., 2017) as a standard way to improve sample-efficiency in the goal-conditional setting.
The training algorithm is summarized in Alg. 1. We first train SCALOR on data collected from a random policy and fit a distribution p(zwhere) to representations zwhere of collected data. Each rollout, we generate a new goal for the agent by picking a random zwhat from the initial observation z1 and sampling a new zwhere from the fitted distribution p(zwhere). 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.
Algorithm 1 SMORL: Self-Supervised Multi-Object RL (Training) Require: SCALOR encoder qφ, goal-conditional policy πθ , goal-conditional SAC trainer, number of training
episodes K. 1: Train SCALOR on sequences uniformly sampled from D using loss described in Eq. 4. 2: Fit prior p(zwhere | zwhat) to the latent encodings of observations. 3: for n = 1, ...,K episodes do 4: Sample goal zg = ( ẑwhereg , z what g ) . 5: Collect episode data with policy πθ(at | zt, zg) and SCALOR representations of observations qφ(zt | z<t,o≤t). 6: Store transitions (zt,at, zt+1, zg) into replay bufferR. 7: Sample transitions from replay buffer (z,a, z′, zg) ∼ R. 8: Relabel zwhereg goal components to a combination of future states and p(zwhere | zwhat). 9: Compute matching reward signal R = r(z′, zg).
10: Update policy πθ(at | zt, zg) using R with SAC trainer. 11: end for We also refer to Alg. 2 in App. D.2 for a 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 to a single object, we have to extract the state of this object in the current observation in order to compute a reward. One way is to rely on the tracking of 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 a maximal reward that is not connected with environment changes, but rather with internal vision and tracking flaws (details in App. E).
We follow an alternative approach, namely to use the zwhat component of discovered objects and match them with the current goal representation zwhatg . As the z
what space encodes the appearance of objects, two detections corresponding to the same object should be close in this space (we verify that this hypothesis holds in App. A.1). Thus, it is easy to find the object corresponding to the current goal object using the distance mink ||zwhatk − zwhatg ||. In case of failure to discover a close representation, i.e. when all zwhatk have a distance larger than some threshold α to the goal representation z what g , we use a fixed negative reward rno_goal to incentivise the agent to avoid this situation.
Our reward signal is thus
r(z, zg) =
{ −||zwhere
k̂ − zwhereg || if mink ||zwhatk − zwhatg || < α,
rno_goal otherwise, (7)
where k̂ = argmink ||zwhatk − zwhatg ||.

heading: 4.3 COMPOSING INDEPENDENT SUB-GOALS DURING EVALUATION
text: At evaluation time, the agent receives a 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 are independent of each other and that the agent can thus sequentially achieve them by cycling through them until all of them are solved. The evaluation algorithm is summarized in Alg. 3, with more details added in App. D.2.
5 EXPERIMENTS
We have done computational experiments to address the following questions:
• How well does our method scale to challenging tasks with a large number of objects in case when ground-truth representations are provided?
• How does our method perform compared to prior visual goal-conditioned RL methods on image-based, multiobject continuous control tasks?
• How suitable are the representations learned by the compositional generative world model for discovering and solving RL tasks?
To answer these questions, we constructed the Multi-Object Visual Push and Multi-Object Visual Rearrange environments. Both environments are based on MuJoCo (Todorov et al., 2012) and the Multiworld package for image-based continuous control tasks introduced by Nair et al. (2018), and contain a 7-dof Sawyer arm where the agent needs to be controlled to manipulate a variable number 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 to 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 of the algorithms as the average 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.

heading: 5.1 SMORL WITH GROUND-TRUTH (GT) STATE REPRESENTATION
text: 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 as input. We are using a one-hot encoding for object identities zwhat and object and arm coordinates as zwhere components. With such a representation, the matching task becomes trivial, so our main focus in this experiment is 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 3 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.

heading: 5.2 VISUAL RL METHODS COMPARISON
text: 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 our method and RIG need to train a 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, once trained, SCALOR only adds little overhead compared to RIG’s VAE during RL training, and compared to Skew-Fit, our method is still faster to train as Skew-Fit needs to continuously retrain its VAE.
We show the results in Fig. 4. For the simpler Multi-Object Visual Push environment, the performance of SMORL is comparable to the best performing baseline, while for the more challenging MultiObject 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, our 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
partially observable from the point of view of the agent. On top of this, we 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 or improving SCALOR itself could help close the gap to an agent with perfect information.

heading: 5.3 OUT-OF-DISTRIBUTION GENERALIZATION FOR DIFFERENT NUMBER OF OBJECTS
text: One important advantage of structured policies is that they could potentially still be applicable for observations 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 et al., 2018). To see how our algorithm can generalize to a changing environment, we tested our SMORL agent trained on observations of the Rearrange environment with 2 objects on the environment with 1 object. As can be seen from Fig. 5, the performance of such an agent increases during training up to a performance comparable to a SMORL agent that was trained on the 1 object environment.

heading: 6 CONCLUSION AND FUTURE WORK
text: 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 a powerful way to acquire useful skills such as object manipulation. Our SMORL agent learns how to control different entities in the environment and can then combine the learned skills to 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 learned sub-tasks with a 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 et 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 et al., 2018) is an interesting direction for future research.

heading: ACKNOWLEDGEMENTS
text: 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 of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039B).

heading: A ANALYSIS OF REPRESENTATIONS LEARNED BY SCALOR
text: A.1 CLUSTERING OF zWHAT COMPONENTS
In this section, we analyze the representations learned by SCALOR. First, we looked at how well different detections of the same object cluster together in the zwhat space SCALOR learns. This is important in order to find out whether we can use distances in zwhat space to match corresponding objects which is necessary to compute rewards for the agent (see Sec. 4.2). A well separated zwhat space also indicates the usefulness of SCALOR’s representations for other potential downstream tasks such as classification. In Fig. 6, we plot the first and second principal component of points in zwhat space, and color each point according to 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 as it is observed in more different poses than the objects. There is also a small cluster of misdetections (center top), with the gray color of the table. Overall, this shows that the zwhat space is well suited for the purpose of matching.
A.2 DISENTANGLEMENT ANALYSIS OF REPRESENTATIONS LEARNED BY SCALOR AND VAE
After seeing that SCALOR representation can be successfully used for object classification, we further examined the quality of the object location information learned by SCALOR by evaluating how disentangled they are. For this, we computed Mutual Information Gap (MIG) (Chen et al., 2018) scores for SCALOR and VAE components. As SCALOR representations are unordered sets of vectors, we used the clusters obtained from the cluster analysis (see App. A.1) to produce a vector zwherevec that has consistent dimension ordering by matching z
what components to clusters. In the case of an object not being recognized in an image, we imputed zeros values to its part in the vector zwherevec .
We estimated MIG by adapting the disentanglement_lib (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 zwhere components are more disentangled and thus are better suited for the construction of independent RL sub-tasks. In addition, it can be seen that the VAE disentanglement score is quite low, potentially because different factors of variation (object coordinates) have the same variance and thus could be more difficult to disentangle (Rolinek et al., 2019).
A.3 SCALOR TRAJECTORY TRAVERSALS
One of the ways to evaluate the quality of learned representations is to show how it reconstructs the scene. To this end, Fig. 8 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.

heading: B ABLATION ANALYSIS OF GOAL-CONDITIONED ATTENTION POLICY
text: 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, we test two more variants of our attention mechanism: one where we use only goal-conditional attention heads, and one where we use only goal-unconditional 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 a consistent fixed length vector; however, it removes the ability to flexibly select parts of the inputs based on the task at hand. Finally, we also implemented the DeepSet method (Zaheer et al., 2017) as an alternative approach to process inputs of sets of vector representation. In our case, we instantiate DeepSets by transforming each component embedding with a one hidden layer MLP with ReLU activation to feature vectors of dimensionality 128 and then summing up these vectors.
The results in Fig. 9 show that both types of attention heads are necessary to achieve the best results, with the goal-conditional heads having a larger impact on the final performance. Without the goalconditional heads, the SMORL algorithm performs significantly worse. In addition, we observe that SMORL with DeepSets can also perform competitively on the two objects tasks, however, it is significantly worse on the out-of-distribution task with one object.

heading: C LONGER TRAINING FOR VISUAL REARRANGE WITH TWO OBJECTS
text: 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, our SMORL agent is continuing to improve performance. This result hints at that with even more training steps, SMORL might eventually reach the performance of a SAC agent that has privileged information of the ground-truth state.
D IMPLEMENTATION DETAILS
D.1 DETAILS OF ATTENTION MECHANISM
As discussed in Sec. 4.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 by running attention with a goal-dependent query Q(zg) = zgW
q . 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 is 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(P q) to the goal-dependent query. P q ∈ RM×de is simply a matrix of parameters that is trained via backpropagation; we initialize P q by sampling from N (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 q,W v,W k. 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 (Paszke et al., 2019) torch.nn.MultiheadAttention module to implement the attention mechanism. In practice, we use two separate instantiations of these modules to implement
goal-conditional and goal-unconditional heads. In accordance to the original transformer attention formulation (Vaswani et al., 2017), this module also includes a linear transformation that mixes the outputs of the different heads together. As we think this transformation is not strictly necessary, we have omitted it for notational clarity. Moreover, note that we also linearly embed the policy inputs zg and zt,n’s into a 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 a fully detailed version of the training algorithm in Alg. 2. In addition, Alg. 3 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 it for a fixed number of time steps. Following this, the agent sequentially moves on to the next object in the goal image 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 a goal as solved when the zwhere component of the best matching object from the observation is closer to the zwhere component of the sub-goal than some threshold.
Algorithm 2 SMORL: Self-Supervised Multi-object RL (Training with Details) Require: SCALOR encoder qφ, goal-conditioned policy πθ , goal-conditioned value function Qw, number of
data points from random policy N , number of training episodes K, number of time steps in the episode H 1: Collect D = {oi}Ni=1 using random initial policy. 2: Train SCALOR on sequences data uniformly sampled from D using loss described in Eq. 4. 3: Fit prior p(zwhere | zwhat) to the latent encodings of observations {( zwherei , z what i )}N i=1
obtained using qφ(zt | z<t,o≤t).
4: for n = 1, ...,K episodes do 5: for t = 1, ..., H steps do 6: if t = 1 then 7: Generate goal zg = ( ẑwhereg , z what g ) using SCALOR and initial observation o1 (pick random detected object k and substitute zwhere by sampled from prior ẑwhereg ∼ p(zwhere | zwhat)). 8: end if 9: Encode zt using qφ(zt | z<t,o≤t). 10: Get action at ∼ πθ(at | zt, zg). 11: Execute at and get next state observation ot+1 from environment. 12: Encode zt+1 using qφ(zt+1 | z≤t,o≤t+1). 13: Store (zt,at, zt+1, zg) into replay bufferR. 14: With probability 0.5, replace ẑwhereg with a sample p(zwhere | zwhat). B Sample “imagined” goals 15: Sample transition (z,a, z′, zg) ∼ R. 16: Compute matching reward signal R = r(z′, zg) using Eq. 7. 17: Minimize Bellman Error using (z,a, z′, zg, R). 18: end for 19: for l = t, ..., H steps do 20: Sample future state ohi that has matching component in observation representation set zhi to the original goal zg , l < hi ≤ H − 1. B Sample HER “future” goals 21: Store (zl,al, zl+1, zhi,k) intoR (for k such that zhi,k is matching the original goal zg). 22: end for 23: end for
D.3 SCALOR
We are using the SCALOR implementation from the original authors1. The parameters that were modified from the default settings can be found in Table 1. In particular, we are using zwhat dimension equal to 8 for 1 object and equal to 4 for two objects. We observed that using smaller dimensionalities for zwhat makes the training more stable, if it is possible to train SCALOR with it. As for our purpose, the background model is not important and our environments have stable background, for this work, we are modeling the background with small zbg = 1.
1https://github.com/JindongJiang/SCALOR
Algorithm 3 SMORL (Evaluation) Require: Trained SMORL agent πθ , goal image og , SCALOR encoder qφ, evaluation episode length L,
sub-goal episode length l 1: Get goal representation zg = {zm}Nm=1 = qφ(og) where N is the number of recognized objects. 2: Get the number of attempts K = L
l .
3: Initialize goal index m = 1. 4: Initialize evaluation step t = 1. 5: for k = 1, ...,K steps do 6: Obtain initial observation o1 and pick sub-goal zm. 7: for s = 1, ..., l steps do 8: Encode zt using qφ(zt | z<t,o≤t). 9: Get action at ∼ πθ(at | zt, zm).
10: Execute at and get next observation ot+1 from environment. 11: Set t = t+ 1. 12: end for 13: if all sub-goals zm are solved then 14: Stop evaluation. 15: end if 16: Set m = (m+ 1) mod N . 17: while zm is solved do 18: Set m = (m+ 1) mod N . 19: end while 20: end for
During RL training, we process the first observation o1 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.
D.4 SMORL
We refer to Table 2 for general hyper-parameters of SMORL and to Table 3 for environment specific hyper-parameters of SMORL.
D.5 PRIOR WORK
For the baselines, i.e. SAC, RIG and Skew-Fit, we started from standard settings and made environment-specific tweaks to tune them for best performance. In particular, significant hyperparameter search effort (>500 runs) was spent on finding the best SAC parameters for Multi-Object Visual Rearrange 2, 3, and 4 objects.

heading: E PROBLEMS WITH SCALOR TRACKING DURING RL TRAINING
text: During our experimentation with the reward specification, we first considered SCALOR’s internal tracking of objects. SCALOR assigns each discovered object an ID, and these IDs are in principle propagated over time steps. By matching IDs, one can easily compute distances to the goal zg in the space of the zwhere component (because we pick the episode goal from the objects discovered in the first observation during RL training). However, with such a reward specification, the agent was easily finding ways to exploit the biases towards a position in the propagation of the representation to the next time step.
In particular, one underlying assumption of SCALOR is that "two objects cannot coexist in the same position" (Jiang et al., 2019). However, due to 2D-projecting 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 goal. This shows the importance of evaluating learned representations in downstream tasks.

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

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.

The 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.