Learning Representations in Reinforcement Learning: an Information Bottleneck Approach

1 Introduction . In training a reinforcement learning algorithm , an agent interacts with the environment , explores the ( possibly unknown ) state space , and learns a policy from the exploration sample data . In many cases , such samples are quite expensive to obtain ( e.g. , requires interactions with the physical environment ) . Hence , improving the sample efficiency of the learning algorithm is a key problem in RL and has been studied extensively in the literature . Popular techniques include experience reuse/replay , which leads to powerful off-policy algorithms ( e.g. , ( Mnih et al. , 2013 ; Silver et al. , 2014 ; Van Hasselt et al. , 2015 ; Nachum et al. , 2018a ; Espeholt et al. , 2018 ) ) , and model-based algorithms ( e.g. , ( Hafner et al. , 2018 ; Kaiser et al. , 2019 ) ) . Moreover , it is known that effective representations can greatly reduce the sample complexity in RL . This can be seen from the following motivating example : In the environment of a classical Atari game : Seaquest , it may take dozens of millions samples to converge to an optimal policy when the input states are raw images ( more than 28,000 dimensions ) , while it takes less samples when the inputs are 128-dimension pre-defined RAM data ( Sygnowski & Michalewski , 2016 ) . Clearly , the RAM data contain much less redundant information irrelevant to the learning process than the raw images . Thus , we argue that an efficient representation is extremely crucial to the sample efficiency . In this paper , we try to improve the sample efficiency in RL from the perspective of representation learning using the celebrated information bottleneck framework ( Tishby et al. , 2000 ) . In standard deep learning , the experiments in ( Shwartz-Ziv & Tishby , 2017 ) show that during the training process , the neural network first ” remembers ” the inputs by increasing the mutual information between the inputs and the representation variables , then compresses the inputs to efficient representation related to the learning task by discarding redundant information from inputs ( decreasing the mutual information between inputs and representation variables ) . We call this phenomena ” information extraction-compression process ” ( information E-C process ) . Our experiments shows that , similar to the results shown in ( Shwartz-Ziv & Tishby , 2017 ) , we first ( to the best of our knowledge ) observe the information extraction-compression phenomena in the context of deep RL ( we need to use MINE ( Belghazi et al. , 2018 ) for estimating the mutual information ) . This observation motivates us to adopt the information bottleneck ( IB ) framework in reinforcement learning , in order to accelerate the extraction-compression process . The IB framework is intended to explicitly enforce RL agents to learn an efficient representation , hence improving the sample efficiency , by discarding irrelevant information from raw input data . Our technical contributions can be summarized as follows : 1 . We observe that the ” information extraction-compression process ” also exists in the context of deep RL ( using MINE ( Belghazi et al. , 2018 ) to estimate the mutual information ) . 2 . We derive the optimization problem of our information bottleneck framework in RL . In order to solve the optimization problem , we construct a lower bound and use the Stein variational gradient method developed in ( Liu et al. , 2017 ) to optimize the lower bound . 3 . We show that our framework can accelerate the information extraction-compression process . Our experimental results also show that combining actor-critic algorithms ( such as A2C , PPO ) with our framework is more sample-efficient than their original versions . 4 . We analyze the relationship between our framework and MINE , through this relationship , we theoretically derive an algorithm to optimize our IB framework without constructing the lower bound . Finally , we note that our IB method is orthogonal to other methods for improving the sample efficiency , and it is an interesting future work to incorporate it in other off-policy and model-based algorithms . 2 Related Work . Information bottleneck framework was first introduced in ( Tishby et al. , 2000 ) . They solve the framework by iterative Blahut Arimoto algorithm , which is infeasible to apply to deep neural networks . ( Shwartz-Ziv & Tishby , 2017 ) tries to open the black box of deep learning from the perspective of information bottleneck , though the method they use to compute the mutual information is not precise . ( Alemi et al. , 2016 ) derives a variational information bottleneck framework , yet apart from adding prior target distribution of the representation distribution P ( Z|X ) , they also assume that P ( Z|X ) itself must be a Gaussian distribution , which limits the capabilities of the representation function . ( Peng et al. , 2018 ) extends this framework to variational discriminator bottleneck to improve GANs ( Goodfellow et al. , 2014 ) , imitation learning and inverse RL . As for improving sample-efficiency , ( Mnih et al. , 2013 ; Van Hasselt et al. , 2015 ; Nachum et al. , 2018a ) mainly utilize the experience-reuse . Besides experience-reuse , ( Silver et al. , 2014 ; Fujimoto et al. , 2018 ) tries to learn a deterministic policy , ( Espeholt et al. , 2018 ) seeks to mitigate the delay of off-policy . ( Hafner et al. , 2018 ; Kaiser et al. , 2019 ) learn the environment model . Some other powerful techniques can be found in ( Botvinick et al. , 2019 ) . State representation learning has been studied extensively , readers can find some classic works in the overview ( Lesort et al. , 2018 ) . Apart from this overview , ( Nachum et al. , 2018b ) shows a theoretical foundation of maintaining the optimality of representation space . ( Bellemare et al. , 2019 ) proposes a new perspective on representation learning in RL based on geometric properties of the space of value function . ( Abel et al. , 2019 ) learns representation via information bottleneck ( IB ) in imitation/apprenticeship learning . To the best of our knowledge , there is no work that intends to directly use IB in basic RL algorithms . 3 Preliminaries . A Markov decision process ( MDP ) is a tuple , ( X , A , R , P , µ ) , where X is the set of states , A is the set of actions , R : X ×A×X → R is the reward function , P : X ×A×X → [ 0 , 1 ] is the transition probability function ( where P ( X ′ |X , a ) is the probability of transitioning to state X ′ given that the previous state is X and the agent took action a in X ) , and µ : X → [ 0 , 1 ] is the starting state distribution . A policy π : X → P ( A ) is a map from states to probability distributions over actions , with π ( a|X ) denoting the probability of choosing action a in state X . In reinforcement learning , we aim to select a policy π which maximizes K ( π ) = Eτ∼π [ ∑∞ t=0 γ tR ( Xt , at , Xt+1 ) ] , with a slight abuse of notation we denote R ( Xt , at , Xt+1 ) = rt . Here γ ∈ [ 0 , 1 ) is a discount factor , τ denotes a trajectory ( X0 , a0 , X1 , a1 , ... ) . Define the state value function as V π ( X ) = Eτ∼π [ ∑∞ t=0 γ trt|X0 = X ] , which is the expected return by policy π in state X . And the state-action value function Qπ ( X , a ) = Eτ∼π [ ∑∞ t=0 γ trt|X0 = X , a0 = a ] is the expected return by policy π after taking action a in state X. Actor-critic algorithms take the advantage of both policy gradient methods and valuefunction-based methods such as the well-known A2C ( Mnih et al. , 2016 ) . Specifically , in the case that policy π ( a|X ; θ ) is parameterized by θ , A2C uses the following equation to approximate the real policy gradient ∇θK ( π ) = ∇θĴ ( θ ) : ∇θĴ ( θ ) ≈ ∞∑ t=0 ∇θ [ log π ( at|Xt ; θ ) ( Rt − b ( Xt ) ) + α2H ( π ( ·|Xt ) ) ] = ∞∑ t=0 ∇θĴ ( Xt ; θ ) ( 1 ) where Rt = ∑∞ i=0 γ irt+i is the accumulated return from time step t , H ( p ) is the entropy of distribution p and b ( Xt ) is a baseline function , which is commonly replaced by V π ( Xt ) . A2C also includes the minimization of the mean square error between Rt and value function V π ( Xt ) . Thus in practice , the total objective function in A2C can be written as : J ( θ ) ≈ ∞∑ t=0 log π ( at|Xt ; θ ) ( Rt − V π ( Xt ) ) − α1 ∥Rt − V π ( Xt ) ∥2 + α2H ( π ( ·|Xt ) ) = ∞∑ t=0 J ( Xt ; θ ) ( 2 ) where α1 , α2 are two coefficients . In the context of representation learning in RL , J ( Xt ; θ ) ( including V π ( Xt ) and Qπ ( Xt , at ) ) can be replaced by J ( Zt ; θ ) where Zt is a learnable low-dimensional representation of state Xt . For example , given a representation function Z ∼ Pϕ ( ·|X ) with parameter ϕ , define V π ( Zt ; Xt , ϕ ) |Zt∼Pϕ ( ·|Xt ) = V π ( Xt ) . For simplicity , we write V π ( Zt ; Xt , ϕ ) |Zt∼Pϕ ( ·|Xt ) as V π ( Zt ) . 4 Framework . 4.1 Information Bottleneck in Reinforcement Learning The information bottleneck framework is an information theoretical framework for extracting relevant information , or yielding a representation , that an input X ∈ X contains about an output Y ∈ Y . An optimal representation of X would capture the relevant factors and compress X by diminishing the irrelevant parts which do not contribute to the prediction of Y . In a Markovian structure X → Z → Y where X is the input , Z is representation of X and Y is the label of X , IB seeks an embedding distribution P ⋆ ( Z|X ) such that : P ⋆ ( Z|X ) = arg max P ( Z|X ) I ( Y , Z ) − βI ( X , Z ) = arg max P ( Z|X ) H ( Y ) −H ( Y |Z ) − βI ( X , Z ) = arg max P ( Z|X ) −H ( Y |Z ) − βI ( X , Z ) ( 3 ) for every X ∈ X , which appears as the standard cross-entropy loss1 in supervised learning with a MI-regularizer , β is a coefficient that controls the magnitude of the regularizer . Next we derive an information bottleneck framework in reinforcement learning . Just like the label Y in the context of supervised learning as showed in ( 3 ) , we assume the supervising signal Y in RL to be the accurate value Rt of a specific state Xt for a fixed policy π , which can be approximated by an n-step bootstrapping function Yt = Rt = ∑n−2 i=0 γ irt+i + γn−1V π ( Zt+n−1 ) in practice . Let P ( Y |Z ) be the following distribution : P ( Yt|Zt ) ∝ exp ( −α ( Rt − V π ( Zt ) ) 2 ) ( 4 ) .This assumption is heuristic but reasonable : If we have an input Xt and its relative label Yt = Rt , we now have Xt ’ s representation Zt , naturally we want to train our decision function V π ( Zt ) to approximate the true label Yt . If we set our target distribution to be C · exp ( −α ( Rt − V π ( Zt ) ) 2 ) , the probability decreases as V π ( Zt ) gets far from Yt while increases as V π ( Zt ) gets close to Yt . For simplicity , we just write P ( R|Z ) instead of P ( Yt|Zt ) in the following context . With this assumption , equation ( 3 ) can be written as : P ⋆ ( Z|X ) = arg max P ( Z|X ) EX , R , Z∼P ( X , R , Z ) [ logP ( R|Z ) ] − βI ( X , Z ) = arg max P ( Z|X ) EX∼P ( X ) , Z∼P ( Z|X ) , R∼P ( R|Z ) [ −α ( R− V π ( Z ) ) 2 ] − βI ( X , Z ) ( 5 ) The first term looks familiar with classic mean squared error in supervisd learning . In a network with representation parameter ϕ and policy-value parameter θ , policy loss Ĵ ( Z ; θ ) in equation ( 1 ) and IB loss in ( 5 ) can be jointly written as : L ( θ , ϕ ) = EX∼P ( X ) , Z∼Pϕ ( Z|X ) [ Ĵ ( Z ; θ ) + ER [ −α ( R− V π ( Z ; θ ) ) 2 ] ︸ ︷︷ ︸ J ( Z ; θ ) ] − βI ( X , Z ; ϕ ) ( 6 ) where I ( X , Z ; ϕ ) denotes the MI between X and Z ∼ Pϕ ( ·|X ) . Notice that J ( Z ; θ ) itself is a standard loss function in RL as showed in ( 2 ) . Finally we get the ultimate formalization of IB framework in reinforcement learning : Pϕ∗ ( Z|X ) = arg max Pϕ ( Z|X ) EX∼P ( X ) , Z∼Pϕ ( Z|X ) [ J ( Z ; θ ) ] − βI ( X , Z ; ϕ ) ( 7 ) The following theorem shows that if the mutual information I ( X , Z ) of our framework and common RL framework are close , then our framework is near-optimality . Theorem 1 ( Near-optimality theorem ) . Policy πr = πθr , parameter ϕr , optimal policy π⋆ = πθ⋆ and its relevant representation parameter ϕ⋆ are defined as following : θr , ϕr = argmin θ , ϕ EPϕ ( X , Z ) [ log Pϕ ( Z|X ) Pϕ ( Z ) − 1 β J ( Z ; θ ) ] ( 8 ) θ⋆ , ϕ⋆ = argmin θ , ϕ EPϕ ( X , Z ) [ − 1 β J ( Z ; θ ) ] ( 9 ) Define Jπr as EPϕr ( X , Z ) [ J ( Z ; θr ) ] and Jπ ⋆ as EPϕ⋆ ( X , Z ) [ J ( Z ; θ⋆ ) ] . Assume that for any ϵ > 0 , |I ( X , Z ; ϕ⋆ ) − I ( X , Z ; ϕr ) | < ϵβ , we have |Jπ r − Jπ⋆ | < ϵ . 4.2 Target Distribution Derivation and Variational Lower Bound Construction In this section we first derive the target distribution in ( 7 ) and then seek to optimize it by constructing a variational lower bound . 1Mutual information I ( X , Y ) is defined as ∫ dXdZP ( X , Z ) log P ( X , Z ) P ( X ) P ( Z ) , conditional entropy H ( Y |Z ) is defined as − ∫ dY dZP ( Y , Z ) logP ( Y |Z ) . In a binary-classification problem , − logP ( Y |Z ) = − ( 1− Y ) log ( 1− Ŷ ( Z ) ) − Y log ( Ŷ ( Z ) ) . We would like to solve the optimization problem in ( 7 ) : max Pϕ ( Z|X ) EX∼P ( X ) , Z∼Pϕ ( Z|X ) [ J ( Z ; θ ) − β logPϕ ( Z|X ) ︸ ︷︷ ︸ L1 ( θ , ϕ ) +β logPϕ ( Z ) ︸ ︷︷ ︸ L2 ( θ , ϕ ) ] ( 10 ) Combining the derivative of L1 and L2 and setting their summation to 0 , we can get that Pϕ ( Z|X ) ∝ Pϕ ( Z ) exp ( 1 β J ( Z ; θ ) ) ( 11 ) We provide a rigorous derivation of ( 11 ) in the appendix ( A.2 ) . We note that though our derivation is over the representation space instead of the whole network parameter space , the optimization problem ( 10 ) and the resulting distribution ( 11 ) are quite similar to the one studied in ( Liu et al. , 2017 ) in the context of Bayesian inference . However , we stress that our formulation follows from the information bottleneck framework , and is mathematically different from that in ( Liu et al. , 2017 ) . In particular , the difference lies in the term L2 , which depends on the the distribution Pϕ ( Z | X ) we want to optimize ( while in ( Liu et al. , 2017 ) , the corresponding term is a fixed prior ) . The following theorem shows that the distribution in ( 11 ) is an optimal target distribution ( with respect to the IB objective L ) . The proof can be found in the appendix ( A.3 ) . Theorem 2 . ( Representation Improvement Theorem ) Consider the objective function L ( θ , ϕ ) = EX∼P ( X ) , Z∼Pϕ ( Z|X ) [ J ( Z ; θ ) ] − βI ( X , Z ; ϕ ) , given a fixed policy-value parameter θ , representation distribution Pϕ ( Z|X ) and state distribution P ( X ) . Define a new representation distribution : Pϕ̂ ( Z|X ) ∝ Pϕ ( Z ) exp ( 1 βJ ( Z ; θ ) ) . We have L ( θ , ϕ̂ ) ≥ L ( θ , ϕ ) . Though we have derived the optimal target distribution , it is still difficult to compute Pϕ ( Z ) . In order to resolve this problem , we construct a variational lower bound with a distribution U ( Z ) which is independent of ϕ . Notice that ∫ dZPϕ ( Z ) logPϕ ( Z ) ≥ ∫ dZPϕ ( Z ) logU ( Z ) . Now , we can derive a lower bound of L ( θ , ϕ ) in ( 6 ) as follows : L ( θ , ϕ ) = EX , Z [ J ( Z ; θ ) − β logPϕ ( Z|X ) ] + β ∫ dZPϕ ( Z ) logPϕ ( Z ) ≥ EX , Z [ J ( Z ; θ ) − β logPϕ ( Z|X ) ] + β ∫ dZPϕ ( Z ) logU ( Z ) = EX∼P ( X ) , Z∼Pϕ ( Z|X ) [ J ( Z ; θ ) − β logPϕ ( Z|X ) + β logU ( Z ) ] = L̂ ( θ , ϕ ) ( 12 ) Naturally the target distribution of maximizing the lower bound is : Pϕ ( Z|X ) ∝ U ( Z ) exp ( 1 β J ( Z ; θ ) ) ( 13 ) 4.3 Optimization by Stein Variational Gradient Descent Next we utilize the method in ( Liu & Wang , 2016 ) ( Liu et al. , 2017 ) ( Haarnoja et al. , 2017 ) to optimize the lower bound . Stein variational gradient descent ( SVGD ) is a non-parametric variational inference algorithm that leverages efficient deterministic dynamics to transport a set of particles { Zi } ni=1 to approximate given target distributions Q ( Z ) . We choose SVGD to optimize the lower bound because of its ability to handle unnormalized target distributions such as ( 13 ) . Briefly , SVGD iteratively updates the “ particles ” { Zi } ni=1 via a direction function Φ⋆ ( · ) in the unit ball of a reproducing kernel Hilbert space ( RKHS ) H : Zi ← Zi + ϵΦ⋆ ( Zi ) ( 14 ) where Φ∗ ( · ) is chosen as a direction to maximally decrease2 the KL divergence between the particles ’ distribution P ( Z ) and the target distribution Q ( Z ) = Q̂ ( Z ) C ( Q̂ is unnormalized 2In fact , Φ∗ is chosen to maximize the directional derivative of F ( P ) = −DKL ( P ||Q ) , which appears to be the ” gradient ” of F distribution , C is normalized coefficient ) in the sense that Φ⋆ ← argmax ϕ∈H { − d dϵ DKL ( P [ ϵϕ ] ||Q ) s.t . ∥Φ∥H ≤ 1 } ( 15 ) where P [ ϵΦ ] is the distribution of Z + ϵΦ ( Z ) and P is the distribution of Z . ( Liu & Wang , 2016 ) showed a closed form of this direction : Φ ( Zi ) = EZj∼P [ K ( Zj , Zi ) ∇Ẑ log Q̂ ( Ẑ ) |Ẑ=Zj +∇ẐK ( Ẑ , Zi ) |Ẑ=Zj ] ( 16 ) where K is a kernel function ( typically an RBF kernel function ) . Notice that C has been omitted . In our case , we seek to minimize DKL ( Pϕ ( ·|X ) || U ( · ) exp ( 1β J ( · ; θ ) ) C ) |C=∫ dZU ( Z ) exp ( 1β J ( Z ; θ ) ) , which is equivalent to maximize L̂ ( θ , ϕ ) , the greedy direction yields : Φ ( Zi ) = EZj∼Pϕ ( ·|X ) [ K ( Zj , Zi ) ∇Ẑ ( 1 β J ( Ẑ ; θ ) + logU ( Ẑ ) ) |Ẑ=Zj +∇ẐK ( Ẑ , Zi ) |Ẑ=Zj ] ( 17 ) In practice we replace logU ( Ẑ ) with ζ logU ( Ẑ ) where ζ is a coefficient that controls the magnitude of ∇Ẑ logU ( Ẑ ) . Notice that Φ ( Zi ) is the greedy direction that Zi moves towards L̂ ( θ , ϕ ) ’ s target distribution as showed in ( 13 ) ( distribution that maximizes L̂ ( θ , ϕ ) ) . This means Φ ( Zi ) is the gradient of L̂ ( Zi , θ , ϕ ) : ∂L̂ ( Zi , θ , ϕ ) ∂Zi ∝ Φ ( Zi ) . Since our ultimate purpose is to update ϕ , by the chain rule , ∂L̂ ( Zi , θ , ϕ ) ∂ϕ ∝ Φ ( Zi ) ∂Zi ∂ϕ . Then for L̂ ( θ , ϕ ) = EPϕ ( X , Z ) [ L̂ ( Z , θ , ϕ ) ] : ∂L̂ ( θ , ϕ ) ∂ϕ ∝ EX∼P ( X ) , Zi∼Pϕ ( ·|X ) [ Φ ( Zi ) ∂Zi ∂ϕ ] ( 18 ) Φ ( Zi ) is given in equation ( 17 ) . In practice we update the policy-value parameter θ by common policy gradient algorithm since : ∂L̂ ( θ , ϕ ) ∂θ = EPϕ ( X , Z ) [ ∂J ( Z ; θ ) ∂θ ] ( 19 ) and update representation parameter ϕ by ( 18 ) . 4.4 Verify the information E-C process with MINE This section we verify that the information E-C process exists in deep RL with MINE and our framework accelerates this process . Mutual information neural estimation ( MINE ) is an algorithm that can compute mutual information ( MI ) between two high dimensional random variables more accurately and efficiently . Specifically , for random variables X and Z , assume T to be a function of X and Z , the calculation of I ( X , Z ) can be transformed to the following optimization problem : I ( X , Z ) = max T EP ( X , Z ) [ T ] − log ( EP ( X ) ⊗P ( Z ) [ expT ] ) ( 20 ) The optimal function T ⋆ ( X , Z ) can be approximated by updating a neural network T ( X , Z ; η ) . With the aid of this powerful tool , we would like to visualize the mutual information between input state X and its relative representation Z : Every a few update steps , we sample a batch of inputs and their relevant representations { Xi , Zi } ni=1 and compute their MI with MINE , every time we train MINE ( update η ) we just shuffle { Zi } ni=1 and roughly assume the shuffled representations { Zshuffledi } ni=1 to be independent with { Xi } ni=1 : I ( X , Z ) ≈ max η 1 n n∑ i=1 [ T ( Xi , Zi ; η ) ] − log ( 1 n n∑ i=1 [ expT ( Xi , Z shuffled i ; η ) ] ) ( 21 ) Figure ( 1 ) is the tensorboard graph of mutual information estimation between X and Z in Atari game Pong , x-axis is update steps and y-axis is MI estimation . More details and results can be found in appendix ( A.6 ) and ( A.7 ) . As we can see , in both A2C with our framework and common A2C , the MI first increases to encode more information from inputs ( ” remember ” the inputs ) , then decreases to drop irrelevant information from inputs ( ” forget ” the useless information ) . And clearly , our framework extracts faster and compresses faster than common A2C as showed in figure ( 1 ) ( b ) . After completing the visualization of MI with MINE , we analyze the relationship between our framework and MINE . According to ( Belghazi et al. , 2018 ) , the optimal function T ∗ in ( 20 ) goes as follows : expT ∗ ( X , Z ; η ) = C Pϕ ( X , Z ) P ( X ) Pϕ ( Z ) s.t . C = EP ( X ) ⊗Pϕ ( Z ) [ exp T∗ ] ( 22 ) Combining the result with Theorem ( 2 ) , we get : expT ∗ ( X , Z ; η ) = C Pϕ ( Z|X ) Pϕ ( Z ) ∝ exp ( 1 β J ( Z ; θ ) ) ( 23 ) Through this relationship , we theoretically derive an algorithm that can directly optimize our framework without constructing the lower bound , we put this derivation in the appendix ( A.5 ) .

Summary: In this paper, the authors proposed to utilize variational lower bound of mutual information for learning representations in Reinforcement learning. To optimize the proposed variational lower bound with a more flexible encoding network, the author proposed to utilize stein variational gradient descent (or amortized svgd). Instead of learning the representation separately, the authors incorporate the framework into PPO or A2C, which yields a joint training framework for policy optimization.