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

Abstract: Mutual information maximization provides an appealing formalism for learning representations of data. In the context of reinforcement learning, such representations can accelerate learning by discarding irrelevant and redundant information, while retaining the information necessary for control. Much of the prior work on these methods has addressed the practical difficulties of estimating mutual information from samples of high-dimensional observations, while comparatively less is understood about \emph{which} mutual information objectives are sufficient for RL from a theoretical perspective. In this paper we identify conditions under which representations that maximize specific mutual-information objectives are theoretically sufficient for learning and representing the optimal policy. Somewhat surprisingly, we find that several popular objectives can yield insufficient representations given mild and common assumptions on the structure of the MDP. We corroborate our theoretical results with deep RL experiments on a simulated game environment with visual observations.

Under review as a conference paper at ICLR 2021
WHICH MUTUAL-INFORMATION REPRESENTATION LEARNING OBJECTIVES ARE SUFFICIENT FOR CONTROL?
Anonymous authors Paper under double-blind review
ABSTRACT
Mutual information maximization provides an appealing formalism for learning representations of data. the context of reinforcement learning, such representa- tions can accelerate learning by discarding irrelevant and redundant information, while retaining the information necessary for control. Much of the prior work these methods has addressed the practical difficulties of estimating mutual information from samples of high-dimensiona. observations, while comparatively less is understood about which mutual information objectives sufficient for inforcement learning (RL. from a theoretical perspective. In this paper we identify conditions under which representations that maximize specific mutual-ing mutual-information objectives are theoretically sufficient for learning and representing the optimal policy. Somewhat surprisingly. we find that several popular objectives can yield insufficient representations given mild and common assumptions on the structure of the MDP. We corroborate our theoretical results with empirical results experiments on a simulated game environment with visual observations.
INTRODUCTION
While deep reinforcement learning (RL) algorithms are capable of learning policies from high- dimensional observations, such as images (Mnih et al. 2013] Lee et al., (2019: Kalashnikov et al.) 2018). practice policy learning faces a bottleneck in acquiring useful representations of the observation space (Shelhamer et al., (016). State representation learning approaches aim to remedy this issue by learning structured and compact representations on which to perform RL. While a wide range of representation learning objectives have been proposed (Lesort et al.) 2018), particularly appealing class of methods that is amenable to rigorous analysis is based on maximizing mutual information (MI) between variables. In the unsupervised learning setting, this is often realized as the InfoMax principle (Linsker 1988 Bell & Sejnowski 1995), which maximizes the mutual information between the input and its latent representation subject to domain-specifie constraints. This approach has been widely applied in unsupervised learning in the domains of image, audio, and natural language understanding (Oord al. 2018 Hjelm et al. 2018 Ravanelli & Bengio 2019). In RI., the variables of interest for MI maximization are sequential states, actions, and rewards (see Figure [). As we will discuss, several popular methods for representation learning RL involve mutual information maximization with different combinations of these variables (Anand et al. 2019
Oord et 2018 Pathak al. 2017 Shelhamer et al. 2016).
A useful representation should retain the factors of variation that are necessary to learn and represent the optimal policy or the optimal value function, and discard irrelevant and redundant information. While much prior work has focused on the problem of how to optimize various mutual information objectives in high dimensions (Song& & Ermon 2019 Belghazi al. 2018 Oord et al.) 2018 HHjem et al. 2018), we focus instead on whether the representations that maximize these objectives are actually theoretically sufficient for learning and representing the optimal policy or value function. We find that some commonly used objectives are insufficient given relatively mild and common assumptions on the structure of the MDP, and identify other objectives which are sufficient. We show these results theoretically and illustrate the analysis empirically in didactic examples in which MI can be computed exactly. Our results provide some guidance to the deep RL practitioner on when and why objectives may be expected work well fail, and also provide framework to analyze newly

---

Under review as a conference paper at ICLR 2021
proposed representation learning objectives based on MI. To investigate how theoretical results pertain to deep RL, we compare the performance of RL agents in a simulated game trained with state representations learned by maximizing the MI objective given visual inputs. The experimental results corroborate our theoretical findings, and demonstrate that the sufficiency of a representation can have substantial impact the performance of an RI. agent that uses that representation.
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-basee unsupervised learning- Mutual information-basee methods are par- ticularly appealing for representation learning as they admit both rigorous analysis and intuitive interpretation. Tracing its roots to the InfoMax principle (Linsker) 1988 [Bell & Sejnowski 1995), a common technique is maximize the MI between the input and its latent representation subject to domain-specific constraints (Becker & Hinton| 1992). This technique has been applied to learn repre- sentations for natural language (Devlin et al. 2019), video (Sun et al. 2019), and images (Bachman iet al. 2019 Hjelm et al. 2018). A major challenge to using MI maximization methods in practice is the difficulty of estimating MI from samples (McAllester & Statos 2018) and with high-dimensional
inputs (Song & Ermon, 2019). Much recent work has focused on improving MI estimation via variational methods (Song & Ermon 2019: et al., 2019) Oord et al. [2018] Belghazi et al., 2018). In this work we are concerned with analyzing the MI objectives, and not the estimation method, our experiments with image observations, we make use of noise contrastive estimation methods (Gutmann & Hyvärinen 2010), though other choices could also suffice.
Mutual information objectives in RL. Reinforcement learning 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 works omit the action, maximizing the MI between current and future states (Anand eet al. 2019; Oord et al. 2018; Stooke et al. 2020). Much prior work learns latent forward dynamics models (Watter et al. 2015; Karl et al. 2016 Zhang et al. 2018b Hafner al., 2019 Lee al..
2019), related to the forward information objective we introduce in Section 4 Multi-step inverse models, closely related to the inverse information objective (Section 4 have been used learn
control-centric representations (Yu et al. [211 Gregor et al.) (0166). Single-step inverse models have been deployed as regularization of forward models (Zhang et al.) 2018a; Agrawal al.) 2016) and as an auxiliary loss for policy gradient RL [Shelhamer et al. (2016); Pathak et al. (2017). The MI objectives that we study have also been used as reward bonuses improve exploration, without impacting the representation, in the form of empowerment (Klyubin al. 2008 2005| Mohamed & Rezende 2015 Leibtried et al., 2019) and information-theoree curiosity (Still & Precup) 20112).
Representation learning for reinforcement learning. In RI. the problem of finding compact state space has been studied as state aggregation or abstraction (Bean et al. 1987: Li et al.) 2006). Abstraction schemes include bisimulation (Givan al., 2003), homomorphism (Ravindran & Barto 2003), utile distinction (McCallum 1996), and policy irrelevance (Jong & Stone 2005). While efficient algorithms exist for MDPs with known transition models for some abstraction schemes such as bisimulation (Ferns et al. 2006 Givan et al. 2003), in general obtaining error-free abstractions is highly impractical for most problems of interest. For approximate abstractions prior work has bounded the sub-optimality of the policy (Bertsekas et al. 1988; Dean & Givan 1997| Abel et al. 2016) as well as the sample efficiency (Lattimore & Szepesvari 2019; Van Roy & Dong 2019 9
Du et al. 2019). with some results extending to the deep learning setting (Gelada et al. 2019
Nachum et al. (0118. In this paper, we focus on whether a representation can used to learn the optimal policy, and not the tractability of learning. Alternatively, priors based on the structure of the physical world can be used to guide representation learning (Jonschkowski & Brock 2015). In deep RL, many auxiliary objectives distinct from the objectives that we study have been proposed, including leta-ng general value functions (Veeriah et al. 2019),, predicting multiple value functions (Bellemare et al. 2219 Fedus et al.) 2019 Jaderberg et al. (0166) and predicting domain- specific measurements (Mirowski, Dosovisskyy & Koltun 2016). We restrict our analysis to objectives that can be expressed Ml-maximization.
2

---

Under review as a conference paper at ICLR 2021
3 REPRESENTATION LEARNING FOR RL
The goal of representation learning for RL is to learn a compact representation of the state space that discards irrelevant and redundant information. In this section formalize each part of this statement, starting with defining the RL problem and representation learning in the context of RL. We then propose and define the metric of sufficiency to evaluate the usefulness of representation.
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.A, Tr). where S the set of states, A the set of actions, T A [0. 1] the state transition distribution, and S x A x S + R the reward function. We will use capital letters to refer to random variables and lower case letters to refer to values of those variables (e.g., the random variable for the state analysis
and is specific state). Throughout 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 given reward function and denote this optimal policy as के = arg max, n E, E. ('r, Ac)] for discount factor 7/- We also define the optimal Q-function as (s, at) = Ex 1= r As) ay]. The optimal Q-function satisfies the recursive Bellman equation, (5 ay) r(s, a,) + >Ep(s 11 mn.) arg max ($ a(+1). The optimal policy and the optimal Q-function are related according to x =s(ss arg max (s a).
will often be interested in multiple reward
Mutual information. In information theory, the mutual information (MI) between two random variables, X and Y, is defined as (Cover 1999):
p(z, v) = (X) H(XI)).
I(x;Y) = Ep(a,) log
(I)
p(z)p(v)
The first definition indicates that MI can be understood as a relative entropy (or KL-divergence), 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 RI. The goal of representation learning for RL is to find compact representation of the state space that discards details in the state that are not relevant for representing the policy or value function, while preserving task-relevant informa- tion (see Figure T). While state aggregation methods typically define deterministic rules to group states in the representation (Bean et al. 1987 Li al. 2006), MI-based representation learning methods used for deep RL treat the representation as random variable (Nachum ct al. 2018 Oord al. 2018 Pathak et al.) (0017). Accordingly, we formalize a representation as stochastic mapping between original state space and representation space.
A
Ac11
S₂
S
Z+
Z
Figure 1: Probabilistic graphi- cal model illustrating the state representation learning prob- lem, estimating state represen- tation Z from original state S.
A stochastic representation oz (s) is mapping from states & E S to a probability distribution pl S 8) over elements of new representation space E Z.
Definition 1.
In this work consider learning state representations from data by maximizing an objective J. Given an objective J. we define the set of representations that maximize this objective as Φ (Øz) s.t. oz arg max J(Ø).
Unlike problem formulations for partially observed settings (Watter et al., 2015 Hafner et al. 2019 Lee al. 2019). we assume that S is Markovian state; therefore the representation for given state is conditionally independent of the past states, a common assumption in the state aggregation literature (Bean al. 1987 Li et al. 2006). See Figure T for depiction of the graphical model.
3

---

Under review as a conference paper at ICLR 2021
3.2 SUFFICIENT REPRESENTATION FOR REINFORCEMENT LEARNING
We now turn to the problem of evaluating stochastic representations for RL. Intuitively, we expect useful state representation be capable of representing the optimal policy in the original state space.
Definition 2. A representation Oz is T-' -"-sufficient with respect set of reward functions R if Vr E R. (s1) = (22) = म. (A|si) = (A(s2).
When a stochastic representation Φ produces the same distribution over the representation space for two different states $1 and $2 we say it aliases these states. Unfortunately, as already proven in Theorem of Li et al. (2006) for the more restrictive case 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 Q-learning to the optimal policy in the original state space (refer to Theorem 4 of Li et al. (2066) for the proof of this).
Definition 3. A representation 62 is Q"-sufficient with respect to a set of reward functions R if Vrr E R, (si) = 22((2) Va, (a, $1) (a,
Note that Q*-sufficiency implies x -sffficiienyy since the optimal policy and the optimal Q-function are directly related via (s) arg max, (5, a) (Sutton Barto 2018); however the converse is not true. We emphasize that while Q²" -sufficiency guarantees convergence, it does not guarantee tractability, which has been explored in prior work (Lattimore & Szepesvari 2019 Du al. 2019).
We will further say that an objective J is sufficient with respect to some set of reward functions RR if 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 even when R contains single known reward function.
MUTUAL INFORMATION FOR REPRESENTATION LEARNING IN RL
In our study, consider several MI objectives proposed in the literature.
Forward information: A commonly sought characteristic of state representation is to ensure it retains maximum predictive power over future state representations. This property satisfied by representations maximizing the following MI objective,
J Jfwd (( Ze, At) H(Z++)) (Z Zz, At).
We suggestively name this objective "forward information" due to second term, which the entropy of the forward dynamics distribution. This objective is related to that proposed in Nachum ct al. 2018), where they consider a sequence of actions.
State-only transition information: Several popular methods (Oord al.) 2018 Anand al. 2019 Stooke et al. 2020) optimize a similar objective, but do not include the action
Istate T(7 Z,) H(*++)) H(?++al7.))
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 1,+1- Such representations can be learned by maximizing the following information theoretic objective:
Jino I(A; Ze+x = H (A, Ze+k)
We suggestively name this objective "inverse information" due to the second term, which is the entropy of the inverse dynamics. A wide range of prior work learns representations by optimizing
It is common to history of states, optimizing objective like 7 T(7++t; 71,.., Ze). In our setting assume the given state is Markovian, in which case the objectives equivalent.
4

---

Under review as a conference paper at ICLR 2021
closely related objectives (Gregor et al. ct al., 2017 Yu et al. 2019 Zhang et al.
2016 Shelhamer et al. 2016 Agrawal et al. 2016 Pathak
2018a). Intuitively, inverse models allow the representation to capture only the elements of the state that necessary to predict the action, allowing the discard of potentially irrelevant information.
5 SUFFICIENCY ANALYSIS
In this section we analyze the sufficiency for control of representations obtained by maximizing each objective presented in Section[4 4 To focus on the representation learning problem, we decouple from RL by assuming access with some probability, which can then be used to learn the desired representation. We also assume that distributions, such as the dynamics or inverse dynamics, can be modeled with arbitrary accuracy. and that the maximizing set of representations for a given objective can be computed. While these assumptions might be relaxed in any practical RL. algorithm, and exploration plays a confounding role, studying these objectives under such simplifying assumptions allows to compare them in terms of sufficiency on an equal playing field, isolating the role of representation learning from other confounding components of complete RL algorithm.
dataset of transitions collected with a policy that reaches all states
5.1 FORWARD INFORMATION
In this section show that representation that maximizes J fud is sufficient for optimal control under any reward function. This result aligns with intuition that a representation that captures forward dynamics can represent everything predictable in the state space, and can thus be used learn the optimal policy for any task. Note that this strength can also be a 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.
Proposition 1. f is sufficient for all reward functions.
Proof. (Sketch) We first show that if Z,, Ar are maximally informative of Ze+k. they are also maximally informative of the return Rg Due to the Markov structure, =, Ze. Ar) p( RR s, Ar). In other words, given Φz, additionally knowing S doesn't change belief about the future return. The Q-value is the expectation of the return, so has as much information about the Q-value S does. The full proof can be found in Appendix 8.1
5.2 STATE-ONLY TRANSITION INFORMATION
While Jstate is closely related to J fwdd we now show that it is not sufficient.
Proposition Istate not sufficient for all reward functions.
Suppose that the two actions a. and aı equally likely under the policy distribution. Each state gives no information about which the two possible next states is more likely; this depends the action. Therefore, representation maximizing Istate free to alias states with the same next-state distribution, such as So and $3- An alternative view that such representation can maximize Jstate (Z+++)) H (Ze+ Zz) by reducing both terms in equal amounts aliasing So and S3 decreases the marginal entropy as well as the entropy of capable of representing the optimal policy which must distinguish So and Sy in order to choose the correct action to reach 82, which yields reward.
Proof. Consider the counter-examplee in Figure 2
predicting the next state starting from S or $2. However, this aliased representation
5.3 INVERSE INFORMATION
Here we show that representations that maximize Jinn are not sufficient for control in all MDPs. Intuitively, way that the representation be insufficient is by retaining only controllable state
5

---

Under review as a conference paper at ICLR 2021
Jee
a0
Value of MI objective
12
Asst
return 1.0 (optimal)
a1
10-
RL return 0.5
50
S2
04
02
$1
S3
0.0
0.6 0.8
10
1.2 14
State representations, ordered by I(Z: S)
Figure 2: (left) representation that aliases the states So and $3 into single state maximizes Jstate but is not sufficient to represent the optimal policy which must choose different actions in So and 83 reach $2 which yields reward. (right) Values Jattt and Juwd for a few representative state representations, ordered by increasing T(Z:)) S). The representation that aliases So and 83 (plotted with diamond) maximizes Jatates but the policy learned with this representation may not optimal (as shown here). The original state representation (plotted with star) is sufficient.
elements, while the reward funetion depends on state elements outside the agent's control. We then show that additionally representing the immediate reward is not enough to resolve this issue.
Jine is not sufficient for all reward functions. Additionally, adding [R2 Zi) the objective does not make it sufficient.
Proposition 3.
Proof. Consider the MDP illustrated in Figure 3 and the representation that aliases the states so and $1. The same actions taken from these states lead to different next states which may have different rewards (ag leads to the reward from Sp while aj leads to 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 immediate rewards. The same counterexample holds since assumed So and have the same reward.
a0
a1
S2
Value of MI objective
16
Jav
80
$3
RL return 1.0 (optimal)
1.4
return 0.50
1.-
S1
S4
0.8
85
0.
1's 1
State representations, ordered by
Figure 3: (left) In this MDP, representation that aliases the states and S into single state maximizes Jines yet is not sufficient represent the optimal policy, which must distinguish between So and $1 in order to take a different action (towards the high-reward states outlined in green). (right) Values of Jine and J wd for few selected state representations, ordered by increasing [ZZ S). The representation that aliases and (plotted with diamond) maximizes Jine. but is not sufficient to learn the optimal policy. Note that this counterexample holds also for Jino + 7 [RR Z).

---

Under review as a conference paper at ICLR 2021
6 EXPERIMENTS
In this section, present experiments studying MI-based representation learning with image observations, to analyze whether the conclusions of our theoretical analysis hold in practice. Our goal is not to show that any particular method is necessarily better worse, but rather to illustrate that the sufficiency arguments that we presented translate into quantifiable performance differences in the deep RL setting.
6.1 EXPERIMENTAL SETUP
To separate representation learning from RL, we first optimize each representation learning objective on a dataset of offline data consisting of 50k transitions collected from a uniform random policy. We then freeze the weights of the state encoder learned in the first phase and train RL agents with the representation as state input. To clearly illustrate the characteristics of each objective, we use the simple pygame (Shinners 2011 I) video game catcher, in which the agent controls a paddle that can move back and forth to catch fruit that falls from the top of the screen (see Figure 4. A positive reward given when the fruit caught and a negative reward when the fruit is not caught. The episode terminates after one piece of fruit falls. We optimize I fard and Jstate with noise contrastive estimation (Gutmann & Hyvärinen| 2010), and Jine by training an inverse model via maximum likelihood. For the RL algorithm, modified slightly for the discrete action distribution. Please see Appendix 8.2 for full experimental details.
the Soft Actor-Critt algorithm Haarnoja et al. (2018),
6.2 COMPUTATIONAL RESULTS
In principle, expect that representation learned with Jine may not sufficient to solve the catcher game. Because the agent does not trol the fruit, representation maximizing J Jine might discard that information, thereby making impossible to represent the optimal policy. We observe in Figure 5 (top left) that indeed rep- resentations trained to maximize Jive result in RL agents that converge slower and lower asymptotic expected return. Further, attempting to learn a decoder from the learned representa- tion to the position of the falling fruit incurs high error (Figure bottom left), indicating that the fruit is not precisely captured by the repre- sentation. We argue that this type of problem setting is not contrived, and is representative of many situations in realistic tasks. Consider, for instance, an autonomous vehicle that stopped at a stoplight. Because the agent does not control the color of the stoplight, it may not be captured in the representation learned by Jane and the resulting RL policy may choose run the light.
Figure 4: (left) Original catcher game in which the agent (grey paddle) moves left or right catch fruit (yellow square) that falls from the top of the screen. (right) Variation catcher-grip in which agent is instantiated as gripper, and must open the gripper to catch fruit.
In the second experiment, we consider a failure mode of J Jetate- We augment the paddle with a gripper that the agent controls and must be open order properly catch the fruit. 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 Istate might alias states where the gripper is open with states where the gripper is closed. In our experiment, we see that the error predicting the state of the gripper from the representation learned via Jsstate is chance (Figure 5 bottom right). This degrades the performance of 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 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 learn when properly the headlights.
7

---

Under review as a conference paper at ICLR 2021
catcher
catcher-gra
1.0
1,00
0.75
0.5
0.50
Return
Return
0.25
0.0
0.00
0.25
00.00
-1.0
0.00
0.15
0.30
0.45
0.60
0.00
0.15
0.30
0.45
0.60
Env Steps (Millions)
Env Steps (Millions)
end-to-end
J food
Jane
Jatate
Obj.
Agent Err.
Fruit Err.
Obj.
Agent Err.
Fruit Err.
Grip Err.
J fwd
0.1
0.1
J
0.1
0.1
0.3
Dino
0.15
0.47
J Jstate
0.1
0.1
0.5
Figure 5: (top) Policy performance using learned representations as state inputs to RL, for the catcher and catcher-grij environments. (bottom) 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 Istate need not represent the gripper, leading these representations to perform poorly in catcher and catcher-grip respectively.
J produces useful representations in all cases, and is equally or more effective than learning representations purely from the RL objective alone (as in Figure 5). We experiment with more visual complexity by adding background distractors; these results are presented in Appendix 8.4 We find that in this setting representations learned with J fud to yield even larger gains over learning representations end-to-end via RL. We also analyze the learned representations by evaluating how well they predict the optimal Q" in Appendix 8.3
DISCUSSION
In this work, aimed to analyze mutual information representation learning objectives for control from a theoretical perspective. In contrast to much prior work that studies how these objectives can be effectively optimized given high-dimensional observations, we analyze which objectives are guaranteed to yield representations that are actually sufficient for learning the optimal policy, Surprisingly, show that common objectives yield representations that are theoretically insufficient, and provide a proof of sufficiency for third. We validate our theoretical results with an empirical investigation on a simple video game environment, and show that the insufficiency of these objectives can degrade the performance of deep RI agents.
We view this investigation as a step forward in understanding the theoretical characteristics of representation learning techniques commonly used in deep RL. We see many exciting avenues for future work. First, identifying more restrictive MDP classes which insufficient objectives are in fact sufficient, and relating these to realistic applications. Second, investigating if sample complexity bounds can be established the case of a sufficient objective. Third, extending our analysis to the partially observed setting, which is more reflective of practical applications. We these directions as fruitful in providing a deeper understanding of the learning dynamics of deep RL, and potentially yielding novel algorithms for provably accelerating RL. with representation learning.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
David Abel, D) Ellis Hershkowitz, and Michael I. Littman. Near optimal behavior via approximate state abstraction. In Proceedings of the 33rd International Conference on International Conference on Machine Bearning-Volume 48, pp. 2915-2923, 2016.
Pulkit Agrawal, Ashvin V Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Learning to poke by poking: Experiential learning of intuitive physics. In Advances neural information processing systems, pp. 5074-5082, 2016.
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, pp. 3766-8779, 2019.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems, pp. 15509,
James C Bean, John R Birge, and Robert I. Smith. Aggregation in dynamic programming. Operations Research, 35(2):215-22 220, 1987.
Suzanna Becker and Geoffrey E Hinton. Self-organizing neural network that discovers surfaces in random-dot stereograms. Nature, 355(6356):661-163 1992.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint orXiv:/801.0400 2018.
Anthony J Bell and Terrence J Sejnowski. An www.fommation-maxiiizzzite approach to blind separation and blind deconvolution. Neural computation, 7(6):
1995.
Mare Bellemare, Will Dabney, Robert Dadashi, Adrien Ali Taiga, Pablo Samuel Castro, Nicolas Le Roux, Dale Schuurmans, Tor Lattimore, and Clare Lyle. A geometric perspective on optimal representations for reinforcement learning. In Advances in Neural Information Processing Systems, pp. 4360-4371, 2019. 2019.
Dimitri P Bertsekas, David A Castanon, al. Adaptive aggregation methods for infinite horizon dynamic programming. 1988.
Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
Thomas Dean and Robert Givan. Model minimization in markov decision processes. In AAAI/IAAI, pp. 106-111, 1997.
Jacob Devlin, Ming-Was Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguisties: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy and Vladlen Koltun. Learning to act by predicting the future. In International Conference Learning Representations, 2016.
Simon S Du, Sham M Kakade, Ruosong Wang, and Lin F Yang. Is good representation sufficient 2019.
for sample efficient reinforcement learning? arXiv preprint
William Fedus, Carles Gelada, Yoshua Bengio, Marc G Bellemare, and Hugo Larochelle. Hyperbolic discounting and learning over multiple horizons. arXiv preprint arXiv: 100.100022222 06865, 2019.
Norm Ferns, Pablo Samuel Castro, Doina Precup, and Prakash Panangaden. Methods for computing state similarity in markov decision processes. In Proceedings of the Twenty-Secone Conference on Uncertainty in Artificial Intelligence, pp. 174-188 181, 2006.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference Machine Learning, pp. 2170-2179, 2019.

---

Under review as a conference paper at ICLR 2021
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence. 147(1-2):163-2 2003.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint 1rXiv:161 1.07507, 2016.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp. 297-304, 2010.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with a stochastic actor. arXiv preprint 1XXiv:8001.012900 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference Machine Learning, pp. 2555-2565, 2019.
R Devon Hjelm, Alex Fedorov, Samuel -ariie-Mrr,,,,oon,, Karan Grewal, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprinl arXiv:
2018.
Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint arXiv:1611.05397, 2016.
Nicholas K Jong and Peter Stone. State abstraction discovery from irrelevant state variables. In IJCAI, volume 8, pp. 752-757, 2005.
Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Autonomous Robots, 39(3):077-428, 2015.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Scalable deep reinforcement learning for dision-based robotic manipulation. Conference on Robot Learning, pp. 173, 2018.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. In International Conference on Learning Representations, 2016.
Alexander S Klyubin, Daniel Polani, and Chrystopher Nehaniv. Empowerment: A universal agent-centric measure of control. In 2005 IEEE Congress Evolutionary Computation, volume 1, pp. 128 135. IEEE, 2005.
Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. Keep your options open: An information- driving principle for sensorimotor systems. PloS one, 3(12), 2008
Tor Lattimore and Csaba Szepesvari. Learning with good feature representations in bandits and in rl with generative model. arXiv preprint 1rXiv:1911.0700 2019.
Alex X. Lee, Anusha Nagabandi, Picter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint (90000090000055333 2019.
Felix Leibfried, Sergio Pascual- Díaz, and Jordi Grau-Moya. A unified bellman optimality principle combining reward maximization and empowerment. In Advances in Neural Information Processing Systems, pp. 7867-7878, 2019.
Timothée Lesort, Natalia Díaz- Rodríguez. Jean-Franois Goudou, and David Filliat. State representa- tion learning for control: An overview. Neural Newworks, 108:379-392, 2018
Lihong Li, Thomas Walsh, and Michael Littman. Towards a unified theory state abstraction for mdps. In ISAIM, 2006.
10

---

Under review as a conference paper at ICLR 2021
Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):105-11 117, 1988.
David McAllester and Karl Statos. Formal limitations on the measurement of mutual information. arXiv preprint arXiv:/811.0420 2018.
Andrew Kachites McCallum, Reinforcement Learning with Selective Perception and Hidden State. PhD thesis, University of Rochester, 1996.
Piotr Mirowski. Learning to navigate. In 1t International Workshop on Multimodal Understanding and Learning for Embodied Applications, pp. 25-25, 2019.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprini arXiv: 12.55602, 2013.
Shakir Mohamed and Danilo Jimenez. Rezende. Variational information maximisation for intrinsically motivated reinforcement learning. In Advances neural information processing systems, pp- 2125-2133. 2015.
Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine. Near-optimal representation learning for hicrarchical reinforcement learning. 2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint
2018. 2018.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-superrvied prediction. In Proceedings of the 34th International Conference on Machine 70, pp. 7778-2787. JMLR. org, 2017.
Ben Poole, Sherjil Ozair, Aaron van den Oord, Alexander A Alemi, and George Tucker. On variational bounds of mutual information. arXiv preprint 1905.09900.0092 2019.
Mirco Ravanelli and Yoshua Bengio. Learning speaker representations with mutual information, 2019.
Balaraman Ravindran and Andrew G Barto. Smdp homomorphisms: an algebraic approach to abstraction in cemi-markov decision processes. In Proceedings of the 18th international joint conference on Artificial intelligence, pp. 1011-1016, 2003.
Evan Shelhamer, Parsa Mahmoudieh, Max Argus, and Trevor Darrell. Loss is its own reward: Self-supervision for reinforcement leannin arXiv preprint 17Xiv:1612.07307,2000 2016.
Pete Shinners. Pygame. http::/ org/ 2011.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information
estimators. arXiv preprint arXiv:1910.06222, 2019.
Susanne Still and Doina Precup. An information-theerriicc approach to curiosity-driven reinforcement learning. Theory Biosciences, 131(3):139-144 148, 2012.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. Technical report, UC Berkeley, 2020.
Chen Sun, Fabien Baradel, Kevin Murphy, and Cordelia Schmid. Contrastive bidirectional transformer for temporal representation learning. arXiv preprint
2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Benjamin Van Roy and Shi Dong. Comments on the lu-kakade-wang-yaann lower bounds. arXiv preprint arXiv:1911.078 2019.
Vivek Veeriah, Matteo Hessel, Zhongwen Xu, Janarthanan Rajendran, Richard L Lewis, Junhyuk Oh, Hado van Hasselt, David Silver, and Satinder Singh. Discovery of useful questions as auxiliary tasks. In Advances in Neural Information Processing Systems, pp. 2019.

---

Under review as a conference paper at ICLR 2021
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In Advances in neural information processing systems, pp. 2746-2756 2015.
Tianhe Yu, Gleb Shevchuk, Dorsa Sadigh, and Chelsea Finn. Unsupervised visuomotor control
through distributional planning networks. In Robotics Science and Systems, 2019.
Amy Zhang, Harsh Satija, and Joelle Pineau. Decoupling dynamics and reward for transfer learning. arXiv preprint arXiv: 1804. 10689, 2018a.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured latent representations for model-based reinforcement learning. arXiv preprint 180.: 2018b.
12

---

Under review as a conference paper at ICLR 2021
8 APPENDIX
8.1 SUFFICIENCY OF y PROOF OF PROPOSITION
We describe the proofs for the sufficiency results from Section here. We begin by providing a set of lemmas, before proving the sufficiency of Jl fwdd
A
S
SELa
X
R₁
Rest
Zr
Zook
Figure 6: Graphical model for Lemma depicting true states S, states in the representation Z, actions A, rewards R, and the variable X (which we will interpret as the sum of future rewards the proof of Proposition
Lemma 1. Let X be random variable dependent on SitK with the conditional independence assumptions implied by the graphical model in Figure 66 (In the main proof of Proposition | we will let X be the sum of rewards from time + k onwards.) If (Ze4ki Ze, A.) 1(S, S,, A. Vk, then I(X; Z,, A.) [(X; St, A.)yk.
Proof: For proof by contradiction, assume there is some Øz and some r such that 7(X; A,) <
T(X; St, At) and that I ( Z,, A,) = I(S++1; S,, At). Now we know that because
S
SiLk + Zi+k form Markov chain, by the data processing inequality (DPI) I(2+4:i Ze, Ar) S
I[S:+N; Ze, At) S 1(Se4R St, At). We will proceed by showing that that [(X; ; Ze: At)
I(X; St, At)
I(S++ki Zı, At) I(S++A: St, At). which gives the needed contradiction.
Using chain rule, we can expand the following expression in two different ways.
((x St, A.) T(X: Z,|S, 4,) + I(X; S, A,) = 0 I(X: Ad)
(5)
T(X: Z,, S, = IXX; (4) T(X;
(6)
Note that the first term in Equation Equating the expansions, we can see that to satisfy our assumption that I(X; Ze, Ad) I(X: A.). we must have that S, A,) > O.
5 is zero by the conditional independence assumptions in Figure 6
Now we follow a similar procedure to expand the following expression:
(S Ze, Se, Av) [(S+4t; Z, Ar) + (S A.) 0 I(Sr44; A,)
+ +T I(S Z,, A,) = ((() Ze, AA)
(8)
The first term in Equation 7 is zero by the conditional independence assumptions in Figure 6
Comparing the first term in Equation 8 with the first term in Equation 6 we see because S, S444 + form a Markov chain, by the DPI that /(5++ S Zr A.) > [X; Z., At). Therefore must have T(S++1; S, Ze, At) O. Combining Equations and
T(S Sr, A,) I(S+44: 5y Z,, A.) I(Se+k; Z,, A,)
Since I(S++R; Se A.) 0, [S5++t; Z,, Ax) < (S Ar), which is exactly the contradiction we set out to show.
Lemma If I(Y; 2) = I(Y;X) and Y 1 Z|X, then 3p(Z s.t. Vz, tY = 2) =
MY Z z)dz.
13

---

Under review as a conference paper at ICLR 2021
Proof. First note that the statement is not trivially true. Without any assumption regarding MI, can write,
x) =x))== /
z)dz pYY|Z, 1 =2)xK\\)p-Z= = z)d2
(10)
Comparing this with the statement we'd d like to prove, we can see that the key idea to show that the MI equivalence implies that p(Y|Z, X =) 2)) To begin, consider I(Y; Z) I(Y; We can re-write this equality using the entropy definition of MI.
H(Y) (Y H(Y X
(11)
Note that the HY) cancel and substituting the definition of entropy we have:
p 2) log (Z) = (ocyx.x]loopp(]a.
(12)
Note that on the right-hand side, we can use the Tower property to re-write the expectation as
Ex(rx) (orxxx))lopp((xx)) = Bp(z) Ep(xx [log ((Y = (x(r)x)((x22) (log (l)
(13)
Now we can use Tower property again to re-write the expectation on both sides.
ex 2x
Ex(x) 1 p
Bx(x) [ [ogpp [ htwp://BBBxxxnx///x////xxxxXXX11
(14)
(x) 2
log p(Y [X]] = 0
Ep(x) E (axx))/|/x(xxxx00lgggg
loy X)]] = 0
Log probabilities are always - 0, therefore for the sum to equal zero, each term must be zero.
log
(15)
By Jensen's 's inequality,
lYg S log E (121
(16)
By the monotonicity of the logarithm:
p(Y S
(17)
there exists some x and some such that p(Y = X = x) < |M must be some other y' for the same x where p(Y = y1lx = x) > Ex(z|x=o)/p( = p(r|X= = z) must sum to 1.
by = y/|)], then there
because
x(Y| EX - - =2)dz= / p(Y, Z|X r)
(18)
Where the last equality follows by conditional independence of Y and Z given X.
Given the lemmas stated above, we can then use them to prove the sufficiency of y
Proposition! 1. (Sufficiency I fwa) Let (S, A,T.A,TT)) be an MDP with dynamics p(S++1 At). Let the policy distribution p p(A{S) S) and steady-state state occupancy p(S) have full support on the action and state alphabets A and S respectively. See Figure a graphical depiction of the conditional independence relationships between variables.
For a representation Φz and set of reward functions functions R, if if Zr, Ar) is maximized Ak > 0, 00 Va, V, s1) (a.
then Vr E and Vs1 S2 S, (2 oz (S2)
Proof. Note that (Z++ki Z,, Ad) is maximized if the representation Φ2 is taken to be the identity, In other words max, I(Z++a; Ze, Ad) = (St+k:Se At).
Define the random variable R2 to the discounted return starting from state 8t.
カ-t
R =
R
(19)
14

---

Under review as a conference paper at ICLR 2021
Plug in Rs for the random variable X in Lemma 1
I( 2. Ar) = [(S(+t: Ar)
I(Re+ki Z,) =
( I(Re+s:: St)
(20)
Now let X [5),A)]: Y R2, and Z = Zz, and note that by the structure of the graphical model in
Figure Y 1 Z|X. Plugging into Lemma
(S=ss) AA
(21)
Now the Q-function given a reward function r and a state-action pair (s, a) can be written as an expectation of this random variable R, given St = and A a. (Note that p( Re s, A, a) can be calculated from the dynamics, policy, and reward distributions.)
(s a) = Ept =,)
222)
Since (z((s) 22 p(z,|S, $1) piz: Sy 22). Therefore by Equation 21 p S.
$1, A,) = p( R4 S = 52 A.). Plugging this result into Equation 22 Q (a, 81 = Q a, S2). Because this reasoning holds for all Q-functions it also holds for the optimal Q, therefore (a, (a, s2).
8.2 EXPERIMENTAL DETAILS
8.2.1 DIDACTIC EXPERIMENTS
The didactic examples are computed as follows. Given the list of states in the MDP, compute the possible representations, restricting our focus to representations that group states into "blocks." We do this because there are infinite stochastic representations and the MI expressions are not convex in the parameters of p(Z]S), making searching over these representations difficult. compute the value of the MI objective as well as the optimal
consider
Given each state representation, value function using exact value iteration. In these examples, we assume that the policy distribution is uniform, and that the environment dynamies deterministic. Since we consider the infinite horizon setting, we use the steady-state state occupancy in our calculations.
8.2.2 DEEP RL EXPERIMENTS
The deep RL experiments with the catcher game are conducted as follows. First, we use a uniform random policy to collect 50k transitions in the environment. In this simple environment, the uniform random policy suffices to visit all states (the random agent is capable of accidentally catching the fruit, for example). Next, each representation learning objective is maximized this dataset. For all objectives, the images are pre-processed in the same manner (resized to 64x64 pixels and normalized) and embedded with convolutional network. The convolutional encoder consists of five convolutional layers with ReLU activations and produces a latent vector with dimension 256. We use the latent vector estimate each mutual information objective, as described below.
We interpret the latent embeddings of the images St and S as the parameters of Gaussian distributions p( S.) and p(Z]S(+1). We obtain a single sample from each of these two distributions, concatenate them and pass them through single linear layer to predict the action. The objective we maximize is the cross-entropy of the predicted actions with the true actions, as in Agrawal et al. 2016 and Shelhamer et al. 2016. To prevent recovering the trivial solution of preserving all the information in the image, we add an information bottleneck to the image embeddings. We tune the Lagrange multiplier on this bottleneck such that the action prediction loss remains the same value as when trained without the bottleneck. This approximates the objective ming 7 [ZZ; I max I Imv- current timestep and take the mean of the predicted distribution as the state for the RL agent.
Inverse information:
To use the learned encoder for RL, embed the image from the
Note this result is stronger than what we needed: it means that representations that maximize J are guaranteed to be able to represent even sub-optimal Q-functions. This makes sense in light of the fact that the proof holds for all reward functions the sub-optimal Q under reward is the optimal under another.
15

---

Under review as a conference paper at ICLR 2021
State-only information: We follow the Noise Contrastive Estimation (NCE) approach presented in CPC (Oord et al. 2018). Denoting Z and Zt + 1 as the latent embedding vectors from the convolutional encoders, use log-bilinear model as in CPC to compute the score: f(,,, Ze+1) exp(Z? W Ze+1) for the cross-entropy loss. We also experimented with an information bottleneck as described above, but found that it wasn't needed to obtain insufficient representations. To the learned encoder for RL, we embed the image from the current timestep and use this latent vector as the state for the RL agent.
We follow the same NCE strategy as for state-only information, with the difference that we concatenate the action to Z, before computing the contrastive loss.
Forward information:
We then freeze the state encoder learned via Ml-maximization and use the representation as the state input for RL. The RL agent is trained using the Soft Actor- Critic algorithm Haarnoja et al. (2018). modified slightly for the discrete action distribution (the Q-function outputs Q-values for all actions rather than taking action as input, the policy outputs the action distribution rather than parameters of distribution, and we can directly compute the expectation in the critic loss rather than sampling). The policy and critic networks consist of two hidden linear layers 200 units each. We use ReLU activations.
8.3 ANALYSIS: PREDICTING Q" FROM THE REPRESENTATION
In Section 6 we evaluated the learned representations by running temporal difference RL algorithm with the representation as the state input. In this section, instead of using the bootstrap to learn the Q-function, we instead regress the Q-function to the optimal Q'. To do this, we first compute the (roughly) optimal Q" by running RI. with ground truth game state as input and taking the learned Q as Q" Then, we instantiate new RI agent and train it with the learned image representation as input, regressing the Q-function directly onto the values of Q" We evaluate the policy derived from this new Q-function, and plot the results for both the catcher and catcher-grip environments in Figure 7 We find that similar to the result achieved using the bootstrap, the policy performs poorly when using representations learned by insufficient objectives (Jine in catcher and Jatate in catcher-grip). Interestingly, we find that the error between the learned Q-values and the Q" -"-vllues is roughly the same for sufficient and insufficient representations. We hypothesize that this discrepancy between Qaluue eтrоr and policy performance is due to the fact that small differences in Q-values on a small set of states can result in significant behavior differences in the policy.
catcher-grip
catcher
1.0
0.5
0.5
Return
Return
0.0
(approx.)
(approx.) Q"
-0.5
J fuol
J fust
Itat
-1.0
0.00
0.08
0.16
0.24
0.312
0.90
Env Steps (Millions)
Env Steps
(Millions)
Figure 7: Performance of policies obtained from a Q-function trained to predict Q", given state representations learned by each MI objective, the (left) catcher environment and (right) catcher- grip environment. Insufficient objectives Jine and Jetate respectively perform worse than sufficient objective Jywd-
8.4
DEEP RL. EXPERIMENTS WITH BACKGROUND DISTRACTORS
In this section we repeat the experiments from Section 6 with added visual complexity in the form of background distractors. We randomly generate images of 10 circles of different colors and replace
16

---

Under review as a conference paper at ICLR 2021
the black background of the game with these images. Examples of the agent's observations shown in Figure
fruit
We plot the results for both the catcher and catcher-grip games with distractors in Figure g As in S forming RL with the frozen representation as input (top), well the error of decoding true state elements from the representation (bottom). In both environments, end-to-end RL from im- ages performs poorly, demonstrating the need for representation learning to aid in solving the task. As predicted by the theory, the represen- tation learned by Jine fails both games, and the representation learned by Jastte fails in the catcher-grip game. find that the difference in performance between sufficient and insufficient
show both the result of per-
agent
Figure 8: Example 64x64 pixel observations with background distractors.
objectives is even more pronounced this setting than in the plain background setting.
catcher
catcher-grip
1.0
0.4
0.2
0.0
Return
Return
end-to-end
end-to-end
-0.5
-066
Jyud
fust
Jas
-08
J state
-1.0
1.0
0.15
0.30
0.60
0.15
0.30
0.45
0.00
Steps (Millions)
Env Steps (Millions)
Obj.
Agent Err.
Fruit Err.
Obj.
Agent Err.
Fruit Err.
Grip Err.
I Fwdd
0.1
0.1
Jodd
0.1
0.1
0.3
Jino
0.15
0.47
J Jstote
0.1
0.1
0.5
Figure (top) Policy performance using learned representations as state inputs RL, for the catcher and catcher-gre environments with background distractors. (bottom) Error in predicting the positions of ground truth state elements from each learned representation. Representations maximizing J need not represent the fruit, while representations maximizing Jstate need not represent the gripper, leading these representations to perform poorly in catcher and catcher-grip respectively.
17


Reviews:
[{"review_id": "rVvlHcP8y9U", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Update**\nI 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.\n\n- Prop 1 statement, $Q(a, s)$ appears again. It does not type check and not fixed in Def. 3.\n- 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?\n- 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.\n\n**Overview**\nThe 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. \n\n**Strengths**\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.). It can be used to discuss which MI based objective is adequate for RL in RL.\n\n**Weaknesses**\nI think the negative argument about state-only transition information is a bit flawed. The problem that I see here are two fold.  \n- 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?)\n- Two, the proof does not discuss timesteps beyond $t=0$.\n\nMore 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 \n- At time = 1, I have 0.8 probability on $s_1$, 0.2 probability on $s_2$.\n- 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)\n\nThen 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.\n\nMinor comments:\n- Definition 3, Q(a, s) does not type check.\n- In Lemma 1, what is equation 8?\n- 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.", "extra_scores": {}, "extra_text": {}}, {"review_id": "o8FAdXc9w27", "reviewer": "AnonReviewer3", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\n\nThere 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. \n\nOne 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?\n\nMinor: there’s a reference in the last paragraph on page 6 to Figure 5 which I think should be to Figure 4.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "hx8YlI3mIxu", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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. \n\nl 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.\n\nHowever, l still have the following concerns: \n\nEmpirically 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.     \n\nRegarding 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.  \n\n\nRegarding 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.  \n\nRegarding 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? \n\nFrom 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.  \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "eorCr9ZfMe", "reviewer": "AnonReviewer5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**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.\n\n**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.\n\nAlong 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.\n\nOverall, the paper is well written. Key mathematical concepts are defined clearly in sections 3 and 4. The descriptions and definitions are clear and concise.\n\nAlthough 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. \n\n**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.\n\nIt’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.\n\nThe 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.\n\nThis 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.\n\n**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.\n\n**Questions**:\n\nCould you please provide more details on the training scheme in the experiments section. How is the representation used in practice for downstream tasks?\n\nWhere do other unsupervised representation learning schemes (e.g. VAEs and normalizing flows) fit within this framework?\n\nFor 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?\n\n**Additional Feedback**:\n\nRelated Work:  \nEquations 2 and 4 are referenced far before they are defined. I generally avoid referencing equations this far in advance.  \nI would consider citing Mohamed & Rezende, 2015.\n\nRepresentation Learning in RL:  \nMissing the discount factor in the preliminaries section.  \n\nSufficiency Analysis:  \nFigure 2: J_inv —> J_state\n\nExperiments:  \nFigure 6: should be J_state in the table.", "extra_scores": {}, "extra_text": {}}, {"review_id": "bnyoQaU3FAz", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n\nThis paper studies 3 mutual information (MI) optimization objectives for learning a latent representation Z in a sequential decision-making task of the MDP style:\n\nforward information: MI between Z at time t+k, and Z at time t concatenated with the action at time t,\n\nstate-only transition information: MI between Z at time t+k, and Z at time t,\n\ninverse information: MI between the action at time t, and Z at time t+k, conditioned on Z at time t.\n\nAll 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.\n\nQuality and Details\n\nThe 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.\n\nClarity\n\nThe paper is clearly written and easy to follow.\n\nOriginality and Significance\n\nThe research question is interesting and original, but the experiments only deal with 2 simple environments.\n\nPros\n\nThe paper has a good motivation and a systematic approach to answer the raised question.\n\nCons\n\nThe empirical evaluation is limited.\n\nMinor\n\nI 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?\n\nMaybe I skipped it, but what was k in the experiments?", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

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