Mutual Information State Intrinsic Control

1 INTRODUCTION . Reinforcement learning ( RL ) allows an agent to learn meaningful skills by interacting with an environment and optimizing some reward function , provided by the environment . Although RL has achieved impressive achievements on various tasks ( Silver et al. , 2017 ; Mnih et al. , 2015 ; Berner et al. , 2019 ) , it is very expensive to provide dense rewards for every task we want the robot to learn . Intrinsically motivated reinforcement learning encourages the agent to explore by providing an “ internal motivation ” instead , such as curiosity ( Schmidhuber , 1991 ; Pathak et al. , 2017 ; Burda et al. , 2018 ) , diversity ( Gregor et al. , 2016 ; Haarnoja et al. , 2018 ; Eysenbach et al. , 2019 ) and empowerment ( Klyubin et al. , 2005 ; Salge et al. , 2014 ; Mohamed & Rezende , 2015 ) . Those internal motivations can be computed on the fly when the agent is interacting with the environment , without any human engineered reward . We hope to extract useful “ skills ” from those internally motivated agents , which could later be used to solve downstream tasks , or simply augment the sparse reward with those intrinsic rewards to solve a given task faster . Most of the previous works in RL model the environment as a Markov Decision Process ( MDP ) . In an MDP , we use a single state vector to describe the current state of the whole environment , without explicitly distinguishing the agent itself from its surrounding . However , in the physical world , there is a clear boundary between an intelligent agent and its surrounding . The skin of any mammal is an example of such boundary . The separation of the agent and its surrounding also holds true for most of the man-made agents , such as any mechanical robot . This agent-surrounding separation has been studied for a long time in psychology under the concept of self-consciousness . Self-consciousness refers that a subject knows itself is the object of awareness ( Smith , 2020 ) , effectively treating the agent itself differently from everything else . Gallup ( 1970 ) has shown that self-consciousness widely exists in chimpanzees , dolphins , some elephants and human infants . To equally emphasize the agent and its surrounding , we name this separation as agent-surrounding separation in this paper . The widely adopted MDP formulation ignores the natural agent-surrounding separation , but simply stacks the agent state and its surrounding state together as a single state vector . Although this formulation is mathematically concise , we argue that it is over-simplistic , and as a result , it makes the learning harder . With this agent-surrounding separation in mind , we are able to design a much more efficient intrinsically motivated RL algorithm . We propose a new intrinsic motivation by encouraging the agent to ∗Correspondence to : Rui Zhao { zhaorui.in.germany @ gmail.com } . perform actions such that the resulting agent state should have high Mutual Information ( MI ) with the surrounding state . Intuitively , the higher the MI , the more control the agent could have on its surrounding . We name the proposed method “ MUtual information-based State Intrinsic Control ” , or MUSIC in short . With the proposed MUSIC method , we are able to learn many complex skills in an unsupervised manner , such as learning to pick up an object without any task reward . We can also augment a sparse reward with the dense MUSIC intrinsic reward , to accelerate the learning process . Our contributions are three-fold . First , we propose a novel intrinsic motivation ( MUSIC ) that encourages the agent to have maximum control on its surrounding , based on the natural agentsurrounding separation assumption . Secondly , we propose scalable objectives that make the MUSIC intrinsic reward easy to optimize . Last but not least , we show MUSIC ’ s superior performance , by comparing it with other competitive intrinsic rewards on multiple environments . Noticeably , with our method , for the first time the pick-and-place task can be solved without any task reward . 2 PRELIMINARIES . For environments , we consider four robotic tasks , including push , slide , pick-and-place , and navigation , as shown in Figure 2 . The goal in the manipulation task is to move the target object to a desired position . For the navigation task , the goal is to navigate to a target ball . In the following , we define some terminologies . 2.1 AGENT STATE , SURROUNDING STATE , AND REINFORCEMENT LEARNING SETTINGS . In this paper , the agent state sa means literally the state variable of the agent . The surrounding state ss refers to the state variable that describes the surrounding of the agent , for example , the state variable of an object . For multi-goal environments , we use the same assumption as previous works ( Andrychowicz et al. , 2017 ; Plappert et al. , 2018 ) , which consider that the goals can be represented as states and we denote the goal variable as g. For example , in the manipulation task , a goal is a particular desired position of the object in the episode . These desired positions , i.e. , goals , are sampled from the environment . The division between the agent state and the surrounding state is naturally defined by the agentsurrounding separation concept introduced in Section 1 . From a biology point of view , a human can naturally distinguish its own parts , like hands or legs from the environments . Analog to this , when we design a robotic system , we can easily know what is the agent state and what is its surrounding state . In this paper , we use upper letters , such as S , to denote random variables and the corresponding lower case letter , such as s , to represent the values of random variables . We assume the world is fully observable , including a set of states S , a set of actionsA , a distribution of initial states p ( s0 ) , transition probabilities p ( st+1 | st , at ) , a reward function r : S ×A → R , and a discount factor γ ∈ [ 0 , 1 ] . These components formulate a Markov Decision Process represented as a tuple , ( S , A , p , r , γ ) . We use τ to denote a trajectory , which contains a series of agent states and surrounding states . Its random variable is denoted as T . 3 METHOD . We focus on agent learning to control its surrounding purely by using its observations and actions without supervision . Motivated by the idea that when an agent takes control of its surrounding , then there is a high MI between the agent state and the surrounding state , we formulate the problem of learning without external supervision as one of learning a policy πθ ( at | st ) with parameters θ to maximize intrinsic MI rewards , r = I ( Sa ; Ss ) . In this section , we formally describe our method , mutual information-based state intrinsic control ( MUSIC ) . 3.1 MUTUAL INFORMATION REWARD FUNCTION . Our framework simultaneously learns a policy and an intrinsic reward function by maximizing the MI between the surrounding state and the agent state . Mathematically , the MI between the surround- ing state random variable Ss and the agent state random variable Sa is represented as follows : I ( Ss ; Sa ) = KL ( PSsSa || PSs ⊗ PSa ) ( 1 ) = sup T : Ω→R EPSsSa [ T ] − log ( EPSs⊗PSa [ e T ] ) ( 2 ) ≥ sup φ∈Φ EPSsSa [ Tφ ] − log ( EPSs⊗PSa [ e Tφ ] ) : = IΦ ( S s ; Sa ) , ( 3 ) where PSsSa is the joint probability distribution ; PSs ⊗ PSa is the product of the marginal distributions PSs and PSa ; KL denotes the Kullback-Leibler ( KL ) divergence . MI is notoriously difficult to compute in real-world settings ( Hjelm et al. , 2019 ) . Compared to the variational information maximizing-based approaches ( Barber & Agakov , 2003 ; Alemi et al. , 2016 ; Chalk et al. , 2016 ; Kolchinsky et al. , 2017 ) , the recent MINE-based approaches have shown superior performance ( Belghazi et al. , 2018 ; Hjelm et al. , 2019 ; Velickovic et al. , 2019 ) . Motivated by MINE ( Belghazi et al. , 2018 ) , we use a lower bound to approximate the MI quantity I ( Ss ; Sa ) . First , we rewrite Equation ( 1 ) , the KL formulation of the MI objective , using the Donsker-Varadhan representation , to Equation ( 2 ) ( Donsker & Varadhan , 1975 ) . The input space Ω is a compact domain of Rd , i.e. , Ω ⊂ Rd , and the supremum is taken over all functions T such that the two expectations are finite . Secondly , we lower bound the MI in the Donsker-Varadhan representation with the compression lemma in the PAC-Bayes literature and derive Equation ( 3 ) ( Banerjee , 2006 ; Belghazi et al. , 2018 ) . The expectations in Equation ( 3 ) are estimated by using empirical samples from PSsSa and PSs ⊗ PSa . The statistics model Tφ is parameterized by a deep neural network with parameters φ ∈ Φ , whose inputs are the empirical samples . 3.2 EFFECTIVELY COMPUTING THE MUTUAL INFORMATION REWARD IN PRACTICE . Lemma 1 . There is a monotonically increasing relationship between Iφ ( Ss ; Sa | T ) and EPT ′ [ Iφ ( S s ; Sa | T ′ ) ] , mathematically , Iφ ( S s ; Sa | T ) n EPT ′ [ Iφ ( S s ; Sa | T ′ ) ] , ( 4 ) where Ss , Sa , and T denote the surrounding state , the agent state , and the trajectory , respectively . The trajectory fractions are defined as the adjacent state pairs , namely T ′ = { St , St+1 } . The symbol n denotes a monotonically increasing relationship between two variables and φ represents the parameter of the statistics model in MINE . Proof . See Appendix A . We define the reward for each transition at a given time-step as the mutual information of the pair of adjacent states at that time-step , see Equation ( 4 ) Right-Hand Side ( RHS ) . However , in practice , we find that it is not very efficient to train the MI estimator using state pairs . To counter this issue , we use all the states in the same trajectory in a batch to train the MI estimator , see Equation ( 4 ) LeftHand Side ( LHS ) , since more empirical samples help to reduce variance and therefore accelerate learning . In Lemma 1 , we prove the monotonically increasing relationship between Equation ( 4 ) RHS and Equation ( 4 ) LHS . In more detail , we divide the process of computing rewards into two phases , i.e. , the training phase and the evaluation phase . In the training phase , we efficiently train the MI estimator with a large batch of samples from the whole trajectory . For training the MI estimator network , we first randomly sample the trajectory τ from the replay buffer . Then , the states sat used for calculating the product of marginal distributions are sampled by shuffling the states sat from the joint distribution along the temporal axis t within the trajectory . We use back-propagation to optimize the parameter ( φ ) to maximize the MI lower bound , see Equation ( 4 ) LHS . For evaluating the MI reward , we use a pair of transitions to calculate the transition reward , see Equation ( 4 ) RHS and Equation ( 5 ) , instead of using the complete trajectory . Each time , to calculate the MI reward for the transition , the reward is calculated over a small fraction of the complete trajectory τ ′ , namely r = Iφ ( Ss ; Sa | T ′ ) . The trajectory fraction , τ ′ , is defined as adjacent state pairs , τ ′ = { st , st+1 } , and T ′ represents its corresponding random variable . The derived Lemma 1 brings us two important benefits . First , it enables us to efficiently train the MI estimator using all the states in the same trajectory . And a large batch of empirical samples reduce the variance of the gradients . Secondly , it allows us to estimate the MI reward for each transition with only the relevant state pair . This way of estimating MI enables us to assign rewards more accurately at the transition level . Based on Lemma 1 , we calculate the transition reward as the MI of each trajectory fraction , namely rφ ( at , st ) : = Iφ ( S s ; Sa|T ′ ) = 0.5 ∑t+1 i=tTφ ( s s i , s a i ) − log ( 0.5 ∑t+1 i=t e Tφ ( s s i , s̄ a i ) ) , ( 5 ) where ( ssi , s a i ) ∼ PSsSa|T ′ , s̄ai ∼ PSa|T ′ , and τ ′ = { st , st+1 } . In case that the estimated MI value is particularly small , we scale the reward with a hyper-parameter α and clip the reward between 0 and 1 . MUSIC can be combined with any off-the-shelf reinforcement learning methods , such as deep deterministic policy gradient ( DDPG ) ( Lillicrap et al. , 2016 ) and soft actor-critic ( SAC ) ( Haarnoja et al. , 2018 ) . We summarize the complete training algorithm in Algorithm 1 and in Figure 1 . MUSIC Variants with Task Rewards : The introduced MUSIC method is an unsupervised reinforcement learning approach , which is denoted as “ MUSIC-u ” , where “ -u ” stands for unsupervised learning . We propose three ways of using MUSIC to accelerate learning . The first method is using the MUSIC-u pretrained policy as the parameter initialization and then fine-tuning the agent with the task rewards . We denote this variant as “ MUSIC-f ” , where “ -f ” stands for fine-tuning . The second variant is to use the MI intrinsic reward to help the agent to explore more efficiently . Here , the MI reward and the task reward are added together . We name this method as “ MUSIC-r ” , where “ -r ” stands for reward . The third approach is to use the MI quantity from MUSIC to prioritize trajectories for replay . The approach is similar to the TD-error-based prioritized experience replay ( PER ) ( Schaul et al. , 2016 ) . The only difference is that we use the estimated MI instead of the TD-error as the priority for sampling . We name this method as “ MUSIC-p ” , where “ -p ” stands for prioritization . Skill Discovery with MUSIC and DIAYN : One of the relevant works on unsupervised RL , DIAYN ( Eysenbach et al. , 2019 ) , introduces an information-theoretical objective FDIAYN , which learns diverse discriminable skills indexed by the latent variable Z , mathematically , FDIAYN = I ( S ; Z ) + H ( A | S , Z ) . The first term , I ( S ; Z ) , in the objective , FDIAYN , is implemented via a skill discriminator , which serves as a variational lower bound of the original objective ( Barber & Agakov , 2003 ; Eysenbach et al. , 2019 ) . The skill discriminator assigns high rewards to the agent , if it can predict the skill-options , Z , given the states , S. Here , we substitute the full state S with the surrounding state Ss to encourage the agent to learn control skills . DIAYN and MUSIC can be combined as follows : FMUSIC+DIAYN = I ( Sa ; Ss ) +I ( Ss ; Z ) +H ( A | S , Z ) . The combined version enables the agent to learn diverse control primitives via skill-conditioned policy ( Eysenbach et al. , 2019 ) in an unsupervised fashion . Comparison and Combination with DISCERN : Another relevant work is Discriminative Embedding Reward Networks ( DISCERN ) ( Warde-Farley et al. , 2019 ) , whose objective is to maximize the MI between the state S and the goal G , namely I ( S ; G ) . While MUSIC ’ s objective is to maximize the MI between the agent state Sa and the surrounding state Ss , namely I ( Sa ; Ss ) . Intuitively , DISCERN attempts to reach a particular goal in each episode , while our method tries to manipulate the surrounding state to any different value . MUSIC and DISCERN can be combined as FMUSIC+DISCERN = I ( Sa ; Ss ) + I ( S ; G ) . Optionally , we can replace the full states S with Ss , since it performs better than with S empirically . Through this combination , MUSIC helps DISCERN to learn its discriminative objective .

Summary: The paper propose MUSIC, an RL algorithm for learning controllers in a unsupervised way. They key idea of the proposed algorithm is to separate the state of the robot like joint angles from the state of the environments such as location of an external object and optimize the mutual information between the two set of states. By maximizing this mutual information, the resulting policy learns to better control the environment and can be used to train downstream tasks. The paper experimented with different ways of training the downstream tasks and demonstrated favorable results compared to prior methods.