Combining Behaviors with the Successor Features Keyboard
Authors: Wilka Carvalho
Andre Saraiva
Angelos Filos
Andrew Kyle Lampinen
Loic Matthey
Richard Lewis
Honglak Lee
Satinder Singh
Danilo Jimenez Rezende
Daniel Zoran
Venue: NeurIPS 2023 poster (2023)
Keywords: deep reinforcement learning
successor features
transfer
generalization
feature-discovery
TL;DR: We are the first method to enable transfer with successor features in a complex 3D environment when all needed representations are discovered

Abstract: The Option Keyboard (OK) was recently proposed as a method for transferring behavioral knowledge across tasks. OK transfers knowledge by adaptively combining subsets of known behaviors using Successor Features (SFs) and Generalized Policy Improvement (GPI).
However, it relies on hand-designed state-features and task encodings which are cumbersome to design for every new environment.
In this work, we propose the "Successor Features Keyboard" (SFK), which enables transfer with discovered state-features and task encodings.
To enable discovery, we propose the "Categorical Successor Feature Approximator" (CSFA), a novel learning algorithm for estimating SFs while jointly discovering state-features and task encodings.
With SFK and CSFA, we achieve the first demonstration of transfer with SFs in a challenging 3D environment where all the necessary representations are discovered.
We first compare CSFA against other methods for approximating SFs and show that only CSFA discovers representations compatible with SF&GPI at this scale.
We then compare SFK against transfer learning baselines and show that it transfers most quickly to long-horizon tasks.

Combining Behaviors with the Successor Features Keyboard
Wilka Carvalho*
Andre Saraiva² Loic Matthey²
Angelos Filos?
Andrew Kyle Lampinen²
Richard L. Lewis
Honglak Lee
Satinder Singh Danilo J. Rezende
Daniel Zoran
Harvard University
Google DeepMind
University of Michigan
Abstract
The Option Keyboard (OK) was recently proposed as a method for transferring be- havioral knowledge across tasks. OK transfers knowledge by adaptively combining subsets of known behaviors using Successor Features (SFs) and Generalized Policy Improvement (GPI). However, it relies on codings which are cumbersome to design for every new environment. this work, we propose the "Successor Features Keyboard" (SFK), which enables transfer with discovered state-features and task encodings. To enable discovery, we propose the "Categorical Successor Feature Approximator" (CSFA), a novel learning algorithm for estimating SFs while jointly discovering state-features and task encodings. With SFK and CSFA, achieve the first demonstration of transfer with SFs in a challenging 3D environment where all the necessary representations are discovered. We first compare CSFA against other methods for approximating SFs and show that only CSFA discovers representations compatible with SF&GPI at this scale. We then compare SFK against transfer learning baselines and show that it transfers most quickly to long- long-horizon tasks.
state-features and task
Introduction
Transfer Task: Cook tomato and potato, then place them on plate
Quary
Query
Cook tomato
Place held objects
Consider household robot that learns tasks for interacting with objects such as finding and moving them around. When this robot is deployed to house and needs to perform combinations of these tasks, collecting data for reinforcement learning (RL) will be expensive. Thus, ideally this robot can effectively transfer its knowledge to efficiently learn these novel tasks with minimal interactions in the environment. We study this form of transfer in Deep RL.
and potato
on plate
Legend
Query: w
Representation
"keys" for known
GPI
GPI
SF-Key:
Beook CLT
behaviors
"values"
Valuere) ① w
used to
generate behavior a1, a2,
Higheet
posato
Vplece.plast
generate behavlor
Initial state
Later state
Figure 1: Diagram of transfer with the Successor Features Keyboard (SFK). SFK uses SF-based "keys" represent behaviors by how much of some feature (known as "cumulant") they obtain. SFK dynamically selects from behaviors with GPT by generating prefer- ence "queries" over these features. The "value" each behavior then computed with a dot-product and the highest value is used to generate behavior. Prior work has hand-designed the features that define queries and SFs. Here, we study the problem of transferring with SFK while discovering all necessary representations.
One promising method for transfer is the Option Keyboard (OK) 22, which trans- fers to new tasks by adaptively combining subsets of known behaviors. OK combines known behaviors by leveraging Successor Features (SFs) and Generalized Policy Im- provement (GPI) 3 4 SFs are predictive
"Contact author: wcarvalho @ h.harvard.edu Work done during internship.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

representations for behaviors. They represent behaviors with estimates of how much state-features (known "cumulants") will be experienced given that behavior. GPI can be thought of as a query- key-value system that, given feature preferences, selects from behaviors that obtain those features.
While OK is a promising transfer method, it relies on hand-designed representations for the "cumu- lants" and task feature- ceaturr-preferencee (i.e. task encodings). There work that has discovered either cumulants or task encodings G. However, either (a) they only showed transfer with SF&GPI using a fixed (not dynamic) transfer query, (b) they only demonstrated results simple grid-worlds where an agent combines short-horizon "goto" tasks (e.g. goto A and B), or (c) they leveraged separate networks for the SFs of each task-policy and for each task-encoding. This means parameter count scales with the number of tasks and limits representation re-use across tasks.
Ideally, we can transfer with dynamic query while leveraging discovered representations for cumulants and feature preferences. Further, a transfer method should scale to sparse-reward long- horizon tasks such as those found in complex 3D environments. However, jointly discovering cumulants and feature preferences while estimating SFs is challenging in this setting. First, jointly learning cumulants and SFs involves estimating boot-strapped returns over a non-stationary target with high-variance and a shifting magnitude [ To maximize knowledge sharing across tasks, we can share our SF-estimator and task encoder across tasks. However, no work has yet achieved transfer results doing so.
To transfer with dynamic query that leverages discovered representations while sharing functions across tasks, we propose two novel methods, the Successor Features Keyboard (SFK) and the Categor- ical Successor Feature Approximator (CSFA). We present high-level overview of in SFK in Figure]] SFK leverages CSFA to learn SF-estimates over discovered cumulants and task-preferences in pretraining phase. Afterwards, finetuning phase. SFK learns to generate dynamic queries which are linear combinations of CSFA -SFA-discovered pref-rrefesssces CSFA addresses challenges with estimating non-stationary return by estimating SFs with a variant of the categorical two-hot repre- sentation introduced by MuZero 77. We discretize the space of cumulant-return values into bins and learns probability mass function (pmf) over them. Modelling cumulant-returns with a pmf is more robust to outliers and can better accomodate shifting magnitude. In contrast, standard methods estimate SFs with regression of a single point-estimate varying scales [9]
S which is succeptible to outliers and
We study SFK and CSFA in Playroom [0] a challenging 3D environment with high- high-dimensional pixel observations and hang-horioon tasks defined by sparse rewards. Prior work on transfer with SF&GPI has mainly focused on transfer in simpler 2D environments [I] B. While Borsa et al. 8 studied transfer in a 3D environment, they relied on hand-designed cumulants and task encodings g and only transferred to combinations of "Goto" tasks. We discover cumulants and task encodings while studying transfer to combinations of long-horizon, sparse-reward "Place near" tasks.
Contributions. (1) We propose the Successor Features Keyboard, a novel method that transfers with SP&GPI using a dynamic query, discovered representations, and a task encoder and ST-approximator that are shared across tasks. (2) To enable discovery when sharing a task encoder and ST-approximator cross tasks, we propose a novel learning algorithm, the Categorical Successor Feature Approximator. (3) We present the first demonstration of transfer with successor features in a complex 3D environment where all the necessary representations are discovered.
2 Related work on Transfer in Deep RL
Several avenues exist to transfer knowledge in Deep RL. We can transfer an agent's representations (how they represent situations), their control policy (how they act in situations), or their value function (how they evaluate situations). To transfer representations, one can learn mapping from source domains to target domains
[2] learn disentangled representations [3 14, or learn a modular architecture [T5] [6]. To transfer a control policy, some methods distill knowledge from a source policy to a target policy [17], others exploit policy improvement [18) and a third set transfer low-level policies by learning meta-controller [19] Finally, to transfer value functions, some approaches learn universal value functions 20 and others exploit SFs 22. Below we review approaches most
closely related to ours.

---

Transferring policies. One strategy to transfer a policy is to leverage multi-task RL (MTRL) training where you learn and transfer coal-condi policy [21]. Another stratgy is to distill knowledge from one policy to another, as Distral does [22]. Distral works by learning two policies: one goal- conditioned policy and another goal-agnostic "centroid policy". The action- likelloiiklhhoods of each policy are then distilled into the other by minimizing KL-divergences. Distral has strong performance in multi-task settings but it relies on the utility of "centroid" policy for sharing knowledge across tasks. When we transfer to longer horizon tasks with sparse rewards, neither MTRL nor Distral may provide poliey with good jumpstart performance [23). In this work, study jumpstart performance and exploit successor features (SFs) with generalized policy improvement (GPI) 3
Method
disc. Φ
disc. w
query
share w
share the
3D
transfer
Barreto et al.
3
static X X
X
goto-n
Zhu et al. [24]
/
/
static x
x
/
goto-n
Barreto et al. [4]
static
x
/
goto-c
Filos et [5]
static
x
goto-c
Borsa et al.
8
x
static
x /
/
goto-c
Carvalho et al. [6
X
static
7
x
goto-c
Barreto et al.
dynamic
x
goto-c
SFK (ours)
dynamic
place-c
Table 1: Related methods for transfer with SF&GPI. SFK is the first method to transfer with
dynamic query, discover cumulants & and task encodings w, while sharing a task encoder wh
and SF approximator the across tasks. Each of these is important to transfer with SF&GPI in large-scale multi-task setting. Together, this enables the first SF method to transfer to combinations of long-horizon place tasks in 3D environment with discovered & and w. Note: 3 refers to methods which learn from demonstration data, which do not. goto-n is "goto new goal state", goto-c is "goto object combo", and place-c is "place object combo".
Successor Features are useful because they enable computing of action-values for new task encod- ings [3). When combined with GPI, prior work has shown strong zero-shot or few-shot transfer to combinations of tasks. To accomplish this, GPI can evaluate known SFs with "query" transfer task encoding. SF&GPI relies on "cumulants" (which SFs predict returns over) and task encodings that respect a dot-product relationship. Prior work has had one of three limitations. Some work has discovered representations but only shown transfer with static transfer query and did not share SF-estimators or task-encoders across tasks B 24 5) Other work has shared SF-estimators across tasks but exploited hand-designed task encodings with static GPI queries 8 6 The Option Keyboard [2] transferred using a dynamic query; however, they hand-designed cumulants and task encodings and didn't share functions across tasks. In this work, we present the Successor Features Keyboard, where we address all three limitations. We transfer with a dynamic query, discover cumulants and task encodings, and learn both a task-encoder and SF-estimator that shared across tasks. Additionally, prior work has only studied transfer to combinations of short-horizon "go to" tasks whereas include longer horizon "place" tasks and do so in a 3D) environment. We summarize these differences in Table
3 Background
We study RL agent's ability to transfer knowledge from a set of nh training tasks Ttrain (KI...., kn,) } to a set of n/, transfer tasks Tnew = ........,;; k²; }. During training, tasks are sampled from distribution Purain (T Trrsin). ). At transfer, tasks are sampled from distribution Peranater nTTex). Each task is specified as a Partially Observable Markov Decision Process (POMDP, (25)), M; (S" AS X'A R. P. fa), where S", A and A' are the environment state, action, and observation spaces. Rewards are parameterized by task description K, i.e. rft1 R(s at, s.41. k) is the reward for transition (ef a, 1. When the agent takes action in state s; sh+1 is sampled according to p(-1sf, (x), an observation F++1 is generated via fa (*f+1)), and the agent gets reward rê+ We assume the agent learns recurrent state function that maps histories to agent state representations, Given this learned state, we aim to obtain behavior policy ((,,, that maximises the expected reward when taking action o, in state s; i.c. that maximizes
A,
8, = sp (x,, St-1, .t-1), 04-1).
3

---

QT = as; [o 2 We study agents that continue learning during transfer and aim to maximize jump-start performance [23].
Transfer with SF&GPI requires two things: (1) state- features known as "cumulants" Φt+1
@ (St, dz, S1+1), which are useful "descriptions" of state- transtrannition, and (2) a task encoding w (K),,, which define "preferences" over said transitions. Reward is then defined as r = de w Successor Features y are then value functions that describe the discounted sum of future & that will be experienced under policy x.
8
2 (s,, 0t) E
Σ
7 detti1
Given 77, we can obtain action-values for r as Q 1/7 w.
The linear decomposition of QThN is interesting because it can be exploited re-evaluate 1h"t for new tasks with GPI. Assume we have learned SFs {1b (s, a) a) for tasks Terain. Given a new task n', can obtain new policy (ss; k²) with GPI in two steps: (1) re-evaluate each SF with the task's query encoding to obtain new Q-values (2) select an action using the highest Q-value. In summary,
I (s R E arg max
(QUin)) ( "wo(k')) arg max max aEAA VE(I....m.)
aEA iE(I....mg)
where (k')) is static transfer query for transfer task k².
Option Keyboard. One benefit of equation i that enables transfer to linear combinations of training task encodings. However, it has two limitations. First, the feature "preferences" (x())) are fixed across time. When n' is a complex task (e.g. avoiding an object at some time-points but going towards it at others), we may want something that is state-dependent. Second, if we learn wy with a nonlinear function approximator such as a neural network, there is no guarentee that wy(k') is in the span of training task encodings. The "Option Keyboard" [[] can circumvent these issues by learning transfer policy that maps states and tasks a dynamic transfer query go g0(8, K'):
M (1 K²) E arg max max aEA iE(I....,,,)
T (o(s),k))}
Learning. In the most general setting, we learn the Φø, wy and ge from experience. Rewards r" and their task encodings w wK(K) reference deterministic task policies To that maximize them. Borsa et al. [ showed that this allows us parameterize an SF-approximator for a policy Tw with an encoding of the task that defines it, i.e. we can approximate vh" with wo(s, a, w). This is known as a Universal Successor Feature Approximator (USFA) and can be learned with TD-learnig with cumulants as pseudo-rewards. To discover On and wp, we can match their dot dot-product to the experienced reward 3 Defining de+1 Φo at, 8t+1 ) we summarize this as:
C. = ll + (//0 4,41, w) (S a,, w)
C- ||r __,1o
No prior work has jointly learned task encoder up and USFA v/((s, a, w) while discovering cumulants dg. this work, we introduce Successor Features Keyboard to address these limitations to enable transfer with SF&GPI and discovered representations in large-scale 3D environment.
4 Method
We propose a novel method for transfer, the Successor Features Keyboard (SFK), where necessary representations are discovered. To discover representations, propose the Categorical Successor Feature Appreximator for jointly learning SFs the cumulants Φe. and task encodings Wg. The rest of this section is structured as follows. In $4.1 describe CSFA and how to leverage it for pretraining. Finally, in $4.2
we describe how to leverage SFK for transfer. We provide background in $A
4.1 Pretraining with a Categorical Successor Feature Approximator
We propose a novel learning algorithm, Categorical Successor Feature Appreximator (CSFA), com- posed of a novel architecture, shown in Figure and a novel learning objective (equation (6) Challenge: when jointly learning Universal SF-approximator th and a cumulant-network

---

Categorical Approximator gi (s, .)
Transfer behavior with Successor Features Keyboard
Training task behaviors we want to transter VI
F(A,7oux) Toux)
argmaax, max
(sfo(s; w)
go(s, Mone)
Plad
Hond
+ M
Structured, Categorical Network Via (,,0, A))
SF
Q
GPI
Task Encoder
Network (8,
State fn
State In
New State In
Transter policy
(5,
Σ kay
Task
Observation
Observation Encoder
Legend
New Task Encoder
Frozen
parameters
Ft
Zt
Transfer Task
Figure 2: Left: Categorical Successor Feature Approximator (CSFA) estimates SFs for a task encoding with structured, categorical network. Right: Successor Features Keyboard (SFK) transfers by dynamically selecting combinations of known task behaviors Werain = (Wi,.... Wn ). SFK accomplishes this by learning policy ga(s, (8, Wnex) that generates linear combinations of known task encodings Werain. SFK then leverages GPI to compute Q-values for known behaviors, Qui.. (s, a, w:) (s,s, Wneu) and acts using the highest Q-value.
for long-horizon tasks, the needs to fit @o-generated returns that are potentially high-variance, non- stationary, and changing in magnitude ( CSFA addresses this challenge by modelling SFs with a probability mass function (pmf) over discretized range of continous values. CSFA then fits this data by leveraging categorical cross-entropy loss, enabling our estimator to give probabil- ity mass to different ranges of values. This is in contrast to prior work that models SFs with a point-estimate that fit via regression modelling a stationto target with changing magnitude
g] Leveraging point-estimate can be unstable for
( (we show evidence in CC
M Pin Dren socen
Example probabiities and bin values
Architecture. CSFA represents SFs with a pmf (IB, p*)), where R = (bI,.... bbx) apriori defined set of bin values and t probabilities for each bin value. Specifically, CSFA estimates * a) E R" with tha a, w) and repre- sents the k-th dimension as v (s, a, w)
(s, 10)
P
pd softmax((olf (la(s (s, u. ex
n-dimensional SF-vector
learned probabilities
bin values
25
0
5
01
bM
Task Encoder
Dimension Encoder
Em pllbbm. At each time-step, we update state function sø with an encoding of the current observation
State fn
Observation
fooc(xx), the previ- ous action a,-1, and the previous state repre- sentation 8--1, i.e.
so(2;; (1-1,88-1) Each set of probabilities is computed as n = softmax(lo (St, w, ex)), where eg is an embed- ding for the current SF dimension. In summary,
Figure 3: Diagram of how successor features can be computed with probability mass function.
(0 a, w) (s, a, w)) (s,
Σ Pm Dm p t = so (5)
We provide diagram of representing SF with pmf in Figure Using pmf allows us re-use the same network to estimate SFs across cumulants. We hypothesize that this provides a stronger learning signal to stabilize learning across more challenging return estimates. We show evidence in F
Learning objective. We learn to generate behavior by employing variant of Q-learning. In particular, we generate Q-values using learned SFs, and use these Q-values create targets for both estimating Q-values and for estimating SFs. For both, targets correspond to the action which
5

---

maximized future Q-estimates. We learn SFs with categorical cross-entropy loss where we obtain targets from scalars with the twohot(-) operator. Intuitively, this represents a scalar with likelihoods across the two closest bins. In summary,
Di r++1 + 7o0o 18 a", wa (k)) woo (R)
La (1vo de; wo(k)) (x(()) v2. 11
(6)
c 1 log(p"t) (opg(pt))twwhootieee [twohot(y)" n Σ
de (e) + vill (sc41, wr (1))
a" = arg max, ( w) J wt where p softmax(1o w. et), w is stop-gradient op- eration w. Like prior work (26), mitigate non non-stationary in the return-target, y and Ut by having target parameters 8 that update a slower rate than 0. The overall loss is L= BQLQ + B + B,Lr.
Important implementation details. Estimating SFs while jointly learning a cumulant function and task encoder can be unstable in practice (3). No work has jointly learned all three functions while sharing them across tasks. Here, we detail important implementation challenges that prohibited us from discovering representations that worked with SF&GPI in our large-scale setting. D1. We found that passing gradients to we through ( (8, a, w) or through ( a, w) wo during Q-learning lead to dimensional collapse (27] and induces small angle between task encodings (see IC.T). We hypothesize that this makes ( a, w) unstable and manifests as poor GPI performance. D2. When |||||
is large, it can make (s, a, w) unstable. We hypothesize that this is because magnifies errors due to the ST-approximation error that loosen previously found bounds on GPI performance [8 We discuss this in more detail in $B To mitigate this, bound w by enforcing lie unit sphere,
<x((x) This makes the SF-error be consistent across tasks. [[]6(
ie. w-
4.2 Transfer with the Successor Features Keyboard
The original Option Keyboard (equation 3) learned a policy that mapped states 8, to queries w, w' = =g(st, x²). However, they used hand-designed wi and thus hand-designed space for In our setting, we learn Wg. However, GPI performance is bound by the distance of a transfer query w' to known preference vectors (3] (we discuss this in more detail in (B). Thus, we want to sample w' that are not too "far" from known w. To accomplish this, we shift from learning a policy that samples preference vectors to a policy that samples coefficients fa, Vi for known preference vectors Warain = { [W1.... Wn.) }. The GPI query is then computed as a weighted sum. Below we describe this policy in more detail along with how to learn it.
At each time-step, the agent uses a pretrained CSFA to compute SFs ypt (se, a) for w. E Werain' {1lTTO = the (8,, a, w:)) w.W..........
2e = fe(()))
8, = s0(2;: at-1; 8t-1)
In our experiments, we freeze the observation encoder, state function, and task encoder and learn a new state function and task encoder at transfer time with parameters 0n. Given a new state representation sp, we sample coefficients independently:
-/= = go, (5f, wo, (xx*))
arwy
where
at pa, (ap,(("'ss (a"1sf, W9, (r²))
(9)
We find that a Bernoulli distribution performs well. We learn this a-coefficient policy with policy gradients (28 by performing gradient ascent with gradients A, V log py (a st, w). where A, Rg - Ve, (T) is the "advantage" of the coefficient chosen at time t. Here, R; = Dimo Tt+i+1 is the experienced return and Ve, (s7) is the predicted return. Optimizing this increases the likelihood of choosing coefficients in proportion to Ag,
5 Experiments
We study transfer with sparse-reward long-horizon tasks in the complex 3D Playroom environ- ment [10). To transfer behavioral knowledge, we propose SFK for combining behaviors with SF&GPI, and CSFA for discovering the necessary representations. In $5.1 we study the utility of CSFA for discovering representations that are compatible with SF&GPI. In $.2 we study the utility of SFK for transferring to sparse-rewe long-horizon tasks.

---

Environment setup. We conduct our experiments in the 3D play- room environment. Observations are partial and egocentric pixel- based images. The agent gets no other information. Actions. The agent rotate its body and look up or down. To pick up object must move its point of selection on the screen. When it picks up an object, it must continuously hold it in order move it elsewhere. To accomplish this, the agent has 46 actions. Training tasks. The agent experiences n,, 32 training tasks T Train {KI,-... Kn., composed of "Find A" and "Place A near B". |A] = 8 and [ = 3. All tasks provide reward of 1 upon-completion. We provide more details in F
Figure 4: Playroom.
5.1 Evaluating the utility of discovered representations for SF&GPI
Our first experiments are sanity check for the utility of discovered representations for SF&GPI, We train and evaluate agents on the same set of tasks. However, during evaluation the agent has access to all training tasks and must select the appropriate one with SF&GPI; given task encodings Wergin {wi,..,, Wn, J. the agent acts according to policy <(, w;) = arg max, max ( (v5s,, a, wa) wi)- When wa ws, (s, wx) = arg max Q(s, a, w.). This will fail if the agent hasn't 't learned repre- sentations that support GPI. If an agent cannot perform GPI on training tasks, then will probably fail with novel transfer tasks. We expand on this in (B) Metric. We evaluate agents with average success rate. Challenges. Most prior work has leveraged SF&GPI for combining "Find tasks" where an agent simply navigates to objects 4 8 6.. We add significantly longer horizon "Place Near" task where the agent must select an object and hold it as moves it to another object. This tests the utility of discovered representations for learning SFs that enable SF&GPI over long horizons.
Research questions. Q1. How does CSFA compare against baseline methods that share their SF estimatar &p across tasks while discovering & and w/ Q2. Is each piece of CSFA necessary?
Baselines. (1) Universal Successor Feature Approximators (USFA) [8] the only method that shares an SF estimator across tasks and has shown results in 3D environment. (2) Mod- ular Successor Feature Approxima- tors (MSFA) 6 showed that lever- aging modules improved SF estima- tion and enabled cumulant discovery. However, they did not discover task encodings and only showed results in simple grid- worlds. Both basclines es- timate SFs with point-estimates. Com- paring to them tests (a) the utility of our categorical representation and (b) CSFA's ability to discover both cumu- lants and task encodings that enable GPI in large-scale setting.
Find (GPI)
Place near (GPI)
3.4
3.0
0.9
58
Success Rate
0.4
Caurs)
(Borsa et 2019)
6.2
MSFA (Canvalho et al, 2023)
58
Num learner frames 1e9
Num leamnr frames 1e9
Figure 5: CSFA discovers representations compatible with GPI across short and long task-horizons. Both USFA and MSFA degrade in performance on Find tasks when paired with GPI. Neither are able to learn our longer horizon Place task. We hypothesize that this is because they approximate SFs with a point-estimate. (4 seeds)
CSFA discovers SFs compatible with SF&GPI while baseline methods cannot. For fair com- parison, we train each baseline with the same learning algorithm (except for SF-losses), enforce w lie on a unit sphere, and stop gradients from Q-lerning. We found that large-capacity networks were needed for discovering In particular,
parameterized ée with a 8-layer residual network (ResNet) (29] for all methods. One important difference is that we leverage a set of ResNet modules for MSFA since Carvalho et al. [6] showed that modules facilitate cumulant discovery. Figure shows that USEA and MSFA can perform GPI for Find tasks; however, neither learn place tasks in computational budget. Given that controlled for how and w are learned, we hypothesize that the key limitation of USFA and MSFA is their reliance on scalar SF estimates.
A categorical representation is necessary. One difference between MSFA/USFA and CSFA is that CSFA shares an estimator parameters across individual SFs. Figure T shows that when CSFA shares estimator but produces scalar estimates (CSFA- categorical), GPI performance degrades
7

---

below MSFA/USFA. We hypothesize that a network producing point-estimates has trouble estimating returns for cumulants of varying magnitude. In C we show evidence that CSFA has more stable SF-errors compared to scalar methods.
Sharing our estimator across cumu- lants necessary. we keep our categorical representation but don't t share it across cumulants (CSFA independent), GPI performance de- grades on our long-horizon place near task. Stopping gradients from Q- learing is necessary. Interestingly, when pass gradients to the task- encoder from Q-learning (CSFA no stop grad), we get perfect train performance on place tasks but highly unstable GPI performance. We found that passing gradients leads to dimen- sional collapse (27) (see $C We hy- pothesize that this makes Wi unsta- ble. Likewise, if don't bound tasks (CSFA- - no [|w]],, we find degraded GPI performance compared train- ing but for even simpler tasks. While other methods may work, enforcing w lie on a unit-sphere is simple solution.
Find (GPI)
Place A near (GPI)
Success Rate
26
14
24
(Surs)
CSPA no [[w]] (D2)
CSFA indeperiiatt
CSFA
no stop grad (D1)
Num learner frames
Num learner frames
Figure 6: CSFA Ablations. Left. Not bounding ||]]| (D2) degrades GPI performance. Interestingly, removing our categorical representation (CSFA categorical) gets perfect training performance but terrible GPI performance. Right. Passing gradients to the task-encoder from the SF- approximator leads to unstable GPI performance (D1) de- spite good training performance. If CSFA does share its SF-estimator, it seems to learn more slowly. Thankfully, each addition is simple and together enables GPI with long- horizon place tasks. (3 seeds)
We present full results for these plots in $D
5.2 Transferring to combinations of long horizon tasks
Our second experiments test the utility of SFK for transferring combinations of long horizon, sparse-reward tasks. Transfer tasks are conjunctions of known tasks. Subtasks can be completed in any order but reward is only provided at task-completion. Find tasks contribute reward of 1 and place tasks contribute a reward of 2. We provide more details in the $F
Research questions. Q3. How do compare against baseline transfer methods? Q4. How important is to use CSFA and to sample coefficents over known task encodings?
Avcrgge
100 million
Deep OK (ours)
MTRL
Frames
10B
Tremes
165
Trames
Figure 7: SFK transfer most quickly to longer horizon tasks. Distral and MTRL learn at about the same speed, though Distral slightly faster. For put where the agent needs do 3 place tasks (A near B, C near D, and E near F) SFK learns with 100+ fewer samples. (9 seeds)
Baselines. (1) Multitask RL (MTRL). We train an Impala [30] agent on all training tasks and see how this enables faster learning for our transfer tasks. We select Impala because it is well studied
in the Playroom environment B1 32 10 (2) Distral (22 is a common transfer learning method
which leams centroid policy that is distilled training task policies. Comparing against MTRL and Distral tests the utility of combining behaviors with SP&GPI.
Q3: SFK has better jumpstart performance. Figure 7 shows that all methods get similar per- formance by 500 million frames. However, for longer task combinations (Put 2 times or 3 times), SFK gets to similar performance with far fewer frames. When remove our curriculum (Figure

---

this gap further increases. No method does well for our longest task (Put x 4) which involves objects. We conjecture that one challenge that methods face is holding an object over prolonged periods of time. If the agent selects the wrong action, it will drop the object it's holding. This may make it challenging when there's some noise from either (a) centroid task, as with Distral, (b) SF&GPI as with SFK. Despite not reaching optimal performance, SFK provides a good starting point for long-horizon tasks.
Q4: Leveraging CSFA and sampling coeffi- cents over known task encodings is critical to jumpstart performance. Figure[]showe us that if we don't it leverage CSFA to estimate SFs and instead use USFA, SFK fails to transfer. We hypothesize that this is because of the USFA uses point-estimates for SFs, which shows poor GPI on training tasks (see Figure (5) its not surprising it fails on novel transfer tasks. Di- rectly sampling from the encoding space does not transfer as quickly. Empirically, we find that learned task encodings have a high cosine simi- larity, indicating that they occupy a small portion of the encoding space (see $C. We hypothesize that this makes it challenging to directly sample in this embedding space produce meaningful behavior.
Put 3 (no curriculum)
Put 4 (no curriculum)
Arcrge Reward
200 million
Deep OK (ours)
Oistral (lch et 2017!
C5
Tremres 100
Trames
105
Figure 8: SFK transfers most quickly with no curriculum of shorter tasks. SFK and Distral reach the same performance but SFK is 200+ mil- lion frames faster. MTRL fails to transfer here. For "Put x 4", SFK is suboptimal but achieves some success. (8 seeds)
Find A & Put B near
Put A near B & near D & near
6 Discussion and conclusion
Average Rowerd
5 12
We have presented SFK, a novel method for transfer that adaptively combines known behav- iors using SF&GPI. To discover representations that are compatible with SF&GPI, SFK esti- mates SFs with novel algorithm, CSFA. CSFA constitutes both novel architecture (of the same name) which approximates SFs with a pmf over a discretized continous values and a novel learn- ing objective which estimates SI's for discovered cumulants with categorical cross-entropy loss.
Deep OK
binary
Deep OK - CSFA
Deep (Klssrs
Framos
1eB
Frames
208
Figure 9: SFK ablation. SFK fails to transfer if we remove CSFA. We find that sampling directly in task-encoding space instead of sampling binary coefficents is both slower to learn and has higher performance variance. We hypothesize that this is due to concentration in the task encoding space.
We first showed that CSFA is able to discover SF&GPI-coppatible cumulants and task encod- ings for long-horizon sparse-reward tasks in the 3D Playroom environment ($5.1). We compared CSFA to other methods which share their approximator across tasks: (1) USFA, which showed results in 3D environment but hand-designed representations, and (2) MSFA, which discovered cumulants but did so in gridworlds with hand hand-designed task encodings. We additionally compared against an ablation which removed our categorical representation and or did not share it across cumulants. Our results show that a categorical representation over discretized values can better handle estimating SF-retumns when discovering cumulants and task encodings for long-horizon sparse-reward tasks.
We built on these results for our second set of experiments and showed that SFK provides strong jumpstart performance for transfer to combinations of training tasks ($5.2). We compared SFK to (1) Mulitask RL (MTRL) pretraining and finetuning, and (2) Distral, which distills knowledge back and forth between task-specific policy and "centroid" policy. Our results showed that, for long-horizon tasks, SFK could transfer with 100 million+ fewer samples when curriculum was present, and 200 million+ fewer samples when no curriculum was present. We found that simply using Bernoulli distribution for sampling task coefficents performed well because it facilitates exploiting GP& SF&GPI enable transfer to linear combinations of task encodings. By leveraging a Bernoulli distribution, the dynamic transfer query was simply a linear combination of learned task encodings. This is precisely the type of transfer task encoding that SF&GPI has been shown to work well with.

---

Limitations. While we demonstrated discovery cumulants and task encodings that enabled transfer with dynamic SF&GPI query, we relied on generating queries as weighted sums of known task encodings. A more general method would directly sample queries from the task encoding space. We found that the task encodings discovered were fairly concentrated. Future work may mitigate this with contrastive learning [3]. This respects the dot-produs relationship of cumulants and task encodings while enforcing they be spread in latent space. Finally, while demonstrated good jumpstart performance for long-horizon tasks, we did not reach optimal performance in our sample budget. Despite building the Option Keyboard, we did not employ options sampled new query at each time-step. Future work may improve performance by sampling queries more sparsely and by implementing latent "initiation set" 34
This paper did not study the upper limit on the number main tasks our method can handle. However, it was a significant increase in complexity from prior work 4 1148660 6. For example, if one combines 2 "go to "tasks, this leads to C44, 2) = 6 combinations. We have 8 find tasks and place near tasks, where Cnn, k) denotes n choose combinations. If combine just 2 place near tasks, this leads C(24. 2) = 276 possible transfer combinations. We are optimistic that one can scale the capacity of sufficiently expressive task encoder with the number and complexity of tasks being learned.
Conclusion. These results present the first demontration of transfer with the Option Keyboard (SF&GPI with dynamic query) when all representations are discovered and when the task-encoder and SF-approximator are shared across tasks. This may enable other methods that leverage SFs in multi-task settings to leverage discovered representations (e.g. for exploration (35) or for multi-agent
RL [5 36)). More broadly, this work may also empower neuroscience theories that leverage SFs as cognitive theories to better incorporate discovered representations (e.g. for multitask transfer (37]
for learning cognitive maps (38). We are hopeful that SFK and CSFA will enable SFs to be adopted more broadly with less hand-engineering.
7
Acknowledgements
The authors would like to thank the anonymous reviews for helpful comments in improving the paper and its accessibility. We also thank members of Google DeepMind for their helpful feedback.
10

---

References
[1] André Barreto, Diana Borsa, Shaobo Hou, Gheorghe Comanici, Eser Aygün, Philippe Hamel, Daniel Toyama, Shibl Mourad, David Silver, Doina Precup, et al. The option keyboard: Combining skills in reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019.
[2] André Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. PNAS, 117748):3077--30087, 2020.
[3] André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. NIPS, 30, 2017.
[4] André Barreto, Diana Borsa, John Quan, Tom Schaul, David Silver, Matteo Hessel, Daniel Mankowitz, Augustin Zidek, and Remi Munos. Transfer in deep reinforcement learning using successor features and generalised policy improvement. In ICML, pages 501-510. PMLR. 2018.
[5] Angelos Filos, Clare Lyle, Yarin Gal, Sergey Levine, Natasha Jaques, and Gregory Farquhar. Psiphi-learning Reinforcement learning with demonstrations using successor features and inverse temporal difference learning. In International Conference on Machine Learning, pages -317 PMLR, 2021.
[6] Wilka Carvalho, Angelos Filos, Richard L Lewis, Satinder Singh, et al. Composing task knowledge with modular successor feature approximators. arXiv preprint arXiv:3301.12305. 2023.
[7] Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604 609, 2020.
[8] Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado Van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. ICLR, 2019.
[9] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Emesttu, and Noah Dormann. Stable-baselines:: Reliable reinforcement learning implementations. The Journal of Machine Learning Research, 22(1):12348 12355, 2021.
[10] Josh Abramson, Arun Ahuja, lain Barr, Arthur Brussee, Federico Carnevale, Mary Cassin, Rachita Chhaparia, Stephen Clark, Bogdan Damoc, Andrew Dudzik, et al. Imitating interactive intelligence. arXiv preprint arXiv:2012.0555 (tt::20/2.0000 2020.
[11] Majid Abdolshah, Hung Le, Thommen Karimpanal George, Sunil Gupta, Santu Rana, and Svetha Venkatesh. new representation of successor features for transfer across dissimilar environments. In International Conference on Machine Learning, pages 1-9. PMLR, 2021.
[12] Matthew E Taylor, Peter Stone, and Yaxin Liu. Transfer learning via inter-task mappings for temporal difference learning. Journal of Machine Learning Research, 8(9), 2007.
[13] Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, pages 1480-1490, PMLR, 2017.
[14] Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerch- ner. Cobra: Data-efficient model-based through unsupervised object discovery and curiosity- driven exploration. arXiv preprint arXiv: (0t.:/99000000.. 2019.
[15] Andrei A RusU, Neil C Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell. Progressive neural networks. arXiv preprint arXiv:1606.0460 2016.

---

[16] Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha, Andrei Rusu, Alexander Pritzel, and Daan Wierstra. Pathnet: Evolution channels gradient descent in super neural networks. arXiv preprint 17X1v:1701.08734, 2017.
[17] Andrei A Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. arXiv preprint arXiv: 1 1.060299 2015.
[18] Richard E Bellman. Dynamic programming. Princeton university press, 2010.
[19] Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2): 181-211, 1999.
[20] Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function ap- proximators. In International conference on machine learning, pages 1312-1320. PMLR, 2015.
[21] Junhyuk Oh, Satinder Singh, Honglak Lee, and Pushmeet Kohli. Zero-shot task generalization with multi-task deep reinforcement learning. In International Conference on Machine Learning. pages 2661-2670. PMLR, 2017.
[22] Yee Teh, Victor Bapst, Wojciech Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning, Advances in neural information processing systems, 30, 2017.
[23] Zhuangdi Zhu, Kaixiang Lin, Anil Jain, and Jiayu Zhou. Transfer learning in deep reinforce- ment learning: A survey. arXiv preprint arXiv:2009.0788 2020.
[24] Yuke Zhu, Daniel Gordon, Eric Kolve, Dieter Fox, Li Fei-Fei, Abhinav Gupta, Roozbeh Mottaghi, and Ali Farhadi. Visual semantic planning using deep successor representations, In Proceedings of the IEEE international conference on computer vision, pages 483 492,449221 2017.
[25] Leslie Pack Kaelbling, Michael I. Littman, and Anthony Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134, 1998.
[26] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51877540):599-533,, 2015.
[27] Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.093488 2021.
[28] Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
[29] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[30] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning, pages 1407-1416. PMLR, 2018
[31] Stephanie CY Chan, Andrew Kyle Lampinen, Pierre Harvey Richemond, and Felix Hill. Zipfian environments for reinforcement learning. In Conference on Lifelong Learning Agents, pages 406-429. PMLR, 2022.
12

---

[32] Felix Hill, Andrew Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James McClelland, and Adam Santoro. Environmental drivers of systematicity and generalization in situated agent. arXiv preprint arXiv: 2019.
[33] Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages $15-823, 2015.
[34] Vivek Veeriah, Zeyu Zheng, Richard Lewis, and Satinder Singh. Grasp: affordance selection for planning. arXiv preprint arXiv:2202.04772 2022.
maadient-based
[35] David Janz, Jiri Hron, Przemyslaw Mazur, Katja Hofmann, José Miguel Herándezz-oobtoo. and Sebastian Tschiatschcck. Successor uncertainties: exploration and uncertainty in temporal difference learning. Advances Neural Information Processing Systems, 32, 2019.
[36] Tarun Gupta, Anuj Mahajan, Bei Peng, Wendelin Böhmer, and Shimon Whiteson. Uneven: Universal value exploration for multi-agent reinforcement learning. In International Conference Machine Learning, pages 3930-3940 1941. PMLR, 2021.
[37] Momchil S Tomov, Eric Schulz, and Samuel J Gershman. Multi-task reinforcement learning in humans. Nature Human Behaviour, 5(6):764-773, 2021.
[38] William De Cothi, Nils Nyberg, Eva-Maria Griesbauer, Carole Ghanamé, Fiona Zisch, Julie M Lefort, Lydia Fletcher, Coco Newton, Sophie Renaudineau, Daniel Bendor, et al. Predictive maps in rats and humans for spatial navigation. Current Biology, 32(17)336763368, 2022.
[39] Ehsan Imani and Martha White. Improving regression performance with distributional losses. In International conference on machine learning, pages 2157-2166. PMLR, 2018.
[40] Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dedieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, lurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiee, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. http:// com/deepminn
[41] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.604 2014.
[42] Felix Hill, Olivier Tieleman, Tamara Von Glehn, Nathaniel Wong, Hamza Merzic, and Stephen Clark. Grounded language learning fast and slow. arXiv preprint urXiv:2009.07719, 2020.
[43] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforce- ment learning. International conference on machine learning, pages 1928-1937. PMLR. 2016.
[44] Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefraneq, Simon Green, Victor Valdés, Amir Sadik, et al. Deepmind lab. arXiv preprint 1rXiv:1612.0380), 2016.
13

---

A Learning to estimate scalar values with distributional losses
We consider the setting where we are learning function that maps input & E X to scalar y E RR
Representing scalars with probability mass functions (pmfs). One can represent any scalar y as weighted average of some of fixed bins B = {bi...., bar) and learnable weights p {pI..... PM [
as
M
~
Pm bm
(10)
where bm R and Pm E Rt. (B. jointly define pmf. We consider the setting where IB is defined apriori.
Learning pmf representations of scalars. We can learn the weights with neural network fo that computes p softmax (x)) E RM by using a negative log- likelihoodd objective 39. MuZero showed that one can construct an effective target from a scalar y with the two-hot representation twohot(y) E TRM, The resultant objective then LELL = twohot(y) log(p). In summary,
p softmax (fotmax((f(())) E R
(11)
CULL twohot(y) log(p)
(12)
Prior literature has shown that leveraging this strategy for representing value-functions can improve learning (e.g. by stabilizing gradients) [39 77.
The two-hot encoding twohot(y) is a generalization of the one-hot encoding where all elements are O except for the two entries closest y at positions and + 1. These two entries sum up with more weight given to the entry that is closer to y:
1bm+1 y// 16m+1 6ml
if
twohot(y): = I0m /1 1bm+1 6ml
if +1
(13)
else
Representing a successor feature with a pmf. In our setting, we approximate an n-dimensional SF- vector dh (s, a) E R" with (p(s, a, w). We learn different set of weights for each dimension and assume all SF values lie between by and bas of apriori defined set of bin values TB (bi...., by }- We'll reference the weights for the k-th SF dimension as th
{pr"' ,...,,,. "p..,...... PM }. For the k-th SF dimension, we compute these weights as p = softmax (aammax(lof (lo (s, w. ex)), where & is a state representation, is task representation, and ex is embedding of that SF dimension. We can then compute the k-th SF as visi a, w) MM pil bm. We present a schematic of this in Figure
B Challenges of SF&GPI with learned representations
Our goal is to learn successor feature (SF) estimator (ha a, w) that (a) uses the output of cumulant function 6 = ( [s. a) as its prediction target and (b) is parameterized by the output of task encoder
we (x). In this section, we will first describe how these terms show up in the performance bounds for GPT (equation 15). Afterwards, we will describe how these terms show up in empirical challenges faced when trying to learn these values for sparse-reward long-horizon tasks (...T). We will conclude with how these challenges manifest in the simplified setting of $5.1 where we "transfer" known training tasks (BB.)..
GPI performance bound. Given learned va, da, and Wa. we want to leverage GPI to transfer to task encoding w' (K()). In this setting we have learned n policies { in for tasks W = (w and can leverage GPI to transfer to w' as follows
n(,, w') arg max max (b(s,, a, w) w')
(14)
WEW
Define x² as the optimal policy for task w'. Filos et al. 5 showed that the gap between a tasks's optimal Q-value (* w') and the Q-value for the GPI-policy Q a, (")) can be bound by (1) the SF approximation emoor dy -| " loo.. (2) the reward approximation error ó,
14

---

Legend
Unbounded Wg
Design choice
Empiriaal Phenomena
Large
[
(4)
GPI bound term
Long-horizon
(Borsa et al, 2019)
Environment factors
Observations we make
Weaker GPI Bound
Error in ?/
(Borsa et al, 2019)
(Filos et al, 2021)
Small 0 (wi W))
Unstable 4/8 (s, a, w)
Error In Φ
/(1)
w
Passing Q-learning gradients Wgg
(Barreto al, 2018) -
Sparse reward
Learning Øg
Scalar we
Figure 10: Summary of GPI challenges when jointly learning SFs ths cumulants de, and task encodings wg. (1) Estimating SFs over non-stationary cumulants with point-estimate can make the unstable (see Figure 13 (2) Large task encoding magnitude ||w||| amplifies the effect of the SF-error the GPI bound (see equation (5] (3) Task encodings exhibit dimensional collapse when they get Q-learning gradients (see Figure 15b) (4) Sparse reward, long-horizon tasks exacerbate errors in
approximating rewards and SFs. See [B. for more.
||"" - T oooo (3) the magnitude of the transfer encoding [|w/)), (4) distance of w the closest training task, Aw = minje (l.......) ||w2 - wyll. That is, Vs,a
2 Q"(s,a, - (s, S 1- i I + [|w0'l by + (2 7),
(15)
(1 7)
Intuitively, this says that the error on the esimated Q-values for task encoding w' will grow in proportion (1) the distance of the task encoding w' to known tasks A (2) the product of the SF-approximation error and task encoding magnitude || and (3) the reward-approximation error 8,
B.1 Empirical observations
We found some design choices can weaken GPI performance and connect then to equation (c.g. see Figure 6 for a summary).
Challenge 1. The first challenge, and the main focus of this paper, comes from estimating SFs with a scalar point-estimate while jointly discovering cumulants. Empirically, we found that cumulants changing in magnitude
the course of learning. As cumulants change in magnitude, their corresponding return can quickly change (see Figure 12p. Empirically, we found that estimators that scalar point-estimates have oscillating SF-errors du (see Figure While they learn training tasks, they failed to produce SF-estimates that work with GPI (see Figure 14).
Challenge 2. GPI performance is bounded by product of the SF-approximation error and task encoding magnitude ||w'||6p Empirically, we found that if we did not bound [|w]], I. can have relatively large magnitude (see Figure 15ap. In figure II provide intuition for why this can inhibit GPI performance. We confirm this empiniaaly in Figure
Challenge 3. In this work, we focus on learning an SF-approximator parameterized by a task-encoder (0(8, a, w). When passing gradients to the task-encoder from Q-learning with SF-approximator, we found that this lead to dimensional collapse 27 (see Figure 15b). We conjecture that dimensional collapse can make the SF-estimator unreliable differentiating between training tasks since an input. This may manifest with poor GPI performance shown in Figure 21
15

---

B.2 Why might SF&GPI fail for transferring to training tasks?
Our first experiments showed that jointly learned SFs the cumulants Φg. and task encodings we can exhibit poor GPI performance for transferring to training tasks ($5.1) Here, we provide intuition for why this can happen.
Legend
Successorr Feature
Successor Feature approximation
Task Encociing
10
4/2
Figure 11: Illustration of why SF&GPI can fail for transferring to training tasks. Define 4/s = (ss a, w:). vi = ti e. When transferring to training tasks, the correct task encoding for yh is wi. Assume we are transferring to w2- GPI will select the wrong task if th we happens due (a) error in the SF-approximatin or (b) the magnitude of w2-
> 1h w2. This
When we transfer to a train task, the agent simply needs to act with the SF corresponding to the task policy, i.c. we want task selected the inner max Wsal max, (i(s, a, w) ww w;. In this setting, the distance to the closest training task is 0, i.e. A. minnje(.......) 1| wyll = 0. This leads to
the following change to equation 15
2 1- I 8v + (2- (1 7)
(16)
This tells that, even when we are transferring to a known task, performance is bound by (I) the product of the SF-approximation error and task encoding magnitude / and (2) the reward- approximation error. The second term is intuitive-if we have not learned cumulants that predict reward, then cannot learn appopooosssssors that predict action-values for task reward. Below, unpack the first term.
Define ub (b(s, w.). Now assume that approximations for SFs are bounded by an error e, i.c. J. db + Now consider simplified setting with only two tasks wi and w2 and their SFs th and 222 shown Figure II] Consider computing GPI for w2. First, simply due to the error e, we can have that TpT u22
2b W2. Second, if [w|] is large, can dominate the dot poo-product and lead to
TT w/ > J Wg-
C Additional analysis
Here we present some additional analysis. In SC.1 we show additional analysis simplified curriculum of only "Find" tasks. In C.2 show more ablations for our method.
16

---

C.1 Analysis demonstrating observed learning challenges
We conduct additional analysis on n4 = 8 training tasks Ttrain = (K1,.... (K1,.... Knn, composed of "Find A" where |A| 8. All tasks provide reward of upon-completion. This a simplified version of the experiments in $5..1 where we study some of the challenges observed during training. Specifically, we present evidence for the following:
1. The magnitude of cumulants changes during learning (see Figure 12
2. The SF TD-error can be more unstable for methods that estimate SFs with point-estimate (see Figure l3p.
3. ||/||| can be relatively large if don't bound (see Figure 15ap.
4. Passing gradients the task encoder leads to dimensional collapse (see Figure LSb
Cumulant Average
Cumulant
Return for different Cumulant Magnitudes
0.30
CSFA (ours)
CSFA (ours)
0.25
0.0
0.20
0.02
0.01
0.05
0.00
0.4
0.8
0.0
0.4
0.5
Frames
1e9
Frames
1e9
Average Cumuliant time-step
Figure 12: Evidence that the magnitude of cumulants and cumulant-returns changes over the course of learning. The left and middle plots show that the cumulant average and the cumulant L1 (ō and ||]| respectively) both increase during training. The right-most most plot shows the discounted return that would be calculated over 100 sime-steps for different values of Φ E (10-1, 10-3 10-2, 10-1).
We calculate return as R Σ - y where y .99. We emphasize that this is not reflective of the true returns experienced by the agent but instead demonstrates how the return magnitude can change dramatically as the cumulant average changes. In Figure [3] we show that methods which estimate this return with scalar point-estimate exhibit unstable learning dynamics.
SF TD error [6wl
0.10
0.06
0.04
CSFA (ours)
0.02
CSFA no categorical
USFA (Borsa al, 2019)
0.00
0.0
0.2
0.6
0.6
Frames
1e9
Figure 13: Evidence that the SF-TD error is less stable for methods that use point-estimate. While both scalar and CSFA methods show an increase in error during training, only point-estimate methods (CSFA- no categorical, and USFA) exhibit oscillatory behavior. In contrast, CSFA, which utilizes probability mass function for return estimation, demonstrates a consistently more stable SF-TD error. We note that all methods increase TD-error as the agent learns because the error is over larger values. In Figure [4
we show that oscillatory behavior correlates with poor GPI performance.
17

---

(GPI)
1.0
CSFA
CSFA no catesgoriall
0.4
al, 2019)
0.2
07
Figure 14: While all methods can learn Find Tasks, only CSFA can do GPI with find tasks.
|
Task Similarity
1.00
10
0.98
0.96
CSFA (ours)
CSFA (ours)
6
CSFA
||w|| (D2)
CSFA no stop grad (D1)
0.88
0.25
0.50
0.75
1.25
1.50
0.0
0.2
0.4
0.8
Frames
1e9
Frames
1e9
(a) When we do not bound w, it can have relatively large magnitude.
(b) When we don't stop gradients to we through Q-learning, task encodings can collapse single angle. We measure the average cosine similarity between pairs tasks. If we pass gradients from Q-learning (CSFA no stop grad), tasks collapse onto single line.
Figure 15: Analysis of task encoding for two different conditions.
18

---

C.2 More ablations
Here we present additional ablation results using the full curriculum in $5.1 We ablate the number of bins used represent our probability mass function (Figure T6p and we ablate the depth of the residual network used to parameterize cumulants (Figure [7p. For both, networks with more capacity (i.e. more bins or more residual blocks) are better able to perform GPI. To be clear, this is not saying that higher capacity networks get better train performance, instead this says that higher-capacity networks enable representations that better support SF&GPI.
Find
Find
Place
Place
(ours) CST 101 Bins
Success Pate
CSTA (ouns) CSRA: 101 birs
leg
Num learnrr frames
learner frames 1c9
Num frames 1c9
Figure 16: Ablation on number of bins. We see that as we decrease the number bins (i.e. the capacity), the GPPperformance degrades. We hypothesize that more bins provide more capacity for modelling a wider range of cumulant-returns.
Find
Find (GPI))
Place near
Place near (GPI)
1.0
CSRA (ours)
Success Rabe
0.0
CSTA- - blccts
CSRA 1 Block
0.9
0.5
0.6
05
02
02
Num learner trames
Num leamer frames
2
Num learner frames
7
Num leamer frames 1e9
Ablation on number of ResBlocks. We see that as we decrease the number of residual blocks, the GPI performance degrades for hong-horizon tasks, With 4-blocks, performance degrades from 100% to 90%; with block, performance degrades from 95% to 85%. This suggests that use more residuals blocks better supports cumulant discovery for long-horizon sparse-reward tasks.
Figure 17:
19

---

Find A & Put B near C
Put A near B & near D & E near F
2.5
4
Average Reward
1.5
2
1.0
Deep OK (ours)
1
0.5
Deep OK - reuse We
Deep OK reuse Sg
0.0
Frames
1e8
Frames
1e8
Figure 18: Ablation on learning a new state function and task encoder. We see that re-using our state function still enables learning, albeit more slowly than learning new state encoder for the transfer policy. We see that re-using the task encoder prohibits transfer. We conjecture that this is because this leads new representations for the train tasks. Since train task encodings parameterize the SF-estimates, this may induce error in the SF-estimates over the course of learning.
20

---

D Full results
We present both training and GPI performance for the ablation results in $5.1 The original plots only showed GPI performance.
Find
Find (GPI)
Place
Place (GPI)
Success Rate
CSFA [ours)
USPA (Dona et el, 2018)
MSFA (Canvatho et 2023)
Num learner frames leg
Num learner fromes 1e9
Num learnar frames
1e9
Num learner fromes 1e9
Figure 19: CSFA discovers representations compatible with GPI across short and long task- horizons. We that while most USFA seeds learn Find tasks, smaller subset perform GPI. MSFA can do so consistently. Neither are able to learn our longer horizon task. We hypothesize that this is because they approximate SFs with point-estimate. (4 seeds)
Find A
Find A (GPI)
0.8
Success Rate
0.4
0.4
0.2
0.2
no || (D2)
CSFA no categorical
0.0
0.0
Num learner frames
1e9
Num learner frames 1e9
Figure 20: Full results for short-horizon task ablation. This plot makes clear that CSFA with scalar estimator (CSFA no categorical) gets perfect training performance but terrible GPI perfor- mance. also show that when we don't bound ||ы/]] may sometimes get good train performance but there still a drop in GPI performance.
21

---

Find
Find (GPI)
Place
Place near (GPI)
independent
Success Rate
CSFA nootop grad (D1)
Num learner frames 1c9
Num leamer frames
1c9
Num leamerr frames
1c9
Num leamer frames
'ce
Figure 21: Full results for long-horizon task ablation. This plot makes clear that using separate categorical SF-estimator for each cumulant (CSFA-independent) leads to a drop in GPI performance for long-horizon tasks when successfull (90% to 80%) but often unsuccessful in learning long- horizon tasks. This is even more striking when enable gradients to pass from Q-learning into our task encoder (CSA no stop grad), where multiple seeds will get perfect train performance for long-horizon tasks but exhibit very unstable GPI performance.
22

---

E
Implementation Details
The code is proprietary so we cannot release it but we included implementation details below to
enable reproducibility.
Architectures were implemented with the haiku library. Optimizers were implemented with the optax library. RL algorithms were implemented with the rlax library. All ResNets Blocks correspond use the "BlockV1" definition found in the haiku library. All activation functions used were the relu activation unless otherwise stated.
Everything was implemented with the jax ecosystem 40
The rest of this section is structed as follows. In $E.1 we discuss the compute we used in our experiments. In $E.2 we detail SF-based architectures: MSFA, USFA, CSFA. In SE.3 we detail transfer architectures: SFK, Impala (used for MTRL), and Distral. In $E.4 we describe learning for we describe hyperparameter search.
each method. Finally, in $E.5
E.1 Compute resources
All experiments were run using Google Dragonfish TPU with x 2 topology devices. Experiments for training tasks ($5.1) lasted about 1.5 days. Our simplified experiments in | 5-6 hours. Transfer experiments ($5.2) lasted 2-3 hours. We ran a learner, actors, and evaluators on TPUs. All experiments carried out using a distributed A3C setup [30] with discrete actions. Each experiment had learner, 512 actors and evaluators per task, each running in parallel.
E.2 SF-based architectures: USFA, MSFA, & CSFA
All SF-based architectures had more-or-less the same architecture with the exception of the SF- Approximator. We present diagram of this general architecture in Figure 22 Below we detail this architecture along with differences across methods.
Successor Feature Approximators
approsimator (be(s, w)
Linear (16) + L.2 norm
ResNet blocks,
MLP
LSTM
LSTM
encodor
Lineer
Function
de (*,a)
my
Word Embed (b4))
Concat
ResNet
Task
Previous action
Ct
Figure 22: Architecture schematic used by SF-based methods.
Observation function. All architectures used a ResNet as their observation encoder with channel dimensions (16, 32, 32) for each ResNet block. Each ResNet block had two Resdiual layers and used a stride of 2 throughout. The output of this was then flattened, concatenated with the previous action, and passed through projection of dimension 256.
Task encoder. All architectures embedded words in task descriptions using standard word embed- dings of dimension 64. The sequence then fed through LSTM of dimension 64. The outputs of each LSTM time-step were then summed and passed through projection to the cumulant-diment (16 across all experiments). THe output was then divided by its L.2 norm.
State function. All architectures used an L.STM dimension 512. MSFA differed in that used 4 LSTMs. We tried having their individual size be 512/4=128 and 1024/4=256 and found much
23

---

better performance with a smaller LSTM. They shared information with a transformer attention block. Please see Carvalho et al. [6] for details.
Cumulant function. All architectures used a MLP-based ResNet with blocks of dimension 256 each. To create these ResNet blocks, we use the "BlockV1" definition in haiku but replace all convolution layers with linear layers. We then pass the output of the ResNet to a 4-layer MLP where each layer has dimension 256. The MLP outputs scalar cumulants. MSFA differed in that it used a separate ResNet + MLP netwrok for each state module output. The outputs were then concatenated. This again follows Carvalho et al. [6]
SF Approximators
CSFA used an MLP with dimensions [512, 512] and output size N We converted back and forth from the two-hot representation using fraxx library functions) In order re-use this MLP across cumulants, we embed the cumulant dimension with a word embedding of dimension 256. The input the SF-MLP is then the concatenation of (a) cumulant embedding, (b) the task encoding, and (c) the current state.
USFA used an MLP with dimensions [512, 512] and output size A*C. The input to the SF-MLP is the concatenation of (a) the task encoding, and (b) the current state.
MSFA used 4 MLPs with dimensions [512. 512] and output size **//4. The outputs were then concatenated to the form the full set of SFs. The input to each SF-MLP is the concatenation of (a) the task encoding, and (b) the current module-state. We refer to 6 for a detailed implementation notes.
A=number of actions, C=number of cumulants, N=number of bins. CSFA used N 301 bins evenly spaced between [-5,5]. All methods used C = 16 cumulants and Playroom has A = 46 actions.
E.3 Transfer architectures
All transfer methods were based off of an Impala architecture because it has been used in previous Impala and Distral have more-or-less the same architecture, except that Distral has an additional centroid policy head. Below we share functions shared across meth-
Playroom resuts [1
32 10
ods. In IE..I we detail implementation specific Impala and Distral and $E..2 we detail implementation specific to SFK. We present diagrams of these architectures in Figure 23
Task encoder. We embed words in task descriptions using word embeddings of dimension 64. The sequence was then fed through an LSTM of dimension 64. The outputs of each LSTM time-step were then summed. This sum is the output.
Value function. All methods use an MLP with 512 hidden dimension and scalar output.
E.3.1 Impala & Distral
Observation function. All architectures used a ResNer as their observation encoder with channel dimensions (16, 32,322 32) for each ResNet block. Each ResNet block had two Resdiual layers and used stride of 2 throughout. The output of this was then flattened, concatenated with (1) the previous action and the task. This then passed through projection of dimension 256.
State function. All architectures used an LST of dimension 512.
Policy
Impala uses an MLP with dimensions [1024] and output size A. The policy is parameterized with softmax.
Distral learns two functions: one centroid function produces logits h((as)) and task- specific function produces logits fi (a(s). The task-policy is then parameterized as (alss) x (a Iı(a)s) + (a(()))) Following Teh et [22, we set a .5 and / for acting. During learning we use different 3 to encourage distillation. We detail this in S..4
24

---

Impala/ Distral
Deep Option Keyboard
Coeffiiiett-poliy Head [ai.... Ci.y)
(Only Distral) Centrold Policy head
8,
Value head
Task Policy head MILP (1024)
head
MLP (512)
MLP (1024)
MLP (512)
(1024)
0
approximator (8, a, w)
Linear (16)
LSTM (51.2)
LSTM (512)
LSTM (512)
LSTM
Linear (256)
(64)
Linear (258)
Previous coefficients
LSTM (64)
LSTM
(64)
Concat
Word Embed (64)
Concat
Previous action R(-1
ReslNet (18, 32, 32)
ResNet 32,
Word Embed (64)
Word Embed (84)
Known Tasks
Previous
Legend
a1-1
Frozen parameters
Zt
Zt
New
Task
Figure 23: Architecture schematic used by transfer methods.
E.3.2 SFK
SFK uses the CSFA architecture but freezes all parameters. It then learns state function, task encoder, policy head, and value function. In principle, more could be re-used and potentially enable even faster learning but we opted for simplicity.
New state function. We use LSTM of dimension 512 that takes in the state from CSFA's LSTM, the output of the observation encoder, and previous coefficients produced by the coefficient-policy..
Policy, SFK uses an MLP with dimensions [1024] and output size W, where W="number training tasks" to use for GPI. In our experiments, W was 32.
E,4 Learning
All trainable parameters were updated with the adam optimizer 41 For each method, we set the following adam hyperparameters: B1 0, 32 .95, -e-6.
E.4.1 CSFA, MSFA, & USFA
Following Mnih et al. (26), all architectures had replay buffer that held up to 100,000 trajectories of length 30. This is off-policy setting. Following Mnih al. (26), all methods used target network parameters 0 that were updated smoothly using the online network parameters 0 with a coefficents of .9. That at iteration 0f 0% + O,, All methods used learning rate of 3 and set the max gradient norm to 55. We used the following coefficients to balance loss terms: 3Q = le3, 3p 1e4. Since CSFA uses a categorical loss Lest to learn SFs, which has a different magnitude from the scalar loss used learn Q-values and, found that needed to make this smaller with Bg 3 whereas we used Bg = 5e2 for scalar-based methods.
E.4.2 SFK, Impala, & Impala
We collected training data using FIFO replay queue with 36 trajectories of length 128. Batches of 32 trajectories were sampled uniformly. This was on-pol setting. SFK used learning rate of Se-5. SFK learns to select actions with the Impala learning algorithm, but with an augmented actions-space which produced task-coefficents instead of environment actions. Impala used learning rate 4 for training and 8 for transfer. Distral used a learning rate Se-6 for training and 3e-5 for transfer. All methods set the max gradient norm to 40.
25

---

For Distral, we set 1/3 = 3e-3. We set a = .5. Since Distral is a policy-gradient method, we added its policy distillation loss Laist to an Impala loss Lipplla- The full loss was Ldistral Limpata + Cais We used Oais = 5.0 for training tasks and Oaii = .1 for transfer tasks.
E,5 Hyperparameters
All methods and their corresponding hyper-parameters are based on prior work that has learned in Playroom environment [32 42 31. Specifically, we mainly used hyperparameters from [1
since they also study how Impala-based MTRL. can transfer in Playroom (though they do so in continual learning setting with skewed distributions). Starting from their hyperparameters, we swept the following.
Impala (MTRL). Prior hyperparameters worked well here since we used them for training tasks. For transfer, we swept learning rates [3e- 4. 8e-5, 3e-5, 6,-6, 3e- 3e-6], value hidden sizes [200, 512, 1024], and policy hidden sizes [200, 512, 1024].
Distral was originally built on top of the asynchronous advantage actor critic (A3C) algorithm 43 We built top of Impala since it is ied in playroom. As such, we kept hyper-parameters and tuned the following: (1) learning rates 5, Se 3 6]. (2) Odia. [10, 1, I,, 01, 0001), (3) Impala entropy coefficent [9.4e-3, 9.4e,4.44444 9.4e 5, 9.4e- 9.4e-9.7 7]. (4) 1/3 [3e-3, le 1e-3,3e-44
advanced actor-critic algorithm that has been stud-
Value-based methods (as SF-based methods) are more challenging to learn for Playroom. To facilitate search, we concentrated our search on simpler "Find" tasks. First, we searched for MSFA which has shown good performance with discovered cumulants. Reusing optimizer settings from Impala, Bg = 3 = 0 and searched over (1) SQ: [1. lel, 1e2, 1e3, le4] (2) max grad norm: [40, 5, .5.. Once agents could learn Find tasks, searched over (3) resblocks: [1, 4, 8] (4) Bg: [1, lel, 1e2, 1e3, 1e4] and (5) 3 [1, le1, 5e1 le2, 5e2, Ie3). We found the same hyper-paramenn parametrrs worked well for USFA and CSFA. For CSFA, searched Sy: [.8, .,, .08, .03, .088, .003, .0008] and bins [101, 301]. We began with 3. values which led the SF-loss to be approximately equal to the Q-value loss and then went downward from there since the SF-loss can be thought regularizer on the Q-value loss.
set
For SFK, we were able to re-use Impala HPs but searched over max grad norm [40, 5, 55., value hidden sizes [200, 512, 1024]. policy hidden sizes [200, 512, 1024], and Impala entropy coefficent [9.4e-3, 9.4e 4, 9.4c 5, 9.4e-9.666.. 9.4
26

---

F
Environment Details
This work is significant increase in complexity from prior successor feature literature 3 8 The most complex task studied by prior work were "go to" tasks the DMLab environment by 8 Here,
6
is required. (1) Playroom "place near" tasks add object-interaction. while increasing the task horizon and reward sparsity (2) prior work only combined 4 "go to" tasks while
combine 8 Find and 24 place nears tasks (3) prior work used DMLab with & actions 44 while Playroom has 46 actions. We have added this discussion to the main text.
F.1 Observation space
The agent experiences (1) observations pixel-based RGB images that have dimensions 72 x 96 x 3 and (2) task descriptions as string descriptions in a synthetic language.
F.2 Action Space
The agent can rotate its body and look up or down. To pick up object it must move its point of selection on the screen. When it picks up an object, it must continuously hold it in order to move it elsewhere. To accomplish this, the agent has 46 actions. Objects are helf as long as the agent is emitting a "GRAB" action, and dropped the first instance that a GRAB action is not emitted. We describe these actions in Table An action-repeat of 4 was applied (i.e. each action was repeated times).
Table 46 actions available to the agent. The number next to each action name indicates how many meters the agent will move with that action. For example, move forward() 1) indicates that agent will move forward 1 meter.
Type of action
actions
movement without grip
noop right(1)
move forward(1)
move forward(-1)
right(-1)
look right(1)
look right(-1)
look down(1)
look down(-
fine-grained movements without grip
move right(0.05)
move right(-0.00 -0.05)
look down(0.03)
look down(-O
look right(0.2)
look right( -0.))
look right(0.05)
look right(-0.05)
movement with grip
grab
grab move (orward(1)
grab move forward(-]]
grab + move right(1)
grab move right(-1)
grab look right(1)
grab + look righ((-1)
grab look down(1)
grab look down(-1)
fine-grained movements with grip
grab
r0ght(0.05)
grab
right(-0.05)
grab look d0wn(0.03)
grab + look down(-0.03)
grab look right(0.2)
grab look right
grab + look right(0.05)
grab look right( -0.05)
object manipulation
grab + spin right(1)
grab spin right(-1)
grab + spin up(T) grab + spin forward(-1)
grab spin up(-1)
grab + spin forward(1)
grab + pull(1)
grab pull(-1)
fine-grained object
grab pull(0.5)
grab pull(-00.
(0l)(
manipulation
pul)(-0.5)
F33
Training Tasks
The agent experiences n. 72 training tasks Tersin {KI,..., (Ki....,Kn.. Kn.. composed of "Find A" and "Place near B". |A| = 18 and [ 3.
The set A comprises 18 items that can be picked up: boat, bus, car, helicopter, keyboard, plane, robot, rocket, train, racket, candle, mug, hairdryer, picture frame, plate, potted plant, roof block, and rubber duck.
The set B comprises items that can be stacked: book, cube block and sponge.
"Find" tasks follow the form "Find a A", where the placeholder "A" represents the object to found. "Put near" tasks are of the form "Put a A near B", where "A" represent the object to be picked
27

---

up and "B" represents the object that A should be placed near. All methods experienced the same training curriculum. For all tasks, the agent only get reward upon completing all subtasks. All tasks provide a reward of upon-completion.
F.4 Transfer Tasks
Transfer tasks were made using a subset of 8 pickupable objects A': boat, bus, car, helicopter, keyboard, plane, robot, rocket.
Transfer tasks where then combinations of known train tasks that used A'. An "and" strong was placed into the task description to indicate combinations. For example "Place A near and C near D and E near F" would be "place boat near a book and place bus near cube block and place keyboard near sponge". Objects were sampled with replacement, indicating that object may appear multiple times but would have a different color. For all tasks, the agent only get reward upon completing all subtasks. "Find" subtasks provide reward of and "Place" subtasks provide reward of 2. Returning to our example, since it has 3 subtasks, it would provide reward of 6 but only upon completion of all subtasks.
28


Reviews:
[{"review_id": "O2FEfVBu10", "reviewer": "Reviewer_z7ra", "summary": "First, I must emphasize that I was new to the RL community before this NeurIPS. This is the first RL paper I have reviewed. I have been struggling to understand this paper. I tried my best. I hope the AC, reviewers, and authors may think me helpful.\n\nThis paper studies how to transfer a pre-learned RL system to new tasks using successor features and the optional keyboard. Compared with previous methods, new things are as follows: 1) the Q-value is no longer linear about the task encoding, here, the authors use a network $g_{\\theta}$ to predict the weight of SF features based of task encoding and state representation. 2) The SR approximation is modeled as the expectation of a discrete distribution, the authors claim this can help stabilize the training. 3) The learning objective is modified. The new objective consists of three parts a) a term to decrease the prediction error of reward using $\\phi$, b) a term to decrease the prediction error of reward using the difference of $\\psi$, and c) a term to maximize the possibility of current $\\psi(a_t,s_t,w)$ to occur using the twohot operation. 4) A new transfer scheme that learns a new task encoder and state function at transfer time, where the task encoding is the weighted sum of a Bernoulli random variable and is aware of the time. The proposed method achieves better performance than USFA and MSFA in 3D playroom dataset, showing a better success rate. \n\n", "questions": "I have some puzzles about the paper, plz help me clarify them.\n\n1. Eq. (5), what is the SF dimension? Is $B$ a random variable and what does it stand for? How $b_k$ take values?\n\n2. I don't understand why you can write SF $\\psi_{\\theta}^k$ as the Expectation of some distribution. Do you mean that $b_i$ should be $\\gamma^i$ and $\\mathrm{Exp}(l_{\\theta}(s_t,w,e_k))$ should be $\\phi_{t+1+I}$?\n\n3. Is there any difference between saying $\\psi_{\\theta}^k=B^T l_{\\theta}^{\\psi}$ is an Linear combination  of all $\\mathrm{Exp}(l_{\\theta}(s_t,w,e_k))$ and saying it is an expectation of $p_{\\psi^k}? Why do you prefer the latter?\n\n4. What does Line 144 mean? What is the difference between $\\tilde{\\psi}^{\\pi_k},\\psi^{\\pi_k},\\tilde{\\Psi}^{w}$? I look into the \"Universal Successor Feature Approximator\" paper. It seems the tilde stands for approximation. What confuses me is why it only approximates the neural network $\\psi_{\\theta}$? Shouldn't they be the same thing?\n\n5. Do you lose a $t+1$ in $\\phi_{\\theta}(s_t,a_t)$ in Eq. (4)?\n\n", "limitations": "Limited scenarios. The computation burden seems to be increased compared with previous methods. The networks can be harder to train due to complicated designs. Computing approximations from sampling of distributions can also heavy the burden of computation both in training and inference.", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "# Novelty\nAbove all, I may not be able to judge the novelty of an RL paper. So a comment based on my experience in my familiar domain.\n\n1. The new assumption of SF and Q-Value: It is good to introduce $g_{\\theta}$, but the Q-value is still a linear combination of the SF. So I think this is still somewhat narrow. \n2. New SF approximation: In my opinion, the key advantage is that the distribution distributes the difference between $\\psi(a_t,s_t,w)$ and $\\psi(a_{t+1},s_{t+1},w)$ to $m$ different points. So each point only needs to change modestly when the time progress, and the similar features among different times can be better stored by those points of high density among times. Correct me if I am wrong here. This seems a good point, but the authors lack enough explanation here.\n3. New Learning Objective: I don't understand why the authors omit $\\mathcal{L}_{\\psi}$ in training, the new objective seems to have no term to align the difference of $\\psi$ and the cumulant $\\phi$. The authors also do not provide enough explanation to why this objective is designed, nor any ablations to explain this.I have a hard time understanding this part.\n4. New transfer scheme: This seems a novel point. The authors again use distributions to model the task encoding. This form may be of stronger approximation ability.\n\n\n# Clarity\n1. I like the Summary of Challenges in Fig.8, it helps quickly figure out the clues of this work.\n2. The authors give enough reference to key concepts, this helps me a lot in reviewing this paper.\n\n\n\n# Significance\nGood in the showed case, but narrow in the scenario.", "weaknesses": "\n\n\n\n\n\n\n\n# Clarity:\nOne of my major concerns is the clarity of this work. It seems like it was done in a hurry. The notations are used carelessly, usually not unified. I can see the same symbols with different superscripts or subscripts, or arguments. The components of the proposed method lack enough explanation of why they are helpful. Those make me restless, can't sleep overnights and recall them in dreams. I suspect several places use wrong notations, correct me if I am wrong. Some issues of clarity are as follows:\n1. Line 112 you miss an \"of\".\n2. It is $r_t^{\\tau}$ or $r_t$ in the second eq. of Eq (4)?\n3. I recommend to write Eq. (7) as $\\mathrm{Twohot}(y_t^{\\psi_k})\\log p_{\\psi_k}^T$ for clarity.\n4. Not every fig of this paper has an error bar.\n5. Notations are a bit confusing and frequently unexplained. I have a hard time understanding them.\n6. What does $\\theta^o$ mean in Eq. (6\\&7)?\n7. Righthand side of Eq. (7) seems not a scalar?\n8. Line 213, what does the loss function mean? Why do you design it? Also, is the state function also trained by it?\n\n# Significance\n1. The improvements are not significant.\n2. Environment is a bit narrow. Why only in 3D cases?\n3. The authors only study one dataset, making the result not strong enough.\n4. Limited ablation studies. No ablation about the loss, the new task encoding scheme, or the new SF approximation. The showed two ablations, in my opinion, are about minor issues.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to read this paper and for asking these clarifying questions. We have used your comments to improve the text. Due to space constraints, we respond to the most pressing comments and questions below.\n\n# Environment is a bit narrow. Why only 3D case? Using only one dataset makes the result not strong enough.\n\nWe emphasize that Playroom is a significant increase in complexity over prior SF work and discuss this in more detail in the shared response. We also significantly improve over the standard Playroom baseline, IMPALA. We acknowledge that applying our work to other domains would be interesting. We have added this discussion to the main text.\n\n# The improvements are not significant\n\nCan you please provide specific examples for why this is not significant? In 5.1, baseline methods fail at our tasks. In 5.2, we learn with 100 million+ fewer samples. Additionally, this is the first demonstration of transfer in a 3D environment with SF&GPI where all necessary representations are discovered. \n\n# Is computation burden increased?\n\nPlease see the joint response.\n\n# Clarifying how SFs are computed\n\nFirst note that we can represent any scalar $y$ as a weighted average of some of fixed bins $\\mathbb{B}=\\\\{b_1, \\ldots, b_M\\\\}$ and learnable weights $\\mathbf{p} = \\\\{p_1, \\ldots, p_M \\\\}$ as $ y \\approx \\sum^M_{m=1} p_m b_m$ where $b_m \\in \\mathbb{R}$ and $p_m \\in \\mathbb{R}^+$. $(\\mathbb{B},\\mathbf{p})$ define a pmf. We can learn $\\mathbf{p} = \\operatorname{softmax}(f_{\\theta}(x)) \\in \\mathbb{R}^M$ for inputs $x$ with a negative log-likelihood objective [1]. \n\nNow, consider approximating an n-dimensional SF-vector $\\psi^{\\pi_w}(s,a)\\in \\mathbb{R}^n$ with $\\psi_{\\theta}(s,a, w)$ using a pmf. Assuming all SF values can be represented with bins $\\mathbb{B}$, we can compute $\\psi_{\\theta}(s,a, w)$ as\n\n$$\n\\psi_{\\theta}(s,a, w) = \\\\{\\psi^k\\_{\\theta}(s,a, w)\\\\}^n_{k=1}\n\\quad\n\\psi^k\\_{\\theta}(s,a, w) = \\sum^M\\_{m=1} p^{\\psi\\_k}\\_m b\\_m\n\\quad\n\\mathbf{p}^{\\psi\\_k} = \\operatorname{softmax}(l\\_{\\theta}(\\ldots))\n$$\n\nwhere $\\mathbf{p}^{\\psi_k}$ are the weights for the $k$-th SF pmf. Leveraging this strategy for learning value-functions is known to improve learning by stabilizing gradients [1]. We highlight that SFs are types of value functions.\n\n[1] Improving Regression Performance with Distributional Losses\n\n# Eq 5, what is the SF dimension?\n\n$\\psi_{\\theta}(s,a, w) \\in \\mathbb{R}^n$ and $\\psi^k_{\\theta}(s,a, w) \\in \\mathbb{R}$. \n\n# What is B?\n\nWe defined $\\mathbb{B}=\\\\{b_1, \\ldots, b_M\\\\}$ with $M$ real numbers evenly spaced between $b_1$ and $b_M$.\n\n# What does it mean to write SFs as the expectation of some distribution?\n\nIt's the same as $\\psi^k_{\\theta}(s,a, w) = \\sum^M_{m=1} p^{\\psi_k}_m b_m$. We have changed the text to use this instead.\n\n# Righthand side of Eq 7 seems not a scalar?\n\nBoth $\\operatorname{twohot}\\left({y_t^{\\psi_k}}\\right)$ and $p^{\\psi_k}$ are length-$M$ vectors, so the righthand side of eq 7 is a scalar. We have clarified this in the text.\n\n# What does $\\theta^o$ mean in Eq 6 & 7?\n\nDuring Q-learning, we have a target $y_t = r_{t+1} + \\gamma Q_{\\theta}(s_{t+1}, a^*)$ that is *non-stationary* because $\\theta$  is changing. We mitigate this with *target parameters* $\\theta^o$ that update at a lower frequency. We have clarified this in the text.\n\n\n\n# Why do we omit ${L}_{\\psi}$? Are we aligning SFs and cumulants?\n\nThe original learning objective for USFA/MSFA has $L_\\psi$, $L_r$, and $L_Q$. Here, $L_\\psi$ aligns $\\psi_{\\theta}(s,a, w)$ with $\\mathbb{E}\\_{\\pi\\_w}\\left[\\sum^{\\infty}\\_{i=0} \\gamma^i \\phi_{t+i+1}\\right]$ via a regression loss. We **replace** $L_\\psi$ with $L^{\\tt cat}_{\\psi}$, a negative log-likelihood loss. Both estimate the same expectation.\n\n# Line 213, what does the loss function mean? Does this train the state function?\n\nThis loss indicates that we are performing gradient ascent with gradients $A_t \\nabla \\log p_{\\theta}\\left(\\alpha \\mid s_t, w\\right)$, where $A_t = R_t - V_{\\theta}\\left(s_t\\right)$ is the \"advantage\" of the coefficient chosen $\\alpha$ at time $t$. Here, $R_t = \\sum_{i=0}^{\\infty} r_{t+i+1}$  is the experienced return and $V_{\\theta}\\left(s_t\\right)$ is the predicted return. Optimizing this increases the likelihood of choosing coefficients in proportion to the \"advantage\" $A_t$.\n\nYes, this trains the new state function. We have changed the text to clarify this.\n\n# Only 2 ablations\n\nCould you please point out what ablations you believe are missing? We provided the following:\n\nFigure 4:\n\n1. no categorical: replaces pmf representation and negative-log-likelihood loss with scalar representation and regression loss.\n2. no stop gradient: ablates stopping gradients in loss.\n3. no $||w||$: ablates bounding norm of task encoding.\n4. independent: replaces sharing parameters across SF-dimensions with using different parameters for each dimension.\n\nFigure 7:\n\n1. Deep OK - binary: replaces bernoulli distribution with gaussian distribution.\n2. Deep OK - CSFA: replaces CSFA with USFA for estimating SFs during transfer.\n\nFigure 14: number of pmf bins\n\nFigure 15: number of resblocks in cumulant network\n\n# Q-value as linear combination of SFs is somewhat narrow\n\nIf SFs capture rich features, we think this dot-product has great potential for generalization. CLIP has demonstrated that a dot-product between visual features and language encodings can provide powerful generalization. This is analogous to our dot product between SFs (visual features) and task encodings (language encodings). \n\n# networks can be harder to train due to complicated designs.\n\nCan you please specify what part of our method might complicate training?\n\n# What does Line 144 mean?\n\nWe were trying to say that one can parameterize an SF-approximator for a policy $\\pi_w$ with an encoding of the task that defines it, i.e. we can approximate $\\psi^{\\pi_w}$ with $\\psi_\\theta(s, a, w)$. We have simplified the text.\n\n# Do you lose a $t+1$ in eq. 4?\nYes we did."}, {"review_id": "4J3Kjd7wvG", "reviewer": "Reviewer_RZN8", "summary": "This work focuses on attempting something which no previous work has done before;  learning task encodings, cumulants and successor features (SFs). The author(s) proposed a solution based on learning these three functions jointly. They considered a model that learns the SFs as probability mass function using the two-hot representations inspired from MuZero. They named this model as the Categorical Successor Feature Approximator (CSFA). With the learned SFs, the authors will then be able to perform Generalised Policy Iteration (GPI) to obtain the policy for the relevant task at evaluation phase. In order to find a better policy for this task, the authors also introduced a method named Deep Option Keyboard (DOK), which relies on a parameterised function to output coefficients which are then used as preference vectors. The models are evaluated in a 3D environment, using pixel observations. The tasks included going to certain locations and collecting and then dropping objects at different locations. In particular, DOK is used to study the effects of transfer in sparse reward and long horizon tasks. The proposed model is able to perform better than the baselines such as Universal Successor Features and Modular Successor Feature Approximators.\n", "questions": "Q1. The authors discretize the cumulant returns values using bins. Is there a study done to show how the different number of bins impact the performance of the agent? \n\nQ2. I understand that Bernoulli distribution was used as the parameterized probability function. What are the other distributions that were considered? Why did Bernoulli distribution do better than the others? Do you have any insight on this?\n\nQ3. On line 273, the authors mentioned that sharing an estimator across cumulants is necessary. What do you mean by keeping categorical representation? Is this successor features or state function or something related to cumulant network (with reference to figure 2). \n\n\n(Potential typos):\nLine 289: Known known tasks?\n\nCaption of figure 5: Distral and MTRL learn an ??? about the same speed. My guess is “an optimal or policy”?\n", "limitations": "None", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "I find the work interesting as this is an attempt at a challenging problem of learning task encodings, cumulants and successor features jointly. The figures, in particular, figure 2, is clear in aiding the reader to understand the architecture of their model. The writing overall was mostly clear and easy to read. The experiment's descriptions are concise and helped me to understand the setup a lot even though I am not familiar with the playroom environment. The authors also performed ablation studies which showed that the key to their model is learning the SFs using a probability mass function, rather than the canonical expected values.\n\n", "weaknesses": "Although there are no major weaknesses in my opinion, I do have some points for thought.\n\nFirstly,I am not sure it is a good idea to call the mechanism for learning the preferences, “Deep Option Keyboard” when it seems that there are no options involved. \n\nSecondly, it would also perhaps be a good idea to see how this model compares with the standard Q-learning algorithm, such as Rainbow or even just a distributional RL model since the proposed model uses a probability mass function to get an idea how much of the transfer benefits one could achieve from the utilization of the successor features.\n\nThirdly, there was some mention of static vs dynamic query. Perhaps it will be helpful to give a clear definition of what it means by static and dynamic query. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this positive review. We are happy you appreciated this submission. Below we respond in detail to your questions.\n\n**Changing name to not give false impression of using options.**\n\nThank you for pointing this out. We have changed the name to \"Successor Feature keyboard\".\n\n**Adding a Q-learning baseline**\n\nThank you for this suggestion. USFA, MSFA have both been shown to outperform Q-learning baselines for transfer and IMPALA-based MTRL has been shown to outperform Q-learning baselines in Playroom [1]; given that our approach outperforms all of them, we believe it would outperform Q-learning even more strongly. Unfortunately, we do not have the resources to implement and tune this baseline during this rebuttal period. We have included a discussion of this in the main text.\n\n**Making explicit the difference between static and dynamic query**\n\nGreat suggestion. We have made this more explicit in text. We summarize this here: a \"static query\" for transfer task $\\kappa'$ corresponds to the output of $w_{\\theta}(\\kappa')$ whereas a \"dynamic query\" is the output of the state-dependent function $g_{\\theta}(s, \\kappa')$.\n\n**Is there a study done to show how the different number of bins impact the performance of the agent?**\n\nWe ablate the number of bins in Figure 14 of the appendix. We find that using more bins is helpful, though fewer bins does reasonably well.\n\n**What kinds of distributions were considered? Why did a Bernoulli distribution do better?**\n\nIn addition to a Bernoulli distribution, we also considered a gaussian distribution as the original Option Keyboard did. We showed results for this in Figure 7 (Deep OK - binary). We did not consider other distributions since a Bernoulli distribution is simple to implement and it provided good results. We also studied auto-regressive variants of both distributions and found no improvements. \n\nWe believe that a Bernoulli distribution performed well because it facilitates exploiting SF&GPI. SF&GPI enable transfer to linear combinations of task encodings. By leveraging a Bernoulli distribution for sampling task encoding coefficients, the dynamic query is a linear combination of learned task encodings. This is precisely the type of transfer task encoding that SF&GPI has been shown to work well with. We have added a discussion on this to the \"discussion and conclusion\" section.\n\n**On line 273, the authors mentioned that sharing an estimator across cumulants is necessary. What do you mean by keeping categorical representation? Is this successor features or state function or something related to cumulant network**\n\nThanks for this clarifying question. We consider estimating $n$ SF with pmfs. Each pmf has a domain of $M$ values $\\mathbb{B}=\\\\{b_1, \\ldots, b_M\\\\}$. We have two choices for the $k$-th SF:\n\n1. CSFA: Use the same parameters but have an embedding of the dimension $e_k$ be input, i.e. $p_{\\psi_k} \\propto \\exp \\left(l_\\theta^\\psi\\left(s_t, w, e_k\\right)\\right) \\in \\mathbb{R}^M$\n2. Ablation: Use separate parameters $\\theta_k$, i.e. $p_{\\psi_k} \\propto \\exp \\left(l_{\\theta_k}^\\psi\\left(s_t, w\\right)\\right) \\in \\mathbb{R}^M$\n\nThe setting of this ablation is what is typically done when modeling SFs. When we say \"keeping the categorical representation\", we simply mean that we still model each SF with a pmf as apposed to using a scalar representation. We have clarified this in the text.\n\n**Caption of figure 5: Distral and MTRL learn an ??? about the same speed.** \n\nWe meant \"Distral and MTRAL learn **at** about the same speed\""}, {"review_id": "MuCRrhu0y2", "reviewer": "Reviewer_xth2", "summary": "This paper focuses on the Option Keyboard (OK) task which was recently proposed as a method for transferring behavioral knowledge across tasks. They propose a new method of Deep Option Keyboard (Deep OK), which enables transfer with discovered state features and task encodings. To enable discovery, they further propose the Categorical Successor Feature Approximator (CSFA), a novel learning algorithm for estimating SFs while jointly discovering state features and task encodings. With Deep OK and CSFA, this paper achieves the first demonstration of transfer with SFs in a challenging 3D environment where all the necessary representations are discovered. The experiment results show that CSFA can discover SFs compatible with SFGPI and Deep OK achieves strong jumpstart performance for transfer to combinations of training tasks.\n", "questions": "In Line 364, the author claims that  “This may enable other methods that leverage SFs in a multi-task setting”. Since Deep OK can’t get the optimal performance, how can leverage SFs for discovering representations and get the optimal performance simultaneously？", "limitations": "See the weakness and question sections.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tTransfering with SFGPI in a large-scale multi-task setting, discovering cumulants and task encodings, while sharing a task encoder and SF approximator across tasks gives some insights to the community and seems novel.\n2.\tThe paper’s writing is generally clear and the figures were helpful for understanding.\n3.\tThe problem formulation is clear and the approach is well-reasoned.\n4.\tThe experiments are sufficient and convincing to some extent.\n", "weaknesses": "1.  Though SFGPI is dynamic and can share a task encoder and SF approximator across tasks, its generalization ability is not clear and no more discussion can be found in the paper. Maybe a more general method that sampling queries from the task encoding space that is more fairly concentrated could be designed and at least provide some discussions.\n2. In Figure 3, CSFA discovers representations compatible with GPI across short and long task horizons. The longer horizon is defined by just one “Place Near” task. Will CSFA outperform MSFA and USFA on other general tasks?\n3. Deep Ok doesn’t reach an optimal performance but uses far fewer frames compared with Distral.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are glad that you found the writing clear and approach well-reasoned. Below we respond to specific questions.\n\n**Maybe a more general method than sampling queries from the task encoding space that is more fairly concentrated could be designed and at least provide some discussions.**\n\nWhile our analysis indicated that the learned task encoding space was somewhat concentrated, at the scale and complexity of our problem, we did not find this to be a bottleneck. In Figure 7, we compare againt sampling from a gaussian distribution (Deep OK - binary) and we show that this does not generally learn as quickly. \n\nIf future work does encounter this challenge, one solution might be to leverage a contrastive learning framework similar to CLIP [1] to learn task encodings that are better dispersed (for example, on a hyper-sphere). CLIP has been successfully applied to learning observation representations that are aligned with text representations. An agent's state representation is a representation of their stream of visual observations, and when goals are given via text, they are text representations. Thus, given the success of CLIP, it is reasonable that future work can leverage such a contrastive learning method to improve the quality of the task encoding space. We have added this to our discussion section.\n\n**The longer horizon is defined by just one “Place Near” task.** **Will CSFA outperform MSFA and USFA on other general tasks?**\n\nWe believe that the benefits of our approach will be general; indeed, we focus on \"place near\" tasks because they are substantially more complex and difficult than the tasks considered in prior work, which focused on simpler \"go to\" tasks [3-7]. We discuss this in more detail in the shared response above. Given the performance improvement of CSFA on these more complex tasks, nothing suggests it shouldn't continue to improve in other settings with similar difficulty.\n\n**Since Deep OK can’t get the optimal performance, how can leverage SFs for discovering representations and get the optimal performance simultaneously?**\n\nOne could improve/build on Deep OK to get optimal performance by integrating it with a distillation method. The main benefit of Deep OK (and SF&GPI more generally) is that it enables relatively good performance, relatively quickly. One can leverage the policy learned by Deep OK $\\pi^{\\tt SF}$ as a source policy in a policy distillation algorithm that transfers its knowledge to a fully parametric action policy $\\pi^{\\tt target}$. For example, by minimizing the KL-divergence between the two policies as the target policy learns to maximize reward, as with Kickstarting RL [2]. One could design this algorithm so that it trades off between selecting actions from $\\pi^{\\tt SF}$ to selecting actions from $\\pi^{\\tt target}$ over the course of learning. We have added this to our discussion section.\n\nWe emphasize that SF&GPI is not intended to provide optimal transfer performance. Instead, they are intended to provide reasonably good, fast transfer performance. See [3-5] for more. \n\n[1] Learning Transferable Visual Models From Natural Language Supervision\n\n[2] Kickstarting deep reinforcement learning\n\n[3] Successor Features for Transfer in Reinforcement Learning\n\n[4] Transfer in Deep Reinforcement Learning Using Successor Features and Generalised Policy Improvement\n\n[5] Universal Successor Features Approximators\n\n[6] Composing Task Knowledge with Modular Successor Feature Approximators\n\n[7] The Option Keyboard Combining Skills in Reinforcement Learning"}, {"review_id": "6DSs8FFYkl", "reviewer": "Reviewer_cEsx", "summary": "The author(s) present a novel and efficient approach to knowledge transfer in long-horizon tasks, aiming to minimize interactions with the environment (less frames required in training). Their proposed method leverages Categorical Successor Feature Approximator (CSFA) within an enhanced deep OK network, resulting in several tasks (e.g., Find, Place).", "questions": "see weakness and limitations", "limitations": "- This paper could benefit from an expanded evaluation of the transfer learning quality beyond just the success rate metric, e.g. adding visualizations of the baseline task encoding and dynamic query encoding to show the benefit of dynamic query is scalable.\n- This study lacks of discussion about the upper bound of the number of tasks it can effectively handle during sharing coding. The author primarily tests the approach on two main tasks (\"find\" and \"place\") with several sub-tasks or steps required to complete each main task. However, it remains unclear what the upper limit number of main task and sub-tasks?", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Proposed  a dynamic transfer query that enables sharing among task-encoding and SF-approximators;\n- Implemented this methodology in long-horizon tasks.\n", "weaknesses": "1. The authors did not clearly position their contribution in relation to other approaches commonly used in long-horizon tasks, such as hierarchical approaches, hindsight approaches and meta-learning approach [1-3].\n2. While the overall results of the experiments are significant (e.g., Figure 3 improve 0.2 ~ 0.8 in success rate), the description and documentation of some experiments might be unclear or insufficient. 1) USFA or MSFA might require fewer computational resources compared to MSFA, as their calculations appear to be simpler (line 246). However, without further detailed information, it is difficult to make a definitive judgment. I suggest that the authors include memory and computational load comparisons among these methods. 2) The author's comparisons are focused on MTRL, but there is a lack of comparison with other approaches mentioned in references [1-3]. It would be beneficial to include comparisons with other relevant approaches mentioned in reference [3].\n3. From the paper, it is unclear why MSFA is considered sharable or suitable for fitting different tasks. How it compare with other pre-training encodings or other method representations.\n4. Some typos, e.g., \"Transfer tasks are conjunctions of known known tasks\", line (289)\n\n[1] Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors, NeurIPS 2020\n\n[2] Learning to Reach Goals via Iterated Supervised Learning, NeurIPS 2020\n\n[3] Skill-based Meta-Reinforcement Learning, ICLR 2022", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for this review. Below we respond to specific questions.\n\n**The authors did not clearly position their contribution in relation to other approaches commonly used in long-horizon tasks, such as hierarchical approaches, hindsight approaches and meta-learning approach [1-3]. The author's comparisons are focused on MTRL, but there is a lack of comparison with other approaches mentioned in references [1-3]. It would be beneficial to include comparisons with other relevant approaches mentioned in reference [3]**\n\nThank you for pointing out this work. We have included a discussion on methods for learning long-horizon tasks in the paper and include these 3 papers. In summary, [1,3] assume access to offline datasets with useful (though potentially sub-optimal) behavior, and [2,3] demonstrate results on environments with hand-designed state-representations. In our work we don't assume offline datasets and aim to learn state and goal representations from partial observations of a pixel-based 3D environment. We do not have the resources to implement and compare to these methods in this rebuttal phase. \n\n[1] learns long horizon tasks via hiearchical planning, whereas we focus on transferring value functions via successor features. [1] assumes access to diverse (but potentially sub-optimal) demonstration data that is used to learn a dynamics model used for goal-directed planning. Our method does not assume any demonstration data. Thus, [1] is not applicable to our setting without augmenting it with a method to actively explore the environment to produce data for learning its dynamics model.\n\n[2] learns a goal-conditioned policy for data that is relabeled such that the final state in a trajectory is a re-defined as a goal state for a self-imitation learning objective. While this method is applicable to our setting, there are several differences between settings we study. First, they assume access to a hand-designed state-representation which also defines the goal-space. This is in contrast to our setting where an agent must learn state and goal representations from high-dimensional, pixel-based, partial observation of a 3D environment. This challenge is exacerbated by the environment's large state space that grows exponentially with the number of objects. Our tasks each have task objects and at least 2 distractor objects. In summary, it is not clear that [2] is applicable to our setting without modifications to accodomate (1) learning state- and goal-representations in a (2) partially observable, pixel-based 3D environment with (3) a large state-space induced by task and distractor objects.\n\n[3] learns long-horizon tasks by integrating meta-reinforcement learning with a hierarchical policy that leverages skills that are obtained from an offline dataset. In our work, we do not assume access to a dataset that has useful skills to be learned. Additionally, [3] studies environments with hand-designed, symbolic state representations whereas we focus on learning state representations for a partially observable, pixel-based, 3D environment. \n\n**USFA or MSFA might require fewer computational resources compared to CSFA, as their calculations appear to be simpler (line 246). However, without further detailed information, it is difficult to make a definitive judgment. I suggest that the authors include memory and computational load comparisons among these methods.**\n\nOur interpretation of this comment is that you are referring to the loss function of CSFA (eqs 6 & 7) vs. the loss function of USFA/MSFA (eq 4), (line 146 not 246). Please see the joint response for this.\n\n**From the paper, it is unclear why MSFA is considered sharable or suitable for fitting different tasks. How it compare with other pre-training encodings or other method representations.**\n\nWe consider MSFA because it (and USFA) are the only methods designed to share an SF-estimator across different task encodings, i.e. to learn $\\psi_{\\theta}(s,a,w)$. Other methods learn separate estimators $\\psi^{\\pi_i}_{\\theta_i}(s,a)$ with different parameters $\\theta_i$ for each task encoding $w_i$. We have clarified this in the text.\n\n\n**This study lacks of discussion about the upper bound of the number of tasks it can effectively handle during sharing coding. The author primarily tests the approach on two main tasks (\"fìnd\" and \"place\") with several sub-tasks or steps required to complete each main task. However, it remains unclear what the upper limit number of main task and subtasks?** \n\nWhile we do now know the upper limit on the number of main tasks our method can handle, we emphasize that this work is a significant increase in complexity from prior work in the literature[4-8]. We discuss this in more detail in our shared response above. We note that prior work focused on combining 4 \"go to\" tasks. If one combines 2 \"go to\" tasks, this leads to $\\left(\\begin{array}{l}4 \\\\\\\\ 2\\end{array}\\right)=6$ combinations. We have 8 find tasks and 24 place near tasks. If we combine just 2 place near tasks, this leads $\\left(\\begin{array}{l}24 \\\\\\\\ 2\\end{array}\\right)=276$ possible transfer combinations. While we did not explicitly study increasing the number of find and place tasks, we are optimistic that one can scale the capacity of a sufficiently expressive task encoder with the number and complexity of tasks being learned.\n\n\n\n[1] Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors\n\n[2] Learning to Reach Goals via Iterated Supervised Learning\n\n[3] Skill-based Meta-Reinforcement Learning\n\n[4] Successor Features for Transfer in Reinforcement Learning\n\n[5] Transfer in Deep Reinforcement Learning Using Successor Features and Generalised Policy Improvement\n\n[6] Universal Successor Features Approximators\n\n[7] Composing Task Knowledge with Modular Successor Feature Approximators\n\n[8] The Option Keyboard Combining Skills in Reinforcement Learning"}]

Author rebuttal:
We thank the reviewers for their very constructive feedback which we have used to improve the clarity and accessibility of our paper. We are encouraged that reviewers found that we tackle a novel challenge [R2, R3] with a well-reasoned approach [R3], have significant [R4] and convincing results [R2, R3] on a challenging problem [R2, R3, R4], with figures that help in understanding the method [R1, R2, R3] and writing that is clear and easy to understand [R2, R3]. Note that we have labled reviewers as follows: (R1: z7ra, R2: RZN8, R3: xth2, R4: cEsx). 

We respond to individual comments below but provide some common responses here. We have also added a figure to explain how we compute a successor feature with a learned probability mass function.

**Does CSFA use more compute or have more complex calculations compared to USFA or MSFA?**

We clarify that the calculations of CSFA are only slightly more complex than USFA/MSFA. Prior work that has learned Q-values, SFs, and cumulants has had 3 terms, $L=L\_\psi + L\_r + L\_Q$ [1-3] (we omit coefficients for simplicity). In our new objective, we **replace** $L\_\psi$ with $L^{\tt cat}\_{\psi}$. We accidentally omitted $L\_Q$ when describing the losses used by prior work in section 3. $\mathcal{L}^{\tt cat}\_\psi$ differs in that it estimates SFs with a negative log-likelihood loss whereas $\mathcal{L}\_\psi$ estimates SFs using mean-squared-error. 

We acknowledge that modelling SFs with a pmf does use more parameters, since you need parameters for every bin; however, these computations at the last layer are negligible compared to the rest of the network. We controlled for the number of parameters across SF-based methods by (1) increasing the size of the SF-estimators so that they were approximately equal and (2) keeping all other parts of the architecture the same between methods. CSFA used 7.8M parameters for its SF-estimatr, USFA used 8.1M parameters, and MSFA used 7.5M parameters. Once we do this, the computations in the very last layer where the loss is computed (a softmax followed up a dot-product, eq. 7) is negligible compared to the computations in the rest of the network with millions of parameters. For example, much more computation is done by the ResNets for the agent's vision core and cumulant network. These were shared across all SF-based methods. We have highlighted this in the experiments section of the text.

**Clarifying the contribution from studying playroom relative to prior experiments in the successor feature literature**

Here, we highlight that this work is a significant increase in complexity from prior successor feature literature [4-8]. The most complex task studied by prior work were "go to" tasks in the DMLab environment by [6]. Here, no object-interaction is required. (1) Playroom "place near" tasks add object-interaction, while increasing the task horizon and reward sparsity (2) prior work only combined 4 "go to" tasks while we combine 8 Find and 24 place nears tasks (3) prior work used DMLab with 8 actions [9] while Playroom has 46 actions. We have added this discussion to the main text.

[1] Successor Uncertainties: Exploration and Uncertainty in Temporal Difference Learning

[2] Composing Task Knowledge with Modular Successor Feature Approximators

[3] PsiPhi-Learning: Reinforcement Learning with Demonstrations using Successor Features and Inverse Temporal Difference Learning

[4] Successor Features for Transfer in Reinforcement Learning

[5] Transfer in Deep Reinforcement Learning Using Successor Features and Generalised Policy Improvement

[6] Universal Successor Features Approximators

[7] Composing Task Knowledge with Modular Successor Feature Approximators

[8] The Option Keyboard Combining Skills in Reinforcement Learning

[9] DeepMind Lab

Meta-review: This paper presents a method for transferring behavioral knowledge across tasks by jointly discovering state-features and task encodings while estimating successor features. The paper received 1 strong accept, 2 borderline accept, and 1 borderline reject recommendations from reviewers. The reviewers appreciated the challenging problem of jointly learning task encodings, cumulants and successor features, sound approach, thorough experiments, and good results. Questions and concerns included lack of discussions on some approach components, some related work (such as hierarchical, hindsight, and meta-learning approaches), compute/complexity, and Deep Option Keyboard optimization. Most of these concerns were addressed in the rebuttal and discussion. Some remaining concerns included lack of empirical evidence on how to improve/build on Deep Option Keyboard to obtain optimal performance. The borderline reject reviewer did not follow up on the post-rebuttal discussions, and also indicated their low confidence in their review – the ACs felt that this reviewer’s major questions and concerns were adequately addressed by the rebuttal. Overall, after carefully considering the paper, rebuttal, and discussions, the ACs feel that the positives outweigh the negatives, and recommend acceptance. It is recommended that the revision include a discussion on computation complexity and task limitations.
Decision: Accept (poster)