Learning Representations in Reinforcement Learning: an Information Bottleneck Approach
Authors: Yingjun Pei
Xinwen Hou
Venue:  (2020)
Keywords: representation learning
reinforcement learning
information bottleneck

Abstract: The information bottleneck principle is an elegant and useful approach to representation learning. In this paper, we investigate the problem of representation learning in the context of reinforcement learning using the information bottleneck framework, aiming at improving the sample efficiency of the learning algorithms.We analytically derive the optimal conditional distribution of the representation, and provide a variational lower bound. Then, we maximize this lower bound with the Stein variational (SV) gradient method. 
We incorporate this framework in the advantageous actor critic algorithm (A2C) and the proximal policy optimization algorithm (PPO). Our experimental results show that our framework can improve the sample efficiency of vanilla A2C and PPO significantly. Finally, we study the information-bottleneck (IB) perspective in deep RL with the algorithm called mutual information neural estimation(MINE).
We experimentally verify that the information extraction-compression process also exists in deep RL and our framework is capable of accelerating this process. We also analyze the relationship between MINE and our method, through this relationship, we theoretically derive an algorithm to optimize our IB framework without constructing the lower bound.

Under review as a conference paper at ICLR 2020
Learning Representations in Reinforcement Learning: An Information Bottleneck Approach
Anonymous authors Paper under double-blind review
Abstract
The information bottleneck principle in (Tishby ale 2000) is an elegant and useful approach to representation learning. In this paper, we investi- gate the problem of representation learning in the context of reinforcement learning using the information bottleneck framework, aiming at improving the sample efficiency of the learning algorithms. We analytically derive the optimal conditional distribution of the representation, and provide vari- ational lower bound. Then, we maximize this lower bound with the Stein variational (SV) gradient method (originally developed in (Liu & Wang, 2016; Liu et al 2017)). We incorporate this framework in the advanta- geous actor critic algorithm (A2c)(Mnin)) et al 2016) and the proximal policy optimization algorithm (PPO) (Schulman et al, 2017). Our experi- mental results show that our framework can improve the sample efficiency of vanilla A2C and PPO significantly. Finally, we study the information bottleneck (TB) perspective in deep RI, with the algorithm called mutual information neural estimation (MIN)) (Belghazi et al,, 2018). We exper- imentally verify that the information detraction-oopprssiiinnn process also exists in deep RL and our framework is capable of accelcrating this process. We also analyze the relationship between MINE and our method, through this relationship, we theoretically derive an algorithm to optimize our IB framework without constructing the lower bound.
Introduction
In training a reinforcement lcarning algorithm, an agent interacts with the environment, ex- plores the (possibly unknown) state space, and learns 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 elficiency the learning al- gorithm is key problem in RL and has been studied extensively in the literature. Popular techniques include experience reuse rense/replay, which leads to no--oorerffll offapolicy algorithgs (eg (Miih et, al 2013; Silver et all, 2014; Van Hasselt et 2015 Nachum et al 201 Sag Espcholt et al. 2018)), and model-based algorithms (e.g., (Hafner et al, 2018: Kaiser et al 2019)). Moreover, it known that effective representations can greatly reduce the sample complexity in RL. This can be seen from the following motivating example: In the envi- ronment of classical Atari game: Seaquest, converge to an optimal policy when the input states are raw images (more than 28,000 di- mensions), while while it. takes less samples when the inputs are 128-dimension pre-defined RAA.M data(Sygnowski & Michalewski 2016). Clearly, the RAM data contain much less redundant information irrelevant. the learning process than the raw images. Thus, we argue that an efficient representation is extremely crucial to the sample efficiency.
may take dozens of millions samples to
In this paper, we try to improve the sample efficiency in RL from the perspective of representation learning using the celebrated information bottleneek framework (Tishyy (Shwartz-Ziv & Tishby,
al 2000). In standard deep learning, the experiments 2017] show that during the training process, the neural network first "remembers" the inputs by increasing the mutual information between the inputs and the repre- sentation variables, then compresses the inputs efficient representation related to

---

Under review as a conference paper at ICLR 2020
the learning task by discarding redundant information from inputs (decreasing the mu- tual information between inputs and representation variables). We call this phenomena "information extraction-compresioo process" (informatinn 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-coopprssiin phenomena in the context of deep RL (we need to use MiNgaBelghazi et al,, 2018) for estimating the mu- tual information). This observation motivates us to adopt the information bottleneck (IB) framework in reinforcement learning, in order to accolerate the extraction-compresst cess. The IB framework is intended explicitly enforce RL agents learn an efficient representation, hence improving the sample efficieney, by discarding irrelevant information from raw input data. Our technical contributions can be summarized as follows:
pro-
1. We observe that the "information extraction-comppesiia process" also exists in the context of deep RL (using (Belhhaiii 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- compressoooo process. Our experimental results also show that combining actor-crilic 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 rela- tionship, 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 futuure work to incorporate it in other off-policy and model-based algorithms.
Related Work
Information bottleneck framework was first introduced in (Tishby al,, 2000). They solve the framework by iterative Arimooo algorithm, which is infeasible to apply to deep Arimoto-algorithm, neural networks. (Shwartz-Zii Ziv Tishby, 2017) tries to open the black box 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(ZIX) itself must. be Gaussian distribbtiinn which limits the capabiliiies of the representation function. (Peng et al 2018) extends this framework to variational discriminator bottleneck to improve CAANs(Goodeell et al 2014), imitation learning and inverse RL. As for improving sample-effieinncy, (Mnih et al. 2013: Van Hasselt et al. et all 2018a) mainly, utilize the experience-reuse. Besides experience-reusee
2015 Nachum
(Silver
2014; Fujimoto et al., 2018) tries to learn a dleterministic_polle "olicy,_[Eppholts seeks mitigate the delay of off-policy. (Hafner al., 2018: Kaiser et all 2019 learn
(Espeholt, ot at 2018
the environmmnn 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, s018hssshows.o ) shows a theoretical foundation of maintaining the optimality of representation space. [Bellemare et al., 2019) proposes a new perspective on represenntation learning in RL based on geometric properties of the space of value function. (Abel et al. 2019) learns representation via information bottleneck in imitation/ mmitation/appreaticslll aaprenticeship learning. To the best of our knowledge, there is work that intends to directly use IB in basic RT algorithms.
2

---

Under review as a conference paper at ICLR 2020
3 Preliminaries
A Markov decision procoss(MDP) is tuple, (P,,A pu), where X' is the set of states,
A is the set of actions, R: : X' x A X' + I is the reward function, P: x A x X -10, 1] is the transition probability function(whas
P(X [X, a) is the probability of transitioning to state X" given that the previous state is X and the agent took action in X), and u: X -100 1] is the starting state distribution. A policy т : - 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, Br Rx(,X, at, Xx+1], with a slight abuse of notation we denote R(X,, ar Xx+1) = rs. Here 7 E (0,1) state value function as V*(X)= = Er-- r
aim to select policy which maximizes K(T)
discount factor, denotes a trajectory (Xo, 0o, X0,X1, (I,,),...). Define the
= X]. which is the expected return by policy π in state X. And the state-action value function Q (x, a) = E - X0r
0, =a] a] is the expected return by policy after taking action in state X.
Actor-critic algorithms take the advantage both policy gradient methods and value- function-based methods such the well-known A2C(MMih et al, 2016). Specifically, in the case that policy Taa(X;0) 0) is parameterized by 0, A2C uses the following equation to approximate the real policy gradient VxK(x)= [V]((O)
20
VoJ(0) R Vollog((a/Xx;
+ (z = Σ vol(X:0) (1)
where R2 = Ei= Te is the accumulated return from time step 4, Hy)) is the entropy of distribution p and b(X,) is baseline function, which is commonly replaced by V (X.).
A2C also includes the minimization of the mean square error between R2 and value function V= (x.)). Thus in practice, the total objective function in A2C can be written as:
((0 ~0
log (a, (x) - a1 [R
a2H(((((X))) - Σ ((
(2)
where a1, a2 are two coefficients.
In the context of representation learning in RL, J(X,; (including (*.),) and Q=xX, at))
can be replaced by J(Z;;) 0) where Ze is a learnable low-dimensional representation of state
Xt. For example, given representation function Z ~ P&(-IX) with parameter o, define
V= (Z; X1, Φ) (= VX,,). For simplicity, we write V (T(Z; Xe, Ø) lz x.) as (Z,)..
4 Framework
4.1 Information Bottleneck in Reinforcement Learning
The information bottleneck framework is an information theoretical framework for extract- ing relevant information, or yiclding a representation, that an input X E X contains about an output Y E 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 (Z(X) such that:
(=(X arg max T(Y, 3T(X, arg max H(Y) (Y) 3T((X Z) P(Z(X) arg max P(Z(X)
P(Z(X)
() 31(X, 2)
(3)
3

---

Under review as a conference paper at ICLR 2020
for every X E X, which appears the standard cross-entropy loss in supervised learning with a MI-regularizer, 3 a coefficient that controls the magnitude of the regularizer.
Next we derive information bottleneck framework in reinforcement learning. Just like the label the context of supervised learning as showed in (3), we assume the supervising signal Y in RL to be the accurate value R, of specific state X, for fixed policy x, which can be approximated by n-step bootstrapping function Y, R. Ei= 7 y pe--VVT(Z++a in practice. Let P(Y be the following distribution:
Z() x exp(-a((R, ( (+())))
(4)
This assumption is heuristic but reasonable: we have an imput X, and its relative label Y, R., now have X,'s representation Ze, naturally function (Z,) to approximate the true label Y.. we set our target distribution to be C exp(-oo R2 (Z)))) the probability decreases (* gets far from Y, while increases as V* (Z:) gets close to Y,
want to train decision
For simplicity, just write P(R[Z) instead of P(Y) Z+) in the following context.
With this assumption, equation can be written as: P(R -
(*) = arg max E P(Z[X)
arg max Ex-PNy,--PaW,,B-PPRP -2)) Z)
(5)
P(Z
The first term looks familiar with classic mean squared error in supervisd learning. In network with representation parameter and policy-value parameter 0, policy loss J(Z:0) in equation (1) and IB loss in ( can be jointly written as:
L(0, -x=xxxxxxxxxx--_
((Z:0) EAl ERI-a(BF aa(R V -I
2i (6)
J(Z:0)
where [(X, Φ) denotes the MI between X and Z ~ P)) Notice that (2; itself is
standard loss function in RL as showed in (2). Finally we get the ultimate formalization
of IB framework in reinforcement learning:
Por (2(X) arg max Ex-P(X),--R(( ([2 - P,(2IX)
(7)
The following theorem shows that if the mutual information I(X, 2) of our framework and common RL framework are close, then our framework is near-optimality.
Theorem 1 (Near-optimality theorem). Policy x" Tgr, parameter ø", optimal policy x TT. and its relevant representation parameter 6* are defined as following:
8", 8¹ arg min E (x.z) log Po(2() 1 P2
J(Z:0)
(8)
0.6
3
0 arg min E Pa (x 1 [(:
(9)
0,6
3
Define J=" as Ep (x,-((X,)) and F= as Ep. (x,2 Assume that for any 0, |(XX ((X, f we have J J <E.
4.2 Target Distribution Derivation and Variational Lower Bound Construction
In this section we first derive the target distribution in r and then seek to optimize it by constructing variational lower bound.
MMttal information I(X,Y) is defined f ddxdZP(X, 2) log P(XP(2) conditional tropy H(YlZ) is defined as - f dY dZ P(Y, log P(Y|Z). In binarycdlasification problem, log P(Y|Z) -(1 log(1 f(2)) (l))(
PXX2)
4

---

Under review as a conference paper at ICLR 2020
We would like to solve the optimization problem in (0)
max EX x) (Z; - Blog (x(Z + log (Z(
(10)
Ps (2)x)
(1,0))
Combining the derivative of L1 and L2 and setting their summation to 0, we can get that Pa(Z)X) x Palz) exp( 1 g ((;)0)
(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,
stress that
our formulation follows from the information bottleneck framework, and is mathematically different from that in (Liu et ale 2017)- In particular, the difference lies in the Lerm La, which depends on the the distribution Po XX) 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 (LI) is an optimal target distribuioo (with respoct the IB objective L). The proof can be found in the appendix (A.3).
(A.3).
Theorem 2. (Representation Improvement Theorem) Consider the objective function - (X, Z; σ), given a fixed policy-value parame- ter 0, representation distribution (x( and state distribution P(X). Define a new representation distribution: (z(Z) x (Z( exp( (x;())(()0), We have L 6) 2 o),
L(0, Φ) = =-x
Though we have derived the optimal target distribution, it is still difficult to compute (Z) In order to resolve this problem, we construct a variational lower bound with a distribution U(Z) which independent of é. Notice that d2P((Z) log Pz(Z) 2 dZ dZPp (Z) log U (Z). Now, we can derive lower bound of L(0, Φ) in as follows:
L) Ex (J(Z,) Blog Pa(Z 3 dZP Po(Z)
2 Ex.z(J(Z;0) - 3 log Poog / dZPP(Z)
-Bx-P00,2-P,(2x,, -Blg Ps (Z + U()] (()
(12)
Naturally the target distribution of maximizing the lower bound is:
Pa(ZIx) ((Z))xp) =
(13)
4.3 Optimization by Stein Variational Gradient Descent
Next we utilize the method (Liu & Wang. 2016) et al. 2017 (Haarnoja el. al. 2017) to optimize the lower bound.
Stein variational gradient (SSC)))(SVGD) is a non-parametric variational inference algo- rithm that leverages efficient deterministic dynamics to transport set of particles (
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" "
{ via a direction function (*()) in the unit ball of reproducing kernel Hilbert space (RKHS) H:
Z < Z, + (Z,)
(14)
where (-) is chosen as a direction to maximally decrease the KL divergence between the particles' distribution P(Z) and the target distribution Q(Z) = Q(Z) (Q is unnormalized
c
PIn fact, tx is chosen to maximize the directional derivative of F(P) -Bxx which appears to be the "gradient" of F
5

---

Under review as a conference paper at ICLR 2020
distribution, C' is normalized coefficient) in the sense that
f- arg max dEH
(()) s.t. 1 1
(15)
dc
where P(aa) the distribution of Z + (Z) and is the distribution of Z. (Lu & Wang 2016) showed a closed form of this direction:
L 2.) = (e,,,ez,,,,,,,,... V.)V 2 log Q(2) 22-2, + +V2 12=2,1
(16)
where K. is kernel function(typically omitted.
RBF kernel function). Notice that C' has been
In our case, we seek to minimize (xL(Po(-IX)I U(+) expl exp() which is equivalent to maximize [(0, Φ), the greedy direction yields:
lc-s (xU(
Φ (2.) = EzjoPat X) [Z V J J(; + log logU(2)) 2-2, +V2 12=2,1 (17)
3
In practice we replace log (Z) with Çlog U(2) where is a cocfficient that controls the magnitude of V log (2). Notice that (z) is the greedy direction that Zj moves towards [0, 6)'s target distribution as showed in (13) (distribution that maximizes ((0,0)) This means 4 (2,)
the gradient of L(Z,, 0, Ø): 87.
01(2:000 (Zz)).
Since our ultimate purpose is to update o, by the chain rule, (L. for L(,, Φ) = pp, (x.2) (x,) (Z,
x ((2))
OZ. 76
Then
d6
o0, o)
oz,
x Px
(18)
06
06
((,)) is given in equation 17 In practice update the policy-value parameter O by common policy gradient algorithm since:
OL(0,)) Φ)
0J(Z;0)
(x,(x
(19)
00
00
and update representation parameter o by 18
4.4 Verify the information E-C process with MINE
This section we verify that the information E-C process exists 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 effi- ciently. Specifically, for random variables and 7, assume T to be a function X and Z, the calculation of T(X, Z) can be transformed to the following optimization problem:
I(X, Z) max Eax.z( [T] log(Ep(xysrrz/(expp
(20)
T
The optimal function T*XX can be approximated by updating a neural network T(X, Z;n).
With the aid of this powerful tool, we would like to visualize the mutual information between imput state X and its relative representation Z: Every few update steps, we sample batch of inputs and their relevant representations {X,, 2r and compute their MI with MINE, every time train MINE(update 17) we just shuffle [Z,)2 and roughly assume the shuffled representations (2.hufiled Jel
to be independent with {X,),
I(X, 2) max
Σ T(X,, Zan) - -log( Σ TxxpT(X, T(X, 2. mi))
(21)
72

---

Under review as a conference paper at ICLR 2020
(igure(i) is the tensorboard graph of mutual information estimation between X and Z in Atari game Pong, x-axis is update, steps and y-axis MI estimation. More details and results can be found in appendix(AA6) and (A.7). As we see, in both A2C with our framework and common A2C, the MI first increases to encode more information from inputs("renns "pputs"rrmmmber" the inputs), then decreases to drop irrelevant information from in- puts( fforge" the useless information). And clearly, our, framework extracts faster and compresses faster than common A2C as showed in figure([))()) (b))
mutuul_info_oss
LnftuaLinfo_Joss
2.20
M m
1.00
1.50
1.40
1.00
1000
4.000k
8.000k
12.00k
16.00k
20.00k
0.000
4.000k
8.000k
12.00k
16.00k
20.00k
(a) MT A2C
(b) MT A2C with our framework
Figure 1: Mutual information visualization in Pong
After completing the visualization of MI with MINE, we analyze the relationship between framework and MINE. According to (Belghazi et al., 2018). the optimal function 7 in goes as follows:
20
exp T' (X,7.) =C' (X,
s.t. C= Bp(x)oon(())exp" P(X)P&(Z)
(22)
Combining the result with Theorem we get:
expx" (X, = 2 Po(Z)X) Po(Z)
x exp( 1 exp())((200)
(23)
Through this relationship, we theoretically derive an algorithm that can directly optimize framework without constructing the lower bound, we put this derivation in the ap- pendix A.55)
Experiments
In the experiments we show that our framework can improve the sample efficiency of basic RLL algorit.imsttyppcallyy A2C PPO). Our anonymous code can be found in https:/ attps://githut/ /gthhub. Other results can be found in last two appendices.
com com//moooo AtooryooouSummitttalll/////////
In A2C with our framework, we sample Z by network o(X, E) where ~ (.;00,0.1 and the take more samples
number of samples from each state X 32, readers are encouraged if the computation resources are sufficient. We set. the IB coefficient as B = 0.001. We choose two prior distributions U(Z) of our framework, the first one is uniform distribution, apparently when UZZ) is the uniform distribution, V log ((2) 1ă=z can be omitted. The second one is Gaussian distribution, which is defined as follows: for given state X,, sample a batch of !=? then: U(Z)=1a -N(;;; ヒ Zj,0² = E -z))
We also set < as 0.005| (22 log U(2) 1ž-z to control the magnitude of V log U2) 1z-z. Following (Liu et al., 2017). the kernel function in (17) we used is the Gaussian RBF kernel KC2,, -,) (xpp-1122 med denotes the median of pairwisc distances between the particles {2j)== As for the
12212/h where h med /eed-/2logte lognn 1),
parer--praamterss... in RL, we we_ hem_simply_oee choposycchoose.. choose the default. parameters in A2C of Openai- baseline((https://eg htti:::://gthub...t cecchuu..omm/ppee www/ooooao////eeeeeeooe/eoeeee/e baselines/tenneeeeeeee (trre/mmster, hraaaiee///rrmmmmmttrrraainnnn baselines/a20, implement the following four algorit hms::
summary,

---

Under review as a conference paper at ICLR 2020
A2C with uniform-SV SVIB: Use c) the embedding function, optimize by frame-
work(algorithm (A.4)) with (2) being uniform distribution.
A2C with Gaussian SVIB: Use o(X,e) the embedding function, optimize by our frame-
work(algorithm A.4)) with (Z) being Gaussian distribution.
Regulrr A2C in Openai-baselines with ((X) the embedding function.
A2C with noise(For fairness):A2C with the same embedding function <(X,c) as A2C with our framework.
Figure (a)) show the performance of four A2C-based algorithms in 5 gym Atari games. We can see that A2C with our framework is more sample-efficient than both A2C and A2C with noise in nearly all games.
(a) A2C-Pong
(b) A2C-AirRaid
(c) A2C- Bee-BeemRider
(d) 32C-Qbert
86
(e) 12CC-ppaelnvaders
(f) PPO-Pong
(g) PPO-Qbert
(h) PPO- Breakout
Figure 2: (a)-(e) show the performance of four A2C-based algorithms, x-axis is time steps(2000 update steps for each time step) and y-axis the average reward over 10 episodes, (f)-(h) show the performance of four PPO-based algorithms, x-axis is time steps(300 update steps for each time step). We make exponential moving average of each game to smooth the curve(In PPO-Pong, we add 21 to all four curves in order to make exponential moving average). We can see that our framework improves sample efficiency of basic A2C and PPO.
Notice that in Spacelnvaders, A2C with Gaussian SVIB worse. We suspect that this is because the agent excessively drops information from inputs that it misses some information related to the learning process There is a more detailed experimental discussion about this phenomena in appendix A.7)) We also implement four PPO-based algorithms whose experimental settings are same as A2C except that we set the number of samples as 26 for the sake of compulation efficiency. Results can be found in the in figure (2)(()))
6 Conclusion
We study the information bottleneck principle in RI: We propose optimization problem for learning the representation in RL based on the information-bottlenee derive the optimal form of the target distribution. We construct lower bound and utilize Stein Variational gradient method to optimize it. Finally, we verify that the information extraction and compression process also exists in deep RL, and our framework can accelerate this process. We also theoretically derive an algorithm based on MINE that can directly optimize our framework and plan to study it experimentally in the future work.
framework and
8

---

Under review as a conference paper at ICLR 2020
References
David Abel, Dilip Arumugam, Kavosh Asadi, Yuu Jinnai, Michael L Littman, and Law- son LS Wong. State abstraction compression in apprenticeship learning. In Proceedings of the AAAI Conference on Artificial Intelligence. AAAT Press, pp. 3134-3142, 2019.
Alexander A. Alemi, Ian Fischer, Joshua Dillon, and Kevin Murphy, Deep variational information bottleneck, In Proceedings of the International Conference on Learning Rep- resentations, 2016.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: Mutual information neural cstimation. arXiv preprint arXiv:1801.04000 2018.
Mare G. Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. A geometric perspec- tive on optimal representations for reinforcement learning. International Conference on Learning Representations, 2019.
Mathew Botvinick, Sam Ritter, Jane X Wang, 7eb Kurth-Nelson, Charles Blundell, and Demis Hassabis. Reinforcement learning, fast and slow. Trends in cognitive sciences, 23 (5):408 422, 2019.
Tessler Chen, Daniel J. Mankowitz, and Shie Mannor. Reward constrained policy optimiza- tion. arXiv preprint AXXv:11005.11744, 2018
Lasse Espcholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, and lain Dunning. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint arXiv:1802.01561, 2018
Scott. Fujimoto, Herke Van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv:1802.09477 2018.
lan J Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Xu Bing, David Warde-Parley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In International Conference on Neural Information Processing Systems, pp. 2672-2680) 2014.
Tuomas Haarnoja, Haoran Tang, Picter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. Proceedings of the 34th International Conference on Machine Learning, 70:1352 1361, 2017.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, and James Davidson. Learn- ing latent dynamics for planning from pixels. International Conference Machine Learn- ing, pp. 2565, 2018.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, and Sergey and Levine. based reinforcement learning for atari. arXiv preprint MtXiv:1903.00374 2019.
Timothée Lesort, Natalia Díaz-Rodríguez, Jean François Goudou, and David Filliat. State representation learning for control: An overview. Neural Networks, 108: 508936000118300033 2018.
Qiang Liu and Dilin Wang. Stein variational gradient descent: A general purpose bayesian inference algorithm. Advances in Neural Information Processing Systems 29, pp. 2378- 2386, 2016.
Yang Liu, Prajit Ramachandran, Qiang Liu, and Jian Peng. Stein variational policy gradi- ent. arXiv preprint arXiv:1704.02399, 2017.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. Com- puter Science, 2013.
9

---

Under review as a conference paper at ICLR 2020
Volodymyr Mnih, Adrià Puigdomèncch Badia, Mchdi Mirza, Alex Graves, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International Conference on Machine Learning, pp. 1928-19331 2016.
Ofir Nachum, Shane Gu, Honglak Lee, and Sergey Levine. Data-efficient hicrarchical rein- forcement learning. Neural Information Processing Systems, pp. 3307 3317, 2018a.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hierarchical reinforcement learning. International Conference on Learning Representations, 2018b.
Xue Bin Peng, Angjoo Kanazawa, Sam Toyer, Pieter Abbeel, and Sergey Levine. Vari- ational discriminator bottleneck: Improving imitation learning, inverse rl, and gans by constraining information flow. International Conference on Learning Representations, 2018.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06447 2017.
Ravid Shwartz-Ziv and Naftali Tishby. Opening the black box of deep neural networks via information. arXiv preprint arXiv:1703.00811 2017.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Ried- miller. Deterministic policy gradient algorithms. In International Conference Interna- tional Conference on Machine Learning, pp. 387-395, 2014.
Jakub Sygnowski and Henryk Michalewski. Learning from the memory of atari 2600. arXiv preprint 2Xiv:1605.01335 2016.
Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method. University of Mllinoi, 411(29-30)3668 377, 2000.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. Computer Science, pp. 2094 2100, 2015.
A Appendix
A.1 Proof of Theorem
Theorem. (Theorem 1 restated)Policy x" = Ter, parameter Φ', optimal policy n" Taa and its relevant representation parameter o* are defined as following:
Pa(Z
1
0", 0" arg min E. (x, log Pz(Z)
J(Z:0)
(24)
8
3
0" arg min E 2) B (2;0)
(25)
0.0
Define Jx as Ep (x.- and JH as Pp- (JZ Assume that for any 6 >> 0, [T(X, 7; Φ*) - T(X, 72 Φ-) 3 we have |Jm J" < e. Specifically, in value-based
algorithm, this theorem also holds between expectation of two value functions.
Proof. From equation 24 we can get:
I(X,Z:0*) B J 2 21(X, Z;0') 1 F
(26)
3
From equation 25 can get:
Jx 2
3 n"
(27)
These two equations give the following inequality:
((T(X,2:0") T(X, 7:00) 2 J
_xx 20
(28)
10

---

Under review as a conference paper at ICLR 2020
According to the assumption, naturally we have:
J -ー'
(29)
Notice that if we use our IB framework value-based algorithm, then the objective function J can be defined as:
= V = (1-7)-1 / (X)
= (1--)~1 / (X) (ZPA(Z(X)FF(2))
(30)
where (Z) [xe(X'=((')==== ==())=2)) dx (X) and d= is the discounted future state distribu- tion, readers can find detailed definition of dt in the appendix of (Chen al, 2018). We can get:
(31)
A.2 Target Distribution Derivation
We show the rigorous derivation of the target distribution in 111
Denote P as the distribution of X, (Z) Po(Z) as the distribution of Z. We use Po as the short hand notation for the conditional distribution (Z) Morcover, we write L(O, Φ) L(6, P&) and (p.9)x f (X))). Notice that PZ( (P()).
Take the functional derivative with respect Po of the first term L1:
(0, Po)
= f dzdx (0, ((Z = de
(0, Po + c
8P
SPo(ZIX)
xZ
c=0
de / [XxP(x) (EXP(X)(P((11)--- J(; 8 log( + cD,, ((,,)))))))))))
2
/xxpxx) (,(.))),(::0-- (0::00)-3og--00
Blog Po(()))) (Po(-IX), -8 86 ((-
Po(-IX)
-PP()J((:0) - Blog -loggo(.1)) - ((:)) xz
Hence, we can see that
8L, Po) P(X)]J(Z:0) log (Z( - -3]. (2
Then we consider the second term. By the chain rule of functional derivative, we have that
(L) (2
(l, Pa) (.)
3 (1 1+ log P SP2(.)
(.) (Z)() 2
8P,(2
3 ʃ di(1 dā(1+ ([l+llog log PS (2)))(2- 2))-()))(-----))) SP(X) log P3 (Z))
(32)
Combining the derivative of L1 and La and selling their sunmation to 0, we can get that
(z( x Po(Z)exp( 3 J(Z:8))
(33)
A.3 Proof of Theorem
Theorem. (Theorem restated) L(0, Ex-p11).2--p22m1(((210000 31 22 o),
given fixed policy-value parameter 0, representation distribution Pa(ZIX) and state dis- tribution P(X), define new representation distribution:P Pa(Z) exP() 0)),
have L(O, .) 2 ()
11

---

Under review as a conference paper at ICLR 2020
Proof. Define I(X)
((X) dZPx(ZIX) | dZP exp( 1
(34)
7
L( - Ex (xx - - - (a( exp( (2(0)))
log
}}
((x)P
Po(Z)
[x fog [(x)] {Ez~P:(ZIx)100 log
PZ(
Pa(Z)
3ex flog [(X)] 2 1og
Pz(
Blog [X)) 88 1og Pa(Z)
Pz(Z)
(exx kL (P, (Z)
(35)
1 L( Ex ( exp( J (Z: (2E2_))() log (z(
P2(2
exp())xp(
Ex 2x Pa(ZIX)((X) 1 (elog((X)
1 Pz(Z(X),
P2(Z(X)
(36)
L(0, - (() (Px (2)
(37)
According to the positivity of the KL-divergence, have L(0, 3) > L),
A.4 Algorithm
Algorithm
information-bottleneddee state abstraction in RL
0.6 + initialize network paramoters
B.C initialize hyper-parameters in
6 <- learning rate
M number of samples from (-
repeat
Draw batch of data {X,, az, Re, X Pi-1 from environment for each X, E (x,)== do
Draw M samples { from Po(()X,)
end for
Get the batch of data D = (X,, ) MM 1=11 de, Re, X
Compute the representation gradients V (0, Φ) in D according to (18
Compute the RL gradients VoL(A, σ) in according to (19
Update d: - o + ev ((8,0)
Update 0 + + (. until Convergence
A.5 Integrate MINE to our framework
MINE can also be applied to the problem of minimizing the MI between 7 and X where is generated by neural network (o(+)X):
1 (X, Z) = min max E
" Px (P. [(x,2;)) Bogg)llooo(ppxxxpp2peee0ol T( ((x.2m))
(38)
12

---

Under review as a conference paper at ICLR 2020
Apparently I* [X, Z) is 0 without any constraints, yet if we MINE to optimize our IB framework in (.) (X, Z) might not be O. With the help of MINE, like what people did in supervised learning, the objective function in (6) can be written as:
T(0, 9, n)
min ( (x, (2 (x. (XX, (((,
RL
+ B (o "(
(39)
The key sleps to optimize L(0, 0 n) is to update 0, o, iteratively as follows: (X, m:))
u++1 arg min E (x.2) 2x,
T(X T()))))))
(40)
m
mutual information term
Oc+1, 01-1 t- arg max L(Oe, de, M/+1))
(41)
6.,0e
Yet in our experiment, these updates of f, O, 7) did not work at all. We suspect might be caused by the unstability of T(X, Z; n)) We have also found that the algorithm always tends to push parameter @ to optimize the mutual information term to regardless of the essent RL loss J(Z; 0): In the early stage of training process, policy и is so bad that the agent is unable to get high reward, which means that the RI. loss extremely hard to optimize. Yet. the mutual information term is relatively easier to optimize. Thus the consequence is that these updates tend to push parameter & to optimize the mutual information term to in the carly training stage.
However, MINE's connection with our framework makes it possible to optimize T(X, 7) in more deterministic way and put the RL loss into mutual information term.
According to equation 23 T(X, n) (Z( this implies that we could introduce another function T in place of T(X, Z;m) for the sake of variance reduction:
1
T(X, + 3 J(z:0) T(X,
(42)
in the sense that parameter 1) only needs to approximate the constant C' (T', the optimiza- tion steps turn out to be:
0.1 Pct1, M1+1
arg max arg min (-minnn-E..p. (x [XX, 2:m) He.ds "
log (log(E (2) f(x 2m)) 1(22011)))
(43)
This algorithm has the following three potential advantages in summary:
1. We're able optimize T(X, Z;n) in more deterministic way instead of the form like equation(10) (C), which is hard to converge in reinforcement learning.
2. We prevent excessive optimization in mutual information term by putting RL loss J (Z; 0) into this term
3. We're able to directly optimize the IB framework without constructing a variational lower bound.
Here we just analyze and derive this algorithm theoretically, implementation and experi- ments will be left as the future work.
Notice that we say if we directly optimize our IB framework with MINE, one of the problems that the function might be unstable in RL, yet in section and experiments A.O). we directly use MINE to visualize the MI. This is because when we optimize our framework, we need to start to update every training step. While when we visualize the MI, we start to update T every 2000 training steps. Considering the computation efficiency, every time we start to update T when we use MINE to optimize our framework, we must update in one step few steps. While when visualizing the MI, update T in 256 steps. Besides, we reset parameter 7 every time we begin to update T when we visualize the MI, clearly we can't do this when optimizing our framework since it's min-max optimization problem.
13

---

Under review as a conference paper at ICLR 2020
A.6 Study the cntemmatioo-botleeneee perspective in RL
Now we introduce the experimental settings of MI visualization. And we show that the agent in RL usually tends to follow the information E-C process.
We compare the M ?)) between A2C and A2C with our framework. Every 2000 update steps(2560 frames each step), re-initialize the parameter 17 then sample batch of inputs and their relevant representations {X,, Z, = o(Xi, e) = 64, and compute the MI with MINE. The learning rate of updating 7) is same as openai-baselines A2C: 0.0007, training steps is 256 and the network architecture can be found in our code file "policy.py"
Figure(3) is the MI visualization in game Qbert. Note that there is a certain degree of fluctuations in the curve. This is because that unlike supervised learning, the distribution of datasets and learning signals R" keep changing in reinforcement learning: R (X) changes with policy and when и gets better, the agent might. find new states, in this case, T(X, Z) might increase again because the agent needs to encode information from new states in order to learn better policy. Yet finally, the MI always tends to decrease. Thus can say that the agent RL usually tends to follow the information E-C process.
mutuaLinfo_jos
mutual_info_ooss
My Mmh M
2.50
2.50
Wowl Mm Mas nu
2.30
2.30
2.30
2.10
190
190
0.000
4.000k
8.000k
12.00k
18.00k
20.00k 24.0k
0.000
4,000k
8.000k
12.00k
16.00k
20.00k
24.00
(a) MI in A2C
(b) MI in A2C with our framework
Figure 3: Mutual information visualization in Qbert. As policy gets better, the agent might find new states, in this case, I(X, Z) might increase again because the agent needs to encode information from new states in order to learn better policy. Yet finally, the MI always tends to decrease. Thus follows the information E-C process.
We argue that it's unnecessary to compute I(Z,Y) like (Shwarte-Ziv Ziv & Tishby, 2017): According to (.) the training loss coalinually decreases supervised learning (Reward continually increases as showed in (ggure()(a) in reinforcement learning), I(2,Y) must increase gradually.
Reward
We also add some additional experimental results of MI visualization in the appendix A.7
A.7 Additional experimental results of performance and MT visualization
This section we add some additional experimental results about our framework.
Notice that in game MsPacman, performance of A2C with framework is worse than regular A2C. According to the MI visualization MsPacman in figure(E)(b), we suspect that this is because A2C with our framework drops the information from inputs excessively that it misses some information relative to the lcarning process. To see it accurately, in (g the orange curve, which denotes A2C with framework, from step(x-axis) 80 to 100, suddenly drops plenty of information. Meanwhile, in figure (b), from step(x- axis) 80 to 100, the rewards of orange curve start to decrease.
As showed in figure D unlike Pong, Breakout, Qbert and some other shooting games, the frame of MsPacman contains much more information related the reward: The walls, the ghosts and the tiny beans everywhere. Thus if the agent drops information too fast, it may hurt the performance.
14

---

Under review as a conference paper at ICLR 2020
(a) A2C-Assault
(b) A2C-MsPaccmn
(c)
Figure 4: Additional results of performance: Average rewards 10 episodes
w A Aamga
(a) MI in Assault
(b) MI in MsPacman
Figure 5: Additional results of MI visualization
(a) MsPacman
Figure 6: The raw frame of MsPacman
15


Reviews:
[{"review_id": "rJgp9Wr3Yr", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper investigates the information bottleneck approach in learning representations in reinforcement learning. The authors propose 1) combining actor-critic objective and IB loss; 2) deriving optimal target distribution of the composite objective; 3) optimizing a lower bound of target function and using SVGD for optimization; 4) empirically show that the proposed method achieves promising results; 5) empirically show that there exists E-C process in RL; 6) discuss the relationship with MINE.\n\nPros:\n1. Well organized paper with a clean presentation.\n2. The technical contribution is enough (composite loss, optimal target distribution, using SVGD to solve lower bound, comparing with MINE)\n3. The experimental results are promising.\n\nCons/Questions:\n1. In Sec A.2, deriving the target distribution by KKT condition seems more rigorous (P_phi should be non-negative and sum up to 1). Using gradient equal to zero gets the same results is because the constraints are not active here (the optimal solution is not on the boundary of simplex).\n\n2. If not use the lower bound of the target distribution (use P rather than U), is there any method to solve this?\n\nOverall, I found this paper is interesting, with good theoretical contributions and promising empirical results.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I carefully checked the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly."}, "extra_text": {}}, {"review_id": "SJlFMnxCtS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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).\n\nIn spite of the interesting theoretical contributions, I have to recommend rejection as the current empirical evaluation of the proposed approach is extremely limited, making it difficult to assess its benefits over more straightforward algorithms.\n\nOn the positive side, the authors derive a sensible approach to IB representation learning in RL, and provide solutions to the optimization challenges it leads to. I did not have time to check all the maths in the Appendix (I only went through the derivations in A.1 and A.2), but they seem to make sense overall (though it is unclear to me if the new algorithm proposed in A.5 is a practical one, so I am not taking it into account in this evaluation).\n\nThe key negative point is definitely the weak empirical evaluation. The main results are from a limited sample of 5 Atari games, when the full Atari benchmark has 10x more games and is known to exhibit high variance among games when comparing RL algorithms (the additional results from the Appendix on 3 additional games also show situations where the proposed method does not seem to help much, confirming that larger scale experiments are needed for a proper evaluation). In addition it seems like each algorithm is run only once (instead of using multiple seeds) and only over ~200K timesteps, which is three orders of magnitude lower than results typically reported on Atari. Another issue is that there is no comparison to other representation learning techniques (like those mentioned in the related work section, or the recent \"Unsupervised State Representation Learning in Atari\"), nor to a natural and more straightforward variant of the proposed method where Z would simply be sampled from a (learned) Gaussian distribution Z ~ N(mu(X), var(X)), which at first sight seems like an easier-to-optimize objective (using the reparameterization trick)… I may be wrong, but then this should probably be explained in the paper (I realize that the proposed approach is more general, but then it should be shown how this extra flexibility can lead to improved results). Finally, the impact on runtime performance is not analyzed: how much slower do A2C / PPO become when optimizing the mutual information term with SVGD? Overall it is really unclear that better RL results can be obtained through this technique.\n\nAnother important issue is that I found the paper rather difficult to follow, due to some inconsistent / unclear notations or equations. Here are the main ones I noted:\n•\tThe discount factor is not accounted for in the derivation of the objectives in eq. 1-2 (I know this is often the case in practice, but the reason for dropping it should at least be mentioned)\n•\tThe jump from V^pi(X) to V^pi(Z) at the end of Section 3 is explained too succintly. It suggests that V^pi(Z) must be constant (equal to V^pi(X)) over all values of Z that may be sampled from X, which as far as I can tell is not the case in the rest of the paper. It is also unclear whether pi depends on Z or X. And the notation J(Z) makes it look like J does not depend on X, but it seems like it does because even if pi depends only on Z, by its definition V^pi(Z) actually still depends on X (this also leads to weird equations like eq. 33 where X does not appear in the right-hand side). Overall this is rather confusing.\n•\tThe « for every X » at the top of p. 4 does not make sense to me, due to the term I(X, Z) in eq. 3 where X is a random variable and thus does not take a specific value.\n•\tThe paragraph below eq. 4 is a bit confusing. It looks like it amounts to saying that Y is a Gaussian around V(Z)?\n•\tEq. 5 suggests that R depends on Z, but shouldn’t it depend on X? If it depended on Z, then wouldn’t it influence the optimization since by modifying P(Z|X) we can control the distribution on Z and thus the distribution on R?\n•\tIn eq. 10 it is unclear whether L1 and L2 contain the expectation, also L2 is defined as a function of both theta and phi but seems to depend only on theta\n•\tBelow eq. 15 it is said that « P is the distribution of Z » but P does not appear in eq. 15\n•\tIn eq. 16 should the phi on the left hand side be phi* as in eq. 15?\n•\tBelow eq. 16 it is said « Notice that C has been omitted », but it is unclear whether it was not included to alleviate notations or because it disappears naturally in the mathematical derivation of eq. 16\n•\tThe motivation for introducing zeta below eq. 17 is unclear, especially since it seems to play an important role considering that zeta = 0.005 << 1 is used in the experiments (with no explanation as to how this specific value was selected)\n•\t\\hat{L}(Z_i, theta, phi) (between eq. 17 and 18) does not seem to be defined\n•\tWhat is the motivation for using the Gaussian U(Z) as described in Section 5? In particular I find it weird that it depends on X_i, while U(Z) is supposed to replace the marginal P(Z) and not the conditional P(Z | X)\n\nMinor points:\n•\tIn the definition of Y_t = R_t = sum_i=0^n-2 … above eq. 4, I think the sum should be up to n-1 for an n-step return\n•\tEq. 4 uses R_t on the right hand term instead of Y_t which looks weird\n•\tTheorem 1 states « Assume that for any epsilon > 0, … »: « for any » should probably be « there exists », since if the inequality was true for any epsilon, it would imply both mutual informations are equal (also the formulation of the theorem does not make it clear that the last inequality is the main result)\n•\tFootnote 1 p. 4: I(X, Y) should be I(X, Z)\n•\tIn eq. 15 the phi below the argmax should be in bold\n•\tPlease add a reference that the reader can refer to in order to understand where eq. 20 is coming from\n•\t« Apparently » is used in a couple of places but should probably be replaced with another word\n•\tWhen U is uniform, how do you choose its support?\n•\tAfter « median of pairwise distances between the particles », I think i=1 should be j=1\n•\tIt is unclear to me what « A2C with noise is »: it is said that the same phi(X, epsilon) is used « as A2C with our framework », but is it the same phi as A2C with uniform SVIB or A2C with Gaussian SVIB? And whichever it is, how does it differ from the one it is equal to?\n•\t« we add 21 to all four curves in order to make exponential moving average »: I do not understand that sentence\n•\t« we set the number of samples as 26 for the sake of computation efficiency »: I fail to see how going from 32 to 26 is going to make a major difference in computational efficiency\n\nUpdate: thank you for your response, but in the absence of a revised version addressing my concerns (as well as those from the other reviewers), I cannot increase my rating", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "HJxLVnA3FB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThe authors discussed the information bottleneck in the setting of RL thoroughly, and proposed a novel target Y (\\log P) in the scenario of policy optimization. Further the authors derive the variational lower bound under the setting of RL, and proposed to utilize amortized SVGD to learn the proposed variational lower bound. The authors also examine the mutual information using recent proposed MINE, which also indicates the effectiveness of the proposed method.\n\nUtilizing recent advanced techniques in representation learning and applying it to RL tasks which have high dimensional images state spaces can accelerate the training processing. Though the problem that the authors studied is important, the proposed method lacks novelty and has not been well evaluated comparing with other representation techniques. Specifically, I have several concerns about the paper.\n\n- Variational lower bound for RL. I admit the proposed \\log p(y|z) is novel (although heuristic), but the rest of the results can be easily derived from variational lower bound framework. I suggest the author clarify and discuss about the derivation of this part with techniques introduced in variational bounds of mutual information (Poole et al. 2019).\n\n- The derivation of updating of $\\phi$ (Equation 18). The derivation is totally unnecessary since it is exactly the same as the updating formulation of Amortized SVGD (Feng et. al 2017). See Sec 3 in Feng et. al 2017 for details. I think the author should motivate the usage of implicit models (the embedding function $\\phi(x, \\epsilon)$) and then show that we should optimize the implicit encoders using amortized SVGD.\n\n- The motivation of using amortized SVGD is unclear. If the authors parameterized the embedding function $\\phi(z)$ with a tractable density model such as Gaussian distributions, it is still possible to optimize the proposed lower bound. Ablation studies should be conducted to demonstrate why we can enjoy the benefits of implicit models in the setting of RL. \n\n- Baselines including other representation learning methods should be compared. The authors only compared with naive A2C or PPO, while there are many other methods which can be directly utilized to learn the representation during the policy optimization process. One example is  Contrastive Predicting Coding (Oord et. al 2018) (I noticed that the author mentioned MINE does not work well in this setting, but at least the variational lower bound using explicit models such as Gaussian encoders should be compared).\n\n- (Minor things) The plots in the figures are hard to read.  Figure (2) shows the advantages of the proposed methods which are really hard to read and distinguish several proposed methods. It would be nice to have a table to summarize the final performance (From the plots It is hard to conclude which one is the best). In addition, figure 1&3 is unreadable (Taking screenshots of tensorboard plots is really bad). \n\nOverall I think the paper studied an interesting problem while the motivation and the advantage of the proposed method are still unclear, which requires more discussion and comparison with other methods which can be utilized for representation learning. \n\n\nReference Papers: \n\nPoole, Ben, et al. \"On variational bounds of mutual information.\" arXiv preprint arXiv:1905.06922 (2019).\n\nFeng, Yihao, Dilin Wang, and Qiang Liu. \"Learning to draw samples with amortized stein variational gradient descent.” UAI 2017.\n\nOord, Aaron van den, Yazhe Li, and Oriol Vinyals. \"Representation learning with contrastive predictive coding.\" arXiv preprint arXiv:1807.03748 (2018).\n\n", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose to use the information bottleneck to learn state representations for RL. They optimize the IB objective using Stein variational gradient descent and combine it with A2C and PPO. On a handful of Atari games, they show improved performance.

The reviewers primary concerns were:
*Limited evaluation. The method was only evaluated on a handful of the Atari games and no comparison to other representation learning methods was done.
*Using a simple Gaussian embedding function would eliminate the need for amortized SVGD. The authors should compare to that alternative to demonstrate the necessity of their approach.

The ideas explored in the paper are interesting, but given the issues with evaluation, the paper is not ready for acceptance at this time.
Decision: Reject