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