Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 6: Marginally above acceptance threshold
review_content: **Summary**: The paper discusses three mutual information (MI) objectives for representation learning in RL, referred to as forward, state, and inverse. The forward MI objective models latent dependencies given the action. The state MI objective models latent dependencies alone. And the inverse MI objective models dependencies between actions and future states (empowerment). The paper shows that of these three common objectives, only the forward objective is sufficient for learning the optimal policy / value function. This is demonstrated using simple examples and experiments on a simple game environment.

**Strong Points**: This paper attempts to provide an overarching view of the many previous works that have used MI objectives for representation learning in RL. Multiple approaches are compared under one generic formulation. This can be helpful in connecting disparate research areas, e.g., connecting contrastive representation learning with empowerment and other task-agnostic policy learning objectives. Papers that present more general interpretations, such as this one, help to build consensus in the research community and standardize techniques, which can be more useful that proposing modifications of existing techniques.

Along similar lines, to the best of my knowledge, this paper presents a somewhat novel idea: analyzing MI objectives in terms of whether they are sufficient for (optimally) performing downstream tasks. Although representation learning and policy/value learning are not always performed separately, many works do not theoretically interrogate whether their representation learning objective is sufficient to perform the task. Having a clearer theoretical grounding can be helpful for deciding which representation learning objectives are worth pursuing.

Overall, the paper is well written. Key mathematical concepts are defined clearly in sections 3 and 4. The descriptions and definitions are clear and concise.

Although the experiments are rather limited, they do help to isolate key aspects of the analysis. In particular, I found the regression of components of the environment (fruit error and gripper error) to be compelling and helpful. 

**Weak Points**: I found the technical formulation somewhat unclear. The paper presents representations as stochastic mappings from states (s) to latent variables (z). Sufficiency is defined as having the same optimal policies / value functions when the representations are the same. However, it is unclear, in practice, how such representations are intended to be used for value/policy learning. From the proof sketch in proposition 1, it seems as though the value is estimated by integrating over Z, however, this may not be feasible in practice, and we may need to use samples. Regarding this point, it is also unclear what is performed in the actual experiments of the paper.

It’s unclear whether analyzing the sufficiency of representations for downstream tasks is impactful for future work. Sufficiency (see definitions in section 3) describes whether the learned representation removes any task-relevant information. However, there are infinitely many representations that are sufficient. While sufficiency is certainly important for learning representations, it is only half of the consideration. In practice, one would want a *minimal* sufficient representation that is amenable for learning. Likewise, there are cases where representation learning and task-based learning occur together, in which case, task-based learning may be able to overcome the insufficiency of representation learning.

The paper is almost entirely lacking in experiment details. The descriptions for the examples in figures 2 and 3 could be improved in clarity, along with details on the estimation of each of the MI terms. Likewise, beyond basic descriptions of the tasks in the experiments section, very few details are present. While I understand that the focus of this paper is more theoretical in nature, such details are essential for reproducible results and ensuring technical rigor.

This paper may not be entirely representative of previous work. Mutual information objectives of differing types are applied in various contexts, in many cases trained during data collection. This paper explores a limited setting, in which these objectives are trained on data collected from a uniform policy and (I assume) the policy / value function is estimated only from the representation. While the authors claim that this enables a fair comparison between the objectives, it also somewhat limits the scope/impact of the paper, as it is less realistic. Further, experiments are performed on fairly limited, fully-observed environments. While I understand that these simple environments are meant to capture the essential differences in these objectives, focusing only on these relatively “toy” environments could limit the impact of this paper. For instance, in the two experiments in the experiments section, representation learning provides fairly marginal improvements in performance over just directly training end-to-end. Even just sticking with these environments, the paper could be improved with further analyses on the types of representations that are learned in each case.

**Accept / Reject**: Given the relative lack of details and the limited experimental investigation, I would lean slightly toward rejection. While I agree with the direction of this paper, I feel that these aspects would need to be improved for the paper to have substantial impact on the rest of the research community. With a more in-depth discussion of the experiment details and perhaps further analysis of what is and is not captured by each objective, this paper could reach the bar for acceptance. Experiments and analysis with previously published approaches would help to further improve the paper.

**Questions**:

Could you please provide more details on the training scheme in the experiments section. How is the representation used in practice for downstream tasks?

Where do other unsupervised representation learning schemes (e.g. VAEs and normalizing flows) fit within this framework?

For J_state, it is stated that previous values of Z can be ignored due to the Markov state. However, if the mapping from S to Z discards dynamics information, Z will not be Markovian. Is this formulation assumption still valid?

**Additional Feedback**:

Related Work:  
Equations 2 and 4 are referenced far before they are defined. I generally avoid referencing equations this far in advance.  
I would consider citing Mohamed & Rezende, 2015.

Representation Learning in RL:  
Missing the discount factor in the preliminaries section.  

Sufficiency Analysis:  
Figure 2: J_inv —> J_state

Experiments:  
Figure 6: should be J_state in the table.
review_title: Interesting and useful direction, but perhaps limited in clarity and applicability

final_score: rating: 7: Good paper, accept
review_content: This paper studies which commonly-used mutual information objectives for learning state representations are sufficient for reinforcement learning. In particular, they provide counterexamples to show that state-only and inverse MI objectives are not Q*-sufficient, while proving that forward MI is Q*-sufficient. They validate their findings empirically with experiments in a simple RL domain.

There has been a lot of work recently in reinforcement learning that uses mutual information objectives resulting in performance gains, so it’s very fascinating to see a finding that these objectives may be theoretically insufficient, despite their empirical success. The counterexamples shown are simple and the authors do a good job of explaining the intuition. I think this paper will be of great interest to the ICLR community. 

One question: while J_state and J_inv are not sufficient, are there conditions in which they can be? If these conditions are limited to just the reward, could this somehow give insight on how to design better reward functions?

Minor: there’s a reference in the last paragraph on page 6 to Figure 5 which I think should be to Figure 4.

review_title: Conditions in which mutual information objectives are sufficient for reinforcement learning

final_score: rating: 5: Marginally below acceptance threshold
review_content: Summary of the work: This work studies which mutual-information representation learning objectives (1. forward information, 2. state-only transition information, 3. inverse information) are sufficient for control in terms of representing the optimal policy, in the context of reinforcement learning (RL). As a result, they find a representation that maximizes 1 is sufficient for optimal control under any reward function, but 2 and 3 fails to provide that guarantee in some MDP cases. They provide both proof and interesting counter examples to justify the findings. Besides, they conduct some empirical studies on a video game (i.e. Catcher) and show that the sufficiency of a representation can have a substantial impact on the performance of an RL agent that uses that representation. 

l like the idea of trying to understand the recent popular mutual information objectives in RL. To the best of my knowledge, Q^*-sufficiency analysis for mutual information objectives is novel. The counterexamples in sufficiency analysis are interesting. The paper is well written.

However, l still have the following concerns: 

Empirically mutual information objectives often play the role of auxiliary losses to improve the sample efficiency of RL. Thus, it is very useful for us to theoretically understand which objective is better in terms of improving sample efficiency. The property of sufficiency (i.e. the ability to represent the optimal policy) is important. However, only this property may not strong enough, because there may exist a trade-off between minimizing information loss and maximizing state space reduction (Li et al. 2006). Can we add one more perspective, such as whose representation is finer (like Definition 2 in Li et al. 2006) to better understand these MI objectives, if possible? Intuitively, a coarser representation results in a larger reduction in the state space, which in turn translates into the efficiency of solving the problem.     

Regarding the experiment results, the authors give some intuitive descriptions to show that state-only transition objective and inverse objective may be insufficient, but forward objective works in the catcher game. To enhance its solidness, l strongly suggest that we may conduct an experiment on the predictability of optimal Q-function by the representations trained by various MI objectives. For example, we can spend a long time to train a good enough policy and treat it as an optimal policy.  


Regarding the modeling of representation, the representation is modeled as a random variable in this paper for analysis. However, in practice, the representation is always built upon neural networks and thus deterministic. Therefore, l am a little curious if the derived conclusion will still hold for deterministic cases.  

Regarding the proof of Proposition 4, l am sorry that l do not fully understand the derivation from Q(s,a) to Q^* (s,a). Can the authors provide more details on that? 

From the proof of proposition 4 in appendix,  I(Z_t,A_t;Z_(t+k)) seems to be maximized for ∀k>0,t>0. l suggest we make that clearer in the definition of those MI objectives (e.g. Equation 2), since some practical algorithms are based on the fixed k, not all k.  

review_title: Official Blind Review #2

final_score: rating: 5: Marginally below acceptance threshold
review_content: Summary

This paper studies 3 mutual information (MI) optimization objectives for learning a latent representation Z in a sequential decision-making task of the MDP style:

forward information: MI between Z at time t+k, and Z at time t concatenated with the action at time t,

state-only transition information: MI between Z at time t+k, and Z at time t,

inverse information: MI between the action at time t, and Z at time t+k, conditioned on Z at time t.

All of the 3 objectives above are theoretically analysed with the result that only forward information can lead to a sufficient representation for reinforcement learning (based on the concept of Q*-sufficiency). This is experimentally confirmed in 2 simple games: catcher and gripper. In these 2 games, a latent representation based on forward information for reinforcement learning will eventually recover an end-to-end trained agent, while the other latent representations won't.

Quality and Details

The motivation of the paper is interesting and a systematic approach is taken to answer the question which representation learning objective is best, starting from a theoretical analysis and ending in an empirical study. However, I am a bit concerned about the experimental side of the paper which is a bit scarce.

Clarity

The paper is clearly written and easy to follow.

Originality and Significance

The research question is interesting and original, but the experiments only deal with 2 simple environments.

Pros

The paper has a good motivation and a systematic approach to answer the raised question.

Cons

The empirical evaluation is limited.

Minor

I have the feeling that the experiments in Figure 5 and 6 are missing baselines? In Figure 5, I can't see a state-only ablation, while in Figure 6, I can't see an inverse-information ablation?

Maybe I skipped it, but what was k in the experiments?
review_title: Interesting research question but limited empirical evaluation

final_score: rating: 5: Marginally below acceptance threshold
review_content: **Update**
I have briefly checked the updated paper and the corresponding Proposition 1. While I currently do not find any issues with the counterexample, there are still many tiny issues in the proofs that prevent me from recommending acceptance.

- Prop 1 statement, $Q(a, s)$ appears again. It does not type check and not fixed in Def. 3.
- Lemma 2 statement, the last quantity, what is small $z$ in the integral? I also don't see why the $\exists$ symbol is there, isn't $Z$ as a random variable already defined, so $p(Z | X)$ is simply the conditional distribution?
- The step around (21) and (22) are very unclear. I can see what you are trying to do here, but I think it should be laid out step by step.

**Overview**
The paper discuss on a theoretical level what mutual information 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 validity of the theory. 

**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.). It can be used to discuss which MI based objective is adequate for RL in RL.

**Weaknesses**
I think the negative argument about state-only transition information is a bit flawed. The problem that I see here are two fold.  
- One, arguments about the action distribution seems missing here, which should affect the representation distribution (also, it seems that the action distribution is not restricted to depend on $Z$, although the optimal policy clearly does?)
- Two, the proof does not discuss timesteps beyond $t=0$.

More concretely, if we assume Figure 2 (left), then $s_0$ gives no information about $s_1$ or $s_2$, so the representation learning has the freedom to compress $s_1$ and $s_2$ together. But if we consider the next timestep, then the compressed representation might not be optimal. Suppose 
- At time = 1, I have 0.8 probability on $s_1$, 0.2 probability on $s_2$.
- At time = 2, my action distribution ensures that I have 0.5 probability of landing on either $s_1$ and $s_2$ when I start from either $s_1$ or $s_2$ (the graph for a Figure 2 has a bug for s1 and s2 actions)

Then if I use compressed representation, state-only information is zero (because I compressed the representation to 1 state anyways), but if I used original representation, then $H(S_2 | S_1) = 0.5 \log 2$, and $H(S_1)$ is smaller because it is Bernoulli(0.2). Therefore, the counterexample does not show that state-only information is not adequate. I am inclined to believe that the claim is true with added assumptions over the action distribution, but the proof presented here is not correct.

Minor comments:
- Definition 3, Q(a, s) does not type check.
- In Lemma 1, what is equation 8?
- Proof of lemma 2 seems problematic: why is $H(A+B) - H(A | Y) - H(B | Y) = I(Y; A+B)$ true? If A and B are near opposites of each other, then $A+B \approx 0$, and $H(A + B | Y) \approx \infty > H(A | Y) + H(B | Y)$. I think you need additional independence assumptions.
review_title: Interesting perspective, but the state-only argument seems problematic.

Meta-review: **Overview**: 
The paper tries to answer which mutual information (MI) objective is sufficient  for representation learning (repL) in reinforcement learning (RL). Three common objectives are considered: forward, state, and inverse. The paper shows that only the forward objective is sufficient for learning, i.e., sufficient for learning of optimal policy/value function. The authors also demonstrate this phenomena using empirical experiments.

**Quality, Clarity, Originality and Significance**: 
All the reviewers believe this paper is novel in terms of methodology, i.e., evaluate the sufficiency of the repL in terms of down stream tasks. However, there is a lack of clarity in the experiment sections. The authors have provided more details in the rebuttal phase. The reviewers also have concerns that this paper may be too far from typical experimental settings to have a real impact on the field. An unofficial review pointed out there is a mistake in the proof of the paper. The authors later also confirmed the flaw and claimed it is fixed.

**Recommendation**:
The paper is indeed interesting and novel. However, the impact to the practice community might not be significant. That being said, the paper should warrant publication eventually. However, the authors changed large amount of text about the proofs before and after rebuttal, which also introduced some additional typos, confusions, and also technique sloppiness or flaws. The reviewers are concerned about this. Overall I believe that the paper is not in a state to be published yet.

Decision: Reject