Year: 2021

Abstract:
The deep reinforcement learning (RL) framework has shown great promise to tackle sequential decision-making problems, where the agent learns to behave optimally through interactions with the environment and receiving rewards. The ability of an RL agent to learn different reward functions concurrently has many benefits, such as the decomposition of task rewards and promoting skill reuse. One obstacle for achieving this, is the amount of data required as well as the capacity of the model for solving multiple tasks. In this paper, we consider the problem of continuous control for robot manipulation tasks with an explicit representation that promotes skill reuse while learning multiple tasks with similar reward function. Our approach relies on two key concepts: successor features (SF), a value function representation that decouples the dynamics of the environment from the rewards, and an actor-critic framework that incorporates the learned SF representations. We first show how to learn a decomposable representation required by SF. Our proposed methods, is able to learn decoupled state and reward feature representations. We empirically study this approach on non-trivial continuous control problems with compositional structure built into the reward functions of the tasks.

heading: 1 INTRODUCTION
text: 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).

heading: 2 BACKGROUND

heading: 2.1 REINFORCEMENT LEARNING
text: 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.

heading: 2.2 MULTI-TASK RL & TRANSFER LEARNING
text: 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.

heading: 3 ACTOR-CRITIC SUCCESSOR FEATURES
text: 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.

heading: 3.1 SUCCESSOR FEATURES DECOMPOSITION
text: 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.

heading: 3.2 UNIVERSAL SUCCESSOR FEATURES AND TASK INFERENCE
text: We start our discussion by assuming the existence of appropriate representations φ and w that are a decomposition of the reward function. We then turn our attention to the policy learning phase that will make use of these representations as a substitute for the reward signal. For an RL agent to generalise to unseen tasks, the agent needs to be able to identify and exploit some common structure underlying the task. Two possible sources of structure in this scenario are: i) similarity in the space of tasks, i.e. reward functions, and ii) the shared dynamics of the environment. One framework that enables exploiting structures are universal value function approximators (UVFAs) (Schaul et al., 2015). UVFAs extend the notion of value functions to also include the description of a task, thus directly exploiting the common structure in the associated optimal value functions. As discussed in (Borsa et al., 2018), UVFAs and SF & GPI address the transfer problem in quite different ways. With UVFAs, one trains an approximator Q̃(s, a,w) by solving the training tasks w ∈Mφ using any RL algorithm of choice. Borsa et al. (2018) proposed universal successor features approximators (USFAs) with GPI. They show that by combining USFAs and GPI can outperform SF&GPI on a grid-world navigation problem. Since majority of existing work explores SFs in the context of GPI or Q-learning under a discrete action setting, in this paper we will examine a more practical UVFA style framework targeted at solving more challenging continuous action control problems.
In this work, we ask the following questions: Can we incorporate SFs and train a policy which can yield better generalization across multiple tasks? Do SFs introduce implicit composability at the policy level? Composability of controllers is specially important for real-world applications, where reuse of past experience can greatly improve sample efficiency for tasks that can naturally be decomposed into simpler sub-problems. For instance, a policy for pick-and-place task can be decomposed into (1) reaching specific target, (2) grabbing an object, (3) avoiding certain obstacles. Such decomposable policies can be learned offline without the need to interact with the environment. Although we are not using an explicit notion of compositionality as in (Haarnoja et al., 2018a), we count on such properties emerging from the representation itself.

heading: 3.2.1 ALGORITHM DERIVATION
text: In this work we will be extending Soft Actor-Critic (Haarnoja et al., 2018b), by two means of generalisation. That is by combining, UVFA and USFA to learn to generalize values over multiple tasks. The structure of the RL problem creates a number of challenges. Without having access to the target function, we generally use the approximation itself to build targets and many of the techniques like replay buffer and target networks and optimizing two target networks by selecting the minimum which leads to lower variance, are just strategies to remedy this instability (Mnih et al., 2015; Fujimoto et al., 2018). Unfortunately since we are dealing with a multi-dimensional ψ target functions, these instabilities are exacerbated. We elude this problem by taking the average of ψ along the dimension of the features and the minimum is determined based on this average value.
This causes further instability during the policy learning, and longer delayed critic target updates were required. Similar to Hunt et al. (2018), we define action-dependent SF to include the entropy of the policy
ψπ(st, at) = φt + Eπ [ ∞∑ τ=i+1 γτ−t ( φτ + αH[π(.|s)] )] (5)
The max-entropy action-value of π for any convex combination of rewards w is then given by Qπw(s, a) = ψ π(s, a)w. For more details see Appendix A.1

heading: 3.2.2 LEARNING SUCCESSOR FEATURES IN CONTINUOUS CONTROL
text: We now describe how to build a decomposed representation required for φ and w. To make our discussion more concrete, consider a simple 2D goal reaching environment. We formulate the problem in Eq. 2 as computing an approximate φ̃ as a multi-task problem, solving the following approximation:
φ̃(s, a, s′)Tw ≈ ri(s, a, s′), for i = 1, 2, ..., D (6)
We can start our assumption by considering the scenario where such representations φ ∈ Rn and w exists that satisfy Eq. 2 exactly. For instance, we can define φi as an indicator function signaling important event in the state features e.g. whether an object of type i has been picked up by the agent (Barreto et al., 2020). Analogously, in a continuous control setting, we can consider a goal reaching task with a non-linear reward function. A common reward function for a simple reach task can be defined as the euclidean distance between two vectors, that is, the position of the agent’s link and the goal location. Given the state vector µ = (x, y) and goal vector g = (x, y), we define the following reward function,
r(µ, g) = −||µ− g||2 = ||µ||2 + 2gTµ− ||g||2 (7)
It is trivial to see the corresponding φ(µ) and w(g) can be recovered as follows,
φ(µ) =  1µ ||µ||2  ,w(g) = −||g||22g −1  (8) Thus given such decomposition, we can recover the reward function as the linear combination of φ(µ)Tw(g). Such representation, for this particular reacher task, immediately gives a solution to compute ψπ . And by changing the target goal location associated with each reacher task, one can create different tasks, where the only part of the representation i.e. goal location, changing is w, relying on the assumption that the representation captures any task-related features.
Now that we have defined φ, we turn to the question of how to determine an appropriate policy π to solve the task. We will start with what is perhaps the simplest case, for a single-goal reacher task. In our experiments, Figure 2, we compare the performance of SF integrated with SAC. As a sanity check, we alternate between keeping either φ or w fixed and regressing the other. This experiment ensures that our objective is able to learn the decomposed representation for this reward function for each component individually. These results confirm that, given the appropriate successor features, our policy learning mechanism should be able to solve the task, and perform as close to the optimal policy learned via the original reward signal. However most RL tasks have complex reward functions and it is not feasible to always hand-engineer these features. Instead, we are interested in learning these as a pre-training stage.
Barreto et al. (2020) proposed an approach to learn such representation. Given a task r, we are looking for w ∈ Rd that leads to good performance of the generalized policy πψ(s; w) . Suppose we have a set of m sample transitions from a given task, {(si, ai, r′i, s′i)}mi=1. Then based on Eq.2, we can infer w by solving the following minimization:
min w̃ = m∑ i=1 |φ(si, ai, s′i)T w̃− r′i|p, (9)
where p ≥ 1 and we may want to consider the inclusion of a regularization term in this objective function. If we handcraft φ, we can obtain a solution w̃, which can be plugged in during the representation learning phase. It turns out that Eq.9 can be generalized to also allow φ to be inferred from data. Given sample transitions from k tasks, {(sij , aij , r′ij , s′ij)} mj i=1, with j = 1, 2, ..., k, we can formulate the problem as the search for a function φ̃ : S ×A× S 7→ Rc and k vectors w̃i ∈ Rc satisfying,
min φ̃ k∑ j=1 min w̃j mj∑ i=1 |φ(sij , aij , s′ij)T w̃j − r′ij |p (10)
where p ≥ 1. Note that the features φ̃ can be arbitrary nonlinear functions of their inputs. As discussed in (Barreto et al., 2020) this objective can be decomposed into single objective for learning φ and w stage-wise where the initial learned representation φ can be re-plugged in to learn w. More generally, the problem in Eq.10 can be solved as a multi-task regression (Caruana, 1997).
Barreto et al. (2020) proposed the following procedure for learning φ̃ and w̃. First, a random policy π is used to collect data from tasks w1,w2, ...,wk. They propose to use a neural network with the number of output units or “heads”, corresponding to the number of tasks. The top layer of this joint network represents φ, and hence each output unit represents ri = φTwi. Given a sample transition (s, a, ri, s
′), only the i-th output unit of the network has access to the branch where the sample is used to modify wi and φ. Once this optimization is done, the weights wi are discarded and the same network is used to re-compute w, while keeping the learned representation φ̃ from previous iteration frozen. This process repeats the optimization process solving for equation 9, but this time solving for w. When learning w, “fresh” output units are added to the network i.e. the head branches of the network weights are randomly initialized again and using gradient descent the associated weights w are modified, while keeping φ fixed. Having computed φ and w, they resume to learning the policy in their GPI setting. (Barreto et al., 2020) demonstrated the effectiveness of this procedure using a simple tabular grid world navigation task. In this paper, we show that this procedure does not allow for learning more complex representation such as a simple reacher task with a euclidean distance reward function. Instead, we propose the following procedure which deemed more appropriate for more complex control tasks with continuous state and action spaces.

heading: 3.2.3 JOINT TASK INFERENCE
text: Following (Barreto et al., 2020), we propose optimizing an approximation to the objective in Eq.10 where the representations φ and w are learned jointly, as opposed to the two-stage optimization procedure proposed in (Barreto et al., 2020). Given a set of transitions {(si, ai, r′i, s′i)}mi=1, ideally covering the space of state space for all the tasks seen at training time, we use two neural networks φθ and wθ, where the function φ takes as input the tuple (s, a, s′) and the function w is a function of the task itself, in a multi-goal setting, this can be the position of the target goal or task ids. Given these two networks, using stochastic gradient descent we can approximate Eq.10 by optimizing,
min φ̃θ k∑ j=1 min w̃θ m∑ i=1 |φ̃(sij , aij , s′ij)T w̃(g)− r′ij |p (11)
where φ̃ and w are optimized jointly by considering their output as φ̃Tw ≈ r. This decomposition has two advantages over (Barreto et al., 2020). Firstly, it is easier to decompose the environment dynamics information fed to the network via φ(s, a, s′) and task-specific information g via w(g) where each of these components are fed the relevant information and the composition of these corresponds to the reward approximation. This is not possible with the architecture proposed in (Barreto et al., 2020) since both φ and w share the same input and the decomposition of these representations are made explicit during the stage-wise optimization procedure. In addition, single-stage optimization makes the inference task simpler to train and evaluate. We will illustrate the results of both these approaches in our setup for a simple single-goal and multi-goal reacher task. Figure 3 shows the architecture of the two methods.

heading: 4 EXPERIMENTS
text: Our experiments consist of learning continuous control robot manipulation tasks shown in Figure 4. In the case of 2D link reacher, we fix the goals, separated into a set of training goal targets and test goal targets, and for metaworld, we consider single goal tasks where the tasks are different.
When learning a decomposed representation of the reward function, the function itself should ideally be a dense function. Otherwise it is hard to learn a good representation with a sparse rewards as also pointed out by (Machado et al., 2020)). For this reason, we modify the original reward function for the default 2D Reacher task in DMSuite to a simple L2 norm squared reward function which is easier to learn than a semi-sparse reward function. The original reward function of this task is sparse, where dense reward are given when the agent is in the proximity of the target goal.
For each task, first, we need to learn representations φ̃ and w̃. Barreto et al. (2020) proposed to use a random policy to collect the initial data which could be sufficient for a simple grid world navigation task. In our setting, we found that a simple random policy does not allow full coverage of state space. It is possible that a random policy in a small grid-world setting is enough however we found this to be insufficient in a continuous control setting, allowing to learn a generalizable φ̃ and w̃. Therefore we used a mixture of expert data and data generated by a random policy to train φ̃ and w̃.
Figure 5 shows comparison between the proposed method in (Barreto et al., 2020) and our method where we compare a SAC policy trained with the original reward, our method with ACSF where φ and w are learned jointly using the training scheme, referred to as ACSFJointHeadMLP, proposed in (Barreto et al., 2020) vs our method for learning the joint representation. The task for this experiment is a multi-goal reacher where for the architecture in (Barreto et al., 2020), each goal i.e. task is represented by a separate head of the network. We can see that training a joint architecture becomes harder for two reasons. Since the joint network takes in the input (s, a, s′, g), it is not possible to disentangle the inputs to the module learning φ and the module w. As expected, the performance drops even further, once we remove the goal information from the observations tuple. We prefer to disentangle these observations for the dynamics learning vs any task-related part of the observation. In addition, the training methodology in (Barreto et al., 2020) is a 3-stage training for the representation learning component, meanwhile one must ensure overfitting is not happening at each stage of learning, and this is hard to evaluate in general without having access to any prior ground-truth. Although this is not a direct comparison of the GPI method in (Barreto et al., 2020), the main objective of these experiment is to evaluate the representation learning procedures embedded into our policy learning framework.
Figure 6 shows the generalization capability of our method evaluated on seen goal locations at train time vs unseen goal locations. Finally Figure 7 shows the performance of our method compared to a goal-conditioned SAC policy. Figure 7a shows the performance of both methods on the reacher task where the goal location is perturbed but this is still a single goal task. Figure 7b shows the performance on door close task and finally Figure 7c shows when these tasks are learned simultaneously. These results confirm that our method works in par with a goal-conditioned policy and it has learned the correct features which in turn enable learning robust policies.

heading: 5 RELATED WORK
text: Multi-task RL and representation learning in RL are important topics that have generated a large body of literature and as (Dabney et al., 2021) formally points out, there remains challenging aspect of RL when learning representations. There exist various techniques that inject tasks information directly into the definition of the value function. UVFAs have been used for zero-shot generalization (Mankowitz et al., 2018) and learning a set of fictitious goals previously encountered by the agent (Andrychowicz et al., 2017). Haarnoja et al. (2018a) proposed learning with multiple objectives where they define a compound task expressed in terms of the sum of individual rewards. More recently, (Sodhani et al., 2021) proposed an approach for a contextual multi-task RL as a way to incorporate task metadata, or contextual information. One other common and successful way to approach the representation learning problem is through the use of auxiliary tasks: additional prediction problems that shape the representation used by the agent (Jaderberg et al., 2017; Bellemare et al., 2019).
Another body of work relates to the use of SFs. Kulkarni et al. (2016) presented a generalization of successor representation in a deep reinforcement learning framework, based on DQN with discrete action space. Similar in the spirit, (Zhang et al., 2017), proposes a deep RL architecture built on DQN to solve a robot navigation task. Lehnert & Littman (2019) draws connections between modelbased RL and successor features by analyzing properties of different learned latent state spaces. More related to our work is Ma et al. (2020) which combines goal-conditioning with SFs, where SFs are learned end-to-end using temporal difference learning methods. However in their setup, learning these successor features is done implicitly inside their architecture. In addition, they use the true rewards during policy learning and the additional auxiliary loss for learning ψ demonstrates additional stability over DDPG baselines.
Barreto et al. (2019) showed how to combine SF&GPI with deep learning using the reward functions themselves as features for future tasks. Barreto et al. (2020) introduced the representation learning component of SFs as opposed to using hand-crafted features as done in prior work. Hansen et al. (2019) introduces Variational Intrinstic Successor FeatuRes (VISR) which shows that a behavioral mutual information (BMI) maximization provides a solution for learning the successor feature representation. They evaluate their method as an unsupervised pre-training technique. Hunt et al. (2018) introduced SF in the context of maximum entropy framework, extending GPI theorem to max-ent objective. Blier et al. (2021) formally derive a temporal difference algorithm for successor state and goal-dependant value function learning with function approximation. Zahavy et al. (2021) considers a specific class of policy composition called set improving policies (SIPs). Gimelfarb et al. (2021) proposes risk-aware successor features (RaSF) integrated into a generalized policy improvement framework to maximize entropic utilities. Touati & Ollivier (2021) introduces the forward-backward (FB) representation of the dynamics of an MDP, learning two representations instead of one in the SF framework. Their setting is also more tailored towards a goal-oriented RL problems with discrete action space.

heading: 6 CONCLUSION
text: Our results suggest that our method can successfully incorporate SFs representation for learning continuous control policies. Learning decoupled representations has been a longstanding challenge in RL, and this work must continue in continuous control tasks, especially for the multi-task setting with continuous goals. Our work sheds light on possible ways of incorporating SFs and their capability to solve more challenging RL control problems. Unlike the framework proposed in (Barreto et al., 2020), which is explicitly tailored for transfer, this is not yet as straightforward in our actorcritic architecture. Future work in this area includes more stable online feature learning and scaling to more complex reward functions. By demonstrating that ACSF performs well on large continuous tasks, we believe our approach is an important step in the direction of composable representations.

heading: A APPENDIX
text: A.1 IMPLEMENTATION DETAILS
The loss function for learning ψ is:
LACSF (ψ) = E(s,a,s′,d)∼D
[ 1
2
( ψ(s, a)− y(φ, s′, d) )2] (12)
where the target is given by
y(φ, s′, d) = φ(s, a, s′) + γ(1− d)(ψ(s, a)− α log πθ(ã′|s′)), ã′ ∼ πθ(.|s′) (13)
where d corresponds to the termination criterion. The gradient with respect to the parameter θ is
OθLACSF (ψ) = E(s,a,s′,d) [( ψ(s, a)− y(s, a, s′)Oθψ(s, a) )] (14)
which is similar to the gradient used in (Haarnoja et al., 2018b), with the distinction that ψ is a multi-dimensional matrix rather than a vector as in the original definition of Q. The dimensionality of ψ is a hyperparameter that needs to be tuned depending on the task observations complexity.
Figure 3 shows the comparison of the setup for our representation learning module compared to (Barreto et al., 2020). First stage requires learning the representations φ and w. Ideally, the data used to train these modules needs to be diverse enough, covering the state space adequately. We found that learning these representation online tangled with policy learning does not perform well since the representation needs to be able to handle regions where policy behavior is optimal. Therefore we resort to using expert policy data. However we still do require the diversity. More intuitively, this is because at the beginning of the agents learning experience, a random policy is exploring, it is important to have good coverage of the state space that can ultimately lead to the highest performing regions. Therefore we use a mixture of expert to random policy data of a 75%− 85% of expert and the remaining data generated from a random policy with some action noise. We found this ratio to be important for training robust representations that will later on perform well during the SF policy learning stage.
We also regularize the w using the layer weight normalization (Salimans & Kingma, 2016) which is a simple reparameterization of the weight vectors in a neural network that decouples the magnitude of a weight tensor from its direction. For the policy learning, we use an open-source SAC implementation 1, extending this implementation to a ACSF policy.
A.2 APPROXIMATING THE REWARD FUNCTION
We demonstrated our method on tasks with more reaching-style reward functions which are still learnable with our objective function. However this simple objective fails to capture more difficult tasks or tasks with more complex state representations. For instance, our method fails to learn good representations for some of the tasks in Meta-World such as picking and placing objects or tasks interacting with an object such as pressing a button. We hypothesize that more complex reward functions are too difficult for the current method to automatically decompose.
A.3 FEW-SHOT TRANSFER
Consider the 3 reaching scenarios presented in Figure 8. In our experiments we found that, an agent that had been trained on the reacher and door close task, struggled to zero-shot transfer to the third reacher task with an obstacle on the way. The performance had higher variance than we expected and this happened with both SAC and ACSF which could entail an underlying issue with the policy learning. Note that we would expect the agent to adapt to the third scenario through finetuning. However in our experiments we found that the continual learning the third behavior led to slower learning and, in some cases, reaching sub-optimal behavior due to instability. We found that it is better to train an agent on the 3 tasks simultaneously or an agent trained per task separately. We hypothesise that an agent starting from scratch with random exploration has a better chance at discovering all the skills compared to an agent who has learned to reach the first two scenarios. Transferring to the third case, it needs to unlearn the previous behavior and adapt.
1https://github.com/denisyarats/pytorch sac
Algorithm 1 ACSF- Plugging in learned φ̃ and w̃ Input: initial policy parameters θ
Ψ parameters corresponding to ψ1 and ψ2 Replay buffer D Set target parameters equal to main parameters ψtar,1 ← ψ1, ψtar,2 ← ψ2
for Time t = 0 to∞ do Observe state s and select action a ∼ πθ(.|s) Step in environment, s′t ∼ p(.|st, at) D ← D ∪ (st, at, φt, s′t, g) if it’s time to update then
for j in range(numUpdates) do Randomly sample a batch of transitions, B = {(s, a, φ, s′, d, g)} from D Compute targets for the D-dimensional Ψ functions:
ΨDψtar,1(s, a, g) = D∑ d=1 Ψψtar,1
ΨDψtar,2(s, a, g) = D∑ d=1 Ψψtar,2
y(φ, s′, d) = φ(s, a, s′) + γ(1− d)(min i=1,2 Ψψtar,i(s, a, g)− α log πθ(ã′|s′, g)),
ã′ ∼ πθ(.|s′, g)
Update Ψ-functions by one step of gradient descent using
Oψi 1 |B| ∑
(s′,a′,φ,s′,d)
∈ B(min i=1,2 Ψ(s, ãθ(s), g)− α log πθ(ãθ(s)|s, g))
Compute the Q−function
Q1← Ψπtar,1(s, a, g)Tw Q2← Ψπtar,2(s, a, g)Tw
Update policy by one step of gradient ascent using
Oθ 1 |B| ∑
(s,g)∈B
(min i=1,2
Q(s, ãθ(s), g)− αlogπθ(ãθ(s)|s, g)),
where ãθ(s) is a sample from πθ(.|s) which is differentiable w.r.t. θ via the reparameterization trick. Update target networks with
ψtarg,i ← ρψtarg,i + (1− ρ)ψi

Review summaries:
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. 

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.