Learning One Representation to Optimize All Rewards

1 Introduction . We consider one kind of unsupervised reinforcement learning problem : Given a Markov decision process ( MDP ) but no reward information , is it possible to learn and store a compact object that , for any reward function specified later , provides the optimal policy for that reward , with a minimal amount of additional computation ? In a sense , such an object would encode in a compact form the solutions of all possible planning problems in the environment . This is a step towards building agents that are fully controllable after first exploring their environment in an unsupervised way . Goal-oriented RL methods [ ACR+17 , PAR+18 ] compute policies for a series of rewards specified in advance ( such as reaching a set of target states ) , but can not adapt in real time to new rewards , such as weighted combinations of target states or dense rewards . Learning a model of the world is another possibility , but it still requires explicit planning for each new reward ; moreover , synthesizing accurate trajectories of states over long time ranges has proven difficult [ Tal17 , KST+18 ] . ⇤Work done during an internship at Facebook Artificial Intelligence Research Paris . 2Code : https : //github.com/ahmed-touati/controllable_agent 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Instead , we exhibit an object that is both simpler to learn than a model of the world , and contains the information to recover near-optimal policies for any reward provided a posteriori , without a planning phase . [ BBQ+18 ] learn optimal policies for all rewards that are linear combinations of a finite number of feature functions provided in advance by the user . This limits applications : e.g. , goal-oriented tasks would require one feature per goal state , thus using infinitely many features in continuous spaces . We reuse a policy parameterization from [ BBQ+18 ] , but introduce a novel representation with better properties , based on state occupancy prediction instead of expected featurizations . We use theoretical advances on successor state learning from [ BTO21 ] . We obtain the following . • We prove the existence of a learnable “ summary ” of a reward-free discrete or continuous MDP , that provides an explicit formula for optimal policies for any reward specified later . This takes the form of a pair of representations F : S ⇥A⇥ Z ! Z and B : S ⇥A ! Z from state-actions into a representation space Z ' Rd , with policies ⇡z ( s ) : = argmaxa F ( s , a , z ) > z . Once a reward is specified , a value of z is computed from reward values and B ; then ⇡z is used . Rewards may be specified either explicitly as a function , or as target states , or by samples as in usual RL setups . • We provide a well-principled unsupervised loss for F and B . If FB training is perfect , then the policies are provably optimal for all rewards ( Theorem 2 ) . With imperfect training , sub-optimality is proportional to the FB training error ( Theorems 8–9 ) . In finite spaces , perfect training is possible with large enough dimension d ( Proposition 6 ) . Explicitly , F and B are trained so that F ( s , a , z ) > B ( s0 , a0 ) approximates the long-term probability to reach s0 from s if following ⇡z . This is akin to a model of the environment , without synthesizing state trajectories . • We provide a TD-like algorithm to train F and B for this unsupervised loss , with function approximation , adapted from recent methods for successor states [ BTO21 ] . No sparse rewards are used : every transition reaches some state s0 , so every step is exploited . As usual with TD , learning seeks a fixed point but the loss itself is not observable . • We prove viability of the method on several environments from mazes to pixel-based MsPacman and a virtual robotic arm . For single-state rewards ( learning to reach arbitrary states ) , we provide quantitative comparisons with goal-oriented methods such as HER . ( Our method is not a substitute for HER : in principle they could be combined , with HER improving replay buffer management for our method . ) For more general rewards , which can not be tackled a posteriori by trained goal-oriented models , we provide qualitative examples . • We also illustrate qualitatively the sub-optimalities ( long-range behavior is preserved but local blurring of rewards occurs ) and the representations learned . 2 Problem and Notation . Let M = ( S , A , P , ) be a reward-free Markov decision process with state space S ( discrete or continuous ) , action space A ( discrete for simplicity , but this is not essential ) , transition probabilities P ( s0|s , a ) from state s to s0 given action a , and discount factor 0 < < 1 [ SB18 ] . If S is finite , P ( s0|s , a ) can be viewed as a matrix ; in general , for each ( s , a ) 2 S⇥A , P ( ds0|s , a ) is a probability measure on s0 2 S. The notation P ( ds0|s , a ) covers all cases . Given ( s0 , a0 ) 2 S ⇥A and a policy ⇡ : S ! Prob ( A ) , we denote Pr ( ·|s0 , a0 , ⇡ ) and E [ ·|s0 , a0 , ⇡ ] the probabilities and expectations under state-action sequences ( st , at ) t 0 starting with ( s0 , a0 ) and following policy ⇡ in the environment , defined by sampling st ⇠ P ( dst|st 1 , at 1 ) and at ⇠ ⇡ ( st ) . For any policy ⇡ and state-action ( s0 , a0 ) , define the successor measure M⇡ ( s0 , a0 , · ) as the measure over S ⇥A representing the expected discounted time spent in each set X ⇢ S ⇥A : M⇡ ( s0 , a0 , X ) : = X t 0 t Pr ( ( st , at ) 2 X | s0 , a0 , ⇡ ) ( 1 ) for each X ⇢ S ⇥A . Viewing M as a measure deals with both discrete and continuous spaces . Given a reward function r : S ⇥ A ! R , the Q-function of ⇡ for r is Q⇡r ( s0 , a0 ) : =P t 0 t E [ r ( st , at ) |s0 , a0 , ⇡ ] . We assume that rewards are bounded , so that all Q-functions are well-defined . We state the results for deterministic reward functions , but this is not essential . We abuse notation and write greedy policies as ⇡ ( s ) = argmaxa Q ( s , a ) instead of ⇡ ( s ) 2 argmaxa Q ( s , a ) . Ties may be broken any way . We consider the following informal problem : Given a reward-free MDP ( S , A , P , ) , can we compute a convenient learnable object E such that , once a reward function r : S ⇥A ! R is specified , we can easily ( with no planning ) compute , from E and r , a policy ⇡ whose performance is close to maximal ? 3 Encoding All Optimal Policies via the Forward-Backward Representation . We first present forward-backward ( FB ) representations of a reward-free MDP as a way to summarize all optimal policies via explicit formulas . The resulting learning procedure is described in Section 4 . Core idea . The main algebraic idea is as follows . Assume , at first , that S is finite . For a fixed policy , the Q-function depends lineary on the reward : namely , Q⇡r ( s , a ) = P s0 , a0 M ⇡ ( s , a , s0 , a0 ) r ( s0 , a0 ) where M⇡ ( s , a , s0 , a0 ) = P t 0 t Pr ( ( st , at ) = ( s0 , a0 ) |s , a , ⇡ ) . This rewrites as Q⇡r = M⇡r viewing everything as vectors and matrices indexed by state-actions . Now let ( ⇡z ) z2Rd be any family of policies parameterized by z . Assume that for each z , we can find d⇥ ( S ⇥ A ) -matrices Fz and B such that M⇡z = F > zB . Then Q⇡zr = F > zBr . Specializing to zR : = Br , the Q-function of policy ⇡zR on reward r is Q ⇡zR r = F > zRzR . So far ⇡z was unspecified ; but if we define ⇡z ( s ) : = argmaxa ( F > z z ) sa at each state s , then by definition , ⇡zR is the greedy policy with respect to F > zRzR . At the same time , F > zRzR is the Q-function of ⇡zR for reward r : thus , ⇡zR is the greedy policy of its own Q-function , and is therefore optimal for reward r. Thus , if we manage to find F , B , and ⇡z such that ⇡z = argmaxF > z z and F > zB = M⇡z for all z 2 Rd , then we obtain the optimal policy for any reward r , just by computing Br and applying policy ⇡Br . This criterion on ( F , B , ⇡z ) is entirely unsupervised . Since F and B depend on ⇡z but ⇡z is defined via F , this is a fixed point equation . An exact solution exists for d large enough ( Appendix , Prop . 6 ) , while a smaller d provides lower-rank approximations M⇡z ⇡ F > zB . In Section 4 we present a well-grounded algorithm to learn such F , B , and ⇡z . In short , we learn two representations F and B such that F ( s0 , a0 , z ) > B ( s0 , a0 ) is approximately the long-term probability M⇡z ( s0 , a0 , s0 , a0 ) to reach ( s0 , a0 ) if starting at ( s0 , a0 ) and following policy ⇡z . Then all optimal policies can be computed from F and B . We think of F as a representation of the future of a state , and B as the ways to reach a state ( Appendix B.4 ) : if F > B is large , then the second state is reachable from the first . This is akin to a model of the environment , without synthesizing state trajectories . General statement . In continuous spaces with function approximation , Fz and B become functions S ⇥A ! Rd instead of matrices ; since Fz depends on z , F itself is a function S ⇥A⇥ Rd ! Rd . The sums over states will be replaced with expectations under the data distribution ⇢ . Definition 1 ( Forward-backward representation ) . Let Z = Rd be a representation space , and let ⇢ be a measure on S ⇥A . A pair of functions F : S ⇥A⇥Z ! Z and B : S ⇥A ! Z , together with a parametric family of policies ( ⇡z ) z2Z , is called a forward-backward representation of the MDP with respect to ⇢ , if the following conditions hold for any z 2 Z and ( s , a ) , ( s0 , a0 ) 2 S ⇥A : ⇡z ( s ) = argmax a F ( s , a , z ) > z , M⇡z ( s0 , a0 , ds , da ) = F ( s0 , a0 , z ) > B ( s , a ) ⇢ ( ds , da ) ( 2 ) where M⇡ is the successor measure defined in ( 1 ) , and the last equality is between measures . Theorem 2 ( FB representations encode all optimal policies ) . Let ( F , B , ( ⇡z ) ) be a forward-backward representation of a reward-free MDP with respect to some measure ⇢ . Then , for any bounded reward function r : S ⇥A ! R , the following holds . Set zR : = Z s , a r ( s , a ) B ( s , a ) ⇢ ( ds , da ) . ( 3 ) assuming the integral exists . Then ⇡zR is an optimal policy for reward r in the MDP . Moreover , the optimal Q-function Q ? for reward r is Q ? ( s , a ) = F ( s , a , zR ) > zR . For instance , for a single reward located at state-action ( s , a ) , the optimal policy is ⇡zR with zR = B ( s , a ) . ( In that case the factor ⇢ ( ds , da ) does not matter because scaling the reward does not change the optimal policy . ) We present in Section 4 an algorithm to learn FB representations . The measure ⇢ will be the distribution of state-actions visited in a training set or under an exploration policy : then zR = E ( s , a ) ⇠⇢ [ r ( s , a ) B ( s , a ) ] can be obtained by sampling from visited states . In finite spaces , exact FB representations exist , provided the dimension d is larger than # S ⇥ # A ( Appendix , Prop . 6 ) . In infinite spaces , arbitrarily good approximations can be obtained by increasing d , corresponding to a rank-d approximation of the cumulated transition probabilities M⇡ . Importantly , the optimality guarantee extends to approximate F and B , with optimality gap proportional to F > B M⇡z/⇢ ( Appendix , Theorems 8–9 with various norms on F > B M⇡/⇢ ) . For instance , if , for some reward r , the error F ( s0 , a0 , zR ) > B ( s , a ) M⇡zR ( s0 , a0 , ds , da ) /⇢ ( ds , da ) is at most `` on average over ( s , a ) ⇠ ⇢ for every ( s0 , a0 ) , then ⇡zR is 3 '' krk1 / ( 1 ) -optimal for r. These results justify using some norm over F > B M⇡z/⇢ , averaged over z 2 Rd , as a training loss for unsupervised reinforcement learning . ( Below , we average over z 2 Rd from a fixed rescaled Gaussian . If prior information is available on the rewards r , the corresponding distribution of zR may be used instead . ) If B is fixed in advance and only F is learned , the method has similar properties to successor features based on B ( Appendix B.4 ) . But one may set a large d and let B be learned : arguably , by Theorem 2 , the resulting features “ linearize ” optimal policies as much as possible . The features learned in F and B may have broader interest .

Summary: This paper considered the representation learning in reward-free MDPs to encode the optimal policy given any reward function without planning. The proposed unsupervised learning algorithm FB offered two features $F$ and $B$ that encoded the information of the transition to extract the optimal policy. They showed that as long as $F^\top B$ approximates the probability of visiting one state-action to another, then we can extract the optimal policy given any reward function directly. Motived by the nice theoretical properties of $F$ and $B$, they showed that the unsupervised learning algorithm can achieve the same level of learning compared with DQN, DQN with HER, etc. on several benchmark environments. They also incorporated some engineering tricks to make the theory-boosted algorithm more effective.