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)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: 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 sequentialdecision making tasks, but still exhibit difﬁculty in generalizing to scenarios notseen 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 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.

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 performanceof online algorithms for generalization in RL can be hindered in the ofﬂine settingdue 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 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 .

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 enoughto recover existing SSL objectives while also improving zero-shot generalizationperformance 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 asequence 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 [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].

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 apotentially distinct sampling distribution during test time [Cobbe et al., 2019, Song et al., 2019, Hjelmet al., 2022].

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

Original: (limited) training a broader set of problems.
Revised: (limited) training 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, thisproblem 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.,
Revised: An example of this setting is any environment with distractor features [Cobbe et al.,

Original: Stone et al., 2021], i.e. features with no dependence on the reward signal nor the agent’s decisions.
Revised: Stone et al., 2021], 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 [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

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

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 withlimited data (i.e. ofﬂine RL) is challenging, since the representation will have a large estimation errorover 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 andextracting the most information out of the data consists in introducing auxiliary learning signals otherthan instantaneous reward [Raileanu and Fergus, 2021].

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

Original: ,2022] can outperform it by leveraging additional auxiliary signals.
Revised: ,2022] 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 theaforementioned methods is their exorbitant reliance on online access to the environment, an impracticalrestriction 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, ﬁxeddataset of experience [Ernst et al., 2005, Lange et al., 2012].

Original: A natural limitation for generalization fromofﬂine data is that policy improvement is dependent on dataset quality.
Revised: A natural limitation for generalization fromofﬂ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-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.

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

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, Murphy, 2000] 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 togeneralize across perceptually different POMDPs without adaptation, observations with similarfuture behavior should be assigned 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 [Sutton et al., 2011] to capture future behavior with respect to any instantaneoussignal (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 choice of cumulant determines the natureof the behavioral similarity that is induced into state representations.

Original: For example, using rewardsimilarity leads to learning bisimulation metrics [Ferns et al., 2004, Li et al., 2006, Castro, 2020,
Revised: For example, using rewardsimilarity leads to learning bisimulation metrics [Ferns et al., 2004, Li et al., 2006, Castro, 2020,

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

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 learningalgorithm for reinforcement learning that aggregates latent representations of observationsby their future behavior (or generalized value function).2.

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 [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.

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 16distinct games; the second consists of 1M transitions from 3 variations of 4 distincts tasks.3.

Original: We evaluate performance of GSF and other baseline methods on both benchmarks, and show
Revised: We evaluate performance of GSF and other baseline methods on both benchmarks, and show

Original: that GSF outperforms both previous state-of-the-art ofﬂine RL and representation learningbaselines on the entire distribution of levels.
Revised: that GSF outperforms both previous state-of-the-art ofﬂine RL and representation learningbaselines on the entire distribution of levels.

Original: We analyze the theoretical properties of GSF and describe the impact of hyperparameters and
Revised: We analyze the theoretical properties of GSF and describe the impact of hyperparameters and

Original: cumulant functions on empirical behavior
Revised: cumulant functions on empirical behavior

Original: in both ofﬂine Procgen and the ofﬂine Distracting
Revised: in both ofﬂine Procgen and the ofﬂine Distracting

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

Original: Generalization in RL has received a lot of attention: extrapolation to unseen rewards [Barreto et al.
Revised: Generalization in RL has received a lot of attention: extrapolation to unseen rewards [Barreto et al.

Original: ,2016, Misra et al., 2020], observations [Zhang et al., 2020, Raileanu and Fergus, 2021, Liu et al.
Revised: ,2016, Misra et al., 2020], observations [Zhang et al., 2020, Raileanu and Fergus, 2021, Liu et al.

Original: ,2020a, Agarwal et al., 2021a, Mazoure et al., 2022] and transition dynamics
Revised: ,2020a, Agarwal et al., 2021a, Mazoure et al., 2022] and transition dynamics

Original: [Ball et al., 2021].
Revised: [Ball et al., 2021].

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

Original: ,2020, Stooke et al., 2021] or data augmentation [Ball et al., 2021, Sinha and Garg, 2021].
Revised: ,2020, Stooke et al., 2021] or data augmentation [Ball et al., 2021, Sinha and Garg, 2021].

Original: Dataaugmentation is a promising technique, but typically relies on handcrafted domain information, whichmight not be available a priori .
Revised: Dataaugmentation is a promising technique, but typically relies on handcrafted domain information, whichmight 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 theofﬂine RL setting is poor even when using such handcrafted data augmentations, without additionalrepresentation 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 useinstantaneous auxiliary signals in order to prevent overﬁtting onto a unique signal (e.g. reward acrosstasks) and improve generalization performance.

Original: Theoretical generalization guarantees have only beenprovided so far for limited scenarios, mostly for bandits [Swaminathan and Joachims, 2015], linear
Revised: Theoretical generalization guarantees have only beenprovided so far for limited scenarios, mostly for bandits [Swaminathan and Joachims, 2015], linear

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

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

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

Original: Theproblemoflearningmeaningfulstaterepresentations(or abstractions) for planning and control has been extensively studied previously [Jong and Stone, 2005, Li et al., 2006], but saw real breakthroughs only recently, in particular due to advances in self-supervisedlearning(SSL).
Revised: 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).

Original: Outside ofRL, SSLhas achievedspectacular resultsby closingthe gapbetween unsupervised and supervised learning on certain datasets [Hjelm et al., 2018, Oord et al., 2018, Caron et al., 2020,
Revised: 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.

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

Original: Representation learning, and speciﬁcally self-supervised learning, has also been usedto achieve state-of-the-art generalization and sample efﬁciency results in RL on challenging controlproblems such as data efﬁcient Atari 
Revised: 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],

Original: [Schwarzer et al., 2020, 2021], DeepMind Control
Revised: DeepMind Control

Original: [Agarwal et al. ,2021a]andProcgen[Mazoureetal. ,2020,Stookeetal.
Revised: [Agarwal et al., 2021a] and Procgen [Mazoure et al., 2020, Stooke et al., 2021,

Original: ,2021,RaileanuandFergus,2021,Mazoureetal.
Revised: Raileanu and Fergus, 2021, Mazoure et al., 2022].

Original: Noteworthy instances of theoretically-motivated representation learning methods for RL includeheuristic-guided learning [Cheng et al., 2021],and random Fourier features [Nachum and Yang, 2021].
Revised: 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].

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 balanceinterpolation and extrapolation errors, while ensuring proper diversity of actions (i.e. prevent collapseto most frequent action in the data).

Original: Popular ofﬂine RL algorithms such as BCQ [Fujimoto et al., 2019],
Revised: Popular ofﬂine RL algorithms such as BCQ [Fujimoto et al., 2019],

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

Original: Some methods, such as F-BRC [Kostrikov et al.
Revised: Some methods, such as F-BRC [Kostrikov et al.

Original: ,2021a] are deﬁned only for continuous action spaces while others, such as MOReL [Kidambi et al.
Revised: ,2021a] are deﬁned only for continuous action spaces while others, such as MOReL [Kidambi et al.

Original: ,2020] estimate a pessimistic transition model.
Revised: ,2020] 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 algorithmssuch as CQL is that they are perhaps overly pessimistic for generalization purposes, i.e. CQL and

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

Original:  Problem setting
Revised:  Problem setting

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

Original: A (inﬁnite-horizon) partially-observable Markov decision process [POMDP, Murphy, 2000]
Revised: A (inﬁnite-horizon) partially-observable Markov decision process [POMDP, Murphy, 2000]

Original: 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: 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 =

Original: P [ s 0 ] is the starting 126 state distribution,
Revised: P [ s 0 ] is the starting 127 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 is 127 an observation space, p
Revised:  ( S ) is a transition function, O is 128 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 ) 1 is an observation function, r : S ⇥ A !

Original: [ r min ,r max ] 128 is a reward function and  2 [0 , 1) is a discount factor.
Revised: [ r min ,r max ] 129 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 statess

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 by 130 a policy ⇡ :
Revised: At every timestep t = 1 , 2 , 3 ,.. , the agent, parameterized by 131 a policy ⇡ :

Original: O !
Revised: 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 state 131 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 132 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 overthe 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-actionvalue 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 executing 136 ⇡ 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 137 ⇡ 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 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.

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 assume 142 that both
Revised: We assume 143 that both

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

Original: O .
Revised: O .

Original: Importantly, since we perform control from pixels, we are in the
Revised: Importantly, since we perform control from pixels, we are in the

Original: POMDP setting [see Yarats et al., 2019] and therefore emphasize the difference between observationso t and corresponding states s t throughout the paper.
Revised: POMDP setting [see Yarats et al., 2019] and therefore emphasize the difference between observationso 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 generalizationperformance.

Original: Among others, Liu et al.
Revised: Among others, Liu et al.

Original: [2020a], Agarwal et al.
Revised: [2020a], Agarwal et al.

Original: [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: [2021a] deﬁne the similarity of twoobservations to depend on the distance between action sequences rolled out from that observation undertheir respective optimal policies.

Original: They achieve this by ﬁnding a latent space Z ✓ S in which the distance 152 d
Revised: They achieve this by ﬁnding a latent space Z ✓ S in which the distance 153 d

Original: 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 ; 153 the aforementionned works learn
Revised: 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

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 fromobservation 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 thelatent 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 implementationof 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 statesfrom 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: 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.

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 [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 .

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

Original: For instance, two trajectoriesmight visit an identical set of latent states, but have drastically different actions.
Revised: For instance, two trajectoriesmight 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 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).

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  hence 
Revised: 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

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 abetterstaterepresentationthanreasoningaboutactionsimilarities.

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

Original: Sohowdoweunifythesesimilarity metrics under a single framework?
Revised: Sohowdoweunifythesesimilarity 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 generalway 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 valuefunctions, 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 determineswhich components of the future trajectory are most relevant for generalization.

Original: some cumulant function c , whose quantiles are then used to label each observation inthe
Revised: 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 otherauxiliary signals, for example, by its future observation occupancy measure, known as successorrepresentation

Original: [Dayan, 1993, Barreto et al., 2016].
Revised: [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 iterationmeasures the properties the agent will exhibit in the future, such as accumulated returns, actions,or observation visitation density.

Original: See Thm.
Revised: See Thm.

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

Original: Following the work of Sutton et al.
Revised: Following the work of Sutton et al.

Original: [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: [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 !

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

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

Original: While typically cumulants are scalar-valued functions 193 (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 194 (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.

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 amountof 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 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

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,i.e.

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

Original: To avoid this problem, we suggestperforming a comparison based on order statistics.
Revised: To avoid this problem, we suggestperforming 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 lookingat 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 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]).

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: i

Original: [Equation]
Revised: [Equation]

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

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 as
Revised: 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 as

Original: I i ( k ) = {
Revised: I i ( k ) = {

Original: [Equation]
Revised: [Equation]

Original: M i : F 
Revised: M i : F

Original: [Equation]
Revised: [Equation]

Original: i (kK
Revised: i (

Original: [Equation]
Revised: [Equation]

Original: G µi
Revised: G µi

Original: [Equation]
Revised: [Equation]

Original: o )  F  1 i ( k +1 K ) } and its aggregated version as I ( k ) =
Revised: 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 ) 3 .

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 indexof 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,
Revised: 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,

Original: Song and Ermon, 2020] and this forms the basis of our self-supervised learning objective.
Revised: Song and Ermon, 2020] 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 simpleself-supervised learning procedure to minimize distance in the latent representation space betweenobservations 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 alatent state representation z

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

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-linearprojection network h ✓ : Z !

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

Original: The projection h ✓ ( z ) is then used in a multiclass InfoNCE loss [Oord et al., 2018, Song and
Revised: The projection h ✓ ( z ) is then used in a multiclass InfoNCE loss [Oord et al., 2018, Song and

Original: Ermon, 2020] where a linear classiﬁer W 2 R |Z| ⇥ K aims to correctly predict the observation labels(i.e.
Revised: Ermon, 2020] where a linear classiﬁer W 2 R |Z| ⇥ K aims to correctly predict the observation labels(i.e.

Original: quantile bins k
Revised: quantile bins k

Original: =1 , 2 ,..,K ) from h ✓ ( z ) , for temperature ⌧ > 0 :
Revised: =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 totask-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 quantileinto 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 withrespect to the encoder parameters.

Original: To learn the value function Q ✓ , we use CQL [Kumar et al.
Revised: To learn the value function Q ✓ , we use CQL [Kumar et al.

Original: , 2020],which is trained using a linear combination of Q-learning [Watkins and Dayan, 1992, Ernst et al.
Revised: , 2020],which is trained using a linear combination of Q-learning [Watkins and Dayan, 1992, Ernst et al.

Original: , 2005]and behavior-regularization:
Revised: , 2005]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 5 and LSE being the log-sum-exp operator 6 .

Original: For domains withcontinuous actions, we also decode the Bolztmann policy ⇡ from Q ✓ .
Revised: For domains withcontinuous 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 exactlearning 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 torecover 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, 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.

Original: This approach can be approximated in ourframework by picking c ( o t ,a t )
Revised: This approach can be approximated in ourframework 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 distancebetween local optimal behaviors and long-term dependencies in the transitions, notablyusing d TV .
Revised: • 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 .

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 aPOMDP-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 effectivelycompute 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 (readlabels) 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 any 265 k =1 , 2 ,..,K and " > 0 the following holds without loss of generality:
Revised: Then, for any 266 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
Revised: I ( k ) | G 1 ( o 1 )  G 2 ( o 2 ) | > 3 " ò  2 e  2 n 1 " 2 / 4 + P ï sup k

Original: =1 , 2 ,..,K  ˆ F  1 k +1 / K  
Revised: =1 , 2 ,..,K  ˆ F  1 k +1 / K

Original: ˆ F  1 k / K  > " ò + p ( n 1 ,K, " )
Revised: ˆ F  1 k / K  > " ò + p ( n 1 ,K, " )

Original: +  (8)
Revised: +  (8)

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 monotonicallywith increasing bin number K (second term) but the variance of bin labels depends on the numberof 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 binassignment.

Original: Hence, choosing a large K will amount to pairing states by rankings, but results in high https://en.
Revised: Hence, choosing a large K will amount to pairing states by rankings, but results in high https://en.

Original: variance, as orderings are estimated from data and each bin will have n =1 .
Revised: 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 willgroup 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 compoundingextrapolation 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 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.

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

Original: [2021], or assessperformance on benchmarks that do not evaluate generalization across observation functions
Revised: [2021], or assessperformance on benchmarks that do not evaluate generalization across observation functions

Original: [e.g., Yuet al., 2020].
Revised: [e.g., Yuet al., 2020].

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 (discreteactions) and ofﬂine Distracting Suite (continuous actions) - two ofﬂine RL datasets to directly testfor generalization of RL agents across observation functions 7 .

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 PPOwith 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 inparallel 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 thenpass 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 cumulantfunction’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 theirrespective 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 adaptedto an online setting via a similar procedure, except that all GVF estimators are trained simultaneouslyand 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 [Cobbe et al., 2020], which is widely used to evaluate zero-shot generalizationacross 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 procedurallygenerated 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"
Revised: Ofﬂine Procgen results We compare the zero-shot performance on the entire distribution of "easy"

Original: 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: 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],

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,
Revised: 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,

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

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 befound 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 thanother 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 differentsimilarity 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,
Revised: o t ( o ) , the successor representation 8 [Dayan, 1993,

Original: Barreto et al., 2016]
Revised: Barreto et al., 2016]

Original: s.t.
Revised: s.t.

Original: GSF learns the distribution induced by µ over D µ [Machado et al.
Revised: GSF learns the distribution induced by µ over D µ [Machado et al.

Original: 2020]and 3) a t ( a ) , action counts, s.t.
Revised: 2020]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 featurecumulant 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 .
Revised: 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.
Revised: 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.

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

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 andreward-based cumulants, which were then used in conjunction with CQL to learn a single multi-taskpolicy.

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 notseen during training, normalized by per-environment median CQL scores.

Original: The results are consistent with ﬁndings of Agarwal et al.
Revised: The results are consistent with ﬁndings of Agarwal et al.

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

Original: Online Procgen results We additionally compare performance of GSF to that of PPO and PPO with
Revised: Online Procgen results We additionally compare performance of GSF to that of PPO and PPO with

Original: PSEs in the classical online Procgen benchmark [Cobbe et al., 2019].
Revised: PSEs in the classical online Procgen benchmark [Cobbe et al., 2019].

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 for20M frames obtained on the entire distribution of easy levels while training on 200 easy levels for PPO,

Original: PPO+PSEs
Revised: PPO+PSEs

Original: [Agarwal et al., 2021a] and our PPO+GSFs.
Revised: [Agarwal et al., 2021a] 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 PPOand PPO+PSE on most environments, showing that GSF can be efﬁciently combined with both ofﬂineand 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 combinesreinforcement learning with representation learning to improve zero-shot generalization performanceon 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 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.

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.