Improving Zero-Shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions
Authors: Bogdan Mazoure
Ilya Kostrikov
Ofir Nachum
Jonathan Tompson
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
representation learning
self-supervised learning
generalization

Abstract: Reinforcement learning (RL) agents are widely used for solving complex sequential decision-making tasks, but still exhibit difficulty generalizing to scenarios not seen during training. 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 offline RL setting, i.e. learning from a static dataset. We show that the performance of online algorithms for generalization in RL can be hindered in the offline setting due to poor estimation of similarity between observations. We propose a new theoretically-motivated framework called Generalized Similarity Functions (GSF), which uses contrastive learning to train an offline RL agent to aggregate observations based on the similarity of their expected future behavior, where we quantify this similarity using generalized value functions. We show that GSF is general enough to recover existing SSL objectives while improving zero-shot generalization performance on two complex pixel-based offline RL benchmarks.

Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions
Bogdan Mazoure*
Ilya Kostrikov UC Berkeley
McGill University, Quebec AI Institute
bogdan. maazoure@mail.mmgil..
Ofir Nachum
Jonathan Tompson Google Brain
Google Brain
Abstract
Reinforcement learning (RL) agents are widely used for solving complex sequential decision making tasks, but still exhibit difficulty in generalizing to scenarios not seen during training. While prior online approaches demonstrated that using additional signals beyond the reward function can lead to better generalization capabilitiess in RL agents, i.e., using supervupe learning (SSL). they struggle the offline RL setting, i.c.. learning from a static dataset. We show that performance online algorithms for generalization RL can be hindered in the offline setting due to poor estimation of similarity between observations. We propose new eecoreticallymmotivaee
framework called Generalized Similarity Functions (GSF), which uses contrastive learning to train an offline RL agent to aggregate observations based on the similarity of their expected future behavior, where we quantify this similarity using generalized value functions. We show that GSF is general enough to recover existing SSL objectives while also improving zero-shot generalization performance two pixel-based offline RL benchmarks.
Introduction
Reinforcement learning (RL) is powerful framework for solving complex tasks that require sequence decisions. The RL paradigm has allowed for major breakthroughs various fields, e.g. outperforming humans on video games [1, 2]. controlling stratospheric balloons [3] and learning reward functions from robot manipulation videos [4]. More recently, RL agents have been tested in generalization setting, i.c. which training involves finite number of related tasks sampled from some distribution, with a potentially distinct sampling distribution during test time [5, 6, 7]. 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 (limited) training task collection onto broader set of problems. Since the learning agent is given no training data from test-time tasks, this problem is referred to zero shoo generalization. In our work, generalization where the difference between tasks is predominantly due to perceptually distinct observations. An example of this setting any environment with distractor features [8, 9], i.e. features with no dependence on the reward signal nor the agent's 's decisions. This generalization setting has recently received much attention [10, 11. 12], due its particular relevance real-world scenarios. for example deploying the same autonomous driving agent at day or at night.
interested in the problem of zero-shot
Work done while at Google Brain.
36th Conference on Neural Information Processing Systems (NeurIPS 2022).

---

Generalization capabilities of an agent can be analyzed through the prism of representation learning. under which the agent's 's current belief about a rich and high-dimensien environment are summarized in a oowddimenioonal entity, called a representation. Recent work in online RL has shown that learning state representations with specific properties such as disentanglement/ [13] or linear separability [14] can improve zero-shot generalization performance. Achieving this with limited data (i.e. offline RL) is challenging, since the representation will have a large estimation error over regions of low data coverage. A common solution to mitigate this task-specific overfitting and extracting the most information out of the data consists in introducing auxiliary learning signals other than instantaneous reward [15]. As we show later in the paper, many such signals already contained in the dataset can be used to further improve generalization performance. For instance, the generalization performance of PPO on Proegen remains limited even when training on 200M frames, while generalizationnorientee agents [15, 12] can outperform it by leveraging additional auxiliary signals. 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.
In contrast, in many real-world scenarios access to the environment is restricted to offline, fixed dataset of experience [16, 17J. A natural limitation for generalization from offline data is that policy improvement is dependent on dataset quality. Specifically, high-dimensional problems such control from pixels require large amounts of training experience: a standard training of PPO [18] for 25 million frames on Proegen [8] generates more than 300 Gb of data, an impractical amount of data to share for offline RL research. Improving zero-shot generalization performance from an offline dataset of nigh-dimensional observations is therefore a hard problem due to limitations dataset size and quality.
In this work, we are interested in improving zero-shot generalization across a family of Partially- Observable Markov decision processes [POMDPs, 19] in an offline RL setting, i.e. by training agents on a fixed dataset. We hypothesize that in order for an RL agent to be able to generalize across perceptu- ally different POMDPs without adaptation, observations with similar future behavior should be assigned to close representations. 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. Specifically, the choice of cumulant determines the nature of the behavioral similarity that is induced into state representa- tions. For example, using reward similarity leads to learning bisimulation metrics [21, 22, 23,24], 24], while using future state-action visitation counts encourages reward-free behavioral similarity [25, 10, 11, 12].
Our main contributions are as follows:
1. 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).
2. Existing offline RL benchmarks [26, 27] are not well-suited to test zero-shot generalization, and so we devise two new benchmarks: offline Procgen and Distracting Control Suite. The first 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.
3. We evaluate performance of GSF and other baseline methods on both benchmarks, and show that GSF outperforms both previous state-of-the-art offline RL and representation learning basclines on the entire distribution of levels.
4. We analyze the theoretical properties of GSF and describe the impact of hy perparaaeeess and cumulant functions on empirical behavior in both offline Procgen and the offline Distracting Suite benchmarks.
2 Related Works
Generalization in reinforcement learning Generalizing a model's predictions across a variety of unseen, digh-dimensional inputs has been extensively studied in the static supervised learning setting [28, 29, 30, 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]. Each generalization scenario is best solved by their respective set of methods: sufficient exploration [25, 34], auxiliary learning signals [35, 36, 37] or data augmentation 38). Data augmentation is promising technique, but typically relies on handerafted domain information, which might not be available a priori. In
fact, we will show in our experiments that generalization in the offline RL setting is poor even when

---

using such handcrafted data augmentations, without additional representation learning mechanisms. In this work, we posit that representation learning should use instantaneous auxiliary signals in order to prevent overfitting onto unique signal (e.g. reward across tasks) and improve generalization performance. 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].
Representation learning For For simple POMDPs, npar-optimal policies can be found by optimizing for the reward alone. However, more complex settings may require additional auxiliary signals in order to find state abstractions better suited for control. 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). Outside of RL, SSL has achieved spectacular results by closing the gap between unsupervised and supervised learning on certain datasets [46, 47, 48, 49]. Representation learning, and specifically self-supervised learning, has also been used to achieve sate-of-the-art generalization and sample efficiency results in RL on challenging control problems such as data efficient Atari [2, 50], DeepMind Control [1 [1]] and Procgen [36, 37, 15, 12]. Noteworthy instances of theoretically-mottee motiivatdd representation learning methods for RL include heuristic- -uiiied learning [5] I1 random Fourier features [42] and metric learning [52, 53].
Offiine 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). Popular offline RL algorithms such as BCQ [54], MBS [55], and CQL [56] rely on a behavior regularization loss [57] as tool to control the extrapolation error. Some methods, such as F-BRC [58] are defined only for continuous action spaces while others, such as MOReL [59] estimate a pessimistic transition model. The major issue with current offline RL algorithms such as CQL is that they are perhaps overly pessimistic for generalization purposes, i.e. CQL and MBS ensure that the policy improvement is well-supported by the batch of data.
3 Problem setting
3.1 Partially-observable Markov decision processes
(infnnite-horizzon)arriii partiooooo Markov decision process [POMDP, 19] M is defined by the where S is a state space, Po = F(oo] is the starting state distribution, i an action space, pg = IP] S a. J:S A- (SS) is a transition function, O is an observation space. Po :t:: (O) is an observation function, : A + [mmix=]max]] is a reward function and 7 E [0,1) is discount factor. The system starts in one of the initial states so ~~ Po with observation Oo (-(ss) At every timestep t = 1,2,3..., the agent, parameterized by policy π O - (AA). samples an action a ~ ~ (-)oe). The environment transitions into a next state 8141 ~ (ps(-1sp (s, and emits a reward r = =r(s;a;) along with next observation 0++1 ~ (p (-)5s))).
tuple M (S.=
The goal of an RL agent is to maximize the cumulative discounted rewards 2= obtained over the entire episode. Value-based off-policy RL algorithms achieve this by estimating the state-action value function under a target policy n:
2 Q (s) = E E (((44. (,
(1)
for 8 S,S, E and where P... denotes the joint distribution of (8++8 obtained by executing in the environment.
An important distinction from online RL that, in the offline RL setting, we assume access to historical dataset DH (instead of a simulator) collected by logging experience of the policy, μe, in the form (ot,, ii=N.t=T where, for practical purposes, the episode is truncated at T timesteps. Furthermore, we assume that the agent can only be trained on a limited collection of POMDPs Minin {M, and its performance is evaluated on the set of test POMDPs Miss. We assume that both Mirain and Miest were sampled from a common task distribution and that every POMDP M, E M = Mirain U.Miest shares the same transition dynamics and reward function with M but has
(22) denotes the entire set of distributions over the space X.
3

---

different observation function P., Importantly, since we perform control from pixels, we are in the POMDP setting [see 60] and therefore emphasize the difference between observations o, and corresponding states throughout the paper.
3.2 Representation learning
Previous works in the RI. literature have studied the of auxiliary signals improve generalization performance. Among others, [10, 1I] define the similarity of two observations to depend on the distance between action sequences rolled out from that observation under their respective optimal policies, They achieve this by finding latent space Z C S in which the distance dz x') for all 2, z, E Z is equivalent to distance between true latent states dg(s, for all s, S; the aforementionned works learn Z by optimizing action- based similarities between observations. In practice, latent space is decoded from observation o using a latent state decoder : Z from observation o,. Throughout the paper, we assume that all value functions have a linear form in the latent decoded state, i.e. (o(o.a) = - (o) =0_ Zy, which agrees with our practical implementation of all algorithms. Within this model family, the ability of an RI. agent to correctly decode latent states from unseen observations directly affects its policy, and therefore,
generalization capabilities. In the next section, discuss why representation learning is important for offline RL, and how existing action-based similarity metrics fail to recover the true latent states for important families of POMDPs.
Motivating example
Level
Level 2
Actions
8
11
11
7
13
8
0
8
13
Values
8.20
8.05
8.74
9.06
9.29
8.30
8.35
40
8.55
Value ranks
Figure 1: (a) Two levels of the Climber game from the Procgen benchmark [8] with near-identical true latent states and near-identical value functions but drastically different action sequences. (b) Four levels of the jumping task [61] where the constant reward signal makes policy similarity more informative than state value similarity.
Multiple recently proposed self-superviser objectives [10, 11] conjecture that observations 0, E M,,02 E M2 that emit similar future action sequences under optimal policies रं . should be decoded into nearby latent states 22, While this heuristic can correctly group observations with respect to their true latent state simple action spaces, fails to identify similar pairs of trajectories in POMDPs with multiple optimal policies'. For instance, two trajectories might visit an identical set of latent states, but have drastically different actions.
Fig. la 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 same PPO policy (average total variation distance of 0.4423 across episode). 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.
In this scenario, assigning observations to a similar latent state by value function similarity would yield better state representation than reasoning about action similarities. On the other hand, Fig. shows do- main where grouping state representations by action sequences can be optimal. So how do we unify these similarity metries under single framework? the next section, we use this insight design general way of improve representation learning through self-supervised learning of discounted future behavior.
While optimal policies guaranteed to be unique, the optimal value function is unique.

---

Method
We propose measuring a generalized notion future behavior similarity using generalized value functions, defined by the corresponding cumulant function. The choice of cumulant determines which components of the future trajectory are most relevant for generalization.
Duta Ni bos (M()
Daca teck At Iies P'b)
Control LTD
a
Oa
Q(z.
C
Estimate C G3 from [']s] and D
Representation learning
z
h(z)
LGSF (M()), 4 )
09.0 0 00
babelfo) 00
2.00
Minibatch = (o, M(o)
00
2. Label observations. in DP using P-11
3. Train agent using ÊTD + Mfasr
Figure 2: Schematic view of GSF: the offline dataset D is used estimate POMDP-specific GVFs wrt some cumulant function c, whose quantiles are then used to label each observation in the dataset.
5.1
Quantifying future behavior with GVFs
An RL agent's future discounted behavior can be quantified not only by its value function, but other aux- iliary signals, for example, by its future observation occupancy measure, known successor representa- tion [62, 32]. The choice of the signal used during value iteration measures the properties the agent will exhibit in the future, such as accumulated returns, actions, observation visitation density. See Thm. 1 in the Appendix for the connection between successor features and interpolation error in our method.
Following the work of [20], we can broaden the class of value functions to any kind of cumulative discounted signal, as defined by a bounded cumulant function e: e:A RR. s.t. (e(o,a)] S Coux for
Cmax ee A ((o,a). 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.
Definition (Generalized value function) Let c be any bounded function over R4, let [0.1] and fu any policy. The generalized value function is defined, for any timestep t and O,
GH = Ep (cx)oo)
Since, in our case, can learn G for each distinct POMDP M, for the dataset DV, index the GVF using the POMDP index, i.e. G Learn (in practice, learning is parallelized).
Algorithm 1: LLearnGVF(,,Di.p.. (i,(0) .J,0,7): Offline estimation of GVF G
Input :Cumulant function c, dataset DV, POMDP label initial parameters
((0), target parameters latent state decoder f, iterations J, learning rate a, discount for j 1 do 0 // Sample transition from subset corresponding to POMDP C c(o.a)
of random crop(o):
-()) (GGu 1)2) u())-c-7Gao-n(())
Update target parameters 0 with of online parameters 0:
5

---

5.2 Measuring distances between GVFs of different POMDPs
Examining the difference between future behaviors of two observations quantifies the exact amount of expected behavior change between these two observations. Using the GVF framework, we could compute the distance between 01 E M1 and 0g E M2 by first estimating the latent state with f() using a latent state decoder J, and then using the following distance as measure of dissimilarity
( .o I| (Coo))
However, the distance between GVFs from two different POMDPs can have drastically different scales, Lee o0 - (02) 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 more heterogeneous signal, such as the extrinsic reward function. To avoid this problem, we suggest performing comparison based on order statistics.
S.me 2.mas thus making point-wise comparison meaningless. The 1-7
Namely, robust distance estimate between GVF signals across POMDPs be obtained by looking at the cumulative distribution function of G; denoted (g(g) = PP (o.) I g] for all or E O. G; is deterministic GVF with the of discontinuity points of measure 0, and as such F; can be understood through the induced state distribution P. (using continuous mapping theorem from [63]). It can be estimated from independent and identically distributed samples of DH
(g(
1
C₁ max C2,max
(4)
72
1 1-7
and its inverse, the empirical quantile function [64]
(p) - inf di 1-7
:pS
(5)
for p [0.1]. We use the empirical quantile function 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 defined as
= M, F, ( = (o) < F. (*++)) }} and its aggregated version as I [k)) - U (k))
Importantly, augment the dataset D with observation-speeife labels, which correspond to the index of the quantile bin into which the GVF G of an observation o M, falls into 4(0)= = max 1,))
These self-supervised labels are then used in multiclass InfoNCE loss [47], which is a variation of metric learning with respect to the quantile distance defined above [65, 66] and this forms the basis of our self-supervised learning objective.
5.3 Self-supervised learning of GSFs
After augmenting the offline 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.
First, the observation is encoded using a non non-linear encoder fo O + with parameters eb into a latent state representation = (o) The representation is then passed into two separate trunks: 1) a linear matrix 0, which recovers the state-action value function Qa (o.a) = =T= z, and 2) a non-linear projection network he Z + Z with parameters f), to obtain a new embedding, used for self supervupe learning. The projection he (z) then used in a multiclass InfoNCE loss [47, 66] where a linear classifier W e x K aims to correctly predict the observation labels (i.e. quantile bins A = 1, from (2), for temperature 7 0:
K
lasr (osr(O,,).. W) EEgDDa
1 = LooooooLossfftaaaWWWW>
(6)
'A special case of quantile binning occurs when n. in which case the auxiliary task is to predict the rank of the GVF associated to a given observation in the current minibatch.
This encoder is different from the used to evaluate the GVFs.

---

5.4 Full Algorithm
Given m training tasks, GSF first learns GVF estimates G. G2 by applying LearnGVF to task-specific data from the offline dataset D=, Each data point in DD is then labeled with the quantile into which its GVF falls. These labels are then used to jointly optimize Eq.6 and a control loss with respect the encoder parameters. To learn the value function Qo, using a linear combination of Q-lcarning [67, 16] and behavior-regularizatoo::
CQL [56], which is trained
(fo -a..a, ( ymax. Toap://eeellllttna..oooddooooaad 040000000000aa---.../////460e6d0
(7)
for À A20, 0, 0 target network parameters and LSE being the log-sum-exp operator ?. For domains with continuous actions, also decode the Bolztmann policy from Qo.
Fig. provides 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 discrete action space.
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. Below, we highlight which existing algorithms be recovered by GSF.
Cross-State Self-Constrain [CSSC, 10]: In CSSC, observations 01 are considered sim- ilar if they have identical future action sequences of length K under some fixed policy; total of |AA| distinct classes are possible. This approach can be approximated in our framework by picking c(or,at) ) Ia, (a), Va E A. The problem reduces to a |A| -[--t-way classification problem for observations of timestep which GSF approximates using K quantiles.
Policy similarity embedding [PSE, 11]: PSEs balance the distance between local optimal behaviors and long-term dependencies in the transitions, notably using dry. If we consider the space of Boltzmann policies Boltzmano with respect to a POMDP-speciffe value function Q, then choosing c(or,ae) = as)) in GSF will effectively compute the distance between unnormalized policies.
The choice of K induces bias- trade-off' How should the number of quantiles K (read labels) be set, and what is the effect smaller/ larger values of on the learned representations? Thm. highlights trade-off when choosing the number of quantile bins empirically.
Theorem 1 Let G1 G2 be generalized value functions with cumulants 01 C2 from respective POMDPs Mj. M2. K be the number of quantile bins, 71, n2 the number of sample transitions from each POMDP. Suppose that P(sup;=1.2... |c1 (oI...)))().))) 02 c2(02.19(0221)) (1->))))=))))3 Then, for any
k= 1,2,..,K and E >> 0 the following holds without loss of generality:
P sup mnrra)) (o,) Gylog)) 2m 44P sup sup
(1T (1 (1
8. (Vx)) K.)) (8)
The proof can be found the Appendix Sec. A.5. For POMDP M,. 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, (first term). The inter-POMDP error (third term) does not affect the bin assignment. Hence, choosing 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 n= 1. Setting K too small will group together unrelated observations, inducing high bias.
Limitations As is the case with all offline RL. methods, GSF limited by the compounding extrapolation error under low data coverage. 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 enviromment-depeeddent searching for an optimal (K,c) pair can be computationally expensive.
A copy of 0 updated solely using an exponential moving average (see Appendix). Phtpps: ://en. wikipedia, htt//ikkkLLoSSumix..
7

---

Experiments
Unlike for single task offline RL [26], most prior work on zero-shot generalization from offline data either come up with an ad hoc solution suiting their needs, e.g. [33], or assess performance on bench- marks that do not evaluate generalization across observation functions [e.g., 68]. To accelerate progress this field, we devised two benchmarks: offline Procgen (discrete actions) and offline Distracting Suite (continuous actions) two offline RL datasets to directly test for generalization of RL agents across observation functions". Moreover, for standard comparison, we provide generalization results on the classical online Procgen simulator, comparing PPO to PPO with GSF and PPO with PSE, respectively.
GVF training In the offinn setting, the training dataset is used to learn set of task-specific GVFs in parallel as follows. First, we project each observation into a latent representation == for (o); we then pass through non-linear network hø 2 + Rd xm where de is the dimensionality of the cumulant function's output. The output of hg is then split into disjoint chunks, which are in turn used in their respective temporal difference losses ÉTD in place of value functions. The procedure can be adapted an online setting via similar procedure, except that all GVF estimators are trained simultaneously and independently in separate simulators.
Offiine Procgen benchmark We evaluate the proposed approach on offline version of the Procgen benchmark [8], which widely used to evaluate zero-shot generalization across complex visual perturbations. Given a random seed, Procgen supports sampling procedurally generated level configurations for 16 games under various complexity modes: "casy", "hard" and "exploration" More details can be found in Appendix.
Offine 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 offline benchmarks which balances RL and BC objectives, CURL [35], CTRL [12],
DeepMDP [69] which learns metric closely related 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 Celf-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 optimal policy space.
2.0
2.0
COL
0.0
1..0
-2.0
3.0
DeepMDP
CSSC
PSE
CURL
GSF
Algorithm
Figure 3: Returns the offline Procgen benchmark [8] after 1M training steps. Boxplots are constructed over random seeds and 16 games; each method is normalized by the per-game median CQL performaccc.. White dots represent average of distribution.
Fig. 3 shows the performance of all methods over random seeds and all 16 games on the offline Procgen benchmark after 1 million gradient steps. Per-game average scores for all methods can be found Tab. 2 (Appendix). The scores are standardized per-game using the downstream task's (offline RL) performance, in this case implemented by CQL. It can be seen that GSF performs better than other offline RL and representation learning baselines.
CCode be found at htttp:://ggthiii httpa/uu...oommmmuuuee//eeeee cctm//mazzourrrgssppub....

---

Using different cumulant functions can lead to different label assignments and hence different similarity groups. Fig. 3 examines the performance GSF with respect to cumulants: 1) (s))) rewards s.t. GSF learns the policy's Q"-value, 2) 1o, (o), the successor representation" [62, 32] s.t. GSF learns the distribution induced by μ over Di [71] and 3) 1a, (a), action counts, s.t. GSF learns discounted policy. While rewards and successor feature cumulant choices leads to similar performance, using action-baed baseed distance leads to larger variance.
Offline Distracting Control Suite results Following the same procedure as for offline Procgen results, we first formed an offline dataset from the challenging Distracting Control Suite [9] by saving the replay buffer of Soft Actor-Critic [72] trained for 1M frames background perturbations. Next, we pre-trained Gf G22 with action and reward-based cumulants, which were then used in conjunction with CQL to learn single multi-task policy. Fig shows the online performance of GSF evaluated on 10 background perturbations not seen during training, normalized by per-enviromment median CQL scores.
environments with different
Median
KJM
Mean
Optimality Sap
6SF (action count)
0.04
0.06
0.58
O ORA O 1...0.000
0.375 0.350 0.375
0.71 0.77 0.73
0.74
Ingrovennent over COL improveneent, inprovereent)
Figure 4: Improvement of GSF over CQL with action and reward- -asedd similarity functions aggregated using performance metrics seeds and 4 environments of the Distracting Control Suite reported as suggested by [73] with 95% confidence intervals. We can see that using action counts results in higher mean, median and interquartile mean (IQM) statistics and lower optimality gap (i.e. fraction of scores falling under certain threshold), than when using rewards, or when comparing with the performance of CQL.
The results are consistent with findings 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.
Online Procgen results We additionally compare performance of GSF that of PPO and PPO with PSEs in the classical online Procgen benchmark [5]. Figure 8 shows test returns for 20M frames obtained on the entire distribution of easy levels while training on 200 easy levels for PPO, PPO+PSEs [1]] and our PPO+GSFs. PPO+GSF outperforms or matches both PPO and PPO+PSE on most environments, showing that GSF can be efficiently combined with both offline and online algorithms. See Appendix A.6.1 for detailed results.
7 Discussion
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. GSF relies computing the similarity between observation pairs with respect to any instantancous accumulated signal, which leads to improved empirical performance two newly introduced benchmarks, offline Procgen and offline Distracting Suite. Theoretical results suggest that GSF's hyperparameter choice depends on a bias- Saasinaiance trade-off.
References
[1] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, (altur,558877 518(7540))529-5333 2015.
2015.
[2] Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. International Conference on Learning Representations, 2020.
'In the continuous observation space, we learn a d' d-dimensional successor feature vector zy via TD and compute the quantiles over ||

---

[3] Mare G Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C Machado, Subhodeep Moitra, Sameera S Ponda, and Ziyu Wang. Autonomous navigation of stratospheric balloons using reinforcement learning. Nature, 588(7833))777-82, 2020.
[4] Annie S Chen, Suraj Nair, and Chelsea Finn. Learning generalizable robotic reward functions from" in-the-wild" human vidcos. arXiv preprint drXiv:2103.16817, 2021.
[5] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generalization in reinforcement learning. In International Conference Machine Learning. pages 1282-1289. PMLR, 2019.
[6] Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational over- fitting in reinforcement learning. International Conference on Learning Representations, 2019.
[7] Devon Hjelm, Bogdan Mazoure, Florian Golemo, Felipe Frujeri, Mihai Jalobeanu, and Andrey Kolobov. The sandbox environment for generalizable agent research (segar). arXiv preprint arXiv:2203. 10351, 2022.
[8] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference machine learning, pages 2048-2056 PMLR, 2020.
(9] Austin Stone, Oscar Ramirez, Kurt Konolige, and Rico Jonschkowski. The distracting control suite-a challenging benchmark for reinforcement learning from pixels. arXiv preprint WrXiv:2101002722 2021.
[10] Guan Ting Liu, Pu-Jen Cheng, and Guan Yu Lin. Cross-state self-constraint for feature generalization in deep reinforcement learning. 2020.
[11] Rishabh Agarwal, Marlos Machado, Pablo Samuel Castro, and Mare Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. arXiv preprint arXiv:2101.0 05265,, 2021.
[12] Bogdan Mazoure, Ahmed M Ahmed, Patrick MacAlpine, R Devon Hjelm, and Andrey Kolobov. Cross-trajectory representation learning for zero-shot generalization in rl. International Conference on Learning Representations, 2022.
[13] Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerehner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, pages 1490-1490. PMLR, 2017.
[14] Zichuan Lin, Derek Yang. Li Zhao, Tao Qin, Guangwen Yang, and Tie- Van Liu. RdZ: Reward decomposition with representation decomposition. Advances in Neural Information Processing Systems, 33, 2020.
[15] Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement learning. International Conference on Machine Learning. 2021.
[16] Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503-556, 2005.
[17] Sascha Lange. Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning, In Reinforcement learning, pages 45-73. Springer, 2012.
[18] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms, arXiv preprint arXiv: 2017.
[19] Kevin P Murphy. A survey of pomdp solution techniques. environment, 2:X3, 2000.
[20] Richard Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Systems- Volume 2, pages 761-768, 2011.
[21] Norm Ferns. Prakash Panangaden, and Doina Precup. Metrics for finite markov decision processes. In UAI, volume 4, pages 162-169, 2004.
[22] Lihong Li, Thomas Walsh, and Michael I. Littman. Towards unified theory of state abstraction fdps. ISAIM, 4:5, 2006.
10

---

[23] Pablo Samuel Castro.
Scalable methods for computing state similarity in deterministic markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10069-10076, 2020.
[24] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforeement learning without reconstruction. International Conference on Learning Representations, 2020.
[25] Dipendra Misra, Mikael Henaff, Akshay Krishnamurthy, and John Langford. Kinematic state abstraction and provably efficient tich-observation reinforcement learning. In International conference on machine learning, pages 6961-6971 PMLR. 2020.
[26] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
[27] Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Tom Le Paine, Sergio Gómez Colmenarejo, Konrad Zolna, Rishabh Agarwal, Josh Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI unplugged: Benchmarks for offine reinforcement learning. 2020.
[28) Peter Bartlett. The sample complexity of pattern classification with neural networks: the size of the weights is more important than the size of the network. IEEE transactions on Information Theerry,442) 44(2):525-553 1998.
[29] Eleni Triantafillou, Tyler Zhu, Vincent Dumoulin, Pascal Lamblin, Utku Evci, Kelvin Xu, Ross Goroshin, Carles Gelada, Kevin Swersky, Pierre-Antoine Manzagol, et al. Meta-dataset: A dataset of datasets for learning to learn from few examples. International Conference on Learning Representations, 2019.
[30] Guillermo Valle-Pérez and Ard A Louis. Generalization bounds for deep learning, arXiv preprint arXiv:2012.0 04 15, 2020.
[31] Liu, William Hamilton, Guodong Long, Jing Jiang, and Hugo Larochelle. A universal represennn tation transformer layer for few-shot image classification. arXiv preprint arXiv:r006. 1102, 2020.
[32] André Barreto, Will Dabney, Rémi Munos, Jonathan Hunt, Tom Schaul, Hado Van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. arXiv preprint 160iv:1606.030 05312. 2016.
[33] Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models facilitate zero-shot dynamics generalization from a single offline environment. International Conference Machine Learning. 2021.
[34] Alekh Agarwal, Mikael Henaff, Sham Kakade, and Wen Sun. Pc-pg: Policy cover directed exploration for provable policy gradient learning. Neural Information Processing Systems, 2020.
[35] Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised repre- sentations for reinforcement learning. International Conference on Machine Learning, 2020.
[36] Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. Neural Information Processing Systems, 2020.
[37] Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879 PMLR, 2021.
[38] Samarth Sinha and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning. arXiv preprint S.ti::2003003222 2021.
[39] Adith Swaminathan and Thorsten Joachims. Batch learning from logged bandit feedback through counterfactual risk minimization. The Journal of Machine Learning Research, 16(1):1731-1755.
2015.
[40] Justin Boyan and Andrew W Moore. Generalization in reinforcement learning: Safely approximat- ing the value function. Advances in neural information processing systems, pages 369-376, 1995.
[41] Yuanhao Wang. Ruosong Wang. and Sham M Kakade. An exponential lower bound for linearly- realizable mdps with constant suboptimality gap. arXiv preprint arXiv: /2690, 2021.
[42] Ofir Nachum and Mengjiao Yang. Provable representation learning for imitation with contrastive fourier features. Neural Information P Processing Systems, 2021.
[43] Ruosong Wang, Yifan Wu, Ruslan Salakhutdinov, and Sham M Kakade. Instabilities of offline rl with pre-trained neural representation. arXiv preprint arXiv: 103.0494.74 2021.
11

---

[44] Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. arXiv preprint arXiv:2103.07990 2021.
[45] Nicholas K Jong and Peter Stone. State abstraction discovery from irrelevant state variables. In IJCA,, volume 8, pages 752-757. Citeseer, 2005.
[46] Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Phil Bachman, Adam
Trischler, and Yoshua Bengio. Learning deep repressentations by mutual information estimation and maximization. arXiv preprint arXiv:1888.066770 2018.
[47] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 03748, 2018.
[48] Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. arXiv preprint arXiv:2006. 09882, 2020.
[49] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07773 2020.
[50] Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, Devon Hjelm, Philip Bachman, and Aaron Courville. Pretraining representations for data-efficient reinforcement learning. arXiv preprint vrXiv:2106.04700 2021.
[51] Ching-An Cheng, Andrey Kolobov, and Adith Swaminathan. Heuristic-guides reinforcement learning. arXiv preprint arXiv:2 027577 2021.
[52] Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinffrrcement learning via distance metric learning and behavior regularization. arXiv preprint arXiv:2010.0// 2020.
[53] Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Christensen, and Hao Su. Multi-task batch reinforcement learning with metric learning. Advances in Neural Information Processing Systems, 33:6197-6210, 2020.
[54] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062. PMLR, 2019.
[55] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. Provably good batch reinforcement learning without great exploration. NeurlPS, 2020.
[56] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv: 2020.
[57] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 133i 2019.
[58] Ilya Kostrikov, Jonathan Tompson. Rob Fergus, and Ofir Nachum. Offline reinforcement learning with fisher divergence critic regularization. arXiv preprint arXiv:2103 080500 2021.
[59] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. arXiv preprint arXiv:2005. 059551 2020.
[60] Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. arXiv preprint 1910.01 2019.
[61] Remi Tachet, Philip Bachman, and Harm van Seijen. Learning invariances for policy generalization. arXiv preprint 1=Xiv:1809002555 2018.
[62] Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613-624, 1993.
[63] Henry Mann and Abraham Wald. On stochastic limit and order relationships. The Annals of Mathematical Statistics. 14(3):217-226, 1943,
[64] A. W. van der Vaart. Asymptotic statistics. cambridge series in statistical and probabilistic mathematics, 1998.
[65] Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Neural Information Processing Systems, 2020.
12

---

[66] Jiaming Song and Stefano Ermon. Multi-label contrastive predictive coding. Neural Information Processing Systems, 2020.
[67]
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3-4):279 1992.
[68]
Tianhe Yu. Deirdre Quillen, Zhampeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning. pages 1094-1100. PMLR, 2020.
[69] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning. pages 2170-2179. PMLR, 2019.
[70] Junhyuk Oh, Satinder Singh, and Honglak Lee. arXiv:1707.03497 2017.
Value prediction network. arXiv preprint
[71] Marlos C Machado, Mare G Bellemare, and Michael Bowling, Count-basee exploration with the successor representation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5125-5133, 2020.
[72] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[73] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 34, 2021.
[74] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. International Conference on Learning Representations, 2021.
[75] Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings the AAAI Conference on Artificial Intelligence, volume 33, pages 3796-38666 -8003, 2019.
[76] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, al. Impala: Scalable distributed deep-rl with importance weighted learner architectures. In International Conference on Machine Learning, pages 1407-1416. PMLR, 2018
[77] Sergio Guadarrama, Anoop Korattikara, Oscar Ramirez, Pablo Castro, Ethan Holly, Sam Fishman, Ke Wang, Ekaterina Gonina, Neal Wu, Efi Kokiopoulou, Luciano Sbaiz, Jamie Smith, Gábor Bartók, Jesse Berent, Chris Harris, Vincent Vanhoucke, and Eugene Brevdo. TF-Agents: library for reinforcement learning in tensorflow. https: //g //ttttttt../..e// httg//ghuu..cmmteesrrrreeee com/tensorffoww aagets, 2018. [Online; accessed 4- 2-0tt.ee--2021.
[78] Aryeh Dvoretzky, Jack Kiefer, and Jacob Wolfowitz. Asymptotic minimax character of the sample distribution function and of the classical multinomial estimator. The Annals of Mathematical Statistics, pages 642-669, 1956.
[79] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729-9738 2020.
[80] Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint arXiv: / 8802. 2018.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See end of Section 5
(c) Did you discuss any potential negative societal impacts of your work? [Yes] We include brief discussion in Appendix Section A.
13

---

(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix Section A.5
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplcmental material or as a URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [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)? [Yes] See Appendix Section A.4
4. 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? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as URL? [Yes] Due to the large size the offline datasets, we include the code which can re-generate the data from scratch.
(d) Did you discuss whether and how consent obtained from people whose data you's 'e using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information offensive content? [N/A]
5. 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? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "qx4SG2RFWa9", "reviewer": "Reviewer_EVg5", "summary": "This paper proposes a framework called Generalized Similarity Functions (GSF), which estimates the similarity between observations by the various future behaviors. The previous works tried to calculate the similarity based on the future action sequences, state distribution, value, etc; but the adequateness of those methods depends on the problem, and no method is optimal in all circumstances. To this end, GSF learns from some arbitrary c(s,a) function and results in a value functions of those functions. Authors argue that the proposed framework can include the previous methods by using appropriate c(s,a). After learning GSF, the paper augments the data by additional labels, i.e. quantile of GSF, and learns the latent space by minimizing InfoNCE loss.", "questions": "- It seems $c(s,a)$ is proposed to be a vector-valued function to recover some cases, e.g. successor features, CSSC. $G$ is a cumulative sum of $c$, so it will be vector-valued when $c$ is vector-valued. However, below definition 1, it is explained that $G_i$ is learned for each distinct POMDP. Is it only the case when $G$ is learned based on reward? Or is $G$ matrix-valued when $c$ is vector-valued? Either way, the description does not seem very clear.\n\n- The aggregated version $I(k)=\\cup_{i=1}^m I_i(k)$ does not seem to be used. The definition of $l_i(o)=\\max_k \\mathbb{1}_{o\\in I_i(k)}$ is also strange; is it $\\arg \\max$? Based on a number of quantiles based on different dimension of $G$, how exactly is $l_i$ computed?\n\n- It is hard to understand how CSSC is recovered (mainly due to the lack of understanding on $l_i$). Can it be exactly recovered when we have large enough $K$? Can you also elaborate more on how to recover PSE precisely?\n\n- In the experiment section, the proposed framework shows much performance improvement on the Procgen benchmark (almost double the performance of CQL), whereas other algorithms perform much worse than CQL, or are barely comparable. Is GSF the only algorithm that is based on CQL? What happens if we use other representation learning algorithms with CQL? (or is it already done in that way?)\n\n- On the other hand, in distracting control suite, GSF only shows a very small improvement (~=4%). What is the main reason for being this bad compared to the Procgen benchmark?\n\n- What happens if we use all these different similarity metrics at the same time? Can the algorithm find out the best one to use? Can we know what similarity metric to use in advance?", "limitations": "The authors adequately addressed the limitations and potential negative social impact of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "- This paper proposes a framework that can capture any signal and use it to train the representation. It is a generalized framework that can recover a number of previous frameworks. In that sense, it is somewhat original but not something that is completely new.\n\n- The paper seems to have OK-ish technical quality, but some stuff is not clearly written and is hard to follow (e.g. definitions of $c$, $G$, $F$). The experiment session is also somewhat not easy to understand (what is the optimality gap in figure 4?)\n\n- This paper will have some impact on the researchers in this field, due to the good performance and the proposal of the benchmarks.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "LDMUJ7AEC9", "reviewer": "Reviewer_dP54", "summary": "The paper proposes a novel representation learning method Generalized Similarity Functions (GSF) to improve the generalization ability of offline RL. GSF clusters latent states with similar future behaviors quantified by the order statistics of Generalized Value Function (GVF). Empirical results on the offline Procgen benchmark and the offline Distracting Control Suite show that GSF successfully enhances the performance on unseen tasks.", "questions": "* Q1: The authors may need to spend more effort on the related works such as [1][2].\n* Q2: The authors are suggested to add a TSNE plot of the latent states of different training and testing tasks.\n* Q3: A comparison of the computation cost is needed.\n\n\n**Typos:**\n1. line 185: you may need to delete \"the\" in \"by its the value function\". \n2. line 233: two repeated words \"used\".\n", "limitations": "The authors state that the limitation of the proposed method is the search of hyper-parameters.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n* The problem of zero-shot generalization in offline RL is significant, and the idea of GSF seems novel.\n* The choice of hyper-parameters has some theoretical analysis.\n* The paper conducted experiments in both discrete and continuous benchmarks, and the experimental results in both offline and online settings demonstrate the effectiveness of the proposed method.\n\n**Weaknesses:**\n* Some related works are missing. [1] and [2] both introduce representation clustering methods for offline multi-task/meta RL and can generalize to unseen tasks.\n* There is no visualization of the learned representations of different training and testing tasks, which could help readers better understand the effectiveness of GSF.\n* The paper does not discuss the computation cost of GSF.\n\n[1]  Li J, Vuong Q, Liu S, et al. Multi-task batch reinforcement learning with metric learning[J]. NeurIPS 2020.\n[2] Li L, Yang R, Luo D. Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. ICLR 2021.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "5v6LEiCXry", "reviewer": "Reviewer_XVwk", "summary": "This paper proposes a behavior similarity-based representation learning method that improves the generalization performance in an offline RL setting. Two key ideas are presented here: \n- (1) the proposed behavior similarity is based on the distance of GVF predictions between two states.\n- (2) to get a robust estimation of the distance between two GVF signals, it is formulated as a cumulative distribution function under assumptions, which is further converted to K quantile bins and can be learned by self-supervised learning methods in multi-label classification tasks, as defined in Eq. (6).\n\nProcgen and DM Control suite experiments were conducted to demonstrate the effectiveness of the proposed method.", "questions": "- (1) Cumulant specification question as I raised in the Weakness part.\n- (2) The performance gain issue in Fig. 8, as I raised in the Weakness part.\n- (3) Besides the cumulant function definition, we know that the GVF function is also dependent on a target policy, could the authors explain to me how the target policy is defined?\n    - If it is different from the behavior policy that collects samples to compose the offline dataset, can the authors explain to me further how to tackle the off-policy learning issue here when you use the TD loss, as shown in Fig. 2?\n\n- (4) I can see the full loss function contains both TD loss and GSF loss in Fig. 2, but there is no section, nor an algorithm that gives a complete description of how exactly the learning process happens. I can find the TD loss in Algorithm 1, line 6, but I am still wondering what is the complete algorithm solution as you summarized in Fig. 2, \"Train Agent using L_TD + lambda * L_GSF\". Could the authors help me to understand the complete training process?", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n---------------------------------\n- The problem setting of studying generalization using an offline dataset, unlike RAD, CURL, and bi-simulation Metric has significant importance in real-world problems.\n- The authors demonstrate sufficient knowledge in generalizable RL literature. I really enjoy reading the introductory part.\n- The proposed behavior similarity based on GVF predictions makes sense to me that should demonstrate better generalization performance.\n- The conversion of robust GVF distance estimation problem to a multi-label (as K bins) self-supervised learning task, is quite nice.\n\nWeaknesses\n------------------------------\n- The generality of GVF functions. I am concerned about how to specify the cumulant of the GVF function, as this is the key part that determines the generalization performance (as the authors mentioned in lines 314 ~ 319). After checking the paper, specifying the cumulants seem to still require pre-defined using some kind of human heuristic knowledge. Will the authors explain how to resolve the concern of generality of hand-specified cumulants?\n- The performance gain (Fig. 8 in Appendix, page 21) seems not significant. When comparing the proposed method to a vanilla PPO in the Procgen benchmark seems very small, and most tasks do not see a significant performance improvement, except for Caveflyer, Leaper, and plunder. Will the authors help to explain any reasons for that?\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept