Generalizing Successor Features to continuous domains for Multi-task Learning

1 INTRODUCTION . Reinforcement learning ( RL ) tackles sequential decision making problems by defining optimal behavior through a reward function , where the agent learns how to behave through interacting with the environment and receiving rewards . The ability of RL algorithms to generalize across different , yet related reward functions , has a great potential to realize more data-efficient algorithms with the capability to transfer to new reward functions . In this paper we look at one particular type of generalization , where the reward function itself changes , however the underlying dynamics of the environment remain the same . This setup is flexible enough to allow transfer happen across tasks , by appropriately defining the rewards which induce different task decompositions . This type of task decomposition potentially allows the agent to tackle more complex problems than , would be possible were the tasks modeled as a single task . We are interested in a setup where the agent is exposed to multiple tasks i.e . tasks with different reward functions . In a multi-goal setting , different reward function can simply be the difference in the euclidean distance to different target goal locations . In a multi-task setting , the difference can be intricately designed in the reward function . For instance a reward function that determines walking forward vs walking backward . We argue that these differences in the structure of the reward function are difficult to capture within a goal or context-conditioned RL frameworks ( Sodhani et al. , 2021 ) . In the context of robotics , generalization across tasks is crucial . Consider an agent playing ball games with a racket Figure 1 . An agent trained to dribble the ball vs hitting the ball , should be able to quickly learn to play squash , as many of the skills such as approaching and hitting the ball are shared in a more complex task of playing squash . From the learners perspective , all these tasks share the same common properties , the ball falls to the ground due to gravity , depending on heavy it is , and it moves with certain velocity when it is hit by the racket . In other words , all these tasks share common dynamics . What changes is the small details in the reward function . For instance the difference between dribbling a ball vs hitting it against the wall , can be the rotation angle of the racket and the amount of force required . If it was possible to learn a representation that could decouple such discrepancies between the reward functions , i.e . decoupling the task dynamics and task-related dynamics , one could train an agent that could re-use the learned representation and quickly fine-tune itself to the more taskspecific representation and achieve a faster learning . Successor features ( SF ) ( Barreto et al. , 2017 ) is one framework that enables such decomposability of representation , explicitly built into the RL formulation . The main goal of this framework is to promote a desired property where instead of being posed as a decoupled representation learning problem , transfer is instead integrated into the RL framework as much as possible , preferably in a way that is almost transparent to the agent . SFs in theory , enable fast transfer between tasks that differ only in their reward function . The advantage of using an SF framework over model-based RL where one learns models of the reward function , is the ability of dynamics representation re-use which is decoupled from the task-specific representation . Our main contribution is to address the generalization and expensive inference problem in the classical SF frameworks coupled with GPI ( Barreto et al. , 2017 ) . We show that a simple architecture can provide a solution to this feature learning problem and demonstrate the effectiveness of our method compared to ( Barreto et al. , 2020 ) in more challenging continuous state and action setting . The majority of the existing work using SFs , operates under the discrete action setting or under the GPI setting optimizing a set of policies which in practice are hard to apply on real robotics applications . To the best of our knowledge , our method is the first to show the applicability of SFs coupled with an appropriate representation learn- ing mechanism solving challenging continuous control tasks . We show that simple modifications to an actor-critic framework can be easily coupled with SFs and empirically demonstrate the efficacy of our method . To summarize , our contributions are as follows : First we propose a practical implementation of SF framework for continuous state and action domains in the context of actor-critic architecture . Secondly we propose a robust method for learning the representations φ and w with the ability to learn disentangled representations for the state-space vs task-specific representation of complex nonlinear reward functions . Finally we demonstrate the efficacy our method on a ranging from tasks including classical continuous control 2D reacher domain on DM control suite ( Tassa et al. , 2020 ) to more challenging 3D reacher and manipulation tasks with the Sawyer arm on Meta-World benchmark ( Yu et al. , 2019 ) . 2 BACKGROUND . 2.1 REINFORCEMENT LEARNING . We assume the interaction between agent and environment can be modeled as a Markov Decision Process ( MDP ( Puterman , 1994 ) ) . An MDP is defined as a tuple M ≡ 〈S , A , p , R , γ〉 with state space S and action space A . For each s ∈ S and a ∈ A the function p ( .|s , a ) gives the next-state distribution upon taking action a in state s , where p ( .|s , a ) is referred to as the dynamics of the MDP . The random variableR ( s , a , s′ ) determines the reward received in the transition s a−→ s′ . Usually we are interested in the expected value of this variable , which is denoted by r ( s , a , s′ ) , and γ ∈ [ 0 , 1 ) weighs the importance of future rewards . The agent ’ s goal is to find a policy π : S → A , that is , a mapping from states to actions , that maximizes the value of every state-action pairs , defined as Qπ ( s , a ) ≡ Eπ [ ∞∑ i=0 γir ( St+i , At+i , St+i+1 ) |St = s , At = a ] . ( 1 ) where St and At are random variables indicating the state occupied and the action selected by the agent at time step t and Eπ [ . ] denotes expectation over the trajectories induced by π . The function Qπ ( s , a ) is referred to as the “ action-value function ” of policy π. RL algorithms based on dynamic programming build on two fundamental operations , policy evaluation which is the computation of Qπ ( s , a ) , the value function of policy π on task with reward r , and policy improvement theorem ( Bellman , 1957 ) . Once a policy π has been evaluated , we can compute a greedy policy π′ ( s ) ∈ arg maxaQπ ( s , a ) that is guaranteed to perform at least as well as π , that is : Qπ ′ ( s , a ) ≥ Qπ ( s , a ) for any ( s , a ) ∈ S × A . The computation of π′ is referred to as policy improvement . 2.2 MULTI-TASK RL & TRANSFER LEARNING . In practical situations , agents often face multiple related tasks , such as the robot learning numerous skills in Figure 1 . We define tasks Mi drawn from the setM . Then , the goal of multi-task learning is to find π∗i , an optimal policy for each MDP Mi with corresponding optimal value function Q π∗i i . Barreto et al . ( 2017 ) extended the policy improvement theorem to the scenario where the new policy is computed based on the value functions of a set of policies and referred to this as generalized policy improvement ( GPI ) . Suppose the agent has computed n policies corresponding toQπ1 , Qπ2 , ... , Qπn action-value functions . Let Qmax = maxiQπi and define π ( s ) ← arg maxaQmax ( s , a ) for all s ∈ S , then Qπ ( s , a ) ≥ Qmax ( s , a ) for all ( s , a ) ∈ S × A . The only caveat is that it is a waste of computation to compute the value functions of π∗1 , π ∗ 2 , ... , π ∗ n. This approach becomes appealing if we have a way to quickly compute the value functions of the policies πi on the task Mn+1 . 3 ACTOR-CRITIC SUCCESSOR FEATURES . This section will describe Successor Features and their previous use in the discrete action setting . We will then explain our extension , through Universal Value Function Approximators and an ActorCritic approach , to learn useful Successor Features and corresponding policies for high dimensional multi-task continuous control . 3.1 SUCCESSOR FEATURES DECOMPOSITION . Barreto et al . ( 2017 ) proposed a simple reward model which leads to the generalization of successor representation ( SR ) proposed in ( Dayan , 1993 ) . The key assumption is that the reward function can be approximately represented as a linear combination of learned features φ ( s ) . The successor representation ( SR ) ( Dayan , 1993 ) is a representation that generalizes between states using similarity between their successors , that is , the states that follow the current state given the agent ’ s policy . The generalization of SR with function approximation is referred to as Successor Features ( SF ) ( Barreto et al. , 2017 ) of ( s , a ) under policy π . Following ( Barreto et al. , 2017 ; 2018 ; 2020 ) , let φ : S×A×S → Rd be an arbitrary function whose output we will see as “ features ” . We assume that there exist features such that the reward function can be written as r ( s , a , s′ ) = φ ( s , a , s′ ) Tw ( 2 ) where φ ( s , a , s′ ) ∈ Rd are features of ( s , a , s′ ) and w ∈ Rd are weights . Intuitively we can think of φ ( s , a , s′ ) as salient events that may be desirable or undesirable to the agent . Based on Eq.2 we can define an environment Mφ ( S , A , p , γ ) as Mφ ≡ { M ( S , A , p , r , γ ) |r ( s , a , s′ ) = φ ( s , a , s′ ) Tw } , ( 3 ) that is , Mφ is the set of MDPs induced by φ through all possible instantiations of w. SFs make it possible to compute the value of a policy π on any task Mi ∈ Mφ by simply plugging in the representation vector wi defining the task . Specifically , if we substitute Eq.2 in the definition of action-value function of a policy we have Qπ ( s , a ) = Eπ [ rt+1 + γrt+2 + ... |St = s , At = a ] = Eπ [ φTt+1w + φ T t+2w ... |St = s , At = a ] = Eπ [ ∞∑ i=t γi−tφi+1|St = s , At = a ] T w = ψπ ( s , a ) Tw ( 4 ) One benefit of doing so is that if we replace wi with wj in Eq.4 , we immediately obtain the evaluation of π on task Mj . This way only relevant module must be relearned , when either the dynamics or reward changes . The key insight of SFs is that linearity of rewards rw with respect to the features φ which gives us the decomposition of the action value of policy π on task rw . In the GPI setting when the agent is presented with a new task Mn+1 , it needs to compute { Q π∗1 n+1 , Q π∗2 n+1 , ... , Q π∗n n+1 } , that is , the evaluation of each π∗i under the new reward function induced by wn+1 . This in turn would require applying the GPI theorem to the newly-computed set of value functions that will give rise to a policy that performs at least as well as the policy based on any subset of these . Hence ( Barreto et al. , 2017 ) proposed to incorporate SFs where the reward function changes to rn+1 ( s , a , s′ ) = φ ( s , a , s′ ) Twn+1 , as long as we have the correct wn+1 we can compute the value function of π∗i by simply computing Qπ ∗ i n+1 ( s , a ) = ψ π∗i ( s , a ) Twn+1 . This reduces the computation of all Q π∗i n+1 to the simpler supervised problem of approximating wn+1 ( Barreto et al. , 2020 ) . Although at first glance , GPI and SFs seem tangled , this characterization of SF does not depend on GPI framework itself . Thus , SF can be used in any RL framework where such decomposition of reward is viable . Furthermore the combination of SFs and GPI provides an elegant framework for transfer in a multi-task setting . However the computational cost of optimizing policies per task is prohibitive in more complex multi-task setting where in the size of task set Mφ is large .

Summary: This work looks at the use of successor features for solving simple continuous control tasks (in particular, reaching to different locations and door closing). The two contributions they enumerate are a ``practical implementation of SF framework for continuous state and action domains'' and jointly learning $\phi$ and $w$ (that is the successor features and the task weights $w$). They show their approach outperforms a goal-conditioned SAC baseline on these control tasks including generalizing to target locations that were not used during training.