Second-Order Rewards For Successor Features

1 INTRODUCTION . 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 formulation of the reward as the dot product between a state representation vector and a learned parameter vector , that is a linear product . 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 . As the reward function of the successor feature framework is linear , it is fair to question whether the model can always accurately predict the reward . As no assumptions are made about the state representation , theoretically it is possible to enable perfect recovery of any reward function if given predictive state representation Barreto et al . ( 2017 ) . The state representation , within the successor feature framework , is learned end-to-end by a state encoder to perform well in state reconstruction and reward prediction tasks . The state encoder , given a large enough set of parameters , should have enough representational power to disentangle the factors that are useful for reward prediction by a linear model . However , because of how the encoder is trained , its parameters are utilized for both state reconstruction and reward prediction tasks ; while the reward model parameters are only used for reward prediction . If the encoder learns a sub-optimal state representation for reward prediction , say because of a highly complex environment , the reward model might be unable to compensate with its limited set of parameters correctly to predict the reward . Eysenbach et al . ( 2018 ) and Hansen et al . ( 2019 ) have shown , within the successor feature framework , that there is no strong guarantee that the state encoder is always able to learn features that enable accurate modelling of the reward . In this paper , a novel extension to the successor feature framework is proposed , where the rewards are modelled with a second-order function . The second-order function , which follows naturally from the original linear variant , gives a stronger guarantee on performance of the model due to both its representational structure and extra parameters . This is especially true in cases where the encoder learns state representations that are less than optimal where a linear model does not have enough representational power to compensate . While , in cases where the encoder can learn sufficient represents for both reconstruction and reward tasks , the second-order function still provides many added benefits . In particular , the additional parameters of the reward model should lessen the representation load of the encoder with regards to the reward tasks allowing more of its representational capacity to be dedicated towards modelling the environment in a task agnostic manner . Further benefits , via a new term emerging after derivation , include a representational form of environmental stochasticity and the ability to use directed exploration during transfer 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 second-order reward function . This formulation increases the representational power of the reward function while decreasing the representational load on the state encoder providing stronger guarantees on performance . • 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 2 , we introduce the successor feature framework with a non-linear reward function in Section 3 , Section 4 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 5 . 2 BACKGROUND . 2.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 ) . 2.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 , a ; θφ ) > w . Following from this , the Q function can be rewritten as : Q ( s , a ) = Eπ [ rt+1 + γrt+2 + . . . |St = s , At = a ] = Eπ [ φ ( at+1 , st+1 ; θφ ) > w + φ ( at+2 , 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 QLearning ( 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 . 3 MODEL , ARCHITECTURE , AND TRAINING . A natural extension to the Successor Feature framework begins by adjusting the fundamental structure of how the reward is represented . In this work , a second-order extension is proposed that improves the flexibility of the reward function while providing other benefits . This paper shows that by improving the representational power of the reward component , with a non-linear function of the state , it provides 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 the 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 . 3.1 NON-LINEAR REWARD FUNCTION . The successor feature framework builds upon functional representation of the current reward rt as a linear combination of the current state representation φt ( s ) ∈ Rz and a learned reward vector w ∈ Rz , such that rt = φt ( s ) > w . In this paper we extend the reward model by changing this linear reward model to one with the following form : rt = φt ( s ) > o + φt ( s ) > Aφt ( s ) ( 1 ) where { φt ( s ) , 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 φt ( s ) . From here on , we use φt instead of φt ( s ) for brevity . With a slight abuse of notation , we can see the original formulation leads to this if w is replaced with o+Aφ : rt = φ > ( 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 4.5 .

Summary: Successor features decomposes the policy into two components: one modelling the environmental dynamics and the other, the rewards. The paper describes an extension to the successor features framework where the rewards are modelled using a second order function. The extension provides a more robust reward component. Under this formulation a second term appears that models the expected auto-correlation matrix of the state features and can be used for guided exploration during transfer. The authors provide experimental results is three domains with comparison with the linear approach.