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: The paper proposes a method to incorporate Successor Features (SFs)  in domains with continuous state and action spaces. It proposes an actor-critic architecture (a variation of the Soft Actor-Critic method) that learns disentangled representations for the environment dynamics and the tasks. The network architecture guarantees such disentanglement with two independent modules: one for the representation of the dynamics $\phi$ (fed by the current state, action, and next state) and one for the task representation $\boldsymbol{w}$ (fed by task-specific information, such as the goal). These modules are learned jointly in a single-stage training procedure, contrasting prior work [1]. The main contribution of this model is the enablement of the SFs for continuous domains without relying on the costly inference mechanism from the classic SFs framework implementation while enabling generalization among similar tasks.