Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper discusses the sufficiency of several objectives for representation learning. It shows that learning representation by maximizing forward information is sufficient for RL under their definition and other two objectives are not. It also uses experiment to show that forward-information maximization is useful in RL in the sense that using it can train an agent that is at least as good as an end-to-end learner and can be better in some cases.
main_review: Using sufficiency to analyze MI objectives is novel, and can enlighten future research on theory of representation learning and algorithm design. The paper is written clearly. Their theoretical results and experiments clearly show the value of using forward information for representation learning. Although I am not pretty sure whether I have missed something, to my knowledge, their proof, including the proof of Lemma 2, is now corrected.
limitations_and_societal_impact: The authors have addressed the limitations of their paper clearly.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper investigates three popular mutual information (MI) based objectives for learning reinforcement learning (RL) representations through the lens of sufficiency. The authors distinguish between two types of sufficiency: (i) a state representation is $\pi^*$-sufficient if it can be used to represent the optimal policy; and (ii) it is $Q^*$-sufficient if it can additionally guarantee learnability (i.e., convergence of Q-learning to the optimal policy with that representation). The paper argues that the latter is the right type of sufficiency, and show that only one out of the three popular MI-based objectives satisfies is sufficient. For the other two, they provide counterexamples illustrating simple conditions under which these objectives can fail. The authors additionally verify their claims via deep RL experiments on variations of a video game.
main_review: There are several strong points about this work.
- It investigates a very important question. Various MI-based objectives for state representation have been previously proposed in the literature. A straightforward way to compare the different objectives is via empirical comparison on various RL environments. In addition to that though, we also need formalisms that compare the objectives along rigorous criteria. The authors introduce such a formalism based on the concept of sufficiency. This is a concrete way to compare the resulting representations, with a strong intuitive basis: at the very least, we want representations that can guarantee learnability.
- The paper contains several original elements, including the sufficiency proof for the forward information objective, and the illuminating counterexamples for the other two objectives.
- The experiments are smartly designed to make the insufficient objectives fail by 
- The paper is clearly written, and simple to follow. I believe that the proofs are technically sound, except one detail in Lemma 1 that I am not so clear about (see below).
- The related work is covered well. Furthermore, the discussion section contains many interesting points.

My concerns are as follows.
- The authors do an excellent job in highlighting when the state only MI and inverse MI can fail under the lens of sufficiency. Their experiments provide further evidence for that. However, it is far from clear whether the notion of sufficiency is really the right formalism to compare different objectives in practice. The experiments are rather artificial and specifically designed to showcase specific weaknesses of the two objectives without sufficiency. I think a reasonable way to understand whether sufficiency is really important in practice would be to compare empirical performance of the various objectives in many benchmarked problems. If the two objectives without the sufficiency property suffer from poor performance often, this would indicate that the lack of sufficiency makes a big impact in practice. If not, this could suggest that in practice the situation is more complex.
- Related to the point above, it almost appears after reading this work that the forward objective is the clear winner. But is this always the case though in practical settings? One obvious problem with the forward objective is that it does not consider information about the reward, only about the state and action. Furthermore, it lacks a notion of task relevance, as the authors point out. Conversely, what would be the strengths of the state only and inverse objectives, compared to the forward objective. I am getting the feeling that the comparison is kind of one-sided.
- As the authors point out, the notion of sufficiency that is central in their paper has in fact been explored in several past works on state abstraction (e.g., Li, Walsh and Littman 2006). Even though it is true that this work provides a formal framework for comparing the various objectives, the central ideas behind this have already been laid out in previous works.

Additional questions to authors.
- What happens when we combine the various objectives? We could for example consider some weighted combination of the different objectives. One might wonder whether there is any benefit when learning representations that take into account multiple MI objectives. Would that help performance?
- It makes sense why in general we want a uniform random policy to collect data to learn the state representation: this will allow us to cover the entire state space well. I was wondering though: for the forward objective would it help if we collect the data to learn the state representation via standard RL (i.e., via an evolving policy)? If the RL policy only visits a small part of the state space, then this could allow us to learn even smaller representations Z_t, as we would not need to be equally predictive of all predictable elements in all states but only in the states that matter in the RL policy. This may also be much easier than collecting data with a random uniform policy, especially in very high-D spaces.
- In Lemma 1, the authors claim that Z_t->S_t->S_{t+k}->Z_{t+k} form a Markov chain. In the diagram, we see that the dependency is S_t->Z_t. But in the above claim the direction is reversed. Is this valid, and if yes, why? I was thinking that the direction of the arrow could make a difference. For example, assume Y<-X->Z, and consider the trivial case when Y=Z (and both depend on X). In that case, it can be that I(Y;Z)>I(Y;X)=I(Z;X). But if we reverse the arrow, we get the Markov chain Y->X->Z. In that case, the data processing inequality would imply that I(Y;X)>=I(Y;Z), which is a contradiction. This may mean that we cannot always reverse the arrow. Perhaps I am missing something? I believe the claim in Lemma 1 is still correct, but I was wondering about that part of the proof.
-- The forward objective is trivially maximized if Z_t=S_t, as the authors point out. It might make sense to discuss this aspect in more detail, as well as possible ways to regularize Z_t to ensure it will be the smallest possible and not end up equal to S_t. Similar arguments are true for the other two objectives. The MI objectives are not complete, as in reality there have to be additional constraints on Z_t. Maybe this is something the authors would want to cover in the main text.
limitations_and_societal_impact: The author cover the societal impact by providing citations.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper analyzes three mutual information-based objectives in representation learning for RL from the perspective of sufficiency. The main contribution of this work is to theoretically prove that the forward information metric exhibits, while two other metrics, state-only transition information and inverse information, are not sufficient. Giving numerical examples demonstrating the impact of insufficiency on the RL performance is also a meaningful result presented in this work. 
main_review: -  ​The presented theoretical analysis seems valid and correct. While the sufficiency/insufficiency results themselves are, in some sense, intuitively straightforward, the theoretical proof for the sufficiency of forward information and the counterexamples for the other two metrics are very interesting. 
-This paper is generally well-written, delivering the main message clearly. However, the reviewer suggests to include more details of the proof of proposition 1 in the main text rather than appendix. 
- Given that many recent RL schemes use various mutual information metric, the theoretical findings in this paper would provide some insight on how these work would be improved and enhanced. 

limitations_and_societal_impact: - While the relation between the sufficiency and the learning performance is demonstrated in the presented example, how much generalizable this observation is not clearly explained. It would be good, if the authors would be able to give more systematic explanation on what classes of RL problems the notion of sufficiency makes difference. 
- In the example with distraction, it is not clearly explained why the end-to-end RL does not perform well. It may not be the case that the end-to-end approach has a fundamental limitation in representation, as then the first example would look different. That said, the limited performance of the end-to-end scheme is due to rather practical/numerical issues. This may be okay, but this could bring some question if such a practical issue would affect the other representations. In other words, how could the other reasons be discarded that may affect the performance degradation of insufficient mutual information metrics?
- The result of this paper reads like we should consider the forward information metric. Would there be any disadvantages and/or costs of this metric compared to the other two? Is the takeaway that we should use the forward information, or we need to be more careful about the choice of metric? Such a discussion is needed in this paper.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper discusses on a theoretical level what MI-related quantities are well-suited for representation learning in the context of reinforcement learning. The paper argues that for reinforcement learning, forward information is the only one that is well suited for learning representations by a certain sufficiency definition, and state-only transition information and inverse information do not satisfy the sufficiency property; thus only the forward information is a well-suited principle for representation learning in reinforcement learning. Empirical results on simple offline data suggest the usefulness of the idea. 
main_review: **Update**: I have read the rebuttal. The authors clarified my confusion regarding the correctness of the proofs. I am willing to raise my score to a weak accept. I believe it is crucial to clarify the scope of the statements that are made in the paper, that is "reward depends only on the state" (otherwise the needed representations can be quite different). In retrospect, the choice of the word "control" in the title (and the graphical model in the appendix) might have implied this, but other notations (such as $R(s, a)$) could lead to some confusion. 


-----

**Strengths**

I think this paper is an interesting read, and that it is justified to consider a very simplified view over this setting (mutual information and be well approximated from samples, exploration is a non-issue, etc.). If the theory is true, it can be used to discuss which MI-based objective is adequate for representation learning in RL. The distractor experiment is an interesting setting where there is an incentive to drop this "useless" information (since it cannot be predicted better than random, there is not much that the policy can do about it).

**Weaknesses**: Unfortunately, I think the proof of Proposition 1 is false as it is based on an not always true assumption from Lemma 2.

To make my argument, think about the case where the optimal mutual information in J_fwd is simply zero. For example, each episode is just a sequence of independent contextual bandit problems. So S_{t+1} is independent of S_{t} and A_{t} (which only determines R_t in contextual bandits), mutual information should be zero. Does it seem that even J_fwd would not be incentivized to learn a meaningful representation at all? To have a good representation, at least it seems that the reward should be taken into account?

Let us say in this case, we use the same representation for every state (giving us a mutual information of zero for J_fwd, and that is fine because each timestep is independent anyways now), and it seems that we should obviously not expect Q(s_1, a) = Q(s_2, a) to work for every pair of states? Clearly, it is not Q-sufficient?

This makes me wonder if there is something wrong with the proposition, and I believe is because you applied Lemma 2 where the assumption for Lemma 2 might not hold. In the proof, you use X = [S_t, A_t], Y = [R_t], Z = [Z_t]. To use Lemma 2, you would have to assume I(X; Y) = I(Y; Z), which is true if Z_t = (S_t, A_t), but clearly not true in the above example (where I(Z_t; R_t) = 0, and I(R_t; S_t, A_t) is not zero). So everything falls apart from here.

**Experiments**
In the experiments, the mutual information is estimated by InfoNCE, which is known to have a high bias when the batch size is small [Poole et al., ICML 2019]. How accurate is your estimated mutual information? I don't think I see enough details in the appendix to gauge this. This could be an important factor, because it would be natural to believe that the forward mutual information would be quite high in deterministic dynamics, where $Z_{t+1}$ is known once $Z_{t}$ and $A_t$ are known. I am curious about this in the catcher environment.

Is there any insight as to why J_state works poorly for catcher-grip despite working decently for catcher?



**Comments**

Statement of Lemma 2: the way it is written now makes it very easy to misunderstand as a "type-check" error. I think a more appropriate quantity on the rhs should be $E_{p(Z | X = x)}[p(Y | Z)]$, then one in Eq 11? 



The type of Q functions are not consistent, sometimes it is $(s, a)$, others it is $(a, s)$. This needs to be fixed.
limitations_and_societal_impact: Yes.
review_title: null

Meta-review: All the reviewers found the paper worthy of acceptance at NeurIPS and I am happy to accept it. I encourage the authors to include in the final version the clarifications that they added in the rebuttal.
Decision: Accept (Poster)