UNIVERSAL SUCCESSOR FEATURES APPROXIMATORS
Year: 2019

Abstract:
The ability of a reinforcement learning (RL) agent to learn about many reward functions at the same time has many potential benefits, such as the decomposition of complex tasks into simpler ones, the exchange of information between tasks, and the reuse of skills. We focus on one aspect in particular, namely the ability to generalise to unseen tasks. Parametric generalisation relies on the interpolation power of a function approximator that is given the task description as input; one of its most common form are universal value function approximators (UVFAs). Another way to generalise to new tasks is to exploit structure in the RL problem itself. Generalised policy improvement (GPI) combines solutions of previous tasks into a policy for the unseen task; this relies on instantaneous policy evaluation of old policies under the new reward function, which is made possible through successor features (SFs). Our proposed universal successor features approximators (USFAs) combine the advantages of all of these, namely the scalability of UVFAs, the instant inference of SFs, and the strong generalisation of GPI. We discuss the challenges involved in training a USFA, its generalisation properties and demonstrate its practical benefits and transfer abilities on a large-scale domain in which the agent has to navigate in a first-person perspective three-dimensional environment.

heading: 1 INTRODUCTION
text: Reinforcement learning (RL) provides a general framework to model sequential decision-making problems with sparse evaluative feedback in the form of rewards. The recent successes in deep RL have prompted interest in increasingly more complex tasks and a shift in focus towards scenarios in which a single agent must solve multiple related problems, either simultaneously or sequentially. This paradigm is formally known as multitask RL (Taylor and Stone, 2009; Teh et al., 2017). One of the benefits of learning about multiple tasks at the same time is the possibility of transferring knowledge across tasks; in essence, this means that by jointly learning about a set of tasks one should be able to exploit their common structure to speed up learning and induce better generalisation (Taylor and Stone, 2009; Lazaric, 2012). A particularly interesting instance of transfer is the generalisation to new, unseen tasks. This potentially allows an agent to perform a task with little or no learning by leveraging knowledge from previously learned tasks. In this paper we will be exploring this scenario.
Consider an RL agent in a persistent environment trying to master a number of tasks. In order to generalise to unseen tasks, the agent needs to be able to identify and exploit some common structure underlying the tasks. Two possible sources of structure in this scenario are: i) some similarity between the solutions of the tasks, either in the policy or in the associated value-function space, and ii) the shared dynamics of the environment (e.g., physics). In this paper we will attempt to build an agent that can make use of both types of structure. For this, we will build on two frameworks that exploit these structures in isolation. The first one are Schaul et al.’s (2015) universal value function approximators (UVFAs). 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. The second framework we build upon exploits the common structure in the environment and capitalises on the power of dynamic programming. Barreto et al.’s (2017) framework is based on two core ∗Corresponding author: borsa@google.com
concepts: successor features (SFs), a representation scheme that allows a policy to be evaluated on any task of a given format, and generalised policy improvement (GPI), a generalisation of dynamic programming’s classic operator that uses a set of policies instead of a single one.
UVFAs and SF&GPI generalise to new tasks in quite different, and potentially complementary, ways. UVFAs aim to generalise across the space of tasks by exploiting structure in the underlying space of value functions. In contrast, SF&GPI strategy is to exploit the structure of the RL problem itself. In this paper we propose a model that exhibits both types of generalisation. The basic insight is to note that SFs are multi-dimensional value functions, so we can extend them in the same way a universal value function extends their unidimensional counterparts. We call the resulting model universal successor features approximators, or USFAs for short. USFA is a strict generalisation of its precursors. Specifically, we show that by combining USFAs and GPI we can recover both UVFAs and SF&GPI as particular cases. This opens up a new spectrum of possible approaches in between these two extreme cases. We discuss the challenges involved in training a USFA and demonstrate the practical benefits of doing so on a large-scale domain in which the agent has to navigate in a three-dimensional environment using only images as observations.

heading: 2 BACKGROUND
text: In this section we present some background material, formalise the scenario we are interested in, and briefly describe the methods we build upon.

heading: 2.1 MULTITASK REINFORCEMENT LEARNING
text: We consider the usual RL framework: an agent interacts with an environment and selects actions in order to maximise the expected amount of reward received in the long run (Sutton and Barto, 1998). As usual, we assume that such an interaction can be modeled as a Markov decision process (MDP, Puterman, 1994). An MDP is defined as a tuple M ≡ (S,A, p, R, γ) where S andA are the state and action spaces, p(·|s, a) gives the next-state distribution upon taking action a in state s, R(s, a, s′) is a random variable representing the reward received at transition s a−→ s′, and γ ∈ [0, 1) is a discount factor that gives smaller weights to future rewards.
As mentioned in the introduction, in this paper we are interested in the multitask RL scenario, where the agent has to solve multiple tasks. Each task is defined by a reward function Rw; thus, instead of a single MDP M , our environment is a set of MDPs that share the same structure except for the reward function. Following Barreto et al. (2017), we assume that the expected one-step reward associated with transition s a−→ s′ is given by
E [Rw(s, a, s ′)] = rw(s, a, s ′) = φ(s, a, s′)>w, (1)
where φ(s, a, s′) ∈ Rd are features of (s, a, s′) and w ∈ Rd are weights. The features φ(s, a, s′) can be thought of as salient events that may be desirable or undesirable to the agent, such as for example picking up an object, going through a door, or knocking into something. In this paper we assume that the agent is able to recognise such events—that is, φ is observable—, but the solution we propose can be easily extended to the case where φ must be learned (Barreto et al., 2017). Our solution also applies to the case where (1) is only approximately satisfied, as discussed by Barreto et al. (2018).
Given w ∈ Rd representing a task, the goal of the agent is to find a policy πw : S 7→ A that maximises the expected discounted sum of rewards, also called the return G(t)w = ∑∞ i=0 γ iR (t+i) w , where R(t)w = Rw(St, At, St+1) is the reward received at the tth time step. A principled way to address this problem is to use methods derived from dynamic programming (DP), which heavily rely on the concept of a value function (Puterman, 1994). The action-value function of a policy π on task w is defined as Qπw(s, a) ≡ Eπ [ G (t) w |St = s,At = a ] , where Eπ[·] denotes expected value when following policy π. Based on Qπw we can compute a greedy policy π ′(s) ∈ argmaxaQπw(s, a); one of the fundamental results in DP guarantees that Qπ ′
w (s, a) ≥ Qπw(s, a) for all (s, a) ∈ S ×A. The computation of Qπw(s, a) and π
′ are called policy evaluation and policy improvement; under certain conditions their successive application leads to the optimal value function Q∗w, from which one can derive an optimal policy for task w as π∗w(s) ∈ argmaxaQ∗w(s, a) (Sutton and Barto, 1998).
As a convention, in this paper we will add a tilde to a symbol to indicate that the associated quantity is an approximation; we will then refer to the respective tunable parameters as θ. For example, the agent computes an approximation Q̃πw ≈ Qπw by tuning θQ.

heading: 2.2 TRANSFER LEARNING
text: Here we focus on one aspect of multitask RL: how to transfer knowledge to unseen tasks (Taylor and Stone, 2009; Lazaric, 2012). Specifically, we ask the following question: how can an agent leverage knowledge accumulated on a set of tasksM⊂ Rd to speed up the solution of a new task w′ /∈M? In order to investigate the question above we recast it using the formalism commonly adopted in learning. Specifically, we define a distribution Dw over Rd and assume the goal is for the agent to perform as well as possible under this distribution. As usual, we assume a fixed budget of sample transitions and define a training setM ∼ Dw that is used by the agent to learn about the tasks of interest. We also define a test setM′ ∼ Dw and use it to assess the agent’s generalisation—that is, how well it performs on the distribution of MDPs induced by Dw. A natural way to address the learning problem above is to use Schaul et al.’s (2015) universal valuefunction approximators (UVFAs). The basic insight behind UVFAs is to note that the concept of optimal value function can be extended to include as one of its arguments a description of the task; an obvious way to do so in the current context is to define the function Q∗(s, a,w) : S ×A× Rd 7→ R as the optimal value function associated with task w. The function Q∗(s, a,w) is called a universal value function (UVF); a UVFA is then the corresponding approximation, Q̃(s, a,w). When we define transfer as above it becomes clear that in principle a sufficiently expressive UVFA can identify and exploit structure across the joint space S ×A×Rd. In other words, a properly trained UVFA should be able to generalise across the space of tasks.
A different way of generalising across tasks is to use Barreto et al.’s (2017) framework, which builds on assumption (1) and two core concepts: successor features (SFs) and generalised policy improvement (GPI). The SFs ψ ∈ Rd of a state-action pair (s, a) under policy π are given by
ψπ(s, a) ≡ Eπ [ ∞∑ i=t γi−tφi+1 |St = s,At = a ] . (2)
SFs allow one to immediately compute the value of a policy π on any task w: it is easy to show that, when (1) holds, Qπw(s, a) = ψ
π(s, a)>w. It is also easy to see that SFs satisfy a Bellman equation in which φ play the role of rewards, so ψ can be learned using any RL method (Szepesvári, 2010).
GPI is a generalisation of the policy improvement step described in Section 2.1. The difference is that in GPI the improved policy is computed based on a set of value functions rather than on a single one. Suppose that the agent has learned the SFs ψπi of policies π1, π2, ..., πn. When exposed to a new task defined by w, the agent can immediately compute Qπiw (s, a) = ψ
πi(s, a)>w. Let the GPI policy be defined as π(s) ∈ argmaxaQmax(s, a), where Qmax = maxiQπi . The GPI theorem states that Qπ(s, a) ≥ Qmax(s, a) for all (s, a) ∈ S × A. The result also extends to the scenario where we replace Qπi with approximations Q̃πi (Barreto et al., 2017).

heading: 3 UNIVERSAL SUCCESSOR FEATURES APPROXIMATORS
text: UVFAs and SF&GPI address the transfer problem described in Section 2.2 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. One can then generalise to a new task by plugging its description w′ into Q̃ and then acting according to the policy π(s) ∈ argmaxaQ̃(s, a,w′). With SF&GPI one solves each task w ∈ M and computes an approximation of the SFs of the resulting policies πw, ψ̃ πw
(s, a) ≈ ψπw(s, a). The way to generalise to a new task w′ is to use the GPI policy defined as π(s) ∈ argmaxamaxw∈M ψ̃ πw (s, a)>w′.
The algorithmic differences between UVFAs and SF&GPI reflect the fact that these approaches exploit distinct properties of the transfer problem. UVFAs aim at generalising across the space of tasks by exploiting structure in the function Q∗(s, a,w). In practice, such strategy materialises in the choice of function approximator, which carries assumptions about the shape of Q∗(s, a,w). For example, by using a neural network to represent Q̃(s, a,w) one is implicitly assuming that
Q∗(s, a,w) is smooth in the space of tasks; roughly speaking, this means that small perturbations to w will result in small changes in Q∗(s, a,w).
In contrast, SF&GPI’s strategy to generalise across tasks is to exploit structure in the RL problem itself. GPI builds on the general fact that a greedy policy with respect to a value function will in general perform better than the policy that originated the value function. SFs, in turn, exploit the structure (1) to make it possible to quickly evaluate policies across tasks—and thus to apply GPI in an efficient way. The difference between the types of generalisation promoted by UVFAs and SF&GPI is perhaps even clearer when we note that the latter is completely agnostic to the way the approximations ψ̃ πw are represented, and in fact it can applied even with a tabular representation.
Obviously, both UVFAs and GPI have advantages and limitations. In order to illustrate this point, consider two tasks w and w′ that are “similar”, in the sense that ||w − w′|| is small (|| · || is a norm in Rd). Suppose that we have trained a UVFA on task w and as a result we obtained a good approximation Q̃(s, a,w) ≈ Q∗(s, a,w). If the structural assumptions underlying Q̃(s, a,w) hold— for example, Q̃(s, a,w) is smooth with respect to w—, it is likely that Q̃(s, a,w′) will be a good approximation of Q∗(s, a,w′). On the other hand, if such assumptions do not hold, we should not expect UVFA to perform well. A sufficient condition for SF&GPI to generalise well from task w to task w′ is that the policy π(s) ← argmaxaψ̃ πw (s, a)>w′ does well on task w′, where πw is a solution for task w. On the downside, SF&GPI will not exploit functional regularities at all, even if they do exist. Let policy πw′ be a solution for tasks w′. In principle we cannot say anything about ψπw′ (s, a), the SFs of πw′ , even if we have a good approximation ψ̃ πw (s, a) ≈ ψπw(s, a).
As one can see, the types of generalisation provided by UVFAs and SF&GPI are in some sense complementary. It is then natural to ask if we can simultaneously have the two types of generalisation. In this paper we propose a model that provides exactly that. The main insight is actually simple: since SFs are multi-dimensional value functions, we can extend them in the same way as universal value functions extend regular value functions. In the next section we elaborate on how exactly to do so.

heading: 3.1 UNIVERSAL SUCCESSOR FEATURES
text: As discussed in Section 2.2, UVFs are an extension of standard value functions defined asQ∗(s, a,w). If πw is one of the optimal policies of task w, we can rewrite the definition as Qπw(s, a,w). This makes it clear that the argument w plays two roles in the definition of a UVF: it determines both the task w and the policy πw (which will be optimal with respect to w). This does not have to be the case, though. Similarly to Sutton et al.’s (2011) general value functions (GVFs), we could in principle define a function Q(s, a,w, π) that “disentangles” the task from the policy. This would provide a model that is even more general than UVFs. In this section we show one way to construct such a model when assumption (1) (approximately) holds.
Note that, when (1) is true, we can revisit the definition of SFs and write Q(s, a,w, π) = ψπ(s, a)>w. If we want to be able to compute Q(s, a,w, π) for any π, we need SFs to span the space of policies π. Thus, we define universal successor features as ψ(s, a, π) ≡ ψπ(s, a). Based on such definition, we call ψ̃(s, a, π) ≈ ψ(s, a, π) a universal successor features approximator (USFA). In practice, when defining a USFA we need to define a representation for the policies π. A natural choice is to embed π onto Rk. Let e : (S 7→ A) 7→ Rk be a policy-encoding mapping, that is, a function that turns policies π into vectors in Rk. We can then see USFs as a function of e(π): ψ(s, a, e(π)). The definition of the policy-encoding mapping e(π) can have a strong impact on the structure of the resulting USF. We now point out a general equivalence between policies and reward functions that will provide a practical way of defining e(π). It is well known that any reward function induces a set of optimal policies (Puterman, 1994). A point that is perhaps less immediate is that the converse is also true. Given a deterministic policy π, one can easily define a set of reward functions that induce this policy: for example, we can have rπ(s, π(s), ·) = 0 and rπ(s, a, ·) = c, with c < 0, for any a 6= π(s). Therefore, we can use rewards to refer to deterministic policies and vice-versa (as long as potential ambiguities are removed when relevant).
Since here we are interested in reward functions of the form (1), if we restrict our attention to policies induced by tasks z ∈ Rd we end up with a conveniently simple encoding function e(πz) = z. From this encoding function it follows that Q(s, a,w, πz) = Q(s, a,w, z). It should be clear that UVFs are a particular case of this definition when w = z. Going back to the definition of USFs, we can
finally write Q(s, a,w, z) = ψ(s, a, z)>w. Thus, if we learn a USF ψ(s, a, z), we have a value function that generalises over both tasks and policies, as promised.

heading: 3.2 USFA GENERALISATION
text: We now revisit the question as to why USFAs should provide the benefits associated with both UVFAs and SF&GPI. We will discuss how exactly to train a USFA in the next section, but for now suppose that we have trained one such model ψ̃(s, a, z) using the training tasks inM. It is then not difficult to see that we can recover the solutions provided by both UVFAs and SF&GPI. Given an unseen task w′, let π be the GPI policy defined as
π(s) ∈ argmaxamax z∈C Q̃(s, a,w′, z) = argmaxamax z∈C ψ̃(s, a, z)>w′, (3)
where C ⊂ Rd. Clearly, if we make C = {w′}, we get exactly the sort of generalisation associated with UVFAs. On the other hand, setting C =M essentially recovers SF&GPI. The fact that we can recover both UVFAs and SF&GPI opens up a spectrum of possibilities in between the two. For example, we could apply GPI over the training set augmented with the current task, C = M∪ {w′}. In fact, USFAs allow us to apply GPI over any set of tasks C ⊂ Rd. The benefits of this flexibility are clear when we look at the theory supporting SF&GPI, as we do next.
Barreto et al. (2017) provide theoretical guarantees on the performance of SF&GPI applied to any task w′ ∈M′ based on a fixed set of SFs. Below we state a slightly more general version of this result that highlights the two types of generalisation promoted by USFAs (proof in Barreto et al.’s (2017) Theorem 2).
Proposition 1 Let w′ ∈M′ and let Qπw′ be the action-value function of executing policy π on task w′. Given approximations {Q̃πzw′ = ψ̃(s, a, z)>w′}z∈C , let π be the GPI policy defined in 3. Then,
‖Q∗w′ −Qπw′‖∞ ≤ 2
1− γminz∈C ‖φ‖∞ ‖w′ − z‖︸ ︷︷ ︸ δd(z) +max z∈C ‖w′‖ · ‖ψπz − ψ̃(s, a, z)‖∞︸ ︷︷ ︸ δψ(z)  , (4) where Q∗w′ is the optimal value of task w
′, ψπz are the SFs corresponding to the optimal policy for task z, and ‖f − g‖∞ = maxs,a |f(s, a)− g(s, a)|.
When we write the result in this form, it becomes clear that, for each policy πz, the right-hand side of (4) involves two terms: i) δd(z), the distance between the task of interest w′ and the task that induced πz, and ii) δψ(z), the quality of the approximation of the SFs associated with πz.
In order to get the tightest possible bound (4) we want to include in C the policy z that minimises δd(z) + δψ(z). This is exactly where the flexibility of choosing C provided by USFAs can come in handy. Note that, if we choose C =M, we recover Barreto et al.’s (2017) bound, which may have an irreducible minz∈Cδd(z) even with a perfect approximation of the SFs inM. On the other extreme, we could query our USFA at the test point C = {w′}. This would result in δd(w′) = 0, but can potentially incur a high cost due to the associated approximation error δψ(w′).

heading: 3.3 HOW TO TRAIN A USFA
text: Now that we have an approximation Q̃(s, a,w, z) a natural question is how to train this model. In this section we show that the decoupled nature of Q̃ is reflected in the training process, which assigns clearly distinct roles for tasks w and policies πz.
In our scenario, a transition at time t will be (st, at,φt+1, st+1). Note that φ allows us to compute the reward for any task w, and since our policies are encoded by z, transitions of this form allow us to learn the value function of any policy πz on any task w. To see why this is so, let us define the temporal-difference (TD) error (Sutton and Barto, 1998) used in learning these value functions.
Given transitions in the form above, the n-step TD error associated with policy πz on task w will be
δt,nwz = t+n−1∑ i=t γi−trw(si, ai, si+1) + γ nQ̃(st+n, πz(st+n),w, z)− Q̃(st, at,w, z)
= [ t+n−1∑ i=t γi−tφ(si, ai, si+1) + γ nψ̃(st+n, at+n, z)− ψ̃(st, at, z) ]> w = (δt,nz ) >w, (5)
where at+n = argmaxbQ̃(st+n, b, z, z) = argmaxbψ̃(st+n, b, z) >z. As is well known, the TD error δt,nwz allows us to learn the value of policy πz on task w; since here δ t,n wz is a function of z and w only, we can learn about any policy πz on any task w by just plugging in the appropriate vectors.
Equation (5) highlights some interesting (and subtle) aspects involved in training a USFA. Since the value function Q̃(s, a,w, z) can be decoupled into two components, ψ̃(s, a, z) and w, the process of evaluating a policy on a task reduces to learning ψ̃(s, a, z) using the vector-based TD error δt,nz showing up in (5). Since δt,nz is a function of z only, the updates to Q̃(s, a,w, z) will not depend on w. How do the tasks w influence the training of a USFA, then? If sample transitions are collected by a behaviour policy, as is usually the case in online RL, a natural choice is to have this policy be induced by a task w. When this is the case the training tasks w ∈M will define the distribution used to collect sample transitions. Whenever we want to update ψ(s, a, z) for a different z than the one used in generating the data, we find ourselves under the off-policy regime (Sutton and Barto, 1998).
Assuming that the behaviour policy is induced by the tasks w ∈M, training a USFA involves two main decisions: how to sample tasks fromM and how to sample policies πz to be trained through (5) or some variant. As alluded to before, these decisions may have a big impact on the performance of the resulting USFA, and in particular on the trade-offs involved in the choice of the set of policies C used by the GPI policy (3). As a form of illustration, Algorithm 1 shows a possible regime to train a USFA based on particularly simple strategies to select tasks w ∈M and to sample policies z ∈ Rd. One aspect of Algorithm 1 worth calling attention to is the fact that the distribution Dz used to select policies can depend on the current task w. This allows one to focus on specific regions of the policy space; for example, one can sample policies using a Gaussian distribution centred around w.
Algorithm 1 Learn USFA with -greedy Q-learning Require: , training tasksM, distributionDz over Rd, num-
ber of policies nz 1: select initial state s ∈ S 2: for ns steps do 3: sample w uniformly at random fromM 4: {sample policies, possibly based on current task} 5: for i← 1, 2, ..., nz do zi ∼ Dz(·|w) 6: if Bernoulli( )=1 then a← Uniform(A) 7: else a← argmaxbmaxi ψ̃(s, b, zi)>w {GPI } 8: Execute action a and observe φ and s′
9: for i← 1, 2, ..., nz do {Update ψ̃} 10: a′ ← argmaxbψ̃(s, b, zi)>zi {a′ ≡ πi(s′)} 11: θ α←− [ φ+ γψ̃(s′, a′, zi)− ψ̃(s, a, zi) ] ∇θψ̃

heading: 12: s← s′
text: 13: return θ Figure 1: USFA architecture

heading: 4 EXPERIMENTS
text: In this section we describe the experiments conducted to test the proposed architecture in a multitask setting and assess its ability to generalise to unseen tasks.

heading: 4.1 ILLUSTRATIVE EXAMPLE: TRIP MDP
text: We start with a simple illustrative example to provide intuition on the kinds of generalisation provided by UVFAs and SF&GPI. We also show how in this example USFAs can effectively leverage both
types of generalisation and outperform its precursors. For this we will consider the simple two-state MDP depicted in Figure 2. To motivate the example, suppose that state s1 of our MDP represents the arrival of a traveler to a new city. The traveler likes coffee and food and wants to try what the new city has to offer. In order to model that, we will use features φ ∈ R2, with φ1 representing the quality of the coffee and φ2 representing the quality of the food, both ranging from 0 to 1. The traveler has done some research and identified the places that serve the best coffee and the best food; in our MDP these places are modelled by terminal states associated with actions ‘C’ and ‘F ’ whose respective associated rewards are φ(·, C) = φ(C) = [1, 0] and φ(F ) = [0, 1]. As one can infer from these feature vectors, the best coffee place does not serve food and the best restaurant does not serve coffee (at least not a very good one). Nevertheless, there are other places in town that serve both; as before, we will model these places by actions Pi associated with features φ(Pi). We assume that ‖φ(Pi)‖2 = 1 and consider N = 5 alternative places Pi evenly spaced on the preference spectrum. We model how much the traveler wants coffee and food on a given day by w ∈ R2. If the traveler happens to want only one of these (i.e. w ∈ {[1, 0], [0, 1]}), she can simply choose actions ‘C’ or ‘F ’ and get a reward r = φ(·)>w = 1. If instead she wants both coffee and food (i.e. if w is not an “one-hot” vector), it may actually be best to venture out to one of the other places. Unfortunately, this requires the traveler to spend some time researching the area, which we model by an action ‘E’ associated with feature φ(E) = [− ,− ]. After choosing ‘E’ the traveler lands on state s2 and can now reach any place in town: C, F , P1, ..., PN . Note that, depending on the vector of preferences w, it may be worth paying the cost of φ(E)>w to subsequently get a reward of φ(Pi)>w (here γ = 1).
In order to assess the transfer ability of UVFAs, SF&GPI and USFAs, we define a training set M = {10, 01} and K = 50 test tasks corresponding to directions in the two-dimensional w-space: M′ = {w′|w′ = [cos( πk2K ), sin( πk 2K )], k = 0, 1, ...,K}. We start by analysing what SF&GPI would do in this scenario. We focus on training task wC = [1, 0], but an analogous reasoning applies to task wF = [0, 1]. Let πC be the optimal policy associated with task wC . It is easy to see that π(s1) = π(s2) = C. Thus, under πC it should be clear that QπCw′ (s1, C) > Q πC w′ (s1, E) for all test tasks w′. Since the exact same reasoning applies to task wF if we replace action C with action F , the GPI policy (3) computed over {ψπC ,ψπF } will be suboptimal for most test tasks inM′. Training a UVFA on the same setM, will not be perfect either due to the very limited number of training tasks. Nonetheless the smoothness in the approximation allows for a slightly better generalisation inM′. Alternatively, we can use Algorithm 1 to train a USFA on the training setM. In order to do so we sampled nz = 5 policies z ∈ R2 using a uniformly random distribution Dz(·|w) = U([0, 1]2) (see line 5 in Algorithm 1). When acting on the test tasks w′ we considered two choices for the candidates set: C = {w′} and C = {zi|zi ∼ U([0, 1]2), i = 1, 2, ..., 5}. Empirical results are provided in Figure 2. As a reference we report the performance of SF&GPI using the true {ψπC ,ψπF } – no approximation. We also show the learning curve of a UVFA. As shown in the figure, USFA clearly outperforms its precursors and quickly achieves near optimal performance. This is due to two factors. First, contrary to vanilla SF&GPI, USFA can discover and exploit the rich structure in the policyspace, enjoying the same generalisation properties as UVFAs but now enhanced by the combination of the off-policy and off-task training regime. Second, the ability to sample a candidate set C that induces some diversity in the policies considered by GPI overcomes the suboptimality associated with the training SFs ψπC and ψπF . We explore this effect in a bit more detail in the suppl. material.

heading: 4.2 LARGE SCALE EXPERIMENTS
text: Environment and tasks. We used the DeepMind Lab platform to design a 3D environment consisting of one large room containing four types of objects: TVs, balls, hats, and balloons (Beattie et al., 2016; Barreto et al., 2018). A depiction of the environment through the eyes of the agent can be seen in Fig. 3. Features φi are indicator functions associated with object types, i.e., φi(s, a, s
′) = 1 if and only if the agent
collects an object of type i (say, a TV) on the transition s a−→ s′. A task is defined by four real numbers w ∈ R4 indicating the rewards attached to each object. Note that these numbers can be negative, in which case the agent has to avoid the corresponding object type. For instance, in task w = [1-100] the agent is interested in objects of the first type and should avoid objects of the second.
Agent architecture. A depiction of the architecture used for the USFA agent is illustrated in Fig. 1. The architecture has three main modules: i) A fairly standard input processing unit composed of three convolution layers and an LSTM followed by a non-linearity (Schmidhuber, 1996); ii) A policy conditioning module that combines the state embedding coming from the first module, s ≡ f(h), and the policy embedding, z, and produces |A| · d outputs corresponding to the SFs of policy πz, ψ̃(s, a, z); and iii) The evaluation module, which, given a task w and the SFs ψ̃(s, a, z), will construct the evaluation of policy πz on w, Q̃(s, a,w, z) = ψ̃(s, a, z)>w.
Training and baselines. We trained the above architecture end-to-end using a variation of Alg. 1 that uses Watkins’s (1989) Q(λ) to apply Q-learning with eligibility traces. As for the distributionDz used in line 5 of Alg. 1 we adopted a Gaussian centred at w: z ∼ N (w, 0.1 I), where I is the identity matrix. We used the canonical vectors of R4 as the training set,M = {1000, 0100, 0010, 0001}. Once an agent was trained onM we evaluated it on a separate set of unseen tasks,M′, using the GPI policy (3) over different sets of policies C. Specifically, we used: C = {w′}, which corresponds to a UVFA with an architecture specialised to (1); C =M, which corresponds to doing GPI on the SFs of the training policies (similar to (Barreto et al., 2017)), and C =M∪ {w′}, which is a combination of the previous two. We also included as baselines two standard UVFAs that do not take advantage of the structure (1); one of them was trained on-policy and the other one was trained off-policy (see supplement). The evaluation on the test tasksM′ was done by “freezing” the agent at different stages of the learning process and using the GPI policy (3) to select actions. To collect and process the data we used an asynchronous scheme similar to IMPALA (Espeholt et al., 2018).

heading: 4.3 RESULTS AND DISCUSSION
text: Fig. 4 shows the results of the agents after being trained onM. One thing that immediately stands out in the figure is the fact that all architectures generalise quite well to the test tasks. This is a surprisingly good result when we consider the difficulty of the scenario considered: recall that the agents are solving the test tasks without any learning taking place. This performance is even more impressive when we note that some test tasks contain negative rewards, something never experienced by the agents during training. When we look at the relative performance of the agents, it is clear that USFAs perform considerably better than the unstructured UVFAs. This is true even for the case where C = {w′}, in which USFAs essentially reduce to a structured UVFA that was trained by decoupling tasks and policies. The fact that USFAs outperform UVFAs in the scenario considered here is not particularly surprising, since the former exploit the structure (1) while the latter cannot. In any case, it is reassuring to see that our model can indeed exploit such a structure effectively. This result also illustrates a particular way of turning prior knowledge about a problem into a favourable inductive bias in the UVFA architecture.
It is also interesting to see how the different instantiations of USFAs compare against each other. As shown in Fig. 4, there is a clear advantage in includingM to the set of policies C used in GPI (3). This suggests that, in the specific instantiation of this problem, the type of generalisation provided by SF& GPI is more effective than that associated with UVFAs. One result that may seem counter-intuitive at
first is the fact that USFAs with C =M+ {w′} sometimes perform worse than their counterparts using C =M, especially on tasks with negative rewards. Here we note two points. First, although including more tasks to C results in stronger guarantees for the GPI policy, strictly speaking there are no guarantees that the resulting policy will perform better (see Barreto et al.’s Theorem 1, 2017). Another explanation very likely in the current scenario is that errors in the approximations ψ̃(s, a, z) may have a negative impact on the resulting GPI policy (3). As argued previously, adding a point to C can sometimes increase the upper bound in (4), if the approximation at this point is not reliable. On the other hand, comparing USFA’s results using C = M + {w′} and C = {w′}, we see that by combining the generalisation of UVFAs and GPI we can boost the performance of a model that only relies on one of them. This highlights the fine balance between the two error terms in (4) and emphasizes how critical selecting low-error candidates in C can be. In the above scenario, SF&GPI on the training set M seems to provide a more effective way of generalising, as compared to UVFAs, even when the latter has a structure specialised to (1). Nevertheless, with less conservative choices of Dz that provide a greater coverage of the z space we expect the structured UVFA (C = {w′}) to generalise better. Note that this can be done without changingM and is not possible with conventional UVFAs. One of the strengths of USFAs is exactly that: by disentangling tasks and policies, one can learn about the latter without ever having to actually try them out in the environment. We exploit this possibility to repeat our experiments now using Dz = N (w, 0.5 I). Results are shown in Fig.5. As expected, the generalisation of the structured UVFA improves considerably, almost matching that of GPI. This shows that USFAs can operate in two regimes: i) with limited coverage of the policy space, GPI overM will provide a reliable generalisation; ii) with a broader coverage of the space structured UVFAs will do increasingly better.1
1Videos of USFAs in action on the links https://youtu.be/Pn76cfXbf2Y and https://youtu.be/0afwHJofbB0.

heading: 5 RELATED WORK
text: Multitask RL is an important topic that has generated a large body of literature. Solutions to this problem can result in better performance on the training set (Espeholt et al., 2018), can improve data efficiency (Teh et al., 2017) and enable generalisation to new tasks. For a comprehensive presentation of the subject please see Taylor and Stone (2009) and Lazaric (2012) and references therein.
There exist various techniques that incorporate tasks directly into the definition of the value function for multitask learning (Kaelbling, 1993; Ashar, 1994; Sutton et al., 2011). UVFAs have been used for zero-shot generalisation to combinations of tasks (Mankowitz et al., 2018; Hermann et al., 2017), or to learn a set of fictitious goals previously encountered by the agent (Andrychowicz et al., 2017).
Many recent multitask methods have been developed for learning subtasks or skills for a hierarchical controller (Vezhnevets et al., 2017; Andreas et al., 2016; Oh et al., 2017). In this context, Devin et al. (2017) and Heess et al. (2016) proposed reusing and composing sub-networks that are shared across tasks and agents in order to achieve generalisation to unseen configurations. Finn et al. (2017) uses meta-learning to acquire skills that can be fine-tuned effectively. Sequential learning and how to retain previously learned skills has been the focus of a number of investigations (Kirkpatrick et al., 2016; Rusu et al., 2016). All of these works aim to train an agent (or a sub-module) to generalise across many subtasks. All of these can be great use-cases for USFAs.
USFAs use a UVFA to estimate SFs over multiple policies. The main reason to do so is to apply GPI, which provides a superior zero-shot policy in an unseen task. There have been previous attempts to combine SFs and neural networks, but none of them used GPI (Kulkarni et al., 2016; Zhang et al., 2016). Recently, Ma et al. (2018) have also considered combining SFs and UVFAs. They propose building a goal-conditioned policy that aims to generalise over a collection of goals. In their work, the SFs are trained to track this policy and only used to build the critic employed in training the goal-conditioned policy. Thus, they are considering the extrapolation in π(s, g) and using the SFs as an aid in training. Moreover, as both the training and SFs and π(s, g) are done on-policy, the proposed system has only seen instances where the SFs, critic and policy are all conditioned on the same goal. In contrast, in this work we argue and show the benefits of decoupling the task and policy to enable generalisation via GPI when appropriate, while preserving the ability to exploit the structure in the policy space. We use the SFs as a way to factorize and exploit effectively the structure in value function space. And we will use these evaluations directly to inform our action selection

heading: 6 CONCLUSION
text: In this paper we presented USFAs, a generalisation of UVFAs through SFs. The combination of USFAs and GPI results in a powerful model capable of exploiting the same types of regularity exploited by its precursors: structure in the value function, like UVFAs, and structure in the problem itself, like SF&GPI. This means that USFAs can not only recover their precursors but also provide a whole new spectrum of possible models in between them. We described the choices involved in training and evaluating a USFA and discussed the trade-offs associated with these alternatives. To make the discussion concrete, we presented two examples aimed to illustrate different regimes of operation. The first example embodies a MDP where the generalisation in the optimal policy space is fairly easy but the number of optimal policies we would want to represent can be large. This is a scenario where UVFAs would strive, while vanilla SF&GPI will struggle due to the large number of policies needed to build a good GPI policy. In this case, we show that USFAs can leverage the sort of parametric generalisation provided by UVFAs and even improve on it, due to its decoupled training regime and the use of GPI in areas where the approximation is not quite perfect. Our second example is in some sense a reciprocal one, where we know from previous work that the generalisation provided via GPI can be very effective even on a small set of policies, while generalising in the space of optimal policies, like UVFAs do, seems to require a lot more data. Here we show that USFAs can recover the type of generalisation provided by SFs when appropriate. This example also highlights some of the complexities involved in training at scale and shows how USFAs are readily applicable to this scenario. Overall, we believe USFAs are a powerful model that can exploit the available structure effectively: i) the structure induced by the shared dynamics (via SFs), ii) the structure in the policy space (like UVFAs) and finally iii) the structure in the RL problem itself (via GPI), and could potentially be useful across a wide range of RL applications that exhibit these properties.

heading: A TWO TYPES OF GENERALISATION: INTUITION
text: In this paper we argue that one of the main benefits provided by USFAs is the ability to combine the types of generalisation associated with UVFAs and GPI. In this section, we will take a close look at a very simple example to illustrate the two of generalisation we are considering and how they are different. This is a very small example where the number of optimal policies are very limited and the induced tasks are not that interesting, but we chose this solely to illustrate the decision process induced by GPI and how it differs from parametric generalisation in w via a functional approximator (FA).
Let us consider a an MDP with a single state s and two actions. Upon executing action a1 the agent gets a reward of 0 and remains in state s; the execution of action a2 leads to a potentially non-zero reward followed by termination. We define unidimensional features φ ∈ R as φ(s, a1) = 0 and φ(s, a2) = 1. A task is thus induced by a scalar w ∈ R which essentially re-scales φ(s, a2) and defines the reward rw = w the agent receives before termination. In this environment, the space of tasks considered are induced by a scalar w ∈ R. In this space of tasks, one can easily see that there are only two optimal policies: taking action a1 and receiving the reward rw = w if w ≤ 0, or taking action a0 and remaining in s0 indefinitely. Thus the space of optimal value functions is very simple. For convenience, we include a depiction of this space in Figure 6.
Suppose now we are in the scenario studied in the paper, where after training on a set of tasks M the agent should generalise to a test task w′. Specifically, let us consider three points in this spaceM = {w1,w2,w3}– three tasks we are going to consider for learning and approximating their optimal policies {Q̃∗w1 , Q̃ ∗ w1 , Q̃ ∗ w1}. Given these three points we are going to fit a parametric function that aims to generalise in the space of w. A depiction of this is included in Figure 7(a). Now, given a new point w′ we can obtain a zero-shot estimate Q̃∗w′ for Q ∗ w′ – see Figure 7(b). Due to approximation error under a very limited number of training points, this estimate will typically not recover perfectly Q∗w′ = 0. In the case of UVFA (and other FAs trying to generalise in task space), we are going to get a guess based on optimal value function we have built, and we are going to take decision based on this estimate Q̃∗w′ .
Given the same base tasks M = {w1,w2,w3} we can now look at what the other method of generalisation would be doing. We are going to denote by πzi the (inferred) optimal policy of task wi. Since we have learnt the SFs corresponding to all of these policies ψπzi , we can now evaluate how well each of these policies would do on the current test task w′: Qπzw′(s, a) = ψ
πz (s, a)Tw′ for all z ∈M. A depiction of this step is included in Figure 8(a). Given these evaluations of previous behaviours, GPI takes the maximum of these values – ”trusting”, in a sense, the most promising value. In our case this corresponds to the behaviour associated with task w3, which in this case happens to
have the same optimal policy as our test task w′. Thus in this particular example, the evaluation of a previously learned behaviour gives us a much better basis for inducing a behaviour in our test task w′. Moreover if the SFs are perfect we would automatically get the optimal value function for w′.
It is worth noting that, in this case, by learning a USFA we can recover both scenarios described above based on our choice of the candidate set for GPI C. In particular, if C = {w′} we recover the mechanism in Figure 7, while for C =M we recover the generalisation in Figure 8. Furthermore, for any choice of trained points that include one positive and one negative choice of w, by relying on GPI we can generalise perfectly to the whole space of tasks, while an approach based exclusively on the sort of generalisation provided by UVFAs may struggle to fit the full function. Analogously, in scenarios where the structure of the optimal space favours (UV)FAs, we expect USFAs to leverage this type of generalisation. An example of such a scenario is given in the first part of the experimental section – Section 4.1, and further details in Section B.
B ILLUSTRATIVE EXAMPLE: TRIP MDP
In this section we provide some additional analysis and results omitted from the main text. As a reminder, this is a two state MDP, where the first state is a root state, the transition from s1 → s2 comes at a cost rw(s1, E) = φ(s1, E)Tw = − (w1 + w2) and all other actions lead to a final positive reward corresponding to how much the resulting state/restaurant alligns with our preferences (our task) right now. For convenience, we provide below the depiction of the Trip MDP introduced in Section 4.1.
Arrived in new city
Explored neighborhood
1
0
0
1
-ε -ε
.8
.6
.6
.8
.7
.7
...
S₁ S₂E
C F
In the experiments run we considered a fixed set of training tasksM = {01, 10} for all methods. The set of outcomes from the exploratory state s2 is defined as φ(s2, a) = [cos(θ), sin(θ)] for θ ∈ {kπ/2N}k=0,N . Note that this includes the binary states for k = 0 and respectively k = N . We ran this MDP with N = 6, and = 0.05. Thus outside the binary outcomes, the agent can select N − 1 = 5 other mixed outcomes and, as argued in the main text, under these conditions there will be a selection of the w-space in which each of these outcomes will be optimal. Thus the space of optimal policies, we hope to recover, is generally N + 1. Nevertheless, there is a lot of structure in this space, that the functional approximators can uncover and employ in their generalization.
B.1 ADDITIONAL RESULTS
In the main paper, we reported the zero-shot aggregated performance over all direction M′ = {w′|w′ = [cos( πk2K ), sin( πk 2K )], k ∈ ZK}. This should cover most of the space of tasks/trade-offs we would be interest in. In this section we include the generalization for other setsM′. First in Fig. 9 we depict the performance of the algorithms considered across the whole w′ spaceM′ = [0, 1]2. Fig. 10 is just a different visualization of the previous plot, where we focus on how far these algorithms are from recovering the optimal performance. This also shows the subtle effect mentioned in the discussion in the main text, induced by the choice of C in the USFA evaluation.
A particularly adversarial choice of test tasks for the vanilla SF&GPIwould be the diagonal in the [0, 1]2 quadrant depicted in the plot above: M′ = {w′|w′1 = w′2, w1 ∈ [0, 1]}. This is, in a sense, maximally away from the training tasks and both of the precursor models are bound to struggle in this portion of the space. This intuition was indeed empirically validated. Results are provided in Fig. 11. As mentioned above, this is an adversarial evaluation, mainly to point out that, in general, there might be regions of the space were the generalization of the previous models can be very bad, but where the combination of them can still recover close to optimal performance.

heading: C LARGE SCALE EXPERIMENTS: DETAILS
text: C.1 AGENT’S ARCHITECTURE
This section contains a detailed description of the USFA agent used in our experimental section. As a reminder, we include the agent’s architecture below (Figure 1 in the main text).
As highlighted in Section 4.2, our agent comprises of three main modules:
• Input processing module: computes a state representation f(ht) from observation ot. This module is made up of three convolutional layers (structure identical to the one used in (Mnih et al., 2015)), the output of which then serves as input to a LSTM (256). This LSTM takes as input the previously executed action at−1. The output of the LSTM is passed through a non-linearity f (chosen here to be a ReLu) to produce a vector of 128 units, f(ht).
• Policy conditioning module: compute the SFs ψ̃(s, a, z), given a (sampled) policy embedding z and the state representation f(ht). This module first produces nz number of z ∼ Dz samples (nz = 30 in our experiments). Each of these is then transformed via a 2-layer MLP(32,32) to produce a vector of size 32, for each sample z. This vector g(z) gets concatenated with the state representation f(ht) and the resulting vector is further processed by a 2-layer MLP that produces a tensor of dimensions d× |A| for each z, where d = dim(φ). These correspond to SFs ψ̃(s, a, z) for policy πz. Note that this computation can be done quite efficiently by reusing the state embedding f(ht), doing the downstream computation in parallel for each policy embedding z.
• Task evaluation module: computes the value function Q(s, a, z,w) = ψ̃(s, a, z)Tw for a given task description w. This module does not have any parameters as the value functions are simply composable from ψ̃(s, a, z) and the task description w via assumption (1). This module with output nz value functions that will be used to produce a behavior via GPI.
An important decision in this design was how and where to introduce the conditioning on the policy. In all experiments shown here the conditioning was done simply by concatenating the two embeddings w and z, although stronger conditioning via an inner product was tried yielding similar performance.
The ’where’ on the other hand is much more important. As the conditioning on the policy happens quite late in the network, most of the processing (up to f(ht)) can be done only once, and we can sample multiple z and compute the corresponding ψ̃
πz at a fairly low computational cost. As mentioned above, these will be combined with the task vector w to produce the candidate action value functions for GPI. Note that this helps both in training and in acting, as otherwise the unroll of the LSTM would be policy conditioned, making the computation of the SFs and the off-policy n-step learning quite expensive. Furthermore, if we look at the learning step we see that this step can also benefit from this structure, as the gradient computation of f can be reused. We will only have a linear dependence on nz on the update of the parameters and computations in the red blocks in Figure 12.
UVFA baseline agents have a similar architecture, but now the task description w is fed in as an input to the network. The conditioning on the task of UVFAs is done in a similar fashion as we did the conditioning on the policies in USFAs, to make the computational power and capacity comparable. The input processing module is the same and now downstream, instead of conditioning on the policy embedding z, we condition on task description w. This conditioning if followed by a 2-layer MLP that computes the value functions Q̃∗(s, a,w), which induces the greedy policy π (UV FA) w = argmaxa Q̃ ∗(s, a,w).
C.2 AGENT’S TRAINING
The agents’ training was carried out using the IMPALA architecture (Espeholt et al., 2018). On the learner side, we adopted a simplified version of IMPALA that uses Q(λ) as the RL algorithm. In our experiments, for all agents we used λ = 0.9. Depending on the sampling distribution Dz, in learning we will be often off-policy. That is, most of the time, we are going to learn about a policy πz1 and update its corresponding SFs approximations ψ̃(s, a, z1), using data generated by acting in the environment according to some other policy πz2 . In order to account for this off-policiness, whenever computing the n-step return required in eq. 5, we are going to cut traces whenever the policies start to disagree and bootstrap from this step on (Sutton and Barto, 1998). Here we can see how the data distribution induce by the choice of training tasksM can influence the training process. If the data distribution Dz is very close to the setM, as in our first experiment, most of the policies we are going to sample will be close to the policies that generated the data. This means that we might be able to make use of longer trajectories in this data, as the policies will rarely disagree. On the other hand, by staying close to the training tasks, we might hurt our ability to generalise in the policy space, as our first experiment suggest (see Figure 4). By having a broader distribution Dz = N (w, 0.5I), we can learn about more diverse policies in this space, but we will also increase our off-policiness. We can see from Figure 5, that our algorithm can successfully learn and operate in both of these regimes.
For the distributed collection of data we used 50 actors per task. Each actor gathered trajectories of length 32 that were then added to the common queue. The collection of data followed an -greedy policy with a fixed = 0.1. The training curves shown in the paper correspond to the performance of the the -greedy policy (that is, they include exploratory actions of the agents).
C.3 AGENT’S EVALUATION
All agents were evaluated in the same fashion. During the training process, periodically (every 20M frames) we will evaluate the agents performance on a test of held out test tasks. We take these intermediate snapshots of our agents and ’freeze’ their parameters to assess zero-shot generalisation. Once a test task w′ is provided, the agent interacts with the environment for 20 episodes, one minute each and the average (undiscounted) reward is recorded. These produce the evaluation curves in Figure 4. Evaluations are done with a small = 0.001, following a GPI policy with different instantiations of C. For the pure UVFA agents, the evaluation is similar: -greedy on the produced value functions Q̃∗(s, a,w), with the same evaluation = 0.001.

heading: D ADDITIONAL RESULTS
text: In our experiments we defined a set of easy test tasks (close toM) and a set of harder tasks, in order to cover reasonably well a few distinct scenarios:
• Testing generalisation to tasks very similar to the training set, e.g. w′ = [0, 0.9, 0, 0.1]; • Testing generalisation to harder tasks with different reward profiles: only positive rewards,
only negative rewards, and mixed rewards.
In the main text, we included only a selection of these for illustrative purposes. Here we present the full results.
D.1 CANONICAL BASIS: ZERO-SHOT GENERALISATION
This section contains the complete results of the first experiment conducted. As a reminder, in this experiment we were training a USFA agent on M = {1000, 0100, 0010, 0001}, with Dz = N (w, 0.1I) and compare its performance with two conventional UVFA agents (one trained on-policy and the other one using all the data generated to learn off-policy) on a range of unseen test tasks. Complete set of result is included below, as follows: Figure 13 includes results on easy tasks, close to the tasks contained in the training setM (generalisation to those should be fairly straightforward); Figure 14 and Figure 15 present results on more challenging tasks, quite far away from the training set, testing out agents ability to generate to the whole 4D hypercube.
D.2 CANONICAL BASIS: USFAS IN DIFFERENT TRAINING REGIMES.
In this section, we include the omitted results from our second experiment. As a reminder, in this experiment we were training two USFA agents on the same set of canonical tasks, but employing different distributions Dz, one will low variance σ = 0.1, focusing in learning policies around the training setM, and another one with larger variance σ = 0.5, that will try to learn about a lot more policies away from the training set, thus potentially facilitating the generalisation provided by the UVFA component. Results are displayed in Figures 16-17 on all tasks in the hard evaluation set.
D.3 LARGER COLLECTION OF TRAINING TASKS
We also trained our USFA agent on a larger set of training tasks that include the previous canonical tasks, as well as four other tasks that contain both positive and negative rewardM = {1000, 0100, 0010, 0001, 1-100, 01-10, 001-1, -1000}. Thus we expect this agent to generalises better as a result of its training. A selection of these results and sample performance in training are included in Fig. 18.

Review summaries:
The goal here is multi-task learning and generalization, assuming that the expected one-step reward for any member of the task family can be written as $\phi(s,a,s')^T w$. The authors propose universal successor features (USF) $\psi$s, such that the action-value functions Q can be written as $Q(s,a,w,z)=\psi(s,a,z)^T w$, generalizing over mutiple tasks each denoted by $w$, and multiple policies each denoted by $z$. Here, $z$ represents the optimal policy induced by a reward specified by $z$ (from the same set as $w$). Using USFs $\psi$-s, the Q values can be interpolated across policies and tasks. Due to the disentangling of reward and policy generalizations, the training sets for $w$ and $z$ can be independently sampled. The authors further generalize a temporal difference error in these USFs $\psi$s, using the TD error to learn to approximate the $\psi$s by a network (USF Approximator i.e USFA). They then test the generalization capabilities of these USFAs on families of a simple task and a DeepMind Lab based task.