Self-supervised Visual Reinforcement Learning with Object-centric Representations

1 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 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 . 2 RELATED WORK . 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 . 3 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 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 ) . 3.1 GOAL-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 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 .

Summary: 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.