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: This paper proposes the use of state control as intrinsic motivation. It does so by separating the overall state into an agent state and a surrounding state. The idea then is to maximize the mutual information between the agent's internal state and the environment state. This mutual information is given as a reward to the agent which it attempts to maximize using policy gradient algorithms. Experiments attempt to validate the usefulness of such an intrinsic reward signal for pre-training an agent, or for augmenting a given task reward.