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: This paper proposes a representation learning algorithm for RL based on the Information Bottleneck (IB) principle. This formulation leads to the observed state X being mapped to a latent variable Z ~ P(Z | X), in such a way that the standard loss function in actor-critic RL methods is augmented with a term minimizing the mutual information between X and Z (which can be seen as a form of regularization). This results in a loss that is difficult to optimize directly in the general case: the authors thus propose to approximate it through a variational bound, using Stein variational gradient descent (SVGD) for optimization, which is based on sampling multiple Z_i’s for a given state X, so as to compute an approximate gradient for the parameters of the function mapping X to Z. Experiments show that when augmenting the A2C algorithm with this technique, (1) the mutual information I(X, Z) decreases more quickly (better « compression » of the information), and (2) better sample efficiency is observed on 5 Atari games (with also encouraging results with PPO on 3 Atari games).