Non-Linear Rewards For Successor Features

Recently , Reinforcement Learning ( RL ) algorithms have achieved superhuman performance in several challenging domains , such as Atari ( Mnih et al. , 2015 ) , Go ( Silver et al. , 2016 ) , and Starcraft II ( Vinyals et al. , 2019 ) . The main driver of these successes has been the use of deep neural networks , which are a class of powerful non-linear function approximators , with RL algorithms ( LeCun et al. , 2015 ) . However , this class of Deep Reinforcement Learning ( Deep RL ) algorithms require immense amounts of data within an environment , often ranging from tens to hundreds of millions of samples ( Arulkumaran et al. , 2017 ) . Furthermore , commonly used algorithms often have difficulty in transferring a learned policy between related tasks , such as where the environmental dynamics remain constant , but the goal changes . In this case , the model must either be retrained completely or fine-tuned on the new task , in both cases requiring millions of additional samples . If the state dynamics are constant , but the reward structure varies between tasks , it is wasteful to retrain the entire model . A more pragmatic approach would be to decompose the RL agent ’ s policy such that separate functions can learn the state dynamics and the reward structure ; doing so enables reuse of the dynamics model and only requires learning the reward component . Successor features ( Dayan , 1993 ) do precisely this ; a model-free policy ’ s action-value function is expressed as the dot product between a vector of expected discounted future state occupancies , the successor features , and another vector representing the immediate reward in each of those successor states . The factorization follows from the assumption that reward can be predicted as the dot product between a state representation vector and a learned reward vector . Therefore , transfer to a new task requires relearning only the reward parameters instead of the entire model and amounts to the supervised learning problem of predicting the current state ’ s immediate reward . This factorization can be limiting because it is assumed that the reward is a linear function of the current state , which might not always be the case as the encoded features might not capture the required quantity for accurate reward modelling ( Eysenbach et al. , 2018 ; Hansen et al. , 2019 ) . Therefore , this paper introduces a new form for the reward function : non-linear with respect to the current state . We assume that the learned features are not optimal and the reward can not be predicted directly from the raw features , which is not a strong assumption . This form increases the reward function ’ s representational power and makes it possible to incorporate the current state into reward estimation ; lessening the burden on the encoder components . Under the new reward formulation , a secondary term emerges , which learns the future expected auto-correlation matrix of the state features . This new secondary term , referred to as Λ , can be exploited as a possible avenue for directed exploration . Exploring the environment using Λ allows us to exploit and reuse learned environmental knowledge instead of relying on a purely random approach for exploration , such as -greedy . Following this , the contributions of this research are as follows : • A novel formulation of successor features that uses a non-linear reward function . This formulation increases the representational power of the reward function . • Under the new reward formulation , a second term appears that models the future expected auto-correlation matrix of the state features . • We provide preliminary results that show the second term can be used for guided exploration during transfer instead of relying on -greedy exploration . After the introduction of relevant background material in Section 1 , we introduce the successor feature framework with a non-linear reward function in Section 2 , Section 3 provides experimental support and provides an analysis of the new term in the decomposition . The paper concludes with a final discussion and possible avenues for future work in Section 4 . 1 BACKGROUND . 1.1 REINFORCEMENT LEARNING . Consider the interaction between an agent and an environment modelled by a Markov decision process ( MDP ) ( Puterman , 2014 ) . An MDP is defined as a set of states S , a set of actionsA , a reward function R : S → R , a discount factor γ ∈ [ 0 , 1 ] , and a transition function T : S ×A → [ 0 , 1 ] . The transition function gives the next-state distribution upon taking action a in state s and is often referred to as the dynamics of the MDP . The objective of the agent in RL is to find a policy π , a mapping from states to actions , which maximizes the expected discounted sum of rewards within the environment . One solution to this problem is to rely on learning a value function , where the action-value function of a policy π is defined as : Qπ ( s , a ) = Eπ [ ∞∑ t=0 γtR ( st ) |St = s , At = a ] where Eπ [ . . . ] denotes the expected value when following the policy π . The policy is learned using an alternating process of policy evaluation , given the action-value of a particular policy and policy improvement , which derives a new policy that is greedy with respect to Qπ ( s , a ) ( Puterman , 2014 ) . 1.2 SUCCESSOR FEATURES . Successor Features ( SF ) offer a decomposition of the Q-value function and have been mentioned under various names and interpretations ( Dayan , 1993 ; Kulkarni et al. , 2016 ; Barreto et al. , 2017 ; Machado et al. , 2017 ) . This decomposition follows from the assumption that the reward function can be approximately represented as a linear combination of learned features φ ( s ; θφ ) extracted by a neural network with parameters θφ and a reward weight vector w. As such , the expected one-step reward can be computed as : r ( s , a ) = φ ( s ; θφ ) > w . Following from this , the Q function can be rewritten as : Q ( s , a ) ≈ Eπ [ rt+1 + γrt+2 + . . . |St = s , At = a ] = Eπ [ φ ( st+1 ; θφ ) > w + φ ( st+2 ; θφ ) > w + . . . |St = s , At = a ] Q ( s , a ) = ψπ ( s , a ) > · w where ψ ( s , a ) are referred to as the successor features under policy π . The ith component of ψ ( s , a ) provides the expected discounted sum of φ ( i ) t when following policy π starting from state s and action a . It is assumed that the features φ ( s ; θφ ) are representative of the state s , such that ψ ( . ) can be turned into a function ψπ ( φ ( st ; θφ ) , at ) . For brevity , φ ( st ; θφ ) is referred to simply as φt and ψπ ( s , a ) as ψ ( s , a ) . The decomposition neatly separates the Q-function into two learning problems , for ψπ and w : estimating the features under the current policy dynamics , and estimating the reward given a state . Because the decomposition still has the same form as the Q-function , the successor features are computed using a Bellman equation update in which the reward function is replaced by φt : ψπ ( φt , at ) = φt + γE [ ψπ ( φt+1 , at+1 ) ] such that approximate successor features can be learned using an RL method , such as Q-Learning ( Szepesvári , 2009 ) . Following from this , the approximation of the reward vector w becomes a supervised learning problem . Often , this weight is learned using ordinary least squares from the sampled environmental data . One benefit of having a decoupled representation is that only the relevant function must be relearned when either the dynamics or the reward changes . Therefore , if the task changes , but the environmental dynamics remain constant , only the reward vector parameters w must be relearned , which are minimal compared to the total number of parameters in the full model . 2 MODEL , ARCHITECTURE , AND TRAINING . The Successor Feature framework has several limitations , primarily stemming from the assumptions around its derivation , such as constant environmental structure between tasks or that the reward can be linearly predicted from state features . Work towards solving the former has been developed by Zhang et al . ( 2017 ) whereby they learn a linear mapping between task state features . The latter assumption , whereby the reward is assumed to be a linear mapping of state features , is not guaranteed and , as we show , the Successor Feature framework fails in such cases . Therefore , the method presented in this section aims to provide a stronger guarantee of the framework ’ s performance in such cases by developing a more robust reward component . This section discusses our change to the successor feature framework , which adjusts reward function , from a linear function , to a non-linear function . First , a discussion of the new decomposition is given with the full derivation provided in Appendix A . Then experimental support for this change will be presented and analyzed to examine what the new term in the decomposition learns . 2.1 NON-LINEAR REWARD FUNCTION . The successor feature framework builds upon the assumption that the current reward rt can be represented by the linear combination of the current state representation φt ∈ Rz and a learned reward vector w ∈ Rz , such that rt = φ > t w. This form is limiting because there is no guarantee that the reward will be a linear combination of the state features or that the required state features can be learned by the encoder ( Eysenbach et al. , 2018 ; Hansen et al. , 2019 ) . In practice the optimal state features are often not learned ; therefore , we build on the basis that the state features are sub-optimal , which in itself is not a strong assumption . To increase the flexibility of this reward model , let us consider the following form : rt = φ > t o + φ > t Aφt ( 1 ) where { φt , o } ∈ Rz , and A ∈ Rz×z . Both o and A are learnable parameters modelling the reward structure of the environment . Equation 1 shows that the formulation introduces a non-linear transformation with respect to φ . Comparing this with the original formulation , it is evident that this is equivalent to setting w = o + Aφ . The state-action value function Q ( s , a ) , under this new reward structure , can be derived to yield : Qπ ( st , a ) = ψ π ( st , a ) > o + βtr ( AΛπ ( st , a ) ) ( 2 ) where β ∈ { 0 , 1 } controls the inclusion of Λ and tr is the trace operator . It can now be shown that ψ and Λ satisfy the Bellman equation ( Bellman , 1966 ) : ψπ ( st , a ) = Eπ [ φt+1 + γψ ( st+1 , π ( st+1 ) ) |St = s , At = a ] ( 3 ) Λπ ( st , a ) = Eπ [ φt+1φ > t+1 + γΛ ( st+1 , π ( st+1 ) ) |St = s , At = a ] ( 4 ) where for ψ and Λ , φ and φφ > respectively play the role of rewards . In addition to ψ , it is now necessary to model Λ , which outputs an Rz×z matrix per action . The quantity φtφ > t can be interpreted as an auto-correlation matrix of the state features . We can see that this form allows the Λ term to model some form of future expected stochasticity of the environment . For example , the diagonal of Λ will model a second order moment capturing each feature ’ s change with respect to itself φ1 . We provide analysis and further discussion of Λ in Section 3.5 .

Summary: Successor representations are an old idea that has seem recent interest in the ML community. The idea is conceptually straightforward, by assuming the rewards are linear in some space $r = \vect{\phi}(s, a) \cdot \vect{w}$ then we learn something analogous to an action-value function for the discounted expected features under a policy so that the action-value on task $\vect{w}$ is $Q^\pi_{\vect{w}}(s, a) = \psi^\pi(s, a) \cdot \vect{w}$. This allows computing the action value for the policy under a new task $\vect{w}'$ straight.