Which Mutual-Information Representation Learning Objectives are Sufficient for Control?
Authors: Kate Rakelly
Abhishek Gupta
Carlos Florensa
Sergey Levine
Venue: NeurIPS 2021 Poster (2021)
Keywords: reinforcement learning
representation learning
mutual information

Abstract: Mutual information (MI) maximization provides an appealing formalism for learning representations of data. In the context of reinforcement learning (RL), such representations can accelerate learning by discarding irrelevant and redundant information, while retaining the information necessary for control. Much prior work on these methods has addressed the practical difficulties of estimating MI from samples of high-dimensional observations, while comparatively less is understood about which MI objectives yield representations that are sufficient for RL from a theoretical perspective. In this paper, we formalize the sufficiency of a state representation for learning and representing the optimal policy, and study several popular MI based objectives through this lens. Surprisingly, we find that two of these objectives can yield insufficient representations given mild and common assumptions on the structure of the MDP. We corroborate our theoretical results with empirical experiments on a simulated game environment with visual observations. 

Which Mutual-Information Representation Learning Objectives are Sufficient for Control?
Kate Rakelly*
Abhishek Gupta
Carlos Florensa
Sergey Levine
University of California, Berkeley frakelly, abhigupta, florensace, svlevine (eeecs.berkeley.duu
Abstract
Mutual information (MI) maximization provides an appealing formalism for learn- ing representations of data. the context reinforcement learning (RL), such representations accelerate learning by discarding irrelevant and redundant in- formation, while retaining the information necessary for control. Much prior work these methods has addressed the practical difficulties of estimating MI from samples of high-dimensional observations, while comparatively less is understood about which MI objectives yield representations that sufficient for RL from a theoretical perspective. In this paper, we formalize the sufficiency of a state representation for learning and representing the optimal policy, and study several popular MI based objectives through this lens. Surprisingly, we find that two of these objectives can yield insufficient representations given mild and common as- sumptions on the structure of the MDP. We corroborate our theoretical results with empirical experiments on simulated game environment with visual observations.
Introduction
Deep reinforcement learning (RL) algorithms in principle capable of learning policies from high-dimensional observations, such as camera images 49 KD 332. However, policy learning in practice faces bottleneck in acquiring useful representations of the observation space [58]. State representation learning approaches aim to remedy this issue by learning structured and compact representations on which to perform RL. useful state representation should be sufficient learn and represent the optimal policy or the optimal value function, while discarding irrelevant and redundant information. Understanding whether or not an objective is guaranteed to yield sufficient representations is important, because insufficient representations make impossible to solve certain problems. For example, an autonomous vehicle would not be able to navigate safely if its state representation did not contain information about the color of the stoplight in front of it. With the increasing interest in leveraging offline datasets to learn representations for RI. 19 35 c the question of sufficiency becomes even more important to understand if the representation is capable of representing policies and value functions for downstream tasks.
While a wide range of representation learning objectives have been proposed in the literature 41 in this paper focus on analyzing representations learned by maximizing the mutual information (MI) between random variables. Prior work has proposed many different MI objectives involving the variables of states, actions, and rewards at different time-steps 52
58
While much prior 52 27].
work has focused on how to optimize these various MI objectives in high dimensions 61 we focus instead on their ability to yield theoretically sufficient representations. We find that two commonly used objectives are insufficient for the general class of MDPs, in the most general case, and prove that another typical objective is sufficient. We illustrate the analysis with both didactic examples in which MI can be computed exactly and deep RL experiments in which we approximately maximize
Correspondence to rakelly wocs.berkeley..d.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

the MI objective to learn representations of visual inputs. The experimental results corroborate our theoretical findings, and demonstrate that the sufficiency of a representation can have a substantial impact on the performance of an RI. agent that uses that representation. This paper provides guidance to the deep RL. practitioner on when and why objectives may work well or fail, and also provides a formal framework to analyze newly proposed representation learning objectives based on MI.
2 Related Work
In this paper, we analyze several widely used mutual information objectives for control. In this section we first review MI-based unsupervised learning. then the application of these techniques to the RL setting, Finally, we discuss alternative perspectives on representation learning in RL.
Mutual information- -basd unsupervised learning. A common technique for unsupervised repre- sentation learning based on the InfoMax principle [3] y is to maximize the MI between the input and its latent representation subject to domain-specific constraints 77. This technique has been applied to learn representations for natural language [4] video 65 and images [5] 2 and even
policy learning via RL in high dimensions [62] To address the difficulties of estimating MI from samples 45 and with high- -igennimeniional inputs [1]. much recent work has focused on improving MI estimation via variational methods [6] 55
In this work we are concerned with analyzing the MI objectives, and not the estimation method. In our experiments with image observations, we use noise contrastive estimation methods [24 though other choices could also suffice.
Mutual information objectives in RL. RL adds aspects of temporal structure and control to the standard unsupervised learning problem discussed above (see Figure D. This structure can be leveraged by maximizing MI between sequential states, actions, or combinations thereof. Some Several maximize MI objectives that closely resemble the / objective we introduce in Section 5 while others optimize related objectives by learning latent forward dynamics
works omit the action, maximizing the MI between current and future states 4
52 64
prior works [51]
57 60,
26 39]. Multi-step inverse models, closely related to the Jine objective (Section I Single-step inverse models have been deployed as regularization of forward models 73 2 and as an auxiliary loss for policy gradient 5 Our result regarding the sufficiency of this objective similar an example explored in Misra et al. 48; however, we relate ours to MI objectives. The MI objectives that we study have also been used as reward bonuses to improve exploration, without impacting the representation, in
models 77 33 4
have been used to learn control-centric representations 71 23
RL 55
the form of empowerment 37 36 50, 40 and tnforratioo-thhorrtiii curiosity [63
Representation learning for reinforcement learning. In RL, the problem of finding a compact
state space has been studied as state aggregation or abstraction 6 42 Abstraction schemes include bisimulation [22], homomorphism [561. utile distinction 46 and policy irrelevance [30]. While
efficient algorithms exist for MDPs with known transition models for some abstraction schemes such as bisimulation 18 22 in general obtaining error-free abstractions is impractical for most problems of interest. For approximate abstractions prior work has bounded the sub-optimality of the policy II 13 M as well as the sample efficiency [38 68 16, with some results in the deep learning setting 111 [1]. In this paper, we focus whether a representation can be used to learn the optimal policy, and not the tractability of learning. Li et al.
42 shares this focus; while they establish convergence properties of Q-learning with representations satisfying different notions of sufficiency, we leverage their -usfffiicie criteria to evaluate representations learned via MI- based objectives. Alternative approaches to representation learning for RL include priors based on the structure of the physical world [II or heuristics such as disentanglement [6], meta-learning general value functions [69] predicting multiple value functions 110 17 29 and predicting domain-
specific measurements 47 15]. We restrict our analysis to objectives that can be expressed as
Ml-maximization. In paper focus the representation learning problem, disentangled from
exploration, a strategy shared by prior works [9
3 Representation Learning for RL
The goal of representation learning for RL is to learn a compact representation the state space that discards irrelevant and redundant information, while still retaining sufficient information to represent policies and value functions needed for learning. In this section we formalize this problem, and propose and define the concept of sufficiency to evaluate the usefulness of representation.
2

---

3.1 Preliminaries
We begin with brief preliminaries of reinforcement learning and mutual information.
Reinforcement learning. A Markov decision process (MDP) is defined by the tuple (S, (S.ATT. T.r). where S the set of states, A the set of actions, T: S Ax8 + [0. 1] the state transition distribution. and r :S + IR the reward function D We will use capital letters to refer to random variables and lower case letters refer to values of those variables (e.g., S is the random variable for the state and s is specific state). Throughout our analysis we will often be interested in multiple reward functions, and denote set of reward functions as R. The objective of RL is to find a policy that maximizes the sum of discounted returns R. for a given reward function and we denote this optimal policy as T = arg max, x E, D (rr(St, Ar)] for discount factor 7. We also define the optimal Q-function as (s,,, a+) = B. |-==
1 7 A,) ar). The optimal Q-function satisfies the recursive Bellman equation, (s,, a.) = r(s,, at) + 7/ppp=+ (y a.) arg max (S a441). An optimal policy and the optimal Q-function are related according to x (s) = arg max (s..a).
Mutual information. In information theory, the mutual information (MI) between two random variables, X and Y, is defined as 12
p(x, - I(X;Y) = Ep(ea)) log p(z)p(y)
The first definition indicates that MI can be understood as relative entropy (or KL-divergencee, while the second underscores the intuitive notion that MI measures the reduction in the uncertainty of one random variable from observing the value of the other.
Representation learning for RL. While state aggregation meth- ods typically define deterministic rules to group states in the repre- sentation 6 42), MI-based representation learning methods used for deep RL treat the representation Accordingly, we formalize a representation as a stochastic mapping between original state space and representation space.
A++1
random variable [51] (22 (33
S₁
S
Z:
Ze+1
A stochastic representation dz is mapping from states s E S to a probability distribution p(Z S = s) over elements of new representation space E Z.
Definition 1.
Figure 1: Probabilistic graphi- cal model illustrating the state representation learning prob- lem: estimate representation Z from original state S.
In this work we consider learning state representations from data by maximizing objective J. Given an objective J, we define the set of representations that maximize this objective as Φ (Øz Øz E arg max J()) Unlike problem formulations for partially observed settings 770 26 39), we assume that S is a Markovian state; therefore the representation for a given state conditionally independent of the past states, a common assumption the state aggregation 42 See Figure ! for depiction of the graphical model.
literature
3.2
Sufficient Representations for Reinforcement Learning
We now turn to the problem of evaluating stochastic representations for RL. Intuitively, we expect useful state representation to be capable of representing an optimal policy in the original state space.
Definition 2. A representation Φz is x" "--sufficient with respect to set of reward functions R. if Vr E R, (s1) = (2((2) = (Als:) .;
When a stochastic representation Φ produces the same distribution over the representation space for two different states 8 and 82 say it aliases these states. Unfortunately, already proven in Theorem of Li et al. [42 for the more restrictive case of deterministic representations, being able to represent the optimal policy does not guarantee that it can be learned via RL in the representation space. Accordingly, we define a stricter notion of sufficiency that does guarantee the convergence of
We restrict our attention to MDPs where the reward can be expressed as a function of the state, which is fairly standard across broad set of real world RL problems
3

---

of reward functions R if Q-learning to the optimal policy in the original state space (refer to Theorem 4 of Li et al. 42 for the proof of this).
Definition 3. A representation Φz is Q' -suffficient with respect vrr E R, (s1) = (2( = (a (1, = (a,ss)
of reward functions R if
Note that -usfffficinncy implies ग -sfficciency since an optimal policy can be recovered from the optimal Q-function via (s) = arg max, (s, a) 666;; however the converse is not true. We emphasize that while -uuffiiciinnyy guarantees convergence, does not guarantee tractability, which has been explored in prior work [38 [6]
We will further say that objective is sufficient with respect
all the representations that maximize that objective Φ are sufficient with respect to every element of R according to the definition above. Surprisingly, we will demonstrate that not all commonly used objectives satisfy this basic qualification.
Mutual Information for Representation Learning in RL
In our study, consider several MI objectives proposed in the literature.
A commonly sought characteristic of state representation is ensure retains maximum predictive power over future state representations. This property is satisfied by representations maximizing the following MI objective,
Forward information:
f (Z Ze, A.) H (Z++1) (Zee A.).
(2)
We suggestively name this objective "forward information" due to the second term, which is the entropy of the forward dynamics distribution. This objective and closely related ones have been used in prior works [5] 57 60
State-only transition information: Several popular methods 52 tive, but do not include the action:
optimize similar objec-
Jstate
I(Z++ki Z.)
(((+
(3)
As will show, the exclusion of the action can have profound effect the characteristics of the resulting representations.
Inverse information: Another commonly sought characteristic of state representations is to retain maximum predictive power of the action distribution that could have generated an observed transition from S; to S1+1- Such representations can be learned by maximizing the following information theoretic objective:
Jino T(Ar; Ze+4 H H (A, 7
We suggestively name this objective "inverse information" due to the second term, which is the entropy of the inverse dynamics. wide range of prior work learns representations by optimizing Intuitively, inverse models allow the representation to capture only the elements of the state that are necessary to predict the action, allowing the discard of potentially irrelevant information.
closely related objectives 24 58
2
53
A
73
5 Sufficiency Analysis
In this section we analyze the sufficiency representations obtained by maximizing each objective presented in Section 4 To focus on the representation learning problem, we decouple it from RL by assuming access
dataset of transitions collected with policy that reaches all states with non-zero probability, which can then be used to learn the desired representation. We also assume that distributions, such as the dynamics inverse dynamics, can be modeled with arbitrary accuracy, and that the maximizing set of representations for given objective can be computed. While these assumptions might relaxed any practical RL algorithm, and exploration plays a confounding role, the ideal assumptions underlying our analysis provide the best-case scenario for objectives to yield provably sufficient representations. other words, objectives found be provably insufficient under ideal conditions will continue to be insufficient under more realistic assumptions.

---

5.1 Forward Information
In this section we show that representation that maximizes f is sufficient for optimal control under any reward function. This result aligns with the intuition that representation that captures forward dynamies can represent everything predictable in the state space, and can thus be used to learn the optimal policy for any task. This strength can also be weakness if there are many predictable elements that are irrelevant for downstream tasks, since the representation retains more information than is needed for the task. Note that the representation can still discard information in the original state, such as independent random noise at each timestep.
Proposition fwa sufficient for all reward functions.
Proof. (Sketch) We first show in Lemma that if Zz, AA are maximally informative of Zt+1, they are also maximally informative of the return R. Thanks to the Markov structure, we then show in
Lemma 2 that (2) R2 Zz. A,) = p( [, s, At). In other words, given áz, additionally knowing 5 doesn't change our belief about the future return. The Q-value the expectation of the return, so Z has as much information about the Q-value S. See Appendix A. for the proof. D
5.2 State-Only Transition Information
While Jatate is closely related to I fudd we now show that Jstate not sufficient.
Proposition 2. Ustate is not sufficient for all reward functions.
Proof. We show this by counter-example with the deterministic-traniitton MDP defined in Figure (left), For all four states, let the two actions an and be equally likely under the policy distribution. In this case, each state gives no information about which the two possible next states more likely; this depends on the action. Therefore, a representation maximizing J Jstate is free to alias states with the same next-state distribution, such as So and Sy. An alternative view that such representation can maximize staate H(Z++)) H H(Z++ 72,) by reducing both terms in equal amounts aliasing Sg and 83 decreases the marginal entropy as well the entropy of predicting the next state starting from
However, this aliased representation is not capable of representing the optimal policy, which must distinguish Sg and $3 in order to choose the correct action to reach $2, which yields reward.
In Figure 2 (right), we illustrate the insufficiency of Jstate computationally. by computing the values of Jsstate and fwd for different state representations of the above MDP, ordered by decreasing compression (increasing S); left to right. At the far left of the plot is the representation that
hed
a0
Value of MI objective
retur (optimal)
---- 01
10
RL retum 0.5
SO
S2
06
02
$
$3
0.0
0.2
0.6 0.8 10
12 14
State representations, ordered by I(Z: S)
Figure 2: (left) representation that aliases the states So and 53 into single state maximizes Jstate but is not sufficient to represent the optimal policy which must choose different actions in So and 83 to reach 82 which yields reward. (right) Values of Istate and J for a few representative state representations, ordered by increasing T(Z; S). The representation that aliases So and sy (plotted with diamond) maximizes Jatates but the policy learned with this representation may not optimal (as shown here). The original state representation (plotted with a star) is sufficient.
5

---

aliases all states, while the original state representation is at the far right (plotted with star). The representation that aliases states So and 53 (plotted with diamond) maximizes states but is insufficient to represent the optimal policy. Value iteration run with this state representation achieves only half the optimal return (0.5 versus 1.0).
5.3 Inverse Information
Since Jinv preserves state elements that the agent can influence with its actions, we might think that good candidate for sufficient objective that is capable of discarding more information than J' fwd- However, here we show with a counterexample that Jan is not sufficient. Intuitively, an insufficient representation can be obtained by maximizing Jine in an MDP when the reward function depends on elements outside the agent's s control. We then show that additionally requiring the representation to represent the immediate reward is not enough to resolve this issue.
Proposition 3. Dimo is not sufficient for all reward Additionally, adding I(2;; Ze) the objective does not make it sufficient.
Proof. We show this by counter-example with the deterministic-trasitiiie MDP defined in Figure B (left). Consider the representation that aliases the states So and Si. This state representation wouldn's be sufficient for control because the same actions taken from these two states lead to different next states, which have different rewards (ag leads the reward from So while aj leads the reward from $1). However, this representation maximizes Jine because, given each pair of states, the action is identifiable. Interestingly, this problem cannot be remedied by simply requiring that the representation also be capable of predicting the reward at each state. Indeed, the same insufficient representation from the above counterexample also maximizes this new objective as long as the reward Sy and $ are the same.
Analogous to the preceding section, in Figure 3 (right), we plot the values of the objectives Jine and fed for state representations ordered by increasing I(Z; S) value (decreasing compression). The representation that aliases states and s (plotted with diamond) maximizes Jinu, but is insufficient to represent the optimal policy; value iteration with this state representation achieves half the optimal return (0.5 versus 1.0).
a0
011
S2
18
Sed
Value of objective
Jav
$0
S3
RL return 1.0 (optimal)
1.4
RL return 0.50
12-
81
S4
08-
0.
85
1.0
1'4 1'6 1.8
State representations, ordered by ((Z; S)
Figure (left) In this MDP. representation that aliases the states So and S; into single state maximizes Jsnv. yet is not sufficient to represent the optimal policy, which must distinguish between Sg and 81 order to take a different action (towards the high-reward states outlined green). (right) Values of line and Jy ww for few selected state representations, ordered by increasing T(Z; S). The representation that aliases So and s (plotted with diamond) maximizes Jine. but learn the optimal policy. Note that this counterexample holds also for Jine + I [(R; Z).
sufficient to

---

Experiments
To analyze whether the conclusions of our theoretical analysis hold in practice, we present experiments studying MI-based representation learning with image observations. We do aim show that any particular method necessarily better or worse, but rather to determine whether the sufficiency arguments that we presented can translate into quantifiable performance differences in deep RI..
6.1 Experimental Setup
To separate representation learning from RL, we first optimize each representation learning ob- jective dataset of offline data, similar to the protocol in Stooke et al. (44] Our datasets consist of 50k transitions collected from uni- form random policy, which is sufficient to cover the state space in environments. We then freeze the weights the state encoder learned in the first phase and train RL agents with the representation as state input. We perform our experiments on variations on the pygame (59] video game catcher, in which the agent controls paddle that it can move back and forth catch fruit that falls from the top of the screen (see Figure D. A positive reward given when the fruit is caught and a negative reward when the fruit is not caught. The episode terminates after one piece of fruit falls. We optimize I fwd and Jatate with noise contrastive estimation [24,, and Jine by training an inverse model via maximum likelihood. We also include the performance of an agent trained "end-to-end" with pixel inputs. For the RL algorithm, we use the Soft Actor- Critic algorithm [25). modified slightly for the discrete action distribution. Please see Appendix A.2 for full experimental details.
Figure 4: (left) Original catcher game: agent (grey paddle) moves left or right to catch fruit (yellow square) that falls, (right) Variation catcher-grip agent must open the gripper to catch fruit.
6.2 Computational Results
In principle, representation learned with Jine may not be sufficient to solve the catcher game. Because the agent does not control the fruit, representation maximizing Jine might discard that information, thereby making it impossible to represent the optimal policy. We observe in Figure 5 (top left) that indeed the representation trained maximize Jine results in RL agents that converge slower and to lower asymptotic expected return. Further, attempting to learn decoder from the learned representation the position of the fruit incurs high error (Figure S bottom left), indicating that the fruit is not precisely captured by the representation. The characteristics of this simulated game are representative of realistic tasks. Consider an autonomous vehicle that is stopped stoplight. Because the agent does not control the state of the stoplight, it may not be captured in the representation learned by Jine and the resulting RL policy may choose run the light.
In the second experiment, we consider a failure mode of Jstate- We augment the paddle with a gripper that the agent controls and must be open in order to properly catch the fruit (see Figure + right). Since the change in the gripper is completely controlled by a single action, the current state contains no information about the state of the gripper in the future. Therefore, a representation maximizing J Jatate might alias states where the gripper is open with states where the gripper is closed. Indeed, we see that the error predicting the state of the gripper from the representation learned via Jstate is about chance (Figure 5 bottom right). This degrades the performance of an RL agent trained with this state representation since the best the agent can do is move under the fruit and randomly open or close the gripper (Figure 5] top right). In the driving example, suppose turning on the headlights incurs positive reward if it's raining but negative reward if it's sunny. The representation could fail to distinguish the state of the headlights, making it impossible to learn when to properly use them. J Jywd produces useful representations in all cases, and equally or more effective than learning representations purely from the RL objective alone (Figure 5)
7

---

catcher
catcher-grip
1.0
1,00
0.75
0.5
0.30
Return
Return
0.25
0.0
0.00
0.25
-0.00
0.00
0.15
0.30
0.45
0.00
0.00
0.15
0:30
0.45
0.00
Env Steps (Millions)
Env Steps (Millions)
end-to-end
J
Jnn
Istate
Obj.
Agent Err.
Fruit Err.
Obj.
Agent Emr.
Fruit Err.
Grip Err.
J fwd
0.1
0.1
0.1
0.1
0.3
Jino
0.15
0.37
Istate
0.1
0.1
0.46
Figure 5: (left) Policy performance using learned representations as state inputs to RL, for the catcher and catcher-grip environments. (right) Error in predicting the positions of ground truth state elements from each learned representation. Representations maximizing Jine need not represent the fruit, while representations maximizing state need represent the gripper, leading these representations to perform poorly in catcher and calcher-grip respectively.
6.3 Increasing visual complexity via background distractors
Here test whether sufficiency of representation can impact agent performance in more visually complex environments by adding background distractors to the agent's 's observations. We replace the background of the game with randomly generated images of 10 circles of different colors (Figure 6
Analogous to Section 6 in Figure 7 we show the performance of an RL agent trained with the frozen representation as input (top), as well as the enror of decoding true state elements from the representation (bottom). end-to-end RL from images performs poorly, demonstrating the need for representation learn- ing. As predicted by the theory, the represen- tation learned by Jine fails both games, and the representation learned by Istate fails in the catcher- grip game. The difference in perfor- mance between sufficient and insufficient objec- tives is even more pronounced in this setting. With more information present in the form of the distractors, insufficient objectives that do not optimize for representing all the required state information may be "distracted" by representing the background objects instead, resulting in low performance. In Appendix A.3 we experiment with visual distractors that are temporally correlated across time. We also consider variations on our analysis, evaluating how well the representations pre- dict the predict the optimal Q" Appendix A.3 and experimenting with different data distribution for collecting the representation learning dataset in Appendix A.4 These results demonstrate that using insufficient representation learning objectives can degrade performance of an RL agent, not that they necessarily will. Our aim is to provide an illustration of the potential impact of insufficient representations, to underscore the utility of sufficiency as a tool in designing representation learning objectives and debugging system failures.
fruit
both games,
agent
Figure 6: Example 64x64 pixel agent observations with background distractors (circles). randomly generated at each time step, that increase the diffi- culty of learning good representation.

---

catcher
catcher-grip
1.0
0.5
0.0
Return
Return
-0.2
0.0
0.4
-0.5
0.6
-1.0
-1.0-
0.00
0.15
0.30
0.45
0.00
0.15
0.30
0.46
0.60
Env Steps (Millions)
Env Steps (Millions)
end-to-end
J
J
Istate
Obj.
Agent Err.
Fruit Err.
Obj.
Agent Emr.
Fruit Err.
Grip Err.
fwd
0.13
0.14
0.12
0.1
0.25
Jene
0.15
0.47
Istate
0.1
0.14
0.5
Figure 7: With background distractors added to the observations, the state representation learned via Jine fails capture the fruit object accurately in catcher (left), and the representation learned via Jstate continue to perform poorly at capturing the gripper state catcher- grip (right). The performance of the insufficient representations even lower than in the clean background experiment.
7 Discussion
In this work, we analyze which common MI-based representation learning objectives are guaranteed to yield representations provably sufficient for lcarning the optimal policy. We show that two common objectives Jstate and Jine yield theoretically insufficient representations, and provide proof of sufficiency for J fwd- We then show that insufficiency of representations the performance of deep RL agents with experiments on a simulated environment with visual observations. While insufficient representation learning objective can work well for training RL agents on simulated benchmark environments, the same objective may fail for real-world system with different characteristics. We believe that encouraging a focus evaluating the sufficiency of newly proposed representation learning objectives can help better predict potential failures.
degrade
While sufficiency is critical criterion for representations, compression is also highly important. Our results thus highlight an important open problem in unsupervised learning for RL: defining representation learning objectives that provably sufficient and able to discard more information than J Is the "smallest" representation that maximizes J fwd the smallest sufficient representation possible without making further assumptions assumptions on the MDP structure or reward function would suffice to make Jine and J state sufficient? For example, Istate is trivially sufficient when the environment dynamics and the agent's policy are deterministic. However, we hypothesize there may be more interesting MDP classes, related to realistic applications, in which generally insufficient objectives may be sufficient.
the MDP structure
reward function? What
Additionally, extending our analysis to the partially observed setting would be more reflective of practical applications. An interesting class of models to consider in this context are generative models such as variational auto-encoders [4]] Prior work has shown that maximizing the ELBO alone cannot control the content of the learned representation 28 maximizer of the ELBO would be sufficient, while other solutions would not necessarily be. We see these directions fruitful in providing deeper understanding of the learning dynamics of deep RL, and potentially yielding novel algorithms for provably accelerating RL with representation learning.
We conjecture that the zero-distortion
Acknowledgements. We would like to thank Ignasi Clavera, Chelsea Finn, and Ben Poole for insightful conversations at different stages this work. This research was supported by the DARPA Assured Autonomy program and ARL DCIST CRA W91 INF-12222--

---

References
[1] D. Abel, D. E. Hershkowitz, and M. L. Littman. Near optimal behavior via approximate state abstraction. In Proceedings of the 33rd International Conference International Conference on Machine Learning- Volume 48, pages 2915-2923, 2016.
[2] Agrawal, A. V. Nair, P. Abbeel, J. Malik, and S. Levine. Learning poke by poking: Experiential learning of intuitive physics. In Advances in neural information processing systems,
pages 5074-001 5082, 2016.
[3] A. A. Alemi, B. Poole, I. Fischer, J. V. Dillon, R. A. Saurous, and K. Murphy, Fixing a broken elbo. In International Conference on Machine Learning (ICML), 2018.
[4] A. Anand, E. Racah, S. Ozair, Y. Bengio, M.-A. Côté, and Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pages 8766 9766-8779, 2019.
[5] P. Bachman, R. D. Hjelm, and W. Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems, pages 15509-15519, 2019.
[6] J. C. Bean, J. R. Birge, and R. L. Smith. Aggregation in dynamic programming. Operations Research, 35(2):215
1987.
[7] S. Becker and G. E. Hinton. Self-organizing neural network that discovers surfaces in random- dot stereograms. Nature, 355(6356):661-663 1992.
[8] M. I. Belghazi, A. Baratin, S. Rajeswar, S. Ozair, Y. Bengio, A. Courville, and D. Hjelm.
Mine: mutual information neural estimation. arXiv preprint arXiv: 1801: 880100062, 2018.
[9] A. J. Bell and T. J. Sejnowski. An informatiion-maxiiizzztooo approach to blind separation and blind deconvolution. Neural computation, 7(6):1 121-112991111 29519 1995.
[10] M. Bellemare, W. Dabney, R. Dadashi, A. A. Taiga, P. S. Castro, N. Le Roux, D. Schuurmans, Lattimore, and C. Lyle. A geometric perspective on optimal representations for reinforcement learning. In Advances in Neural Information Processing Systems, pages 4360-4371, 2019.
[11] D. P. Bertsekas, D. A. Castanon, et al. Adaptive aggregation methods for infinite horizon dynamic programming. 1988.
[12] T. M. Cover. Elements of information theory. John Wiley & Sons, 1999.
[13] T. Dean and R. Givan. Model minimization in markov decision processes. In AAAI/IAAI, pages
106-111, 1997.
[14] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume (Long and Short Papers), pages 4171-4186, 2019.
[15] A. Dosovitskiy and V. Koltun. Learning to act by predicting the future. In International Conference on Learning Representations, 2016.
[16] S. S. Du, S. M. Kakade, R. Wang, and L. F. Yang. Is a good representation sufficient for sample efficient reinforcement learning? arXiv preprint arXiv: wwXvv:110.03016. 2019.
[17] W. Fedus, C. Gelada, Y. Bengio, M. G. Bellemare, and H. Larochelle. Hyperbolic discounting and learning over multiple horizons. arXiv preprint arXiv: 066600, 2019.
[18] N. Ferns, P. S. Castro, D. Precup, and P. Panangaden. Methods for computing state similarity in markov decision processes. Proceedings of the Twenty-Secom Conference on Uncertainty in Artificial Intelligence, pages 174-181, 2006.
[19] C. Finn, X. Y. Tan, Y. Duan, T. Darrell, S. Levine, and P. Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pages 512-519. IEEE, 2016.
[20] T. garage contributors. Garage: A toolkit for reproducible reinforcement learning research. https://gitthu: wttptttuww..cmmtloooorrroopppppp /gargge 2019.
[21] C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning continuous latent space models for representation learning. International Conference on Machine Learning, pages 2170-2179, 2019.
10

---

[22] R. Givan, T. Dean, and M. Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1-2): 2003.
[23] K. Gregor, D. Rezende, and D. Wierstra. Variational intrinsic control. arXiv preprini 1.07507, 2016.
[24] M. Gutmann and A. Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the Thirleenth International Conference Artificial Intelligence and Statistics, pages 297-304, 2010.
[25] Haarnoja, A. Zhou, P. Abbeel, and Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprint
2018. [26] D. Hafner, T. Lillicrap, Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Leamiing pages 2555-2565, 2019.
[27] R. D. Hjelm, A. Fedorov, Lavoie-Marchildonn, K. Grewal, A. Trischler, and Y. Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint arXiv: 1808. 06670, 2018.
[28] Huszár. Is maximum likelihood useful for representation learning?, 2017.
[29] M. Jaderberg, V. Mnih, W. M. Czarnekki T. Schaul, J. Z. Leibo, D. Silver, and K. Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397
2016.
[30] N. K. Jong and P. Stone. State abstraction discovery from irrelevant state variables. In IJCAI, volume 8, pages 752-757, 2005.
[31] R. Jonschkowski and O. Brock. Learning state representations with robotic priors, Autonomous Robots, 39(3):407-428, 2015.
[32] D. Kalashnikov, A. Irpan, Pastor, Ibarz, A. Herzog, E. Jang, D. Quillen, Holly, M. Kalakr- ishnan, V. Vanhoucke, et al. Scalable deep reinforcement learning for vision-based robotic manipulation. In Conference on Robot Learning, pages 651-673, 2018.
[33] M. Karl, M. Soelch, J. Bayer, and P. van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference on Learning Representations, 2016.
[34] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arKiv: 1312.6114, 2013.
[35] Kipf, E. van der Pol, and M. Welling. Contrastive learning of structured world models. arXiv preprint arXiv: /911:/991..224 2019.
[36] A. S. Klyubin, D. Polani, and C. L. Nehaniv. Empowerment: A universal agent-centric measure of control. 2005 IEEE Congress on Evolutionary Computation, volume pages 135-135. IEEE, 2005.
[37] A. S. Klyubin, D. Polani, and C. L. Nehaniv. Keep your options open: An information based driving principle for sensorimotor systems. PloS one, 3(12), 2008.
[38] T. Lattimore and C. Szepesvari. Learning with good feature representations in bandits and in rl with generative model. arXiv preprint 1rXiv:1911.07766 1.07676, 2019.
[39] A. X. Lee, A. Nagabandi, P. Abbeel, and S. Levine. Stochastic latent actor-critie: Deep reinforcement learning with a latent variable model. arXiv preprint 1907000997 2019.
[40] Leibfried, Pascual-Díaz, and Grau-Moya. unified bellman optimality principle combining reward maximization and empowerment. In Advances in Neural Information Processing Systems, pages 7867-7878, 2019.
[41] T. Lesort, N. Rodriguz J.-F. Goudou, and D. Filliat. State representation learning for control: An overview. Neural Nenworks, 108:379-392, 2018.
[42] L. Li, J. Walsh, and M. L. Littman. Towards a unified theory of state abstraction for mdps. In International Symposium Artificial Intelligence and Mathematics, ISAIM 2006, Fort Lauderdale, Florida, USA, January 4-6, 2006, 2006. Utl.bttp://anyttiin cs umass .mass.edu/ Sehathh0//prceedinnge
pddf
[43] R. Linsker. Self-organization in perceptual network. Computer, 21(3):105-117, 1988.
11

---

[44] B. Mazoure, R. T. d. Combes, T. Doan, P. Bachman, and R. D. Hjelm. Deep reinforcement and infomax learning. arXiv preprint 2Xiv:22006.07217, 2020.
[45] D. McAllester and K. Statos. Formal limitations the measurement of mutual information. arXiv preprini 0rXiv:/8/1.0425/1 2018.
[46] A. K. McCallum. Reinforcement Learning with Selective Perception and Hidden State. PhD thesis, University of Rochester, 1996.
[47] P. Mirowski. Learning to navigate. In 1st International Workshop on Multimodal Understanding and Learning for Embodied Applications, pages 25-25, 2019.
[48] D. Misra, M. Henaff, A. Krishnamurthy, and Langford. Kinematic state abstraction and provably efficient rich-observation reinforeement learning. In International conference on machine learning. pages 6961-6971. PMLR, 2020.
[49] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1rXiv:13/2.5000) 2013.
[50] Mohamed and D. J. Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Advances in neural information processing systems, pages 2125- 2133, 2015.
[51] O. Nachum, S. Gu, H. Lee, and S. Levine. Near-optimal representation learning for hierarchical reinforcement learning. 2018.
[52] A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: arXiv:/807.00 880.007748 2018
[53] D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity driven exploration by self- supervised prediction. In Proceedings of the 34th International Conference on Machine Learning- Vollue 70, pages 1778-7777. JMLR. org, 2017.
[54] M. Phuong. M. Welling, N. Kushman, R. Tomioka, and S. Nowozin. The mutual autoencoder: Controlling information in latent code representations. 2018.
[55] B. Poole, S. Ozair, A. v. d. Oord, A. Alemi, and G. Tucker. On variational bounds of mutual information. arXiv preprint (rXXiv:905.06990 2019.
[56] B. Ravindran and A. G. Barto. Smdp homomorphisms: an algebraic approach to abstraction in semi-markov decision processes. In Proceedings of the 18th international joint conference on Artificial intelligenee, pages 1011-1016, 2003.
[57] M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman. Data-cfficient reinforcement learning with self-predictiv representations.
[58] E. Shelhamer, P. Mahmoudich, M. Argus, and T. Darrell. Loss is its own reward: Self- supervision for reinforcement learning. arXiv preprint 111007766, 2016.
[59] P. Shinners. Pygame. httt:://yggarrs . org/ 2011.
[60] R. Shu, T. Nguyen, Y. Chow, T. Pham, K. Than, M. Ghavamzadeh. S. Ermon, and H. Bui. Predictive coding for locally-linear control. In International Conference on Machine Learning. pages 1. PMLR, 2020.
[61] J. Song and S. Ermon. Understanding the limitations of variational mutual information estima- tors. arXiv preprint (alo:/99000025 2019.
[62] A. Srinivas, M. Laskin, and Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint arXiv:2004.04133 2020.
[63] S. Still and D. Precup. An information-theoon learning. Theory in Biosciences, 131(3):139-143 2012.
approach to curiosity- -riiven reinforcement
[64] A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. Technical report, UC Berkeley, 2020.
[65] C. Sun, F. Baradel, K. Murphy, and Schmid. Contrastive bidirectional transformer for temporal representation learning. arXiv preprint arXiv: 05743, 2019.
[66] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[67] V. Thomas, J. Pondard, F. Bengio, M. Sarfati, Beaudoin, M.-J. Meurs, J. Pineau, D. Precup, and Y. Bengio. Independently controllable factors. arXiv preprint 1708 2017.
12

---

[68] B. Van Roy and S. Dong. Comments on the du-kakade-wang- lower bounds. arXiv preprint 191100191 2019.
[69] V. Veeriah, M. Hessel, Z. Xu, Rajendran, L. Lewis, Oh, H. P. van Hasselt, D. Silver, and
S. Singh. Discovery useful questions as auxiliary tasks. In Advances in Neural Information Processing Systems, pages 9306-9317, 2019.
[70] M. Watter, J. Springenberg. J. Boedecker, and M. Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pages 2746-2754 2015.
[71] T. Yu, G. Shevchuk, D. Sadigh, and C. Finn. Unsupervised visuomotor control through distributional planning networks. Robotics Science and Systems, 2019.
[72] A. Zhan, P. Zhao, L. Pinto, P. Abbeel, and M. Laskin. A framework for efficient robotic
manipulation. arXiv preprint arXiv:2012.07975 2020.
[73] A. Zhang. Satija, and J. Pineau. Decoupling dynamics and reward for transfer learning. arXiv preprint
2018.
[74] M. Zhang, S. Vikram, L. Smith, P. Abbeel, M. J. Johnson, and S. Levine. Solar: Deep structured latent representations for model-based reinforcement learning. arXiv preprint arXiv: 1808, 09105, 2018.
13


Reviews:
[{"review_id": "mSSg9CAD1p1", "reviewer": "Reviewer_UFXP", "summary": "This paper discusses the relation between the sufficiency of representation in RL and the associated info-theoretic quantities. The authors formalize the sufficiency of the representation. The authors further consider three commonly adopted mutual information objectives for obtaining representations, namely, (i) forward information that maximizes $I(Z_{t+1}; Z_t, A_t)$, (ii) state-only information that maximizes $I(Z_{t+1}; Z_t)$, and (iii) inverse information that maximizes $I(A_{t}; Z_{t+k}, Z_t)$. The authors show that (i) maximizing the forward information leads to sufficient representation, and (ii) maximizing neither state-only information nor inverse information does not yield sufficient information. The authors further conduct experiments to verify their theoretical findings empirically.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "Strength:\n\nThe question proposed by this research, namely, what type of information in representations allows for successful RL, is very interesting and important. Successfully answering such a question brings new insight into both the theory and the empirical study of RL.  The authors provide proof and counterexamples to support their theory and validate the theoretical findings with extensive experiments. The usefulness of the forward information is intuitive and consistent with some of the previous studies in Block MDPs [1][2].\n\n\nLimitations and questions:\n\n1. Since all the information-theoretic objectives are typically utilized as auxiliary loss and are rarely trained separately, I wonder if incorporating this information helps RL? As an example, the sufficiency of the forward information is shown only at the maximizer of the forward objective. Is it helpful for features to contain a certain amount of forward information while not maximizing it? Similarly, while the maximizers of other types of information are insufficient, does incorporate them in the representation helps reinforcement learning? The questions may be out of the scope of this research (and are minor points to me), but I believe could strengthen the impact of this work greatly.\n\n2. The notation is a bit confusing to me. In Lemma 2, is $P(Y|X = x)$ the conditional density or a random variable (namely, the same random variable we used to define differential entropy H(X) = \\EE[-log p(X)])? The computation from (12)-(17) is also a bit confusing to me. As a suggestion, it would be nicer if the authors could write every expectation into integrations with density functions properly written. In the proof of Lemma 1, line 565, why does the Markov chain take the form $Z_t \\rightarrow S_t \\rightarrow S_{t+k} \\rightarrow Z_{t+k}$? If $S_{t+k} \\rightarrow Z_{t+k}$, then why do we have $Z_t \\rightarrow S_t$ instead of $S_t \\rightarrow Z_t$?\n\n\nMost of my concerns lay on the correctness of the results (point 2. of Limitations and questions). I will raise my score if the authors could address my concerns properly.\n\n\n[1] Dipendra Misra et al., Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning. (2019)\n[2] Simon S. Du et al., Provably efficient RL with Rich Observations via Latent State Decoding. (2019)", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See Limitations and questions in the Main review. I do not have concerns regarding the societal impact of this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "Pq3OBnlUkKw", "reviewer": "Reviewer_Vgsv", "summary": "This work investigates three mutual information objectives to understand which of these would yield a representation that is sufficient for reinforcement learning. The three being : forward information( learn next state representation given the current state and action); state-only transaction ( learn next state representation given current state only) and inverse information (learn action representation given current state and next state).  Intuitively,  as well as they show, that the later two representations wouldn’t be sufficient for reinforcement learning. ", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "Originality and Significance: \nLearning a sufficient representation of the state based on the given observations is an important issue in reinforcement learning. We usually rely on end-to-end learning using Deep neural networks to represent a policy that feeds upon raw observations and thereby tries to capture an abstract representation of the state within the deep net. This also makes it difficult to learn a policy as the representation learning is only guided by policy learning objectives. This work dives into the notion of having an objective that would give a sufficient representation for policy learning. Also, a representation is referred to as sufficient if an optimal policy based on states could be represented correctly over the learned representation. Experimentally and Theoretically, they show a latent representation learnt over the constraint of doing one-step dynamics i.e. to predict the next state given the current state and action; is the only representation that would give sufficient representation for  RL as compared to the other two representation learning candidates.  \nThey also show the limitations of end-to-end learning by adding additional distractors in the observations which are not relevant for policy. The representation learned for noisy observations under the constraint of forward dynamics is better for policy learning as shown by their experiments. Overall, the significance of this work is to make RL researchers aware of different representation learning constraints and which of these constraints are insufficient.\n\nQuality:  They support their approach via a brief theoretical analysis and experiments over a set of classic environments, which I believe are sufficient to endorse their theoretical claims for the given set of candidate representation objectives. As per my understanding, this work assumes the availability of Markov states and their representation majorly focuses on eliminating irrelevant aspects of the states for the policy. I would appreciate it if the authors also introduce partially observable environments and include temporal representation learning as well as part of candidates for learning representations. This would involve retention (along with elimination) of information across multiple time-steps.\n\nAs per my understanding, they learn the representation first using the pre-generated dataset which supposedly covers the entire state-space. This is a strong assumption and they may want to consider cases in which enough data is not present to learn a good representation for policy learning. In this case, they would have to learn the representation along with the policy. And, the dataset would evolve over the training period. It would be interesting to induce this scenario as part of their experiments. They may have to constrain the representation on both the policy learning objective as well as MI objectives as just doing the MI objective may not be sufficient for a good representation, especially in the early stages of learning.\n\nClarity: Yes, the ideas are conveyed in a clear manner and the given equations could be easily adopted in any representation learning work.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors didn’t necessarily emphasize any of their limitations but did mention partially observable environments as part of future investigation. In my opinion, another limitation of this work may emerge in highly stochastic environments where a representation may not be easily learned.  Also, the computational cost of learning the representation may be more than the end-to-end learning.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": ["I don’t know"], "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "JkzwfKXKjVZ", "reviewer": "Reviewer_jykz", "summary": "This paper analyzes three commonly used information-theoretic objectives for learning state representations. Several claims are made regarding which mutual-information-based objective functions lead to sufficiently powerful learned representations (i.e. representations from which an optimal policy could be learned). The most impactful result seems to be that maximizing the mutual information of state representations without taking into account action information can yield learned representations which are not sufficient for learning an optimal policy. Experimental evidence of the claims is shown for the PyGame \"catcher\".", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "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": "\nOriginality:\nThe analysis performed in this paper (or really, in the supplement), appears to be a new and useful perspective on mutual information (MI) based representation learning, and its effects when keeping it as a separate process from exploration / policy learning. The related work is well cited. \n\nEntirely optional, but if the authors are looking for additional avenues for future work: \nOne related applied work that might be of interest is https://arxiv.org/pdf/1605.09735.pdf. Here, an MI-based intrinsic reward was used to assist in learning coordinated robot motion. The MI-based reward is slightly different than the 3 shown in this work.\nSome theoretical work analyzing this MI-based reward was done, see https://arxiv.org/pdf/1906.05460.pdf, though a sufficiency analysis such as the one shown in this paper is absent. A sufficiency analysis for this more efficiently computed MI-based reward would be interesting for future work.\n\nClarity:\nThe submission is clearly written and well organized. No major comments here (for a minor one, see nit-picks section below). Releasing the code in the supplemental material is useful for reproducibility.\n\nSignificance: \nThe sufficiency analysis here is an interesting contribution to information-based representation learning. It adds evidence to the current knowledge based on which mutual information based objective functions may be useful for reinforcement learning tasks. In the personal opinion of the reviewer,  other researchers and practitioners will likely build upon and extend these ideas.\n\nQuality: \nProposition 2 might be flawed. Here is why:\n\t-- Doing the actual MI computation for the \"aliased\" representation, the MI is ln(2). Since Z is a random variable over 3 states, the maximum MI attainable is ln(3). So saying this is maximizing the MI objective function is quite imprecise.\n\t-- Learning the representation presumes that a fixed uniform policy is used to generate data samples. This need not be the case, and in fact, if a different policy is used, then better representations would be learned by maximizing MI. Therefore, I tend to conclude that the proposition statement is imprecise, since there is no notion of how the fixed initial policy plays a role, and its role is essential.\n\t-- As for the proof, if the policy starts out perfectly uniform, then the \"aliased\" representation is learned, then the policy is trained using that representation, I agree that the true globally optimal policy can not be learned. However, the trial representation deterministically associating z_i with s_i could also be learned by maximizing the MI under this objective function, since it too has the mutual information of ln(2), same as the aliased one. Note that in this case, it does not appear to be possible to learn a representation which truly maximizes MI (i.e. attain the value ln(3)). In that case, the representation is sufficient. So again, something is imprecisely written for this example and proposition.\n\nThe maximum value of MI attainable increases when the size of the state space increases. See section 4.2 of https://arxiv.org/pdf/math-ph/0702002.pdf for a discussion on the maximum of mutual information between two random variables). It seems like this fact would suggest that any amount of aliasing would lower the maximum MI attainable. If the representation's state space isn't fixed ahead of time, then truly maximizing MI may lead to the least aliasing possible, meaning the counter-example's learned representation would likely naturally be the trivial one with no aliasing, though I too have no formal proof for this.\n\nA statement on whether or not MI-maximizing representations can be learned in the first place would make the work more complete.\n\nAll in all, the theoretical results in conjunction with the experimental results seem acceptable, pending some minor revisions such as improving the precision of proposition 2 (and likely, proposition 3, although I did not look as deeply into Prop. 3).\n\n\nOther comments:\nPersonally, it took several readings to clearly figure out where the MI learning objective comes in to the RL problem. Many different setups exist, such as learning the representation entirely before learning the policy, interweaving the tasks (learn policy, learn representation, learn policy, ...), or even treating high MI in the representation as an additional reward in the course of training the agent. \nMy understanding is that, the agent interacts with the environment under an untrained random policy, the state representation is learned by maximizing MI, then the agent is trained using this learned state representation. If this is correct, this should be communicated much clearer.\n\nnit-pick -- The notion of a \"set of reward functions\" was introduced, but then seemingly never used in the main text, or otherwise its usage in the presented definitions could be avoided easily. The propositions only discuss sufficiency in terms of all possible reward functions. As such, it seems superfluous to introduce this and make the definitions epsilon-harder to read.\n\nIn figure 5 and 7 -- \"end-to-end\" is used as a label and its meaning is never introduced anywhere in the text, nor is it intuitive as end-to-end has different meanings in different fields (i.e. ML software development).\n\nFigure 2 -- I have no idea what the x-axis is here. The label seems to indicate \"state representations\", but what does it mean for state representation to be 0, 0.55, 1.1, 1.3?\n\nProposition 2 -- If the claim is that a certain representation maximizes MI, a short calculation would be nice to see. Additionally, indicating what the maximum value of MI is would be nice, since this is a known value.\n\nThe intro example -- The example about traffic lights could be better since red lights are positioned in predictable places (I suspect so fully color blind people can understand them). Hence the color of the light isn't necessarily essential information.\n\nnit-pick notation -- Using \\mathbb{J} for an objective function is strange. Usually this style font indicates a set in mathematics, rather than a plain function.\n\nline 62, spell-checkable typo -- \"whiel\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, it seems of limited foreseeable potential to negatively harm society.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "HUfEysS5K87", "reviewer": "Reviewer_mtq5", "summary": "The paper isolates and studies the role of representations used for control (reinforcement learning), specifically a variety of mutual-information (MI) based objectives for representation learning. A theoretical proof of sufficiency or insufficiency is presented for three types of MI-based objectives, along with specific experiments demonstrating the inadequacy of each insufficient MI objective. Out of the three objectives presented, the “forward information” is the only probably sufficient objective for learning representations necessary for control. The paper serves as a guide for consideration of various MI objectives in future research.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "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": "Overall, the paper presents valuable insights into our understanding of MI-based objectives in learning representations necessary for control. The sufficiency and insufficiency proofs along with the experiments designed to demonstrate inadequacies of each insufficient representation are important results that will inform future research that uses these objectives. However, I am concerned about the limited scope of the paper given that it disentangles representation learning from the control problem, thus disallowing any objectives from learning control-dependent representations. Most image-as-state reinforcement learning research involves the joint optimization of representation learning modules and control modules, making these insights less applicable.\n \nGiven below are more specific strengths and weaknesses.\n \n## Strengths:\n- The proof of sufficiency of the forward information objective and the counter-examples for proof of insufficiency of the state-only and inverse information objectives are novel insights in terms of usage of these representation learning objectives for control.\n- The motivation for learning good representations is made clear throughout the paper through examples of representations needed for autonomous vehicles.\n- The experiments are carefully designed to demonstrate the failure modes of the insufficient MI objectives, both in terms of failure to decode the relevant state information and failure to learn a good policy. The latter experiments on increasing visual complexity further highlight the performance differences from using the various MI objectives.\n\n\n## Weaknesses:\n- As addressed in the paper, the forward information objective has an inherent weakness that it retains more information than what may be needed for solving the task. However, this alone is not a weakness of the paper -- the weakness is in the limited scope of the paper. I feel that the paper could have expanded its scope to consider entangled objectives (joint representation learning + reinforcement learning) with some choice of bottleneck to restrict the information to only contain representations necessary for control. This would have led to representations which can be “sufficient for control” (see my comment below on the misuse of “sufficient” in place of “necessary”) as the title of the paper states.\n\n\n## Other Comments:\n- Minor comment about the use of “sufficient” vs “necessary” throughout the paper. In my view, the MI objectives themselves are referred to as “sufficient” when they are sufficient for learning a good representation (one that doesn’t discard useful information). A good representation then in turn is “necessary” for control i.e. learning a good policy (since no control-relevant information was discarded) but not sufficient (the RL algorithm may fail to find the optimal policy). In some parts of the paper, these distinct concepts seem to be mixed in statements which claim that something is “sufficient for control”. Since this paper disentangles representation learning from control (policy learning) and focuses only on the former, it seems incorrect to make any claims about sufficiency for control since no theoretical proof is provided that guarantees a solution of the control problem given a perfect representation. Instead, sufficiency of the MI objective for representation learning is the accurate statement.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have addressed the limitations of their work accurately, from scope of the work, assumptions made and the limitations of the sufficient forward information MI objective.\n \nThe potential negative societal impacts have been addressed adequately.\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper is very nicely written and the reviewers all agreed that there is merit to the work. The primary concerns including the following.
1. There are limitations in the scope, particularly in that the sufficiency of the representation does not discuss other important properties like compactness. The theory (Proposition 1) seems to assume that the identity transformation is analyzed, and it is not too surprising that this representation is sufficient. The theory is clearly presented, but considering a more practical setting with some bottleneck that removes irrelevant information would be beneficial (even if it is assumed that only irrelevant information is removed).
2. The choice of these three MI objectives and then the clear superiority of the forward objective in the experiments speaks to (a) limitations in the objectives considered and (b) some insufficiency in the experimental setup to make these relatively strong claims that perfectly match the theory. The experiments were run with only 5 seeds, with default hyperparameters and some insufficiently explicit hyperparameter choices (one approach uses a bottleneck, another does not). The beauty of these simple environments is that you can more exhaustively run experiments (more runs, more hyperparameter investigation) to provide sufficient evidence to make these strong claims. These default hyperparameters might be good for one representation and not for the other; results could be very different if these were tuned. As it stands, the experiments highlight results for SAC with default parameters plus these representations with somewhat inconsistent hyperparameter choices, which his relatively weak evidence.  I do not particularly doubt the results, though, because for the three objectives investigated, it is intuitive that forward MI objective is effective, and that the other two remove important information. This speaks to the first insufficiency about the limited scope in objectives. 

Despite these concerns, the clarity of writing and the importance of the topic make this a useful paper for the community. If there is room in the program, this paper could be accepted. This paper, though, could be improved and be a much stronger contribution; I would highly recommend that the authors seriously consider improving the strength of their empirical claims. 

The reviewers had concerns also about the problem setting, where the representations are learned beforehand from a batch of data. I do not share this concern, and believe that studying this problem setting is of interest to better understand learning representations that then are used to accelerate learning downstream. The authors may likely continue to get push back on this setting, but at least know you have one supporter here.

Minor point: though not pointed out by the reviewers, the proof in Lemma 2 has a few steps that are needlessly complex. For example, introducing the tower property twice is not necessary; you can directly jump to p(Y|X) p(X, Z) = P(Y| X) P(Z | X) P(X) for Eq 16, rather than doing the slightly odd thing of writing P(X, Z|X) = P(Z | X). It also seems like this result is quite general: three RVs with pairwise equal MI and conditional independence implies that we have this equality in probability. Of course, as pointed out by a reviewer, this proof as written has several pretty critical typos.
Decision: Accept (Poster)