Xi-learning: Successor Feature Transfer Learning for General Reward Functions

1 INTRODUCTION . Reinforcement Learning ( RL ) successfully addressed many complex problems such as playing computer games , chess , and even Go with superhuman performance ( Mnih et al. , 2015 ; Silver et al. , 2018 ) . These impressive results are possible thanks to a vast amount of interactions of the RL agent with its environment/task . Such strategy is unsuitable in settings where the agent has to perform and learn at the same time . Consider , for example , a care giver robot in a hospital that has to learn a new task , such as a new route to deliver meals . In such a setting , the agent can not collect a vast amount of training samples but has to adapt quickly instead . Transfer learning aims to provide mechanisms quickly to adapt agents in such settings ( Taylor and Stone , 2009 ; Lazaric , 2012 ; Zhu et al. , 2020 ) . The rationale is to use knowledge from previously encountered source tasks for a new target task to improve the learning performance on the target task . The previous knowledge can help reducing the amount of interactions required to learn the new optimal behavior . For example , the care giver robot could reuse knowledge about the layout of the hospital it learned in previous source tasks ( e.g . guiding a person ) to learn to deliver meals . The Successor Feature ( SF ) and General Policy Improvement ( GPI ) framework ( Barreto et al. , 2020 ) is a prominent transfer learning mechanism for tasks where only the reward function differs . Its basic premise is that the rewards which the RL agent tries to maximize are defined based on a lowdimensional feature descriptor φ ∈ Rn . For our care-giver robot this could be ID ’ s of beds or rooms that it is visiting , in difference to its high-dimensional visual state input from a camera . The rewards are then computed not based on its visual input but on the ID ’ s of the beds or rooms that it visits . The expected cumulative discounted successor features ( ψ ) are learned for each behavior that the robot learned in the past . It represents the dynamics in the feature space that the agent experiences for a behavior . This corresponds to the rooms or beds the care-giver agent would visit if using the behavior . This representation of feature dynamics is independent from the reward function . A behavior learned in a previous task and described by this SF representation can be directly re-evaluated for a different reward function . In a new task , i.e . for a new reward function , the GPI procedure re-evaluates the behaviors learned in previous tasks for it . It then selects at each state the behavior of a previous task if it improves the expected reward . This allows to reuse behaviors learned in previous source tasks Source code at https : //tinyurl.com/3xuzxff3 for a new target task . A similar transfer strategy can also be observed in the behavior of humans ( Momennejad et al. , 2017 ; Momennejad , 2020 ; Tomov et al. , 2021 ) . The classical SF & GPI framework ( Barreto et al. , 2017 ; 2018 ) makes the assumption that rewards r are a linear composition of the features φ ∈ Rn via a reward weight vector wi ∈ Rn that depends on the task i : ri = φ > wi . This assumption allows to effectively separate the feature dynamics of a behavior from the rewards and thus to re-evaluate previous behaviors given a new reward function , i.e . a new weight vector wj . Nonetheless , this assumption also restricts successful application of SF & GPI only to problems where such a linear decomposition is possible . We investigate the application of the SF & GPI framework to general reward functions : ri = Ri ( φ ) over the feature space . We propose to learn the cumulative discounted probability over the successor features , named ξ-function , and refer to the proposed framework as ξ-learning . Our work is related to Janner et al . ( 2020 ) ; Touati and Ollivier ( 2021 ) , and brings two important additional contributions . First , we provide mathematical proof of the convergence of ξ-learning . Second , we demonstrate how ξ-learning can be used for meta-RL , using the ξ-function to re-evaluate behaviors learned in previous tasks for a new reward function Rj . Furthermore , ξ-learning can also be used to transfer knowledge to new tasks using GPI . The contribution of our paper is three-fold : • We introduce a new RL algorithm , ξ-learning , based on a cumulative discounted probability of successor features , and two variants of its update operator . • We provide theoretical proofs of the convergence of ξ-learning to the optimal policy and for a guarantee of its transfer learning performance under the GPI procedure . • We experimentally compare ξ-learning in tasks with linear and general reward functions , and for tasks with discrete and continuous features to standard Q-learning and the classical SF framework , demonstrating the interest and advantage of ξ-learning . 2 BACKGROUND . 2.1 REINFORCEMENT LEARNING . RL investigates algorithms to solve multi-step decision problems , aiming to maximize the sum over future rewards ( Sutton and Barto , 2018 ) . RL problems are modeled as Markov Decision Processes ( MDPs ) which are defined as a tuple M ≡ ( S , A , p , R , γ ) , where S and A are the state and action set . An agent transitions from a state st to another state st+1 using action at at time point t collecting a reward rt : st at , rt−−−→ st+1 . This process is stochastic and the transition probability p ( st+1|st , at ) describes which state st+1 is reached . The reward function R defines the scalar reward rt = R ( st , at , st+1 ) ∈ R for the transition . The goal in an MDP is to maximize the expected return Gt = E [ ∑∞ k=0 γ kRt+k ] , where Rt = R ( St , At , St+1 ) . The discount factor γ ∈ [ 0 , 1 ) weights collected rewards by discounting future rewards stronger . RL provides algorithms to learn a policy π : S → A defining which action to take in which state to maximise Gt . Value-based RL methods use the concept of value functions to learn the optimal policy . The stateaction value function , called Q-function , is defined as the expected future return taking action at in st and then following policy π : Qπ ( st , at ) = Eπ { rt + γrt+1 + γ 2rt+2 + . . . } = Eπ { rt + γmax at+1 Qπ ( St+1 , at+1 ) } . ( 1 ) The Q-function can be recursively defined following the Bellman equation such that the current Q-value Qπ ( st , at ) depends on the maximum Q-value of the next state Qπ ( st+1 , at+1 ) . The optimal policy for an MDP can then be expressed based on the Q-function , by taking at every step the maximum action : π∗ ( s ) ∈ argmaxaQ∗ ( s , a ) . The optimal Q-function can be learned using a temporal difference method such as Q-learning ( Watkins and Dayan , 1992 ) . Given a transition ( st , at , rt , st+1 ) , the Q-value is updated according to : Qk+1 ( st , at ) = Qk ( st , at ) + αk ( rt + max at+1 Qk ( st+1 , at+1 ) −Qk ( st , at ) ) , ( 2 ) where αk ∈ ( 0 , 1 ] is the learning rate at iteration k . 2.2 TRANSFER LEARNING AND THE SF & GPI FRAMEWORK . We are interested in the transfer learning setting where the agent has to solve a set of tasksM = { M1 , M2 , . . . , Mm } , that in our case differ only in their reward function . The Successor Feature ( SF ) framework provides a principled way to perform transfer learning ( Barreto et al. , 2017 ; 2018 ) . SF assumes that the reward function can be decomposed into a linear combination of features φ ∈ Φ ⊂ Rn and a reward weight vector wi ∈ Rn that is defined for a task Mi : ri ( st , at , st+1 ) ≡ φ ( st , at , st+1 ) > wi . ( 3 ) We refer to such reward functions as linear reward functions . Since the various tasks differ only in their reward functions , the features are the same for all tasks inM . Given the decomposition above , it is also possible to rewrite the Q-function into an expected discounted sum over future features ψπi ( s , a ) and the reward weight vector wi : Qπii ( s , a ) = E { rt + γ 1rt+1 + γ 2rt+2 + . . . } = E { φ > t wi + γ 1φ > t+1wi + γ 2φ > t+2wi + . . . } = E { ∑∞ k=0 γ kφt+k } > wi ≡ ψπi ( s , a ) > wi . ( 4 ) This decouples the dynamics of the policy πi in the feature space of the MDP from the expected rewards for such features . Thus , it is now possible to evaluate the policy πi in a different taskMj using a simple multiplication of the weight vector wj with the ψ-function : Qπij ( s , a ) = ψ πi ( s , a ) > wj . Interestingly , the ψ function also follows the Bellman equation : ψπ ( s , a ) = E { φt+1 + γψπ ( st+1 , π ( st+1 ) ) |st , at } , ( 5 ) and can therefore be learned with conventional RL methods . Moreover , ( Lehnert and Littman , 2019 ) showed the equivalence of SF-learning to Q-learning . Being in a new task Mj the Generalized Policy Improvement ( GPI ) can be used to select the action over all policies learned so far that behaves best : π ( s ) ∈ argmax a max i Qπij ( s , a ) = argmax a max i ψπi ( s , a ) > wj . ( 6 ) ( Barreto et al. , 2018 ) proved that under the appropriate conditions for optimal policy approximates , the policy constructed in ( 6 ) is close to the optimal one , and their difference is upper-bounded : ||Q∗ −Qπ||∞ ≤ 2 1− γ ( ||r − ri||∞ + min j ||ri − rj ||∞ + ) , ( 7 ) where ‖f − g‖∞ = maxs , a |f ( s , a ) − g ( s , a ) | . For an arbitrary reward function r the result can be interpreted in the following manner . Given the arbitrary task M , we identify the theoretically closest possible linear reward task Mi with ri . For this theoretically closest task , we search the linear task Mj in our set of taskM ( from which we also construct the GPI optimal policy ( 6 ) ) which is closest to it . The upper bound between Q∗ and Q is then defined by 1 ) the difference between task M and the theoretically closest possible linear task Mi : ||r− ri||∞ ; and by 2 ) the difference between theoretical task Mi and the closest task Mj : minj ||ri − rj ||∞ . If our new task M is also linear then r = ri and the first term in ( 7 ) would vanish . Very importantly , this result shows that the SF framework will only provide a good approximation of the true Q-function if the reward function in a task can be represented using a linear decomposition . If this is not the case then the error in the approximation increases with the distance between the true reward function r and the best linear approximation of it ri as stated by ||r − ri||∞ . 3 METHOD : ξ-LEARNING 3.1 DEFINITION AND FOUNDATIONS OF ξ-LEARNING The goal of this paper is to investigate the application of SF & GPI to tasks with general reward functions R : Φ 7→ R over state features φ ∈ Φ : r ( st , at , st+1 ) ≡ R ( φ ( st , at , st+1 ) ) = R ( φt ) , ( 8 ) where we define φt ≡ φ ( st , at , st+1 ) . Under this assumption the Q-function can not be linearly decomposed into a part that describes feature dynamics and one that describes the rewards as in the linear SF framework ( 4 ) . To overcome this issue , we propose to define the expected cumulative discounted probability of successor features or ξ-function , which is going to be the central mathematical object of the paper , as : ξπ ( s , a , φ ) = ∞∑ k=0 γkp ( φt+k = φ|st = s , at = a ; π ) , ( 9 ) where p ( φt+k = φ|st = s , at = a ; π ) , or in short p ( φt+k = φ|st , at ; π ) , is the probability density function of the features at time t+ k , following policy π and conditioned to s and a being the state and action at time t respectively . Note that ξπ depends not only on the policy π but also on the state transition ( constant through the paper ) . With the definition of the ξ-function , the Q-function rewrites ( this is compatible with SFQL in the linear reward case , see Appendix A.6 ) : Qπ ( st , at ) = ∞∑ k=0 γkEp ( φt+k|st , at ; π ) { R ( φt+k ) } = ∞∑ k=0 γk ∫ Φ p ( φt+k = φ|st , at ; π ) R ( φ ) dφ = ∫ Φ R ( φ ) ∞∑ k=0 γkp ( φt+k = φ|st , at ; π ) dφ = ∫ Φ R ( φ ) ξπ ( st , at , φ ) dφ . ( 10 ) Depending on the reward function R , there are several ξ-functions that correspond to the same Q function . Formally , this is an equivalence relationship , and the quotient space has a one-to-one correspondence with the Q-function space . Proposition 1 . ( Equivalence between functions ξ and Q ) Let Q = { Q : S ×A → R s.t . ‖Q‖∞ < ∞ } . Let ∼ be defined as ξ1 ∼ ξ2 ⇔ ∫ Φ Rξ1 = ∫ Φ Rξ2 . Then , ∼ is an equivalence relationship , and there is a bijective correspondence between the quotient space Ξ∼ and Q. Corollary 1 . The bijection between Ξ∼ and Q allows to induce a norm ‖ · ‖∼ into Ξ∼ from the supremum norm in Q , with which Ξ∼ is a Banach space ( since Q is Banach with ‖ · ‖∞ ) : ‖ξ‖∼ = sup s , a ∣∣∣∣∫ Φ R ( φ ) ξ ( s , a , φ ) dφ ∣∣∣∣ = sup s , a |Q ( s , a ) | = ‖Q‖∞ . ( 11 ) Similar to the Bellman equation for the Q-function , we can define a Bellman operator for the ξ-function , denoted by Tξ , as : Tξ ( ξ π ) = p ( φt = φ|st , at ) + γEp ( st+1 , at+1|st , at ; π ) { ξ π ( st+1 , at+1 , φ ) } . ( 12 ) As in the case of the Q-function , we can use Tξ to construct a contractive operator : Proposition 2 . ( ξ-learning has a fixed point ) The operator Tξ is well-defined w.r.t . the equivalence ∼ , and therefore induces an operator T∼ defined over Ξ∼ . T∼ is contractive w.r.t . ‖ · ‖∼ . Since Ξ∼ is Banach , T∼ has a unique fixed point and iterating T∼ starting anywhere converges to that point . In other words , successive applications of the operator T∼ converge towards the class of optimal ξ functions [ ξ∗ ] or equivalently to an optimal ξ function defined up to an additive function k satisfying∫ Φ k ( s , a , φ ) R ( φ ) dφ = 0 , ∀ ( s , a ) ∈ S ×A ( i.e . k ∈ Ker ( ξ → ∫ Φ Rξ ) ) . While these two results state ( see Appendix A for the proofs ) the theoretical links to standard Q-learning formulations , the Tξ operator defined in ( 12 ) is not usable in practice , because of the expectation . In the next section , we define the optimisation iterate , prove its convergence , and provide two variants to perform the ξ updates . 3.2 ξ-LEARNING ALGORITHMS In order to learn the ξ-function , we introduce the ξ-learning update operator , which is an offpolicy temporal difference method analogous to Q-learning . Given a transition ( st , at , st+1 , φt ) the ξ-learning update operator is defined as : ξπk+1 ( st , at , φ ) ← ξπk ( st , at , φ ) + αk [ p ( φt = φ|st , at ) + γξπk ( st+1 , āt+1 , φ ) − ξπk ( st , at , φ ) ] , ( 13 ) where āt+1 = argmaxa ∫ Φ R ( φ ) ξπ ( st+1 , a , φ ) dφ . The following is one of the main results of the manuscript , stating the convergence of ξ-learning : Theorem 1 . ( Convergence of ξ-learning ) For a sequence of state-action-feature { st , at , st+1 , φt } ∞t=0 consider the ξ-learning update given in ( 13 ) . If the sequence of state-action-feature triples visits each state , action infinitely often , and if the learning rate αk is an adapted sequence satisfying the Robbins-Monro conditions : ∞∑ k=1 αk =∞ , ∞∑ k=1 α2k < ∞ ( 14 ) then the sequence of function classes corresponding to the iterates converges to the optimum , which corresponds to the optimal Q-function to which standard Q-learning updates would converge to : [ ξn ] → [ ξ∗ ] with Q∗ ( s , a ) = ∫ Φ R ( φ ) ξ∗ ( s , a , φ ) dφ . ( 15 ) The proof is provided in Appendix A and follows the same flow as for Q-learning . The previous theorem provides convergence guarantees under the assumption that either p ( φt = φ|st , at ; π ) is known , or an unbiased estimate can be constructed . We propose two different ways to approximate p ( φt = φ|st , at ; π ) from a given transition ( st , at , st+1 , φt ) so as to perform the ξ-update ( 13 ) . The first instance is a model-free version and detailed in the following section . A second instance uses a one-step SF model , called One-Step Model-based ( MB ) ξ-learning , which is further described in Sec . B. Model-free ( MF ) ξ-Learning : MF ξ-learning uses the same principle as standard model-free temporal difference learning methods . The update assumes for a given transition ( st , at , st+1 , φt ) that the probability for the observed feature is p ( φ = φt|st , at ) = 1 . Whereas for all other features ( ∀φ′ ∈ Φ , φ′ 6= φt ) the probability is p ( φ′ = φt|st , at ) = 0 , see Appendix D for continuous features . The resulting updates are : φ = φt : ξ π ( st , at , φ ) ← ( 1− α ) ξπ ( st , at , φ ) + α ( 1 + γξπ ( st+1 , āt+1 , φ ) ) φ′ 6= φt : ξπ ( st , at , φ′ ) ← ( 1− α ) ξπ ( st , at , φ′ ) + αγξπ ( st+1 , āt+1 , φ′ ) . ( 16 ) Due to the stochastic update of the ξ-function and if the learning rate α ∈ ( 0 , 1 ] discounts over time , the ξ-update will learn the true probability of p ( φ = φt|st , at ) . A potential problem with the MF procedure is that it might induce a high variance when the true feature probabilities are not binary . 3.3 META ξ-LEARNING After discussing ξ-learning on a single task and showing its theoretical convergence , we can now investigate how it can be applied in transfer learning . Similar to the linear SF framework the ξ-function allows to reevaluate a policy learned for task Mi , ξπi , in a new environment Mj : Qπij ( s , a ) = ∫ Φ Rj ( φ ) ξ πi ( s , a , φ ) dφ . ( 17 ) This allows us to apply GPI in ( 6 ) for arbitrary reward functions in a similar manner to what was proposed for linear reward functions in ( Barreto et al. , 2018 ) . We extend the GPI result to the ξ-learning framework as follows : Theorem 2 . ( Generalised policy improvement in ξ-learning ) LetM be the set of tasks , each one associated to a ( possibly different ) weighting function Ri ∈ L1 ( Φ ) . Let ξπ ∗ i be a representative of the optimal class of ξ-functions for task Mi , i ∈ { 1 , . . . , I } , and let ξ̃πi be an approximation to the optimal ξ-function , ‖ξπ∗i − ξ̃πi‖Ri ≤ ε , ∀i . Then , for another task M with weighting function R , the policy defined as : π ( s ) = arg max a max i ∫ Φ R ( φ ) ξ̃πi ( s , a , φ ) dφ , ( 18 ) satisfies : ‖ξ∗ − ξπ‖R ≤ 2 1− γ ( min i ‖R−Ri‖p ( φ|s , a ) + ε ) , ( 19 ) where ‖f‖g = sups , a ∫ Φ |f · g| dφ . The proof is provided in Appendix A .

Summary: The authors extend the successor features (SF) linear formulation to general non-linear formulation. The new method allows the reward to be arbitrary composition of features.   The propose method is based on learning cumulative discounted probability of SF rather than cumulative discounted sum of SF as in linear SF framework.  The authors provide theoretical proofs about convergence and propose two practical methods. Experiments are conducted in two toy environments.