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: This paper addresses a limitation of the successor features (SF) which were introduced as a mechanism for transfer learning in reinforcement learning when the reward function changes across tasks. The authors claim that the original 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, and propose to address this linear decomposition of the reward function requirement by introducing a novel SF mechanism, Xi-learning, based on learning a cumulative discounted probability of successor features. The paper includes theoretical proofs of the convergence of xi-learning as well as transfer learning guarantees under generalized policy improvement (GPI). The authors compare the performance of their methods against standard Q-learning and SFQL in two different domains.