Learning Controllable Elements Oriented Representations for Reinforcement Learning 

1 INTRODUCTION . Deep Reinforcement Learning ( deep RL ) has proven its ability to solve difficult sequential decisionmaking problems such as Dota2 ( Berner et al. , 2019 ) and StarCraft ( Vinyals et al. , 2019 ) . However , it is still a challenge to apply deep RL to many real-world tasks , because the observations in these tasks are often high dimensional ( e.g . pixels ) and include much task-irrelevant noise . To address these issues , representation learning is introduced in deep RL to provide information-dense descriptions of the original observations , and thus reduce the difficulty of solving RL problems . Many representations learning algorithms are proposed with the design goal to make the representations be predictive of properties of future states . Such a goal is typically achieved by imposing the representations to reflect the dynamics of the environment ( Mazoure et al . ; Lee et al. , 2020b ; Schwarzer et al. , 2020 ; Gelada et al. , 2019 ) . However , the dynamics of the environment are not always task-relevant ( e.g . the noisy and varying background in control tasks ) . This can be problematic because those task-irrelevant dynamics will entice the representations into encoding task-irrelevant information ( e.g . the noisy background ) . Since the capacity of representations is limited , the obtained representations can miss important task-relevant information ( Gelada et al. , 2019 ) or encode too much unimportant task-irrelevant information ( Zhang et al. , 2020 ) . Therefore , it is essential to learn representations with more task-relevant environmental dynamics . We envision that action-related dynamics are likely to be more task-relevant because actions are the only feedback from the agent to the environment . For example , in control tasks and robotics particularly , the agent can complete its task only if it can realize which part of the environment can be influenced by its actions . For simplicity , we use controllable elements to denote those environment elements that can be influenced by the agent ’ s actions and thus reflect the action-related dynamics . We argue that these controllable elements should be explicitly retained for better representations . Following the intuition above , we develop an efficient and flexible framework named Learning Controllable Elements oriented Representations ( LCER ) for reinforcement learning to capture the controllable elements of the environment . Specifically , we formally define a metric to measure the amount of controllable elements that representations capture , and then propose a surrogate objective that is derived from this metric and can make the representations encode controllable elements as well as sufficient information . By emphasizing the importance of the controllable elements , LCER can capture more task-relevant information automatically . We evaluate our method LCER on the DMControl Suite ( Tassa et al. , 2018 ) , which has been widely used in recent years to benchmark sample efficiency ( Kostrikov et al. , 2020b ; Hafner et al. , 2019 ; Laskin et al . ; Lee et al. , 2020b ) . Our experiments show that LCER combined with SAC ( Haarnoja et al. , 2018 ) can outperform other representation learning methods on a majority of tasks , and even achieve better performance than the state-based SAC on the DMControl 100K benchmark . We also evaluate LCER on the Distracting Control Suite ( Stone et al. , 2021 ) , showing LCER ’ s ability to filter out distracting factors . The implementation of our work is available at https : //anonymous . 4open.science/r/LCER-FF40 2 RELATED WORK . In this section , we review the common approaches of representation learning in RL briefly . Predictive representations in RL . Predictive representations are typically trained to predict the properties of future states . One commonly used approach is to predict the representation of future states given the current state ’ s representation and actions . Such an approach is widely used in modelbased RL algorithms . The representations are prone to collapse if without further constraints ( Gelada et al. , 2019 ) , thus a reconstruction loss is usually applied to ensure non-trivial representations . However , this approach tries to encode everything , unable to filter out task-irrelevant information . SPR ( Schwarzer et al. , 2020 ) introduces a target encoder to avoid collapse and reconstruction , but its effect lacks theoretical guarantee , and it does not emphasize the controllable elements as LCER does . There are some methods that train predictive representations from information-theoretic perspectives , such as ( van den Oord et al. , 2018 ; Hjelm et al. , 2019 ; Anand et al. , 2019 ; Lee et al. , 2020b ; Mazoure et al. ) . These methods do not give a bias explicitly towards what kind of environment ’ s dynamics should the representations reflect , while our method LCER focuses more on the actionrelated dynamics which is more likely to be task-relevant . Our experiments in Section 5.2 shows that encoding predictive information does not ensure representations of high quality in complex environments , and it is crucial to retain controllable elements explicitly . There are some prior works that also emphasize controllable elements . Typically , ( Zhang et al. , 2018 ; pok , 2016 ; Badia et al. , 2020 ) learn an inverse model that predicts the action given successive states . However , these approaches typically rely on a reconstructive loss , which will significantly harm their performance in complex environments . There are also some works that emphasize the importance of actions in representation learning by introducing various action-related objectives such as policy selectivity ( Bengio et al. , 2017 ) and contingency awareness ( Bellemare et al. , 2012 ) . However , their methods are often built on heuristics and do not give a clear definition of controllable elements . Compared with these works , LCER firstly introduces a novel measure of controllable elements ( Eq . ( 3 ) ) and then develops an efficient method to optimize it without the aid of reconstruction loss . Bisimulation-based Representations in RL . Bisimulation-based representations group those states into a cluster which are indistinguishable w.r.t . future return , given any action sequence tested . The indistinguishability of two states is typically measured by the bisimulation metrics , which are firstly introduced in ( Ferns & doina Precup , 2014 ) . Recently , ( Zhang et al. , 2020 ) proposes DBC , a representation learning algorithm based on bisimulation that presents impressive results in distracting environments . However , these methods rely on a well-defined reward function , which is probably unavailable in real-world tasks . Besides , they often fail to solve simple tasks efficiently such as ones from DMControl Suite ( Zhang et al. , 2020 ) . Prior-knowledge-based Representations in RL There are other methods that use prior knowledge to constraint the representations ( Jonschkowski & Brock , 2015 ; Jonschkowski et al. , 2017 ; Thomas et al. , 2018 ) . For example , ( Jonschkowski & Brock , 2015 ) forces the representations to satisfy multiple priors such as proportionality and repeatability , which generally describe how the optimal representations should look like . However , such priors need to be specially designed for each environment , and much domain knowledge is required to form good representations . These drawbacks significantly limit the applications of these methods . 3 PRELIMINARIES . 3.1 NOTATION . We assume the underlying environment is a Markov decision process ( MDP ) , described by the tuple M = ( S , A , P , R , γ ) , where S is the state space , A the action space , P : St × At × St+1 → [ 0 , 1 ] the transition probability function which determines the distribution of next state given current state and action , and γ ∈ [ 0 , 1 ] a discount factor . Given the current state s ∈ S , an agent choose its action a ∈ A according to a policy function a ∼ π ( ·|s ) . This action will update the system state to a new state s′ according to the transition function P , and then a reward r = R ( s , a , s′ ) ∈ R is given to the agent . The goal of the agent is to maximize the expected cumulative rewards by learning a policy π . We denote the state at time step t by st. We use the upper letter of s ( i.e . S ) to refer to the random variable of state if there is no ambiguity . We use st : t+k to present the state sequence from t to t+ k , i.e . st : t+k : = ( st , st+1 , ... , st+k ) . These notations can be extended to other variables such as actions and rewards . We use φ : S → Z to denote the embedding function of LCER , which maps s ∈ S to a latent space Z . 3.2 MUTUAL INFORMATION . The mutual information ( MI ) of two random variables X , Y is defined as : I ( X ; Y ) = DKL ( PXY ‖PXPY ) , ( 1 ) where DKL is the Kullback–Leibler divergence that measures the difference of two distributions . The conditional mutual information ( CMI ) I ( X ; Y |Z ) is defined as the expected value of I ( X ; Y ) given the value of Z : I ( X ; Y |Z ) = Ez∼ZI ( X|z ; Y |z ) . ( 2 ) 4 LEARNING CONTROLLABLE ELEMENTS ORIENTED REPRESENTATIONS . Given an embedding function : φ : S → Z , we argue that the number of controllable elements that φ contains can be measured by the conditional mutual information ( CMI ) : Ik ( φ ) : = I ( φ ( St+k ) ; At : t+k−1|φ ( St ) ) ( 3 ) Intuitively , Ik measures the information that φ ( St+k ) and At : t+k−1 shares given φ ( St ) . Note that the non-controllable elements captured by φ do not contribute to Ik because they are not related to the actions . Training representations to maximize Eq . ( 3 ) is an efficient way to make the representations capture controllable elements . However , representations obtained in this way may be too compressive1 to provide sufficient information for solving tasks in some cases , because not all task-relevant elements are controllable . For example , in navigation tasks , the position of the goal is task-relevant but not controllable by the agent . Therefore , we need to utilize Eq . ( 3 ) in a more delicate manner . Specifically , in Section 4.1 we propose a surrogate objective derived from Eq . ( 3 ) which additionally provides a mechanism to control the degree of compression . In Section 4.2 , we describe how to optimize this objective , leading to our final algorithm LCER . 4.1 THE OPTIMIZATION OBJECTIVE FOR LCER . By the chain rule of mutual information , we can break the Ik defined in Eq . ( 3 ) into the subtraction of two terms : Ik ( φ ) = I ( [ φ ( St ) , At : t+k−1 ] ; φ ( St+k ) ) − I ( φ ( St ) ; φ ( St+k ) ) . The first term tries to capture predictive information , while the second term serves as a compression term that aims to filter out action-irrelevant information . To make φ capture the controllable elements and maintain sufficient information at the same time , we introduce a hyper-parameter β to control the degree of compression : Ik ( φ ; β ) : = I ( [ φ ( St ) , At : t+k−1 ] ; φ ( St+k ) ) − βI ( φ ( St ) ; φ ( St+k ) ) . ( 4 ) 1By saying “ compressive ” , we mean that these representations only encode controllable elements and filter out any other information . By choosing different β when training φ to maximize Eq . ( 4 ) , we can control the degree of compression , making it possible to adapt to different environments by choosing different β . However , directly maximizing Ik ( φ ; β ) is intractable , because maximizing the first term and minimizing the second term in Eq . ( 4 ) at the same time have an antagonistic effect to each other , thus it is difficult to tune in practice . In fact , we find this will often lead to unstable training process in our early experiments . To tackle this problem , we substitute I ( φ ( St ) ; φ ( St+k ) ) with Ik ( φ ; β = 0 ) − Ik ( φ ; β = 1 ) : Ik ( φ ; β ) = I ( [ φ ( St ) , At : t+k−1 ] ; φ ( St+k ) ) − β ( Ik ( φ ; β = 0 ) − Ik ( φ ; β = 1 ) ) = I ( [ φ ( St ) , At : t+k−1 ] ; φ ( St+k ) ) ︸ ︷︷ ︸ J1k · ( 1− β ) + I ( φ ( St+k ) ; At : t+k−1|φ ( St ) ) ︸ ︷︷ ︸ J2k ·β . ( 5 ) In Eq . ( 5 ) , we maximize two terms ( J1k and J 2 k ) at the same time , and their coefficients are 1 − β and β respectively . The first term J1k is closely related to the predictive information encoded by φ . The second term J2k is related to the controllable elements , which will play an important role in distracting environments as shown in Section 5.3 . In practice , we sum over different choices of k , leading to the final objective to maximize : IK ( φ ; β ) : = K∑ k=1 Ik ( φ ; β ) . ( 6 )

Summary: This paper proposes a novel representation learning flamework, LCER, which enables the agent to extract task-relevant representations from image-based observations. It optimizes the surrogate loss of conditional mutual information (CMI) between future-state embedding and action sequences given past-state embedding. The key contribution seems the balancing between InfoNCE objective (similar to CURL) and JSD-based CMI. The experimental comparison shows that LCER performs best among some image-based RL algorithms (PI-SAC, SLAC. CURL, PlaNet); in 100k DM Control benchmark (6 tasks), and distracted-background settings.