Which Mutual-Information Representation Learning Objectives are Sufficient for Control?

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. 


Reviews and author responses:

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.  Strength:

The 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].


Limitations and questions:

1. 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.

2. 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$?


Most 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.


[1] Dipendra Misra et al., Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning. (2019)
[2] Simon S. Du et al., Provably efficient RL with Rich Observations via Latent State Decoding. (2019) See Limitations and questions in the Main review. I do not have concerns regarding the societal impact of this work.

 Dear authors,

Thanks for the detailed explanation on the proof of lemmas. The proofs now make sense to me. I do not have additional concerns.

 Dear reviewer,

Please let us know if our response has adequately addressed your concerns, and if there are any further clarifications we can provide! Thanks to your review, we will update the proof of Lemma 2 to correct typos and more clearly explain the derivation. We will also edit the text to better motivate isolating the study of representation learning from other aspects of RL.

 Dear reviewer,

Please let us know if our response adequately addresses your concerns, and if you have any further questions! Thanks to your review, we will edit the text of the paper to motivate our study of Markovian environments, discuss further the assumption of full state support for the dataset, and extend the discussion of the limitations of our analysis.

 Thank you for your review! We believe the main theoretical result (Proposition 1) is correct, and we detail our reasoning below. Please let us know if this fully addresses your primary concern about the paper.

“The notation of Lemma 2 is confusing”: Indeed $P(Y|X=x)$ is the conditional density, not a random variable. We apologize for the typo in the statement of Lemma 2. The RHS of the statement should be $E_{p(Z | X=x)} [p(Y|Z)]$. We also made a typo in the Markov chain, it should indeed be $Z_t \leftarrow S_t$. Note that the DPI holds because it relies on conditional independence, which is the same regardless of the direction of this arrow in this case, since the resulting conditional independencies are equivalent. We will make the proof of Lemma 2 more clear by writing out the expectations with integrated density functions. The idea of this section is to show the KL-divergence is 0 by showing that it is upper-bounded by 0.  Here we provide a line-by-line explanation:

Eq. 13: An assumption of the lemma is that $I(Y ;Z) = I(Y ; X)$. We can write this as $H(Y) - H(Y|Z) = H(Y) - H(Y|X)$, therefore $H(Y|Z) = H(Y|X)$.

Eq. 14: $\sum_{Y, Z} p(Y, Z) \log p(Y|Z) = \sum_{Y, X} p(Y, X) \log p(Y|X)$. This is the definition of conditional entropy.

Eq. 15 and 16: Here we use the Tower property to rewrite both sides. 

We start with the RHS. By the Tower property, $\sum_{Y, X} p(Y, X) \log p(Y|X) = \sum_Z p(Z) \sum_{Y, X} p(Y, X | Z) \log p(Y|X)$. Because $Y \perp Z | X$, we can write this as $\sum_{X, Z} p(X, Z) \sum_Y p(Y | X) \log p(Y|X)$. Now we use the Tower property again: $\sum_{X} p(X) \sum_{X, Z} p(X, Z | X) \sum_Y p(Y | X) \log p(Y|X)$, which we can rewrite as $\sum_{X} p(X) \sum_Z p(Z | X) \sum_Y p(Y | X) \log p(Y|X)$.

Now we turn to the LHS. $\sum_{Y, Z} p(Y, Z) \log p(Y|Z) = \sum_X p(X) \sum_{Y, Z} p(Y, Z | X) \log p(Y|Z)$ by Tower property. By conditional independence, this is equivalent to $\sum_X p(X) \sum_{Y, Z} p(Y| X) p(Z|X)  \log p(Y|Z)$.

Subtracting the RHS from the LHS, we have: $\sum_X p(X) \sum_Y p(Y|X) \[\log p(Y|X) - \sum_Z p(Z | X) \log p(Y|Z)\] = 0$. This is the final line of Eq. 16.


“Using representation learning objectives as an auxiliary loss for RL”: From a theoretical perspective, fully minimizing the RL loss in addition to a representation learning loss will always result in a sufficient state representation, since the RL objective will preserve all the information needed to obtain the optimal policy. We consider the case where a state representation is learned and then fixed prior to performing RL. This setting allows us to study representation learning isolated from other aspects of RL, and is also a practical paradigm in scenarios where data collection is dangerous or expensive (Kipf et al. 2019, Zhan et al. 2020). On the empirical side, prior work has compared adding different representation learning objectives similar to those considered in our work as auxiliary objectives to RL (e.g., Stooke et al. 2020, Shelhamer et al. 2016), but our focus is not on auxiliary objectives, but rather on representation learning.

“Are representations that don’t fully maximize the objective still useful?”: Our primary contribution is the theoretical sufficiency analysis of MI-based objectives. We agree it would be interesting future work to try to relate the sub-optimality of a solution for representation learning to the sub-optimality of the resulting RL policy.


 Thank you for your thoughtful review of our paper and the suggestions for future work!

“Proposition 2 clarification”: 
- “the maximum MI attainable is ln(3)”: The maximum attainable MI for the objective depends on the MDP dynamics, not only the size of the state space. Since the MI of any representation is less than or equal to the MI of the identity representation, we use “maximizing the objective” to mean achieving the value achieved by the identity representation. We will make this more clear.
- “Fixed uniform policy”: We use a fixed uniform policy in this counterexample to demonstrate insufficiency. Throughout the paper, we assume a fixed policy distribution to remove the confounding factor of exploration from the analysis.
- “The identity representation could be learned by maximizing $J_{state}$”: This is correct; in this example, there are two representations that maximize $J_{state}$, the identity representation and the representation that aliases $s_0$ and $s_3$. The point is that while the former is sufficient, the latter is *not* sufficient, but since both have the same value of $J_{state}$, they would be indistinguishable to an algorithm maximizing $J_{state}$. It is not to say that sufficient representations cannot be learned, it is simply to say that sufficient representations *need not* be learned. 

“any amount of aliasing would lower the maximum MI attainable”: Aliasing may not lower the MI, because MI can be written as the sum of two entropy terms that can balance each other. $I(X; Y) = H(X) - H(X | Y)$. While $H(X)$ increases with the state space of $X$, $H(X | Y)$ may also increase. See our calculation below of the $J_{state}$ objective for two different state representations to see an example of this.

“Where the MI objective learning comes into the RL problem”: You are correct that the setup we consider consists of first learning a representation by maximizing the MI objective on a dataset, and then performing RL with that learned representation. We do this in our experiments to mirror the setting in which we derive our theoretical results, where we separate representation learning from exploration to simplify the analysis. Additionally, the practice of performing representation learning on a fixed offline dataset is a common and practical one (Zhan et al. 2020, Stooke et al. 2020, Kipf et al. 2019), particularly in applications where data collection is expensive or dangerous (e.g., robotics, health care). It is very difficult to provide guarantees on states and actions that are out of support, so we assume complete coverage. 

“End-to-end label in Figures 5 and 7”: We will add a clear definition of this algorithm to the text of the paper. We use the term to refer to performing RL from the raw image inputs without any additional representation learning. We use the same convolutional encoder for this RL agent and for the representation learning algorithms.

“Figure 2 x-axis”: The x-axis is the value of I(S; Z) for the representation. This gives a notion of how much information is retained in the representation - the identity representation is the furthest right.

“Prop. 2 example MI calculation”: As an example, we calculate the value of $J_{state}$ manually for both the identity representation and the one that aliases $s_0$ and $s_3$ (these correspond to the representations plotted with diamond and star in the right side of Figure 2). We will add this calculation to the appendix of the paper. 
We can expand MI in terms of entropies: $J_{state} = I(Z’; Z) = H(Z’) - H(Z’ | Z)$.

For the identify representation, Z = S: The steady-state occupancy distribution is $p(S) = \frac{1}{4}$, so $H(S) = 4 * \sum \frac{1}{4} \log(4) = \log(4) = 2 \log(2)$. We assume a uniform policy distribution, so $p(S’|S) = \frac{1}{2}$, and therefore $H(S’ | S) = - 8 * \frac{1}{2} * \frac{1}{4} * \log(\frac{1}{2}) = \log(2)$. Subtracting the two values, we get $J_{state} = \log(2)$ for the identity representation.

For the representation that aliases $s_0$ and $s_3$: Let’s call the aliased state $z_0$. The steady-state occupancy is $p(Z=s_1) = p(Z=s_2) = \frac{1}{4}$ and $p(Z=z_0) = \frac{1}{2}$. Therefore $H(Z) = \frac{3}{2}\log(2)$. The conditional distributions are: $p(Z=s_2|z_0) = p(Z=s_1|z_0) = \frac{1}{2}$, while $p(Z=z_0| s_1) = p(Z=z_0|s_2) = 1$. Therefore $H(Z’|Z) = -2 * (1) * \frac{1}{4} \log(1) - 2 * \frac{1}{2} * \frac{1}{2} \log(\frac{1}{2}) = \frac{1}{2} \log(2)$. Subtracting the second term from the first, we see that for this representation as well, $J_{state} = \log(2)$.
Since both of these representations have the same value of $J_{state}$, an algorithm maximizing $J_{state}$ could not distinguish between them, potentially yielding the insufficient representation that aliases $s_0$ and $s_3$.

“Intro example with stoplight”: Our point with that example is that the information the stoplight communicates could be discarded - as you point out, since the information can be encoded by spatial position as well as by color, both these channels could be discarded by the insufficient representation.


 Thank you for your thoughtful review of our paper! We respond to your concerns below.

“Expand the scope to consider entangled objectives”: From a theoretical perspective, fully minimizing the RL loss in addition to a representation learning loss will always result in a sufficient state representation, since the RL objective will preserve all the information needed to obtain the optimal policy. We consider the case where a state representation is learned and then fixed prior to performing RL. This setting allows us to study representation learning isolated from other aspects of RL, and is also a practical paradigm in scenarios where data collection is dangerous or expensive (Kipf et al. 2019, Zhan et al. 2020). On the empirical side, prior work has compared adding different representation learning objectives similar to those considered in our work as auxiliary objectives to RL (e.g., Stooke et al. 2020, Shelhamer et al. 2016); however, our focus is not on auxiliary objectives, but rather on representation learning. We do make use of an information bottleneck of the form of minimizing $I(Z; S)$ in the experiments in order to approximately select the smallest representation that maximizes the objective. Note that adding this information bottleneck does not change the theoretical sufficiency results. In the case of insufficient objectives $J_{inv}$ and $J_{state}$, the insufficient representations that maximize these objectives discard needed information and thus have a lower $I(Z; S)$. The bottleneck thus acts to select the insufficient representations from all the representations that maximize the objective.

“‘Sufficient’ versus ‘necessary’ terminology”: We define a representation to be sufficient when it contains all the information necessary to learn and represent the optimal policy. We use the word “sufficient” because the representation may contain additional information that is not needed, as indeed is likely to be the case in practice with $J_{fwd}$ in many cases. While recent work has shown that a poor representation may result in slow convergence (e.g., Du et al. 2019), RL is still guaranteed to converge if the representation is sufficient (Li et al. 2006, Theorem 4). Please see lines 137-145 of our paper for a discussion of this. We will further expand the discussion of this terminology to make this clear in the final paper.


 Thank you for your review. In the response below, we respond to your concerns by motivating our reasoning for studying representation learning in Markovian environments, explaining the assumption of full state support in the dataset used for representation learning, and discussing the limitations of our analysis. Please let us know if there are any other concerns we can address!

“Partially observed environments”: We agree that analyzing partially observed environments is an interesting future direction; however, the tools required to do so would likely be quite different from those used in our paper. Representation learning in Markovian environments is a long-studied problem in the state aggregation literature that we draw on for our analysis (Bean et al. 1987, Givan et al. 2003, Li et al. 2006). It is a practical problem as many realistic environments can be approximated as Markovian. For example, mild partial observability is commonly handled by concatenating together short windows of observations to form a single observation (frame-stacking). In this situation, inferring a more compact and less redundant representation space in which to perform RL can greatly improve the tractability of learning. 

“Full state support assumption for the dataset”: While the assumption that the dataset has full support over the state space may not hold in practice due to the limited size of datasets and difficulty of exploration, this assumption is standard in RL theory (Sutton & Barto 1998) and can be achieved, for example, by an epsilon-greedy policy. Without this assumption, we cannot say anything at all about the representation of states outside the support, and so it is not possible to guarantee sufficiency for any representation learning method. We carry out the experiments in the same setting for consistency. Note that our goal is not to propose a new method, but rather to analyze whether existing MI-based representation learning methods are sufficient -- any method that is not sufficient even under full state support will also not be sufficient without full state support.

“Limitations of the analysis”: We discuss a few limitations of our work in Section 7, including the question of whether in the general case $J_{fwd}$ yields representations that are minimal (i.e., most compressed) in addition to sufficient, characterizing specific classes of MDPs in which $J_{inv}$ and $J_{state}$ can be proven sufficient, and extending the analysis to partial observability. Additionally, we make several assumptions for our theoretical analysis that may be relaxed in practice -- that the MI objectives can be perfectly maximized, the reward is a function solely of the state, and the dataset used for representation learning has full coverage of the state space. We will extend the discussion of limitations in the text.


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.   Originality and Significance: 
Learning 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.  
They 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.

Quality:  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.

As 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.

Clarity: Yes, the ideas are conveyed in a clear manner and the given equations could be easily adopted in any representation learning work.
 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.


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.  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.
 
Given below are more specific strengths and weaknesses.
 
## Strengths:
- 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.
- The motivation for learning good representations is made clear throughout the paper through examples of representations needed for autonomous vehicles.
- 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.


## Weaknesses:
- 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.


## Other Comments:
- 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.
 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.
 
The potential negative societal impacts have been addressed adequately.


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".  
Originality:
The 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. 

Entirely optional, but if the authors are looking for additional avenues for future work: 
One 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.
Some 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.

Clarity:
The 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.

Significance: 
The 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.

Quality: 
Proposition 2 might be flawed. Here is why:
	-- 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.
	-- 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.
	-- 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.

The 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.

A statement on whether or not MI-maximizing representations can be learned in the first place would make the work more complete.

All 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).


Other comments:
Personally, 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. 
My 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.

nit-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.

In 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).

Figure 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?

Proposition 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.

The 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.

nit-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.

line 62, spell-checkable typo -- "whiel" Yes, it seems of limited foreseeable potential to negatively harm society.

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