Original: Does Zero-Shot Reinforcement Learning Exist?
Revised: Does Zero-Shot Reinforcement Learning Exist?

Original: Abstract
Revised: Abstract

Original: 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.
Revised: 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.

Original: This marks a shift from the reward-centric RL paradigm towards “controllable” agents that can follow arbitrary instructions in an environment.
Revised: This marks a shift from the reward-centric RL paradigm towards “controllable” agents that can follow arbitrary instructions in an environment.

Original: Current RL agents can solve families of related tasks at best, or require planning anew for each task.
Revised: Current RL agents can solve families of related tasks at best, or require planning anew for each task.

Original: Strategies for approximate zero-shot RL have been suggested using successor features (SFs)
Revised: Strategies for approximate zero-shot RL have been suggested using successor features (SFs)

Original: (Borsa et al., 2018) or forward-backward (FB) representations (Touati & Ollivier, 2021), but testing has been limited.
Revised: (Borsa et al., 2018) or forward-backward (FB) representations (Touati & Ollivier, 2021), but testing has been limited.

Original: 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).
Revised: 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).

Original: To disentangle universal representation learning from exploration, we work in an ofﬂine setting and repeat the tests on several existing replay buffers.
Revised: To disentangle universal representation learning from exploration, we work in an ofﬂine setting and repeat the tests on several existing replay buffers.

Original: SFs appear to suffer from the choice of the elementary state features.
Revised: SFs appear to suffer from the choice of the elementary state features.

Original: SFs with
Revised: SFs with

Original: 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.
Revised: 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.

Original: In contrast, FB representations jointly learn the elementary and successor features from a single, principled criterion.
Revised: In contrast, FB representations jointly learn the elementary and successor features from a single, principled criterion.

Original: They perform best and consistently across the board, reaching 85% of supervised
Revised: They perform best and consistently across the board, reaching 85% of supervised

Original: RL performance with a good replay buffer, in a zero-shot manner.
Revised: RL performance with a good replay buffer, in a zero-shot manner.

Original:  Introduction
Revised:  Introduction

Original: For breadth of applications, reinforcement learning (RL) lags behind other ﬁelds 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).
Revised: For breadth of applications, reinforcement learning (RL) lags behind other ﬁelds 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).

Original: The RL paradigm itself may be in part to blame: RL agents are usually trained for only one reward function or a small family of related rewards.
Revised: The RL paradigm itself may be in part to blame: RL agents are usually trained for only one reward function or a small family of related rewards.

Original: 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 to do, reacting instantly to such commands as “fetch this object while avoiding that area”.
Revised: 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 to do, reacting instantly to such commands as “fetch this object while avoiding that area”.

Original: The promise of zero-shot RL is to train without rewards or tasks, yet immediately perform well on any reward function given at test time, with no extra training, planning, or ﬁnetuning, and only a minimal amount of extra computation to process a task description (Section 2 gives a more precise deﬁnition).
Revised: The promise of zero-shot RL is to train without rewards or tasks, yet immediately perform well on any reward function given at test time, with no extra training, planning, or ﬁnetuning, and only a minimal amount of extra computation to process a task description (Section 2 gives a more precise deﬁnition).

Original: How far away are such zero-shot agents?
Revised: How far away are such zero-shot agents?

Original: In the RL paradigm, a new task (reward function) means re-training the agent from scratch, and providing many reward samples.
Revised: In the RL paradigm, a new task (reward function) means re-training the agent from scratch, and providing many reward samples.

Original: Model-based RL trains a reward-free, task-independent world model, but still requires heavy planning when a new reward function is speciﬁed (e.g, Chua et al., 2018; Moerland et al., 2020).
Revised: Model-based RL trains a reward-free, task-independent world model, but still requires heavy planning when a new reward function is speciﬁed (e.g, Chua et al., 2018; Moerland et al., 2020).

Original: Model-free RL is reward-centric from start, and produces specialized agents.
Revised: Model-free RL is reward-centric from start, and produces specialized agents.

Original: Multi-task agents generalize within a family of related tasks only.
Revised: Multi-task agents generalize within a family of related tasks only.

Original: Reward-free, unsupervised skill pre-training (e.g, Eysenbach et al., 2018) still requires substantial downstream task adaptation, such as training a hierarchical controller.
Revised: Reward-free, unsupervised skill pre-training (e.g, Eysenbach et al., 2018) still requires substantial downstream task adaptation, such as training a hierarchical controller.

Original: Is zero-shot RL possible?
Revised: Is zero-shot RL possible?

Original: If one ignores practicality, zero-shot RL is easy: make a list of all possible rewards up to precision 𝜀 , then pre-learn all the associated optimal policies.
Revised: If one ignores practicality, zero-shot RL is easy: make a list of all possible rewards up to precision 𝜀 , then pre-learn all the associated optimal policies.

Original: Scalable zero-shot RL must somehow exploit the relationships between policies for all tasks.
Revised: Scalable zero-shot RL must somehow exploit the relationships between policies for all tasks.

Original: Learning to go from 𝑎 to 𝑐 is not independent from going from 𝑎 to 𝑏 and 𝑏 to 𝑐 , and this produces rich, exploitable algebraic
Revised: Learning to go from 𝑎 to 𝑐 is not independent from going from 𝑎 to 𝑏 and 𝑏 to 𝑐 , and this produces rich, exploitable algebraic

Original: Suggested strategies for generic zero-shot RL so far have used successor representations (Dayan, 1993), under two forms: successor features (SFs)
Revised: Suggested strategies for generic zero-shot RL so far have used successor representations (Dayan, 1993), under two forms: successor features (SFs)

Original: (Barreto et al., 2017) as in (Borsa et al., 2018;
Revised: (Barreto et al., 2017) as in (Borsa et al., 2018;

Original: Hansen et al., 2019; Liu & Abbeel, 2021); and forward-backward (FB) representations (Touati &
Revised: Hansen et al., 2019; Liu & Abbeel, 2021); and forward-backward (FB) representations (Touati &

Original: Ollivier, 2021).
Revised: Ollivier, 2021).

Original: Both SFs and FB lie in between model-free and model-based RL, by predicting features of future states, or summarizing long-term state-state relationships.
Revised: Both SFs and FB lie in between model-free and model-based RL, by predicting features of future states, or summarizing long-term state-state relationships.

Original: Like model-based approaches, they decouple the dynamics of the environment from the reward function.
Revised: Like model-based approaches, they decouple the dynamics of the environment from the reward function.

Original: Contrary to world models, they require neither planning at test time nor a generative model of states or trajectories.
Revised: Contrary to world models, they require neither planning at test time nor a generative model of states or trajectories.

Original: Yet SFs heavily depend on a choice of basic state features.
Revised: Yet SFs heavily depend on a choice of basic state features.

Original: To get a full zero-shot RL algorithm, a representation learning method must provide those.
Revised: To get a full zero-shot RL algorithm, a representation learning method must provide those.

Original: 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.
Revised: 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.

Original: Meanwhile, FB is a standalone method with no task prior and good theoretical backing, but testing has been limited to goal-reaching in a few environments.
Revised: Meanwhile, FB is a standalone method with no task prior and good theoretical backing, but testing has been limited to goal-reaching in a few environments.

Original: Here:
Revised: Here:

Original: • We systematically assess SFs and FB for zero-shot RL, including many new models of SF basic features, and improved FB loss functions.
Revised: • We systematically assess SFs and FB for zero-shot RL, including many new models of SF basic features, and improved FB loss functions.

Original: 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.
Revised: 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.

Original: • We systematically study the inﬂuence of basic features for SFs, by testing SFs on features from ten RL representation learning methods.
Revised: • We systematically study the inﬂuence of basic features for SFs, by testing SFs on features from ten RL representation learning methods.

Original: such as latent next state prediction, inverse curiosity module, contrastive learning, diversity, various spectral decompositions...
Revised: such as latent next state prediction, inverse curiosity module, contrastive learning, diversity, various spectral decompositions...

Original: • We expose new mathematical links between SFs, FB, and other representations in RL.
Revised: • We expose new mathematical links between SFs, FB, and other representations in RL.

Original: • We discuss the implicit assumptions and limitations behind zero-shot RL approaches.
Revised: • We discuss the implicit assumptions and limitations behind zero-shot RL approaches.

Original:  Problem and Notation; Deﬁning Zero-Shot RL
Revised:  Problem and Notation; Deﬁning Zero-Shot RL

Original: Let ℳ =
Revised: Let ℳ =

Original: ( 𝑆, 𝐴, 𝑃, 𝛾 ) be a reward-free Markov decision process (MDP) with state space 𝑆 , action space 𝐴 , transition probabilities 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) from state 𝑠 to 𝑠 ′ given action 𝑎 , and discount factor 0 <
Revised: ( 𝑆, 𝐴, 𝑃, 𝛾 ) be a reward-free Markov decision process (MDP) with state space 𝑆 , action space 𝐴 , transition probabilities 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) from state 𝑠 to 𝑠 ′ given action 𝑎 , and discount factor 0 <

Original: 𝛾 < 1 (Sutton & Barto, 2018).
Revised: 𝛾 < 1 (Sutton & Barto, 2018).

Original: If 𝑆 and 𝐴 are ﬁnite, 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) can be viewed as a stochastic matrix 𝑃 𝑠𝑎𝑠 ′ ∈ R ( | 𝑆 |×| 𝐴 | ) ×| 𝑆 | ; in general, for each ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 , 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) is a probability measure on 𝑠 ′ ∈ 𝑆 .
Revised: If 𝑆 and 𝐴 are ﬁnite, 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 ) can be viewed as a stochastic matrix 𝑃 𝑠𝑎𝑠 ′ ∈ R ( | 𝑆 |×| 𝐴 | ) ×| 𝑆 | ; in general, for each ( 𝑠, 𝑎 ) ∈ 𝑆 × 𝐴 , 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) is a probability measure on 𝑠 ′ ∈ 𝑆 .

Original: The notation 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) covers all cases.
Revised: The notation 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) covers all cases.

Original: Given ( 𝑠 0 , 𝑎 0 ) ∈
Revised: Given ( 𝑠 0 , 𝑎 0 ) ∈

Original: 𝑆 × 𝐴 and a policy 𝜋 : 𝑆 → Prob( 𝐴 ) , we denote Pr( ·| 𝑠 0 , 𝑎 0 , 𝜋 ) and E [ ·| 𝑠 0 , 𝑎 0 , 𝜋 ] the probabilities and expectations under state-action sequences ( 𝑠 𝑡 , 𝑎 𝑡 ) 𝑡 ≥ 0 starting at ( 𝑠 0 , 𝑎 0 ) and following policy 𝜋 in the environment, deﬁned by sampling 𝑠 𝑡 ∼ 𝑃 (d 𝑠 𝑡 | 𝑠 𝑡 − 1 , 𝑎 𝑡 − 1 ) and 𝑎 𝑡 ∼ 𝜋 (d 𝑎 𝑡 | 𝑠 𝑡 ) .
Revised: 𝑆 × 𝐴 and a policy 𝜋 : 𝑆 → Prob( 𝐴 ) , we denote Pr( ·| 𝑠 0 , 𝑎 0 , 𝜋 ) and E [ ·| 𝑠 0 , 𝑎 0 , 𝜋 ] the probabilities and expectations under state-action sequences ( 𝑠 𝑡 , 𝑎 𝑡 ) 𝑡 ≥ 0 starting at ( 𝑠 0 , 𝑎 0 ) and following policy 𝜋 in the environment, deﬁned by sampling 𝑠 𝑡 ∼ 𝑃 (d 𝑠 𝑡 | 𝑠 𝑡 − 1 , 𝑎 𝑡 − 1 ) and 𝑎 𝑡 ∼ 𝜋 (d 𝑎 𝑡 | 𝑠 𝑡 ) .

Original: We deﬁne 𝑃 𝜋 (d 𝑠 ′ , d 𝑎 ′ | 𝑠, 𝑎 ) :
Revised: We deﬁne 𝑃 𝜋 (d 𝑠 ′ , d 𝑎 ′ | 𝑠, 𝑎 ) :

Original: = 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) 𝜋 (d 𝑎 ′ | 𝑠 ′ ) and 𝑃 𝜋 (d 𝑠 ′ | 𝑠 ) :
Revised: = 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) 𝜋 (d 𝑎 ′ | 𝑠 ′ ) and 𝑃 𝜋 (d 𝑠 ′ | 𝑠 ) :

Original: = ∫︀ 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) 𝜋 (d 𝑎 | 𝑠 ) , the state-action transition probabilities and state transition probabilities induced by 𝜋 .
Revised: = ∫︀ 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) 𝜋 (d 𝑎 | 𝑠 ) , the state-action transition probabilities and state transition probabilities induced by 𝜋 .

Original: Given a reward function 𝑟 : 𝑆 → R , the 𝑄 -function of 𝜋 for 𝑟 is 𝑄 𝜋𝑟 ( 𝑠 0 , 𝑎 0 ) :
Revised: Given a reward function 𝑟 : 𝑆 → R , the 𝑄 -function of 𝜋 for 𝑟 is 𝑄 𝜋𝑟 ( 𝑠 0 , 𝑎 0 ) :

Original: = ∑︀ 𝑡 ≥ 0 𝛾 𝑡 E [ 𝑟 ( 𝑠 𝑡 +1 ) | 𝑠 0 , 𝑎 0 , 𝜋 ] .
Revised: = ∑︀ 𝑡 ≥ 0 𝛾 𝑡 E [ 𝑟 ( 𝑠 𝑡 +1 ) | 𝑠 0 , 𝑎 0 , 𝜋 ] .

Original: For simplicity, we assume the reward 𝑟 depends only on the next state 𝑠 𝑡 +1 instead on the full triplet ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) , but this is not essential.
Revised: For simplicity, we assume the reward 𝑟 depends only on the next state 𝑠 𝑡 +1 instead on the full triplet ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) , but this is not essential.

Original: We focus on ofﬂine unsupervised RL, where the agent cannot interact with the environment.
Revised: We focus on ofﬂine unsupervised RL, where the agent cannot interact with the environment.

Original: The agent only has access to a static dataset of logged reward-free transitions in the environment, 𝒟 = { ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑠 ′ 𝑖 ) } 𝑖 ∈ℐ with 𝑠 ′
Revised: The agent only has access to a static dataset of logged reward-free transitions in the environment, 𝒟 = { ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑠 ′ 𝑖 ) } 𝑖 ∈ℐ with 𝑠 ′

Original: 𝑖 ∼ 𝑃 (d 𝑠 ′ 𝑖
Revised: 𝑖 ∼ 𝑃 (d 𝑠 ′ 𝑖

Original: | 𝑠 𝑖 , 𝑎 𝑖 ) .
Revised: | 𝑠 𝑖 , 𝑎 𝑖 ) .

Original: These can come from any exploration method or methods.
Revised: These can come from any exploration method or methods.

Original: The ofﬂine setting disentangles the effects of the exploration method and representation and policy learning: we test each zero-shot method on several training datasets from several exploration methods.
Revised: The ofﬂine setting disentangles the effects of the exploration method and representation and policy learning: we test each zero-shot method on several training datasets from several exploration methods.

Original: We denote by 𝜌 (d 𝑠 ) and 𝜌 (d 𝑠, d 𝑎 ) the (unknown) marginal distribution of states and state-actions in the dataset 𝒟 .
Revised: We denote by 𝜌 (d 𝑠 ) and 𝜌 (d 𝑠, d 𝑎 ) the (unknown) marginal distribution of states and state-actions in the dataset 𝒟 .

Original: We use both E 𝑠 ∼𝒟 [ · ] and E 𝑠 ∼ 𝜌 [ · ] for expectations under the training distribution.
Revised: We use both E 𝑠 ∼𝒟 [ · ] and E 𝑠 ∼ 𝜌 [ · ] for expectations under the training distribution.

Original: Zero-shot RL: problem statement.
Revised: Zero-shot RL: problem statement.

Original: The goal of zero-shot RL is to compute a compact representation ℰ of the environment by observing samples of reward-free transitions ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) in this environment.
Revised: The goal of zero-shot RL is to compute a compact representation ℰ of the environment by observing samples of reward-free transitions ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) in this environment.

Original: Once a reward function is speciﬁed later, the agent must use ℰ to immediately produce a good policy, via only elementary computations without any further planning or learning.
Revised: Once a reward function is speciﬁed later, the agent must use ℰ to immediately produce a good policy, via only elementary computations without any further planning or learning.

Original: Ideally, for any downstream task, the performance of the returned policy should be close to the performance of a supervised RL baseline trained on the same dataset labeled with the rewards for that task.
Revised: Ideally, for any downstream task, the performance of the returned policy should be close to the performance of a supervised RL baseline trained on the same dataset labeled with the rewards for that task.

Original: Reward functions may be speciﬁed at test time either as a relatively small set of reward samples ( 𝑠 𝑖 , 𝑟 𝑖 ) , or as an explicit function 𝑠 ↦→ 𝑟 ( 𝑠 ) (such as 1 at a known goal state and 0 elsewhere).
Revised: Reward functions may be speciﬁed at test time either as a relatively small set of reward samples ( 𝑠 𝑖 , 𝑟 𝑖 ) , or as an explicit function 𝑠 ↦→ 𝑟 ( 𝑠 ) (such as 1 at a known goal state and 0 elsewhere).

Original: The method will be few-shot, zero-planning in the ﬁrst case, and truly zero-shot in the second case.
Revised: The method will be few-shot, zero-planning in the ﬁrst case, and truly zero-shot in the second case.

Original:  Related work
Revised:  Related work

Original: [Table] [Table]
Revised: [Table] [Table]

Original: requires unsupervised learning and the absence of any planning or ﬁne-tuning at test time.
Revised: requires unsupervised learning and the absence of any planning or ﬁne-tuning at test time.

Original: The proposed strategies for zero-shot RL discussed in Section 1 ultimately derive from successor representations (Dayan, 1993) in ﬁnite spaces.
Revised: The proposed strategies for zero-shot RL discussed in Section 1 ultimately derive from successor representations (Dayan, 1993) in ﬁnite spaces.

Original: In continuous spaces, starting with a ﬁnite number of features 𝜙 , successor features can be used to produce policies within a family of tasks directly related to 𝜙 (Barreto et al., 2017; Borsa et al., 2018; Zhang et al., 2017; Grimm et al., 2019), often using hand-crafted 𝜙 or learning 𝜙 that best linearize training rewards.
Revised: In continuous spaces, starting with a ﬁnite number of features 𝜙 , successor features can be used to produce policies within a family of tasks directly related to 𝜙 (Barreto et al., 2017; Borsa et al., 2018; Zhang et al., 2017; Grimm et al., 2019), often using hand-crafted 𝜙 or learning 𝜙 that best linearize training rewards.

Original: 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 et al., 2016).
Revised: 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 et al., 2016).

Original: We include APS among our baselines, as well as many new criteria to build 𝜙 automatically.
Revised: We include APS among our baselines, as well as many new criteria to build 𝜙 automatically.

Original: Successor measures (Blier et al., 2021) avoid the need for 𝜙 by directly learning models of the distribution of future states: doing this for various policies yields a candidate zero-shot RL method, forward-backward representations (Touati & Ollivier, 2021), which has been tested for goal-reaching in a few environments with discrete actions.
Revised: Successor measures (Blier et al., 2021) avoid the need for 𝜙 by directly learning models of the distribution of future states: doing this for various policies yields a candidate zero-shot RL method, forward-backward representations (Touati & Ollivier, 2021), which has been tested for goal-reaching in a few environments with discrete actions.

Original: FB uses a low-rank model of long-term state-state relationships reminiscent of the state-goal factorization from Schaul et al.
Revised: FB uses a low-rank model of long-term state-state relationships reminiscent of the state-goal factorization from Schaul et al.

Original: Model-based RL (surveyed in Moerland et al.
Revised: Model-based RL (surveyed in Moerland et al.

Original: (2020)) misses the zero-planning requirement of zero-shot RL.
Revised: (2020)) misses the zero-planning requirement of zero-shot RL.

Original: 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).
Revised: 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).

Original: 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; Andrychowicz et al., 2017).
Revised: 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; Andrychowicz et al., 2017).

Original: A parametric family of tasks must be deﬁned in advance (e.g., reaching arbitrary goal states).
Revised: A parametric family of tasks must be deﬁned in advance (e.g., reaching arbitrary goal states).

Original: New rewards cannot be set a posteriori: for example, a goal-state-oriented method cannot handle dense rewards.
Revised: New rewards cannot be set a posteriori: for example, a goal-state-oriented method cannot handle dense rewards.

Original: Zero-shot task transfer methods learn on tasks and can transfer to related tasks only (e.g, Oh et al., 2017; Sohn et 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.
Revised: Zero-shot task transfer methods learn on tasks and can transfer to related tasks only (e.g, Oh et al., 2017; Sohn et 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.

Original: Instead, we aim at not having any predeﬁned family of tasks.
Revised: Instead, we aim at not having any predeﬁned family of tasks.

Original: Unsupervised skill and option discovery methods, based for instance on diversity (Eysenbach et al., 2018; Gregor et al., 2016) or eigenoptions (Machado et al., 2017) can learn a variety of behaviors without rewards.
Revised: Unsupervised skill and option discovery methods, based for instance on diversity (Eysenbach et al., 2018; Gregor et al., 2016) or eigenoptions (Machado et al., 2017) can learn a variety of behaviors without rewards.

Original: Downstream tasks require learning a hierarchical controller to combine the right skills or options for each task.
Revised: Downstream tasks require learning a hierarchical controller to combine the right skills or options for each task.

Original: Directly using unmodiﬁed skills has limited performance without heavy ﬁnetuning (Eysenbach et al., 2018).
Revised: Directly using unmodiﬁed skills has limited performance without heavy ﬁnetuning (Eysenbach et al., 2018).

Original: Still, these methods can speed up downstream learning.
Revised: Still, these methods can speed up downstream learning.

Original: The unsupervised aspect of some of these methods (including DIAYN and APS) has been disputed, because training still used end-of-trajectory signals, which are directly correlated to the downstream task in some common environments: without this signal, results drop sharply (Laskin et al., 2022).
Revised: The unsupervised aspect of some of these methods (including DIAYN and APS) has been disputed, because training still used end-of-trajectory signals, which are directly correlated to the downstream task in some common environments: without this signal, results drop sharply (Laskin et al., 2022).

Original:  Successor Representations and Zero-Shot RL
Revised:  Successor Representations and Zero-Shot RL

Original: For a ﬁnite MDP, the successor representation (SR) (Dayan, 1993) 𝑀 𝜋 ( 𝑠 0 , 𝑎 0 ) of a state-action pair ( 𝑠 0 , 𝑎 0 ) under a policy 𝜋 , is deﬁned as the discounted sum of future occurrences of each state:
Revised: For a ﬁnite MDP, the successor representation (SR) (Dayan, 1993) 𝑀 𝜋 ( 𝑠 0 , 𝑎 0 ) of a state-action pair ( 𝑠 0 , 𝑎 0 ) under a policy 𝜋 , is deﬁned as the discounted sum of future occurrences of each state:

Original: [Equation]
Revised: [Equation]

Original: In matrix form, SRs can be written as 𝑀 𝜋 = 𝑃 ∑︀ 𝑡 ≥ 

Revised: In matrix form, SRs can be written as 𝑀 𝜋 = 𝑃 ∑︀ 𝑡 ≥

Original: 𝛾 𝑡 𝑃 𝑡𝜋 = 𝑃 (Id − 𝛾𝑃 𝜋 ) − 1 , where 𝑃 𝜋 is the state transition probability.
Revised: 𝛾 𝑡 𝑃 𝑡𝜋 = 𝑃 (Id − 𝛾𝑃 𝜋 ) − 1 , where 𝑃 𝜋 is the state transition probability.

Original: 𝑀 𝜋 satisﬁes the matrix Bellman equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 .
Revised: 𝑀 𝜋 satisﬁes the matrix Bellman equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 .

Original: Importantly, SRs disentangle the dynamics of the MDP and the reward function: for any reward 𝑟 and policy 𝜋 , the 𝑄 -function can be expressed linearly as 𝑄 𝜋𝑟 =
Revised: Importantly, SRs disentangle the dynamics of the MDP and the reward function: for any reward 𝑟 and policy 𝜋 , the 𝑄 -function can be expressed linearly as 𝑄 𝜋𝑟 =

Original: 𝑀 𝜋 𝑟 .
Revised: 𝑀 𝜋 𝑟 .

Original: Successor features and successor measures.
Revised: Successor features and successor measures.

Original: Successor features (SFs)
Revised: Successor features (SFs)

Original: (Barreto et al., 2017) extend
Revised: (Barreto et al., 2017) extend

Original: SR to continous MDPs by ﬁrst assuming we are given a basic feature map 𝜙 : 𝑆 → R 𝑑 that embeds states into 𝑑 -dimensional space, and deﬁning the expected discounted sum of future state features:
Revised: SR to continous MDPs by ﬁrst assuming we are given a basic feature map 𝜙 : 𝑆 → R 𝑑 that embeds states into 𝑑 -dimensional space, and deﬁning the expected discounted sum of future state features:

Original: [Equation]
Revised: [Equation]

Original: SFs have been introduced to make SRs compatible with function approximation.
Revised: SFs have been introduced to make SRs compatible with function approximation.

Original: For a ﬁnite MDP, the original deﬁnition (1) is recovered by letting 𝜙 be a one-hot state encoding into R | 𝑆 | .
Revised: For a ﬁnite MDP, the original deﬁnition (1) is recovered by letting 𝜙 be a one-hot state encoding into R | 𝑆 | .

Original: Alternatively, successor measures (SMs) (Blier et al., 2021) extend SRs to continuous spaces by treating the distribution of future visited states as a measure 𝑀 𝜋 over the state space 𝑆 ,
Revised: Alternatively, successor measures (SMs) (Blier et al., 2021) extend SRs to continuous spaces by treating the distribution of future visited states as a measure 𝑀 𝜋 over the state space 𝑆 ,

Original: [Equation]
Revised: [Equation]

Original: SFs and SMs are related: by construction, 𝜓 𝜋 ( 𝑠 0 ,
Revised: SFs and SMs are related: by construction, 𝜓 𝜋 ( 𝑠 0 ,

Original: [Equation]
Revised: [Equation]

Original: Zero-shot RL from successor features and forward-backward representations.
Revised: Zero-shot RL from successor features and forward-backward representations.

Original: Successor representations provide a generic framework for zero-shot RL, by learning to represent the relationship between reward functions and 𝑄 -functions, as encoded in 𝑀 𝜋 .
Revised: Successor representations provide a generic framework for zero-shot RL, by learning to represent the relationship between reward functions and 𝑄 -functions, as encoded in 𝑀 𝜋 .

Original: Given a basic feature map 𝜙 : 𝑆 → R 𝑑 to be learned via another criterion, universal SFs (Borsa et al., 2018) learn the successor features of a particular family of policies 𝜋 𝑧 for 𝑧 ∈ R 𝑑 ,
Revised: Given a basic feature map 𝜙 : 𝑆 → R 𝑑 to be learned via another criterion, universal SFs (Borsa et al., 2018) learn the successor features of a particular family of policies 𝜋 𝑧 for 𝑧 ∈ R 𝑑 ,

Original: [Equation]
Revised: [Equation]

Original: Once a reward function 𝑟 is revealed, we use a few reward samples or explicit knowledge of the function 𝑟 to perform a linear regression of 𝑟 onto the features 𝜙 .
Revised: Once a reward function 𝑟 is revealed, we use a few reward samples or explicit knowledge of the function 𝑟 to perform a linear regression of 𝑟 onto the features 𝜙 .

Original: Namely, we estimate 𝑧 𝑟 :=
Revised: Namely, we estimate 𝑧 𝑟 :=

Original: arg min
Revised: arg min

Original: 𝑧 E 𝑠 ∼ 𝜌
Revised: 𝑧 E 𝑠 ∼ 𝜌

Original: [( 𝑟 ( 𝑠 ) − 𝜙 ( 𝑠 ) ⊤ 𝑧 ) 2 ]
Revised: [( 𝑟 ( 𝑠 ) − 𝜙 ( 𝑠 ) ⊤ 𝑧 ) 2 ]

Original: =
Revised: =

Original: E 𝜌 [ 𝜙𝜙 ⊤ ] − 1 E 𝜌 [ 𝜙𝑟 ] .
Revised: E 𝜌 [ 𝜙𝜙 ⊤ ] − 1 E 𝜌 [ 𝜙𝑟 ] .

Original: Then we return the policy 𝜋 𝑧 𝑟 .
Revised: Then we return the policy 𝜋 𝑧 𝑟 .

Original: This policy is guaranteed to be optimal for all rewards in the linear span of the features 𝜙 :
Revised: This policy is guaranteed to be optimal for all rewards in the linear span of the features 𝜙 :

Original: Theorem 1 (Borsa et al.
Revised: Theorem 1 (Borsa et al.

Original: Assume that (4) holds.
Revised: Assume that (4) holds.

Original: Assume there exists a weight 𝑤 ∈ R 𝑑 such that 𝑟 ( 𝑠 ) =
Revised: Assume there exists a weight 𝑤 ∈ R 𝑑 such that 𝑟 ( 𝑠 ) =

Original: 𝜙 ( 𝑠 ) ⊤ 𝑤, ∀ 𝑠 ∈ 𝑆 .
Revised: 𝜙 ( 𝑠 ) ⊤ 𝑤, ∀ 𝑠 ∈ 𝑆 .

Original: Then 𝑧 𝑟
Revised: Then 𝑧 𝑟

Original: = 𝑤 , and 𝜋 𝑧 𝑟 is the optimal policy for reward 𝑟 .
Revised: = 𝑤 , and 𝜋 𝑧 𝑟 is the optimal policy for reward 𝑟 .

Original: Forward-backward (FB) representations (Touati & Ollivier, 2021) apply a similar idea to a ﬁnite-rank model of successor measures.
Revised: Forward-backward (FB) representations (Touati & Ollivier, 2021) apply a similar idea to a ﬁnite-rank model of successor measures.

Original: They look for representations 𝐹 : 𝑆 × 𝐴 × R 𝑑
Revised: They look for representations 𝐹 : 𝑆 × 𝐴 × R 𝑑

Original: → R 𝑑 and 𝐵 : 𝑆 → R 𝑑 such that the long-term transition probabilities 𝑀 𝜋 𝑧 in (3) decompose as
Revised: → R 𝑑 and 𝐵 : 𝑆 → R 𝑑 such that the long-term transition probabilities 𝑀 𝜋 𝑧 in (3) decompose as

Original: [Equation]
Revised: [Equation]

Original: In a ﬁnite space, the ﬁrst equation rewrites as the matrix decomposition 𝑀 𝜋 𝑧 =
Revised: In a ﬁnite space, the ﬁrst equation rewrites as the matrix decomposition 𝑀 𝜋 𝑧 =

Original: 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) .
Revised: 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) .

Original: Once a reward function 𝑟 is revealed, we estimate 𝑧 𝑟 :=
Revised: Once a reward function 𝑟 is revealed, we estimate 𝑧 𝑟 :=

Original: E
Revised: E

Original: 𝑠 ∼ 𝜌
Revised: 𝑠 ∼ 𝜌

Original: [ 𝑟 ( 𝑠 ) 𝐵 ( 𝑠 )]
Revised: [ 𝑟 ( 𝑠 ) 𝐵 ( 𝑠 )]

Original: from a few reward samples or from explicit knowledge of the function 𝑟 (e.g. 𝑧 𝑟 =
Revised: from a few reward samples or from explicit knowledge of the function 𝑟 (e.g. 𝑧 𝑟 =

Original: 𝐵 ( 𝑠 ) to reach 𝑠 ).
Revised: 𝐵 ( 𝑠 ) to reach 𝑠 ).

Original: Then we return the policy 𝜋 𝑧 𝑟 .
Revised: Then we return the policy 𝜋 𝑧 𝑟 .

Original: If the approximation (5) holds, this policy is guaranteed to be optimal for any reward function:
Revised: If the approximation (5) holds, this policy is guaranteed to be optimal for any reward function:

Original: Theorem 2 (Touati & Ollivier (2021)) .
Revised: Theorem 2 (Touati & Ollivier (2021)) .

Original: Assume that (5) holds.
Revised: Assume that (5) holds.

Original: Then for any reward function 𝑟 , the policy 𝜋 𝑧 𝑟 is optimal for 𝑟 , with optimal 𝑄 -function 𝑄 ⋆𝑟
Revised: Then for any reward function 𝑟 , the policy 𝜋 𝑧 𝑟 is optimal for 𝑟 , with optimal 𝑄 -function 𝑄 ⋆𝑟

Original: =
Revised: =

Original: 𝐹 ( 𝑠, 𝑎, 𝑧 𝑟 ) ⊤ 𝑧 𝑟 .
Revised: 𝐹 ( 𝑠, 𝑎, 𝑧 𝑟 ) ⊤ 𝑧 𝑟 .

Original: For completeness, we sketch the proofs of Theorems 1–2 in Appendix B. Importantly, both theorems are compatible with approximation: approximate solutions provide approximately optimal policies.
Revised: 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.

Original: Connections between SFs and FB.
Revised: Connections between SFs and FB.

Original: A ﬁrst difference between SFs and FB is that SFs must be provided with basic features 𝜙 .
Revised: A ﬁrst difference between SFs and FB is that SFs must be provided with basic features 𝜙 .

Original: The best 𝜙 is such that the reward functions of the downstream tasks are linear in 𝜙 .
Revised: The best 𝜙 is such that the reward functions of the downstream tasks are linear in 𝜙 .

Original: But for unsupervised training without prior task knowledge, an external criterion is needed to learn 𝜙 .
Revised: But for unsupervised training without prior task knowledge, an external criterion is needed to learn 𝜙 .

Original: We test a series of such criteria below.
Revised: We test a series of such criteria below.

Original: In contrast, FB uses a single criterion, avoiding the need for state featurization by learning a model of state occupancy.
Revised: In contrast, FB uses a single criterion, avoiding the need for state featurization by learning a model of state occupancy.

Original: Second, SFs only cover rewards in the linear span of 𝜙 , while FB apparently covers any reward.
Revised: Second, SFs only cover rewards in the linear span of 𝜙 , while FB apparently covers any reward.

Original: But this difference is not as stark as it looks: exactly solving the FB equation (5) in continuous spaces requires 𝑑 = ∞ , and for ﬁnite 𝑑 , the policies will only be optimal for rewards in the linear span of 𝐵 (Touati & Ollivier, 2021).
Revised: But this difference is not as stark as it looks: exactly solving the FB equation (5) in continuous spaces requires 𝑑 = ∞ , and for ﬁnite 𝑑 , the policies will only be optimal for rewards in the linear span of 𝐵 (Touati & Ollivier, 2021).

Original: Thus, in both cases, policies are exactly optimal only for a 𝑑 -dimensional family of rewards.
Revised: Thus, in both cases, policies are exactly optimal only for a 𝑑 -dimensional family of rewards.

Original: Still, FB can use an arbitrary large 𝑑 without any additional input or criterion.
Revised: Still, FB can use an arbitrary large 𝑑 without any additional input or criterion.

Original: FB representations are related to successor features: the FB deﬁnition (5) implies that 𝜓 ( 𝑠, 𝑎, 𝑧 ) :
Revised: FB representations are related to successor features: the FB deﬁnition (5) implies that 𝜓 ( 𝑠, 𝑎, 𝑧 ) :

Original: =
Revised: =

Original: 𝐹 ( 𝑠, 𝑎, 𝑧 ) are the successor features of 𝜙 ( 𝑠 ) :
Revised: 𝐹 ( 𝑠, 𝑎, 𝑧 ) are the successor features of 𝜙 ( 𝑠 ) :

Original: = ( E 𝜌 𝐵𝐵 ⊤ ) − 1 𝐵 ( 𝑠 ) .
Revised: = ( E 𝜌 𝐵𝐵 ⊤ ) − 1 𝐵 ( 𝑠 ) .

Original: This follows from multiplying (3) and (5) by 𝐵 ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) − 1 on the right, and integrating over 𝑠 ′ ∼ 𝜌 .
Revised: This follows from multiplying (3) and (5) by 𝐵 ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) − 1 on the right, and integrating over 𝑠 ′ ∼ 𝜌 .

Original: Thus, a posteriori, FB can be used to produce both 𝜙 and 𝜓 in SF, although training is different.
Revised: Thus, a posteriori, FB can be used to produce both 𝜙 and 𝜓 in SF, although training is different.

Original: This connection between FB and SF is one-directional: (5) is a stronger condition.
Revised: This connection between FB and SF is one-directional: (5) is a stronger condition.

Original: In particular 𝐹 = 𝐵 = 0 is not a solution: contrary to 𝜓 = 𝜙 = 0 in (4), there is no collapse.
Revised: In particular 𝐹 = 𝐵 = 0 is not a solution: contrary to 𝜓 = 𝜙 = 0 in (4), there is no collapse.

Original: No additional criterion to train 𝜙 is required: 𝐹 and 𝐵 are trained jointly to provide the best rank- 𝑑 approximation of the successor measures 𝑀 𝜋 .
Revised: No additional criterion to train 𝜙 is required: 𝐹 and 𝐵 are trained jointly to provide the best rank- 𝑑 approximation of the successor measures 𝑀 𝜋 .

Original: This summarizes an environment by selecting the features that best describe the relationship 𝑄 𝜋𝑟
Revised: This summarizes an environment by selecting the features that best describe the relationship 𝑄 𝜋𝑟

Original: =
Revised: =

Original: 𝑀 𝜋 𝑟 between rewards and 𝑄 -functions.
Revised: 𝑀 𝜋 𝑟 between rewards and 𝑄 -functions.

Original:  Algorithms for Successor Features and FB Representations
Revised:  Algorithms for Successor Features and FB Representations

Original: We now describe more precisely the algorithms used in our experiments.
Revised: We now describe more precisely the algorithms used in our experiments.

Original: The losses used to train 𝜓 in SFs, and 𝐹, 𝐵 in FB, are described in Sections 5.1 and 5.2 respectively.
Revised: The losses used to train 𝜓 in SFs, and 𝐹, 𝐵 in FB, are described in Sections 5.1 and 5.2 respectively.

Original: To obtain a full zero-shot RL algorithm, SFs must specify the basic features 𝜙 .
Revised: To obtain a full zero-shot RL algorithm, SFs must specify the basic features 𝜙 .

Original: Any representation learning method can be used for 𝜙 .
Revised: Any representation learning method can be used for 𝜙 .

Original: We use ten possible choices (Section 5.3) based on existing or new representations for RL: random features as a baseline, autoencoders, next state and latent next state transition models, inverse curiosity module, the diversity criterion of APS, contrastive learning, and ﬁnally, several spectral decompositions of the transition matrix or its associated Laplacian.
Revised: We use ten possible choices (Section 5.3) based on existing or new representations for RL: random features as a baseline, autoencoders, next state and latent next state transition models, inverse curiosity module, the diversity criterion of APS, contrastive learning, and ﬁnally, several spectral decompositions of the transition matrix or its associated Laplacian.

Original: Both SFs and FB deﬁne policies as an argmax of 𝜓 ( 𝑠, 𝑎, 𝑧 )
Revised: Both SFs and FB deﬁne policies as an argmax of 𝜓 ( 𝑠, 𝑎, 𝑧 )

Original: ⊤ 𝑧 or 𝐹 ( 𝑠, 𝑎, 𝑧 )
Revised: ⊤ 𝑧 or 𝐹 ( 𝑠, 𝑎, 𝑧 )

Original: ⊤ 𝑧 over actions 𝑎 .
Revised: ⊤ 𝑧 over actions 𝑎 .

Original: With continuous actions, the argmax cannot be computed exactly.
Revised: With continuous actions, the argmax cannot be computed exactly.

Original: We train an auxiliary policy network 𝜋 ( 𝑠, 𝑧 ) to approximate this argmax, using the same standard method for SFs and FB (Appendix H.4).
Revised: We train an auxiliary policy network 𝜋 ( 𝑠, 𝑧 ) to approximate this argmax, using the same standard method for SFs and FB (Appendix G.4).

Original: 5.1 Learning the Successor Features 𝜓
Revised: 5.1 Learning the Successor Features 𝜓

Original: The successor features 𝜓 satisfy the R 𝑑 -valued Bellman equation 𝜓 𝜋 =
Revised: The successor features 𝜓 satisfy the R 𝑑 -valued Bellman equation 𝜓 𝜋 =

Original: 𝑃𝜙 + 𝛾𝑃 𝜋 𝜓 𝜋 , the collection of ordinary Bellman equations for each component of 𝜙 .
Revised: 𝑃𝜙 + 𝛾𝑃 𝜋 𝜓 𝜋 , the collection of ordinary Bellman equations for each component of 𝜙 .

Original: The 𝑃 in front of 𝜙 comes from using 𝜙 ( 𝑠 𝑡 +1 ) not 𝜙 ( 𝑠 𝑡 ) in (2).
Revised: The 𝑃 in front of 𝜙 comes from using 𝜙 ( 𝑠 𝑡 +1 ) not 𝜙 ( 𝑠 𝑡 ) in (2).

Original: Therefore, we can train 𝜓 ( 𝑠, 𝑎, 𝑧 ) for each 𝑧 by minimizing the Bellman residuals ⃦⃦ 𝜓 ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑧 ) − 𝜙 ( 𝑠 𝑡 +1 )
Revised: Therefore, we can train 𝜓 ( 𝑠, 𝑎, 𝑧 ) for each 𝑧 by minimizing the Bellman residuals ⃦⃦ 𝜓 ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑧 ) − 𝜙 ( 𝑠 𝑡 +1 )

Original: − 𝛾 ¯ 𝜓 ( 𝑠 𝑡 +1 , 𝜋 𝑧 ( 𝑠 𝑡 +1 ) , 𝑧 ) ⃦⃦ 2 where ¯ 𝜓 is a non-trainable target version of 𝜓 as in parametric 𝑄 -learning.
Revised: − 𝛾 ¯ 𝜓 ( 𝑠 𝑡 +1 , 𝜋 𝑧 ( 𝑠 𝑡 +1 ) , 𝑧 ) ⃦⃦ 2 where ¯ 𝜓 is a non-trainable target version of 𝜓 as in parametric 𝑄 -learning.

Original: This requires sampling a transition ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) from the dataset and choosing 𝑧 .
Revised: This requires sampling a transition ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) from the dataset and choosing 𝑧 .

Original: We sample random values of 𝑧 as described in Appendix H.3.
Revised: We sample random values of 𝑧 as described in Appendix G.3.

Original: This is the loss used in Borsa et al.
Revised: This is the loss used in Borsa et al.

Original: But this can be improved, since we do not use the full vector 𝜓 ( 𝑠, 𝑎, 𝑧 ) :
Revised: But this can be improved, since we do not use the full vector 𝜓 ( 𝑠, 𝑎, 𝑧 ) :

Original: only 𝜓 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 is needed for the policies.
Revised: only 𝜓 ( 𝑠, 𝑎, 𝑧 ) ⊤ 𝑧 is needed for the policies.

Original: Therefore, as in Liu & Abbeel (2021), instead of the vector-valued Bellman residual above, we just use
Revised: Therefore, as in Liu & Abbeel (2021), instead of the vector-valued Bellman residual above, we just use

Original: [Equation]
Revised: [Equation]

Original: for each 𝑧 .
Revised: for each 𝑧 .

Original: This trains 𝜓 ( · , 𝑧 ) ⊤ 𝑧 as the 𝑄 -function of reward 𝜙 ⊤ 𝑧 , the only case needed, while training the full vector 𝜓 ( · , 𝑧 ) amounts to training the 𝑄 -functions of each policy 𝜋 𝑧 for all rewards 𝜙 ⊤ 𝑧 ′ for all 𝑧 ′ ∈ R 𝑑 including 𝑧 ′ ̸
Revised: This trains 𝜓 ( · , 𝑧 ) ⊤ 𝑧 as the 𝑄 -function of reward 𝜙 ⊤ 𝑧 , the only case needed, while training the full vector 𝜓 ( · , 𝑧 ) amounts to training the 𝑄 -functions of each policy 𝜋 𝑧 for all rewards 𝜙 ⊤ 𝑧 ′ for all 𝑧 ′ ∈ R 𝑑 including 𝑧 ′ ̸

Original: = 𝑧 .
Revised: = 𝑧 .

Original: We have found this improves performance.
Revised: We have found this improves performance.

Original: 5.2 Learning FB Representations: the FB Training Loss
Revised: 5.2 Learning FB Representations: the FB Training Loss

Original: The successor measure 𝑀 𝜋 satisﬁes a Bellman-like equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 , as matrices in the ﬁnite case and as measures in the general case (Blier et al., 2021).
Revised: The successor measure 𝑀 𝜋 satisﬁes a Bellman-like equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 , as matrices in the ﬁnite case and as measures in the general case (Blier et al., 2021).

Original: We can learn FB by iteratively minimizing the Bellman residual on the parametric model 𝑀 =
Revised: We can learn FB by iteratively minimizing the Bellman residual on the parametric model 𝑀 =

Original: 𝐹 ⊤ 𝐵𝜌 .
Revised: 𝐹 ⊤ 𝐵𝜌 .

Original: Using a suitable norm ‖·‖ 𝜌 for the Bellman residual (Appendix C) leads to a loss expressed as expectations from the dataset:
Revised: Using a suitable norm ‖·‖ 𝜌 for the Bellman residual (Appendix B) leads to a loss expressed as expectations from the dataset:

Original: [Equation]
Revised: [Equation]

Original: where the constant term does not depend on 𝐹 and 𝐵 , and where as usual ¯ 𝐹 𝑧 and ¯ 𝐵 are non-trainable target versions of 𝐹 and 𝐵 whose parameters are updated with a slow-moving average of those of 𝐹 and 𝐵 .
Revised: where the constant term does not depend on 𝐹 and 𝐵 , and where as usual ¯ 𝐹 𝑧 and ¯ 𝐵 are non-trainable target versions of 𝐹 and 𝐵 whose parameters are updated with a slow-moving average of those of 𝐹 and 𝐵 .

Original: Appendix C quickly derives this loss.
Revised: Appendix B quickly derives this loss.

Original: Contrary to Touati & Ollivier (2021), the last term involves 𝐵 ( 𝑠 𝑡 +1 ) instead of 𝐵 ( 𝑠 𝑡 ) , because we use 𝑠 𝑡 +1 instead of 𝑠 𝑡 for the successor measures (3).
Revised: Contrary to Touati & Ollivier (2021), the last term involves 𝐵 ( 𝑠 𝑡 +1 ) instead of 𝐵 ( 𝑠 𝑡 ) , because we use 𝑠 𝑡 +1 instead of 𝑠 𝑡 for the successor measures (3).

Original: We sample random values of 𝑧 as described in Appendix H.3.
Revised: We sample random values of 𝑧 as described in Appendix G.3.

Original: We include an auxiliary loss (Appendix C) to normalize the covariance of 𝐵 , E 𝜌 𝐵𝐵 ⊤ ≈ Id , as in
Revised: We include an auxiliary loss (Appendix B) to normalize the covariance of 𝐵 , E 𝜌 𝐵𝐵 ⊤ ≈ Id , as in

Original: Touati & Ollivier (2021)
Revised: Touati & Ollivier (2021)

Original: (otherwise one can, e.g., scale 𝐹 up and 𝐵 down since only 𝐹 ⊤ 𝐵 is ﬁxed).
Revised: (otherwise one can, e.g., scale 𝐹 up and 𝐵 down since only 𝐹 ⊤ 𝐵 is ﬁxed).

Original: As with SFs above, only 𝐹 ( · , 𝑧 ) ⊤ 𝑧 is needed for the policies, while the loss above on the vector 𝐹 amounts to training 𝐹 ( · , 𝑧 ) ⊤ 𝑧 ′ for all pairs ( 𝑧, 𝑧 ′ ) .
Revised: As with SFs above, only 𝐹 ( · , 𝑧 ) ⊤ 𝑧 is needed for the policies, while the loss above on the vector 𝐹 amounts to training 𝐹 ( · , 𝑧 ) ⊤ 𝑧 ′ for all pairs ( 𝑧, 𝑧 ′ ) .

Original: The full loss is needed for joint training of 𝐹 and 𝐵 .
Revised: The full loss is needed for joint training of 𝐹 and 𝐵 .

Original: But we include an auxiliary loss ℒ ′ ( 𝐹 ) to focus training on the diagonal 𝑧 ′
Revised: But we include an auxiliary loss ℒ ′ ( 𝐹 ) to focus training on the diagonal 𝑧 ′

Original: =
Revised: =

Original: 𝑧 .
Revised: 𝑧 .

Original: This is obtained by multiplying the Bellman gap in ℒ by 𝐵 ⊤ ( 𝐵𝜌𝐵 ⊤ ) − 1 𝑧 on the right, to make 𝐹 ( · , 𝑧 ) ⊤ 𝑧 appear: ℒ ′ ( 𝐹 ) :
Revised: This is obtained by multiplying the Bellman gap in ℒ by 𝐵 ⊤ ( 𝐵𝜌𝐵 ⊤ ) − 1 𝑧 on the right, to make 𝐹 ( · , 𝑧 ) ⊤ 𝑧 appear: ℒ ′ ( 𝐹 ) :

Original: = E ( 𝑠 𝑡 ,𝑎 𝑡 ,𝑠 𝑡 +1 ) ∼ 𝜌 [︁(︀ 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑧 ) ⊤
Revised: = E ( 𝑠 𝑡 ,𝑎 𝑡 ,𝑠 𝑡 +1 ) ∼ 𝜌 [︁(︀ 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑧 ) ⊤

Original: [Equation]
Revised: [Equation]

Original: − 𝐵 ( 𝑠 𝑡 +1 )
Revised: − 𝐵 ( 𝑠 𝑡 +1 )

Original: ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) −
Revised: ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) −

Original: [Equation]
Revised: [Equation]

Original: − 𝛾 ¯ 𝐹 ( 𝑠 𝑡 +
Revised: − 𝛾 ¯ 𝐹 ( 𝑠 𝑡 +

Original: [Equation]
Revised: [Equation]

Original: 𝜋 𝑧 ( 𝑠 𝑡 +1 )
Revised: 𝜋 𝑧 ( 𝑠 𝑡 +1 )

Original: [Equation]
Revised: [Equation]

Original: 𝑧 ) ⊤ 𝑧 )︀ 2 ]︁ .
Revised: 𝑧 ) ⊤ 𝑧 )︀ 2 ]︁ .

Original: This trains 𝐹 ( · , 𝑧 ) ⊤ 𝑧 as the 𝑄 -function for reward 𝐵 ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) − 1 𝑧 .
Revised: This trains 𝐹 ( · , 𝑧 ) ⊤ 𝑧 as the 𝑄 -function for reward 𝐵 ⊤ ( E 𝜌 𝐵𝐵 ⊤ ) − 1 𝑧 .

Original: Though ℒ = 0 implies ℒ ′
Revised: Though ℒ = 0 implies ℒ ′

Original: = 0 , adding ℒ ′ reduces the error on the part used for policies.
Revised: = 0 , adding ℒ ′ reduces the error on the part used for policies.

Original: This departs from Touati & Ollivier (2021).
Revised: This departs from Touati & Ollivier (2021).

Original: 5.3 Learning Basic Features 𝜙 for Successor Features
Revised: 5.3 Learning Basic Features 𝜙 for Successor Features

Original: SFs must be provided with basic state features 𝜙 .
Revised: SFs must be provided with basic state features 𝜙 .

Original: Any representation learning method can be used to supply 𝜙 .
Revised: Any representation learning method can be used to supply 𝜙 .

Original: We focus on prominent RL representation learning baselines, and on those used in previous zero-shot RL candidates such as APS.
Revised: We focus on prominent RL representation learning baselines, and on those used in previous zero-shot RL candidates such as APS.

Original: We now describe the precise learning objective for each.
Revised: We now describe the precise learning objective for each.

Original: Random Features ( Rand ).
Revised: Random Features ( Rand ).

Original: We use a non-trainable randomly initialized network as features.
Revised: We use a non-trainable randomly initialized network as features.

Original: Autoencoder ( AEnc ).
Revised: Autoencoder ( AEnc ).

Original: We learn a decoder 𝑓 : R 𝑑 → 𝑆 to recover the state from its representation 𝜙 :
Revised: We learn a decoder 𝑓 : R 𝑑 → 𝑆 to recover the state from its representation 𝜙 :

Original: [Equation]
Revised: [Equation]

Original: Inverse Curiosity Module ( ICM ) aims at extracting the controllable aspects of the environment (Pathak et al., 2017).
Revised: Inverse Curiosity Module ( ICM ) aims at extracting the controllable aspects of the environment (Pathak et al., 2017).

Original: The idea is to train an inverse dynamics model 𝑔 : R 𝑑 × R 𝑑 → 𝐴 to predict the action used for a transition between two consecutive states.
Revised: The idea is to train an inverse dynamics model 𝑔 : R 𝑑 × R 𝑑 → 𝐴 to predict the action used for a transition between two consecutive states.

Original: We use the loss
Revised: We use the loss

Original: [Equation]
Revised: [Equation]

Original: Transition model ( Trans ).
Revised: Transition model ( Trans ).

Original: This is a one-step forward dynamic model 𝑓 : R 𝑑 × 𝐴 → 𝑆 that predicts the next state from the current state representation:
Revised: This is a one-step forward dynamic model 𝑓 : R 𝑑 × 𝐴 → 𝑆 that predicts the next state from the current state representation:

Original: [Equation]
Revised: [Equation]

Original: Latent transition model ( Latent ).
Revised: Latent transition model ( Latent ).

Original: This is similar to the transition model but instead of predicting the next state, it predicts its representation:
Revised: This is similar to the transition model but instead of predicting the next state, it predicts its representation:

Original: [Equation]
Revised: [Equation]

Original: A clear failure case of this loss is when all states are mapped to the same representation.
Revised: A clear failure case of this loss is when all states are mapped to the same representation.

Original: To avoid this collapse, we compute 𝜙 ( 𝑠 𝑡 +1 ) using a non-trainable version of 𝜙 , with parameters corresponding to a slowly moving average of the parameters of 𝜙 , similarly to BYOL (Grill et al., 2020).
Revised: To avoid this collapse, we compute 𝜙 ( 𝑠 𝑡 +1 ) using a non-trainable version of 𝜙 , with parameters corresponding to a slowly moving average of the parameters of 𝜙 , similarly to BYOL (Grill et al., 2020).

Original: Diversity methods ( APS ).
Revised: Diversity methods ( APS ).

Original: VISR (Hansen et al., 2019) and its successor APS (Liu & Abbeel, 2021) tackle zero-shot RL using SFs with features 𝜙 built online from a diversity criterion.
Revised: VISR (Hansen et al., 2019) and its successor APS (Liu & Abbeel, 2021) tackle zero-shot RL using SFs with features 𝜙 built online from a diversity criterion.

Original: 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).
Revised: 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).

Original: VISR and APS use, respectively, a variational or nearest-neighbor estimator for the mutual information.
Revised: VISR and APS use, respectively, a variational or nearest-neighbor estimator for the mutual information.

Original: We directly use the code provided for APS, and refer to (Liu & Abbeel, 2021) for the details.
Revised: We directly use the code provided for APS, and refer to (Liu & Abbeel, 2021) for the details.

Original: Contrary to other methods, APS is not ofﬂine: it needs to be trained on its own replay buffer.
Revised: Contrary to other methods, APS is not ofﬂine: it needs to be trained on its own replay buffer.

Original: Laplacian Eigenfunctions ( Lap ).
Revised: Laplacian Eigenfunctions ( Lap ).

Original: Wu et al.
Revised: Wu et al.

Original: (2018) consider the symmetrized MDP graph Laplacian induced by an exploratory policy 𝜋 , deﬁned as ℒ
Revised: (2018) consider the symmetrized MDP graph Laplacian induced by an exploratory policy 𝜋 , deﬁned as ℒ

Original: = Id − 12 ( 𝑃 𝜋 diag( 𝜌 ) − 1 + diag( 𝜌 ) − 1 ( 𝑃 𝜋 ) ⊤ ) .
Revised: = Id − 12 ( 𝑃 𝜋 diag( 𝜌 ) − 1 + diag( 𝜌 ) − 1 ( 𝑃 𝜋 ) ⊤ ) .

Original: They propose to learn the eigenfunctions of ℒ via the spectral graph drawing objective (Koren, 2003): min
Revised: They propose to learn the eigenfunctions of ℒ via the spectral graph drawing objective (Koren, 2003): min

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where the second term is an orthonormality regularization to ensure that E 𝑠 ∼ 𝜌
Revised: where the second term is an orthonormality regularization to ensure that E 𝑠 ∼ 𝜌

Original: [ 𝜙 ( 𝑠 ) 𝜙 ( 𝑠 ) ⊤ ]
Revised: [ 𝜙 ( 𝑠 ) 𝜙 ( 𝑠 ) ⊤ ]

Original: ≈ Id , and 𝜆 > 0 is the regularization weight.
Revised: ≈ Id , and 𝜆 > 0 is the regularization weight.

Original: This is implicitly contrastive, pushing features of 𝑠 𝑡 and 𝑠 𝑡 + closer while keeping features apart overall.
Revised: This is implicitly contrastive, pushing features of 𝑠 𝑡 and 𝑠 𝑡 + closer while keeping features apart overall.

Original: Such eigenfunctions have long been argued to play a key role in RL (Mahadevan & Maggioni, 2007; Machado et al., 2017).
Revised: Such eigenfunctions have long been argued to play a key role in RL (Mahadevan & Maggioni, 2007; Machado et al., 2017).

Original: Low-Rank Approximation of 𝑃 ( LRA-P ): we learn features by estimating a low-rank model of the transition probability densities:
Revised: Low-Rank Approximation of 𝑃 ( LRA-P ): we learn features by estimating a low-rank model of the transition probability densities:

Original: 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) ≈ 𝜒
Revised: 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) ≈ 𝜒

Original: ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )
Revised: ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )

Original: 𝜌 (d 𝑠 ′ ) .
Revised: 𝜌 (d 𝑠 ′ ) .

Original: Knowing 𝜌 is not needed: the corresponding loss on 𝜒 ⊤ 𝜇 − 𝑃/𝜌 is readily expressed as expectations over the dataset,
Revised: Knowing 𝜌 is not needed: the corresponding loss on 𝜒 ⊤ 𝜇 − 𝑃/𝜌 is readily expressed as expectations over the dataset,

Original: [Equation]
Revised: [Equation]

Original: We normalize E 𝜌 [ 𝜇𝜇 ⊤ ]
Revised: We normalize E 𝜌 [ 𝜇𝜇 ⊤ ]

Original: ≈ Id with the same loss used for 𝐵 .
Revised: ≈ Id with the same loss used for 𝐵 .

Original: Then we use SFs with 𝜙 := 𝜇 .
Revised: Then we use SFs with 𝜙 := 𝜇 .

Original: If the model 𝑃 = 𝜒 ⊤ 𝜇 𝜌 is exact, this provides exact optimal policies for any reward (Appendix F, Thm. 3).
Revised: If the model 𝑃 = 𝜒 ⊤ 𝜇 𝜌 is exact, this provides exact optimal policies for any reward (Appendix E, Thm. 3).

Original: This loss is implicitly contrastive: it compares samples 𝑠 𝑡 +1 to independent samples 𝑠 ′ from 𝜌 .
Revised: This loss is implicitly contrastive: it compares samples 𝑠 𝑡 +1 to independent samples 𝑠 ′ from 𝜌 .

Original: It is an asymmetric extension of the Laplacian loss (14) (Appendix G).
Revised: It is an asymmetric extension of the Laplacian loss (14) (Appendix F).

Original: The loss (16) is also a special case of the FB loss (8) by setting 𝛾 = 0 , omitting 𝑧 , and substituting ( 𝜒, 𝜇 ) for ( 𝐹, 𝐵 ) .
Revised: The loss (16) is also a special case of the FB loss (8) by setting 𝛾 = 0 , omitting 𝑧 , and substituting ( 𝜒, 𝜇 ) for ( 𝐹, 𝐵 ) .

Original: Indeed, FB learns a ﬁnite-rank model of 𝑃 (Id − 𝛾𝑃 𝜋 ) − 1 , which equals 𝑃 when 𝛾 = 0 .
Revised: Indeed, FB learns a ﬁnite-rank model of 𝑃 (Id − 𝛾𝑃 𝜋 ) − 1 , which equals 𝑃 when 𝛾 = 0 .

Original: A similar loss was introduced in Ren et al.
Revised: A similar loss was introduced in Ren et al.

Original: However, their loss involves a second unspeciﬁed, arbitrary probability distribution 𝑝 , which must cover the whole state space and whose analytic expression must be known.
Revised: However, their loss involves a second unspeciﬁed, arbitrary probability distribution 𝑝 , which must cover the whole state space and whose analytic expression must be known.

Original: It is unclear how to set a suitable 𝑝 in general.
Revised: It is unclear how to set a suitable 𝑝 in general.

Original: Contrastive Learning ( CL ) methods learn representations by pushing positive pairs (similar states) closer together while keeping negative pairs apart.
Revised: Contrastive Learning ( CL ) methods learn representations by pushing positive pairs (similar states) closer together while keeping negative pairs apart.

Original: Here, two states are considered similar if they lie close on the same trajectory.
Revised: Here, two states are considered similar if they lie close on the same trajectory.

Original: We use a SimCLR-like objective (Chen et al., 2020):
Revised: We use a SimCLR-like objective (Chen et al., 2020):

Original: [Equation]
Revised: [Equation]

Original: where 𝑠 𝑡 + 𝑘 is the state encountered at step 𝑡 + 𝑘 along the subtrajectory that starts at 𝑠 𝑡 , where 𝑘 is sampled from a geometric distribution of parameter (1 − 𝛾 CL ) , and cosine ( 𝑢, 𝑣 ) =
Revised: where 𝑠 𝑡 + 𝑘 is the state encountered at step 𝑡 + 𝑘 along the subtrajectory that starts at 𝑠 𝑡 , where 𝑘 is sampled from a geometric distribution of parameter (1 − 𝛾 CL ) , and cosine ( 𝑢, 𝑣 ) =

Original: 𝑢 ⊤ 𝑣
Revised: 𝑢 ⊤ 𝑣

Original: [Equation]
Revised: [Equation]

Original: 𝑢
Revised: 𝑢

Original: [Equation]
Revised: [Equation]

Original: 𝑣
Revised: 𝑣

Original: [Equation]
Revised: [Equation]

Original: ∀ 𝑢, 𝑣 ∈
Revised: ∀ 𝑢, 𝑣 ∈

Original: R 𝑑 is the cosine similarity function.
Revised: R 𝑑 is the cosine similarity function.

Original: Here 𝛾 CL ∈ [0; 1) is a parameter not necessarily set to the MDP’s discount factor 𝛾 .
Revised: Here 𝛾 CL ∈ [0; 1) is a parameter not necessarily set to the MDP’s discount factor 𝛾 .

Original: CL requires a dataset made of full trajectories instead of isolated transitions.
Revised: CL requires a dataset made of full trajectories instead of isolated transitions.

Original: CL is tightly related to the spectral decomposition of the successor measure ∑︀ 𝑡 𝛾 𝑡 CL 𝑃 𝑡 +1 𝜋 , where 𝜋 is the behavior policy generating the dataset trajectories.
Revised: CL is tightly related to the spectral decomposition of the successor measure ∑︀ 𝑡 𝛾 𝑡 CL 𝑃 𝑡 +1 𝜋 , where 𝜋 is the behavior policy generating the dataset trajectories.

Original: Precisely, assuming that 𝜒 and 𝜙 are centered with unit norm, and expanding the log and exp at second order, the loss (17) becomes
Revised: Precisely, assuming that 𝜒 and 𝜙 are centered with unit norm, and expanding the log and exp at second order, the loss (17) becomes

Original: [Equation]
Revised: [Equation]

Original: (compare (16)).
Revised: (compare (16)).

Original: Now, the law of 𝑠 𝑡 + 𝑘 given 𝑠 𝑡 is given by the stochastic matrix (1 − 𝛾 CL )
Revised: Now, the law of 𝑠 𝑡 + 𝑘 given 𝑠 𝑡 is given by the stochastic matrix (1 − 𝛾 CL )

Original: ∑︀ 𝑡 𝛾 𝑡 CL 𝑃 𝑡 +

Revised: ∑︀ 𝑡 𝛾 𝑡 CL 𝑃 𝑡 +

Original: 𝜋 , the rescaled successor measure of 𝜋 .
Revised: 𝜋 , the rescaled successor measure of 𝜋 .

Original: Then one ﬁnds that (18) is minimized when 𝜒 and 𝜙 provide the singular value decomposition of this matrix in 𝐿 2 ( 𝜌 ) norm (Appendix D).
Revised: Then one ﬁnds that (18) is minimized when 𝜒 and 𝜙 provide the singular value decomposition of this matrix in 𝐿 2 ( 𝜌 ) norm (Appendix C).

Original: Formal links between contrastive learning and spectral methods can be found in Tian (2022); Balestriero & LeCun (2022).
Revised: Formal links between contrastive learning and spectral methods can be found in Tian (2022); Balestriero & LeCun (2022).

Original: R e t u r n maze
Revised: R e t u r n maze

Original: R e t u r n walker_stand R e t u r n walker_walk
Revised: R e t u r n walker_stand R e t u r n walker_walk

Original: R e t u r n walker_run R e t u r n walker_flip
Revised: R e t u r n walker_run R e t u r n walker_flip

Original: R e t u r n cheetah_walk R e t u r n cheetah_walk_backward
Revised: R e t u r n cheetah_walk R e t u r n cheetah_walk_backward

Original: R e t u r n cheetah_run
Revised: R e t u r n cheetah_run

Original: R e t u r n cheetah_run_backward
Revised: R e t u r n cheetah_run_backward

Original: R e t u r n quadruped_stand
Revised: R e t u r n quadruped_stand

Original: R e t u r n quadruped_walk R e t u r n quadruped_run
Revised: R e t u r n quadruped_walk R e t u r n quadruped_run

Original: R e t u r n quadruped
Revised: R e t u r n quadruped

Original: [Figure]
Revised: [Figure]

Original: Low-Rank Approximation of SR ( LRA-SR ).
Revised: Low-Rank Approximation of SR ( LRA-SR ).

Original: The CL method implicitly factorizes the successor measure of the exploration policy in Monte Carlo fashion by sampling pairs ( 𝑠 𝑡 , 𝑠 𝑡 + 𝑘 ) on the same trajectory.
Revised: The CL method implicitly factorizes the successor measure of the exploration policy in Monte Carlo fashion by sampling pairs ( 𝑠 𝑡 , 𝑠 𝑡 + 𝑘 ) on the same trajectory.

Original: This may suffer from high variance.
Revised: This may suffer from high variance.

Original: To mitigate this, we propose to factorize this successor measure by temporal difference learning instead of Monte Carlo.
Revised: To mitigate this, we propose to factorize this successor measure by temporal difference learning instead of Monte Carlo.

Original: This is achieved with an FB-like loss (8) except we drop the policies 𝜋 𝑧 and learn successor measures for the exploration policy only: min
Revised: This is achieved with an FB-like loss (8) except we drop the policies 𝜋 𝑧 and learn successor measures for the exploration policy only: min

Original: [Equation]
Revised: [Equation]

Original: (19) with ¯ 𝜒 and ¯ 𝜙 target versions of 𝜒 and 𝜙 .
Revised: (19) with ¯ 𝜒 and ¯ 𝜙 target versions of 𝜒 and 𝜙 .

Original: We normalize 𝜙 to E 𝜌 [ 𝜙𝜙 ⊤ ]
Revised: We normalize 𝜙 to E 𝜌 [ 𝜙𝜙 ⊤ ]

Original: ≈ Id with the same loss as for 𝐵 .
Revised: ≈ Id with the same loss as for 𝐵 .

Original: Of all SF variants tested, this is the closest to FB.
Revised: Of all SF variants tested, this is the closest to FB.

Original:  Experimental Results on Benchmarks
Revised:  Experimental Results on Benchmarks

Original: 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, ﬂip), Cheetah (walk, run, walk backwards, run backwards), and Quadruped (stand, walk, run, jump); see Appendix H.1.
Revised: 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, ﬂip), Cheetah (walk, run, walk backwards, run backwards), and Quadruped (stand, walk, run, jump); see Appendix G.1.

Original: Each task and method was repeated for 3 choices of replay buffer from ExORL: RND, APS, and Proto (except for the APS method, which can only train on the APS buffer).
Revised: Each task and method was repeated for 3 choices of replay buffer from ExORL: RND, APS, and Proto (except for the APS method, which can only train on the APS buffer).

Original: Each of these 403 settings was repeated with 10 random seeds.
Revised: Each of these 403 settings was repeated with 10 random seeds.

Original: The full setup is described in Appendix H. Representation dimension is 𝑑 = 50 , except for Maze ( 𝑑 = 100 ).
Revised: The full setup is described in Appendix G. Representation dimension is 𝑑 = 50 , except for Maze ( 𝑑 = 100 ).

Original: After model training, tasks are revealed by 10,000 reward samples as in (Liu & Abbeel, 2021), except for Maze, where a known goal is presented and used to set 𝑧 𝑟 directly.
Revised: After model training, tasks are revealed by 10,000 reward samples as in (Liu & Abbeel, 2021), except for Maze, where a known goal is presented and used to set 𝑧 𝑟 directly.

Original: As toplines, we use online TD3 (with task rewards, and free environment interactions not restricted to a replay buffer), and ofﬂine TD3 (restricted to each replay buffer labelled with task rewards).
Revised: As toplines, we use online TD3 (with task rewards, and free environment interactions not restricted to a replay buffer), and ofﬂine TD3 (restricted to each replay buffer labelled with task rewards).

Original: Ofﬂine
Revised: Ofﬂine

Original: TD3 gives an idea of the best achievable performance given the training data in a buffer.
Revised: TD3 gives an idea of the best achievable performance given the training data in a buffer.

Original: In Fig.
Revised: In Fig.

Original: 2 we plot the performance of each method for each task in each environment, averaged over the three replay buffers and ten random seeds.
Revised: 2 we plot the performance of each method for each task in each environment, averaged over the three replay buffers and ten random seeds.

Original: Appendix I contains the full results and more plots.
Revised: Appendix H contains the full results and more plots.

Original: Compared to ofﬂine TD3 as a reference, on the Maze tasks, Lap , LRA-P , and FB perform well.
Revised: Compared to ofﬂine TD3 as a reference, on the Maze tasks, Lap , LRA-P , and FB perform well.

Original: On the
Revised: On the

Original: Walker tasks, ICM , Trans , Lap , LRA-SR , and FB perform well.
Revised: Walker tasks, ICM , Trans , Lap , LRA-SR , and FB perform well.

Original: On the Cheetah tasks, ICM , Trans ,
Revised: On the Cheetah tasks, ICM , Trans ,

Original: Lap , LRA-SR , and FB perform well.
Revised: Lap , LRA-SR , and FB perform well.

Original: On the Quadruped tasks, many methods perform well, including, surprisingly, random features.
Revised: On the Quadruped tasks, many methods perform well, including, surprisingly, random features.

Original: On Maze, none of the encoder-based losses learn anything, contrary to FB and spectral SF methods.
Revised: On Maze, none of the encoder-based losses learn anything, contrary to FB and spectral SF methods.

Original: We believe this is because ( 𝑥, 𝑦 ) is already a good representation of the 2D state 𝑠 , so the encoders do nothing.
Revised: We believe this is because ( 𝑥, 𝑦 ) is already a good representation of the 2D state 𝑠 , so the encoders do nothing.

Original: Yet SFs on ( 𝑥, 𝑦 ) cannot solve the task: planning requires speciﬁc representations.
Revised: Yet SFs on ( 𝑥, 𝑦 ) cannot solve the task: planning requires speciﬁc representations.

Original: Fig.
Revised: Fig.

Original: 1 reports aggregated scores over all tasks.
Revised: 1 reports aggregated scores over all tasks.

Original: To average across tasks, we normalize scores: for each task and replay buffer, performance is expressed as a percentage of the performance of ofﬂine TD on the same replay buffer, a natural supervised topline given the data.
Revised: To average across tasks, we normalize scores: for each task and replay buffer, performance is expressed as a percentage of the performance of ofﬂine TD on the same replay buffer, a natural supervised topline given the data.

Original: These normalized scores are averaged over all tasks in each environment, then over environments to yield the scores in Fig.
Revised: These normalized scores are averaged over all tasks in each environment, then over environments to yield the scores in Fig.

Original: The variations over environments, replay buffers and random seeds are reported in Appendix I.
Revised: The variations over environments, replay buffers and random seeds are reported in Appendix H.

Original: These results are broadly consistent over replay buffers.
Revised: These results are broadly consistent over replay buffers.

Original: Buffer-speciﬁc results are reported in Appendices I.3–I.4.
Revised: Buffer-speciﬁc results are reported in Appendices H.3–H.4.

Original: Sometimes a replay buffer is restrictive, as attested by poor ofﬂine TD3 performance, starkly so for the Proto buffer on all Quadruped tasks.
Revised: Sometimes a replay buffer is restrictive, as attested by poor ofﬂine TD performance, starkly so for the Proto buffer on all Quadruped tasks.

Original: APS does not work well as a zero-shot RL method, but it does work well as an exploration method: on average, the APS and RND replay buffers have close results.
Revised: APS does not work well as a zero-shot RL method, but it does work well as an exploration method: on average, the APS and RND replay buffers have close results.

Original: FB and Lap are the only methods that perform consistently well, both over tasks and over replay buffers.
Revised: FB and Lap are the only methods that perform consistently well, both over tasks and over replay buffers.

Original: Averaged on all tasks and buffers, FB reaches 81% of supervised ofﬂine TD3 performance, and 85% on the RND buffer.
Revised: Averaged on all tasks and buffers, FB reaches 81% of supervised ofﬂine TD3 performance, and 85% on the RND buffer.

Original: The second-best method is Lap with 74% (78% on the Proto buffer).
Revised: The second-best method is Lap with 74% (78% on the Proto buffer).

Original:  Discussion 
Revised:  Discussion and Limitations

Original: Can a few features solve many reward functions?
Revised: Can a few features solve many reward functions?

Original: Why are some features better?
Revised: Why are some features better?

Original: The choice of features is critical.
Revised: The choice of features is critical.

Original: Consider goal-reaching tasks: the obvious way to learn to reach arbitrary goal states via SFs is to deﬁne one feature per possible goal state (one-hot encoding).
Revised: Consider goal-reaching tasks: the obvious way to learn to reach arbitrary goal states via SFs is to deﬁne one feature per possible goal state (one-hot encoding).

Original: This requires | 𝑆 | features, and does not scale to continuous spaces.
Revised: This requires | 𝑆 | features, and does not scale to continuous spaces.

Original: But much better features exist.
Revised: But much better features exist.

Original: For instance, with an size- 𝑛 cycle 𝑆 = { 0 , . . .
Revised: For instance, with an size- 𝑛 cycle 𝑆 = { 0 , . . .

Original: , 𝑛 − 1 } with actions that move left and right modulo 𝑛 , then two features sufﬁce instead of 𝑛 : SFs with 𝜙 ( 𝑠 ) =
Revised: , 𝑛 − 1 } with actions that move left and right modulo 𝑛 , then two features sufﬁce instead of 𝑛 : SFs with 𝜙 ( 𝑠 ) =

Original: (cos(2 𝜋𝑠/𝑛 ) , sin(2 𝜋𝑠/𝑛 )) provides exact optimal policies to reach any arbitrary state.
Revised: (cos(2 𝜋𝑠/𝑛 ) , sin(2 𝜋𝑠/𝑛 )) provides exact optimal policies to reach any arbitrary state.

Original: On a 𝑑 -dimensional grid 𝑆 = { 0 , . . . , 𝑛 − 1 } 𝑑 , just 2 𝑑 features (a sine and cosine in each direction) are sufﬁcient to reach any of the 𝑛 𝑑 goal states via SFs.
Revised: On a 𝑑 -dimensional grid 𝑆 = { 0 , . . . , 𝑛 − 1 } 𝑑 , just 2 𝑑 features (a sine and cosine in each direction) are sufﬁcient to reach any of the 𝑛 𝑑 goal states via SFs.

Original: Goal-reaching is only a subset of possible RL tasks, but this clearly shows that some features are better than others.
Revised: Goal-reaching is only a subset of possible RL tasks, but this clearly shows that some features are better than others.

Original: 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;
Revised: 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;

Original: Machado et al., 2017).
Revised: Machado et al., 2017).

Original: FB-based methods are theoretically known to learn such eigenfunctions (Blier et al., 2021).
Revised: FB-based methods are theoretically known to learn such eigenfunctions (Blier et al., 2021).

Original: Yet a precise theoretical link to downstream performance is still lacking.
Revised: Yet a precise theoretical link to downstream performance is still lacking.

Original: Are these ﬁnite-rank models reasonable?
Revised: Are these ﬁnite-rank models reasonable?

Original: FB crucially relies on a ﬁnite-rank model of ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 , while some SF variants above rely on ﬁnite-rank models of 𝑃 or the corresponding Laplacian.
Revised: FB crucially relies on a ﬁnite-rank model of ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 , while some SF variants above rely on ﬁnite-rank models of 𝑃 or the corresponding Laplacian.

Original: It turns out such approximations are very different for 𝑃 or for ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 .
Revised: It turns out such approximations are very different for 𝑃 or for ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 .

Original: Unfortunately, despite the popularity of low-rank 𝑃 assumptions in the theoretical literature, 𝑃 𝜋 is always close to Id in situations where 𝑠 𝑡 +1 is close to 𝑠 𝑡 , such as any continuous-time physical system (Appendix E).
Revised: Unfortunately, despite the popularity of low-rank 𝑃 assumptions in the theoretical literature, 𝑃 𝜋 is always close to Id in situations where 𝑠 𝑡 +1 is close to 𝑠 𝑡 , such as any continuous-time physical system (Appendix D).

Original: Any low-rank model of 𝑃 𝜋 will be poor; actually 𝑃 𝜋 is better modeled as
Revised: Any low-rank model of 𝑃 𝜋 will be poor; actually 𝑃 𝜋 is better modeled as

Original: Id − low-rank, thus approximating the Laplacian.
Revised: Id − low-rank, thus approximating the Laplacian.

Original: On the other hand, 𝑃 𝑡𝜋 with large 𝑡 gets close to rank one under weak assumptions (ergodicity), as 𝑃 𝑡𝜋 converges to an equilibrium distribution when 𝑡 → ∞ .
Revised: On the other hand, 𝑃 𝑡𝜋 with large 𝑡 gets close to rank one under weak assumptions (ergodicity), as 𝑃 𝑡𝜋 converges to an equilibrium distribution when 𝑡 → ∞ .

Original: Thus the spectrum of the successor measures ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is usually more spread out (details and examples in Appendix E), and a low-rank model makes sense.
Revised: Thus the spectrum of the successor measures ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is usually more spread out (details and examples in Appendix D), and a low-rank model makes sense.

Original: The eigenvalues of 𝑃 𝜋 are close to 

Revised: The eigenvalues of 𝑃 𝜋 are close to

Original: and there is little signal to differentiate between eigenvectors, but differences become clearer over time on 𝑃 𝑡𝜋 .
Revised: and there is little signal to differentiate between eigenvectors, but differences become clearer over time on 𝑃 𝑡 𝜋 .

Original: This may explain the better performance of FB and LRA-SR compared to LRA-P .
Revised: This may explain the better performance of FB and LRA-SR compared to LRA-P .

Original: Conclusions

Original: Zero-shot RL methods avoid test-time planning by summarizing long-term state-state relationships for particular policies.

Original: We systematically tested forward-backward representations and many new models of successor features on zero-shot tasks.

Original: We also uncovered algebraic links between SFs,

Original: FB, contrastive learning, and spectral methods.

Original: Overall, SFs suffer from their dependency to basic feature construction, with only Laplacian eigenfunctions working reliably.

Original: Notably, SFs using generic encoder-type feature learning fail if the initial representation of states is already reasonable but does not provide optimal policies: planning-speciﬁc features are necessary for SFs.

Original: Forward-backward representations were best across the board, and provide reasonable zero-shot RL performance.

Original: A Limitations
Revised: Limitations.

Original: First, these experiments are still small-scale and performance is not perfect, so there is space for improvement.
Revised: First, these experiments are still small-scale and performance is not perfect, so there is space for improvement.

Original: Also, all the environments tested here were deterministic: all algorithms still make sense in stochastic environments, but experimental conclusions may differ.
Revised: Also, all the environments tested here were deterministic: all algorithms still make sense in stochastic environments, but experimental conclusions may differ.

Original: 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 ofﬂine RL method.
Revised: 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 ofﬂine RL method.

Original: These zero-shot RL algorithms learn to summarize the long-term future for a wide range of policies (though without synthesizing trajectories).
Revised: These zero-shot RL algorithms learn to summarize the long-term future for a wide range of policies (though without synthesizing trajectories).

Original: This is a lot: in contrast, world models only learn policyindependent one-step transitions.
Revised: This is a lot: in contrast, world models only learn policyindependent one-step transitions.

Original: So the question remains of how far this can scale.
Revised: So the question remains of how far this can scale.

Original: A priori, there is no restriction on the inputs (e.g., images, state history...).
Revised: A priori, there is no restriction on the inputs (e.g., images, state history...).

Original: Still, for large problems, some form of prior seems unavoidable.
Revised: Still, for large problems, some form of prior seems unavoidable.

Original: For SFs, priors can be integrated in 𝜙 , but rewards must be linear in 𝜙 .
Revised: For SFs, priors can be integrated in 𝜙 , but rewards must be linear in 𝜙 .

Original: For FB, priors can be integrated in 𝐵 ’s input.
Revised: For FB, priors can be integrated in 𝐵 ’s input.

Original: Touati & Ollivier (2021) use FB with pixel-based inputs for 𝐹 , but only the agent’s ( 𝑥, 𝑦 ) position for 𝐵 ’s input: this recovers all rewards that are functions of ( 𝑥, 𝑦 ) (linear or not).
Revised: Touati & Ollivier (2021) use FB with pixel-based inputs for 𝐹 , but only the agent’s ( 𝑥, 𝑦 ) position for 𝐵 ’s input: this recovers all rewards that are functions of ( 𝑥, 𝑦 ) (linear or not).

Original: Breaking the symmetry of 𝐹 and 𝐵 reduces the strain on the model by restricting predictions to fewer variables, such as an agent’s future state instead of the full environment.
Revised: Breaking the symmetry of 𝐹 and 𝐵 reduces the strain on the model by restricting predictions to fewer variables, such as an agent’s future state instead of the full environment.

Revised:  Conclusions

Original: B Sketch of Proof of Theorems 1 and 
Revised: Zero-shot RL methods avoid test-time planning by summarizing long-term state-state relationships for particular policies. We systematically tested forward-backward representations and many new models of successor features on zero-shot tasks. We also uncovered algebraic links between SFs, FB, contrastive learning and spectral methods.

Revised: Overall, SFs suffer from their dependency to basic feature construction, with only Laplacian eigenfunctions working reliably.

Revised: Notably, SFs using generic encoder-type feature learning fail if the initial representation of states is already reasonable but does not provide optimal policies: planning-speciﬁc features are necessary for SFs.

Revised: Forward-backward representations were best across the board, and provide reasonable zero-shot RL performance.

Revised: A Sketch of Proof of Theorems 1 and

Original: To provide an intuition behind Theorems 1 and 2 on SFs and FB, we include here a sketch of proof in the ﬁnite state case.
Revised: To provide an intuition behind Theorems 1 and 2 on SFs and FB, we include here a sketch of proof in the ﬁnite state case.

Original: Full proofs in the general case can be found in the Appendix of Touati & Ollivier (2021).
Revised: Full proofs in the general case can be found in the Appendix of Touati & Ollivier (2021).

Original: For Theorem 1 (SFs), let us assume that the reward is linear in the features 𝜙 , namely, 𝑟 ( 𝑠 ) =
Revised: For Theorem 1 (SFs), let us assume that the reward is linear in the features 𝜙 , namely, 𝑟 ( 𝑠 ) =

Original: 𝜙 ( 𝑠 ) ⊤ 𝑤 for some 𝑤 ∈ R 𝑑 .
Revised: 𝜙 ( 𝑠 ) ⊤ 𝑤 for some 𝑤 ∈ R 𝑑 .

Original: Then by deﬁnition, 𝑧 𝑟 =
Revised: Then by deﬁnition, 𝑧 𝑟 =

Original: 𝑤 since 𝑧 𝑟 is the linear regression of the reward on the features (assuming features are linearly independent).
Revised: 𝑤 since 𝑧 𝑟 is the linear regression of the reward on the features (assuming features are linearly independent).

Original: Using the deﬁnition (4) of the successor features 𝜓 , and taking the dot product with 𝑧 𝑟 , we obtain𝜓 ( 𝑠 0 , 𝑎 0 , 𝑧 𝑟 ) ⊤ 𝑧 𝑟 [Equation] E
Revised: Using the deﬁnition (4) of the successor features 𝜓 , and taking the dot product with 𝑧 𝑟 , we obtain

Original: [Equation]
Revised: [Equation]

Original: since 𝑟 = 𝜙 ⊤ 𝑤 .
Revised: since 𝑟 = 𝜙 ⊤ 𝑤 .

Original: This means that 𝜓 ( 𝑠 0 , 𝑎 0 , 𝑧 𝑟 ) ⊤ 𝑧 𝑟 is the 𝑄 -function of reward 𝑟 for policy 𝜋 𝑧 𝑟 .
Revised: This means that 𝜓 ( 𝑠 0 , 𝑎 0 , 𝑧 𝑟 ) ⊤ 𝑧 𝑟 is the 𝑄 -function of reward 𝑟 for policy 𝜋 𝑧 𝑟 .

Original: At the same time, by the deﬁnition (4), 𝜋 𝑧 𝑟 is deﬁned as the argmax of 𝜓 ( 𝑠 0 , 𝑎 0 , 𝑧 𝑟 )
Revised: At the same time, by the deﬁnition (4), 𝜋 𝑧 𝑟 is deﬁned as the argmax of 𝜓 ( 𝑠 0 , 𝑎 0 , 𝑧 𝑟 )

Original: ⊤ 𝑧 𝑟 .
Revised: ⊤ 𝑧 𝑟 .

Original: Therefore, the policy 𝜋 𝑧 𝑟 is the argmax of its own 𝑄 -function, meaning it is the optimal policy for reward 𝑟 .
Revised: Therefore, the policy 𝜋 𝑧 𝑟 is the argmax of its own 𝑄 -function, meaning it is the optimal policy for reward 𝑟 .

Original: For Theorem 2 (FB), let us assume that FB perfectly satisﬁes the training criterion (5), namely, 𝑀 𝜋 𝑧 =
Revised: For Theorem 2 (FB), let us assume that FB perfectly satisﬁes the training criterion (5), namely, 𝑀 𝜋 𝑧 =

Original: 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) in matrix form.
Revised: 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) in matrix form.

Original: Thanks to the deﬁnition (1) of successor representations 𝑀 𝜋 , for any policy 𝜋 𝑧 , the 𝑄 -function for the reward 𝑟 can be written as 𝑄 𝜋 𝑧 𝑟
Revised: Thanks to the deﬁnition (1) of successor representations 𝑀 𝜋 , for any policy 𝜋 𝑧 , the 𝑄 -function for the reward 𝑟 can be written as 𝑄 𝜋 𝑧 𝑟

Original: =
Revised: =

Original: 𝑀 𝜋 𝑧 𝑟 in matrix form.
Revised: 𝑀 𝜋 𝑧 𝑟 in matrix form.

Original: This is equal to 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) 𝑟 .
Revised: This is equal to 𝐹 ⊤ 𝑧 𝐵 diag( 𝜌 ) 𝑟 .

Original: Thus, if we deﬁne 𝑧 𝑟 := 𝐵 diag( 𝜌 ) 𝑟 = E 𝑠 ∼ 𝜌
Revised: Thus, if we deﬁne 𝑧 𝑟 := 𝐵 diag( 𝜌 ) 𝑟 = E 𝑠 ∼ 𝜌

Original: [ 𝐵 ( 𝑠 ) 𝑟 ( 𝑠 )] , we obtain 𝑄 𝜋 𝑧 𝑟
Revised: [ 𝐵 ( 𝑠 ) 𝑟 ( 𝑠 )] , we obtain 𝑄 𝜋 𝑧 𝑟

Original: =
Revised: =

Original: 𝐹 ⊤ 𝑧 𝑧 𝑟 for any 𝑧 ∈ R 𝑑 .
Revised: 𝐹 ⊤ 𝑧 𝑧 𝑟 for any 𝑧 ∈ R 𝑑 .

Original: In particular, the latter holds for 𝑧 =
Revised: In particular, the latter holds for 𝑧 =

Original: 𝑧 𝑟 as well: 𝐹 ⊤ 𝑧 𝑟 𝑧 𝑟 is the 𝑄 -function of 𝜋 𝑧 𝑟 .
Revised: 𝑧 𝑟 as well: 𝐹 ⊤ 𝑧 𝑟 𝑧 𝑟 is the 𝑄 -function of 𝜋 𝑧 𝑟 .

Original: Again, the policies 𝜋 𝑧 are deﬁned in (5) as the greedy policies of 𝐹 ⊤ 𝑧 𝑧 , for any 𝑧 .
Revised: Again, the policies 𝜋 𝑧 are deﬁned in (5) as the greedy policies of 𝐹 ⊤ 𝑧 𝑧 , for any 𝑧 .

Original: Therefore, 𝜋 𝑧 𝑟 is the argmax of its own 𝑄 -function.
Revised: Therefore, 𝜋 𝑧 𝑟 is the argmax of its own 𝑄 -function.

Original: Hence, 𝜋 𝑧 𝑟 is the optimal policy for the reward 𝑟 .
Revised: Hence, 𝜋 𝑧 𝑟 is the optimal policy for the reward 𝑟 .

Original: C Derivation of the Forward-Backward Loss
Revised: B Derivation of the Forward-Backward Loss

Original: Here we quickly derive the loss (8) used to train 𝐹 and 𝐵 such that 𝑀 𝜋
Revised: Here we quickly derive the loss (8) used to train 𝐹 and 𝐵 such that 𝑀 𝜋

Original: ( 𝑠, 𝑎, d 𝑠 ′ )
Revised: ( 𝑠, 𝑎, d 𝑠 ′ )

Original: ≈
Revised: ≈

Original: 𝐹 ( 𝑠, 𝑎 ) ⊤ 𝐵 ( 𝑠 ′ )
Revised: 𝐹 ( 𝑠, 𝑎 ) ⊤ 𝐵 ( 𝑠 ′ )

Original: 𝜌 (d 𝑠 ′ ) .
Revised: 𝜌 (d 𝑠 ′ ) .

Original: Training is based on the Bellman equation satisﬁed by 𝑀 𝜋 .
Revised: Training is based on the Bellman equation satisﬁed by 𝑀 𝜋 .

Original: This holds separately for each policy parameter 𝑧 , so in this section we omit 𝑧 for simplicity.
Revised: This holds separately for each policy parameter 𝑧 , so in this section we omit 𝑧 for simplicity.

Original: Here 𝜌 (d 𝑠 ′ ) is the distribution of states in the dataset.
Revised: Here 𝜌 (d 𝑠 ′ ) is the distribution of states in the dataset.

Original: Importantly, the resulting loss does not require to know this measure 𝜌 (d 𝑠 ′ ) , only to be able to sample states
Revised: Importantly, the resulting loss does not require to know this measure 𝜌 (d 𝑠 ′ ) , only to be able to sample states

Original: 𝑠 ∼ 𝜌 from the dataset.
Revised: 𝑠 ∼ 𝜌 from the dataset.

Original: The successor measure 𝑀 𝜋 satisﬁes a Bellman-like equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 , as matrices in the ﬁnite case and as measures in the general case (Blier et al., 2021).
Revised: The successor measure 𝑀 𝜋 satisﬁes a Bellman-like equation 𝑀 𝜋 = 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 , as matrices in the ﬁnite case and as measures in the general case (Blier et al., 2021).

Original: We can learn FB by iteratively minimizing the Bellman residual 𝑀 𝜋 − ( 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 ) on the parametric model 𝑀 =
Revised: We can learn FB by iteratively minimizing the Bellman residual 𝑀 𝜋 − ( 𝑃 + 𝛾𝑃 𝜋 𝑀 𝜋 ) on the parametric model 𝑀 =

Original: 𝐹 ⊤ 𝐵𝜌 .
Revised: 𝐹 ⊤ 𝐵𝜌 .

Original: 𝑀 𝜋 ( 𝑠, 𝑎, d 𝑠 ′ ) is a measure on 𝑠 ′ for each ( 𝑠, 𝑎 ) , so it is not obvious how to measure the size of the
Revised: 𝑀 𝜋 ( 𝑠, 𝑎, d 𝑠 ′ ) is a measure on 𝑠 ′ for each ( 𝑠, 𝑎 ) , so it is not obvious how to measure the size of the

Original: Bellman residual.
Revised: Bellman residual.

Original: In general, we can deﬁne a norm on such objects 𝑀 by taking the density with respect to the reference measure 𝜌 ,
Revised: In general, we can deﬁne a norm on such objects 𝑀 by taking the density with respect to the reference measure 𝜌 ,

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: where 𝑀 ( 𝑠,𝑎, d 𝑠 ′ )
Revised: where 𝑀 ( 𝑠,𝑎, d 𝑠 ′ )

Original: 𝜌 (d 𝑠 ′ ) is the density of 𝑀 with respect to 𝜌 .
Revised: 𝜌 (d 𝑠 ′ ) is the density of 𝑀 with respect to 𝜌 .

Original: For ﬁnite states, 𝑀 is a matrix 𝑀 𝑠𝑎𝑠 ′
Revised: For ﬁnite states, 𝑀 is a matrix 𝑀 𝑠𝑎𝑠 ′

Original: and this is just a 𝜌 -weighted Frobenius matrix norm, ‖ 𝑀
Revised: and this is just a 𝜌 -weighted Frobenius matrix norm, ‖ 𝑀

Original: ‖ 2 𝜌 =
Revised: ‖ 2 𝜌 =

Original: ∑︀ 𝑠𝑎𝑠 ′
Revised: ∑︀ 𝑠𝑎𝑠 ′

Original: 𝑀 2 𝑠𝑎𝑠 ′
Revised: 𝑀 2 𝑠𝑎𝑠 ′

Original: 𝜌 ( 𝑠, 𝑎 ) /𝜌
Revised: 𝜌 ( 𝑠, 𝑎 ) /𝜌

Original: ( 𝑠 ′ ) .
Revised: ( 𝑠 ′ ) .

Original: (This is also how we proceed to learn a low-rank approximation of 𝑃 in (16).)
Revised: (This is also how we proceed to learn a low-rank approximation of 𝑃 in (16).)

Original: We deﬁne the loss on 𝐹 and 𝐵 as the norm of the Bellman residual on 𝑀 for the model 𝑀 =
Revised: We deﬁne the loss on 𝐹 and 𝐵 as the norm of the Bellman residual on 𝑀 for the model 𝑀 =

Original: 𝐹 ⊤ 𝐵𝜌 .
Revised: 𝐹 ⊤ 𝐵𝜌 .

Original: As usual in temporal difference learning, we use ﬁxed, non-trainable target networks ¯ 𝐹 and ¯ 𝐵 for the right-hand-side of the Bellman equation.
Revised: As usual in temporal difference learning, we use ﬁxed, non-trainable target networks ¯ 𝐹 and ¯ 𝐵 for the right-hand-side of the Bellman equation.

Original: Thus, the Bellman residual is 𝐹 ⊤ 𝐵𝜌 − (︀ 𝑃 + 𝛾𝑃 𝜋 ¯ 𝐹 ⊤ ¯ 𝐵𝜌 )︀ , and the loss is
Revised: Thus, the Bellman residual is 𝐹 ⊤ 𝐵𝜌 − (︀ 𝑃 + 𝛾𝑃 𝜋 ¯ 𝐹 ⊤ ¯ 𝐵𝜌 )︀ , and the loss is

Original: [Equation]
Revised: [Equation]

Original: When computing the norm ‖·‖ 𝜌 , the denominator 𝜌 cancels out with the 𝐹 ⊤ 𝐵𝜌 terms, but we are left with a 𝑃/𝜌 term.
Revised: When computing the norm ‖·‖ 𝜌 , the denominator 𝜌 cancels out with the 𝐹 ⊤ 𝐵𝜌 terms, but we are left with a 𝑃/𝜌 term.

Original: This term can still be integrated, because integrating 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) /𝜌
Revised: This term can still be integrated, because integrating 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) /𝜌

Original: (d 𝑠 ′ ) under 𝑠 ′ ∼
Revised: (d 𝑠 ′ ) under 𝑠 ′ ∼

Original: 𝜌 is equivalent to directly integrating under 𝑠 ′ ∼ 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) , namely, integrating under transitions ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) in the environment.
Revised: 𝜌 is equivalent to directly integrating under 𝑠 ′ ∼ 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) , namely, integrating under transitions ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 +1 ) in the environment.

Original: This plays out as follows:
Revised: This plays out as follows:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: where Const is a constant term that we can discard since it does not depend on 𝐹 and 𝐵 .
Revised: where Const is a constant term that we can discard since it does not depend on 𝐹 and 𝐵 .

Original: Apart from the discarded constant term, all terms in this ﬁnal expression can be sampled from the dataset.
Revised: Apart from the discarded constant term, all terms in this ﬁnal expression can be sampled from the dataset.

Original: Note that we have a − 2 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 )
Revised: Note that we have a − 2 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 )

Original: ⊤ 𝐵 ( 𝑠 𝑡 +1 ) term where Touati & Ollivier (2021) have a − 2 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 )
Revised: ⊤ 𝐵 ( 𝑠 𝑡 +1 ) term where Touati & Ollivier (2021) have a − 2 𝐹 ( 𝑠 𝑡 , 𝑎 𝑡 )

Original: ⊤ 𝐵 ( 𝑠 𝑡 ) term: this is because we deﬁne successor representations (1) using 𝑠 𝑡
Revised: ⊤ 𝐵 ( 𝑠 𝑡 ) term: this is because we deﬁne successor representations (1) using 𝑠 𝑡

Original: +1 while
Revised: +1 while

Original: Touati & Ollivier (2021) use 𝑠 𝑡 .
Revised: Touati & Ollivier (2021) use 𝑠 𝑡 .

Original: Finally, an auxiliary loss is used to normalize 𝐵 so that E 𝑠 ∼ 𝜌
Revised: Finally, an auxiliary loss is used to normalize 𝐵 so that E 𝑠 ∼ 𝜌

Original: [ 𝐵 ( 𝑠 ) 𝐵 ( 𝑠 ) ⊤ ]
Revised: [ 𝐵 ( 𝑠 ) 𝐵 ( 𝑠 ) ⊤ ]

Original: ≈ Id .
Revised: ≈ Id .

Original: This loss is
Revised: This loss is

Original: [Equation]
Revised: [Equation]

Original: norm
Revised: norm

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 26) (The more complex expression in Touati & Ollivier (2021) has the same gradients up to a factor 4 .)
Revised: 26) (The more complex expression in Touati & Ollivier (2021) has the same gradients up to a factor 4 .)

Original: D Relationship Between Contrastive Loss and SVD of Successor Measures
Revised: C Relationship Between Contrastive Loss and SVD of Successor Measures

Original: Here we prove the precise relationship between the contrastive loss (17) and SVDs of the successor measure of the exploration policy.
Revised: Here we prove the precise relationship between the contrastive loss (17) and SVDs of the successor measure of the exploration policy.

Original: Intuitively, both methods push states together if they lie on the same trajectory, by increasing the dot product between the representations of 𝑠 𝑡 and 𝑠 𝑡 + 𝑘 .
Revised: Intuitively, both methods push states together if they lie on the same trajectory, by increasing the dot product between the representations of 𝑠 𝑡 and 𝑠 𝑡 + 𝑘 .

Original: This is formalized as follows.
Revised: This is formalized as follows.

Original: Let 𝜋 be the policy used to produce the trajectories in the dataset.
Revised: Let 𝜋 be the policy used to produce the trajectories in the dataset.

Original: Deﬁne
Revised: Deﬁne

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: where 𝛾 CL is the parameter of the geometric distribution used to choose 𝑘 when sampling 𝑠 𝑡 and 𝑠 𝑡 + 𝑘 .
Revised: where 𝛾 CL is the parameter of the geometric distribution used to choose 𝑘 when sampling 𝑠 𝑡 and 𝑠 𝑡 + 𝑘 .

Original: ˜ 𝑀 is a stochastic matrix in the discrete case, and a probability measure over 𝑆 in the general case: it is the normalized version of the successor measure (3) with 𝜋 the exploration policy.
Revised: ˜ 𝑀 is a stochastic matrix in the discrete case, and a probability measure over 𝑆 in the general case: it is the normalized version of the successor measure (3) with 𝜋 the exploration policy.

Original: By construction, the distribution of 𝑠 𝑡 + 𝑘 knowing 𝑠 𝑡 with 𝑘 ∼ Geom (1 − 𝛾 CL ) is described by ˜ 𝑀 .
Revised: By construction, the distribution of 𝑠 𝑡 + 𝑘 knowing 𝑠 𝑡 with 𝑘 ∼ Geom (1 − 𝛾 CL ) is described by ˜ 𝑀 .

Original: Therefore, we can rewrite the loss as
Revised: Therefore, we can rewrite the loss as

Original: [Equation]
Revised: [Equation]

Original: Assume that 𝜙 and 𝜇 are centered with unit norm,
Revised: Assume that 𝜙 and 𝜇 are centered with unit norm,

Original: namely, ‖ 𝜙 ( 𝑠 ) ‖ 2 = 1 and E 𝑠 ∼ 𝜌 𝜙 ( 𝑠 ) = 0 and likewise for 𝜇 .
Revised: namely, ‖ 𝜙 ( 𝑠 ) ‖ 2 = 1 and E 𝑠 ∼ 𝜌 𝜙 ( 𝑠 ) = 0 and likewise for 𝜇 .

Original: With unit norm, the cosine becomes just a dot product, and the loss is
Revised: With unit norm, the cosine becomes just a dot product, and the loss is

Original: [Equation]
Revised: [Equation]

Original: A second-order Taylor expansion provides
Revised: A second-order Taylor expansion provides

Original: [Equation]
Revised: [Equation]

Original: and therefore, with E 𝜙 = E 𝜇 = 0 , the loss is approximately
Revised: and therefore, with E 𝜙 = E 𝜇 = 0 , the loss is approximately

Original: [Equation]
Revised: [Equation]

Original: where the constant term does not depend on 𝜙 and 𝜇 .
Revised: where the constant term does not depend on 𝜙 and 𝜇 .

Original: This is minimized when 𝜙 ⊤ 𝜇 is the SVD of ˜ 𝑀/𝜌 in the 𝐿 2 ( 𝜌 ) norm.
Revised: This is minimized when 𝜙 ⊤ 𝜇 is the SVD of ˜ 𝑀/𝜌 in the 𝐿 2 ( 𝜌 ) norm.

Original: E Do Finite-Rank Models on the Transition Matrix and on Successor Measures Make Sense?
Revised: D Do Finite-Rank Models on the Transition Matrix and on Successor Measures Make Sense?

Original: It turns out ﬁnite-rank models are very different for 𝑃 or for ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 .
Revised: It turns out ﬁnite-rank models are very different for 𝑃 or for ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 .

Original: In typical situations, the spectrum of 𝑃 is concentrated around 1 while that of ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is much more spread-out.
Revised: In typical situations, the spectrum of 𝑃 is concentrated around 1 while that of ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is much more spread-out.

Original: Despite the popularity of low-rank 𝑃 in theoretical RL works, 𝑃 is never low-rank in continuoustime systems: then 𝑃 is actually always close to the identity.
Revised: Despite the popularity of low-rank 𝑃 in theoretical RL works, 𝑃 is never low-rank in continuoustime systems: then 𝑃 is actually always close to the identity.

Original: Generally speaking, 𝑃 cannot be low-rank if most actions have a small effect.
Revised: Generally speaking, 𝑃 cannot be low-rank if most actions have a small effect.

Original: By deﬁnition, for any feature function 𝜙 , ( 𝑃 𝜋 𝜙 )( 𝑠 ) =
Revised: By deﬁnition, for any feature function 𝜙 , ( 𝑃 𝜋 𝜙 )( 𝑠 ) =

Original: E [ 𝜙 ( 𝑠 𝑡 +1 ) | 𝑠 𝑡 = 𝑠 ] .
Revised: E [ 𝜙 ( 𝑠 𝑡 +1 ) | 𝑠 𝑡 = 𝑠 ] .

Original: Intuitively, if actions have a small effect, then 𝑠 𝑡 +1 is close to 𝑠 𝑡 , and 𝜙 ( 𝑠 𝑡 +1 )
Revised: Intuitively, if actions have a small effect, then 𝑠 𝑡 +1 is close to 𝑠 𝑡 , and 𝜙 ( 𝑠 𝑡 +1 )

Original: ≈ 𝜙 ( 𝑠 𝑡 ) for continuous 𝜙 .
Revised: ≈ 𝜙 ( 𝑠 𝑡 ) for continuous 𝜙 .

Original: This means that 𝑃 𝜋 𝜙 is close to 𝜙 , so that 𝑃 𝜋 is close to the identity on a large subspace of feature functions 𝜙 .
Revised: This means that 𝑃 𝜋 𝜙 is close to 𝜙 , so that 𝑃 𝜋 is close to the identity on a large subspace of feature functions 𝜙 .

Original: In the theory of continuous-time Markov processes, the time- 𝑡 transition kernel is given by 𝑃 𝑡 = 𝑒 𝑡𝐴 with 𝐴 the inﬁnitesimal generator of the process (Levin et al., 2009, §20.1)
Revised: In the theory of continuous-time Markov processes, the time- 𝑡 transition kernel is given by 𝑃 𝑡 = 𝑒 𝑡𝐴 with 𝐴 the inﬁnitesimal generator of the process (Levin et al., 2009, §20.1)

Original: (Øksendal, 1998, §8.1), hence 𝑃 𝑡 is Id + 𝑂 ( 𝑡 ) for small timesteps 𝑡 .
Revised: (Øksendal, 1998, §8.1), hence 𝑃 𝑡 is Id + 𝑂 ( 𝑡 ) for small timesteps 𝑡 .

Original: In general, the transition matrix 𝑃 is better modeled as Id + low-rank, namely, a low-rank assumption on the
Revised: In general, the transition matrix 𝑃 is better modeled as Id + low-rank, namely, a low-rank assumption on the

Original: Laplacian.
Revised: Laplacian.

Original: On the other hand, though ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is never exactly low-rank (it is invertible), it has meaningful low-rank approximations under weak assumptions.
Revised: On the other hand, though ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is never exactly low-rank (it is invertible), it has meaningful low-rank approximations under weak assumptions.

Original: For large 𝑡 , 𝑃 𝑡𝜋 becomes rank-one under weak assumptions (ergodicity), as it converges to the equilibrium distribution of the transition kernel.
Revised: For large 𝑡 , 𝑃 𝑡𝜋 becomes rank-one under weak assumptions (ergodicity), as it converges to the equilibrium distribution of the transition kernel.

Original: For large 𝛾 , the sum ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is dominated by large 𝑡 .
Revised: For large 𝛾 , the sum ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 is dominated by large 𝑡 .

Original: Most eigenvalues of 𝑃 𝜋 are close to 1 , but taking powers 𝑃 𝑡𝜋 sharpens the differences between eigenvalues: with 𝛾 close to 1 , going from 𝑃 𝜋 to ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 =
Revised: Most eigenvalues of 𝑃 𝜋 are close to 1 , but taking powers 𝑃 𝑡𝜋 sharpens the differences between eigenvalues: with 𝛾 close to 1 , going from 𝑃 𝜋 to ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 =

Original: (Id − 𝛾𝑃 𝜋 ) − 1 changes an eigenvalue 1 − 𝜀 into 1 /𝜀 .
Revised: (Id − 𝛾𝑃 𝜋 ) − 1 changes an eigenvalue 1 − 𝜀 into 1 /𝜀 .

Original: In short, on 𝑃 itself, there is little learning signal to differentiate between eigenvectors, but differences become visible over time.
Revised: In short, on 𝑃 itself, there is little learning signal to differentiate between eigenvectors, but differences become visible over time.

Original: This may explain why FB works better than low-rank decompositions directly based on 𝑃 or the Laplacian.
Revised: This may explain why FB works better than low-rank decompositions directly based on 𝑃 or the Laplacian.

Original: For instance, consider the nearest-neighbor random walk on a length- 𝑛 cycle { 0 , 1 , . .
Revised: For instance, consider the nearest-neighbor random walk on a length- 𝑛 cycle { 0 , 1 , . .

Original: , 𝑛 −mod 𝑛 } , namely, moving in dimension 1 .
Revised: , 𝑛 −mod 𝑛 } , namely, moving in dimension 1 .

Original: (This extends to any-dimensional grids.)
Revised: (This extends to any-dimensional grids.)

Original: The associated 𝑃 𝜋 is not low-rank in any reasonable sense: the corresponding stochastic matrix is concentrated around the diagonal, and many eigenvalues are close to 1 .
Revised: The associated 𝑃 𝜋 is not low-rank in any reasonable sense: the corresponding stochastic matrix is concentrated around the diagonal, and many eigenvalues are close to 1 .

Original: Precisely, the eigenvalues are cos(2 𝑘𝜋/𝑛 ) with integer 𝑘 = { 0 , . .
Revised: Precisely, the eigenvalues are cos(2 𝑘𝜋/𝑛 ) with integer 𝑘 = { 0 , . .

Original: . , 𝑛/ 2 } .
Revised: . , 𝑛/ 2 } .

Original: This is ≈ 1 − 2 𝜋 2 ( 𝑘/𝑛 ) 2 when 𝑘 ≪ 𝑛 .
Revised: This is ≈ 1 − 2 𝜋 2 ( 𝑘/𝑛 ) 2 when 𝑘 ≪ 𝑛 .

Original: Half of the eigenvalues are between √ 2 / 2 and 1 .
Revised: Half of the eigenvalues are between √ 2 / 2 and 1 .

Original: However, when 𝛾 → 1 , ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 has one eigenvalue 1 / (1 − 𝛾 ) and the other eigenvalues are1 − cos(2 𝑘𝜋/𝑛 ) ≈ 𝑛 2 / 2 𝜋 2 𝑘 2 with positive integer 𝑘 : there is one large eigenvalue, then the others decrease like cst /𝑘 2 .
Revised: However, when 𝛾 → 1 , ∑︀ 𝛾 𝑡 𝑃 𝑡𝜋 has one eigenvalue 1 / (1 − 𝛾 ) and the other eigenvalues are1 − cos(2 𝑘𝜋/𝑛 ) ≈ 𝑛 2 / 2 𝜋 2 𝑘 2 with positive integer 𝑘 : there is one large eigenvalue, then the others decrease like cst /𝑘 2 .

Original: With such a spread-out spectrum, a ﬁnite-rank model makes sense.
Revised: With such a spread-out spectrum, a ﬁnite-rank model makes sense.

Original: F Which Features are Analogous Between Low-Rank 𝑃 and SFs?
Revised: E Which Features are Analogous Between Low-Rank 𝑃 and SFs?

Original: Here we explain the relationship between a low-rank model of 𝑃 and successor features.
Revised: Here we explain the relationship between a low-rank model of 𝑃 and successor features.

Original: More precisely, if transition probabilities from ( 𝑠, 𝑎 ) to 𝑠 ′ can be written exactly as 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ ) , then SFs with basic features 𝜙 := 𝜇 will provide optimal policies for any reward function (Theorem 3).
Revised: More precisely, if transition probabilities from ( 𝑠, 𝑎 ) to 𝑠 ′ can be written exactly as 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ ) , then SFs with basic features 𝜙 := 𝜇 will provide optimal policies for any reward function (Theorem 3).

Original: Indeed, under the ﬁnite-rank model 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) = 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )
Revised: Indeed, under the ﬁnite-rank model 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) = 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )

Original: 𝜌 (d 𝑠 ′ ) , rewards only matter via the reward features E 𝑠 ′ ∼ 𝜌 𝜇 ( 𝑠 ′ )
Revised: 𝜌 (d 𝑠 ′ ) , rewards only matter via the reward features E 𝑠 ′ ∼ 𝜌 𝜇 ( 𝑠 ′ )

Original: 𝑟 ( 𝑠 ′ ) : namely, two rewards with the same reward features have the same 𝑄 -function, as the dynamics produces the same expected rewards.
Revised: 𝑟 ( 𝑠 ′ ) : namely, two rewards with the same reward features have the same 𝑄 -function, as the dynamics produces the same expected rewards.

Original: Then 𝑄 -functions are linear in these reward features, and using successor features with 𝜙 :=
Revised: Then 𝑄 -functions are linear in these reward features, and using successor features with 𝜙 :=

Original: 𝜇 provides the correct 𝑄 -functions, as follows.
Revised: 𝜇 provides the correct 𝑄 -functions, as follows.

Original: Theorem 3.
Revised: Theorem 3.

Original: Assume that 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) =
Revised: Assume that 𝑃 (d 𝑠 ′ | 𝑠, 𝑎 ) =

Original: 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )
Revised: 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ )

Original: 𝜌 (d 𝑠 ′ ) .
Revised: 𝜌 (d 𝑠 ′ ) .

Original: Then successor features using the basic features 𝜙 := 𝜇 provide optimal policies for any reward function.
Revised: Then successor features using the basic features 𝜙 := 𝜇 provide optimal policies for any reward function.

Original: This is why we use 𝜇 rather than 𝜒 for the SF basic features.
Revised: This is why we use 𝜇 rather than 𝜒 for the SF basic features.

Original: This is also why we avoided the traditional notation 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 )
Revised: This is also why we avoided the traditional notation 𝑃 ( 𝑠 ′ | 𝑠, 𝑎 )

Original: = 𝜙 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ ) often used for low-rank 𝑃 , which induces a conﬂict of notation with the 𝜙 in SFs, and suggests the wrong analogy.
Revised: = 𝜙 ( 𝑠, 𝑎 ) ⊤ 𝜇 ( 𝑠 ′ ) often used for low-rank 𝑃 , which induces a conﬂict of notation with the 𝜙 in SFs, and suggests the wrong analogy.

Original: Meanwhile, 𝜒 plays a role more analogous to SFs’ 𝜓 , although for one-step transitions instead of multistep transitions as in SFs: 𝑄 -functions are linear combinations of the features 𝜒 .
Revised: Meanwhile, 𝜒 plays a role more analogous to SFs’ 𝜓 , although for one-step transitions instead of multistep transitions as in SFs: 𝑄 -functions are linear combinations of the features 𝜒 .

Original: Proof.
Revised: Proof.

Original: Let 𝜋 be any policy.
Revised: Let 𝜋 be any policy.

Original: On a ﬁnite space in matrix notation, and omitting 𝜌 for simplicity, the assumption 𝑃 = 𝜒 ⊤ 𝜇 implies 𝑃 𝜋 = 𝜒 ⊤ 𝜋 𝜇 where 𝜒 𝜋 ( 𝑠 ) :=
Revised: On a ﬁnite space in matrix notation, and omitting 𝜌 for simplicity, the assumption 𝑃 = 𝜒 ⊤ 𝜇 implies 𝑃 𝜋 = 𝜒 ⊤ 𝜋 𝜇 where 𝜒 𝜋 ( 𝑠 ) :=

Original: E 𝑎 ∼ 𝜋 ( 𝑠 ) 𝜒
Revised: E 𝑎 ∼ 𝜋 ( 𝑠 ) 𝜒

Original: ( 𝑠, 𝑎 ) are the 𝜋 -averaged features.
Revised: ( 𝑠, 𝑎 ) are the 𝜋 -averaged features.

Original: Then,
Revised: Then,

Original: [Equation]
Revised: [Equation]

Original: Thus, 𝑄 -functions are expressed as 𝑄 𝜋
Revised: Thus, 𝑄 -functions are expressed as 𝑄 𝜋

Original: ( 𝑠, 𝑎 ) =
Revised: ( 𝑠, 𝑎 ) =

Original: 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝑤 with 𝑤 =
Revised: 𝜒 ( 𝑠, 𝑎 ) ⊤ 𝑤 with 𝑤 =

Original: (︁∑︀ 𝑡 ≥ 0 𝛾 𝑡 ( 𝜇𝜒 ⊤ 𝜋 ) 𝑡 )︁ 𝜇𝑟 .
Revised: (︁∑︀ 𝑡 ≥ 0 𝛾 𝑡 ( 𝜇𝜒 ⊤ 𝜋 ) 𝑡 )︁ 𝜇𝑟 .

Original: Moreover, rewards only matter via 𝜇𝑟 .
Revised: Moreover, rewards only matter via 𝜇𝑟 .

Original: Namely, two rewards with the same 𝜇𝑟 have the same 𝑄 -function for every policy.
Revised: Namely, two rewards with the same 𝜇𝑟 have the same 𝑄 -function for every policy.

Original: In full generality on continuous spaces with 𝜌 again, the same holds with 𝜇 ( 𝑠 ) 𝜌 (d 𝑠 ) instead of 𝜇 ( 𝑠 ) , and E 𝑠 ∼ 𝜌 𝜇 ( 𝑠 ) 𝑟 ( 𝑠 ) instead of 𝜇𝑟 .
Revised: In full generality on continuous spaces with 𝜌 again, the same holds with 𝜇 ( 𝑠 ) 𝜌 (d 𝑠 ) instead of 𝜇 ( 𝑠 ) , and E 𝑠 ∼ 𝜌 𝜇 ( 𝑠 ) 𝑟 ( 𝑠 ) instead of 𝜇𝑟 .

Original: Now, let 𝑟 be any reward function, and let 𝑟 ′ be its 𝐿 2 ( 𝜌 ) -orthogonal projection onto the space generated by the features 𝜇 .
Revised: Now, let 𝑟 be any reward function, and let 𝑟 ′ be its 𝐿 2 ( 𝜌 ) -orthogonal projection onto the space generated by the features 𝜇 .

Original: By construction, 𝑟 − 𝑟 ′ is 𝐿 2 ( 𝜌 ) -orthogonal to 𝜇 , namely, E 𝜌 𝜇 ( 𝑟 − 𝑟 ′ )
Revised: By construction, 𝑟 − 𝑟 ′ is 𝐿 2 ( 𝜌 ) -orthogonal to 𝜇 , namely, E 𝜌 𝜇 ( 𝑟 − 𝑟 ′ )

Original: = 0 .
Revised: = 0 .

Original: So E 𝜌 𝜇𝑟 =
Revised: So E 𝜌 𝜇𝑟 =

Original: E 𝜌 𝜇𝑟 ′ .
Revised: E 𝜌 𝜇𝑟 ′ .

Original: Therefore, by the above, 𝑟 and 𝑟 ′ have the same 𝑄 -function for every policy.
Revised: Therefore, by the above, 𝑟 and 𝑟 ′ have the same 𝑄 -function for every policy.

Original: By deﬁnition, 𝑟 ′ lies in the linear span of the features 𝜇 .
Revised: By deﬁnition, 𝑟 ′ lies in the linear span of the features 𝜇 .

Original: By Theorem 1, SFs with features 𝜙 = 𝜇 will provide optimal policies for 𝑟 ′ .
Revised: By Theorem 1, SFs with features 𝜙 = 𝜇 will provide optimal policies for 𝑟 ′ .

Original: Since 𝑟 and 𝑟 ′ have the same 𝑄 -function for every policy, an optimal policy for 𝑟 ′ is also optimal for 𝑟 .
Revised: Since 𝑟 and 𝑟 ′ have the same 𝑄 -function for every policy, an optimal policy for 𝑟 ′ is also optimal for 𝑟 .

Original: G Relationship between Laplacian Eigenfunctions and Low-Rank 𝑃 Learning
Revised: F Relationship between Laplacian Eigenfunctions and Low-Rank 𝑃 Learning

Original: The loss (16) used to learn a low-rank model of the transition probabilities 𝑃 is an asymmetric version of the Laplacian eigenfunction loss (14) with 𝜆 = 1 .
Revised: The loss (16) used to learn a low-rank model of the transition probabilities 𝑃 is an asymmetric version of the Laplacian eigenfunction loss (14) with 𝜆 = 1 .

Original: Said equivalently, if we use the low-rank 𝑃 loss (16) constrained with 𝜒 = 𝜇 to learn a low-rank model of 𝑃 𝜋 instead of 𝑃 (with 𝜋 the exploration policy), then we get the Laplacian eigenfunction loss (14) with 𝜆 = 1 .
Revised: Said equivalently, if we use the low-rank 𝑃 loss (16) constrained with 𝜒 = 𝜇 to learn a low-rank model of 𝑃 𝜋 instead of 𝑃 (with 𝜋 the exploration policy), then we get the Laplacian eigenfunction loss (14) with 𝜆 = 1 .

Original: Indeed, set 𝜆 = 1 in (14).
Revised: Indeed, set 𝜆 = 1 in (14).

Original: Assume that the distributions of 𝑠 𝑡 and 𝑠 𝑡 +1 in the dataset are identical on average (this happens, e.g., if the dataset is made of long trajectories or if 𝜌 is close enough to the invariant distribution of the exploration policy).
Revised: Assume that the distributions of 𝑠 𝑡 and 𝑠 𝑡 +1 in the dataset are identical on average (this happens, e.g., if the dataset is made of long trajectories or if 𝜌 is close enough to the invariant distribution of the exploration policy).

Original: Then, in the Laplacian loss (14), the norms from the ﬁrst term cancel those from the second, and the Laplacian loss simpliﬁes to
Revised: Then, in the Laplacian loss (14), the norms from the ﬁrst term cancel those from the second, and the Laplacian loss simpliﬁes to

Original: [Equation]
Revised: [Equation]

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: This is the same as the low-rank loss (16) if we omit actions 𝑎 and constrain 𝜒 = 𝜇 .
Revised: This is the same as the low-rank loss (16) if we omit actions 𝑎 and constrain 𝜒 = 𝜇 .

Original: H Experimental Setup
Revised: G Experimental Setup

Original: In this section we provide additional information about our experiments.
Revised: In this section we provide additional information about our experiments.

Original: H.1 Environments
Revised: G.1 Environments

Original: All the environments considered in this paper are based on the DeepMind Control Suite (Tassa et al., 2018).
Revised: All the environments considered in this paper are based on the DeepMind Control Suite (Tassa et al., 2018).

Original: • Point-mass Maze : a 2-dimensional continuous maze with four rooms.
Revised: • Point-mass Maze : a 2-dimensional continuous maze with four rooms.

Original: The states are 4-dimensional vectors consisting of positions and velocities of the point mass ( 𝑥, 𝑦, 𝑣 𝑥 , 𝑥 𝑦 ) , and the actions are 2-dimensional vectors.
Revised: The states are 4-dimensional vectors consisting of positions and velocities of the point mass ( 𝑥, 𝑦, 𝑣 𝑥 , 𝑥 𝑦 ) , and the actions are 2-dimensional vectors.

Original: At test, we assess the performance of the agents on 20 goal-reaching tasks (5 goals in each room described by their ( 𝑥, 𝑦 ) coordinates).
Revised: At test, we assess the performance of the agents on 20 goal-reaching tasks (5 goals in each room described by their ( 𝑥, 𝑦 ) coordinates).

Original: • Walker : a planar walker.
Revised: • Walker : a planar walker.

Original: States are 24-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors.
Revised: States are 24-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors.

Original: We consider 4 different tasks at test time: walker_stand reward is a combination of terms encouraging an upright torso and some minimal torso height, while walker_walk and walker_run rewards include a component encouraging some minimal forward velocity.
Revised: We consider 4 different tasks at test time: walker_stand reward is a combination of terms encouraging an upright torso and some minimal torso height, while walker_walk and walker_run rewards include a component encouraging some minimal forward velocity.

Original: walker_flip reward includes a component encouraging some mininal angular momentum.
Revised: walker_flip reward includes a component encouraging some mininal angular momentum.

Original: • Cheetah : a running planar biped.
Revised: • Cheetah : a running planar biped.

Original: States are 17-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors.
Revised: States are 17-dimensional vectors consisting of positions and velocities of robot joints, and actions are 6-dimensional vectors.

Original: We consider 4 different tasks at test time: cheetah_walk and cheetah_run rewards are linearly proportional to the forward velecity up to some desired values: 2 m/s for walk and 10 m/s for run .
Revised: We consider 4 different tasks at test time: cheetah_walk and cheetah_run rewards are linearly proportional to the forward velecity up to some desired values: 2 m/s for walk and 10 m/s for run .

Original: Similarly, walker_walk_backward and walker_run_backward rewards encourage reaching some minimal backward velocities.
Revised: Similarly, walker_walk_backward and walker_run_backward rewards encourage reaching some minimal backward velocities.

Original: • Quadruped : a four-leg ant navigating in 3D space.
Revised: • Quadruped : a four-leg ant navigating in 3D space.

Original: States and actions are 78dimensional and 12-dimensional vectors, respectively.
Revised: States and actions are 78dimensional and 12-dimensional vectors, respectively.

Original: We consider 4 tasks at test time: quadruped_stand reward encourages an upright torso.
Revised: We consider 4 tasks at test time: quadruped_stand reward encourages an upright torso.

Original: quadruped_walk
Revised: quadruped_walk

Original: and quadruped_run include a term encouraging some minimal torso velecities.
Revised: and quadruped_run include a term encouraging some minimal torso velecities.

Original: quadruped_walk includes a term encouraging some minimal height of the center of mass.
Revised: quadruped_walk includes a term encouraging some minimal height of the center of mass.

Original: H.2 Architectures
Revised: G.2 Architectures

Original: We use the same architectures for all methods.
Revised: We use the same architectures for all methods.

Original: • The backward representation network 𝐵 ( 𝑠 ) and the feature network 𝜙 ( 𝑠 ) are represented by a feedforward neural network with three hidden layers, each with 256 units, that takes as input a state and outputs a 𝐿 2 -normalized embedding of radius √ 𝑑 .
Revised: • The backward representation network 𝐵 ( 𝑠 ) and the feature network 𝜙 ( 𝑠 ) are represented by a feedforward neural network with three hidden layers, each with 256 units, that takes as input a state and outputs a 𝐿 2 -normalized embedding of radius √ 𝑑 .

Original: • For both successor features 𝜓 ( 𝑠, 𝑎, 𝑧 ) and forward network 𝐹 ( 𝑠, 𝑎, 𝑧 ) , we ﬁrst preprocess separately ( 𝑠, 𝑎 ) and ( 𝑠, 𝑧 ) by two feedforward networks with two hidden layers (each with 1024 units) to 512-dimentional space.
Revised: • For both successor features 𝜓 ( 𝑠, 𝑎, 𝑧 ) and forward network 𝐹 ( 𝑠, 𝑎, 𝑧 ) , we ﬁrst preprocess separately ( 𝑠, 𝑎 ) and ( 𝑠, 𝑧 ) by two feedforward networks with two hidden layers (each with 1024 units) to 512-dimentional space.

Original: Then we concatenate their two outputs and pass it into another 2-layer feedforward network (each with 1024 units) to output a 𝑑 -dimensional vector.
Revised: Then we concatenate their two outputs and pass it into another 2-layer feedforward network (each with 1024 units) to output a 𝑑 -dimensional vector.

Original: • For the policy network 𝜋
Revised: • For the policy network 𝜋

Original: ( 𝑠, 𝑧 ) , we ﬁrst preprocess separately 𝑠
Revised: ( 𝑠, 𝑧 ) , we ﬁrst preprocess separately 𝑠

Original: and ( 𝑠, 𝑧 ) by two feedforward networks with two hidden layers (each with 1024 units) to 512-dimentional space.
Revised: and ( 𝑠, 𝑧 ) by two feedforward networks with two hidden layers (each with 1024 units) to 512-dimentional space.

Original: Then we concatenate their two outputs and pass it into another 2-layer feedforward network
Revised: Then we concatenate their two outputs and pass it into another 2-layer feedforward network

Original: (each with 1024 units) to output to output a 𝑑 𝐴 -dimensional vector, then we apply a Tanh activation as the action space is [ − 1 , 1] 𝑑 𝐴 .
Revised: (each with 1024 units) to output to output a 𝑑 𝐴 -dimensional vector, then we apply a Tanh activation as the action space is [ − 1 , 1] 𝑑 𝐴 .

Original: For all the architectures, we apply a layer normalization (Ba et al., 2016) and Tanh activation in the ﬁrst layer in order to standardize the states and actions.
Revised: For all the architectures, we apply a layer normalization (Ba et al., 2016) and Tanh activation in the ﬁrst layer in order to standardize the states and actions.

Original: We use Relu for the rest of layers.
Revised: We use Relu for the rest of layers.

Original: We also pre-normalized 𝑧 : 𝑧 ← √ 𝑑 𝑧 ‖ 𝑧 ‖ 2 in the input of 𝐹 , 𝜋 and 𝜓 .
Revised: We also pre-normalized 𝑧 : 𝑧 ← √ 𝑑 𝑧 ‖ 𝑧 ‖ 2 in the input of 𝐹 , 𝜋 and 𝜓 .

Original: Empirically, we observed that removing preprocessing and pass directly a concatenation of ( 𝑠, 𝑎, 𝑧 ) directly to the network leads to unstable training.
Revised: Empirically, we observed that removing preprocessing and pass directly a concatenation of ( 𝑠, 𝑎, 𝑧 ) directly to the network leads to unstable training.

Original: The same holds when we preprocess ( 𝑠, 𝑎 ) and 𝑧 instead of ( 𝑠, 𝑎 ) and ( 𝑠, 𝑧 ) , which means that the preprocessing of 𝑧 should be inﬂuenced by the current state.
Revised: The same holds when we preprocess ( 𝑠, 𝑎 ) and 𝑧 instead of ( 𝑠, 𝑎 ) and ( 𝑠, 𝑧 ) , which means that the preprocessing of 𝑧 should be inﬂuenced by the current state.

Original: 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.
Revised: 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.

Original: H.3 Sampling of 𝑧
Revised: G.3 Sampling of 𝑧

Original: We mix two methods for sampling 𝑧 :
Revised: We mix two methods for sampling 𝑧 :

Original: We sample 𝑧 uniformly in the sphere of radius √ 𝑑 in R 𝑑 (so each component of 𝑧 is of size ≈ 1 ).
Revised: We sample 𝑧 uniformly in the sphere of radius √ 𝑑 in R 𝑑 (so each component of 𝑧 is of size ≈ 1 ).

Original: We sample 𝑧 using the formula for 𝑧 𝑟 corresponding to the reward for reaching a random goal state 𝑠 in Theorems 1–2.
Revised: We sample 𝑧 using the formula for 𝑧 𝑟 corresponding to the reward for reaching a random goal state 𝑠 in Theorems 1–2.

Original: Namely, we set 𝑧 =
Revised: Namely, we set 𝑧 =

Original: 𝐵 ( 𝑠 ) for FB and 𝑧 = (︀∑︀ 𝑚𝑖 =1 𝜙 ( 𝑠 𝑖 ) 𝜙
Revised: 𝐵 ( 𝑠 ) for FB and 𝑧 = (︀∑︀ 𝑚𝑖 =1 𝜙 ( 𝑠 𝑖 ) 𝜙

Original: ( 𝑠 𝑖 ) ⊤ )︀
Revised: ( 𝑠 𝑖 ) ⊤ )︀

Original: + 𝜙 ( 𝑠 ) for SFs, where 𝑠 ∼ 𝜌 is a random state sampled from the replay buffer and 𝑠 𝑖 are states in a minibatch.
Revised: + 𝜙 ( 𝑠 ) for SFs, where 𝑠 ∼ 𝜌 is a random state sampled from the replay buffer and 𝑠 𝑖 are states in a minibatch.

Original: For the main series of results, we used a 50% mix ratio for those two methods.
Revised: For the main series of results, we used a 50% mix ratio for those two methods.

Original: Different algorithms can beneﬁt from different ratios: this is explored in Appendix
Revised: Different algorithms can beneﬁt from different ratios: this is explored in Appendix

Original: H.4 Learning the Policies 𝜋 𝑧 : Policy Network
Revised: G.4 Learning the Policies 𝜋 𝑧 : Policy Network

Original: As the action space is continuous, we could not compute the arg max over action in closed form.
Revised: As the action space is continuous, we could not compute the arg max over action in closed form.

Original: Instead, we consider a latent-conditioned policy network 𝜋 𝜂 : 𝑆 × 𝑍 → 𝐴 , and we learn the policy parameters 𝜂 by performing stochastic gradient ascent on the objective E 𝑠,𝑧 [ 𝐹 ( 𝑠, 𝜋 𝜂 ( 𝑠 ) , 𝑧 ) ⊤ 𝑧 ] for
Revised: Instead, we consider a latent-conditioned policy network 𝜋 𝜂 : 𝑆 × 𝑍 → 𝐴 , and we learn the policy parameters 𝜂 by performing stochastic gradient ascent on the objective E 𝑠,𝑧 [ 𝐹 ( 𝑠, 𝜋 𝜂 ( 𝑠 ) , 𝑧 ) ⊤ 𝑧 ] for

Original: FB or E 𝑠,𝑧 [ 𝜓 ( 𝑠, 𝜋 𝜂 ( 𝑠 ) , 𝑧 ) ⊤ 𝑧 ] for SFs.
Revised: FB or E 𝑠,𝑧 [ 𝜓 ( 𝑠, 𝜋 𝜂 ( 𝑠 ) , 𝑧 ) ⊤ 𝑧 ] for SFs.

Original: 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 𝜀 to the actions.
Revised: 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 𝜀 to the actions.

Original: Let 𝜃 1 and 𝜃 2 the parameters of two forward networks and let 𝜔 the parameters of the backward network.
Revised: Let 𝜃 1 and 𝜃 2 the parameters of two forward networks and let 𝜔 the parameters of the backward network.

Original: Let 𝜃 − 1 , 𝜃 − 2 and 𝜔 − be the parameters of their corresponding target networks.
Revised: Let 𝜃 − 1 , 𝜃 − 2 and 𝜔 − be the parameters of their corresponding target networks.

Original: Let { ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑠 𝑖 +1 ) }
Revised: Let { ( 𝑠 𝑖 , 𝑎 𝑖 , 𝑠 𝑖 +1 ) }

Original: 𝑖 ∈ 𝐼 ⊂ 𝒟 a mini-batch of size | 𝐼 | = 𝑏 of transitions and let {
Revised: 𝑖 ∈ 𝐼 ⊂ 𝒟 a mini-batch of size | 𝐼 | = 𝑏 of transitions and let {

Original: 𝑧 𝑖 } 𝑖 ∈ 𝐼 a mini-batch of size | 𝐼 | = 𝑏 of latent variables sampled according to H.3.
Revised: 𝑧 𝑖 } 𝑖 ∈ 𝐼 a mini-batch of size | 𝐼 | = 𝑏 of latent variables sampled according to G.3.

Original: The empirical version of the main FB loss in (8) is:
Revised: The empirical version of the main FB loss in (8) is:

Original: [Equation]
Revised: [Equation]

Original: where 𝜀 𝑖 is sampled from a truncated centered Gaussian with variance 𝜎 2 (for policy smoothing).
Revised: where 𝜀 𝑖 is sampled from a truncated centered Gaussian with variance 𝜎 2 (for policy smoothing).

Original: The empirical version of the auxiliary F loss in (9) is: for 𝑘 = 1 , 2 ,
Revised: The empirical version of the auxiliary F loss in (9) is: for 𝑘 = 1 , 2 ,

Original: [Equation]
Revised: [Equation]

Original: where Cov + is the pseudo-inverse of the empirical covariane matrix Cov =
Revised: where Cov + is the pseudo-inverse of the empirical covariane matrix Cov =

Original: 1 𝑏 ∑︀ 𝑖 ∈
Revised: 1 𝑏 ∑︀ 𝑖 ∈

Original: 𝐼 𝐵 𝜔 ( 𝑠 𝑖 ) 𝐵 𝜔 ( 𝑠 𝑖 ) ⊤ .
Revised: 𝐼 𝐵 𝜔 ( 𝑠 𝑖 ) 𝐵 𝜔 ( 𝑠 𝑖 ) ⊤ .

Original: We use 1 /𝑑 as regularization coefﬁcient in front of L ′ ( 𝜃 𝑘 ) .
Revised: We use 1 /𝑑 as regularization coefﬁcient in front of L ′ ( 𝜃 𝑘 ) .

Original: For policy training, the empirical loss is:
Revised: For policy training, the empirical loss is:

Original: [Equation]
Revised: [Equation]

Original: The same techniques are also used for SFs.
Revised: The same techniques are also used for SFs.

Original: H.5 Hyperparameter
Revised: G.5 Hyperparameter

Original: We summarize the hyperparameters used in our experiemnts in Table 1.
Revised: We summarize the hyperparameters used in our experiemnts in Table 1.

Original: [Table]
Revised: [Table]

Original: I Detailed Experimental Results
Revised: H Detailed Experimental Results

Original: We report the full experimental results, ﬁrst as a table (Section I.1).
Revised: We report the full experimental results, ﬁrst as a table (Section H.1).

Original: In Section I.2 we plot aggregated results for easier interpretation: ﬁrst, aggregated across all tasks with a plot of the variability for each method; second, aggregated over environments (since each environment corresponds to a trained zero-shot model); third, for each individual task but still aggregated over replay buffers.
Revised: In Section H.2 we plot aggregated results for easier interpretation: ﬁrst, aggregated across all tasks with a plot of the variability for each method; second, aggregated over environments (since each environment corresponds to a trained zero-shot model); third, for each individual task but still aggregated over replay buffers.

Original: In Section I.3 we plot all individual results per task and replay buffer.
Revised: In Section H.3 we plot all individual results per task and replay buffer.

Original: In Section I.4 we plot aggregate results split by replay buffer.
Revised: In Section H.4 we plot aggregate results split by replay buffer.

Original: I.1 Full Table of Results
Revised: H.1 Full Table of Results

Original: [Table]
Revised: [Table]

Original: I.2 Aggregate Plots of Results
Revised: H.2 Aggregate Plots of Results

Original: Here we plot the results, ﬁrst averaged over everything, then by environment averaged over the tasks of that environment, and ﬁnally by task.
Revised: Here we plot the results, ﬁrst averaged over everything, then by environment averaged over the tasks of that environment, and ﬁnally by task.

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: t u
Revised: t u

Original: [Figure]
Revised: [Figure]

Original: n
Revised: n

Original: [Figure]
Revised: [Figure]

Original: e t
Revised: e t

Original: [Figure]
Revised: [Figure]

Original: n
Revised: n

Original: [Figure]
Revised: [Figure]

Original: e t u
Revised: e t u

Original: [Figure]
Revised: [Figure]

Original: n
Revised: n

Original: [Figure]
Revised: [Figure]

Original: R e t u r n
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original: R e t u r n
Revised: R e t u r n

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: R e t u r n
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original: R e t u r n
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original: R e t u r n
Revised: R e t u r n

Original: [Figure]
Revised: [Figure]

Original: e t
Revised: e t

Original: [Figure]
Revised: [Figure]

Original: r n
Revised: r n

Original: [Figure]
Revised: [Figure]

Original: e t
Revised: e t

Original: [Figure]
Revised: [Figure]

Original: r n
Revised: r n

Original: [Figure]
Revised: [Figure]

Original: I.4 Inﬂuence of the Replay Buffer
Revised: H.4 Inﬂuence of the Replay Buffer

Original: Here we plot the inﬂuence of the replay buffer, by reporting results separated by replay buffer, but averaged over the tasks corresponding to each environment (Fig 10).
Revised: Here we plot the inﬂuence of the replay buffer, by reporting results separated by replay buffer, but averaged over the tasks corresponding to each environment (Fig 10).

Original: Overall, there is a clear failure case of the Proto buffer on the Quadruped environment: the TD supervised baseline performs poorly for all tasks in that environment.
Revised: Overall, there is a clear failure case of the Proto buffer on the Quadruped environment: the TD supervised baseline performs poorly for all tasks in that environment.

Original: Otherwise, results are broadly consistent on the different replay buffers: with a few exceptions, the same methods succeed or fail on the same environments
Revised: Otherwise, results are broadly consistent on the different replay buffers: with a few exceptions, the same methods succeed or fail on the same environments

Original: [Figure]
Revised: [Figure]

Original: In Fig.
Revised: In Fig.

Original: 11, we plot results aggregated over all tasks but split by replay buffer.
Revised: 11, we plot results aggregated over all tasks but split by replay buffer.

Original: Box plots further show variability within random seeds and environments for a given replay buffer.
Revised: Box plots further show variability within random seeds and environments for a given replay buffer.

Original: Overall method rankings are broadly consistent between RND and APS, except for CL .
Revised: Overall method rankings are broadly consistent between RND and APS, except for CL .

Original: Note that the aggregated normalized score for Proto is largely inﬂuenced by the failure on Quadruped: normalization by a very low baseline (Fig. 10, Quadruped plot for Proto), somewhat artiﬁcially pushes
Revised: Note that the aggregated normalized score for Proto is largely inﬂuenced by the failure on Quadruped: normalization by a very low baseline (Fig. 10, Quadruped plot for Proto), somewhat artiﬁcially pushes

Original: Trans high up (this shows the limit of using scores normalized by ofﬂine TD3 score), while the other methods’ rankings are more similar to RND and APS.
Revised: Trans high up (this shows the limit of using scores normalized by ofﬂine TD3 score), while the other methods’ rankings are more similar to RND and APS.

Original: FB and Lap work very well in all replay buffers, with LRA-SR and Trans a bit behind due to their feailures on some tasks.
Revised: FB and Lap work very well in all replay buffers, with LRA-SR and Trans a bit behind due to their feailures on some tasks.

Original: [Figure]
Revised: [Figure]

Original: J Analysis of The Learned Features J.1 Feature Rank
Revised: I Analysis of The Learned Features I.1 Feature Rank

Original: Here we test the hypothesis that feature collapse for some methods is responsible for some cases of bad performance.
Revised: Here we test the hypothesis that feature collapse for some methods is responsible for some cases of bad performance.

Original: This is especially relevant for Maze, where some methods may have little incentive to learn more features beyond the original two features ( 𝑥, 𝑦 )
Revised: This is especially relevant for Maze, where some methods may have little incentive to learn more features beyond the original two features ( 𝑥, 𝑦 )

Original: [Table]
Revised: [Table]

Original: We report in Table 3 the effective rank of learned features 𝐵 or 𝜙 : this is computed as in (Lyle et al., 2022), as the fraction of eigenvalues of E 𝑠 ∼ 𝜌
Revised: We report in Table 3 the effective rank of learned features 𝐵 or 𝜙 : this is computed as in (Lyle et al., 2022), as the fraction of eigenvalues of E 𝑠 ∼ 𝜌

Original: 𝐵
Revised: 𝐵

Original: 𝑠 ) 𝐵 ( 𝑠 ) ⊤ or E
Revised: 𝑠 ) 𝐵 ( 𝑠 ) ⊤ or E

Original: 𝑠 ∼ 𝜌 𝜙 ( 𝑠 ) 𝜙
Revised: 𝑠 ∼ 𝜌 𝜙 ( 𝑠 ) 𝜙

Original: ( 𝑠 ) ⊤ above a certain threshold.
Revised: ( 𝑠 ) ⊤ above a certain threshold.

Original: [Table]
Revised: [Table]

Original: The Maze environment is a clear exception: on Maze, for the Rand , AEnc , Trans , Latent and ICM methods, only about one third of the eigenvalues are above 10 − 4 .
Revised: The Maze environment is a clear exception: on Maze, for the Rand , AEnc , Trans , Latent and ICM methods, only about one third of the eigenvalues are above 10 − 4 .

Original: The other methods keep 100% of the eigenvalues above 10 − 4 .
Revised: The other methods keep 100% of the eigenvalues above 10 − 4 .

Original: This is perfectly aligned with the performance of each method on Maze.
Revised: This is perfectly aligned with the performance of each method on Maze.

Original: So rank reduction does happen for some methods.
Revised: So rank reduction does happen for some methods.

Original: This reﬂects the fact that two features ( 𝑥, 𝑦 ) already convey the necessary information about states and dynamics, but are not sufﬁcient to solve the problem via successor features.
Revised: This reﬂects the fact that two features ( 𝑥, 𝑦 ) already convey the necessary information about states and dynamics, but are not sufﬁcient to solve the problem via successor features.

Original: In the Maze environment, auto-encoder or transition models can perfectly optimize their loss just by keeping the original two features ( 𝑥, 𝑦 ) , and they have no incentive to learn other features, so the effective rank could have been 2.
Revised: In the Maze environment, auto-encoder or transition models can perfectly optimize their loss just by keeping the original two features ( 𝑥, 𝑦 ) , and they have no incentive to learn other features, so the effective rank could have been 2.

Original: These methods may beneﬁt from auxiliary losses to prevent eigenvalue collapse, similar to the orthonormalization loss used for 𝐵 .
Revised: These methods may beneﬁt from auxiliary losses to prevent eigenvalue collapse, similar to the orthonormalization loss used for 𝐵 .

Original: We did not include such losses, because we wanted to keep the same methods (autoencoders, ICM, transition model...) used in the literature.
Revised: We did not include such losses, because we wanted to keep the same methods (autoencoders, ICM, transition model...) used in the literature.

Original: But even with an auxiliary loss, keeping a full rank could be achieved just by keeping ( 𝑥, 𝑦 ) and then blowing up some additional irrelevant features.
Revised: But even with an auxiliary loss, keeping a full rank could be achieved just by keeping ( 𝑥, 𝑦 ) and then blowing up some additional irrelevant features.

Original: J.2 Embedding Visualization via t-SNE
Revised: I.2 Embedding Visualization via t-SNE

Original: 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.
Revised: 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.

Original: J.2.
Revised: I.2.

Original: AEnc , Trans , FB , LRA-P and CL all recover a picture of the maze.
Revised: AEnc , Trans , FB , LRA-P and CL all recover a picture of the maze.

Original: For the other methods the picture is less clear (notably, Rand gets a near-circle, possibly as an instance of concentration of measure theorems).
Revised: For the other methods the picture is less clear (notably, Rand gets a near-circle, possibly as an instance of concentration of measure theorems).

Original: Whether t-SNE preserves the shape of the maze does not appear to be correlated to performance:
Revised: Whether t-SNE preserves the shape of the maze does not appear to be correlated to performance:

Original: AEnc and Trans learn nice features but perform poorly, while the t-SNE of Lap is not visually interpretable but performance is good.
Revised: AEnc and Trans learn nice features but perform poorly, while the t-SNE of Lap is not visually interpretable but performance is good.

Original: [Figure]
Revised: [Figure]

Original: K Goal-Oriented Baselines
Revised: J Goal-Oriented Baselines

Original: Goal-oriented methods such as universal value functions (UVFs) (Schaul et al., 2015) learn a 𝑄function 𝑄 ( 𝑠, 𝑎, 𝑔 ) indexed by a goal description 𝑔 .
Revised: Goal-oriented methods such as universal value functions (UVFs) (Schaul et al., 2015) learn a 𝑄function 𝑄 ( 𝑠, 𝑎, 𝑔 ) indexed by a goal description 𝑔 .

Original: When taking for 𝑔 the set of all possible target states 𝑠 ′ ∈ 𝑆 , these methods can learn to reach arbitrary target states.
Revised: When taking for 𝑔 the set of all possible target states 𝑠 ′ ∈ 𝑆 , these methods can learn to reach arbitrary target states.

Original: For instance, Ma et al.
Revised: For instance, Ma et al.

Original: (2020) use a mixture of universal SFs and UVFs to learn goal-conditioned agents and policies.
Revised: (2020) use a mixture of universal SFs and UVFs to learn goal-conditioned agents and policies.

Original: In principle, such goal-oriented methods are not designed to deal with dense rewards, which are linear combinations of goals 𝑔 .
Revised: In principle, such goal-oriented methods are not designed to deal with dense rewards, which are linear combinations of goals 𝑔 .

Original: Indeed, the optimal 𝑄 -function for a linear combination of goals is not the linear combination of the optimal 𝑄 -functions.
Revised: Indeed, the optimal 𝑄 -function for a linear combination of goals is not the linear combination of the optimal 𝑄 -functions.

Original: Nevertheless, we may still try to see how such linear combinations perform.
Revised: Nevertheless, we may still try to see how such linear combinations perform.

Original: The linear combination may be applied at the level of the 𝑄 -functions, or at the level of some goal descriptors, as follows.
Revised: The linear combination may be applied at the level of the 𝑄 -functions, or at the level of some goal descriptors, as follows.

Original: Here, as an additional baseline for our experiments, we test a slight modiﬁcation of the scheme from
Revised: Here, as an additional baseline for our experiments, we test a slight modiﬁcation of the scheme from

Original: Ma et al.
Revised: Ma et al.

Original: (2020), as suggested by one reviewer.
Revised: (2020), as suggested by one reviewer.

Original: We learn two embeddings ( 𝜓, 𝑤 ) using state-reaching tasks.
Revised: We learn two embeddings ( 𝜓, 𝑤 ) using state-reaching tasks.

Original: The 𝑄 -function for reaching state 𝑠 ′ (goal 𝑔 = 𝑠 ′ ) is modeled as 𝑄
Revised: The 𝑄 -function for reaching state 𝑠 ′ (goal 𝑔 = 𝑠 ′ ) is modeled as 𝑄

Original: [Equation]
Revised: [Equation]

Original: ) trained for reward 1 𝑠 𝑡
Revised: ) trained for reward 1 𝑠 𝑡

Original: +

Revised: +

Original: =
Revised: =

Original: 𝑠 ′ .
Revised: 𝑠 ′ .

Original: A policy network 𝜋 ( 𝑠, 𝑤 ( 𝑠 ′ )) outputs the action at 𝑠 for reaching goal 𝑠 ′ .
Revised: A policy network 𝜋 ( 𝑠, 𝑤 ( 𝑠 ′ )) outputs the action at 𝑠 for reaching goal 𝑠 ′ .

Original: Thus, the training loss is
Revised: Thus, the training loss is

Original: [Equation]
Revised: [Equation]

Original: Similarly to the other baselines, the policy network 𝜋 ( 𝑠, 𝑤 ( 𝑠 ′ )) is trained by gradient ascent on the policy parameters to maximize
Revised: Similarly to the other baselines, the policy network 𝜋 ( 𝑠, 𝑤 ( 𝑠 ′ )) is trained by gradient ascent on the policy parameters to maximize

Original: [Equation]
Revised: [Equation]

Original: At test time, given a reward 𝑟 , we proceed as in FB and estimate
Revised: At test time, given a reward 𝑟 , we proceed as in FB and estimate

Original: [Equation]
Revised: [Equation]

Original: ) and use policy 𝜋 ( 𝑠, 𝑧 ) .
Revised: ) and use policy 𝜋 ( 𝑠, 𝑧 ) .

Original: This amounts to extending from goal-reaching tasks to dense tasks by linearity on 𝑤 , even though in principle, this method should only optimize for single-goal rewards.
Revised: This amounts to extending from goal-reaching tasks to dense tasks by linearity on 𝑤 , even though in principle, this method should only optimize for single-goal rewards.

Original: We use the same architectures for ( 𝜓, 𝑤 ) as for ( 𝐹, 𝐵 ) , with double networks 𝜓 1 , 𝜓 2 and policy smoothing as in (43).
Revised: We use the same architectures for ( 𝜓, 𝑤 ) as for ( 𝐹, 𝐵 ) , with double networks 𝜓 1 , 𝜓 2 and policy smoothing as in (43).

Original: The sparse reward 1 { 𝑠 𝑡 +1 = 𝑠 ′ } could suffer from large variance in continuous state spaces.
Revised: The sparse reward 1 { 𝑠 𝑡 +1 = 𝑠 ′ } could suffer from large variance in continuous state spaces.

Original: We mitigate this either by:
Revised: We mitigate this either by:

Original: • biasing the sampling of 𝑠 ′ by setting 𝑠 ′ to 𝑠 𝑡 +1 half of the time.
Revised: • biasing the sampling of 𝑠 ′ by setting 𝑠 ′ to 𝑠 𝑡 +1 half of the time.

Original: • replacing the reward by a less sparse one, 1 {‖ 𝑠 𝑡 +1 − 𝑠 ′
Revised: • replacing the reward by a less sparse one, 1 {‖ 𝑠 𝑡 +1 − 𝑠 ′

Original: ‖ 2 ≤ 𝜀 } .
Revised: ‖ 2 ≤ 𝜀 } .

Original: Table 4 reports normalized scores for each domain, for the two variants just described, trained on the
Revised: Table 4 reports normalized scores for each domain, for the two variants just described, trained on the

Original: RND replay buffer, averaged over tasks and over 10 random seeds.
Revised: RND replay buffer, averaged over tasks and over 10 random seeds.

Original: The ﬁrst variant scores 84% on Maze, 62% on Quadruped tasks, 17% on Walker tasks, and 2% on
Revised: The ﬁrst variant scores 84% on Maze, 62% on Quadruped tasks, 17% on Walker tasks, and 2% on

Original: Cheetah tasks.
Revised: Cheetah tasks.

Original: The second variant is worse overall.
Revised: The second variant is worse overall.

Original: So overall, this works well on Maze (as expected, since this is a goal-oriented problem), moderately well on Quadruped tasks (where most other methods work well), and poorly on the other environments.
Revised: So overall, this works well on Maze (as expected, since this is a goal-oriented problem), moderately well on Quadruped tasks (where most other methods work well), and poorly on the other environments.

Original: This is expected, as this method is not designed to handle dense combinations of goals.
Revised: This is expected, as this method is not designed to handle dense combinations of goals.

Original: [Table]
Revised: [Table]

Original: Relationship with FB.
Revised: Relationship with FB.

Original: The use of universal SFs in Ma et al.
Revised: The use of universal SFs in Ma et al.

Original: (2020) is quite different from the use of universal SFs in Barreto et al.
Revised: (2020) is quite different from the use of universal SFs in Barreto et al.

Original: (2017) and Borsa et al.
Revised: (2017) and Borsa et al.

Original: The latter is mathematically related to FB, as described at the end of Section 4.
Revised: The latter is mathematically related to FB, as described at the end of Section 4.

Original: 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).
Revised: 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).

Original: A key difference between FB and goal-oriented methods is the following.
Revised: A key difference between FB and goal-oriented methods is the following.

Original: Above, we use 𝜓 ( 𝑠, 𝑎, 𝑤 ( 𝑔 ))
Revised: Above, we use 𝜓 ( 𝑠, 𝑎, 𝑤 ( 𝑔 ))

Original: ⊤ 𝑤 ( 𝑔 ) as a model of the optimal 𝑄 -function for the policy with goal 𝑔 .
Revised: ⊤ 𝑤 ( 𝑔 ) as a model of the optimal 𝑄 -function for the policy with goal 𝑔 .

Original: This is reminiscent of 𝐹 ( 𝑠, 𝑎, 𝑧 )
Revised: This is reminiscent of 𝐹 ( 𝑠, 𝑎, 𝑧 )

Original: ⊤ 𝐵 ( 𝑔 ) with 𝑧 =
Revised: ⊤ 𝐵 ( 𝑔 ) with 𝑧 =

Original: 𝐵 ( 𝑔 ) , the value of 𝑧 used to reach 𝑔 in FB.
Revised: 𝐵 ( 𝑔 ) , the value of 𝑧 used to reach 𝑔 in FB.

Original: However, even if we restrict FB to goal-reaching by using 𝑧 = 𝐵 ( 𝑔 ) only (a signiﬁcant restriction), then FB learns 𝐹 ( 𝑠, 𝑎, 𝐵 ( 𝑔 ))
Revised: However, even if we restrict FB to goal-reaching by using 𝑧 = 𝐵 ( 𝑔 ) only (a signiﬁcant restriction), then FB learns 𝐹 ( 𝑠, 𝑎, 𝐵 ( 𝑔 ))

Original: ⊤ 𝐵 ( 𝑔 ′ ) to model the successor measure, i.e., the number of visits to each goal 𝑔 ′ for the policy with goal 𝑔 , starting at ( 𝑠, 𝑎 ) .
Revised: ⊤ 𝐵 ( 𝑔 ′ ) to model the successor measure, i.e., the number of visits to each goal 𝑔 ′ for the policy with goal 𝑔 , starting at ( 𝑠, 𝑎 ) .

Original: Thus, FB learns an object indexed by ( 𝑠, 𝑎, 𝑔, 𝑔 ′ ) for all pairs ( 𝑔, 𝑔 ′ ) .
Revised: Thus, FB learns an object indexed by ( 𝑠, 𝑎, 𝑔, 𝑔 ′ ) for all pairs ( 𝑔, 𝑔 ′ ) .

Original: Thus, FB learns more information (it models successor measures instead of 𝑄 -functions), and allows for recovering linear combinations of goals in a principled way.
Revised: Thus, FB learns more information (it models successor measures instead of 𝑄 -functions), and allows for recovering linear combinations of goals in a principled way.

Original: Meanwhile, even assuming perfect neural network optimization in goal-reaching methods, there is no reason the goal-oriented policies would be optimal for arbitrary linear combinations of goals, only for single goals.
Revised: Meanwhile, even assuming perfect neural network optimization in goal-reaching methods, there is no reason the goal-oriented policies would be optimal for arbitrary linear combinations of goals, only for single goals.

Original: L Pseudocode of training losses
Revised: K Pseudocode of training losses

Original: We provide here pytorch snippets for the key losses, notably the FB loss, SF loss as well as the various feauture learning methods for SF .
Revised: We provide here pytorch snippets for the key losses, notably the FB loss, SF loss as well as the various feauture learning methods for SF .

Original: 1def compute_fb_loss(agent , obs , action , next_obs , z, discount): 3# compute target successor measure 5with torch.no_grad(): mu = agent.policy_net(next_obs , z) next_action
Revised: 1def compute_fb_loss(agent , obs , action , next_obs , z, discount): 3# compute target successor measure 5with torch.no_grad(): mu = agent.policy_net(next_obs , z) next_action

Original: =
Revised: =

Original: TruncatedNormal(mu=mu, stddev=
Revised: TruncatedNormal(mu=mu, stddev=

Original: agent.cfg.stddev , clip=agent.cfg.stddev_clip) target_F1 , target_F2 = agent.forward_target_net(next_obs , z, next_action)
Revised: agent.cfg.stddev , clip=agent.cfg.stddev_clip) target_F1 , target_F2 = agent.forward_target_net(next_obs , z, next_action)

Original: # batch x z_dim target_B = agent.
Revised: # batch x z_dim target_B = agent.backward_

Original: # batch x z_dim target_ x batch 
Revised: # batch x z_dim target_M1 , B x batch target M

Original: ) 13# compute the  FB loss 

Revised: ) 13# compute the main FB loss

Original:  B = 
Revised: .forward , z, action) B = agent.backward

Original: = torch
Revised: = torch

Original: , B)  Fi
Revised: , B) for Fi

Original: ]] # batch x batch
Revised: , F2]] # batch x batch

Original: I 
Revised: I =

Original: Any
Revised: Any

Original: fb_loss = fb_offdiag + fb_diag 24# compute the auxiliary loss 26next_Q1 , nextQ2 =
Revised: fb_loss = fb_offdiag + fb_diag 24# compute the auxiliary loss 26next_Q1 , nextQ2 =

Original: [torch.einsum(’sd, sd -> s’, target_Fi , z) for target_Fi in [target_F1 , target_F2]] next_Q = torch.min(next_Q1 , nextQ2) cov = torch.matmul(B.T, B) /
Revised: [torch.einsum(’sd, sd -> s’, target_Fi , z) for target_Fi in [target_F1 , target_F2]] next_Q = torch.min(next_Q1 , nextQ2) cov = torch.matmul(B.T, B) /

Original: B.shape[0] inv_cov = torch.linalg.pinv(cov)
Revised: B.shape[0] inv_cov = torch.linalg.pinv(cov)

Original: implicit_reward =
Revised: implicit_reward =

Original: (torch.matmul(B, inv_cov) * z).sum(dim=1)
Revised: (torch.matmul(B, inv_cov) * z).sum(dim=1)

Original: # batch_size target_Q = implicit_reward.detach
Revised: # batch_size target_Q = implicit_reward.detach

Original: () + discount * next_Q # batch_size 

Revised: () + discount * next_Q # batch_size

Original: Q1, Q2 =
Revised: Q1, Q2 =

Original: [torch.einsum(’sd, sd -> s’, Fi, z) for Fi in [F1, F2]] q_loss = F.mse_loss(Q1, target_Q) + F.mse_loss(Q2, target_Q) q_loss /= agent.cfg.z_dim fb_loss += q_loss 38# compute Orthonormality losss 

Revised: [torch.einsum(’sd, sd -> s’, Fi, z) for Fi in [F1, F2]] q_loss = F.mse_loss(Q1, target_Q) + F.mse_loss(Q2, target_Q) q_loss /= agent.cfg.z_dim fb_loss += q_loss 38# compute Orthonormality losss

Original: Cov = torch.matmul(B, B.T)
Revised: Cov = torch.matmul(B, B.T)

Original: orth_loss_diag
Revised: orth_loss_diag

Original: = - 2 * Cov.diag().mean() orth_loss_offdiag
Revised: = - 2 * Cov.diag().mean() orth_loss_offdiag

Original: = Cov[off_diag].pow(2).mean() orth_loss = orth_loss_offdiag
Revised: = Cov[off_diag].pow(2).mean() orth_loss = orth_loss_offdiag

Original: + orth_loss_diag fb_loss
Revised: + orth_loss_diag fb_loss

Original: +=
Revised: +=

Original: agent.cfg.ortho_coef * orth_loss 46return fb_loss
Revised: agent.cfg.ortho_coef * orth_loss 46return fb_loss

Original: Listing 1: Pytorch for FB training loss def
Revised: Listing 1: Pytorch for FB training loss def

Original: 2# compute target q-value with torch.no_grad(): mu = agent.policy_net(next_obs , z) next_action
Revised: 2# compute target q-value with torch.no_grad(): mu = agent.policy_net(next_obs , z) next_action

Original: =
Revised: =

Original: TruncatedNormal(mu=mu, stddev=
Revised: TruncatedNormal(mu=mu, stddev=

Original: agent.cfg.stddev , clip=
Revised: agent.cfg.stddev , clip=

Original: agent.cfg.stddev_clip) next_F1 , next_F2 = agent.successor_target_net(next_obs , z, next_action)
Revised: agent.cfg.stddev_clip) next_F1 , next_F2 = agent.successor_target_net(next_obs , z, next_action)

Original: # batch x z_dim target_phi = agent.feature_net(next_goal).detach() # batch x
Revised: # batch x z_dim target_phi = agent.feature_net(next_goal).detach() # batch x

Original: z_dim next_Q1 , next_Q2 = [torch.einsum
Revised: z_dim next_Q1 , next_Q2 = [torch.einsum

Original: ’sd, sd -> s’, next_Fi , z) for next_Fi in [next_F1 , next_F2]] next_Q = torch.min
Revised: ’sd, sd -> s’, next_Fi , z) for next_Fi in [next_F1 , next_F2]] next_Q = torch.min

Original: Q next
Revised: Q next

Original: Q2) target
Revised: Q2) target

Original: =
Revised: =

Original: ) + discount *
Revised: ) + discount *

Original: _Q 

Revised: _Q

Original: , action)
Revised: , action)

Original: Q
Revised: Q

Original: ’sd for Fi in [F1, F2]] sf loss =
Revised: ’sd for Fi in [F1, F2]] sf loss =

Original: F. loss
Revised: F. loss

Original: Q1,
Revised: Q1,

Original: Q)
Revised: Q)

Original: F
Revised: F

Original: Q) 16return sf_loss
Revised: Q) 16return sf_loss

Original: Listing loss 1def compute_phi_loss(agent , obs , next obs): 3phi = agent.feature_net(obs) next_phi = agent.feature_net next_obs) loss = (phi - next_phi).pow(2).mean() 7# compute Orthonormality losss 

Revised: Listing loss 1def compute_phi_loss(agent , obs , next obs): 3phi = agent.feature_net(obs) next_phi = agent.feature_net next_obs) loss = (phi - next_phi).pow(2).mean() 7# compute Orthonormality losss

Original: Cov = torch.matmul(phi , phi.
Revised: Cov = torch.matmul(phi , phi.

Original: T)
Revised: T)

Original: I = torch.eye(*Cov.size(), device=Cov.device) off_diag
Revised: I = torch.eye(*Cov.size(), device=Cov.device) off_diag

Original: Listing 4: Pytorch code for the contrastive CL loss def compute_phi_loss(agent , obs , action , next_obs): 2phi = agent.feature_net(obs) next_phi = agent.feature_net(next_obs) predicted_action = agent.inverse_dynamic_net(torch.cat([phi , next_phi], dim=-1)) loss = (action - predicted_action).pow(2).mean() 7return loss
Revised: Listing 4: Pytorch code for the contrastive CL loss def compute_phi_loss(agent , obs , action , next_obs): 2phi = agent.feature_net(obs) next_phi = agent.feature_net(next_obs) predicted_action = agent.inverse_dynamic_net(torch.cat([phi , next_phi], dim=-1)) loss = (action - predicted_action).pow(2).mean() 7return loss

Original: Listing 5: Pytorch code of ICM
Revised: Listing 5: Pytorch code of ICM

Original: def compute_phi_loss(agent , obs , action , next_obs): 2phi = agent.feature_net(obs) predicted_next_obs = agent.forward_dynamic_net(torch.cat([phi , action], dim
Revised: def compute_phi_loss(agent , obs , action , next_obs): 2phi = agent.feature_net(obs) predicted_next_obs = agent.forward_dynamic_net(torch.cat([phi , action], dim

Original: 1)) loss =
Revised: 1)) loss =

Original: (predicted_next obs next_obs
Revised: (predicted_next obs next_obs

Original: .pow
Revised: .pow

Original: .mean
Revised: .mean

Original: ) 6return loss
Revised: ) 6return loss

Original: Listing 6: Pytorch code for Trans loss def compute_phi_loss(agent , obs , action , next_obs): phi = agent.feature_net(obs) with torch.no_grad(): next_phi = agent.target_feature_net(next_obs) predicted_next_obs = agent.forward_dynamic_net( cat( phi , action],
Revised: Listing 6: Pytorch code for Trans loss def compute_phi_loss(agent , obs , action , next_obs): phi = agent.feature_net(obs) with torch.no_grad(): next_phi = agent.target_feature_net(next_obs) predicted_next_obs = agent.forward_dynamic_net( cat( phi , action],

Original: =-1)) loss =
Revised: =-1)) loss =

Original: (predicted_next obs next
Revised: (predicted_next obs next

Original: .detach(
Revised: .detach(

Original: .pow
Revised: .pow

Original: .mean
Revised: .mean

Original: () 7# update target network for param , target_param in zip(
Revised: () 7# update target network for param , target_param in zip(

Original: .feature_net.parameters(), agent.target_feature_net.parameters()): target_param data.copy
Revised: .feature_net.parameters(), agent.target_feature_net.parameters()): target_param data.copy

Original: + tau param
Revised: + tau param

Original: ) 12return loss
Revised: ) 12return loss

Original: Listing 7: Pytorch
Revised: Listing 7: Pytorch

Original: Latent loss def compute_phi return
Revised: Latent loss def compute_phi return

Original: [Table]
Revised: [Table]

Original: Listing
Revised: Listing

Original: [Table]
Revised: [Table]

Original: for AEnc loss loss
Revised: for AEnc loss loss

Original: = loss 

Revised: = loss

Original: 22return loss code for LRA-SR loss
Revised: 22return loss code for LRA-SR loss