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 proposes an algorithm learning a kind of universal value function (UVFA), so that given a buffer with a good coverage of transitions it returns the optimal behavior for any reward function or any goal state by computing a simple average.  The idea comes from the fact that the Q-function for a given policy $\pi$ -- that induces a reaching probability from any ($s,a$) $M^\pi(s, a, s', a')$ -- can be expressed as the scalar produce $Q^\pi(s,a) = \sum_{s',a'} M^\pi(s,a, s', a') r(s',a') = \langle M(s,a); r\rangle$. Then, if we learn $F(s,a,z)$ and $B(z)$ such that $M^{\pi_z}(s,a) = F(s,a, z)^T B(z)$ we obtain $Q^\pi_z = \langle F(s,a, z)^T B(z); r\rangle = F(s,a, z_r)^T z_r$ where $z_r = \langle B(z); r\rangle$.  The greedy policy given such a Q-function should maximize (over action) $F(s,a, z_r)^T z_r$. Now, there is a fix point {$F^*, B^*, \pi^*$} such that for all $z$, $Q^{\pi^*}(s, a) = F^*(s,a, z)^T z$, which coincides with the optimal Q-function for the given $z$.   The paper provides an algorithm that learns these optimal $F^*$ and $B^*$. Then, given any reward function, they compute $z_r = \langle B^*(z); r\rangle$ and return the optimal Q-function $Q^*(s,a) = F^*(s,a, z_r)^T z_r$.