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

1 Introduction . Deep reinforcement learning ( RL ) algorithms are in principle capable of learning policies from high-dimensional observations , such as camera images [ 49 , 39 , 32 ] . However , policy learning in practice faces a 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 . A useful state representation should be sufficient to 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 it 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 RL [ 19 , 35 , 64 ] , 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 we 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 [ 4 , 52 , 53 , 58 ] . While much prior work has focused on how to optimize these various MI objectives in high dimensions [ 61 , 8 , 52 , 27 ] , 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 @ eecs.berkeley.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 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 RL 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-based unsupervised learning . A common technique for unsupervised representation learning based on the InfoMax principle [ 43 , 9 ] is to maximize the MI between the input and its latent representation subject to domain-specific constraints [ 7 ] . This technique has been applied to learn representations for natural language [ 14 ] , video [ 65 ] , and images [ 5 , 27 ] and even policy learning via RL in high dimensions [ 62 ] . To address the difficulties of estimating MI from samples [ 45 ] and with high-dimensional inputs [ 61 ] , much recent work has focused on improving MI estimation via variational methods [ 61 , 55 , 52 , 8 ] . 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 1 ) . 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 [ 4 , 52 , 64 ] . Several prior works [ 51 , 57 , 60 , 44 ] maximize MI objectives that closely resemble the Jfwd objective we introduce in Section 4 , while others optimize related objectives by learning latent forward dynamics models [ 70 , 33 , 74 , 26 , 39 ] . Multi-step inverse models , closely related to the Jinv objective ( Section 4 ) , have been used to learn control-centric representations [ 71 , 23 ] . Single-step inverse models have been deployed as regularization of forward models [ 73 , 2 ] and as an auxiliary loss for policy gradient RL [ 58 , 53 ] . Our result regarding the sufficiency of this objective is similar to 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 the form of empowerment [ 37 , 36 , 50 , 40 ] and information-theoretic 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 [ 56 ] , 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 [ 11 , 13 , 1 ] as well as the sample efficiency [ 38 , 68 , 16 ] , with some results in the deep learning setting [ 21 , 51 ] . In this paper , we focus on 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 Q⇤-sufficiency criteria to evaluate representations learned via MIbased objectives . Alternative approaches to representation learning for RL include priors based on the structure of the physical world [ 31 ] or heuristics such as disentanglement [ 67 ] , meta-learning general value functions [ 69 ] , predicting multiple value functions [ 10 , 17 , 29 ] and predicting domainspecific measurements [ 47 , 15 ] . We restrict our analysis to objectives that can be expressed as MI-maximization . In our paper we focus on the representation learning problem , disentangled from exploration , a strategy shared by prior works [ 19 , 35 , 64 , 72 ] . 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 , 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 a 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 , T , r ) , where S is the set of states , A the set of actions , T : S⇥A⇥S ! [ 0 , 1 ] the state transition distribution , and r : S ! R the reward function 2 . We will use capital letters to refer to random variables and lower case letters to refer to values of those variables ( e.g. , S is the random variable for the state and s is a specific state ) . Throughout our analysis we will often be interested in multiple reward functions , and denote a 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 r , and we denote this optimal policy as ⇡ ⇤ r = arg max ⇡ E⇡ [ P t t r ( St , At ) ] for discount factor . We also define the optimal Q-function as Q ⇤ r ( st , at ) = E⇡⇤ [ P1 t=1 t r ( St , At ) |st , at ] . The optimal Q-function satisfies the recursive Bellman equation , Q⇤ r ( st , at ) = r ( st , at ) + Ep ( st+1|st , at ) arg maxat+1 Q ⇤ r ( st+1 , at+1 ) . An optimal policy and the optimal Q-function are related according to ⇡⇤ ( s ) = arg maxa Q⇤ ( s , a ) . Mutual information . In information theory , the mutual information ( MI ) between two random variables , X and Y , is defined as [ 12 ] : I ( X ; Y ) = Ep ( x , y ) log p ( x , y ) p ( x ) p ( y ) = H ( X ) H ( X|Y ) . ( 1 ) 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 RL . While state aggregation methods typically define deterministic rules to group states in the representation [ 6 , 42 ] , MI-based representation learning methods used for deep RL treat the representation as a random variable [ 51 , 52 , 53 ] . Accordingly , we formalize a representation as a stochastic mapping between original state space and representation space . Definition 1 . A stochastic representation Z ( s ) is a mapping from states s 2 S to a probability distribution p ( Z|S = s ) over elements of a new representation space z 2 Z . In this work we 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 J = { Z | Z 2 arg max J ( ) . Unlike problem formulations for partially observed settings [ 70 , 26 , 39 ] , we assume that S is a Markovian state ; therefore the representation for a given state is conditionally independent of the past states , a common assumption in the state aggregation literature [ 6 , 42 ] . See Figure 1 for a depiction of the graphical model . 3.2 Sufficient Representations for Reinforcement Learning . We now turn to the problem of evaluating stochastic representations for RL . Intuitively , we expect a useful state representation to be capable of representing an optimal policy in the original state space . Definition 2 . A representation Z is ⇡⇤-sufficient with respect to a set of reward functions R if 8r 2 R , Z ( s1 ) = Z ( s2 ) = ) ⇡⇤r ( A|s1 ) = ⇡⇤r ( A|s2 ) . When a stochastic representation Z produces the same distribution over the representation space for two different states s1 and s2 we say it aliases these states . Unfortunately , as already proven in Theorem 4 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 2We restrict our attention to MDPs where the reward can be expressed as a function of the state , which is fairly standard across a broad set of real world RL problems 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⇤-sufficient with respect to a set of reward functions R if 8r 2 R , Z ( s1 ) = Z ( s2 ) = ) 8a , Q⇤r ( a , s1 ) = Q⇤r ( a , s2 ) . Note that Q⇤-sufficiency implies ⇡⇤-sufficiency since an optimal policy can be recovered from the optimal Q-function via ⇡⇤ r ( s ) = arg max a Q ⇤ r ( s , a ) [ 66 ] ; 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 [ 38 , 16 ] . We will further say that an objective J is sufficient with respect to some set of reward functions R if all the representations that maximize that objective J 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 .

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