Does Zero-Shot Reinforcement Learning Exist?
Authors: Ahmed Touati
Jérémy Rapin
Yann Ollivier
Venue: ICLR 2023 notable top 25% (2023)
Keywords: controllable agents
zero-shot RL
self-supervised representation learning
successor representation
offline RL

Abstract: A zero-shot RL agent is an agent that can solve any RL task in a given environment, instantly with no additional planning or learning, after an initial reward-free learning phase. This marks a shift from the reward-centric RL paradigm towards controllable agents that can follow arbitrary instructions in an environment. Current RL agents can solve families of related tasks at best, or require planning anew for each task. Strategies for approximate zero-shot RL have been suggested using successor features (SFs) (Borsa et al., 2018) or forward-backward (FB) representations (Touati & Ollivier, 2021), but testing has been limited. 
After clarifying the relationships between these schemes, we introduce improved losses and new SF models, and test the viability of zero-shot RL schemes systematically on tasks from the Unsupervised RL benchmark (Laskin et al., 2021). To disentangle universal representation learning from exploration, we work in an offline setting and repeat the tests on several existing replay buffers.
SFs appear to suffer from the choice of the elementary state features. SFs with Laplacian eigenfunctions do well, while SFs based on auto-encoders, inverse curiosity, transition models, low-rank transition matrix, contrastive learning, or diversity (APS),  perform unconsistently. In contrast, FB representations jointly learn the elementary and successor features from a single, principled criterion. They perform best and consistently across the board, reaching $85\%$ of supervised RL performance with a good replay buffer, in a zero-shot manner.

Published as a conference paper at ICLR 2023
DOES ZERO-SHOT REINFORCEMENT LEARNING EXIST?
Ahmed Touati, Jérémy Rapin & Yann Ollivier Meta Al Research, Paris, (atouati, jrapin, yol @meta.
ABSTRACT
A zero-shot RL agent is an agent that can solve any RL task in given environ- ment, instantly with no additional planning or learning, after initial reward-free learning phase. This marks a shift from the reward-centric RL paradigm towards "controllable" agents that can follow arbitrary instructions in environment. Cur- rent RL agents can solve families of related tasks at best, or require planning anew for each task. Strategies for approximate zero-shot RI. have been suggested using successor features (SFs) (Borsa et al., 2018) or representations (Touati & Ollivier, 2021), but testing has been limited.
(FB)
After clarifying the relationships between these schemes, we introduce improved losses and new SF models, and test the viability of zero-shot RL schemes sys- tematically on tasks from the Unsupervised RI. benchmark (Laskin er al., 2021). ) To disentangle universal representation learning from exploration, we work in an offline setting and repeat the tests on several existing replay buffers.
SFs appear to suffer from the choice of the elementary state features. SFs with Laplacian eigenfunctions do well, while SFs based on auto-encoders. inverse curiosity, transition models, low-rank transition matrix, contrastive learning, or diversity (APS), perform unconsistently. In contrast, FB representations jointly learn the elementary and successor features from single, principled criterion. They perform best and consistently across the board, reaching 85% of supervised RL performance with good replay buffer, in a zero-shot manner.
INTRODUCTION
For breadth of applications, reinforcement learning (RL) lags behind other fields of machine learning, such as vision or natural language processing, which have effectively adapted to a wide range of tasks, often in almost zero-shot manner, using pretraining on large, unlabelled datasets (Brown et al., 2020), The RL paradigm itself may be in part to blame: RL agents are usually trained for only one reward function or small family of related rewards. Instead, we would like to train "controllable" agents that can be given a description of any task (reward function) in their environment, and then immediately know what do, reacting instantly to such commands "fetch this object while avoiding that area".
The promise of zero-shot RI. is to train without rewards or tasks, yet immediately perform well on any reward function given at test time, with extra training, planning, or finetuning, and only minimal amount of extra computation to process a task description (Section 2 gives the precise definition we use for zero-shot RL). How far away are such zero-shot agents? In the RL paradigm, a new task (reward function) means re- re-training the agent from scratch, and providing many reward samples. Model-based RL trains a reward- -free, eek-indeeendent world model, but still requires heavy planning when a new reward function is specified (e.g, Chua et al., 2018; Moerland et al., 2020). Model-free RL reward-centric from start, and produces specialized agents. Multi-task agents generalize within family of related tasks only. Reward-free, unsupervised skill pre-training (c.g, Eysenbach et al., 2018) still requires substantial downstream task adaptation, such as training a hicrarchical controller.
Is zero-shot RL possible? If one ignores practicality, zero-shot RL is easy: make list of all possible rewards up to precision E, then pre-learn all the associated optimal policies. Scalable zero-shot RL must somehow exploit the relationships between policies for all tasks. Learning go from to not independent from going from a to b and to c, and this produces rich, exploitable algebraic relationships (Blier et al., 2021; Schaul et al., 2015).

---

Published as a conference paper at ICLR 2023
100
SFs with features from:
Random
80
Auto encoder
lnverse curiosity module
Latent transition model
660
Transition model
Laplacian eigenfenctions
Zero-shot
Low rank ba transition mabii
40
Contrastive learning
Low-rank successor me asure APS
26
Forw Horward-backwwrct
representation
Figure 1: Zero-shot scores of ten SF methods and FB, percentage of the supervised score of offline TD3 trained on the same replay buffer, averaged on some tasks, environments and replay buffers from the Unsupervised RL and ExORL benchmarks (Laskin al., 2022; Yarats et al., 2021). FB and SFs with Laplacian eigenfunctions achieve zero-shot scores approaching supervised RL.
Suggested strategies for generic zero-shot RL so far have used successor representations (Dayan, 1993), under two forms: successor features (SFs) (Barreto et al., 2017) as in (Borsa et al., 2018; Hansen et al., 2019; Liu & Abbeel, 2021);; and Ollivier, 2021). Both SFs and FB lie in between model-free and model-based RL, by predicting features of future states, summarizing long-term state-state relationships. Like model-based approaches, they decouple the dynamics of the environment from the reward function. Contrary to world models, they require neither planning at test time nor a generative model of states or trajectories.
(FB) representations (Touati &
Yet SFs heavily depend on a choice of basic state features. To get full zero-shot RL algorithm, representation learning method must provide those. While SFs have been successively applied to transfer between tasks, most of the time, the basic features were handcrafted or learned using prior task class knowledge. Meanwhile, FB a standalone method with no task prior and good theoretical backing, but testing has been limited to goal-reaching in few environments. Here:
We systematically assess SFs and FB for zero-shot RI., including many new models of SF basic features, and improved FB loss functions. We use 13 tasks from the Unsupervised RL benchmark (Laskin et al., 2021), repeated on several EXORL training replay buffers (Yarats et al., 2021) to assess robustness to the exploration method.
We systematically study the influence of basic features for SFs, by testing SFs on features from ten RL representation learning methods. such as latent next state prediction, inverse curiosity module, contrastive learning, diversity, various spectral decompositions...
We expose mathematical links between SFs, FB, and other representations in RI. We discuss the implicit assumptions and limitations behind zero-shot RI. approaches.
2 PROBLEM AND NOTATION; DEFINING ZERO-SHOT RL
Let M (S, A, P,7) ) be a reward-free Markov decision process (MDP) with state space S, action space A, transition probabilities P(s'|s, a) from state to s' given action a, and discount factor 0 <<> < (Sutton & Barto, 2018). If S and A are finite, P(s'|s ss a) can be viewed as a stochastic matrix Peas' E {= in general, for each (s, a) E S x A, P(ds'Is, a) is probability measure s' E'E The notation P(ds' a) covers all cases. Given (5o. 0g) S A and a policy :: S - Prob(A), we denote Pr(-Iso, ag, x) and E[-|so, ao, x] the probabilities and expectations under state-action sequences (St, ac)20 starting at (so, a0) and following policy π in the environment, defined by sampling ~ P(ds, St-1. 44-1) and 0 ~ (dda]s:). We define P= da'ls, a) P(ds'|s (da'd'))s ss)) and P= [s) [s):
P(ds'[s, ()ss))(daa)))) (s) the state-action transition probabilities and state transition probabilities induced by Given reward function S + R. the Q-function of π for is Q, (so, un) := E. 7¹ B[[((+++) |s0, do, T]. For simplicity, we assume the reward r depends only on the next state s++1 instead on the full triplet (st, Ge $t+1), but this is not essential.
We focus on offline unsupervised RL, where the agent cannot interact with the environment. The agent only has access to a static dataset of logged reward-free transitions in the environment, D = {(s,, a,, s/) e)ez with s P(ds'Is;, a:). These can come from any exploration method or methods.
The offline setting disentangles the effects of the exploration method and representation and policy learning: test each zero-shot method on several training datasets from several exploration methods.
2

---

Published as a conference paper at ICLR 2023
We denote by p(ds) and p(ds. da) the (unknown) marginal distribution of states and state- actions in the dataset D. We use both EewD and Egup! for expectations under the training distribution.
Zero-shot RL: problem statement. The goal of zero-shot RL is to compute compact represen- tation & of the environment by observing samples of reward- frwerd-free transitions ($t, S(+1) in this environment. Once reward function is specified later, the agent must use & immediately produce a good policy, via only elementary computations without any further planning or learning. Ideally, for any downstream task, the performance of the returned policy should close to the performance of supervised RI. baseline trained on the same dataset labeled with the rewards for that task.
Reward functions may be specified at test time either as a relatively small set of reward samples (s), rt), or explicit function + r(s) (such as 1 at a known goal state and 0 elsewhere). The method will be few-shot, zero-planning in the first case, and truly zero-shot in the second case.
3 RELATED WORK
Zero-shot RL requires unsupervised learning and the absence of any planning or fine-tuning at test time. The proposed strategies for zero-shot RL discussed in Section ultimately derive from successor representations (Dayan, 1993) in finite spaces. In continuous spaces, starting with a finite number of features ye, successor features can be used produce policies within family of tasks directly related to yp (Barreto et al., 2017; Borsa et al., 2018; Zhang et al., 2017; Grimm et al., 2019), often using hand-crafted y or learning op that best linearize training rewards. VISR (Hansen et al., 2019) and its successor APS (Liu & Abbeel, 2021) use SFs with & automatically built online via diversity criteria (Eysenbach et al., 2018; Gregor al., 2016). We include APS among our baselines, as well as many new criteria to build automatically.
Successor measures (Blier et al., 2021) avoid the need for yp by directly learning models of the distribution of future states: doing this for various policies yields a candidate zero-shot RI. method, heward-backwart representations (Touati & Ollivier, 2021), which has been tested for goal-reaching in few environments with discrete actions. FB uses low-rank model of long-term state-state relationships reminiscent of the state-goal factorization from Schaul et al. (2015).
Model-based RL (surveyed in Moerland et al. (2020)) misses the zero-planning requirement of zero-shot RL. Still, learned models of the transitions between states can be used jointly with SFs to provide zero-shot methods (Trans, Latent, and LRA-P methods below).
Goal-oriented and multitask RL has a long history (e.g, Foster Dayan, 2002; Sutton et al., 2011; da Silva et al., 2012; Schaul et al., 2015; Andrychowiez et al., 2017). A parametric family of tasks must be defined in advance (e.g., reaching arbitrary goal states). New rewards cannot be set a posteriori: for example, goal-state-oriented method cannot handle dense rewards. Zero-shot task transfer methods learn on tasks and can transfer to related tasks only (e.g, Oh et al., 2017; Sohn al., 2018); this can be used, e.g., for sim-to-real transfer (Genc et al., 2020) or slight environment changes, which is not covered here. Instead, we aim at not having any predefined family of tasks.
Unsupervised skill and option discovery methods, based for instance on diversity (Eysenbach et al., 2018; Gregor et al., 2016) or cigenoptions (Machado et al., 2017) can learn a variety of behaviors without rewards. Downstream tasks require learning hierarchical controller to combine the right skills or options for each task. Directly using unmodified skills has limited performance without heavy finetuning (Eysenbach et al., 2018). Still, these methods can speed up downstream learning.
The unsupervised aspect of some of these methods (including DIAYN and APS) has been disputed, because training still used end-of- -otaj---rajectory signals, which are directly correlated to the downstream task in some common environments: without this signal, results drop sharply (Laskin et al., 2022).
4 SUCCESSOR REPRESENTATIONS AND ZERO-SHOT RL
For finite MDP, the successor representation (SR) (Dayan, 1993) M = ao) of state-action pair (Sp, ag) under a policy n, is defined as the discounted sum of future occurrences of each state:
Md MM ao, 8) = -E |E20 7' (e (50, an), π. Vs E S.
(1)
3

---

Published as a conference paper at ICLR 2023
In matrix form, SRs can be written M = P P -yPx) where Pa is the state transition probability. M satisfies the matrix Bellman equation M = P 7P,,,M
Importantly, SRs disentangle the dynamics of the MDP and the reward function: for any reward r and policy т, the Q-function can be expressed linearly as Q M"r.
Successor features and successor measures. Successor featurrss (SFs) (Barreto et al., 2017) extend SR to continous MDPs by first assuming are given basic feature map p S - Rd that embeds states into d-dimensional space, and defining the expected discounted sum of future state features:
"hTh (so. ao) E [E120 T (((((5)1))) so, Co, π
(2)
SFs have been introduced to make SRs compatible with function approximation. For a finite MDP, the original definition (1) is recovered by letting y be one-hot state encoding into g
Alternatively, successor measures (SMs) (Blier et al., 2021) extend SRs to continuous spaces by treating the distribution of future visited states as a measure M over the state space S.
(So. ap. X) == Σis0 7 Pr (8e+1 E X Sp, ao. T)
VX
(3)
SFs and SMs are related: by construction, o (sp, ao) = J3 M (s0, Co, ds') d
Zero-shot RL from successor features and representations provide a generic framework for zero-shot RL. by learning to represent the relationship between reward functiions and Q-functions, as encoded in A
representations. Successor
Given basic feature map yo: S + R to be learned via another criterion, universal SFs (Borsa et al., 2018) learn the successor features of a particular family of policies π, for Rdd
(8 Cu, 2) = E Ei20 1 (($++1)) (So. (o). T2
T=(s) == arg max, (1s 2)_ (4)
Once reward function is revealed, we use few reward samples explicit knowledge of the function r to perform a linear regression of r onto the features ip. Namely, we estimate 2, == arg min, [r<s](r(s) -s)
" =) Epp |- Ey Then we return the policy u.,. This policy is guaranteed to be optimal for all rewards in the linear span of the features yp:
Theorem (Borsa al. (2018)). Assume that (4) holds. Assume there exists a weight w E R.d such that r(s) = ( Vs E S. Then 2, w, and Ra is the optimal policy for reward r.
(FB) representations (Touati & Ollivier, 2021) apply a similar idea to a finite-rank model of successor measures. They look for representations F: S x A x Rd - IR and B: S B such that the long-term transition probabilities = in (3) decompose as
R (so, ag ds') F(so. oo, 2) B(s') p(ds'), (s) == arg max, F(s, a, T)T²
(5)
In a finite space, the first equation rewrites as the matrix decomposition M = F. B diag(p)-
Once reward function is revealed, estimate 2, Eg-p [ B(s) from few reward samples or from explicit knowledge of the function r (e.g. 2, = B(s) to reach s). Then we return the policy T2, If the approximation (5) holds, this policy guaranteed be optimal for any reward function: Theorem (Touati & Ollivier (2021)). Assume that (5) holds. Then for any reward function function r, the policy n., optimal for r, with optimal Q-function Q F(s, a, 2,) 2p-
For completeness, we sketch the proofs of Theorems 1-2 in Appendix A. Importantly, both theorems are compatible with approximation: approximate solutions provide approximately optimal policies.
Connections between SFs and FB. A first difference between SFs and FB is that SFs must be provided with basic features s. The best yp is such that the reward functions of the downstream tasks linear in yp. But for unsupervised training without prior task knowledge, is needed to learn yo. We test a series of such criteria below. In contrast, FB uses a single criterion, avoiding the need for state featurization by learning model of state occupancy.
exteral criterion
Second, SFs only cover rewards in the linear span of (p, while FB apparently covers any reward. But this difference is not stark as it looks: exactly solving the FB equation (5) in continuous spaces
4

---

Published as a conference paper at ICLR 2023
requires d oo, and for finite d, the policies will only be optimal for rewards in the linear span of B (Touati & Ollivier, 2021). Thus, in both cases, policies are exactly optimal only for d-dimensional family of rewards. Still, FB can use arbitrary large d without any additional input or criterion.
FB representations are related to successor features: the FB definition (5) implies that eb(s, a, 2) := F(s, a, z) are the successor features of (p() (EE, -BBB)-B((S This follows from multiplying (3) and (5) by B" (E,, BB") on the right, and integrating over s' p. Thus, posteriori, FB can be used to produce both &p and ve in SF, although training is different.
This connection between FB and SF is one-directional: (5) is a stronger condition. In particular = B = 0 is not a solution: contrary to & yp 0 in (4), there is no collapse. No additional criterion train y is required: F and B trained jointly provide the best rank-d approximation of the successor measures M*. This summarizes an environment by selecting the features that best describe the relationship Q rr between rewards and Q-functions.
5 ALGORIHTMM FOR SUCCESSOR FEATURES AND FB REPRESENTATIONS
We now describe more precisely the algorithms used in our experiments. The losses used to train e in SFs, and F, B in FB, are described Sections 5.1 and 5.2 respectively.
To obtain a full zero-shot RL algorithm, SF' must specify the basic features y. Any representation learning method can be used for yo. We new representations for RL: random features as baseline, autoencoders, next state and latent next state transition models, inverse curiosity module, the diversity criterion of APS, contrastive learning, and finally, several spectral decompositions of the transition matrix or its associated Laplacian.
possible choices (Section 5.3) based on existing or
Both SFs and FB define policies as an argmax of (ss, a, T)T 2 or F(s, a, 2)_ over actions a. With continuous actions, the argmax cannot be computed exactly. We train an auxiliary policy network T (s, z) to approximate this argmax, using the same standard method for SFs and FB (Appendix G.4).
5.1 LEARNING THE SUCCESSOR FEATURES Vi
The successor features e satisfy the -aluul Bellman equation T = Pip 7Pxt the collection of ordinary Bellman equations for each component of y. The P in front of ip comes from using (s(.-1) not (ss.) in (2). Therefore, we can train (iss, a, 2) for each by minimizing the Bellman residuals [[,,, Ot; 2) - (8(+++))==+)))))))))(()))) 7V($ T (8++1), =)] where v is a non-trainable target ver- sion of yo as parametric Q-learning. This requires sampling transition (st, a, 8(+1) from the dataset and choosing 2. We sample random values of 2 as described in Appendix G.3.
This is the loss used in Borsa et al. (2018). But this can be improved, since we do not use the full vector (s,, s, a, z): only 1b(8, 22) 2 is needed for the policies. Therefore, as in Liu & Abbeel (2021), instead of the vector-valued Bellman residual above, we just use
(6) for each z. This trains 2) as the Q-function of reward ip 2, the only case needed, while training the full vector 2)) amounts training the Q-functions of each policy n. for all rewards p for all E Rd including + 2. We have found this improves performance.
[(t)) = ((s (1b( 01, x)_== +($+41)*= 2 ( (8)+1))))
5.2 LEARNING FB REPRESENTATIONS:: THE FB TRAINING LOss
The successor measure M satisfies Bellman- like equation M P + TP,M matrices in the finite case and as measures in the general case (Blier et al., 2021). We can learn FB by iteratively minimizing the Bellman residual on the parametric model M = F" Bp. Using suitable norm ||: for the Bellman residual (Appendix B) leads to loss expressed expectations from the dataset:
C(F, B)
|| Bp (P yPx, Bp)
E Bles,an,etiiim (f 2) =)TB(s)) d'p
B(s') ( (e(84))) B(s))² (()
-2 [F(s), 2) B($41)) + Const
(8)
5

---

Published as a conference paper at ICLR 2023
where the constant term does not depend on F and B, and where usual F₂ and B are non non-trainable target versions of F and B whose parameters are updated with a slow-moving average of those of F and B. Appendix B quickly derives this loss, with pseudocode in Appendix L. Contrary to Touati & Ollivier (2021), the last term involves ( instead of (*+), because we use 8t+1 instead of st for the successor measures (3). We sample random values of as described in Appendix G.3.
We include an auxiliary loss (Appendix B) to normalize the covariance of B, B, BB = Id, as in Touati & Ollivier (2021) (otherwise one can, e.g., scale F up and B down since only F B is fixed).
As with SFs above, only F(,, =)== is needed for the policies, while the loss above the vector F amounts to training =)( 2' for all pairs z'). The full loss is needed for joint training of F and B. But we include auxiliary loss (F) to foeus training on the diagonal 2' 2. This is obtained by multiplying the Bellman gap in C by B (BpB") ) on the right, to make F(,,z)) aappear:
(P)))= == (s [F(s, a, =)= -- B($ (Ep BBT)BB---- y-1= ( T. (s 2) 2)
(9)
This trains F(.,2) x)===T= 2 as the Q-function for reward B (E, B Though C = 0 implies L
0,
adding L' reduces the error on the part used for policies. This departs from Touati & Ollivier (2021).
5.3 LEARNING BASIC FEATURES &p FOR SUCCESSOR FEATURES
SFs must be provided with basic state features sp Any representation learning method can be used to supply yo. We focus on prominent RL representation learning baselines, and on those used in previous zero-shot RL candidates such APS. We now describe the precise learning objective for each.
Random Features (Rand), We use a non-trainable randomly initialized network as features.
Autoencoder (AEnc). We learn decoder f:: RRd - S to recover the state from its representation yo:
min E ((f(((())) s) SiP
(10)
Inverse Curiosity Module (ICM) aims at extracting the controllable aspects of the environ- ment (Pathak al., 2017). The idea to train an inverse dynamics model g: Id x Hd + A to predict the action used for a transition between two consecutive states. We use the loss
min E (1 ((())))) atll²). 9:8
(11)
Transition model (Trans). This is a one-step forward dynamic model f: Rd x A - S that predicts the next state from the current state representation:
min (aleaannnnnnppU((((((a Sip
- 12+1)
(12)
Latent transition model (Latent). This is similar to the transition model but instead of predicting the next state, it predicts its representation:
min E (s Df a.) ()
(13)
8:50
clear failure case of this loss is when all states are mapped to the same representation. To avoid this collapse, we compute (p($ using a trainablee version of yp, with parameters corresponding to a slowly moving average of the parameters of yp, similarly to BYOL (Grill et al., 2020).
Diversity methods (APS). VISR (Hansen et al., 2019) and its successor APS (Liu & Abbeel, 2021) tackle zero-shot RL using SFs with features y built online from diversity criterion. This criterion maximizes the mutual information between a policy parameter and the features of the states visited by a policy using that parameter (Eysenbach et al., 2018; Gregor et al., 2016). VISR and APS use, respectively, a variational or nearest-neighhor estimator for the mutual information. We directly use the code provided for APS, and refer to (Liu & Abbeel, 2021) for the details. Contrary to other methods, APS is not offline: it needs to be trained on its own replay buffer.
Laplacian Eigenfunctions (Lap). Wu et al. (2018) consider the symmetrized MDP graph Laplacian induced by exploratory policy π, defined L = Id - (Px diag(p) -1 + diag(p) (p))) (Px)T) TT. They

---

Published as a conference paper at ICLR 2023
propose to learn the eigenfunctions of L via the spectral graph drawing objective (Koren, 2003):
min (1 [1ıs(s))) +0(+++1)l? + 1 gnD [(p(s)) ((s)))² - 1| - lkp(s') (14)
where the second term is an orthonormality regularization to ensure that [p)))((()))(())] Re Id, and > 0 is the regularization weight. This is implicitly contrastive, pushing features of st and 81+1 closer while keeping features apart overall. Such eigenfunctions have long been argued to play key role in RL (Mahadevan & Maggioni, 2007;
Machado et al., 2017).
Low-Rank Rank Approximation of P (LRA-P): we learn features by estimating low-rank model of the transition probability densities: P(ds'|e (sdd'ls,,i a) ss x(s, a) (u(")) p(ds'). Knowing p is not needed: the corresponding loss on x" J - P/p readily expressed as expectations over the dataset,
( (x(e))a) u(s') P(ds'|se:.u) p(ds')
min E ((naa
(15)
X-M
d'rup =Beenasmollx(sst "r) (()1)))) - 2 Beeeoee ar) p(se+1)) + Const (16)
s'rup
We normalize Ep [uph)] xe Id with the same loss used for B. Then we use SFs with (p := Ju. If the model P = p is exact, this provides exact optimal policies for any reward (Appendix E, Thm. 3).
This loss is implicitly contrastive: it compares samples 9 to independent samples s' from p. It is an asymmetric extension of the Laplacian loss (14) (Appendix F). The loss (16) is also a special case of the FB loss (8) by setting 7 O. omitting 2, and substituting (x (u) for (F. Indeed, FB lears a finite-rank model of P(Id y P₂) -1, which equals P when & = O.
A related loss is introduced in Ren et al. (2022), but involves second unspecified, arbitrary probability distribution p. which must cover the whole state space and whose analytic expression must be known. It is unclear how to set suitable p in general.
Contrastive Learning (CL) methods learn representations by pushing positive pairs (similar states) closer together while keeping negative pairs apart. Here, two states are considered similar if they lie close on the same trajectory. We use a SimCLRR-lik objective (Chen et al., 2020):
exp(ccssime(((s)))ooxx)))
min X-P
Bk-Geen(1- (5., )oD
log
(17)
Eg exp(cosine(((si),p) ((s)))
where S1+ the state encountered at step k+ along the subtrajectory that starts at 8₁, where k is sam-
" pled from a geometric distribution of parameter (1 Tcu), and cos Mosine(u,, v) = T Yu, E
Rdd is the cosine similarity function. Here "ccz. E [0; 1) a parameter not necessarily set to the MDP's discount factor y. CL requires dataset made of full trajectories instead of isolated transitions.
CL is tightly related to the spectral decomposition of the successor measure E, n Pt+1, where æ is the behavior policy generating the dataset trajectories. Precisely, assuming that x and şp are centered with unit norm, and expanding the log and exp at second order, the loss (17) becomes
(17) ~ H B.EDD ((x()) (p()))))) 5' S'ND
[x(*:) (((8))Tp(((++)))
(18)
(compare (16)). Now, the law of 8 given 8t is given by the stochastic matrix (1 (1-7eu) E, Th pe+1, the rescaled successor measure of π. Then one finds that (18) is minimized when x and ip provide the singular value decomposition of this matrix in (p) norm (Appendix C). Formal links between contrastive learning and spectral methods can be found in Tian (2022); Balestriero & LeCun (2022).
Low-Rank Approximation of SR (LRA-SR). The CL method implicitly factorizes the successor measure of the exploration policy in Monte Carlo fashion by sampling pairs (st, st+*) on the same trajectory. This may suffer from high variance. To mitigate this, we propose to factorize this successor measure by temporal difference learning instead of Monte Carlo. This is achieved with an FB-like loss (8) except we drop the policies n and learn successor measures for the exploration policy only:
min B(enery)) (x((x) ((s') - xx(s)+1) =(s))²] X-P
2 ((onAe))) [x(*:] ((((+))))]]]] (19)
with X and S target versions of x and yp. We normalize to E [ Ro Id with the same loss as for B. Of all SF variants tested, this is the closest FB.
7

---

Published as a conference paper at ICLR 2023
102
waleer yord
watler mal
welter pu
valler
Chosse wale
Offine TD3
Adnc
Figure 2: Zero-shot scores for each task, with supervised online and offline TD3 as toplines. Average over 3 replay buffers and 10 random seeds.
EXPERIMENTAL RESULTS ON BENCHMARKS
Each of the 11 methods (FB and 10 SF-based models) has been tested on 13 tasks in 4 environments from the Unsupervised RL and EXORL benchmarks (Laskin et al., 2021; Yarats et al., 2021): Maze (reach 20 goals), Walker (stand, walk, run, flip), Cheetah (walk, run, walk backwards, run backwards), and Quadruped (stand, walk, run, jump); see Appendix G.I. Each task and method was repeated for 3 choices of replay buffer from FXORL: RND, APS, and Proto (except for the APS method, which can only train on the APS buffer). Each of these 403 settings was repeated with 10 random seeds.
The full setup is described in Appendix G. Representation dimension is d = 50, except for Maze (d = 100). After model training, tasks are revealed by 10,000 reward samples as in (Liu Abbeel, 2021), except for Maze, where known goal is presented and used set 2, directly. The code can be found at https:://// / ggittbb.bb.....g//g////.g....h.. com/ facebookre www.kch..eeeccceee.e....a. contt rol roll lable _aall__g_ntt..
As toplines, we online TD3 (with task rewards, and free environment interactions not restricted to a replay buffer), and offline TD3 (restricted to each replay buffer labelled with task rewards). Offline TD3 gives an idea of the best achievable performance given the training data in buffer.
In Fig. we plot the performance of each method for each task in each environment, averaged over the three replay buffers and random seeds. Appendix H contains the full results and more plots.
Compared to offline TD3 as a reference, on the Maze tasks, Lap, LRA-P, and FB perform well. On the Walker tasks, ICM, Trans, Lap, LRA-SR, and FB perform well. On the Cheetah tasks, ICM, Trans, Lap, LRA-SR, and FB perform well. On the Quadruped tasks, many methods perform well, including, surprisingly, random features. Appendix plots some of the learned features on Maze.
On Maze, none of the encoder-based losses learn good policies, contrary to FB and spectral SF methods. We believe this is because (x. y) is already a good representation of the 2D state 8, so the encoders do nothing. Yet SFs on (z, y) cannot solve the task (SFs on rewards of the form az + by don't recover goal-oriented tasks): planning with SFs requires specific representations.
Fig. I reports aggregated scores over all tasks. To average across tasks, we normalize scores: for each task and replay buffer, performance is expressed as a percentage of the performance of offline TD3 on the same replay buffer, a natural supervised topline given the data. These normalized scores are averaged over all tasks in each environment, then over environments to yield the scores in Fig. 1. The variations over environments, replay buffers and random seeds are reported in Appendix H.
These results are broadly consistent over replay buffers. Buffer-specific results are reported in Appendices H.3-H.4. Sometimes replay buffer is restrictive, attested by poor offline TD3 performance, starkly so for the Proto buffer on all Quadruped tasks. APS does not work well as a zero-shot RL method, but it does work well replay buffers have close results.
exploration method: average, the APS and RND
FB and Lap the only methods that perform consistently well, both over tasks and over replay buffers. Averaged on all tasks and buffers, FB reaches 81 of supervised offline TD3 performance, and 85% on the RND buffer. The second-best method is Lap with 74% (78% the Proto buffer).
8

---

Published as a conference paper at ICLR 2023
7 DISCUSSION AND LIMITATIONS
Can a few features solve many reward functions? Why are some features better? The choice of features is critical. Consider goal-reaching tasks: the obvious way learn to reach arbitrary goal states via SFs is to define one feature per possible goal state (one-hot encoding). This requires [S] features, and does not scale to continuous spaces. But much better features exist. For instance, with an size-n cycle S = {0,
n - }} with actions that move left and right modulo n, then two features suffice instead of n: SFs with (p(s) = (co8(2158/ /n, sin(2rs//
provides exact optimal policies to 1]4 just 2d features (a sine and cosine in each direction) are sufficient reach any of the nd goal states via SFs.
reach any arbitrary state. On d-dimensional grid S {0, (
Goal-reaching is only subset of possible RL tasks, but this clearly shows that some features are better than others. The sine and cosine are the main eigenfunctions of the graph Laplacian on the grid: such features have long been argued to play a special role in RL (Mahadevan & Maggioni, 2007; Machado et al., 2017). FB-based methods are theoretically known to learn such eigenfunctions (Blier et al., 2021). Yet precise theoretical link to downstream performance is still lacking.
Are these finite-rank models reasonable? FB crucially relies on a finite-rank model of Σ of while some SF variants above rely on finite-rank models of P the corresponding Laplacian. It turns out such approximations are very different for P or for Ee'
Unfortunately, despite the popularity of low-rank P assumptions in the theoretical literature, P. is always close to Id in situations where St+1 is close to St, such as any continuous-time physical system (Appendix D). Any low-rank model of P5 will be poor; actually P= is better modeled as Id - low-rank, thus approximating the Laplacian. On the other hand, Pf with large t gets close to rank one under weak assumptions (ergodicity), as P converges to an equilibrium distribution when t + 00. Thus the spectrum of the successor measures Σn' P is usually more spread out (details and examples in Appendix D), and a low-rank model makes sense. The eigenvalues of Pa are close to 1 and there is little signal to differentiate between eigenvectors, but differences become clearer over time on p This may explain the better performance of FB and LRA-SR compared to LRA-P.
First, these experiments are still small-scale and performance is not perfect, so there is space for improvement. Also, all the environments tested here were deterministic: all algorithms still make sense in stochastic environments, but experimental conclusions may differ.
Limitations.
Second, even though these methods can be coupled with any exploration technique. this will obviously not cover tasks too different from the actions in the replay buffer, as with any offline RL method.
These zero-shot RL algorithms learn to summarize the long-term future for wide range of policies (though without synthesizing trajectories). This is a lot: in contrast, world models only learn policy- independent one-step transitions. So the question remains of how far this can scale. A priori, there is no restriction on the inputs (e.g., images, state history...). Still, for large problems, some form of prior seems unavoidable. For SFs, priors can be integrated in ş2, but rewards must be linear in yo. For FB, priors can be integrated in B's input. Touati & Ollivier (2021) use FB with pixel-based inputs for F, but only the agent's (x. y) position for B's input: this recovers all rewards that are functions of (xx y) (linear or not). Breaking the symmetry of and B reduces the strain on the model by restricting predictions to fewer variables, such as agent's future state instead of the full environment.
8 CONCLUSIONS
Zero-shot RL methods avoid test-time planning by summarizing long-term state-state relationships for particular policies. We systematically tested models of successor features on zero-shot tasks. We also uncovered algebraic links between SFs, FB, contrastive learning, and spectral methods. Overall, SFs suffer from their dependency to basic feature construction, with only Laplacian eigenfunctions working reliably. Notably, planning with SFs requires specific features: SFs can fail with generic encoder-type feature learning, even if the representations
representations and many new
learned representation of states reasonable (such as (a. y)). were best across the board, and provide reasonable zero-shot RI. performance.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
The authors would like to thank Olivier Delalleau, Armand Joulin, Alessandro Lazaric, Sergey Levine, Matteo Pirotta, Andrea Tirinzoni, and the anonymous reviewers for helpful comments and questions on the research and manuscript.
REFERENCES
Marcin Andrychowiez, Dwight Crow, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. In NIPS, 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. 107.v:160.006450, 2016.
arXiv preprint
Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning recover global and local spectral embedding methods. arXiv preprint 220.:::2005./508 2022.
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, David Silver, and Hado P van Hasselt. Successor features for transfer in reinforcement learning. In NIPS, 2017.
Léonard Blier, Corentin Tallec, and Yann Ollivier. Learning successor states and goal-dependent values: A mathematical viewpoint. arXiv preprint arXiv:2101.07123, 2021.
Diana Borsa, André Barreto, John Quan, Daniel Mankowitz, Rémi Munos, Hado van Hasselt, David Silver, and Tom Schaul. Universal successor features approximators. arXiv preprint arXiv:1812.07626, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp.
1597-1607. PMLR, 2020.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Bruno Castro da Silva, George Konidaris, and Andrew Barto. Learning parameterized skills. In ICML, 2012.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613-624, 1993.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In International Conference on Learning Representations, 2018.
David Foster and Peter Dayan. Structure in the space value functions. Machine Learning, 49(2): 325-346, 2002.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In International conference machine learning, pp. 1587-1596. PMLR, 2018.
Sahika Genc, Sunil Mallya, Sravan Bodapati, Tao Sun, and Yunzhe Tao. Zero-shot reinforcement learning with deep attention convolutional neural networks. arXiv preprint arXiv:2001.0000 006005, 2020.
Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint 161 .07507, 2016.
10

---

Published as a conference paper at ICLR 2023
Jean-Bastiens Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a approach to self-supervised learning. Advances in neural information processing systems, 33:21271-21284. 2020.
Christopher Grimm, Irina Higgins, Andre Barreto, Denis Teplyashin, Markus Wulfmeier, Tim Her- tweck, Raia Hadsell, and Satinder Singh. Disentangled cumulants help successor representations transfer to new tasks. arXiv preprint arXiv: 2019.
Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, Farley, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. arXiv preprint 10XX::1906.000000 2019.
Ychuda Koren. On spectral graph drawing. In International Computing and Combinatorics Confer- ence, pp. 496-508. Springer, 2003.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. Urlb: Unsupervised reinforcement learning benchmark. In Thirty-fifth Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
Michael Laskin, Hao Liu, Xue Bin Peng, Denis Yarats, Aravind Rajeswaran, and Pieter Abbeel. Cic: Contrastive intrinsic control for unsupervised skill discovery. arXiv preprint arXtv:2202.001111 2022.
David A Levin, Yuval Peres, and Elisabeth L Wilmer. Markov chains and mixing times. American Mathematical Soc., Providence, 2009.
Hao Liu and Pieter Abbeel. APS: Active pretraining with successor features. In International Conference on Machine Learning, pp. 6736-6747. PMLR, 2021.
Clare Lyle, Mark Rowland, and Will Dabney. Understanding and preventing capacity loss in reinforcement learning. arXiv preprint arXiv:2204.09560, 2022.
Chen Ma, Dylan R Ashley, Junfeng Wen, and Yoshua Bengio. Universal successor features for transfer reinforcement learning. arXiv preprint :XXx:2001.04022. 2020.
Marlos Machado, Mare G Bellemare, and Michael Bowling. A Laplacian framework for option discovery in reinforcement learning. In International Conference on Machine Learning, pp. 2295 0295-2304. PMLR, 2017.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A Laplacian framework for learning representation and control in Markov decision processes. Journal of Machine Learning Research, 8(10), 2007.
Thomas M Moerland, Joost Broekens, and Catholijn M Jonker. Model-based reinforcement learning: A survey. arXiv preprint arXiv:2006.1677 2020.
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, pp. 2661-2670. PMLR, 2017.
Bernt Øksendal. Stochastic differential equations. Springer, 1998.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference machine learning, pp. 2778-2787. PMLR, 2017.
Tongzheng Ren, Tianjun Zhang, Lisa Lee, Joseph E Gonzalez, Dale Schuurmans, and Dai. Spectral decomposition representation for reinforcement learning. arXiv preprint 200.::200009511/1 2022.

---

Published as a conference paper at ICLR 2023
Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function approximators. In Francis Bach and David Blei (eds.), Proceedings of the 32nd International Conference on Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. 1312-1322 Lille, France, 07-09 Jul 2015. PMLR. URL http:/ http///oooeeddeeee inggsddinggg ml m. w.w.es.s.vv.. schaul15 html.
Sungryull Sohn, Junhyuk Oh, and Honglak Lee. Hierarchical reinforcement learning for zero-shot generalization with subtask dependencies. Advances in Neural Information Processing Systems, 31, 2018.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction, MIT press, 2018 2nd edition.
Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsuper- vised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pp. 761-768, 2011.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 10X0000001.00690, 2018.
Yuandong Tian. Deep contrastive learning is provably (almost) principal component analysis. arXiv preprint arXiv:2201.12688 2022.
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems, 34, 2021.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine
learning research, 911)), 2008.
Yifan Wu, George Tucker, and Ofir Nachum. The Laplacian in RL: Learning representations with efficient approximations. In International Conference on Learning Representations, 2018.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. In International Conference Machine Learning, pp. 11920-11931. PMLR, 2021.
Jingwei Zhang, Jost Tobias Springenberg, Joschka Boedecker, and Wolfram Burgard. Deep rein- forcement learning with successor features for navigation across similar environments. 2017 JEEE/RSJ International Conference Intelligent Robots and Systems (IROS), pp. 2371-2378. IEEE, 2017.
12

---

Published as a conference paper at ICLR 2023
LIST OF APPENDICES
Appendix A sketches a proof the basic properties of SFs and the FB representation, Theorems and 2, respectively from Borsa et al. (2018) and Touati & Ollivier (2021).
Appendix B derives the loss (8) we use to learn and B, and the auxiliary orthonormalization loss.
Appendix C describes the precise mathematical relationship between the contrastive loss (17) and successor measures of the exploration policy.
Appendix D further discusses why low-rank models make more successor measures than the transition matrix P.
Appendix E proves that if P is low-rank given by x" μ. then successor features using y ju (but x) provide optimal policies.
Appendix F proves that the loss (16) used learn low-rank P is asymmetric version of the Laplacian cigenfunction loss (14).
Appendix G describes the detailed experimental setup: environments, architectures, policy training, methods for sampling of z, and hyperparameters. The full code can be found at https: : / giggtuuub com/ Facebookresearch control lable _ageet
Appendix H contains the full table of experimental results, as well as aggregate plots over several variables (per task, per replay buffer, etc.).
Appendix analyzes the learned features. We use feature rank analysis to study the degree of feature collapse for some methods. We also provide t-SNE visualizations of the learned embeddings for all methods (for the Maze environment).
Appendix provides hyperparameter sensitivity plots (latent dimension d, learning rate, batch size, mixing ratio for sampling).
Appendix describes a further baseline where we take goal-oriented method inspired from Ma et al. (2020), and extend it to dense rewards by linearity. Since results were very poor except for Maze (which is goal-oriented), we did discuss in the main text.
Appendix L provides Py snippets for the key losses, notably the FB loss, the SF loss well as the various feature learning methods for SF.
13

---

Published as a conference paper at ICLR 2023
A SKETCH OF PROOF OF THEOREMS 1 AND 2
To provide an intuition behind Theorems and 2 on SFs and FB, we include here a sketch of proof in the finite state case. Full proofs in the general case can be found in the Appendix of Touati & Ollivier (2021).
For Theorem (SFs), let assume that the reward is linear in the features y, namely, r(s) - ( w for some w E R Then by definition, 2, = since 2, is the linear regression of the reward on the features (assuming features are linearly independent). Using the definition (4) of the successor features v and taking the dot product with 2r, we obtain
(30 0o, 2r)
IE Σ (( 2r 10, Go T2r
E
Σ Enr((1++1)
| 80, ao Ter
(20)
since r r= op This means that 2b ao, 2r) is the Q-function of reward for policy Tz. At the same time, by the definition (4), Nz. is defined as the argmax of (/(50. Op. x+) 2,. Therefore, the policy æ, is the argmax of its own Q-function, meaning it is the optimal policy for reward r.
For Theorem 2 (FB), let us assume that FB perfectly satisfies the training criterion (5), namely, R F. B diag(p) in matrix form. Thanks to the definition (1) of successor representations M", for any policy æ2. the Q-function for the reward can be written as Q+= = M in matrix form. This is equal F. diag(p)r. Thus, if define 2, := B diag(p)r Ex~p (())r(s)), we obtain Q F. 2, for any Rd. In particular, the latter holds for 2, as well: F z, the Q-function of π., Again, the policies म, are defined in (5) as the greedy policies of F z, for any z. Therefore, is the argmax of its own Q-function. Hence, Tz is the optimal policy for the reward r.
14

---

Published as a conference paper at ICLR 2023
B DERIVATION OF THE BBRWARD-BACKWANN LOSS
Here we quickly derive the loss (8) used to train F and B such that Ms (,, da,)) a a(s,a)" B(')) p(ds'). Training is based on the Bellman equation satisfied by M". This holds separately for each policy parameter 2, so in this section we omit 2 for simplicity.
Here p(ds') the distribution of states in the dataset. Importantly, the resulting loss does not require to know this measure p(ds'), only to be able to sample states ~ from the dataset.
The successor measure M satisfies Bellman-like equation M P + pP, M matrices in the finite case and as measures in the general case (Blier et al., 2021). We can learn FB by iteratively minimizing the Bellman residual M (P yP M²) on the parametric model M F Bp.
M (s, a, ds') is a measure on s' for each (s, a), so it is not obvious how to measure the size of the Bellman residual. In general, we can define norm on such objects M by taking the density with respect to the reference measure p.
(
(M 2 p(ds')
|| = (s d'mp
(21)
where M(s.a,ds') p(ds") is the density of M with respect to p. For finite states, M is a matrix Mans' and this is just a p-weighted Frobenius matrix norm, M| Emae M p(s, ((s,a)/p(s)) (This is also how we proceed to learn low-rank approximation of P in (16).)
We define the loss F and B the norm of the Bellman residual M for the model M FT Bp. As usual in temporal difference learning, we use fixed, non-trainable target networks F. and B for the right-hand-side of the Bellman equation. Thus, the Bellman residual is f Bp (P PP F Bp). and the loss is
C(F, B) = PPpppPP+7PPp (P Bp)l
(22)
When computing the norm I|-Ilpp- the denominator cancels out with the F Bp terms, but we are left with P/p term. This term can still be integrated, because integrating P(ds'|s, a)/ p(ds")
under s' ~ p is equivalent to directly integrating under s' P(ds'|s, a)) namely, integrating under transitions (st, at 8t-1) in the environment. This plays follows:
( ats B(S) P(ds'(s (s, 7
ln.a a [F(++11 ((5++1)) B) 27
C(F, B) aa) s'
p(ds")
(23)
( (F(st. B(s') 7F (8441))
=Eisraeeeito s'mp
[F(st, at) B($+11) + Const 2 -
(24)
where Const a constant term that we can discard since it does not depend on and B.
Apart from the discarded constant term, all terms in this final expression be sampled from the dataset. Note that we have a (2 aT) B($++) term where Touati Ollivier (2021) have (st a (({:) term: this is because define successor representations (1) using 8141 while Touati & Ollivier (2021) use st.
See also Appendix L for pseudocode (including the orthonormalization loss, and double networks for Q-learning as described in Appendix G).
'This is the dual norm on measures of the L (p) norm functions. It amounts to learning a model of M(s, a, dss) by learning relative densities to reach 8 knowing we start (8, relative to the average density p(ds') in the dataset.
15

---

Published as a conference paper at ICLR 2023
The orthonormalization loss. Id. This loss
An auxiliary loss is used to normalize B so that (B()BB() ~
noor (B) : |B - ld 2
(25)
Frobenius
Egrup s'mup (B(s)) B(s'))² 1|B(s 1Bs') + Const.
(26)
(The more complex expression in Touati & Ollivier (2021) has the same gradients up to factor 4.)
The auxiliary loss (9). Learning F(s, a, 2) is equivalent to learning F a, 2) 2' for all vectors z'. Yet the definition of the policies π, in FB only uses Fss, a, 2)) z. Thus, was done for SFs in Section 5.1, one may wonder there is scalar rather than vector loss to train F, that would reduce the error in the directions of F used define Z--
In the case of FB, the full vector loss needed to train F and B. However, one can add the following auxiliary loss F to reduce the errors in the specific direction F(s,a, a,,a a)²2).22. 2. This is obtained as follows.
Take the Bellman gap in the main FB loss (22): this Bellman gap is F Bp (P+ Bp). To specialize this Bellman gap in the direction of F(s, a, 2)T 2, multiply by B" (BpB")- on the right: this yields F (PB' (BpB) y- yP, F z) (using that compute the loss at B B).
This new loss is the Bellman gap on F" 2, with reward PB (BpB))
This is the loss L' described in (9). It is particular case of the main FB loss: L = 0 implies L = 0, since we obtained L' by multiplying the Bellman gap of C.
We use L' on top of the main FB loss to reduce errors in the direction F(s, a, 2) z. However, in the end, the differences are modest.
16

---

Published as a conference paper at ICLR 2023
C RELATIONSHIP BETWEEN CONTRASTIVE LOSS AND SVD OF SUCCESSOR MEASURES
Here prove the precise relationship between the contrastive loss (17) and SVDs of the successor measure of the exploration policy.
Intuitively, both methods push states together if they lie on the same trajectory, by increasing the dot product between the representations of and 8 This is formalized as follows.
Let be the policy used produce the trajectories in the dataset. Define
M == (1 VCL)
cz P
(27)
120
where "cL is the parameter of the geometric distribution used choose & when sampling 8; and st4k-
M is stochastic matrix in the discrete case, and probability measure over S in the general case: it is the normalized version of the successor measure (3) with the exploration policy.
By construction, the distribution of 8 knowing 8t with ~ Geom -cL.) is described by M. Therefore, rewrite the loss as
E Se-eeon(1-aa) (8)
log Brr.occppccosinne(((a......
(28)
M ( ds')
Esxp, s'rup
log bggcxplcossine((())155))
plds')
+ Eaup [log EOenD exp(cosine((((())(5))
(29)
Assume that oo and j are centered with unit norm, namely, ||x())) and Es- (ps) 0 and likewise for μ. With unit norm, the cosine becomes just dot product, and the loss is
(s, ds') p(ds') ((s) ((6')
- Esnp, strup
+ Eamp [log (xx(((())mmss))]
(30)
second-order Taylor expansion provides
log E exp X = EX + E[X (EX)² +(1x))
(31)
and therefore, with E = E μ = O, the loss approximately
(s, ds')
p(ds') :p(s] umd)) + Bonp, s'ap u(s'))²
Banp, stou
(32)
27
(s, p(ds')
E Egwp, strup
(P() (p(s))((s)) u(s))
Const
(33)
where the constant term does not depend yp and μ.
This is minimized when Tp μ is the SVD of M/p in the (2()) norm.
17

---

Published as a conference paper at ICLR 2023
D DO FINITE-RANK MODELS ON THE TRANSITION MATRIX AND ON SUCCESSOR MEASURES MAKE SENSE?
turns finite-rank models very different for P for Σnep. In typical situations, the spectrum P is concentrated around while that of Σ 7 Px is much more spread-out.
Despite the popularity of low-rank in theoretical RL works, P never close low-rank in continuous-time systems: then P is actually always close to the identity. Generally speaking, P cannot be low-rank if most actions have a small effect. By definition, for any feature function y. s). Intuitively, if actions have small effect, then 8:41 is close to St. and ((+++1) (st) for continuous sp. This means that Px close (,, so that Pa is close to the identity on large subspace of feature functions s. In the theory of continuous-time Markov processes, the time-/ transition kernel given by P. e with the infinitesimal generator the process (Levin et al., 2009, $20.1) (Øksendal, 1998, 18.1), hence Pt is Id (tt) for small timesteps In general, the transition matrix P is better modeled as Id + low-rank, which corresponds to low-rank model of the Markov chain Laplacian P.
(s) = B[(p($(+1)
On the other hand, though Σn'p' x is never exactly low-rank (it is invertible), it has meaningful low-rank approximations under weak assumptions. For large t. Pt becomes rank-one under weak assumptions (ergodicity). as it converges to the equilibrium distribution of the transition kernel. For large 7. the Σn Pd is dominated by large f. Most eigenvalues of Pa close 1, but taking powers P sharpens the differences between eigenvalues: with close to 1, going from Pa to into 1/E.
ΣT (Id ,;,) changes eigenvalue
In short, on P itself, there is little learning signal to differentiate between eigenvectors, but differences become visible over time. This may explain why FB works better than low-rank decompositions directly based on P or the Laplacian.
For instance. consider the neightt random walk on length-n cycle {0, .,...,... . 1
mod n),, namely, moving in dimension 1. (This extends to any-dimensional grids.) The associated Pz not low-rank in any reasonable sense: the corresponding stochastic matrix concentrated around the diagonal, and many eigenvalues are close to 1. Precisely, the eigenvalues are cos(2kz /n) with integer f0, ........ n/2). This is 2x 2x²(kkn)) (k/n)² when k n. Half of the eigenvalues are between 22/2 and 1.
However, when - 1. >'' has one eigenvalue 1/(1 - 7) and the other eigenvalues are [2k0s(2k//n] n²/2m=² with positive integer k: there one large eigenvalue, then the others decrease like cst/k?. With such spread-out spectrum, finite-rac rank model makes sense.
18

---

Published as a conference paper at ICLR 2023
E WHICH FEATURES ARE ANALOGOUS BETWEEN LOW-RANK P AND SFS?
Here we explain the relationship between a low-rank model of P and successor features. More precisely, if transition probabilities from (s, a) to s' be written exactly as x(s, a) u(s',, then SFs with basic features := will provide optimal policies for any reward function (Theorem 3).
Indeed, under the finite-rank model P(ds'|s, |s, a) x(s,a) a)" ((d')p(ds")) rewards only matter via the reward features Eg'op (=s") "rrrsrss))) namely, two rewards with the same reward features have the same Q-function, as the dynamics produces the same expected rewards. Then Q-functions are linear in these reward features, and using successor features with follows.
μs provides the correct Q-functions, as
Theorem 3. Assume that P(ds'|s,o) x(s, a) pls' ((s')) Then successor features using the basic features 9p == provide optimal policies for any reward function..
This is why we use μs. rather than X for the SF basic features. This is also why we avoided the traditional notation P(s'|s. [ss a) = (pss a) pls') often used for low-rank P. which induces conflict of notation with the in SFs, and suggests the wrong analogy.
Meanwhile, X plays role more analogous to SFs' although for one-step transitions instead of multistep transitions as in SFs: Q-functions are linear combinations of the features X. particular, the optimal Q Q-function for reward Q x w, for some Wyy. But contrary successor features, there is no simple correspondence to compute w, from
Proof. Let æ be any policy. On a finite space matrix notation, and omitting p for simplicity, the assumption P = x" μ implies P= Nt where (x) = Besx(s) x(8, a) are the m-averaged features. Then,
O = P Σ
P
(34)
120
7 (xTP)'r
(35)
$20
x (Eizo (1x3T)) pr.
(36)
Thus, Q-functions expressed as (7, x(3,a)(w(ZTr) w(x,r) with w((,r) ((xo))) ur.
Moreover, rewards only matter via pr. Namely, two rewards with the same pr have the same Q-function for every policy.
In full generality on continuous spaces with p again, the same holds with p(s)) instead of y(s).
and E,-p (s))r(s) instead of urr
Now, let be any reward function, and let r' be its (p) projection onto the space generated by the features μ.. By construction, r-r' L (p)))-oothogonall to μ., namely, E, u(r-r') So E, pr E, pr', Therefore, by the above, r and r' have the same Q-function for every policy.
By definition, r' lies in the linear span of the features μ. By Theorem 1, SFs with features & = lr will provide optimal policies for r'. Since and r' have the same Q-function for every policy, optimal policy for r' is also optimal for r.
19

---

Published as a conference paper at ICLR 2023
F RELATIONSHIP BETWEEN LAPLACIAN EIGENFUNCTIONS AND LOW-RANK P LEARNING
The loss (16) used learn a low-rank model of the transition probabilities P is asymmetric version of the Laplacian eigenfunction loss (14) with À = 1.
Said equivalently, if we use the low-rank P loss (16) constrained with X = μ learn a low-rank model of Px instead of P (with ı the exploration policy), then we get the Laplacian eigenfunction loss (14) with X = 1.
Indeed, set À = 1 in (14). Assume that the distributions of 8, and 8141 in the dataset are identical on average (this happens, e.g., if the dataset made of long trajectories if p is close enough to the invariant distribution of the exploration policy). Then, in the Laplacian loss (14), the norms from the first term cancel those from the second, and the Laplacian loss simplifies to
(14) E E(nnert)) [1(p(s)) - (0(841))11
+ E END [(v(s) ((s))) l1((ss) ll(((')) 1|1
R'DD
(37)
BR 1|p(()))11 + B [|=()))11 - 2 (e (pp(0)) p((0+1)]]
(38)
+ E END ((p()) - -1<(ss) [|(((s)) [2]
(39)
S'DD
- (1p(s)) (((++1)]
EaDD [(v(()) (s)))))] iD
(40)
This is the same as the low-rank loss (16) if we omit actions a and constrain X =
20

---

Published as a conference paper at ICLR 2023
Figure 3: Maze, Walker, Cheetah and Quadruped environments used in experiments. In the Mmaze domain (left), we show an example of an initial state (yellow point) and the 20 test goals (red circles).
G EXPERIMENTAL SETUP
In this section we provide additional information about our experiments.
Code snippets for the main losses are given in Appendix L. The full code can be found at http:://hh..la.ccoccccccccortttt
G.1 ENVIRONMENTS
All the environments considered in this paper are based on the DeepMind Control Suite (Tassa et al., 2018).
Point-mass Maze: 2-dimensional continuous maze with four rooms. The states are d-dimensional vectors consisting of positions and velocities of the point mass (æ, y, U2 Zyy. and the actions are 2-dimensional vectors. At test, we assess the performance of the agents 20 goal-reaching tasks (5 goals in each room described by their (æ, y) coordinates).
Walker: planar walker. States are 24-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors. We consider different tasks at test time: walker_stann reward is combination of terms encouraging upright torso and some minimal torso height, while walker_wall and walker_run rewards include component encouraging
minimal forward velocity. walker flip reward includes a component encouraging some mininal angular momentum.
Cheetah: a running planar biped. States 17-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors. We consider 4 different tasks time: cheetal walk and cheetah run rewards are linearly proportional to the forward velecity up some desired values: 2 m/s for wa call k and 10 m/s for run. Sim- ilarly, walker_ wa backward and walker_run, caruuuuacckeeee backward rewards encourage reaching some minimal backward velocities.
Quadruped: a four-leg ant navigating 3D space. States and actions are 78- dimensional and 12-dimensional vectors, respectively. We consider 4 tasks at test time: quadruped staand reward encourages an upright torso. quadruped walk and quadruped run include term encouraging some minimal torso velecities. quadruped_we k includes a term encouraging some minimal height of the center of mass.
G.2 ARCHITECTURES
We use the same architectures for all methods.
The backward representation network (s)) and the feature network (o(s) are represented by a feedforward neural network with three hidden layers, each with 256 units, that takes as input a state and outputs a L2- -2-normalized embedding of radius vd.
For both successor features (5(s, a, a) and forward network F(s, 2), we first preprocess separately (s, a) and (s, 2) by two feedforward networks with two hidden layers (each with 1024 units) to 51 512-dimentional space. Then concatenate their two outputs and pass it
21

---

Published as a conference paper at ICLR 2023
into another 2-layer feedforward network (each with 1024 units) to output d-dimensional vector.
For the policy network π (s,, 2). we first preprocess separately and (s, 2) by two feedforward networks with two hidden layers (each with 1024 units) to 512-dimentional space. Then we concatenate their two outputs and pass it into another 2-layer feedforward network (each with 1024 units) to output to output a da-dimensional vector, then we apply a Tanh activation the action space is [-1, 1jdA.
For all the architectures, we apply a layer normalization (Ba et al., 2016) and Tanh activation in the first layer in order standardize the states and actions. We use Relu for the rest of layers. We also pre-normalized 2: z T in the imput of F, π and vb. Empirically, we observed that removing preprocessing and pass directly a concatenation of (s, z) directly to the network leads to unstable training. The same holds when we preprocess (s, a) and instead of (s, a) and (s, 2), which means that the preprocessing of 2 should be influenced by the current state.
||
For maze environments, we added an additional hidden layer after the preprocessing (for both policy and forward / successor features) as it helped to improve the results.
G.3 SAMPLING OF 2
We mix two methods for sampling
1. We sample uniformly in the sphere of radius vd in IR (so each component of is of size 1).
2. We sample using the formula for 2, corresponding to the reward for reaching ran- dom goal state s Theorems 1-2. Namely, we set z = B(s) for FB and 2 pls) for SFs, where ~ is random state sampled from the x- (ss)) (pss)
m2
replay buffer and 8; are states in a minibatch.
For the main series of results, we used a 50% mix ratio for those two methods. Different algorithms can benefit from different ratios: this is explored in Appendix
G.4 LEARNING THE POLICIES ny POLICY NETWORK
As the action space is continuous, we could not compute the arg max over action in closed form. Instead, we consider latent-conditioned policy network Ty : S x Z - A, and we learn the policy parameters by performing stochastic gradient ascent on the objective [F,ssF(s, (s), =)_ for FB [v, (s), 2) =)]= for SFs.
We also incorporate techniques introduced in the TD3 paper (Fujimoto et al., 2018) to address function approximation error in actor-critic methods: double networks and target policy smoothing, adding noise e to the actions.
Let O2 and O2 the parameters of two forward networks and let the parameters of the backward network. Let OT, Og and -- be the parameters of their corresponding target networks.
Let {(st, ag, sq smex)) ier D mini-batch of size 7 b of transitions and let (ei mini-batch of size / = b of latent variables sampled according to G.3. The empirical version of the main FB loss in (8) is (with an additional 1/2 factor):
2
L(0,.w) = 26(b 1)
Σ
Fo, (s, ai, ((, - min Fo, (spext (spexc))) + Ess 24) Bu (spexe)
1-12
iJER
143
Σ RET
Poz (st, a, 2.) (xprx)) vk = 1,2
(41)
b
where E; is sampled from a truncated centered Gaussian with variance o² (for policy smoothing). The empirical version of the auxiliary F loss in (9) is: for A 1,2,
22

---

Published as a conference paper at ICLR 2023
1 Σ SEI
2
58 (0x)
Fo, (Si: (i) TT)T= 21 (o Cov+ - min Foj (sfext, Tny 1=1,2
(speett 2) + Ess 2
(42)
where Cov't is the pseudo-inverse of the empirical covariane matrix Cov & Eier B. We 1/d regularization coefficient in front of _ (0x). For policy training, the empirical loss is:
L(n) = - Σ
min Fo, (8,, (,(6 + sis 2i)
(43)
NE7 1112
The same techniques are also used for SFs.
G.5 HYPERPARAMETERS
Table
I summarizes the hyperparameters used in our experiments.
A hyperparameter sensitivity analysis for two domains (Walker and Cheetah) is included in Ap-
pendix J.
Table I: Hyperparameters used in experiments.
Hyperparameter
Value
Replay buffer size
5 x 10" (10 x 10 for maze)
Representation dimension
50 (100 for maze)
Batch size
1024
Discount factor ?
0.98 (0.99 for maze)
Optimizer
Adam
Learning rate
10-4
Mixing ratio for sampling
0.5
Momentum coefficient for target networks
0,99
Stddev for policy smoothing
0.2
Truncation level for policy smoothing
0.3
Number of gradient steps
10
Number of reward labels for task inference
104
Discount factor "eu for CL Regularization weight for orthonormality loss (spectral methods)
0.6 (0.2 for maze) 1
Hyperparameter tuning. Since all methods share the same core, we chose to use the same hyperparameters rather than tune per method, which could risk leading to less robust conclusions.
We did not do hyperparameter sweeps for each baseline and task, first because this would have been too intensive given the number of setups, and, second, this would be too close to going back to supervised method for each task.
Instead, to avoid any overfitting, we tuned architectures and hyperparameters by hand on the Walker environment only, with the RND replay buffer, and reused these parameters across all methods and tasks (except for Maze, on which all methods behave differently). We identified some trends by monitoring learning curves and downstream performance on Walker, and fixed configuration that led to overall good performance for all the methods. We avoided a full sweep to avoid overfitting based Walker, focus robustness.
For instance, the learning rate 10-4 seemed to work well with all methods, as seen in Appendix
Some trends were common between all methods: indeed, all the methods share common (training of the successor features ya or F, and training of the policies т) and differ by the training of the basic features yp or B.
23

---

Published as a conference paper at ICLR 2023
For the basic features ip, the various representation learning losses easy to fit and had low hyperparameter sensitivity; we always observed smooth decreasing of losses until convergence. The challenging part was learning & and F and their corresponding policies, which is common across methods.
24

---

Published as a conference paper at ICLR 2023
H DETAILED EXPERIMENTAL RESULTS
We report the full experimental results, first as a table (Section H.1).
In Section H.2 we plot aggregated results for easier interpretation: first, aggregated across all tasks with a plot of the variability for each method; second, aggregated over environments (since each environment corresponds
trained zero-shot model); third, for each individual task but still aggregated over replay buffers.
In Section H.3 we plot all individual results per task and replay buffer.
In Section H.4 we plot aggregate results split by replay buffer.
H.1 FULL TABLE OF RESULTS
Buffer
Domain
Task
Method
Rand
ICM
Latent
Lap
LRA-P
LRA-BR
APS
eboctah
run
145+7
97=3
432=8
49:14
133=15
198:
8+1
2:0
247=10
267
25=3
backward
189420
365:6
404:3
32433
B82
221:14
11:0
14:6
261
23817
98-21
walk
665-b60
104=19
928=544
302=67
75=31
2+1
918=22
walk-backward
653++75
982++0
986-=0
453=100
985-0
937
8+
150+51
983-00
981
452+777
maze
reach
11=35
5+1
8#3
10:4
432
436+16
12=3
145=8
59=7
quadraped
jump
784+:55
727+:12
164=22
554-24
624 14
718 18
309+:50
102++29
632-220
649-223
311
APS
rum
459
19
382
448
196:110
stand
966++3
925
248-446
752-36
890-17 17
963
497+
56=99
872-220
924+ 13
417+ 10
543.111
108
524
22
463-18
walker
flip
158.110
317.131
452.13
23
471 15
45
340.118
29
186.121
4131:16
39
96224
127#8
29029
289:1
115+11
204:266
3461:14
35=3
stand
486.127
617.137
19
568-57
24
895.19
641113
205:151
$911135
822:126
176.117
1774:30
462+58
159:220
671=199
8172215
cheetah
rum
5818
84
133114
2715
14212
14913
0110
209110
2101
run ward
149.1
262.17
129-1
30:55
146-1
13346
230-16
walk
287.138
127:151
961
112-11
929 13
12
770.131
110
860131
908
walk-backward
6641::9
973#33
987=00
101=15
6294:31
151:88
7424:46
maze
reach
27:
7+1
742
15:44
8
16
15
19=5
134
526
quadruped
jump
196::29
185
137=27
309
184
113:122
1834224
Proto
run
134+117
23428
88=13
123=17
191-14
125:
196+19
65=16
99--9
137+:14
stand
413-1:5
121-142
290--29
270-38
436
231-:52
264 34
86::12
215-147
287+-53
walk
148:+21
171+21
122=10
156=24
212 12
135:+16
1724-21
40=13
100=
280
walker
flip
133+-12
346-16
510=13
443-
456 10
$48
2-1
78:-21
551::13
507
run
84.12
234.99
189
347
303
280.122
183.128
45.L5
191
336.19
stand
415-:26
905:19
93
582-62
951
937-45
6874-41
251--49
874-423
902-4-25
walk
125.121
632
17
791 17
832 16
883 30
300.11
92
867.11
91717
cheetah
6423
96:3
183-20
66:44
61:
16
138=17
247
rum buck ward
96.16
162 18
160
60.4
143 17
90
21.0
13
82.18
185
walk
2894-22
38
401+40
567-
286 30
330+66
29-111
28
446-33
4-1
walk. backward
469.1.8
542.170
743
345
572 100
199
14.1
517.66
352.145
166
reach
91:2
41:1
411
824
8::4
12
759
736:13
232=200
810
quadruped
jump
770.1.
474.140
17
26
806
490.448
447
326.1
991
651 1
RND
465:1:4
415::0
981:1
418=10
478-8
299:126
3011:12
263:140
461:16
stand
919:+19
770+228
426=-45
830=31
720=27
552:+32
729
94-=111
815
walk
586:1:88
486:1:37
90
21
552-33
410::30
310::14
210::42
316:130
10
walker
flip
267+:28
352+:18
461=9
34:3
399=19
569-300
512:227
50=9
454=223
578=110
96--8
167-t112
251=111
47+:22
250-6
299-220
325++ 10
38=55
B50 14
388
516.1.36
733.122
813:10
191-62
904
326-141
20
190
walk
152++3
4574:3
518=74
30+2
607-31
7584775
818+
53-100
853
760+:19
Table 2:
Score of each method, split by task and replay buffer. Average over random seeds, with + la estimated standard deviation on this average estimator. For Maze, we report the average over the 20 goals defined in the environment. We highlight the three leading methods (four when confidence intervals overlap) for each task.
25

---

Published as a conference paper at ICLR 2023
H.2 AGGREGATE PLOTS RESULTS
Here we plot the results, first averaged over everything, then by environment averaged over the tasks of that environment, and finally by task.
80
60
60
Norma fead
20
20
Rand
ICM
Trans
LRA-P
LRA-SR
APS
Figure 4: Zero-shot scores of len SF methods and FB, aggregated over tasks using normalized scores as described in the text. To assess variability, the box plot on the right shows the variations of the distribution normalized scores over random seeds, environments, and replay buffers.
quadruped
return
000
600
Online TD3
Rand
ICM
Trans
LRA-P
LRA-SR
FB
APS
Offiine TD3
AEnc
Latent
Lap
CL
Figure 5: Zero-shot scores averaged over tasks for each environment, with supervised online and offline TD3 toplines. Average over replay buffers and 10 random seeds.
250
600m
Figure 6: Zero-shot scores for each task, with supervised online and offline TD3 as toplines. Average over 3 replay buffers and 10 random seeds.
26

---

Published as a conference paper at ICLR 2023
H.3
FULL PLOTS OF RESULTS PER TASK AND REPLAY BUFFER
3040
Chockan
sland
0ffin TD3
Lockett
Figure 7: Per-task results on the RND replay buffer, average over 10 random seeds.
2000
304
Wheetegue Jackward
Rand
Ond
10M
Trara
Figure 8: Per-task results the APS replay buffer, average over 10 random seeds.
2000
300
01
Lackett
Figure 9: Per-task results on the Proto replay buffer, average 10 random seeds.
27

---

Published as a conference paper at ICLR 2023
H.4 INFLUENCE OF THE REPLAY BUFFER
Here we plot the influence of the replay buffer, by reporting results separated by replay buffer, but averaged the tasks corresponding each environment (Fig 10).
Overall, there clear failure case of the Proto buffer on the Quadruped environment: the TD3 supervised baseline performs poorly for all tasks in that environment.
Otherwise, results are broadly consistent on the different replay buffers: with a few exceptions, the same methods succeed or fail the same environments.
cheatah
ped
Online TD3
Rand AEnc
ICM
Trans Lap
LRA-P
LRA-SR
Offline
Latent
maze tasks
w walker tasks
cheatah tasks
quadruper pad tasks
Online
Rand AEnn
LRA-SR
Offline TD3
Latent
Lap
CL
tasks
walker
cheetah tasks
tasks
return
Online TD3
LRA-P
Offlnne TD3
AEnc
Latent
Figure 10: Results each replay buffer: RND (top), APS (middle), Proto (bottom). Average over tasks for the Walker, Cheetah and Quadruped environments, average over 20 goals for Maze.
28

---

Published as a conference paper at ICLR 2023
In Fig. 11, plot results aggregated over all tasks but split by replay buffer. Box plots further show variability within random seeds and environments for a given replay buffer.
Overall method rankings are broadly consistent between RND and APS, except for CL. Note that the aggregated normalized score for Proto is largely influenced by the failure on Quadruped: nor- malization by very low baseline (Fig. 10, Quadruped plot for Proto), somewhat artificially pushes Trans high up (this shows the limit of using scores normalized by offline TD3 score), while the other methods' rankings are more similar RND and APS.
FB and Lap work very well in all replay buffers, with LRA-SR and Trans a bit behind due to their feailures some tasks.
80
40
Rand AEnc
ICM
Trans Lap
LRA-P
LRA-SR
PB
Latintt
CL
80
Rand AEmn
Trans Lap
URA-P
URA
Latent
Rand
ICM
Trans
LRA
URA-SR
PB
AEnc
Latont
Lap
CL
Figure 11: Zero-shot scores by replay buffer, as a percentage of the supervised score of offline TD3 trained on the same buffer, averaged over tasks and environments and random seeds. Top: RND buffer; middle: APS buffer; bottom: Proto buffer.
29

---

Published as a conference paper at ICLR 2023
I ANALYSIS OF THE LEARNED FEATURES
I.1 FEATURE RANK
Here the hypothesis that feature collapse for some methods responsible for some cases of bad performance. This is especially relevant for Maze, where some methods may have little incentive learn more features beyond the original two features (x, y).
We report in Table 3 the effective rank of learned features B or yi this is computed as in (Lyle al., 2022), as the fraction of eigenvalues of Exn. (s)) B(x) ESs (())s((s)) above a certain threshold.
Domain
Method Lap
Rand
AEnc
ICM
Latent
Trans
LRA-P
CL
LRA-SR SR
FB
Maze
0.38
0.31
0.33
0.32
0.15
1.0
1.0
1.0
1.0
1.0
Walker
1,0
0.91
1.0
1.0
0.90
1.0
1.0
0.65
1.0
1.0
Cheetah
1.0
0.96
1.0
1.0
1.0
Quadruped
1.0
0.98
0.30
1.0
0.15
1.0
0.84
1.0
1.0
1.0
Table 3:
Feature rank (Lyle et al., 2022) of o B for each method, computed as # {a E (H E. ((s))((s)) [ > E). trained on RND replay buffer and averaged over 10 random seeds. We use = 100, 000 samples to estimate the covariance, and & 10-4
For most methods and all environments except Maze, more than 90% (often 100%) of eigenvalues are above 10-4, 4, so the effective rank close full.
The Maze environment is clear exception: on Maze, for the Rand, AEnc, Trans, Latent and ICM methods, only about one third of the eigenvalues are above 10- The other methods keep 100% of the eigenvalues above 10-1, This is perfectly aligned with the performance of each method on Maze.
So rank reduction does happen for some methods. This reflects the fact that two features (x, v) already convey the necessary information about states and dynamics, but are not sufficient solve the problem via successor features. In the Maze environment, auto-encoder or transition models can perfectly optimize their loss just by keeping the original two features (z., y), and they have no incentive to learn other features, so the effective rank could have been 2.
These methods may benefit from auxiliary losses to prevent eigenvalue collapse, similar to the orthonormalization loss used for B. We did not include such losses, because wanted to keep the same methods (autoencoders, ICM, transition model...) used in the literature. But even with an auxiliary loss, keeping full rank could be achieved just by keeping (z,, y) and then blowing up some additional irrelevant features.
30

---

Published as a conference paper at ICLR 2023
1.2 EMBEDDING VISUALIZATION VIA T-SNE
We visualize the learned state feature embeddings for Maze by projecting them into 2-dimensional space using t-SNE (Van der Maaten Hinton, 2008) in Fig. I.2.
AEnc, Trans, FB, LRA-P and CL all recover a picture of the maze. For the other methods the picture less clear (notably, Rand gets near-circle, possibly as an instance of concentration of measure theorems).
Whether t-SNE preserves the shape of the maze does not appear be correlated performance: AEnc and Trans learn nice features but perform poorly, while the t-SNE of Lap is not visually interpretable but performance good.
Raw State
Rand
AEnc
+
ICM
Trans
Latent
Lap
LRA-P
LRA-SR
FB
Figure 12: Visualization embedding vectors obtained by each method (p for SE and B for FB) on the maze domain after projecting them in two-dimensional space with t-SNE.
31

---

Published as a conference paper at ICLR 2023
J HYPERPARAMETERS SENSITIVITY
732
032
300
400
100
130
130
0.5
400
200
332
200
leaming
FB
AEEe
CL
Latent
LRA-P
Rand
ICM
Thans
Lap
LRA-SR
Figure 13: Return for each method trained Walker domain on the RND replay buffer, for different choices of hyperparameters. Average over Walker tasks and 5 random seeds.
535
605
302
205
105
136
152
3.50
025
6.56
3.75
3,50
Intent dimension
mixing ratio sampling
200
506
100
202
102
1020
2003
1ec
Latent
LRA-P
LRA-SR
Figure 14: Return for each method trained in Cheetah domain, on the RND replay buffer, for different choices of hyperparameters. Average over Cheetah tasks and 5 random seeds.
32

---

Published as a conference paper at ICLR 2023
K GOAL-ORIENTED BASELINES
Goal-oriented methods such as universal value functions (UVFs) (Schaul ct al., 2015) learn a Q- function Q(s, a, g) indexed by goal description g. When taking for g the set of all possible target states E S, these methods can learn to reach arbitrary target states. For instance, Ma et al. (2020) use mixture of universal SFs and UVFs to learn goal-conditioned agents and policies.
In principle, such goal-oriented methods are not designed to deal with dense rewards, which are linear combinations of goals g. Indeed, the optimal Q-function for linear combination of goals is not the linear combination of the optimal Q-functions.
Nevertheless, may still try may be applied at the level of the Q-functions, or at the level of some goal descriptors, as follows.
how such linear combinations perform. The linear combination
Here, an additional baseline for our experiments, we test slight modification of the scheme from Ma et al. (2020), as suggested by one reviewer. We learn two embeddings (15, w) using state-reaching tasks. The Q-function for reaching state s' (goal g s') is modeled as (ss ar s') = (5 ae, woss)) w(s'). trained for reward 1 s.41= A policy network T (ss, w(s')) outputs the action at for reaching goal 8'. Thus, the training loss is
(44)
[(b,, w) = F( (bb( at, w(s')) w(s') - 1 (s41100) 70($11), T (S141, w(s')) "(()),)())))))))))))))))))
T(s'))_ass'))
d'mp
(45)
Similarly to the other baselines, the policy network T(ss, w(s)) is trained by gradient ascent on the policy parameters to maximize
Eacp [Q(s, π (8, (())),,, 8') = TE amp s'oup
[w)s, x (8, wnss) m(6))) wss))
(46)
s' s'oup
At test time, given reward r, we proceed as in FB and estimate
= (r))pl((s))
(47)
and use policy (ss, 2). This amounts to extending from goal-reaching tasks to dense tasks by linearity on w. even though in principle, this method should only optimize for single-goal rewards.
We use the same architectures for (w), w) as for (F, B), with double networks 4/, t/2 and policy smoothing as (43).
The sparse reward 1 (Nx [1=)) could suffer from large variance in continuous state spaces. We mitigate this either by:
biasing the sampling of s' by setting s' to 81+1 half of the time.
replacing the reward by a less sparse one, 1 [[]e:]1=*(laSe)
Table 4 reports normalized scores for each domain, for the two variants just described, trained on the RND replay buffer, averaged over tasks and over 10 random seeds.
The first variant scores 84% on Maze, 62% on Quadruped tasks, 17% on Walker tasks, and 2% on Cheetah tasks. The second variant worse overall.
So overall, this works well on Maze (as expected, since this is goal-oriented problem), moderately well on Quadruped tasks (where most other methods work well), and poorly on the other environments. This is expected, as this method is not designed to handle dense combinations of goals.
Relationship with FB. The use of universal SFs Ma et al. (2020) is quite different from the use of universal SFs in Barreto et al. (2017) and Borsa al. (2018). The latter is mathematically related to FB, as described at the end of Section 4. The former uses SFs as an intermediate tool for a goal-oriented model, and is more distantly related to FB (notably, it is designed to deal with single goals, not linear combinations of goals such as dense rewards).
A key difference between FB and goal-oriented methods is the following. Above, (iss, a, w(g))" w(g) as a model of the optimal Q-function for the policy with goal g. This is reminis- cent of F(s, a, z) B(g) with B(g), the value of used to reach g in FB.
33

---

Published as a conference paper at ICLR 2023
Reward
Domain Chectah
Maze
Walker
Quadruped
(sma')
83.82
17.06
1.60
61.92
(1a=2"1a54)
87.44
12.42
0.74
20.25
Table 4: Normalized score for each domain of two variants of USF from Ma et al. (2020) trained on RND replay buffer, averaged over tasks and 10 random seeds
However, even we restrict FB to goal-reaching by using : = B(g) only (a significant restriction), then FB learns F(s, B(g))' B(g") model the successor measure, i.e., the number of visits each goal g for the policy with goal g, starting at (s, a). Thus, FB learns an object indexed by (s, a, 9. g') for all pairs (g.9') g").
Thus, FB learns more information (it models successor measures instead of Q-functions), and allows for recovering linear combinations of goals in principled way. Meanwhile, assuming perfect neural network optimization in goal-reaching methods, there no reason the goal-oriented policies would be optimal for arbitrary linear combinations of goals, only for single goals.
34

---

Published as a conference paper at ICLR 2023
L PSEUDOCODE OF TRAINING LOSSES
Here we provide PyTorch snippets for the key losses, notably the FB loss, SF loss as well as the various feature learning methods for SF.
def compute_ Ba_uutt____________________ loss (agent, obs, action, next_obs, 2, discount):
compute target successor measure
with sarch.no_gradiy
mu agent.policy. net mext__bbs
next_action TruncatedNormal (mu-mu, ofodev-agent..fg.scco clip- agent.cfg..tdeev_cl..
target_Fl, target_F2 agent. forerrtfrrwar (next_obs, 2,
next_action # batch x z dim
target_ B agent.backwardd wargrrtttttttttree net
# batch 2 dim target_MI MI, target_ M2 [torch.einsum// sd, tá st', target_Fi, [F1, F2]] batch x batch target_M = torch.min (target_ML, target_M2)
target_B) for target_Fi
12
13
compute the main FB loss
15
16
F1, agent. fowwwrddwwaaeroeee _et (obs, 2, action)
17
B .rcett.bccwwrd..oo net (next_obs obs)
net
18
M1, M2 = storch.einsume (' sd, td -> st', Fi, B) for Fi in [F1, F2] #
batcn batch
19
torch. (+M1 size (), device=M1 device)
20
off_diag ~1. -oobbol ()
fb fb_offdiag: tp.Any 0.5 sum( discount target_M) [off loff_diag] .aaaalpppoo (2) me for M in [Mi, M21)
fb diag: tp. -sum (M.diag() mean for in [M1, M2])
23
fb loss fb_offdiag fb diag
24
25
# compute the auxiliary loss
27
next_01, nextQ2 (torch.einsum sd, sd -> s', target_Fi, 2) for target_Fi in [target_Fi, target_F2]]
28
next_ = terch.min (ne next(2)
merrch.aammi (8.T, B) / 8.shape inv cov corch.linalg.95 (cov)
31
implicit_rewant reward (torch.matmin (B, inv_cov) z) s).sum(dinn]
batch target_0 implicit-.evardddettan
32
discount next_ Il batch_size
34
Q1, Q2 ferrcheeinsum( ('sd, sd, s', Fi, 2) for in [F1, F2]] loss .mse_loss (01, target_0) + F.mse_loss (02, target_0)
35
36
loss /= agent.cfg.. 2ft.c_...aaa..ooo. dim
fo_s += goloss
39
compute Orthonormality losss
40
Cov - norch.matmul (8, B.T) 41
orth_ orth_loss_diag
2 (oo.ddaagl ()ean
orth_loss_offdiag
cov[off_diag)
mean ()
44
orth loss orth loss offdiag + erth loss _oss_diag
45
tb fbss += ofgtt.c.g..rttho.com
orth_looo
return foss
Listing Pytorch code for FB training loss
def compute_st. losssssss__ooseeeoooeeessossssss (agent, obs, action, next_obs, discount):
compute target q-value
35

---

Published as a conference paper at ICLR 2023
with torch. no :
mu
(next_obs,
next_action = TruncatedNormal (mu=mu, stadev-agent.cfg.cco stddov, clip- egent.cfg.stddeey
next_Fl, next_F2 wtt.:tsccccsssr_areeeteee.. (next_cba, 2, next_action) batch z_dim
target_phhi agent. www.ureeteeeeeeeee net (next_goal) .ettcc () batch z_dim
next_0L, next_Q2 = [torch. sd, sd -> s', next_Fi, 2) for next_fi in [next_Ft, next_F2]]
next_ torch. min (next 21, next_Q2)
target_ torch.einsum sd, next_
s', target_phi,
discount
12
F1, F2 = agent. successor net (cbs, action)
Q1, [torch. sd,
, Fi, for Fi in [F1, F2]]
15
sf loss E.mse_loss (Q1, target_O) F f.mse_loss (02, target_0)
16
return sf_loss 17
Listing Pytorch code for SF training loss
def compute_phi_ phtspphhhsss___peeee loss (agent, obs, next_obs):
phi agent .eent.feae feature net (obs)
next phi agent. feaetreeeatre net (next [ loss (phi next_phi) pow (2) mean ()
compute Orthonormality 1osss
Cov = torch. (phi, phi.T) 10
I torch. (*Cov.size (),, callccee-oooddevieee
off_diag ~I.bool
13
orth_loss_diag
Cov. cov.diag() ()
14
lots__osssoffdll = Cov loff_diagll (2) mean ()
15
lrsh
woss_offdiag + erth_loss_ddar ddag
loss loss_loss
18
19 return loss
Listing 3: Pytorch code for Laplacian Eigenfunctions Lap loss
def computc_phi_losee (agent, obs, future_obs):
future_phi agent. feature_net (future_obs)
agent. mygent.mm_neetoooc (obs)
future_phi F.normalize (future_phi, dim=1) mu = F.normal (mu, dim-l)
logits = torch. e'nsum sd, td-> st', future_pai)
batch batch
(t size (), device=Logits.devi.
off off >I.bool
10
bggit_oof_diaa
logits[off_diaat reshape (logits. shape [0], logits.
[h] 1)
loss (cgits.diag() + torch. legsumepg (logits_off_diaee dim=1)
loss loss me
13
14
return loss
Listing 4: Pytorch code for the contrastive CL loss
def compute_phi. loss (agent, obs, action, next_obs)
phi agent. feature_net (obs)
next_phi phi agent. feature_net (next_obs)
36

---

Published as a conference paper at ICLR 2023
predicted action agent. agent.inverse inverse www.rreeddaaaaeeeeee.ce..ma (torch, (pphi, phi], dim=-1))
loss (action - prodictod_action) .pow (2) mean ()
return lcss
Listing Pytorch code of ICM loss
def compute_phi_losse (agent, obs, action, next_obs);
phi = agent. featureentte (obs)
predicted_next_tta = agent. forward_dynamicccoadorede (torch. ([phi, action], dim=-1))
loss (predicted_nextte
next_obs) pow (2) mean
return Ioss
Listing 6: Pytorch code for Trans loss
def compute_pa ppi_ loss (agent, obs, action, next_obs): phi = agent feature_net (obs) with _)__aad()))
next_phi agent tarettttrrtt_ttt
teaagget_fattre_neee (next_obs)
bpecicted_next_db = deretad.rrardddyama......o (toorc (torch.cat (lphi, action), dim=-1)) loss (predicted_next_ooo
next_phi.detab ()) -po (2) m)
I update target network for param, parget_paran in zip(agent. featurteeeaeeeeeeee pa...ttttttete()) agent. Barget_feature peteee_________eeeeeeeeeeeeeeeee parameters [eeature_eettpaareeetess(((
target_param data. copy (tau param. data
11
tau) target_parem. data)
12
13
return loss
Listing 7: Pytorch code for Latent loss
def compute _hcpe_ph________ loss (agent, obs)
phi agent feature_ntt (obs)
predicted_obs agent. decoder (phi)
loss (predicted_ote - obs) .po (2) mean
return csss
Listing 8: Pytorch code for AEnc loss
def compute_phi_lane (agent, obs, action, next_obs)
ph agent. feature_net (next_cbs)
agent mu net (torch. (lobs, action), dim=l))
F torch.einsume (dsd, to st", mu, phi)
I = torch. (ize of off_diap - ~I ~I. boob
Device=P.devices
loss 2 p.diag() () P(off_diag) .po (2) .maan ()
10
compute orthonormality loss Cov torch, morccl (phi, phi.T)
11
I = torch.eye (*Cov..sizi device=Cov.deviice
off_diat ~I.bool
14
orth
Cov.diag() .eea ()
15
orth
covioff .pw.ggl.owwtt mean ()
16
lrss = loss lot + crts_ diag
loss +- orth_loss
37

---

Published as a conference paper at ICLR 2023
return loss
20
Listing 9: Pytorch code for LRA-P loss
def compute_pin phi loss (agent, obs, action, next_oos, discount):
phi agent. net (next_cbs)
mu agent.mu_nnee (obs)
SR - torch.einsum/ (f sd, tå -> st', mu, phi)
with torch torch.no_grad()=
target_phi tanntttarggt_m... mu net (next_obs)
net (next obs)
target_m mu agent. target_$8 SR torch.einsum ("sd, td -> st", target_mu, target_phi.
10
torch. eycch.yyc(o (ASR, size wwvicc-SR.devics.
12
off_diag "I.bool() + () () (SR - discount *
13
loss
14
()) [off_diag] pow (2) .eaa ()
# compute orthonormality loss
17
Cov torch.natmul (phi, phi.T)
18
aeth_osss_diae = 2 Cov. Cov.diag() mean
19
orth_loss_offdiag - cov[off_diag] pow (2) mean ()
orth_loss orth_loss_offdial + orth_loss_diar
loss orth_loss
22
23
return lcss
Listing 10: Pytorch code for LRA-SR loss
38


Reviews:
[{"review_id": "zbw-QY39RL", "reviewer": "Reviewer_UcNw", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies zero-shot RL in the sense of \"RL that does not require optimization of a policy when presented with a task (reward model)\". It covers the literature on successor representations (SR) and forward-backward (FB) representation, proposes an unifying view and improved training losses. Then it empirically compares FB and many formulations of state features for SR, to assess which one provides (offline) representations which are suited for zero-shot deduction of (quasi-)optimal policies.", "summary_of_the_review": "Overall, this paper is a strong contribution to learning successor representations and measures. The key finding is the general applicability of FB representations as successor measures and Laplacian eigenfunctions as successor features. Along the way, the paper provides a very appreciable unified view of learning representations for zero-shot RL, with insightful new formulations for FB representations.", "strength_and_weaknesses": "Disclaimer: given the very short timeline for reviewing all assigned papers, and to my great frustration, I didn't have time to check the appendix. Still, I browsed through it and it seems correct and (very) interesting. I did check the rest of the paper in detail and am confident in my evaluation.\n\n### Strengths\n\nThis paper is very well written, easy to follow despite the technicality. One quality I find in the paper (which the authors do not claim) is that it provides a unified presentation of SR and FB. I believe this is incredibly useful for the community and is one important strength of such a paper. Because of this, I view Section 4 as a valuable contribution in itself, even though its contents may not be drastically novel (some of them actually are: the last paragraph of Section 4 on the link between FB and SF is very important, especially the argument explaining the avoidance of representation collapse).  \nThen, the formulations of Section 5 are useful and partially novel, especially those of Section 5.2 which are interesting reformulations of that of Touati & Ollivier (2021).  \nThe empirical evaluation is thorough, well-conducted and well-discussed.  \nThe main conclusion of the paper is that FB representations and Laplacian eigenfunctions for SR are more suited to zero-shot transfer. Limitations of the present evaluation are explicitly mentioned (although not really discussed). This conclusion, backed by convincing empirical evidence and relevant discussion is a nice contribution to the RL community.\n\n### Weaknesses\n\nThis paper has a catchy name, but the name (although funny and I appreciate the pun) does not necessarily reflect the contribution: it is not really about existence of zero-shot RL but rather about the ability of neural networks to learn successor representations or measures. In the interest of readers in a few years, I would suggest keeping the catchyness but maybe focusing more on the ability to learn successor representations.  \n\nAt the bottom of page 4, the authors present the fact that FB does not need state featurization as an advantage. This seems true from an abstract point of view, but I believe it deserves a bit more discussion. In very unbalanced datasets, or with very large variance transition models, or in the regime of very scarce data, fully relying on data to estimate occupancy measures boils down to discarding whatever expert prior knowledge might be available for state representations. In this case, the advantage of FB representations is less obvious and this begs for deeper evaluation (the discussion at the end of Section 7 somehow discards this question a bit too fast). \n\nThe AE features presented in Section 5.3 incur the risk of representation collapse in the sense that they might map to a set of linearly redundant features. From this perspective, I would have liked to have more insight as to the effective dimension of the learned feature space, for example as was done in the work of Lyle et al. (2022). There is a clear difference here since these features are not learned as solutions to a non-stationary regression problem, but still, the diversity of features could (should?) be made explicit.  \nIn the same line of thought, only the latent transition model incorporates a term which discourages representation collapse (with a BYOL-like auxiliary loss). Maybe the other AE features would have benefited from such a term encouraging feature diversity (like regularizing features towards fixed or slowly moving random projections of the input state).  \nLyle, C., Rowland, M., and Dabney, W. (2022). Understanding and preventing capacity loss in reinforcement learning. In 10th International Conference on Learning Representations.\n\nSomehow connected to the previous remark, I would have appreciated a discussion on the latent spaces learned. Beyond the suggestion to evaluate the effective dimension of features (previous paragraph), I am curious as to what a tSNE of the replay buffer projected in the latent space would look like. Alternatively, the set of singular values from a PCA on the same set would be informative about the ability of the learned representation to shatter points. Since, in the discrete action case, the policy is eventually a linear classifier in the representation space, this latent space could (should) be better illustrated and discussed.\n\nI slightly disagree with the idea that if the encoder-based losses learn nothing, it is because $(x,y)$ is already a good representation and they need to do nothing. They still have the freedom to learn bad (or collapsed) representations, which might be an alternative explanation. Also, in the same paragraph, it seems somehow counter-intuitive that SFs on (encoded) $(x,y)$ cannot learn the task: I am not convinced that planning requires \"specific representations\", I'd rather conclude that planning requires representations which preserve key information; this information is in the state $(x,y)$ but maybe not in its encoded version. The authors reiterate the claim that planning-specific features are necessary, in the paper's conclusion, and while I agree in the general case, I seem to fail to understand why the raw state is not a good planning specific feature since it \"contains\" the full state information necessary to build a policy. Is it because we need features which allow linear classification of actions? Maybe this deserves a clarification.\n\nThe paragraph describing how the scores of Figure 2 are computed is a bit unsettling. Since scores are normalized as a percentage of offline TD3, shouldn't the red dashed line stand at 100 on all figures? Is there a re-multiplication by the (average?) score of offline TD3?\n\nI did not find a mention of the experimental code. Given the amount of experiments, releasing the code, the results and the training logs seems a crucial pre-requisite for fair dissemination of these results.\n\n### Minor discussion elements \n\nNo need to address these in detail but I believe they could contain ideas for future work and I'm happy to share them and discuss them.\n\nI am always a bit skeptical with the inclusion of $\\gamma$ in the definition of an MDP. And even more skeptical when it comes to a reward-free MDP. This might sound a bit knitpicking but an MDP is primarily a stochastic process, whose definition does not include any notion of discount. Even the introduction of a reward model over the Markov chain conditioned by $s$ and $a$ does not require $\\gamma$. It is the definition of the discounted return which involves $\\gamma$ and this is quite independent of the MDP's dynamics or rewards: it is intrinsic to the criterion being maximized. Another criterion (average reward for instance) would not require $\\gamma$ and the definition of the successor measure would not involve it either, while the reward-free MDP remains the same. Besides me being picky, this opens the door to alternative FB schemes for other criteria than the expected discounted sum of rewards.\n\nThere seems to be a striking connection between $z_r$ with successor features and the expression of policy values in LSTD / LSPI. This might sound a bit naive (and there are clear differences too) but I'm not sure this connection was made explicit before in prior work on SR.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "t5yQ3c0atJ", "reviewer": "Reviewer_zaXC", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper conducts an empirical study of existing unsupervised pre-training methods for RL. Specifically, this paper focuses on successor features (SFs) and forward-backward representation (FB). The paper uses a variety of existing unsupervised representation learning methods as a way to provide state features for SFs. The empirical results on Unsupervised RL and ExORL benchmarks show that FB performs best. ", "summary_of_the_review": "Although this paper presents an interesting empirical study, the presentation of the paper needs to be significantly revised due to 1) the inconsistency between the problem definition and the experiment, and 2) the unconventional use of the term \"zero-shot\". In addition, the paper would benefit a lot from more in-depth analysis of existing feature learning methods. ", "strength_and_weaknesses": "[Strength]\n\n* Provides a nice survey of successor features, forward-backward representation, and unsupervised representation learning methods. \n\n[Weakness] \n\n* The definition of zero-shot RL in this paper is quite different from conventional zero-shot RL, which can be confusing.\n* Although the evaluation of many methods is valuable, there is not much analysis or insight from the experiments. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "XM4Uns6tGN", "reviewer": "Reviewer_uRSG", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper focuses on the empirical study of zero-shot RL, whose goal is to learn a set of policies and/or representations during training so that they can be adapted to solve unseen tasks in a zero-shot manner. Specifically, the authors focus on the successor features (SFs) framework and forward-backward (FB) representations, suggesting ten diverse choices for SF representations ($\\phi$) as well as an improved mechanism to train FB representations. They test the SF/FB methods on 13 tasks from URLB with ExORL replay buffers and find that FB representations show consistently superior performances on the benchmark.", "summary_of_the_review": "While there is a point (the distinction between FB and SFs) that I do not fully agree with the author's claim, I believe this paper is of high quality and will be helpful to the community. I thus recommend acceptance.\n\n(Post-rebuttal update) Given the additional results and clarification, I raised my score to 8.", "strength_and_weaknesses": "**[Strengths]**\n- To my knowledge, this is the first paper that systematically studies various choices for SF representations ($\\phi$), which would be especially useful for practitioners who try to use SFs without relying on hand-crafted representations.\n- The experiments are exhaustive and statistically sound (with 10 seeds).\n- The claims in the paper are theoretically well-supported.\n\n**[Weaknesses]**\n- One thing that I do not fully agree with the authors is the distinction between SFs and FB. For me, the FB representation can be interpreted as a special case of USFs [1] with the tasks being all possible goal-reaching tasks (ignoring the subtlety in continuous domains). The authors distinguish FB from SFs in that SFs require predefined features and/or tasks (i.e., $\\phi$), but I believe this is not accurate. For example, Ma et al. [1] use an end-to-end training scheme that learns a feature vector $\\phi$ and a (learnable) task vector $w(g)$ by minimizing both the scalar loss (as in Eq (6)) and the vectorized loss (as in the equation above Eq (6)). If we assume that the source tasks consist of all the goal-reaching tasks to reach any states (and also assume a discrete state space), $w(g)$ becomes $B(s)$ and $\\psi(s, a, z)$ becomes $F(s, a, z)$, and thus we exactly recover the FB objective. Hence, I believe the distinction between FB and SFs lies *not* in the way it automatically builds feature vectors ($\\phi$) but rather in the loss function it uses (especially Eq (8)), which can handle continuous spaces as well, unlike USFs with all the goal-reaching tasks.\n- Due to the reason above, I believe *learned* USFs (i.e., learnable $\\phi(s)$ and $w(g)$ with the Q loss and/or vectorized loss [1]) trained with implicit goal-reaching tasks (from the dataset) should also be included as baselines. Also, some strong statements like \"SFs must be provided with basic state features $\\phi$\" should be adjusted accordingly. I hope that the authors include detailed discussions and comparisons between FB and learned USFs.\n- Since this paper focuses on the empirical study of SF representations, I believe the code should be made publicly available and hope that the authors release their code during the discussion period.\n\n**[Additional questions]**\n- What makes $\\phi$ necessarily learn something in Eq (11) and (12)? It seems that $g$ or $f$ can capture everything with $\\phi$ being the identity function.\n- I didn't understand the following sentence in Sec 7: \"Any low-rank model of $P_\\pi$ will be poor; actually $P_\\pi$ is better modeled as Id − low-rank, thus approximating the Laplacian.\". Why does this amount to approximating the Laplacian? Also, it would be better to use parentheses around \"low-rank\" to avoid confusion.\n\n[1] Chen Ma, Dylan R. Ashley, Junfeng Wen, and Yoshua Bengio. Universal Successor Features for Transfer Reinforcement Learning. 2020.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "4LmelbE68x", "reviewer": "Reviewer_NRQ8", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Paper shows a systematic study of different strategies for approximate zero-shot RL including successor features (SFs) and forward-backward (FB) representations. It introduces the concepts clearly and discusses the advantages and limitations of each method in the two categories. Paper introduces new SF models and compares them in zero-shot RL tasks from the Unsupervised RL benchmark. ", "summary_of_the_review": "The authors study zero-shot RL that pretrained on large unlabeled datasets for different downstream tasks, they conduct a systematical extensive evaluation of various unsupervised RL methods, and reveal algebraic links between SFs, FB, contrastive learning, and spectral methods. Seems to me that the findings in this paper are important for future research. ", "strength_and_weaknesses": "Strength\n- This work systematically tested forward-backward representations and many new models of successor features on zero-shot tasks. The evaluation on the ExoRL dataset and unsupervised RL benchmark is very extensive and across multiple different qualities of datasets. It shows that FB representation performs best and consistently across the board in a zero-shot manner. \n- The writing is well-written and well-organized, the authors introduce prior work and compare them both in both methodology and experimental aspects. The discussion of limitations is thoughtful and contains interesting points. \n- This work systematically evaluates prior zero-shot RL methods in the same setting, and then proposes a novel and effective method that consistently outperforms prior works, reaching 85% of supervised RL performance with a good replay buffer, in a zero-shot manner.\n- The results are significant and will facilitate future research on unsupervised zero-shot RL. \n\nWeaknesses\n- The biggest weakness seems to be that it assumes there is offline diverse data. One one hand, this is a limitation since unlike NLP there is a large unlabeled dataset naturally available, in RL, collecting diverse dataset is still not yet solved. On the other hand, this work focuses on evaluating learning zero-shot RL from pre collected unlabeled dataset, which is an important research question and it’s reasonable to have such assumptions in order to make the evaluation easier. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%