Original: Improving Zero-shot Generalization in Ofﬂine Reinforcement Learning using Generalized Similarity Functions
Revised: Improving Zero-shot Generalization in Ofﬂine Reinforcement Learning using Generalized Similarity Functions

Original: Anonymous Author(s)

Original: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Reinforcement learning (RL) agents are widely used for solving complex sequentialdecision making tasks, but still exhibit difﬁculty in generalizing to scenarios notseen during training.
Revised: Reinforcement learning (RL) agents are widely used for solving complex sequential decision making tasks, but still exhibit difﬁculty in generalizing to scenarios not seen during training.

Original: While prior online approaches demonstrated that usingadditional signals beyond the reward function can lead to better generalizationcapabilities in RL agents, i.e., using self-supervised learning (SSL), they struggle inthe ofﬂine RL setting, i.e., learning from a static dataset.
Revised: While prior online approaches demonstrated that using additional signals beyond the reward function can lead to better generalization capabilities in RL agents, i.e., using self-supervised learning (SSL), they struggle in the ofﬂine RL setting, i.e., learning from a static dataset.

Original: We show that performanceof online algorithms for generalization in RL can be hindered in the ofﬂine settingdue to poor estimation of similarity between observations.
Revised: We show that performance of online algorithms for generalization in RL can be hindered in the ofﬂine setting due to poor estimation of similarity between observations.

Original: We propose a newtheoretically-motivated framework called Generalized Similarity Functions (GSF),which uses contrastive learning to train an ofﬂine RL agent to aggregate observationsbased on the similarity of their expected future behavior, where we quantify thissimilarity using generalized value functions .
Revised: We propose a new theoretically-motivated framework called Generalized Similarity Functions (GSF), which uses contrastive learning to train an ofﬂine RL agent to aggregate observations based on the similarity of their expected future behavior, where we quantify this similarity using generalized value functions .

Original: We show that GSF is general enoughto recover existing SSL objectives while also improving zero-shot generalizationperformance on two pixel-based ofﬂine RL benchmarks.
Revised: We show that GSF is general enough to recover existing SSL objectives while also improving zero-shot generalization performance on two pixel-based ofﬂine RL benchmarks.

Original:  Introduction
Revised:  Introduction

Original: Reinforcement learning (RL) is a powerful framework for solving complex tasks that require asequence of decisions.
Revised: Reinforcement learning (RL) is a powerful framework for solving complex tasks that require a sequence of decisions.

Original: The RL paradigm has allowed for major breakthroughs in various ﬁelds,e g. outperforming humans on video games [Mnih et al., 2015, Schwarzer et al., 2020], controllingstratospheric balloons [Bellemare et al., 2020] and learning reward functions from robot manipulationvideos [Chen et al., 2021].
Revised: The RL paradigm has allowed for major breakthroughs in various ﬁelds, e g. outperforming humans on video games [1, 2], controlling stratospheric balloons [3] and learning reward functions from robot manipulation videos [4].

Original: More recently, RL agents have been tested in a generalization setting, i e.in which training involves a ﬁnite number of related tasks sampled from some distribution, with apotentially distinct sampling distribution during test time [Cobbe et al., 2019, Song et al., 2019, Hjelmet al., 2022].
Revised: More recently, RL agents have been tested in a generalization setting i e. in which training involves a ﬁnite number of related tasks sampled from some distribution, with a potentially distinct sampling distribution during test time [5, 6, 7].

Original: The main issue for designing generalizable agents is the lack of on-policy data fromtasks not seen during training: it is impossible to enumerate all variants of a real-world environment  (limited) training  a broader set of problems.
Revised: The main issue for designing generalizable agents is the lack of on-policy data from tasks not seen during training: it is impossible to enumerate all variants of a real-world environment during training and hence the agent must extrapolate from a (limited) training task collection onto a broader set of problems.

Original: Since the learning agent is given no training data from test-time tasks, thisproblem is referred to as zero-shot generalization.
Revised: Since the learning agent is given no training data from test-time tasks, this problem is referred to as zero-shot generalization.

Original: In our work, we are interested in the problem of zero-shot generalization where the difference between tasks is predominantly due to perceptually distinct observations.
Revised: In our work, we are interested in the problem of zero-shot generalization where the difference between tasks is predominantly due to perceptually distinct observations.

Original: An example of this setting is any environment with distractor features [Cobbe et al., Stone et al., 2021], i.e. features with no dependence on the reward signal nor the agent’s decisions.
Revised: An example of this setting is any environment with distractor features [8, 9], i.e. features with no dependence on the reward signal nor the agent’s decisions.

Original: This generalization setting has recently received much attention [Liu et al., 2020a, Agarwal et al., 2021a, Mazoure et al., 2022], due to its particular relevance to real-world scenarios, for example deploying the same autonomous driving agent at day or at night
Revised: This generalization setting has recently received much attention [10, 11, 12], due to its particular relevance to real-world scenarios, for example deploying the same autonomous driving agent at day or at night.

Original: Generalization capabilities of an agent can be analyzed through the prism of representation learning , under which the agent’s current belief about a rich and high-dimensional environment are summarized in a low-dimensional entity, called a representation.
Revised: Generalization capabilities of an agent can be analyzed through the prism of representation learning , under which the agent’s current belief about a rich and high-dimensional environment are summarized in a low-dimensional entity, called a representation.

Original: Recent work in online RL has shown that learningstate representations with speciﬁc properties such as disentanglement [Higgins et al., 2017] or linearseparability [Lin et al., 2020] can improve zero-shot generalization performance.
Revised: Recent work in online RL has shown that learning state representations with speciﬁc properties such as disentanglement [13] or linear separability [14] can improve zero-shot generalization performance.

Original: Achieving this withlimited data (i.e. ofﬂine RL) is challenging, since the representation will have a large estimation errorover regions of low data coverage.
Revised: Achieving this with limited data (i.e. ofﬂine RL) is challenging, since the representation will have a large estimation error over regions of low data coverage.

Original: A common solution to mitigate this task-speciﬁc overﬁtting andextracting the most information out of the data consists in introducing auxiliary learning signals otherthan instantaneous reward [Raileanu and Fergus, 2021].
Revised: A common solution to mitigate this task-speciﬁc overﬁtting and extracting the most information out of the data consists in introducing auxiliary learning signals other than instantaneous reward [15].

Original: As we show later in the paper, many suchsignals already contained in the dataset can be used to further improve generalization performance.
Revised: As we show later in the paper, many such signals already contained in the dataset can be used to further improve generalization performance.

Original: For instance, the generalization performance of PPO on Procgen remains limited even when trainingon 200M frames, while generalization-oriented agents [Raileanu and Fergus, 2021, Mazoure et al. ,2022] can outperform it by leveraging additional auxiliary signals.
Revised: For instance, the generalization performance of PPO on Procgen remains limited even when training on 200M frames, while generalization-oriented agents [15, 12] can outperform it by leveraging additional auxiliary signals.

Original: However, a major issue with theaforementioned methods is their exorbitant reliance on online access to the environment, an impracticalrestriction for real-world scenarios.
Revised: However, a major issue with the aforementioned methods is their exorbitant reliance on online access to the environment, an impractical restriction for real-world scenarios.

Original: In contrast, in many real-world scenarios access to the environment is restricted to an ofﬂine, ﬁxeddataset of experience [Ernst et al., 2005, Lange et al., 2012].
Revised: In contrast, in many real-world scenarios access to the environment is restricted to an ofﬂine, ﬁxed dataset of experience [16, 17].

Original: A natural limitation for generalization fromofﬂine data is that policy improvement is dependent on dataset quality.
Revised: A natural limitation for generalization from ofﬂine data is that policy improvement is dependent on dataset quality.

Original: Speciﬁcally, high-dimensionalproblems such as control from pixels require large amounts of training experience: a standard trainingof PPO [Schulman et al., 2017] for 25 million frames on Procgen [Cobbe et al., 2020] generatesmore than 300 Gb of data, an impractical amount of data to share for ofﬂine RL research.
Revised: Speciﬁcally, high-dimensional problems such as control from pixels require large amounts of training experience: a standard training of PPO [18] for 25 million frames on Procgen [8] generates more than 300 Gb of data, an impractical amount of data to share for ofﬂine RL research.

Original: Improvingzero-shot generalization performance from an ofﬂine dataset of high-dimensional observations istherefore a hard problem due to limitations on dataset size and quality.
Revised: Improving zero-shot generalization performance from an ofﬂine dataset of high-dimensional observations is therefore a hard problem due to limitations on dataset size and quality.

Original: In this work, we are interested in improving zero-shot generalization across a family of PartiallyObservable Markov decision processes [POMDPs, Murphy, 2000] in an ofﬂine RL setting, i.e.by training agents on a ﬁxed dataset.
Revised: In this work, we are interested in improving zero-shot generalization across a family of PartiallyObservable Markov decision processes [POMDPs, 19] in an ofﬂine RL setting, i.e. by training agents on a ﬁxed dataset.

Original: We hypothesize that in order for an RL agent to be able togeneralize across perceptually different POMDPs without adaptation, observations with similarfuture behavior should be assigned to close representations.
Revised: We hypothesize that in order for an RL agent to be able to generalize across perceptuallydifferentPOMDPswithoutadaptation, observationswithsimilarfuturebehaviorshouldbeassigned to close representations.

Original: We use the generalized value function(GVF) framework [Sutton et al., 2011] to capture future behavior with respect to any instantaneoussignal (called cumulant) at a given state.
Revised: We use the generalized value function (GVF) framework [20] to capture future behavior with respect to any instantaneous signal (called cumulant) at a given state.

Original: Speciﬁcally, the choice of cumulant determines the natureof the behavioral similarity that is induced into state representations.
Revised: Speciﬁcally, the choiceofcumulantdeterminesthenatureofthebehavioralsimilaritythatisinducedintostaterepresentations.

Original: For example, using rewardsimilarity leads to learning bisimulation metrics [Ferns et al., 2004, Li et al., 2006, Castro, 2020, Zhang et al., 2020], while using future state-action visitation counts encourages reward-free behavioralsimilarity [Misra et al., 2020, Liu et al., 2020a, Agarwal et al., 2021a, Mazoure et al., 2022].
Revised: Forexample,usingrewardsimilarityleadstolearningbisimulationmetrics[21,22,23,24],while usingfuturestate-actionvisitationcountsencouragesreward-freebehavioralsimilarity[25,10,11,12].

Original: Our main contributions are as follows:
Revised: Our main contributions are as follows:

Original: We propose Generalized Similarity Functions (GSF), a novel self-supervised learningalgorithm for reinforcement learning that aggregates latent representations of observationsby their future behavior (or generalized value function).2.
Revised: We propose Generalized Similarity Functions (GSF), a novel self-supervised learning algorithm for reinforcement learning that aggregates latent representations of observations by their future behavior (or generalized value function).

Original: Existing ofﬂine RL benchmarks [Fu et al., 2020, Gulcehre et al., 2020] are not well-suitedto test zero-shot generalization, and so we devise two new benchmarks: ofﬂine Procgen andDistracting Control Suite.
Revised: Existing ofﬂine RL benchmarks [26, 27] are not well-suited to test zero-shot generalization, and so we devise two new benchmarks: ofﬂine Procgen and Distracting Control Suite.

Original: The ﬁrst consists of 5M transitions from 200 related levels of 16distinct games; the second consists of 1M transitions from 3 variations of 4 distincts tasks.3.
Revised: The ﬁrst consists of 5M transitions from 200 related levels of 16 distinct games; the second consists of 1M transitions from 3 variations of 4 distincts tasks.

Original: We evaluate performance of GSF and other baseline methods on both benchmarks, and show that GSF outperforms both previous state-of-the-art ofﬂine RL and representation learningbaselines on the entire distribution of levels.
Revised: We evaluate performance of GSF and other baseline methods on both benchmarks, and show that GSF outperforms both previous state-of-the-art ofﬂine RL and representation learning baselines on the entire distribution of levels.

Original: We analyze the theoretical properties of GSF and describe the impact of hyperparameters and cumulant functions on empirical behavior in both ofﬂine Procgen and the ofﬂine Distracting Suite benchmarks.
Revised: We analyze the theoretical properties of GSF and describe the impact of hyperparameters and cumulant functions on empirical behavior in both ofﬂine Procgen and the ofﬂine Distracting Suite benchmarks.

Original:  Related Works
Revised:  Related Works

Original: Generalization in reinforcement learning Generalizing a model’s predictions across a varietyof unseen, high-dimensional inputs has been extensively studied in the static supervised learningsetting [Bartlett, 1998, Triantaﬁllou et al., 2019, Valle-Pérez and Louis, 2020, Liu et al., 2020b].
Revised: Generalization in reinforcement learning Generalizing a model’s predictions across a variety of unseen, high-dimensional inputs has been extensively studied in the static supervised learning setting [28, 29, 30, 31].

Original: Generalization in RL has received a lot of attention: extrapolation to unseen rewards [Barreto et al. ,2016, Misra et al., 2020], observations [Zhang et al., 2020, Raileanu and Fergus, 2021, Liu et al. ,2020a, Agarwal et al., 2021a, Mazoure et al., 2022] and transition dynamics [Ball et al., 2021].
Revised: Generalization in RL has received a lot of attention: extrapolation to unseen rewards [32, 25], observations [24, 15, 10, 11, 12] and transition dynamics [33].

Original: Each generalization scenario is best solved by their respective set of methods: sufﬁcient exploration [Misraet al., 2020, Agarwal et al., 2020], auxiliary learning signals [Srinivas et al., 2020, Mazoure et al. ,2020, Stooke et al., 2021] or data augmentation [Ball et al., 2021, Sinha and Garg, 2021].
Revised: Each generalization scenario is best solved by their respective set of methods: sufﬁcient exploration [25, 34], auxiliary learning signals [35, 36, 37] or data augmentation [33, 38].

Original: Dataaugmentation is a promising technique, but typically relies on handcrafted domain information, whichmight not be available a priori .
Revised: Data augmentation is a promising technique, but typically relies on handcrafted domain information, which might not be available a priori .

Original: In fact, we will show in our experiments that generalization in theofﬂine RL setting is poor even when using such handcrafted data augmentations, without additionalrepresentation learning mechanisms.
Revised: In fact, we will show in our experiments that generalization in the ofﬂine RL setting is poor even when using such handcrafted data augmentations, without additional representation learning mechanisms.

Original: In this work, we posit that representation learning should useinstantaneous auxiliary signals in order to prevent overﬁtting onto a unique signal (e.g. reward acrosstasks) and improve generalization performance.
Revised: In this work, we posit that representation learning should use instantaneous auxiliary signals in order to prevent overﬁtting onto a unique signal (e.g. reward across tasks) and improve generalization performance.

Original: Theoretical generalization guarantees have only beenprovided so far for limited scenarios, mostly for bandits [Swaminathan and Joachims, 2015], linear MDPs [Boyan and Moore, 1995, Wang et al., 2021b, Nachum and Yang, 2021] and across rewardfunctions [Castro and Precup, 2010, Barreto et al., 2016, Wang et al., 2021a, Touati and Ollivier, 2021].
Revised: Theoretical generalization guarantees have only been provided so far for limited scenarios, mostly for bandits [39], linear MDPs [40, 41, 42] and across reward functions [32, 43, 44].

Original: Representation learning For simple POMDPs, near-optimal policies can be found by optimizingfor the reward alone.
Revised: Representation learning For simple POMDPs, near-optimal policies can be found by optimizing for the reward alone.

Original: However, more complex settings may require additional auxiliary signals inorder to ﬁnd state abstractions better suited for control.
Revised: However, more complex settings may require additional auxiliary signals in order to ﬁnd state abstractions better suited for control.

Original: The problem of learning meaningful staterepresentations (or abstractions) for planning and control has been extensively studied previously [Jongand Stone, 2005, Li et al., 2006], but saw real breakthroughs only recently, in particular due to advancesin self-supervised learning (SSL).
Revised: The problem of learning meaningful state representations (or abstractions) for planning and control has been extensively studied previously [45, 22], but saw real breakthroughs only recently, in particular due to advances in self-supervised learning (SSL).

Original: Outside of RL, SSL has achieved spectacular results by closing thegap between unsupervised and supervised learning on certain datasets [Hjelm et al., 2018, Oord et al.
Revised: Outside of RL, SSL has achieved spectacular results by closing the gap between unsupervised and supervised learning on certain datasets [46, 47, 48, 49].

Original: ,2018, Caron et al., 2020, Grill et al., 2020].

Original: Representation learning, and speciﬁcally self-supervisedlearning, has also been used to achieve state-of-the-art generalization and sample efﬁciency resultsin RL on challenging control problems such as data efﬁcient Atari [Schwarzer et al., 2020, 2021], DeepMind Control [Agarwal et al., 2021a] and Procgen [Mazoure et al., 2020, Stooke et al., 2021,
Revised: Representation learning, and speciﬁcally self-supervised learning, has also been used to achieve state-of-the-art generalization and sample efﬁciency results in RL on challenging control problems such as data efﬁcient Atari [2, 50], DeepMind Control [11] and Procgen [36, 37, 15, 12].

Original: Raileanu and Fergus, 2021, Mazoure et al., 2022]. Noteworthy instances of theoretically-motivatedrepresentation learning methods for RL include heuristic-guided learning [Cheng et al., 2021], random Fourier features [Nachum and Yang, 2021] and metric learning [Li et al., 2020b,a].
Revised: Noteworthy instances of theoretically-motivated representation learning methods for RL include heuristic-guided learning [51], random Fourier features [42] and metric learning [52, 53].

Original: Ofﬂine reinforcement learning When learning from a static dataset, agents should balanceinterpolation and extrapolation errors, while ensuring proper diversity of actions (i.e. prevent collapseto most frequent action in the data).
Revised: Ofﬂine reinforcement learning When learning from a static dataset, agents should balance interpolation and extrapolation errors, while ensuring proper diversity of actions (i.e. prevent collapse to most frequent action in the data).

Original: Popular ofﬂine RL algorithms such as BCQ [Fujimoto et al., 2019], MBS [Liu et al., 2020c], and CQL [Kumar et al., 2020] rely on a behavior regularization loss [Wuet al., 2019] as a tool to control the extrapolation error.
Revised: Popular ofﬂine RL algorithms such as BCQ [54], MBS [55], and CQL [56] rely on a behavior regularization loss [57] as a tool to control the extrapolation error.

Original: Some methods, such as F-BRC [Kostrikov et al. ,2021a] are deﬁned only for continuous action spaces while others, such as MOReL [Kidambi et al. ,2020] estimate a pessimistic transition model.
Revised: Some methods, such as F-BRC [58] are deﬁned only for continuous action spaces while others, such as MOReL [59] estimate a pessimistic transition model.

Original: The major issue with current ofﬂine RL algorithmssuch as CQL is that they are perhaps overly pessimistic for generalization purposes, i.e.CQL and
Revised: The major issue with current ofﬂine RL algorithms such as CQL is that they are perhaps overly pessimistic for generalization purposes, i.e.

Original:  MBS ensure that the policy improvement is well-supported by the batch of data.
Revised: CQL and MBS ensure that the policy improvement is well-supported by the batch of data.

Original:  Problem setting 
Revised:  Problem setting 3.1 Partially-observable Markov decision processes

Original: 3.1 Partially-observable Markov decision processes

Original: A (inﬁnite-horizon) partially-observable Markov decision process [POMDP, Murphy, 2000] M isdeﬁned by the tuple M = h S ,p 0 , A ,p S , O ,p O ,r,  i , where S is a state space, p 0 ={+
+}
Revised: A (inﬁnite-horizon) partially-observable Markov decision process [POMDP, 19] M is deﬁned by the tuple M = h S ,p 0 , A ,p S , O ,p O ,r,  i , where S is a state space, p {+
+}

Original:  P [ s 0 ] is the starting127 state distribution,
Revised: = P [ s 0 ] is the starting state distribution,

Original: A is an action space, p S = P
Revised: A is an action space, p S = P

Original: [ ·| s t ,a t ]: S ⇥ A !
Revised: [ ·| s t ,a t ]: S ⇥ A !

Original:  ( S ) is a transition function, O is128 an observation space, p
Revised:  ( S ) is a transition function, O is an observation space, p

Original: O = P [ ·| s t ] : S !
Revised: O = P [ ·| s t ]: S !

Original:  ( O ) 1 is an observation function, r : S ⇥ A !
Revised:  ( O ) 2 is an observation function, r : S ⇥ A !

Original: [ r min ,r max ]129 is a reward function and  2 [0 , 1) is a discount factor.
Revised: [ r min ,r max ] is a reward function and  2 [0 , 1) is a discount factor.

Original: The system starts in one of the initial statess 

Revised: The system starts in one of the initial states s

Original: ⇠ p 0 with observation o 

Revised: ⇠ p 0 with observation o

Original: ⇠ p O ( ·| s 0 ) .
Revised: ⇠ p O ( ·| s 0 ) .

Original: At every timestep t = 1 , 2 , 3 ,.. , the agent, parameterized by131 a policy ⇡ : O !
Revised: At every timestep t = 1 , 2 , 3 ,.. , the agent, parameterized by a policy ⇡ : O !

Original:  ( A ) , samples an action a t ⇠ ⇡ ( ·| o t ) .
Revised:  ( A ) , samples an action a t ⇠ ⇡ ( ·| o t ) .

Original: The environment transitions into a next state132 s t +1 ⇠ p S ( ·| s t ,a t ) and emits a reward r t = r ( s t ,a t ) along with a next observation o t +1 ⇠ p O ( ·| s t +1 ) .
Revised: The environment transitions into a next state s t +1 ⇠ p S ( ·| s t ,a t ) and emits a reward r t = r ( s t ,a t ) along with a next observation o t +1 ⇠ p O ( ·| s t +1 ) .

Original: The goal of an RL agent is to maximize the cumulative discounted rewards
Revised: The goal of an RL agent is to maximize the cumulative discounted rewards

Original: P 1 t =

Revised: P 1 t =

Original:  t r t obtained overthe entire episode.
Revised:  t r t obtained over the entire episode.

Original: Value-based off-policy RL algorithms achieve this by estimating the state-actionvalue function under a target policy ⇡ :
Revised: Value-based off-policy RL algorithms achieve this by estimating the state-action value function under a target policy ⇡ :

Original: [Equation]
Revised: [Equation]

Original: for s t 2 S ,a t 2 A and where P ⇡ t denotes the joint distribution of { s t + k ,a t + k } 1 k =1 obtained by executing137 ⇡ in the environment.
Revised: for s t 2 S ,a t 2 A and where P ⇡ t denotes the joint distribution of { s t + k ,a t + k } 1 k =1 obtained by executing ⇡ in the environment.

Original: An important distinction from online RL is that, in the ofﬂine RL setting, we assume access to ahistorical dataset D µ (instead of a simulator) collected by logging experience of the policy, µ , in theform { o i,t ,a i,t ,r i,t } i = N,t = T i =1 ,t =1 where, for practical purposes, the episode is truncated at T timesteps.
Revised: An important distinction from online RL is that, in the ofﬂine RL setting, we assume access to a historical dataset D µ (instead of a simulator) collected by logging experience of the policy, µ , in the form { o i,t ,a i,t ,r i,t } i = N,t = T i =1 ,t =1 where, for practical purposes, the episode is truncated at T timesteps.

Original: Furthermore, we assume that the agent can only be trained on a limited collection of POMDPs
Revised: Furthermore, we assume that the agent can only be trained on a limited collection of POMDPs

Original: M train =
Revised: M train =

Original: { M i } mi =1 , and its performance is evaluated on the set of test POMDPs M test .
Revised: { M i } mi =1 , and its performance is evaluated on the set of test POMDPs M test .

Original: We assume143 that both
Revised: We assume that both

Original: M train and M test were sampled from a common task distribution and that every POMDP144 M i 2 M = M train [ M test shares the same transition dynamics and reward function with M but has145 a different observation function p i, O .
Revised: M train and M test were sampled from a common task distribution and that every POMDP M i 2 M = M train [ M test shares the same transition dynamics and reward function with M but has a different observation function p i, O .

Original: Importantly, since we perform control from pixels, we are in the POMDP setting [see Yarats et al., 2019] and therefore emphasize the difference between observationso t and corresponding states s t throughout the paper.
Revised: Importantly, since we perform control from pixels, we are in the POMDP setting [see 60] and therefore emphasize the difference between observations o t and corresponding states s t throughout the paper.

Original: 3.2 Representation learning
Revised: 3.2 Representation learning

Original: Previous works in the RL literature have studied the use of auxiliary signals to improve generalizationperformance.
Revised: Previous works in the RL literature have studied the use of auxiliary signals to improve generalization performance.

Original: Among others, Liu et al. [2020a], Agarwal et al. [2021a] deﬁne the similarity of twoobservations to depend on the distance between action sequences rolled out from that observation undertheir respective optimal policies.
Revised: Among others, [10, 11] deﬁne the similarity of two observations to depend on the distance between action sequences rolled out from that observation under their respective optimal policies.

Original: They achieve this by ﬁnding a latent space Z ✓ S in which the distance153 d Z ( z,z 0 ) for all z,z 0 2 Z is equivalent to distance between true latent states d S ( s,s 0 ) for all s,s 0 2 S ;154 the aforementionned works learn
Revised: They achieve this by ﬁnding a latent space Z ✓ S in which the distance d Z ( z,z 0 ) for all z, z 0 2 Z is equivalent to distance between true latent states d S ( s, s 0 ) for all s, s 0 2 S ; the aforementionned works learn

Original: Z by optimizing action-based similarities between observations.
Revised: Z by optimizing action-based similarities between observations.

Original: In practice, latent space z is decoded from observation o using a latent state decoder f : O !
Revised: In practice, latent space z is decoded from observation o using a latent state decoder f : O !

Original: Z fromobservation o t .
Revised: Z from observation o t .

Original: Throughout the paper, we assume that all value functions have a linear form in thelatent decoded state, i.e. Q ✓ ( o,a )=
Revised: Throughout the paper, we assume that all value functions have a linear form in the latent decoded state, i.e. Q ✓ ( o,a )=

Original: ✓ > a f ( o )=
Revised: ✓ > a f ( o )=

Original: ✓ > a z , which agrees with our practical implementationof all algorithms.
Revised: ✓ > a z , which agrees with our practical implementation of all algorithms.

Original: Within this model family, the ability of an RL agent to correctly decode latent statesfrom unseen observations directly affects its policy, and therefore, its generalization capabilities.
Revised: Within this model family, the ability of an RL agent to correctly decode latent states from unseen observations directly affects its policy, and therefore, its generalization capabilities.

Original: Inthe next section, we discuss why representation learning is important for ofﬂine RL, and how existingaction-based similarity metrics fail to recover the true latent states for important families of POMDPs.
Revised: In the next section, we discuss why representation learning is important for ofﬂine RL, and how existing action-based similarity metrics fail to recover the true latent states for important families of POMDPs.

Original:  Motivating example
Revised:  Motivating example

Original: Multiple recently proposed self-supervised objectives [Liu et al., 2020a, Agarwal et al., 2021a]conjecture that observations o 1 2{+
+} M 1 ,o 2 2 M 2 that emit similar future action sequences under optimalpolicies ⇡ ⇤ 1 , ⇡ ⇤ 2 should be decoded into nearby latent states z 1 ,z 2 .
Revised: Multiple recently proposed self-supervised objectives [10, 11] conjecture that observations o 1 {+
+} M 1 ,o 2 2 M 2 that emit similar future action sequences under optimal policies ⇡ ⇤ 1 , ⇡ ⇤ 2 should be decoded into nearby latent states z 1 ,z 2 .

Original: While this heuristic can correctlygroup observations with respect to their true latent state in simple action spaces, it fails to identifysimilar pairs of trajectories in POMDPs with multiple optimal policies 2 .
Revised: While this heuristic can correctly group observations with respect to their true latent state in simple action spaces, it fails to identify similar pairs of trajectories in POMDPs with multiple optimal policies 3 .

Original: For instance, two trajectoriesmight visit an identical set of latent states, but have drastically different actions.
Revised: For instance, two trajectories might visit an identical set of latent states, but have drastically different actions.

Original: Fig.
Revised: Fig.

Original: 1a shows one such example: two levels of the Climber game have a near-identical true latentstate (see Appendix) and value function (average normalized mean-squared error of 0.0398 acrossepisode), while having very different action sequences from a same PPO policy (average total variationdistance of 0.4423 across episode).
Revised: 1a shows one such example: two levels of the Climber game have a near-identical true latent state (see Appendix) and value function (average normalized mean-squared error of 0.0398 across episode), while having very different action sequences from a same PPO policy (average total variation distance of 0.4423 across episode).

Original: The problem is especially acute in Procgen, since the PPO policyis high-entropy for some environments (see Fig. 5), i.e. various levels can have multiple drasticallydifferent near-optimal policies, and hence fail to properly capture observation similarities
Revised: The problem is especially acute in Procgen, since the PPO policy is high-entropy for some environments (see Fig. 5), i.e. various levels can have multiple drastically different near-optimal policies, and hence fail to properly capture observation similarities.

Original: In this scenario, assigning observations to a similar latent state by value function similarity would yield abetterstaterepresentationthanreasoningaboutactionsimilarities.
Revised: In this scenario, assigning observations to a similar latent state by value function similarity would yield a betterstaterepresentationthanreasoningaboutactionsimilarities.

Original: Ontheotherhand, Fig.1bshowsadomainwheregroupingstaterepresentationsbyactionsequencescanbeoptimal.
Revised: Ontheotherhand, Fig.1bshowsadomainwheregroupingstaterepresentationsbyactionsequencescanbeoptimal.

Original: Sohowdoweunifythesesimilarity metrics under a single framework?
Revised: Sohowdoweunifythese similarity metrics under a single framework?

Original: In the next section, we use this insight to design a generalway of improve representation learning through self-supervised learning of discounted future behavior.
Revised: In the next section, we use this insight to design a general way of improve representation learning through self-supervised learning of discounted future behavior.

Original:  Method
Revised:  Method

Original: We propose measuring a generalized notion of future behavior similarity using generalized valuefunctions, as deﬁned by the corresponding cumulant function.
Revised: We propose measuring a generalized notion of future behavior similarity using generalized value functions, as deﬁned by the corresponding cumulant function.

Original: The choice of cumulant determineswhich components of the future trajectory are most relevant for generalization.
Revised: The choice of cumulant determines which components of the future trajectory are most relevant for generalization.

Original: some cumulant function c , whose quantiles are then used to label each observation in

Original: 5.1 Quantifying future behavior with GVFs
Revised: 5.1 Quantifying future behavior with GVFs

Original: An RL agent’s future discounted behavior can be quantiﬁed not only by its value function, but otherauxiliary signals, for example, by its future observation occupancy measure, known as successorrepresentation
Revised: An RL agent’s future discounted behavior can be quantiﬁed not only by its value function, but other auxiliary signals, for example, by its future observation occupancy measure, known as successor representation [62, 32].

Original: [Dayan, 1993, Barreto et al., 2016].

Original: The choice of the signal used during value iterationmeasures the properties the agent will exhibit in the future, such as accumulated returns, actions,or observation visitation density.
Revised: The choice of the signal used during value iteration measures the properties the agent will exhibit in the future, such as accumulated returns, actions, or observation visitation density.

Original: See Thm. 1 in the Appendix for the connection between successorfeatures and interpolation error in our method.
Revised: See Thm.in the Appendix for the connection between successor features and interpolation error in our method.

Original: Following the work of Sutton et al. [2011], we can broaden the class of value functions to any kindof cumulative discounted signal, as deﬁned by a bounded cumulant function c : O ⇥ A !
Revised: Following the work of [20], we can broaden the class of value functions to any kind of cumulative discounted signal, as deﬁned by a bounded cumulant function c : O ⇥ A !

Original: R d , s.t.| c ( o,a ) |  c max for c max =sup 
Revised: R d , s.t. | c ( o,a ) |  c max for c max =sup o,a 2 O ⇥

Original: o,a 2 O ⇥ A c ( o,a ) .
Revised: A c ( o,a ) .

Original: While typically cumulants are scalar-valued functions194 (e.g. reward), we also make use of the vector-valued case for learning the successor features [Barretoet al., 2016], in which case the norm of c ( o,a ) is bounded.
Revised: While typically cumulants are scalar-valued functions (e.g. reward), we also make use of the vector-valued case for learning the successor features [32], in which case the norm of c ( o,a ) is bounded.

Original: Deﬁnition 1 (Generalized value function) Let c be any bounded function over R d , let  2 [0 , 1] andµ any policy.
Revised: Deﬁnition 1 (Generalized value function) Let c be any bounded function over R d , let  2 [0 , 1] and µ any policy.

Original: The generalized value function is deﬁned, for any timestep t  1 and o t 2 O , as
Revised: The generalized value function is deﬁned, for any timestep t  1 and o t 2 O , as

Original: [Equation]
Revised: [Equation]

Original: Since, in our case, we can learn G µ for each distinct POMDP M i for the dataset D µ , we index the
Revised: Since, in our case, we can learn G µ for each distinct POMDP M i for the dataset D µ , we index the

Original: GVF using the POMDP index, i.e. G µi =
Revised: GVF using the POMDP index, i e. G µi =

Original: LearnGVF ( c, D µ ,i ) (in practice, learning is parallelized).
Revised: LearnGVF ( c, D µ ,i ) (in practice, learning is parallelized)

Original: 5.2 Measuring distances between GVFs of different POMDPs
Revised: 5.2 Measuring distances between GVFs of different POMDPs

Original: Examining the difference between future behaviors of two observations quantiﬁes the exact amountof expected behavior change between these two observations.
Revised: Examining the difference between future behaviors of two observations quantiﬁes the exact amount of expected behavior change between these two observations.

Original: Using the GVF framework, we couldcompute the distance between o 1 2 M 1 and o 2 2 M 2 by ﬁrst estimating the latent state with z = f ( o )using a latent state decoder f , and then using the following distance as a measure of dissimilarity
Revised: Using the GVF framework, we could compute the distance between o 1 2 M 1 and o 2 2 M 2 by ﬁrst estimating the latent state with z = f ( o ) using a latent state decoder f , and then using the following distance as a measure of dissimilarity

Original: [Equation]
Revised: [Equation]

Original: However, the distance between GVFs from two different POMDPs can have drastically different scales,i.e.
Revised: However, the distance between GVFs from two different POMDPs can have drastically different scales, [Equation] [Equation]

Original: [Equation]
Revised: [Equation]

Original: , thus making point-wise comparison meaningless.
Revised: , thus making point-wise comparison meaningless.

Original: Theissue is less acute for cumulants which are homogenous between different POMDPs (e.g. indicatorfunctions for successor representation), and more problematic when the cumulant incorporates amore heterogeneous signal, such as the extrinsic reward function.
Revised: The issue is less acute for cumulants which are homogenous between different POMDPs (e.g. indicator functions for successor representation), and more problematic when the cumulant incorporates a more heterogeneous signal, such as the extrinsic reward function.

Original: To avoid this problem, we suggestperforming a comparison based on order statistics.
Revised: To avoid this problem, we suggest performing a comparison based on order statistics.

Original: Namely, a robust distance estimate between GVF signals across POMDPs can be obtained by lookingat the cumulative distribution function of G i denoted F i ( g ) =
Revised: Namely, a robust distance estimate between GVF signals across POMDPs can be obtained by looking at the cumulative distribution function of G i denoted F i ( g ) =

Original: P [ G i ( o t )  g ] for all o t 2 O .
Revised: P [ G i ( o t )  g ] for all o t 2 O .

Original: G i is adeterministic GVF with the set of discontinuity points of measure 0, and as such F i can be understoodthrough the induced state distribution P µt (using continuous mapping theorem from Mann and Wald[1943]).
Revised: G i is a deterministic GVF with the set of discontinuity points of measure 0, and as such F i can be understood through the induced state distribution P µ t (using continuous mapping theorem from [63]).

Original: It can be estimated from n independent and identically distributed samples of D µ as
Revised: It can be estimated from n independent and identically distributed samples of D µ as

Original: [Equation]
Revised: [Equation]

Original: i
Revised: and its inverse, the empirical quantile function [64]

Original: [Equation]
Revised: [Equation]

Original: and its inverse, the empirical quantile function [van der Vaart, 1998]
Revised: p

Original: [Equation]
Revised: [Equation]

Original: for p 2 [0 , 1] .
Revised: for p 2 [0 , 1] .

Original: We use the empirical quantile function to partition the range of all GVFs into K quantilebins, i.e. disjoint sets with identical size where the set corresponding to quantile k is deﬁned asI i ( k ) = { [Equation] M i : F 
Revised: We use the empirical quantile function to partition the range of all GVFs into K quantile bins, i.e. disjoint sets with identical size where the set corresponding to quantile k is deﬁned as

Original: [Equation] i ( 
Revised: I i ( k ) =

Original: [Equation] G µi [Equation] o )  F  1 i ( k +1 K ) } and its aggregated version as I ( k ) =
Revised: { o 2 M i : F  1 i ( kK )  G µi ( o )  F  1 i ( k +1 K ) } and its aggregated version as I ( k ) =

Original: [ mi =1 I i ( k ) 3 .
Revised: [ mi =1 I i ( k ) 4 .

Original: Importantly, we augment the dataset D µ with observation-speciﬁc labels, which correspond to the indexof the quantile bin into which the GVF G of an observation o 2 M i falls into l i ( o )=
Revised: Importantly, we augment the dataset D µ with observation-speciﬁc labels, which correspond to the index of the quantile bin into which the GVF G of an observation o 2 M i falls into l i ( o )=

Original: max k
Revised: max k

Original: o 

Revised: o

Original: I i ( k ) .
Revised: I i ( k ) .

Original: These self-supervised labels are then used in a multiclass InfoNCE loss [Oord et al., 2018], whichis a variation of metric learning with respect to the quantile distance deﬁned above [Khosla et al., 2020, Song and Ermon, 2020] and this forms the basis of our self-supervised learning objective.
Revised: These self-supervised labels are then used in a multiclass InfoNCE loss [47], which is a variation of metric learning with respect to the quantile distance deﬁned above [65, 66] and this forms the basis of our self-supervised learning objective.

Original: 5.3 Self-supervised learning of GSFs
Revised: 5.3 Self-supervised learning of GSFs

Original: After augmenting the ofﬂine dataset with observation labels as described above, we use a simpleself-supervised learning procedure to minimize distance in the latent representation space betweenobservations with identical labels.
Revised: After augmenting the ofﬂine dataset with observation labels as described above, we use a simple self-supervised learning procedure to minimize distance in the latent representation space between observations with identical labels.

Original: First, the observation o is encoded using a non-linear encoder f : O !
Revised: First, the observation o is encoded using a non-linear encoder f : O !

Original: Z with parameters into alatent state representation z
Revised: Z with parameters into a latent state representation z

Original: = f ( o ) 4 .
Revised: = f ( o ) 5 .

Original: The representation z is then passed into two separate trunks:1) a linear matrix ✓ a which recovers the state-action value function Q ✓ ( o,a )=
Revised: The representation z is then passed into two separate trunks: 1) a linear matrix ✓ a which recovers the state-action value function Q ✓ ( o,a ) =

Original: ✓ > a z , and 2) a non-linearprojection network h ✓ : Z !
Revised: ✓ > a z , and 2) a non-linear projection network h ✓ : Z !

Original: Z with parameters ✓ h to obtain a new embedding, used for self-supervised233 learning.
Revised: Z with parameters ✓ h to obtain a new embedding, used for self-supervised learning.

Original: The projection h ✓ ( z ) is then used in a multiclass InfoNCE loss [Oord et al., 2018, Song and Ermon, 2020] where a linear classiﬁer W 2 R |Z| ⇥ K aims to correctly predict the observation labels(i.e. quantile bins k =1 , 2 ,..,K ) from h ✓ ( z ) , for temperature ⌧ > 0 :
Revised: The projection h ✓ ( z ) is then used in a multiclass InfoNCE loss [47, 66] where a linear classiﬁer W 2 R |Z| ⇥ K aims to correctly predict the observation labels (i.e. quantile bins k =1 , 2 ,..,K ) from h ✓ ( z ) , for temperature ⌧ > 0 :

Original: [Equation]
Revised: [Equation]

Original: 5.4 Full Algorithm
Revised: 5.4 Full Algorithm

Original: Given m training tasks, GSF ﬁrst learns GVF estimates G µ 1 , .., G µ 2 by applying LearnGVF totask-speciﬁc data from the ofﬂine dataset D µ .
Revised: Given m training tasks, GSF ﬁrst learns GVF estimates G µ 1 , .., G µ 2 by applying LearnGVF to task-speciﬁc data from the ofﬂine dataset D µ .

Original: Each data point in D µ is then labeled with the quantileinto which its GVF falls.
Revised: Each data point in D µ is then labeled with the quantile into which its GVF falls.

Original: These labels are then used to jointly optimize Eq.6 and a control loss withrespect to the encoder parameters.
Revised: These labels are then used to jointly optimize Eq.6 and a control loss with respect to the encoder parameters.

Original: To learn the value function Q ✓ , we use CQL [Kumar et al. , 2020],which is trained using a linear combination of Q-learning [Watkins andDayan, 1992, Ernst et al. , 2005]and behavior-regularization:
Revised: To learn the value function Q ✓ , we use CQL [56], which is trained using a linear combination of Q-learning [67, 16] and behavior-regularization:

Original: [Equation]
Revised: [Equation]

Original: for   0 , ˜ ✓ target network parameters 5 and LSE being the log-sum-exp operator 6 .
Revised: for   0 , ˜ ✓ target network parameters 6 and LSE being the log-sum-exp operator 7 .

Original: For domains withcontinuous actions, we also decode the Bolztmann policy ⇡ from Q ✓ .
Revised: For domains with continuous actions, we also decode the Bolztmann policy ⇡ from Q ✓ .

Original: Fig.
Revised: Fig.

Original: 2 provides a schematic view of the algorithm, while Alg. 2 in the Appendix presents the exactlearning procedure for GSF as implemented on top of a CQL agent for a discrete action space.
Revised: 2 provides a schematic view of the algorithm, while Alg. 2 in the Appendix presents the exact learning procedure for GSF as implemented on top of a CQL agent for a discrete action space.

Original: Recovering existing self-supervised objectives The generality of our framework allows it torecover existing objectives such as CSSC and PSEs by carefully designing the cumulant function.
Revised: Recovering existing self-supervised objectives The generality of our framework allows it to recover existing objectives such as CSSC and PSEs by carefully designing the cumulant function.

Original: Below, we highlight which existing algorithms can be recovered by GSF.
Revised: Below, we highlight which existing algorithms can be recovered by GSF.

Original: • Cross-State Self-Constraint [CSSC, Liu et al., 2020a] : In CSSC, observations o 1 ,o 2 areconsidered similar if they have identical future action sequences of length K under some ﬁxedpolicy; a total of |A| K distinct classes are possible.
Revised: • Cross-State Self-Constraint [CSSC, 10] : In CSSC, observations o 1 ,o 2 are considered similar if they have identical future action sequences of length K under some ﬁxed policy; a total of |A| K distinct classes are possible.

Original: This approach can be approximated in ourframework by picking c ( o t ,a t )
Revised: This approach can be approximated in our framework by picking c ( o t ,a t )=

Original: = a t ( a ) , 8 a 2 A .
Revised: a t ( a ) , 8 a 2 A .

Original: The problem reduces to a |A| T  t -way classiﬁcation problem for observations of timestep t , which GSF approximates using K quantiles.
Revised: The problem reduces to a |A| T  t -way classiﬁcation problem for observations of timestep t , which GSF approximates using K quantiles.

Original: • Policy similarity embedding [PSE, Agarwal et al., 2021a] : PSEs balance the distance between local optimal behaviors and long-term dependencies in the transitions, notablyusing d TV .
Revised: • Policy similarity embedding [PSE, 11] : PSEs balance the distance between local optimal behaviors and long-term dependencies in the transitions, notably using d TV .

Original: If we consider the space of Boltzmann policies ⇡ Boltzmann with respect to aPOMDP-speciﬁc value function Q , then choosing c ( o t ,a t )=
Revised: If we consider the space of Boltzmann policies ⇡ Boltzmann with respect to a POMDP-speciﬁc value function Q , then choosing c ( o t ,a t )=

Original: r ( s t ,a t ) in GSF will effectivelycompute the distance between unnormalized policies.
Revised: r ( s t ,a t ) in GSF will effectively compute the distance between unnormalized policies.

Original: The choice of K induces a bias-variance trade-off How should the number of quantiles K (readlabels) be set, and what is the effect of smaller/ larger values of K on the learned representations?
Revised: The choice of K induces a bias-variance trade-off How should the number of quantiles K (read labels) be set, and what is the effect of smaller/ larger values of K on the learned representations?

Original: Thm.
Revised: Thm.

Original: 1 highlights a trade-off when choosing the number of quantile bins empirically.
Revised: 1 highlights a trade-off when choosing the number of quantile bins empirically.

Original: Theorem 1 Let G 1 , G 2 be generalized value functions with cumulants c 1 ,c 2 from respective POMDPs
Revised: Theorem 1 Let G 1 , G 2 be generalized value functions with cumulants c 1 ,c 2 from respective POMDPs

Original: M 1 , M 2 , K be the number of quantile bins, n 1 , n 2 the number of sample transitions from each
Revised: M 1 , M 2 , K be the number of quantile bins, n 1 , n 2 the number of sample transitions from each

Original: POMDP.
Revised: POMDP.

Original: Suppose that P [sup t =1 , 2 ,.. | c 1 ( o 1 ,t ,µ ( o 1 ,t ))
Revised: Suppose that P [sup t =1 , 2 ,.. | c 1 ( o 1 ,t ,µ ( o 1 ,t ))

Original:  c 2 ( o 2 ,t ,µ ( o 2 ,t )) | > (1   ) " /  ]   .
Revised:  c 2 ( o 2 ,t ,µ ( o 2 ,t )) | > (1   ) " /  ]   .

Original: Then, for any266 k =1 , 2 ,..,K and " > 0 the following holds without loss of generality:
Revised: Then, for any k =1 , 2 ,..,K and " > 0 the following holds without loss of generality:

Original: [-P ï sup o 1 ,o 2 
-]
Revised: [-P ï sup o 1 ,o 2 
-]

Original: I ( k ) | G 1 ( o 1 )  G 2 ( o 2 ) | > 3 " ò  2 e  2 n 1 " 2 / 4 + P ï sup k =1 , 2 ,..,K  ˆ F  1 k +1 / K   ˆ F  1 k / K  > " ò + p ( n 1 ,K, " ) +  (8)
Revised: [Equation]

Original: The proof can be found in the Appendix Sec.
Revised: The proof can be found in the Appendix Sec.

Original: A.5.
Revised: A.5.

Original: For POMDP M 1 , the error decreases monotonicallywith increasing bin number K (second term) but the variance of bin labels depends on the numberof sample transitions n 1 (ﬁrst term).
Revised: For POMDP M 1 , the error decreases monotonically with increasing bin number K (second term) but the variance of bin labels depends on the number of sample transitions n 1 (ﬁrst term).

Original: The inter-POMDP error (third term) does not affect the binassignment.
Revised: The inter-POMDP error (third term) does not affect the bin assignment.

Original: Hence, choosing a large K will amount to pairing states by rankings, but results in highhttps://en. variance, as orderings are estimated from data and each bin will have n =1 .
Revised: Hence, choosing a large K will amount to pairing states by rankings, but results in high variance, as orderings are estimated from data and each bin will have n =1 .

Original: Setting K too small willgroup together unrelated observations, inducing high bias.
Revised: Setting K too small will group together unrelated observations, inducing high bias.

Original: Limitations As is the case with all ofﬂine RL methods, GSF is limited by the compoundingextrapolation error under low data coverage.
Revised: Limitations As is the case with all ofﬂine RL methods, GSF is limited by the compounding extrapolation error under low data coverage.

Original: We hypothesize that wise choices of K and c can mitigatethe extrapolation error by learning observation groups with low intra-group variance, but, since theyare environment-dependent, searching for an optimal ( K,c ) pair can be computationally expensive.
Revised: We hypothesize that wise choices of K and c can mitigate the extrapolation error by learning observation groups with low intra-group variance, but, since they are environment-dependent, searching for an optimal ( K,c ) pair can be computationally expensive.

Original:  Experiments
Revised:  Experiments

Original: Unlike for single task ofﬂine RL [Fu et al., 2020], most prior work on zero-shot generalization fromofﬂine data either come up with an ad hoc solution suiting their needs, e.g. Ball et al. [2021], or assessperformance on benchmarks that do not evaluate generalization across observation functions [e.g., Yuet al., 2020].
Revised: Unlike for single task ofﬂine RL [26], most prior work on zero-shot generalization from ofﬂine data either come up with an ad hoc solution suiting their needs, e.g. [33], or assess performance on benchmarks that do not evaluate generalization across observation functions [e.g., 68].

Original: To accelerate progress in this ﬁeld, we devised two benchmarks: ofﬂine Procgen (discreteactions) and ofﬂine Distracting Suite (continuous actions) - two ofﬂine RL datasets to directly testfor generalization of RL agents across observation functions 7 .
Revised: To accelerate progress in this ﬁeld, we devised two benchmarks: ofﬂine Procgen (discrete actions) and ofﬂine Distracting Suite (continuous actions) - two ofﬂine RL datasets to directly test for generalization of RL agents across observation functions 8 .

Original: Moreover, for a standard comparison,we provide generalization results on the classical online Procgen simulator, comparing PPO to PPOwith GSF and PPO with PSE, respectively.
Revised: Moreover, for a standard comparison, we provide generalization results on the classical online Procgen simulator, comparing PPO to PPO with GSF and PPO with PSE, respectively.

Original: GVF training In the ofﬂine setting, the training dataset is used to learn a set of task-speciﬁc GVFs inparallel as follows.
Revised: GVF training In the ofﬂine setting, the training dataset is used to learn a set of task-speciﬁc GVFs in parallel as follows.

Original: First, we project each observation o into a latent representation z
Revised: First, we project each observation o into a latent representation z

Original: = f 

Revised: = f

Original: ( o ) ; we thenpass z through a non-linear network h ✓ 0 : Z !
Revised: ( o ) ; we then pass z through a non-linear network h ✓ 0 : Z !

Original: R d c ⇥ m where d c is the dimensionality of the cumulantfunction’s output.
Revised: R d c ⇥ m where d c is the dimensionality of the cumulant function’s output.

Original: The output of h ✓ 0 is then split into m disjoint chunks, which are in turn used in theirrespective temporal difference losses ` TD in place of value functions.
Revised: The output of h ✓ 0 is then split into m disjoint chunks, which are in turn used in their respective temporal difference losses ` TD in place of value functions.

Original: The procedure can be adaptedto an online setting via a similar procedure, except that all GVF estimators are trained simultaneouslyand independently in separate simulators.
Revised: The procedure can be adapted to an online setting via a similar procedure, except that all GVF estimators are trained simultaneously and independently in separate simulators.

Original: Ofﬂine Procgen benchmark We evaluate the proposed approach on an ofﬂine version of the
Revised: Ofﬂine Procgen benchmark We evaluate the proposed approach on an ofﬂine version of the

Original: Procgen benchmark [Cobbe et al., 2020], which is widely used to evaluate zero-shot generalizationacross complex visual perturbations.
Revised: Procgen benchmark [8], which is widely used to evaluate zero-shot generalization across complex visual perturbations.

Original: Given a random seed, Procgen supports sampling procedurallygenerated level conﬁgurations for 16 games under various complexity modes: “easy”, “hard” and“exploration”.
Revised: Given a random seed, Procgen supports sampling procedurally generated level conﬁgurations for 16 games under various complexity modes: “easy”, “hard” and “exploration”.

Original: More details can be found in Appendix.
Revised: More details can be found in Appendix.

Original: Ofﬂine Procgen results We compare the zero-shot performance on the entire distribution of "easy" POMDPs for GSF against that of strong RL and representation learning baselines: behavioral cloning(BC) - to assess the quality of the PPO policy, CQL [Kumar et al., 2020] - the current state-of-the-arton multiple ofﬂine benchmarks which balances RL and BC objectives, CURL [Srinivas et al., 2020],
Revised: Ofﬂine Procgen results We compare the zero-shot performance on the entire distribution of "easy" POMDPs for GSF against that of strong RL and representation learning baselines: behavioral cloning (BC) - to assess the quality of the PPO policy, CQL [56] - the current state-of-the-art on multiple ofﬂine benchmarks which balances RL and BC objectives, CURL [35], CTRL [12],

Original: CTRL [Mazoure et al., 2022], DeepMDP [Gelada et al., 2019] - which learns a metric closely relatedto bisimulation across the MDP, Value Prediction Network [VPN, Oh et al., 2017] - which combinesmodel-free and model-based learning of values, observations, next observations, rewards and discounts, Cross-State Self-Constraint [CSSC, Liu et al., 2020a] - which boosts similarity of observations withidentical action sequences, as well as Policy Similarity Embeddings [Agarwal et al., 2020], whichgroups observation representations based on distance in optimal policy space.
Revised: DeepMDP [69] - which learns a metric closely related to bisimulation across the MDP, Value Prediction Network [VPN, 70] - which combines model-free and model-based learning of values, observations, next observations, rewards and discounts, Cross-State Self-Constraint [CSSC, 10] - which boosts similarity of observations with identical action sequences, as well as Policy Similarity Embeddings [34], which groups observation representations based on distance in optimal policy space.

Original: Fig.
Revised: Fig.

Original: 3 shows the performance of all methods over 5 random seeds and all 16 games on the ofﬂine
Revised: 3 shows the performance of all methods over 5 random seeds and all 16 games on the ofﬂine

Original: Procgen benchmark after 1 million gradient steps.
Revised: Procgen benchmark after 1 million gradient steps.

Original: Per-game average scores for all methods can befound in Tab. 2 (Appendix).
Revised: Per-game average scores for all methods can be found in Tab. 2 (Appendix).

Original: The scores are standardized per-game using the downstream task’s (ofﬂine
Revised: The scores are standardized per-game using the downstream task’s (ofﬂine

Original: RL) performance, in this case implemented by CQL.
Revised: RL) performance, in this case implemented by CQL.

Original: It can be seen that GSF performs better thanother ofﬂine RL and representation learning baselines.
Revised: It can be seen that GSF performs better than other ofﬂine RL and representation learning baselines.

Original: Using different cumulant functions can lead to different label assignments and hence differentsimilarity groups.
Revised: Using different cumulant functions can lead to different label assignments and hence different similarity groups.

Original: Fig.
Revised: Fig.

Original: 3 examines the performance of GSF with respect to 3 cumulants: 1) r ( s t ,a t ) ,rewards s.t.
Revised: 3 examines the performance of GSF with respect to 3 cumulants: 1) r ( s t ,a t ) , rewards s.t.

Original: GSF learns the policy’s Q µ -value, 2)
Revised: GSF learns the policy’s Q µ -value, 2)

Original: o t ( o ) , the successor representation 8 [Dayan, 1993, Barreto et al., 2016] s.t.
Revised: o t ( o ) , the successor representation 9 [62, 32] s.t.

Original: GSF learns the distribution induced by µ over D µ [Machado et al. 2020]and 3) a t ( a ) , action counts, s.t.
Revised: GSF learns the distribution induced by µ over D µ [71] and 3) a t ( a ) , action counts, s.t.

Original: GSF learns discounted policy.
Revised: GSF learns discounted policy.

Original: While rewards and successor featurecumulant choices leads to similar performance, using action-based distance leads to larger variance.
Revised: While rewards and successor feature cumulant choices leads to similar performance, using action-based distance leads to larger variance.

Original: In the continuous observation space, we learn a d -dimensional successor feature vector z via TD and compute the quantiles over || z || 1 .

Original: Ofﬂine Distracting Control Suite results Following the same procedure as for ofﬂine Procgenresults, we ﬁrst formed an ofﬂine dataset from the challenging Distracting Control Suite [Stone et al. ,2021] by saving the replay buffer of Soft Actor-Critic [Haarnoja et al., 2018] trained for 1M frames on 4environments with 2 different background perturbations.
Revised: Ofﬂine Distracting Control Suite results Following the same procedure as for ofﬂine Procgen results, we ﬁrst formed an ofﬂine dataset from the challenging Distracting Control Suite [9] by saving the replay buffer of Soft Actor-Critic [72] trained for 1M frames on 4 environments with 2 different background perturbations.

Original: Next, we pre-trained G µ 1 ,G µ 2 with action andreward-based cumulants, which were then used in conjunction with CQL to learn a single multi-taskpolicy.
Revised: Next, we pre-trained G µ 1 ,G µ 2 with action and reward-based cumulants, which were then used in conjunction with CQL to learn a single multi-task policy.

Original: Fig 4 shows the online performance of GSF evaluated on 10 background perturbations notseen during training, normalized by per-environment median CQL scores.
Revised: Fig 4 shows the online performance of GSF evaluated on 10 background perturbations not seen during training, normalized by per-environment median CQL scores.

Original: The results are consistent with ﬁndings of Agarwal et al. [2021a], in that 1) learning policy-basedsimilarity improves generalization capabilities of state representations, and 2) unlike in Procgen,policy-based similarity provides a better learning signal than value-based similarity.
Revised: The results are consistent with ﬁndings of [11], in that 1) learning policy-based similarity improves generalization capabilities of state representations, and 2) unlike in Procgen, policy-based similarity provides a better learning signal than value-based similarity.

Original: Online Procgen results We additionally compare performance of GSF to that of PPO and PPO with PSEs in the classical online Procgen benchmark [Cobbe et al., 2019].
Revised: Online Procgen results We additionally compare performance of GSF to that of PPO and PPO with PSEs in the classical online Procgen benchmark [5].

Original: Figure 8 shows test returns for20M frames obtained on the entire distribution of easy levels while training on 200 easy levels for PPO,
Revised: Figure 8 shows test returns for 20M frames obtained on the entire distribution of easy levels while training on 200 easy levels for PPO,

Original: PPO+PSEs [Agarwal et al., 2021a] and our PPO+GSFs.
Revised: PPO+PSEs [11] and our PPO+GSFs.

Original: PPO+GSF outperforms or matches both PPOand PPO+PSE on most environments, showing that GSF can be efﬁciently combined with both ofﬂineand online algorithms.
Revised: PPO+GSF outperforms or matches both PPO and PPO+PSE on most environments, showing that GSF can be efﬁciently combined with both ofﬂine and online algorithms.

Original: See Appendix A.6.1 for detailed results.
Revised: See Appendix A.6.1 for detailed results.

Original:  Discussion
Revised:  Discussion

Original: In this work we proposed Generalized Similarity Functions, a novel framework which combinesreinforcement learning with representation learning to improve zero-shot generalization performanceon challenging, pixel-based control tasks.
Revised: In this work we proposed Generalized Similarity Functions, a novel framework which combines reinforcement learning with representation learning to improve zero-shot generalization performance on challenging, pixel-based control tasks.

Original: GSF relies on computing the similarity between observationpairs with respect to any instantaneous accumulated signal, which leads to improved empiricalperformance on two newly introduced benchmarks, ofﬂine Procgen and ofﬂine Distracting Suite.
Revised: GSF relies on computing the similarity between observation pairs with respect to any instantaneous accumulated signal, which leads to improved empirical performance on two newly introduced benchmarks, ofﬂine Procgen and ofﬂine Distracting Suite.

Original: Theoretical results suggest that GSF’s hyperparameter choice depends on a bias-variance trade-off.
Revised: Theoretical results suggest that GSF’s hyperparameter choice depends on a bias-variance trade-off Checklist

Revised: For all authors...

Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Revised: [Yes] (b) Did you describe the limitations of your work?

Revised: [Yes] See end of Section(c)

Revised: Did you discuss any potential negative societal

Revised: impacts of your work?

Revised: [Yes] We include a brief discussion in Appendix Section A.

Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Revised: [Yes] 2.

Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Revised: [Yes] See Appendix Section A. 3.

Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Revised: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Revised: [Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Revised: [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Revised: [Yes] See Appendix Section A. 4.

Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Revised: [N/A] (b) Did you mention the license of the assets?

Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Revised: [Yes] Due to the large size of the ofﬂine datasets, we include the code which can re-generate the data from scratch.

Revised: (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?

Revised: [N/A] (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Revised: [N/A] 5.

Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Revised: ] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Revised: [N/A]