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": "pW2CezfZxZa", "reviewer": "Reviewer_qUv5", "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. ", "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": "-  ​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. \n-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. \n- 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. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "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. \n- 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?\n- 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_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}, {"review_id": "n60jYXSyF0D", "reviewer": "Reviewer_nwZo", "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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "**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. \n\n\n-----\n\n**Strengths**\n\nI 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).\n\n**Weaknesses**: Unfortunately, I think the proof of Proposition 1 is false as it is based on an not always true assumption from Lemma 2.\n\nTo 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?\n\nLet 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?\n\nThis 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.\n\n**Experiments**\nIn 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.\n\nIs there any insight as to why J_state works poorly for catcher-grip despite working decently for catcher?\n\n\n\n**Comments**\n\nStatement 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? \n\n\n\nThe type of Q functions are not consistent, sometimes it is $(s, a)$, others it is $(a, s)$. This needs to be fixed.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}, {"review_id": "Sm-goSPZsTi", "reviewer": "Reviewer_Ncmr", "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.", "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": "There are several strong points about this work.\n- 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.\n- The paper contains several original elements, including the sufficiency proof for the forward information objective, and the illuminating counterexamples for the other two objectives.\n- The experiments are smartly designed to make the insufficient objectives fail by \n- 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).\n- The related work is covered well. Furthermore, the discussion section contains many interesting points.\n\nMy concerns are as follows.\n- 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.\n- 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.\n- 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.\n\nAdditional questions to authors.\n- 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?\n- 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.\n- 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.\n-- 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.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The author cover the societal impact by providing citations. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "LQ1lq6PhHo0", "reviewer": "Reviewer_w9XU", "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.", "questions": "", "limitations": "", "rating": 7, "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": "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.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have addressed the limitations of their paper clearly.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

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)