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: There has been a school thought lately in RL that good representations for control should focus on abstracting out on only the task-relevant features. In high-dimensional tasks, this is especially useful in making optimal use of a limited neural network capacity and may even help the agent be robust to distractors in the environment. The authors motivate an approach to learning task-relevant representations by having the representations capture only the controllable elements of the environment. Their approach is based on a mutual information loss between actions sequences and final states, and they validate their approach on a series of DMControl task where both high-dimensional pixel observations as well as ground truth latent state representations are available. The idea is to compare the performance of agents that attempt to learn representations of latent states from pixels and compare against an agent that has direct access to the ground truth latent state. In their empirical results they show that: - their new approach (called LCER) is able to outperform alternative approaches on a set of these tasks. On some tasks, it even outperforms an SAC agent with access to the ground truth states. - LCER is able to outperform alternatives when distracting visual features are introduced into these tasks - LCER's representations learned on a source task are able to transfer better to downstream tasks.