Information Prioritization through Empowerment in Visual Model-based RL

1 INTRODUCTION . Model-based reinforcement learning ( RL ) provides a promising approach to accelerating skill learning : by acquiring a predictive model that represents how the world works , an agent can quickly derive effective strategies , either by planning or by simulating synthetic experience under the model . However , in complex environments with high-dimensional observations ( e.g. , images ) , modeling the full observation space can present major challenges . While large neural network models have made progress on this problem ( Finn & Levine , 2017 ; Ha & Schmidhuber , 2018 ; Hafner et al. , 2019a ; Watter et al. , 2015 ; Babaeizadeh et al. , 2017 ) , sample-efficient learning necessitates some mechanism to prioritize modeling latent representations from observations such that functionally-relevant factors for the task can be captured . This needs to be done without wasting effort and capacity on irrelevant distractors , and without detailed reconstruction . Several recent works have proposed contrastive objectives that maximize mutual information between observations and latent states ( Hjelm et al. , 2018 ; Ma et al. , 2020 ; Oord et al. , 2018 ; Srinivas et al. , 2020 ) . While such objectives avoid reconstruction , they still do not distinguish between relevant irrelevant factors of variation . We thus pose the question : can we devise non-reconstructive representation learning methods that explicitly prioritize information that is most likely to be functionally relevant to the agent ? In this work , we derive a model-based RL algorithm from a combination of representation learning via mutual information maximization ( Poole et al. , 2019 ) and empowerment ( Mohamed & Rezende , 2015 ) . The latter serves to drive both the representation and the policy toward exploring and representing functionally relevant factors of variation . By integrating an empowerment-based term into a ∗Work done during Homanga ’ s research internship at Google . hbharadh @ cs.cmu.edu mutual information framework for learning state representations , we effectively prioritize information that is most likely to have functional relevance , which mitigates distractions due to irrelevant factors of variation in the observations . By integrating this same term into policy learning , we further improve exploration , particularly in the early stages of learning in sparse-reward environments , where the reward signal provides comparatively little guidance . Our main contribution is InfoPower , a model-based RL algorithm for high-dimensional systems with image observations that integrates empowerment into a mutual information based , nonreconstructive framework for learning state space models . Our approach explicitly prioritizes information that is most likely to be functionally relevant , which significantly improves performance in the presence of time-correlated distractors ( e.g. , background videos ) , and also accelerates exploration in environments when the reward signal is weak . We evaluate the proposed objectives on a suite of simulated robotic control tasks with explicit video distractors , and demonstrate up to 20 % better performance in terms of cumulative rewards at 1M environment interactions with 30 % higher sample efficiency at 100k interactions . 2 PROBLEM STATEMENT AND NOTATION . A partially observed Markov decision process ( POMDP ) is a tuple ( S , A , T , R , O ) that consists of states s ∈ S , actions a ∈ A , rewards r ∈ R , observations o ∈ O , and a state-transition distribution T ( s′|s , a ) . In most practical settings , the agent interacting with the environment doesn ’ t have access to the actual states in S , but to some partial information in the form of observations O . The underlying state-transition distribution T and reward distribution R are also unknown to the agent . In this paper , we consider the observations o ∈ O to be high-dimensional images , and so , the agent should learn a compact representation space Z for the latent state-space model . The problem statement is to learn effective representations from observations O when there are visual distractors present in the scene , and plan using the learned representations to maximize the cumulative sum of discounted rewards , J = E [ ∑ t γ t−1rt ] . The value of a state V ( Zt ) is defined as the expected cumulative sum of discounted rewards starting at state Zt . We use q ( · ) to denote parameterized variational approximations to learned distributions . We denote random variables with capital letters and use small letters to denote particular realizations ( e.g. , zt denotes the value of Zt ) . Since the underlying distributions are unknown , we evaluate all expectations through Monte-Carlo sampling with observed state-transition tuples ( ot , at−1 , ot−1 , zt , zt−1 , rt ) . 3 INFORMATION PRIORITIZATION FOR THE LATENT STATE-SPACE MODEL . Our goal is to learn a latent state-space model with a representation Z that prioritizes capturing functionally relevant parts of observations O , and devise a planning objective that explores with the learned representation . To achieve this , our key insight is integration of empowerment in the visual model-based RL pipeline . For representation learning we maximize MI maxZ I ( O , Z ) subject to a prioritization of the empowerment objective maxZ I ( At−1 ; Zt|Zt−1 ) . For planning , we maximize the empowerment objective along with reward-based value with respect to the policy maxA I ( At−1 ; Zt|Zt−1 ) + I ( Rt ; Zt ) . In the subsequent sections , we elaborate on our approach , InfoPower , and describe lower bounds to MI that yield a tractable algorithm . 3.1 LEARNING CONTROLLABLE FACTORS AND PLANNING THROUGH EMPOWERMENT . Controllable representations are features of the observation that correspond to entities which the agent can influence through its actions . For example , in quadrupedal locomotion , this could include the joint positions , velocities , motor torques , and the configurations of any object in the environment that the robot can interact with . For robotic manipulation , it could include the joint actuators of the robot arm , and the configurations of objects in the scene that it can interact with . Such representations are denoted by S+ in Fig . 2 , which we can formally define through conditional independence as the smallest subspace of S , S+ ≤ S , such that I ( At−1 ; St|S+t ) = 0 . This conditional independence relation can be seen in Fig . 2 . We explicitly priori- tize the learning of such representations in the latent space by drawing inspiration from variational empowerment ( Mohamed & Rezende , 2015 ) . The empowerment objective can be cast as maximizing a conditional information term I ( At−1 ; Zt|Zt−1 ) = H ( At−1|Zt−1 ) −H ( At−1|Zt , Zt−1 ) . The first term H ( At−1|Zt−1 ) encourages the chosen actions to be as diverse as possible , while the second term−H ( At−1|Zt , Zt−1 ) encourages the representations Zt and Zt+1 to be such that the action At for transition is predictable . While prior approaches have used empowerment in the model-free setting to learn policies by exploration through intrinsic motivation ( Mohamed & Rezende , 2015 ) , we specifically use this objective in combination with MI maximization for prioritizing the learning of controllable representations from distracting images in the latent state-space model . We include the same empowerment objective in both representation learning and policy learning . For this , we augment the maximization of the latent value function that is standard for policy learning in visual model-based RL ( Sutton , 1991 ) , with maxA I ( At−1 ; Zt|Zt−1 ) . This objectives complements value based-learning and further improves exploration by seeking controllable states . We empirically analyze the benefits of this in sections 4.3 and 4.5 . In the Appendix A.1 we next describe two theorems regarding learning controllable representations . We observe that the max ∑ t I ( At−1 ; Zt|Zt−1 ) objective alone for learning latent representations Z , along with the planning objective provably recovers controllable parts of the observationO , namely S+ . This result in Theorem 1 is important because in practice , we may not be able to represent every possible factor of variation in a complex environment . In this situation , we would expect that when |Z| |O| , learning Z under the objective max ∑ t I ( At−1 ; Zt|Zt−1 ) would encode S+ . We further show through Theorem 2 that the inverse information objective alone can be used to train a latent-state space model and a policy through an alternating optimization algorithm that converges to a local minimum of the objective max ∑ t I ( At−1 ; Zt|Zt−1 ) at a rate inversely proportional to the number of iterations . In Section 4.3 we empirically show how this objective helps achieve higher sample efficiency compared to pure value-based policy learning . 3.2 MUTUAL INFORMATION MAXIMIZATION FOR REPRESENTATION LEARNING . Algorithm 1 : Information Prioritization in Visual Model-based RL ( InfoPower ) Initialize dataset D with random episodes . Initialize model parameters φ , χ , ψ , η. Initialize dual variable λ. while not converged do for update step c = 1 .. C do // Model learning Sample data { ( at , ot , rt ) } k+Lt=k ∼ D. Compute latents zt ∼ pφ ( zt|zt−1 , at−1 , ot ) . Calculate L based on section 3.4 . ( φ , χ , ψ , η ) ← ( φ , χ , ψ , η ) +∇φ , χ , ψ , ηL λ← λ−∇λL // Behavior learning Rollout latent plan , S ← S ∪ { zt , at , rt } V ( zt ) ≈ Eπ [ ln qη ( rt|zt ) + ln qψ ( at−1|zt , zt−1 ) ] Update policy π and value model end // Environment interaction for time step t = 0 .. T − 1 do zt ∼ pφ ( zt|zt−1 , at−1 , ot ) ; at ∼ π ( at|zt ) rt , ot+1 ← env.step ( at ) . end Add data D ← D ∪ { ( ot , at , rt ) Tt=1 } . end For visual model-based RL , we need to learn a representation space Z , such that a forward dynamics model defining the probability of the next state in terms of the current state and the current action can be learned . The objective for this is∑ t−I ( it ; Zt|Zt−1 , At−1 ) . Here , it denotes the dataset indices that determine the observations p ( ot|it ) = δ ( ot − ot′ ) . In addition to the forward dynamics model , we need to learn a reward predictor by maximizing ∑ t I ( Rt ; Zt ) , such that the agent can plan ahead in the future by rolling forward latent states , without having to execute actions and observe rewards in the real environment . Finally , we need to learn an encoder for encoding observations O to latents Z . Most successful prior works have used reconstructionloss as a natural objective for learning this encoder ( Babaeizadeh et al. , 2017 ; Hafner et al. , 2019b ; a ) . A reconstruction-loss can be motivated by considering the objective I ( O , Z ) and computing its BA lower bound ( Agakov , 2004 ) . I ( ot ; zt ) ≥ Ep ( ot , zt ) [ log qφ′ ( ot|zt ) ] +H ( p ( ot ) ) . The first term here is the reconstruction objective , with qφ′ ( ot|zt ) being the decoder , and the second term can be ignored as it doesn ’ t depend on Z . However , this reconstruction objective explicitly encourages encoding the information from every pixel in the latent space ( such that reconstructing the image is possible ) and hence is prone to not ignoring distractors . In contrast , if we consider other lower bounds to I ( O , Z ) , we can obtain tractable objectives that do not involve reconstructing high-dimensional images . We can obtain an NCE-based lower bound ( Hjelm et al. , 2018 ) : I ( ot ; zt ) ≥ Eqφ ( zt|ot ) p ( ot ) [ log fθ ( zt , ot ) − log ∑ t′ 6=t fθ ( zt , ot′ ) ] , where qφ ( zt|ot ) is the learned encoder , ot is the observation at timestep t ( positive sample ) , and all observations in the replay buffer ot′ are negative samples . fθ ( zt , ot′ ) = exp ( zTt Wθzt′ ) The lower-bound is a form of contrastive learning as it maximizes compatibility of zt with the corresponding observation ot while minimizing compatibility with all other observations across time and batch . Although prior work has explored NCE-based bounds for contrastive learning in RL ( Srinivas et al. , 2020 ) , to the best of our knowledge , prior work has not used this in conjunction with empowerment for prioritizing information in visual model-based RL . Similarly , the Nguyen-Wainwright-Jordan ( NWJ ) bound ( Nguyen et al. , 2010 ) , which to the best our knowledge has not been used by prior works in visual model-based RL , can be obtained as , I ( ot ; zt ) ≥ Eqφ ( zt|ot ) p ( ot ) [ fθ ( zt , ot ) ] − e −1Eqφ ( zt|ot ) p ( ot ) e fθ ( zt , ot ) , where fθ is a critic . There exists an optimal critic function for which the bound is tightest and equality holds . We refer to the InfoNCE and NWJ lower bound based objectives as contrastive learning , in order to distinguish them from a reconstruction-loss based objective , though both are bounds on mutual information . We denote a lower bound to MI by I ( ot , zt ) . We empirically find the NWJ-bound to perform slightly better than the NCE-bound for our approach , explained in section 4.5 .

Summary: This paper tackles the problem of prioritizing functionally relevant information from complex observations for model-based RL. To this end, previous work has proposed to replace the reconstruction loss with contrastive loss. Building on that, this paper introduces an additional empowerment objective that is used for both representation learning and policy learning. Experiments show that the proposed model outperforms baselines on a set of DeepMind Control tasks with custom background distractions, which include other visually similar but uncontrollable agents. It is also shown that the similarity between learned states can match well to the similarity between groundtruth simulator states, according to the proposed metric.