Title: Combining Behaviors with the Successor Features Keyboard
Venue: NeurIPS 2023
Accepted: yes
Keywords: deep reinforcement learning

successor features

transfer

generalization

feature-discovery
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

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.

1 Introduction

Consider a household robot that learns tasks for interacting with objects such as finding and moving them around. When this robot is deployed to a 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.

One promising method for transfer is the Option Keyboard (OK)[1, 2], which transfers to new tasks by adaptively combining subsets of known behaviors. OK combines known behaviors by leveraging Successor Features (SFs) and Generalized Policy Improvement (GPI) [3, 4]. SFs are predictive

TransferTask:Cooktomatoandpotato,thenplacethemonaplate

Query_1

Query_2

Cooktomato

Placeheldobjects

andpotato

onplate

Legend

w1

Query:

w

"keys"forknown,

Representation

GPI

Highest

GPI

SF-Key:

behaviors

“values"

Value:w

日

used to

cook-potato

a1, a2,...

generate

behavior

"values"

Highest

lplace-plate

used to

generate

Initialstate

Laterstate

behavior

→t

2

representations for behaviors. They represent behaviors with estimates of how much state-features (known as 'cumulants') will be experienced given that behavior. GPI can be thought of as a querykey-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 'cumulants' and task feature-preferences (i.e. task encodings). There is work that has discovered either cumulants or task encodings [3-6]. However, either (a) they only showed transfer with SF&GPI using a fixed (not dynamic) transfer query, (b) they only demonstrated results in 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 a dynamic query while leveraging discovered representations for cumulants and feature preferences. Further, a transfer method should scale to sparse-reward longhorizon 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[4]. 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 a dynamic query that leverages discovered representations while sharing functions across tasks, we propose two novel methods, the Successor Features Keyboard (SFK) and the Categorical Successor Feature Approximator (CSFA) . We present a high-level overview of in SFK in Figure1. SFK leverages CSFA to learn SF-estimates over discovered cumulants and task-preferences in a pretraining phase. Afterwards, in a finetuning phase, SFK learns to generate dynamic queries which are linear combinations of CSFA-discovered task-preferences. CSFA addresses challenges with estimating a non-stationary return by estimating SFs with a variant of the categorical two-hot representation introduced by MuZero [7]. We discretize the space of cumulant-return values into bins and learns a probability mass function (pmf) over them. Modelling cumulant-returns with a pmf is more robust to outliers and can better accomodate a shifting magnitude. In contrast, standard methods estimate SFs with regression of a single point-estimate [3, 8] which is succeptible to outliers and varying scales [9].

We study SFK and CSFA in Playroom [10], a challenging 3D environment with high-dimensional pixel observations and long-horizon tasks defined by sparse rewards. Prior work on transfer with SF&GPI has mainly focused on transfer in simpler 2D environments [11, 3-6]. While Borsa et al. [8] studied transfer in a 3D environment, they relied on hand-designed cumulants and task encodings [8] 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 SF&GPI using a dynamic query, discovered representations, and a task encoder and SF-approximator that are shared across tasks. (2) To enable discovery when sharing a task encoder and SF-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 a mapping from source domains to target domains [12], learn disentangled representations [13, 14], or learn a modular architecture [15, 16]. 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 a meta-controller [19]. Finally, to transfer value functions, some approaches learn universal value functions [20] and others exploit SFs [2]. 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 a goal-conditioned policy [21]. Another stratgy is to distill knowledge from one policy to another, as Distral does [22]. Distral works by learning two policies: one goalconditioned policy and another goal-agnostic 'centroid policy'. The action-likelihoods 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 a 'centroid' policy for sharing knowledge across tasks. When we transfer to longer horizon tasks with sparse rewards, neither MTRL nor Distral may provide a policy with good jumpstart performance [23]. In this work, we study jumpstart performance and exploit successor features (SFs) with generalized policy improvement (GPI) [3].

Method

disc.

ϕ

disc.

w

query

share

w

θ

share

ψ

θ

3D

transfer

✓

✓

static

✗

✗

✗

goto-n

β

✓

✓

static

✗

✗

✓

goto-n

✓

✗

static

✗

✗

✓

goto-c

β

✓

✓

static

✗

✗

✗

goto-c

Borsa et al. [8]

✗

✗

static

✗

✓

✓

goto-c

✓

✗

static

✗

✓

✗

goto-c

✗

✗

dynamic

✗

✗

✗

goto-c

SFK (ours)

✓

✓

dynamic

✓

✓

✓

place-c

Successor Features are useful because they enable computing of action-values for new task encodings [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 a '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 a static transfer query and did not share SF-estimators or task-encoders across tasks [3, 24, 3, 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 are shared across tasks. Additionally, prior work has only studied transfer to combinations of short-horizon 'go to' tasks whereas we include longer horizon 'place' tasks and do so in a 3D environment. We summarize these differences in Table1.

3 Background

We study an RL agent's ability to transfer knowledge from a set of n κ training tasks T train = { κ 1 , . . . , κ n κ } to a set of n ′ κ transfer tasks T new = { κ n 1 , . . . , κ n n ′ κ } . During training, tasks are sampled from distribution p train ( T train ) . At transfer, tasks are sampled from distribution p transfer ( T new ) . Each task is specified as a Partially Observable Markov Decision Process (POMDP, [25]), M i = ⟨S e , A , X , R, p, f x ⟩ , where S e , A and X are the environment state, action, and observation spaces. Rewards are parameterized by a task description κ , i.e. r κ t +1 = R ( s e t , a t , s e t +1 , κ ) is the reward for transition ( s e t , a t , s e t +1 ) . When the agent takes action a t in state s e t , s e t +1 is sampled according to p ( ·| s e t , a t ) , an observation x t +1 is generated via f x ( s e t +1 ) , and the agent gets reward r κ t +1 . We assume the agent learns a recurrent state function that maps histories to agent state representations, s t = s θ ( x t , s t -1 , a t -1 ) . Given this learned state, we aim to obtain a behavior policy π ( s t , κ ) that maximises the expected reward when taking an action a t in state s t : i.e. that maximizes

Q π,κ t = Q π,κ ( s t , a t ) = E π [∑ ∞ t =0 γ t r κ t +1 ] . 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 = ϕ θ ( s t , a t , s t +1 ) , which are useful 'descriptions' of a state-transition, and (2) a task encoding w = w θ ( κ ) , which define 'preferences' over said transitions. Reward is then defined as r κ t = ϕ ⊤ t w [3]. Successor Features ψ π are then value functions that describe the discounted sum of future ϕ that will be experienced under policy π :

Given ψ π t , we can obtain action-values for r κ t as Q π,κ t = ψ π t ⊤ w .

The linear decomposition of Q π,κ is interesting because it can be exploited to re-evaluate ψ π for new tasks with GPI. Assume we have learned SFs { ψ π i ( s, a ) } n κ i =1 for tasks T train . Given a new task κ ′ , we can obtain a new policy π ( s t ; κ ′ ) 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,

where w θ ( κ ′ ) is a static transfer query for transfer task κ ′ .

Option Keyboard . One benefit of equation 2 is that it enables transfer to linear combinations of training task encodings. However, it has two limitations. First, the feature 'preferences' w θ ( κ ′ ) are fixed across time. When κ ′ 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 w θ with a nonlinear function approximator such as a neural network, there is no guarentee that w θ ( κ ′ ) is in the span of training task encodings. The 'Option Keyboard' [1, 2] can circumvent these issues by learning a transfer policy that maps states and tasks to a dynamic transfer query g θ ( s, κ ′ ) :

Learning . In the most general setting, we learn ψ θ , ϕ θ , w θ and g θ from experience. Rewards r κ and their task encodings w = w θ ( κ ) reference deterministic task policies π w that maximize them. Borsa et al. [8] showed that this allows us to parameterize an SF-approximator for a policy π w with an encoding of the task that defines it, i.e. we can approximate ψ π w with ψ θ ( s, a, w ) . This is known as a Universal Successor Feature Approximator (USFA) and can be learned with TD-learning with cumulants as pseudo-rewards. To discover ϕ θ and w θ , we can match their dot-product to the experienced reward [3]. Defining ϕ t +1 = ϕ θ ( s t , a t , s t +1 ) , we summarize this as:

No prior work has jointly learned a task encoder w θ and USFA ψ θ ( s, a, w ) while discovering cumulants ϕ θ . In this work, we introduce the Successor Features Keyboard to address these limitations to enable transfer with SF&GPI and discovered representations in a 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, we propose the Categorical Successor Feature Approximator for jointly learning SFs ψ θ , cumulants ϕ θ , and task encodings w θ . The rest of this section is structured as follows. In §4.1, we 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 Approximator (CSFA) , composed of a novel architecture, shown in Figure 2, and a novel learning objective (equation 6). Challenge : when jointly learning a Universal SF-approximator ψ θ and a cumulant-network ϕ θ

CategoricalSFApproximator

TransferbehaviorwithSuccessorFeaturesKeyboard

πw(s,a)

Training task

π(s, Tnew) = argmaxa u

[be(s, a, w)  go(s, Wnew)}

behaviorswe

wEWtrain

Shared

wanttotransfer

Head

params

Head

n

w1

network

SF

θ(s,a,w1)

+0→Qw1,g

Structured,Categorical

Network

→Action

θ(s,a,w)

Cumulant

Wn

network

SF

θ(s,a,wn)

→0→Qwn,g

GPI

Task

Network

Encoder

Φe(s,a)

Statefn

Statefn

New

Transfer

ge(s, Wnew)

w

Statefn

policy

Task

dt

2

Wnew

Observation

Observation

Encoder

Encoder

Legend

NewTask

Encoder

Frozen

+C

Ct

parameters

TransferTask

for long-horizon tasks, ψ θ needs to fit ϕ θ -generated returns that are potentially high-variance, nonstationary, and changing in magnitude [9]. CSFA addresses this challenge by modelling SFs with a probability mass function (pmf) over a discretized range of continous values. CSFA then fits this data by leveraging a categorical cross-entropy loss, enabling our estimator to give probability mass to different ranges of values. This is in contrast to prior work that models SFs with a point-estimate that is fit via regression [3, 8]. Leveraging a point-estimate can be unstable for modelling a non-stationary target with changing magnitude [9] (we show evidence in §C.1).

W

Exampleprobabilities

(s,a,w) =

M

ph bm

andbinvalues

m=1

PM

4n

p

4n

pi

n

B

learned

4

binvalues

-.5

0

5

probabilities

b1

bM

B

S

Statefn

Encoder

Task

Dimension

Encoder

↑

Observation

Encoder

口口

2t

Task

one-hotrep of

dimensionk

Architecture . CSFA represents SFs with a pmf ( B , p ψ k ) , where B = { b 1 , . . . , b M } are an apriori defined set of bin values and p ψ k are probabilities for each bin value. Specifically, CSFA estimates an n-dimensional SF-vector ψ π w ( s, a ) ∈ R n with ψ θ ( s, a, w ) and represents the k -th SF dimension as ψ k θ ( s, a, w ) = ∑ M m =1 p ψ k m b m . At each time-step, we update a state function s θ with an encoding of the current observation z t = f enc ( x t ) , the previous action a t -1 , and the previous state representation s t -1 , i.e. s t = s θ ( z t , a t -1 , s t -1 ) . Each set of probabilities is computed as p ψ k = softmax( l θ ( s t , w, e k )) , where e k is an embedding for the current SF dimension. In summary,

We provide a diagram of representing an SF with a pmf in Figure 3. Using a pmf allows us to 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 Figure 6.

Learning objective . We learn to generate behavior by employing a variant of Q-learning. In particular, we generate Q-values using learned SFs, and use these Q-values to create targets for both estimating Q-values and for estimating SFs. For both, targets correspond to the action which

maximized future Q-estimates. We learn SFs with a 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,

a ∗ = arg max a ψ ( s t +1 , a, w ) ⊤ w , where p ψ k = softmax( l θ ( s t , w, e k ) , w is a stop-gradient operation on w . Like prior work [26], we mitigate non-stationary in the return-targets y Q t and y ψ k t by having target parameters θ o that update at a slower rate than θ . The overall loss is L = β Q L Q + β ψ L cat ψ + β r L r .

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 w θ through ψ θ ( s, a, w ) or through ψ θ ( s, a, w ) ⊤ w during Q-learning lead to dimensional collapse [27] and induces a small angle between task encodings (see §C.1). We hypothesize that this makes ψ θ ( s, a, w ) unstable and manifests as poor GPI performance. D2 . When || w || is large, it can make ψ θ ( s, a, w ) unstable. We hypothesize that this is because it magnifies errors due to the SF-approximation error that loosen previously found bounds on GPI performance [8, 5]. We discuss this in more detail in §B. To mitigate this, we bound w by enforcing it lie on a unit sphere, i.e. w = w θ ( κ ) || w θ ( κ ) || . This makes the SF-error be consistent across tasks.

4.2 Transfer with the Successor Features Keyboard

The original Option Keyboard (equation 3) learned a policy that mapped states s t to queries w , w ′ = g ( s t , κ n ) . However, they used a hand-designed w θ and thus hand-designed space for w . In our setting, we learn w θ . However, GPI performance is bound by the distance of a transfer query w ′ to known preference vectors w [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 { α i } n κ i =1 for known preference vectors W train = { w 1 , . . . , w n κ } . 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 ψ π w ( s t , a ) for w i ∈ W train :

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 θ n . Given a new state representation s n t , we sample coefficients independently:

We find that a Bernoulli distribution performs well. We learn this α -coefficient policy with policy gradients [28] by performing gradient ascent with gradients A t ∇ log p θ ( α | s t , w ) , where A t = R t -V θ n ( s n t ) is the 'advantage' of the coefficient α chosen at time t . Here, R t = ∑ ∞ i =0 r t + i +1 is the experienced return and V θ n ( s n t ) is the predicted return. Optimizing this increases the likelihood of choosing coefficients in proportion to A t .

5 Experiments

We study transfer with sparse-reward long-horizon tasks in the complex 3D Playroom environment [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 §5.2, we study the utility of SFK for transferring to sparse-reward long-horizon tasks.

Environment setup . We conduct our experiments in the 3D playroom environment. Observations are partial and egocentric pixelbased images. The agent gets no other information. Actions . The agent can rotate its body and look up or down. To pick up an 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. Training tasks . The agent experiences n κ = 32 training tasks T train = { κ 1 , . . . , κ n κ } composed of 'Find A' and 'Place A near B'. | A | = 8 and | B | = 3 . All tasks provide a reward of 1 upon-completion. We provide more details in §F.

5.1 Evaluating the utility of discovered representations for SF&GPI

Our first experiments are a 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 W train = { w 1 , . . . , w n κ } , the agent acts according to policy π ( s, w i ) = arg max a max w k { ψ ( s, a, w k ) ⊤ w i } . When w k = w i , π ( s, w i ) = arg max a Q ( s, a, w i ) . This will fail if the agent hasn't learned representations that support GPI. If an agent cannot perform GPI on training tasks, then it 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 a significantly longer horizon 'Place Near' task where the agent must select an object and hold it as it 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 ψ across tasks while discovering ϕ and w ? Q2 . Is each piece of CSFA necessary?

Find A (GPI)

PlaceAnearB (GPI)

1.0

1.0

0.8

0.8

Rate

0.6

0.6

Success

0.4

0.4

CSFA(ours)

USFA(Borsa etal,2019)

0.2

MSFA(CarValhoetal,2023)

0.2

0.0

0.0

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

Numlearnerframes1e9

Numlearnerframes

51e9

Baselines . (1) Universal Successor Feature Approximators (USFA) [8] is the only method that shares an SF estimator across tasks and has shown results in a 3D environment. (2) Modular Successor Feature Approximators (MSFA) [6] showed that leveraging modules improved SF estimation and enabled cumulant discovery. However, they did not discover task encodings and only showed results in simple grid-worlds. Both baselines estimate SFs with point-estimates. Comparing to them tests (a) the utility of our categorical representation and (b) CSFA's ability to discover both cumulants and task encodings that enable GPI in a large-scale setting.

CSFA discovers SFs compatible with SF&GPI while baseline methods cannot . For fair comparison, 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, we parameterized ϕ θ 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 5 shows that USFA and MSFA can perform GPI for Find tasks; however, neither learn place tasks in our computational budget. Given that we 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 6 shows that when CSFA shares an estimator but produces scalar estimates (CSFA - no categorical), GPI performance degrades

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 cumulants is necessary . If we keep our categorical representation but don't share it across cumulants (CSFA independent), GPI performance degrades on our long-horizon place near task. Stopping gradients from Qlearing is necessary . Interestingly, when we pass gradients to the taskencoder from Q-learning (CSFA - no stop grad), we can get perfect train performance on place tasks but highly unstable GPI performance. We found that passing gradients leads to dimensional collapse [27] (see §C). We hypothesize that this makes ψ θ unstable. Likewise, if we don't bound tasks (CSFA - no || w || ), we find degraded GPI performance compared to training but for even simpler tasks. While other methods may work, enforcing w

Find A (GPI)

Place AnearB (GPI)

1.0

1.0

0.8

0.8

Rate

0.6

0.6

Success

0.4

0.4

CSFA (ours)

CSFA (ours)

0.2

CSFA-noIIwlI(D2)

0.2

CSFA-independent

CSFA-nocategorical

CSFA-no stop grad（D1)

0.0

0.0

1

Numlearnerframes

2

3

4

5

6

7

1e9

8

0

1

Numlearnerframes

2

3

4

5

6

7

1e9

8

lie on a unit-sphere is a simple solution. 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 to 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 a reward of 1 and place tasks contribute a reward of 2 . We provide more details in the §F.

Research questions . Q3 . How do we compare against baseline transfer methods? Q4 . How important is it to use CSFA and to sample coefficents over known task encodings?

FindA&PutBnearC

PutAnearB&CnearD

Put × 3

5

2.8

3.5

Reward

3.0

4

2.6

2.5

3

Average

2.4

2.0

100

million

2.2

Deep OK (ours)

1.5

Distral (Teh et al, 2017)

1

MTRL

1.0

2.0

0

0

1

2

3

4

5

0

1

2

m

4

5

0

1

2

3

4

5

Frames

1e8

1e8

Frames

1e8

Frames

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 [31, 32, 10]. (2) Distral [22] is a common transfer learning method which learns a centroid policy that is distilled to training task policies. Comparing against MTRL and Distral tests the utility of combining behaviors with SF&GPI.

Q3: SFK has better jumpstart performance . Figure 7 shows that all methods get similar performance 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 we remove our curriculum (Figure 8)

this gap further increases. No method does well for our longest task (Put x 4) which involves 8 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, or (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 coefficents over known task encodings is critical to jumpstart performance . Figure 9 shows us that if we don't 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) so its not surprising it fails on novel transfer tasks. Directly sampling from the encoding space does not transfer as quickly. Empirically, we find that learned task encodings have a high cosine similarity, 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 to produce meaningful behavior.

6 Discussion and conclusion

We have presented SFK, a novel method for transfer that adaptively combines known behaviors using SF&GPI. To discover representations that are compatible with SF&GPI, SFK estimates SFs with a novel algorithm, CSFA. CSFA constitutes both a novel architecture (of the same name) which approximates SFs with a pmf over a discretized continous values and a novel learning objective which estimates SFs for discovered cumulants with a categorical cross-entropy loss.

We first showed that CSFA is able to discover SF&GPI-compatible cumulants and task encod-

Putx3（no curriculum)

Putx4（nocurriculum)

2.5

4

Reward

2.0

m

200

Deep OK(ours)

Average

million

1.5

Distral (Teh et al, 2017)

2

MTRL

1.0

1

0.5

0

0.0

0

1

2

3

4

5

0

1

2

3

4

5

Frames

1e8

Frames

1e8

FindA&PutBnearC

PutAnearB&CnearD&EnearF

3.0

2.5

Reward

2.0

3

Average

1.5

2

1.0

Deepok-binary

0.5

DeepOK-CSFA

1

Deep OK(ours)

0.0

0

1

2

3

4

5

0

1

2

3

4

5

Frames

1e8

Frames

1e8

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 a 3D environment but hand-designed representations, and (2) MSFA, which discovered cumulants but did so in gridworlds with 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-returns 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 a task-specific policy and a 'centroid' policy. Our results showed that, for long-horizon tasks, SFK could transfer with 100 million+ fewer samples when a curriculum was present, and 200 million+ fewer samples when no curriculum was present. We found that simply using a Bernoulli distribution for sampling task coefficents performed well because it facilitates exploiting SF&GPI. 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 of cumulants and task encodings that enabled transfer with a 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 we discovered were fairly concentrated. Future work may mitigate this with contrastive learning [33]. This respects the dot-product relationship of cumulants and task encodings while enforcing they be spread in latent space. Finally, while we demonstrated good jumpstart performance for long-horizon tasks, we did not reach optimal performance in our sample budget. Despite building on the Option Keyboard, we did not employ options as we sampled a new query at each time-step. Future work may improve performance by sampling queries more sparsely and by implementing a latent 'initiation set' [34].

This paper did not study the upper limit on the number of main tasks our method can handle. However, it was a significant increase in complexity from prior work [3, 4, 8, 6, 1]. For example, if one combines 2 'go to'tasks, this leads to C (4 , 2) = 6 combinations. We have 8 find tasks and 24 place near tasks, where C ( n, k ) denotes n choose k combinations. If we 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 a 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 a 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 a 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] or for learning cognitive maps [38]). We are hopeful that SFK and CSFA will enable SFs to be adopted more broadly with less hand-engineering.

References

Andr´ e Barreto, Diana Borsa, Shaobo Hou, Gheorghe Comanici, Eser Ayg¨ un, 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.

Andr´ e Barreto, Shaobo Hou, Diana Borsa, David Silver, and Doina Precup. Fast reinforcement learning with generalized policy updates. PNAS , 117(48):30079-30087, 2020.

Andr´ e Barreto, Will Dabney, R´ emi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. NIPS , 30, 2017.

Andr´ e 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.

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 3305-3317. PMLR, 2021.

Wilka Carvalho, Angelos Filos, Richard L Lewis, Satinder Singh, et al. Composing task knowledge with modular successor feature approximators. arXiv preprint arXiv:2301.12305 , 2023.

Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

Diana Borsa, Andr´ e Barreto, John Quan, Daniel Mankowitz, R´ emi Munos, Hado Van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. ICLR , 2019.

Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. The Journal of Machine Learning Research , 22(1):12348-12355, 2021.

Josh Abramson, Arun Ahuja, Iain Barr, Arthur Brussee, Federico Carnevale, Mary Cassin, Rachita Chhaparia, Stephen Clark, Bogdan Damoc, Andrew Dudzik, et al. Imitating interactive intelligence. arXiv preprint arXiv:2012.05672 , 2020.

Majid Abdolshah, Hung Le, Thommen Karimpanal George, Sunil Gupta, Santu Rana, and Svetha Venkatesh. A new representation of successor features for transfer across dissimilar environments. In International Conference on Machine Learning , pages 1-9. PMLR, 2021.

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.

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.

Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiositydriven exploration. arXiv preprint arXiv:1905.09275 , 2019.

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.04671 , 2016.

Chrisantha Fernando, Dylan Banarse, Charles Blundell, Yori Zwols, David Ha, Andrei A Rusu, Alexander Pritzel, and Daan Wierstra. Pathnet: Evolution channels gradient descent in super neural networks. arXiv preprint arXiv:1701.08734 , 2017.

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:1511.06295 , 2015.

Richard E Bellman. Dynamic programming . Princeton university press, 2010.

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.

Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In International conference on machine learning , pages 1312-1320. PMLR, 2015.

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.

Yee Teh, Victor Bapst, Wojciech M 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.

Zhuangdi Zhu, Kaixiang Lin, Anil K Jain, and Jiayu Zhou. Transfer learning in deep reinforcement learning: A survey. arXiv preprint arXiv:2009.07888 , 2020.

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, 2017.

Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence , 101(1-2):99-134, 1998.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. arXiv preprint arXiv:2110.09348 , 2021.

Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 1999.

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.

Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain 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.

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.

Felix Hill, Andrew Lampinen, Rosalia Schneider, Stephen Clark, Matthew Botvinick, James L McClelland, and Adam Santoro. Environmental drivers of systematicity and generalization in a situated agent. arXiv preprint arXiv:1910.00571 , 2019.

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 815-823, 2015.

Vivek Veeriah, Zeyu Zheng, Richard Lewis, and Satinder Singh. Grasp: Gradient-based affordance selection for planning. arXiv preprint arXiv:2202.04772 , 2022.

David Janz, Jiri Hron, Przemysław Mazur, Katja Hofmann, Jos´ e Miguel Hern´ andez-Lobato, and Sebastian Tschiatschek. Successor uncertainties: exploration and uncertainty in temporal difference learning. Advances in Neural Information Processing Systems , 32, 2019.

Tarun Gupta, Anuj Mahajan, Bei Peng, Wendelin B¨ ohmer, and Shimon Whiteson. Uneven: Universal value exploration for multi-agent reinforcement learning. In International Conference on Machine Learning , pages 3930-3941. PMLR, 2021.

Momchil S Tomov, Eric Schulz, and Samuel J Gershman. Multi-task reinforcement learning in humans. Nature Human Behaviour , 5(6):764-773, 2021.

William De Cothi, Nils Nyberg, Eva-Maria Griesbauer, Carole Ghanam´ e, 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):3676-3689, 2022.

Ehsan Imani and Martha White. Improving regression performance with distributional losses. In International conference on machine learning , pages 2157-2166. PMLR, 2018.

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, Iurii 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, Wojciech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http://github.com/deepmind .

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Felix Hill, Olivier Tieleman, Tamara Von Glehn, Nathaniel Wong, Hamza Merzic, and Stephen Clark. Grounded language learning fast and slow. arXiv preprint arXiv:2009.01719 , 2020.

Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning , pages 1928-1937. PMLR, 2016.

Charles Beattie, Joel Z Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich K¨ uttler, Andrew Lefrancq, Simon Green, V´ ıctor Vald´ es, Amir Sadik, et al. Deepmind lab. arXiv preprint arXiv:1612.03801 , 2016.

A Learning to estimate scalar values with distributional losses

We consider the setting where we are learning a function that maps an input x ∈ X to a scalar y ∈ R .

Representing scalars with probability mass functions (pmfs) . One can represent any scalar y as a weighted average of some of fixed bins B = { b 1 , . . . , b M } and learnable weights p = { p 1 , . . . , p M } as

where b m ∈ R and p m ∈ R + . ( B , p ) jointly define a pmf. We consider the setting where B is defined apriori.

Learning pmf representations of scalars . We can learn the weights p with a neural network f θ that computes p = softmax( f θ ( x )) ∈ R M by using a negative log-likelihood objective [39]. MuZero [7] showed that one can construct an effective target from a scalar y with the two-hot representation twohot( y ) ∈ R M . The resultant objective is then L NLL = twohot( y ) ⊤ log( p ) . In summary,

Prior literature has shown that leveraging this strategy for representing value-functions can improve learning (e.g. by stabilizing gradients) [39, 7].

The two-hot encoding twohot( y ) is a generalization of the one-hot encoding where all elements are 0 except for the two entries closest to y at positions m and m +1 . These two entries sum up to 1, with more weight given to the entry that is closer to y :

Representing a successor feature with a pmf . In our setting, we approximate an n-dimensional SFvector ψ π w ( s, a ) ∈ R n with ψ θ ( s, a, w ) . We learn a different set of weights p for each dimension and assume all SF values lie between b 1 and b M of an apriori defined set of bin values B = { b 1 , . . . , b M } . We'll reference the weights for the k -th SF dimension as p ψ k = { p ψ k 1 , . . . , p ψ k M } . For the k -th SF dimension, we compute these weights as p ψ k = softmax( l θ ( s, w, e k )) , where s is a state representation, w is a task representation, and e k is an embedding of that SF dimension. We can then compute the k-th SF as ψ k θ ( s, a, w ) = ∑ M m =1 p ψ k m b m . We present a schematic of this in Figure 3.

B Challenges of SF&GPI with learned representations

Our goal is to learn a successor feature (SF) estimator ψ θ ( s, a, w ) that (a) uses the output of a cumulant function ϕ = ϕ θ ( s, a ) as its prediction target and (b) is parameterized by the output of a task encoder w = w θ ( κ ) . In this section, we will first describe how these terms show up in the performance bounds for GPI (equation 15). Afterwards, we will describe how these terms show up in empirical challenges we faced when trying to learn these values for sparse-reward long-horizon tasks (§B.1). We will conclude with how these challenges manifest in the simplified setting of §5.1 where we 'transfer' to known training tasks (§B.2).

GPI performance bound . Given learned ψ θ , ϕ θ , and w θ , we want to leverage GPI to transfer to task encoding w ′ = w θ ( κ ′ ) . In this setting we have learned n policies { π w i } n i =1 for n tasks W = { w i } n i =1 , and can leverage GPI to transfer to w ′ as follows

Define π ∗ as the optimal policy for task κ ′ . Filos et al. [5] showed that the gap between a tasks's optimal Q-value Q ∗ ( s, a, w ′ ) and the Q-value for the GPI-policy ˜ Q π ( s, a, w ′ ) can be bound by (1) the SF approximation error δ Ψ = || ψ π w -˜ ψ π w || ∞ , (2) the reward approximation error δ r =

Legend

Unbounded WA

Design choice

Empirical Phenomena

(2)

Large

I/wl|

(4)

GPIbound term

Long-horizon

(Borsa et al, 2019)

Environment factors

Observationswemake

Weaker GPI

Error in

-(Borsa et al,2019)→

Bound

(3)

Smal θ(w, W)

Unstable

Me(s,a, w)

Error in

(1)

=

w

Passing Q-learning

gradients to Wθ

(Barreto et al, 2018)

/

(4)

Sparse reward

Learning ΦeI

Scalar θ l

L

|| r κ ′ -ϕ ⊤ w ′ || ∞ , (3) the magnitude of the transfer encoding || w ′ || , (4) distance of w ′ to the closest training task, ∆ w = min j ∈{ 1 ,...,n κ } || w ′ -w j || . That is, ∀ s, a

Intuitively, this says that the error on the esimated Q-values for task encoding w ′ will grow in proportion to (1) the distance of the task encoding w ′ to known tasks ∆ w , (2) the product of the SF-approximation error and task encoding magnitude || w ′ || δ Ψ , and (3) the reward-approximation error δ r .

B.1 Empirical observations

We found some design choices can weaken GPI performance and connect then to equation 15 (e.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 over the course of learning. As cumulants change in magnitude, their corresponding return can quickly change (see Figure 12). Empirically, we found that estimators that use scalar point-estimates have oscillating SF-errors δ Ψ (see Figure 13). While they can 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 ′ || δ Ψ . Empirically, we found that if we did not bound || w || , it can have a relatively large magnitude (see Figure 15a). In figure 11, we provide intuition for why this can inhibit GPI performance. We confirm this empirically in Figure 20.

Challenge 3 . In this work, we focus on learning an SF-approximator parameterized by a task-encoder ψ θ ( s, 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 for differentiating between training tasks since w is an input. This may manifest with poor GPI performance as shown in Figure 21.

B.2 Why might SF&GPI fail for transferring to training tasks?

Our first experiments showed that jointly learned SFs ψ θ , cumulants ϕ θ , and task encodings w θ can exhibit poor GPI performance for transferring to training tasks (§5.1). Here, we provide intuition for why this can happen.

Legend

Successor Feature

SuccessorFeatureapproximation

口

Task Encoding

w

?

42

1

w2

E

When we transfer to a train task, the agent simply needs to act with the SF corresponding to the task policy, i.e. we want task selected in the inner max w sel = max w ψ ( s, a, w ) ⊤ w i = w i . In this setting, the distance to the closest training task is 0 , i.e. ∆ w = min j ∈{ 1 ,...,n κ } || w ′ -w j || = 0 . This leads to the following change to equation 15.

This tells that, even when we are transferring to a known task , performance is bound by (1) the product of the SF-approximation error and task encoding magnitude || w ′ || δ Ψ , and (2) the rewardapproximation error. The second term is intuitive-if we have not learned cumulants that predict reward, then we cannot learn SF-approximators that predict action-values for task reward. Below, we unpack the first term.

Define ψ i = ψ ( s, a, w i ) . Now assume that approximations for SFs are bounded by an error ϵ , i.e. ˜ ψ i = ψ i ± ϵ . Now consider a simplified setting with only two tasks w 1 and w 2 and their SFs ψ 1 and ψ 2 , shown in Figure 11. Consider computing GPI for w 2 . First , simply due to the error ϵ , we can have that ˜ ψ ⊤ 1 w 2 > ˜ ψ ⊤ 2 w 2 . Second , if || w || is large, it can dominate the dot-product and lead to ˜ ψ ⊤ 1 w ′ 2 > ˜ ψ ⊤ 2 w ′ 2 .

C Additional analysis

Here we present some additional analysis. In §C.1, we show additional analysis in a simplified curriculum of only 'Find' tasks. In §C.2, we show more ablations for our method.

C.1 Analysis demonstrating observed learning challenges

We conduct additional analysis on n κ = 8 training tasks T train = { κ 1 , . . . , κ n κ } composed of 'Find A' where | A | = 8 . All tasks provide a reward of 1 upon-completion. This is 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:

The magnitude of cumulants changes during learning (see Figure 12).

The SF TD-error can be more unstable for methods that estimate SFs with a point-estimate (see Figure 13).

|| w || can be relatively large if we don't bound it (see Figure 15a).

Passing gradients to the task encoder leads to dimensional collapse (see Figure 15b)

CumulantAverage

CumulantLl

ReturnfordifferentCumulantMagnitudes

0.05

0.30

CSFA(ours)

CSFA(ours)

6

0.04

0.25

Return

0.03

0.20

Discounted

4

0.02

0.15

3

0.01

0.10

2

0.00

0.05

-0.01

0.00

0.0

0.2

0.4

0.6

0.8

1.0

0.0

0.2

0.4

0.6

0.8

1.0

10~4

103

10-2

101

Frames

1e9

Frames

1e9

AverageCumulantValueateachtime-step

SF TD error |6yl

0.10

0.08

0.06

0.04

CSFA (ours)

0.02

CSFA-no categorical

USFA (Borsa et al, 2019)

0.00

0.0

0.2

0.4

0.6

0.8

1.0

Frames

1e9

Find A

Find A (GPI)

1.0

1.0

0.8

0.8

Rate

0.6

CSFA(ours)

0.6

Success

CSFA-no categorical

0.4

USFA(Borsa et al,2019)

0.4

0.2

0.2

0.0

0.0

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

0.0

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Frames

1e9

Frames

1e9

llwll

12

10

8

CSFA (ours)

6

CSFA-no ||wlI (D2)

4

2

0.00

0.25

0.50

0.75

1.00

1.25

1.50

Frames

1e9

Task Similarity

1.00

0.98

0.96

CSFA (ours)

0.94

CSFA - no stop grad (D1)

0.92

0.90

0.88

0.0

0.2

0.4

0.6

0.8

1.0

Frames

1e9

(b) When we don't stop gradients to w θ through Q-learning, task encodings can collapse to a single angle . We measure the average cosine similarity between pairs of tasks. If we pass gradients from Q-learning (CSFA - no stop grad), tasks collapse onto a single line.

C.2 More ablations

Here we present additional ablation results using the full curriculum in §5.1. We ablate the number of bins used to represent our probability mass function (Figure 16) and we ablate the depth of the residual network used to parameterize cumulants (Figure 17). 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 A

Find A (GPI)

PlaceAnearB

PlaceAnearB(GPI)

1.0

1.0

1.0

1.0

CSFA (ours)

0.8

CSFA-101bins

0.8

0.8

0.8

Rate

0.6

0.6

0.6

0.6

Success

CSFA (ours)

CSFA -101 bins

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0.0

0.0

0.0

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

Num learner frames1e9

Numlearnerframes1e9

Num learnerframes1e9

Num learnerframes1e9

Find A

Find A (GPI)

PlaceAnearB

Place AnearB (GPI)

1.0

1.0

1.0

1.0

CSFA (ours)

CSFA-4blocks

0.8

0.8

0.8

0.8

Rate

CSFA-1block

0.6

0.6

0.6

0.6

Success

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

1

2

3

4

5

6

7

8

0.0

0

1

2

3

4

5

6

7

8

0.0

0

1

2

3

4

5

6

7

8

0.0

0

1

2

3

4

5

6

7

8

Num learner frames1e9

Numlearnerframes1e9

Numlearnerframes1e9

Numlearnerframes

1e9

Find A & Put B near C

Put A near B & C near D & E near F

2.5

4

Reward

2.0

3

1.5

Average

2

1.0

Deep OK (ours)

1

0.5

Deep OK - reuse We

Deep OK - reuse Se

0.0

0

0

1

2

3

4

5

0

1

2

3

4

5

Frames

1e8

Frames

1e8

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 A

Find A (GPI)

PlaceAnearB

PlaceAnearB(GPI)

1.0

1.0

1.0

1.0

0.8

0.8

0.8

0.8

Rate

0.6

CSFA (ours)

0.6

0.6

0.6

Success

USFA (Borsa etal,2019)

0.4

MSFA（Carvalho etal,2023)

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0

1

2

3

4

5

6

7

8

0.0

0

1

2

3

4

5

6

7

8

0.0

0

1

2

3

4

5

6

7

8

0.0

Q

1

2

3

4

5

6

7

8

Numlearnerframes1e9

Numlearnerframes1e9

Numlearnerframes1e9

Numlearnerframes

s1e9

Find A

Find A (GPI)

1.0

1.0

0.8

0.8

Rate

0.6

0.6

Success

0.4

0.4

CSFA (ours)

0.2

0.2

CSFA-no IIwll(D2)

CSFA-no categorical

0.0

0.0

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

Num learner frames

1e9

Num learner frames

1e9

Find A

Find A (GPI)

PlaceAnearB

Place AnearB (GPI)

1.0

1.0

1.0

1.0

CSFA (ours)

0.8

CSFA-independent

0.8

0.8

0.8

Rate

CSFA- no stop grad (D1)

0.6

0.6

0.6

0.6

Success

0.4

0.4

0.4

0.4

0.2

0.2

0.2

0.2

0.0

0.0

0.0

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

0

1

2

3

4

5

6

7

8

0

1

2

E

4

5

6

7

8

Num learner frames1e9

Numlearnerframes

1e9

Numlearnerframes1e9

Numlearnerframes

1e9

E Implementation Details

The code is proprietary so we cannot release it but we included implementation details below to enable reproducibility .

Everything was implemented with the jax ecosystem [40]. 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.

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 §E.3 we detail transfer architectures: SFK, Impala (used for MTRL), and Distral. In §E.4 we describe learning for each method. Finally, in §E.5 we describe our hyperparameter search.

E.1 Compute resources

All experiments were run using Google Dragonfish TPU with 2 × 2 topology devices. Experiments for training tasks (§5.1) lasted about 1.5 days. Our simplified experiments in §C.1 lasted 5-6 hours. Transfer experiments (§5.2) lasted 2-3 hours. We ran a learner, actors, and evaluators on TPUs. All experiments were carried out using a distributed A3C setup [30] with discrete actions. Each experiment had 1 learner, 512 actors and 2 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 SFApproximator. We present a diagram of this general architecture in Figure22. Below we detail this architecture along with differences across methods.

SuccessorFeatureApproximators

πw(s,α)

SF-

w

approximator

be(s,a,w)

Φt+1

Linear (16)

+ L2 norm

St

8blocks,dim=256

ResNet

[256] x 4

MLP

LSTM

(512)

D

Task

LSTM

(64)

Linear

CumulantFunction

encoder

(y)0m

(256)

Φe(s,a)

Word

Concat

Embed

(64)

ResNet

Task

Previous

(16, 32, 32)

action

at-1

Ct

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 a projection of dimension 256.

Task encoder . All architectures embedded words in task descriptions using standard 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 and passed through a projection to the cumulant-dimension (16 across all our experiments). THe output was then divided by its L2 norm.

State function . All architectures used an LSTM of dimension 512. MSFA differed in that it used 4 LSTMs. We tried having their individual size be 512/4=128 and 1024/4=256 and found much

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 8 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 A*N. We converted back and forth from the two-hot representation using rlax library functions. In order to re-use this MLP across cumulants, we embed the cumulant dimension with a word embedding of dimension 256. The input to 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 A*C/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 Playroom resuts [31, 32, 10]. 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 methods. In §E.3.1, we detail implementation specific to Impala and Distral and in §E.3.2, we detail implementation specific to SFK. We present diagrams of these architectures in Figure23.

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 1 scalar output.

E.3.1 Impala & Distral

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 (1) the previous action and (2) the task. This is then passed through a projection of dimension 256.

State function . All architectures used an LSTM of dimension 512.

Policy

Impala uses an MLP with dimensions [1024] and output size A. The policy is parameterized with a softmax.

Distral learns two functions: one centroid function produces logits h ( a | s ) and a taskspecific function produces logits f i ( a | s ) . The task-policy is then parameterized as π ( a | s ) ∝ ( α ∗ h ( a | s ) + βf i ( a | s )) . Following Teh et al. [22], we set α = . 5 and β = 1 for acting. During learning we use a different β to encourage distillation. We detail this in §E.4.

Impala/Distral

Deep Option Keyboard

Coefficient-policy

(Only Distral)

7πw (s, a)

Value

Head

Value

Task

Centroid

head

(@1,...Qn)

head

Policy head

Policy head

SF-

MLP

MLP

MLP

MLP

MLP

(512)

(1024)

(512)

(1024)

(1024)

w-

approximator

4θ(s,a,w)

Linear (16)

+ L2 norm

S

St

LSTM

LSTM

LSTM

(512)

(512)

(512)

LSTM

Linear

(64)

Linear

(256)

(256)

Previous

coefficients

LSTM

LSTM

(64)

→w-

Concat

Embed

Word

Concat

(64)

Word

Previous

action

at-1

(16,32,32)

ResNet

Known

(64)

Previous

(16,32,32)

ResNet

Embed

Word

Legend

Embed

Tasks

action

(64)

(64)

at-1

↑

Frozen

Ct

+c

New Task

parameters

Task

E.3.2 SFK

SFK uses the CSFA architecture but freezes all parameters. It then learns a new 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 of 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: β 1 = 0 , β 2 = . 95 , ϵ = 6 e -6 .

E.4.1 CSFA, MSFA, & USFA

Following Mnih et al. [26], all architectures had a replay buffer that held up to 100,000 trajectories of length 30. This is an off-policy setting. Following Mnih et al. [26], all methods used target network parameters θ o that were updated smoothly using the online network parameters θ with a coefficents of . 9 . That is, at iteration i , θ o i = θ o i -1 ∗ . 1 + . 9 ∗ θ i . All methods used a learning rate of 3 e -4 and set the max gradient norm to . 5 . We used the following coefficients to balance loss terms: β Q = 1 e 3 , β ϕ = 1 e 4 . Since CSFA uses a categorical loss L cat ψ to learn SFs, which has a different magnitude from the scalar loss used to learn Q-values and , we found that we needed to make this smaller with β ψ = 8 e -3 whereas we used β ψ = 5 e 2 for scalar-based methods.

E.4.2 SFK, Impala, & Impala

We collected training data using a FIFO replay queue with 36 trajectories of length 128. Batches of 32 trajectories were sampled uniformly. This was an on-policy setting. SFK used a learning rate of 8 e -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 a learning rate of 3 e -4 for training and 8 e -5 for transfer. Distral used a learning rate of 8 e -6 for training and 3 e -5 for transfer. All methods set the max gradient norm to 40 .

For Distral, we set 1 /β = 3 e -3 . We set α = . 5 . Since Distral is a policy-gradient method, we added its policy distillation loss L dist to an Impala loss L impala . The full loss was L distral = L impala + α dis L dist . We used α dis = 5 . 0 for training tasks and α dis = . 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 [31] since they also study how Impala-based MTRL can transfer in Playroom (though they do so in a 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 [ 3 e -4 , 8 e -5 , 3 e -5 , 8 e -6 , 3 e -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 on top of Impala since it is a more advanced actor-critic algorithm that has been studied in playroom. As such, we kept hyper-parameters and tuned the following: (1) learning rates [3 e -4 , 8 e -5 , 3 e -5 , 8 e -6 , 3 e -6] , (2) α dis , [10 , 5 , 1 , . 1 , , 01 , . 001] , (3) Impala entropy coefficent [9 . 4 e -3 , 9 . 4 e -4 , 9 . 4 e -5 , 9 . 4 e -6 , 9 . 4 e -7] , (4) 1 /β [3 e -3 , 1 e -3 , 3 e -4] .

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, we set β ψ = β ϕ = 0 and searched over (1) β Q : [1, 1e1, 1e2, 1e3, 1e4] (2) max grad norm: [40 , 5 , . 5] . Once agents could learn Find tasks, we searched over (3) resblocks: [1, 4, 8] (4) β ϕ : [1, 1e1, 1e2, 1e3, 1e4] and (5) β ψ : [1, 1e1, 5e1 1e2, 5e2, 1e3]. We found the same hyper-parameters worked well for USFA and CSFA. For CSFA, we searched β ψ : [.8, .3, .08, .03, .008, .003, .0008] and bins [101, 301]. We began with β ψ 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 as a regularizer on the Q-value loss.

For SFK, we were able to re-use Impala HPs but searched over max grad norm [40 , 5 , . 5] , value hidden sizes [200, 512, 1024], policy hidden sizes [200, 512, 1024], and Impala entropy coefficent [9 . 4 e -3 , 9 . 4 e -4 , 9 . 4 e -5 , 9 . 4 e -6 , 9 . 4 e -7] .

F Environment Details

This work is a significant increase in complexity from prior successor feature literature [3, 4, 8, 6, 1]. The most complex task studied by prior work were 'go to' tasks in the DMLab environment by [8]. 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 [44] while Playroom has 46 actions. We have added this discussion to the main text.

F.1 Observation space

The agent experiences (1) observations as pixel-based RGB images that have dimensions 72 × 96 × 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 an 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 in the first instance that a GRAB action is not emitted. We describe these actions in Table2. An action-repeat of 4 was applied (i.e. each action was repeated 4 times).

F.3 Training Tasks

The agent experiences n κ = 72 training tasks T train = { κ 1 , . . . , κ n κ } composed of 'Find A' and 'Place A near B'. | A | = 18 and | B | = 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 3 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 be found. 'Put near' tasks are of the form 'Put a A near a B', where 'A' represent the object to be picked

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 1 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 B and C near D and E near F' would be 'place a boat near a book and place a bus near a cube block and place a keyboard near a sponge'. Objects were sampled with replacement, indicating that an 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 a reward of 1 and 'Place' subtasks provide a reward of 2 . Returning to our example, since it has 3 subtasks, it would provide a reward of 6 but only upon completion of all 3 subtasks.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\n- Proposed  a dynamic transfer query that enables sharing among task-encoding and SF-approximators;\n- Implemented this methodology in long-horizon tasks.\n\nWEAKNESSES:\n1. 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\n\nQUESTIONS:\nsee weakness and limitations\n\nLIMITATIONS:\n- 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?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n1.\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\nWEAKNESSES:\n1.  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\nQUESTIONS:\nIn 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？\n\nLIMITATIONS:\nSee the weakness and question sections.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\nI 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\nWEAKNESSES:\nAlthough 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\nQUESTIONS:\nQ1. 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\nLIMITATIONS:\nNone", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFirst, 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\nSTRENGTHS:\n# 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.\n\nWEAKNESSES:\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\nQUESTIONS:\nI 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\nLIMITATIONS:\nLimited 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}