Improving zero-shot generalization in offline reinforcement learning using generalized similarity functions
Authors: Bogdan Mazoure
Ilya Kostrikov
Ofir Nachum
Jonathan Tompson
Venue: ICLR 2022 Submitted (2022)
Keywords: reinforcement learning
representation learning
self-supervised learning
offline RL
generalized value function
generalization

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 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 also improving zero-shot generalization performance on a complex offline RL benchmark, offline Procgen.

Under review as a conference paper at ICLR 2022
IMPROVING ZERO-SHOT GENERALIZATION IN OF- FLINE REINFORCEMENT LEARNING USING GENERAL- IZED SIMILARITY FUNCTIONS
Anonymous authors Paper under double-biee blind review
ABSTRACT
Reinforcement learning (RL) agents are widely used for solving complex sequen- tial decision making tasks, but still exhibit difficulty in generalizing to scenarios not during training. While prior online approaches demonstrated that using additional signals beyond the reward function can lead to better generalization ca- pabilities 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 perfor- mance 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-motival
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 on a complex offline RL benchmark, offline Proegen.
INTRODUCTION
Reinforcement learning (RL) is a powerful framework for solving complex tasks that require a quence of decisions. The RL paradigm has allowed for major breakthroughs in various fields, e.g. outperforming humans on video games (Mnih al., 2015; Schwarzer et al., 2020). controlling strato- spheric balloons (Bellemare et al., 2020) and learning reward functions from robot manipulation videos (Chen et al., 2021). More recently, RL agents have been tested in a generalization setting, iee in which training involves a finite number of related tasks sampled from some distribution, with a potentially distinct sampling distribution during test time (Cobbe al., 2019; Song et al., 2019). The main issue for designing generalizable agents is the lack of on-policy data from tasks not seen during training: is impossible to enumerate all variations of a real-world environment during training and hence the agent must extrapolate from a (limited) training task collection onto a broader set of problems. Since the learning agent is given no training data from test-time tasks, this problem is referred of zero-shot generalization where the difference between tasks is predominantly due to perceptually distinct observations. An example this setting any environment with distractor features (Cobbe et al., 2020; Stone et al., 2021), i.c. features with no dependence on the reward signal nor the agent's 's decisions. This generalization setting has recently received much attention (Liu al., 2020a; Agar- wal et al., 2021; Mazoure et al., 2021), due to its particular relevance to real-world scenarios, for example deploying single autonomous driving agent at day or at night.
zero-shot generalization. In our work,
interested in the problem
Generalization capabilities of an agent can be analyzed through the prism of representation learning, under which the agent's current belief about rich and high-dimensional environment are summa- rized in low-dimensional entity, called representation. Recent work in online RL has shown that learning state representations with specific properties such as disentanglement (Higgins et al., 2017) or linear separability (Lin et al., 2020) can improve zero-shot generalization performance. Achiev- ing this with limited data (i.e. offline RL) is challenging, since the representation will have a large estimation error regions of low data coverage. common solution to mitigate this task task-specific overfitting and extracting the most information out of the data consists in introducing auxiliary learn- ing signals other than instantaneous reward (Raileanu and Fergus, 2021). ). As we show later in the

---

Under review as a conference paper at ICLR 2022
paper, many such signals already contained in the dataset can be used to further improve generaliza- tion performance. For instance, the generalization performance of PPO on Proegen remains limited even when training 200M frames, while anneralizaaion-orientee agents (Raileanu and Fergus, 2021; Mazoure et al., 2021) ) can outperform it by leveraging additional auxiliary signals. However, 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 an offline, fixed dataset of experience (Ernst et al., 2005; Lange et al., 2012). A natural limitation for generalization from offline data is that policy improvement is dependent on dataset quality. Specifically, high- dimensional problems such as control from pixels require large amounts of training experience: a standard training of PPO (Schulman et al., 2017) for 25 million frames on Procgen (Cobbe et al., 2020) generates more than 300 Gb of data, an impractical amount of data to share for offline RL re- search. Improving zero-shot generalization performance from an offline dataset of high-dimensional observations is therefore hard problem due to limitations on dataset size and quality.
In this work, we are interested in improving zero-shot generalization across family of Partially- Observable Markov decision processes (MDPs, Puterman, 1990) in 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 gen- eralize across perceptually different POMDPs without adaptation, observations with similar future behavior should be assigned to close representations. We use the generalized value function (GVF) framework (Sutton al., 201 1) to capture future behavior with respect to arbitrary instantaneous signal (called cumulant) for a given state. The choice of cumulant then determines the nature of the behavioral similarity that is encouraged for generalization. For example, using reward as the signal gives rise to of reward-aware behavioral similarity such as bisimulation (Ferns et al., 2004; Li et al., 2006: Castro, 2020; Zhang et al., 2020); using future state-action counts encourages reward-free behavioral similarity (Misra et al., 2020; Liu et al., 2020a; Agarwal et al., 2021; Mazoure et al., 202 1).
Our main contributions are as follows:
1. We propose Generalized Similarity Functions (GSF), a novel self-supervised learning al- gorithm for reinforcement learning, that aggregates latent representations by the future be- havior (or generalized value function) under their respective observations.
2. We devise new benchmark constructed to test zero-shot generalization of offline RL algo- rithms: offline Procgen. It consists of 5M transitions from 200 related levels of 16 distinct games.
3. We evaluate performance of GSF and other baseline methods offline Procgen, and show that GSF outperforms both previous state-of-the-arr offline RL and representation learning baselines the entire distribution of levels.
4. We analyze the theoretical properties of GSF and describe the impact of hyperparameters and cumulant functions on empirical behavior.
RELATED WORKS
Generalization in reinforcement learning Generalizing a model's predictions across variety of unseen, high-dimensional inputs has been extensively studied in the static supervised learning set- ting (Bartlett, 1998; Triantafillou et al., 2019; Valle-Pérez and Louis, 2020; Liu al., 2020b). Gener- alization in RL has received a lot of attention: extrapolation to unseen rewards (Barreto et al., 2016; Misra et al., 2020), observations (Zhang al., 2020; Raileanu and Fergus, 2021: Liu et al., 2020a; Agarwal et al., 2021; Mazoure et al., 2021) and transition dynamics (Ball et al., 202 1). Each gen- eralization scenario is best solved by their respective set of methods: sufficient exploration (Misra et al., 2020; Agarwal et al., 2020), auxiliary learning signals (Srinivas et al., 2020; Mazoure et al., 2020; Stooke et al., 2021) or data augmentation (Ball et al., 2021; Sinha and Garg, 2021). Data aug- mentation is a promising technique, but typically relies handcrafted 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
2

---

Under review as a conference paper at ICLR 2022
across tasks) and improve generalization performance. Theoretical generalization guarantees have only been provided so far for limited scenarios, mostly for bandits (Swaminathan and Joachims, 2015), linear MDPs (Boyan and Moore, 1995; Wang et al., 2021b; Nachum and Yang, 2021) and across reward functions (Castro and Precup, 2010; Barreto et al., 2016; Wang et al., 2021a; Touati and Ollivier, 2021).
Representation learning For simple POMDPs, near-optimal policies can be found by optimiz- ing 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 previ- ously (Jong and Stone, 2005; Li et al., 2006), but saw real breakthrought 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 certain datasets (Hjelm et al., 2018; Oord et al., 2018; Caron et al., 2020; Grill et al., 2020). Representation learning, and specifically self-supervised learning, has also been used to achieve state-of-the-art generaliza- tion and sample efficiency results in RL on challenging control problems such as data efficient Atari (Schwarzer et al., 2020; 2021), I). DeepMind Control (Agarwal et al., 2021) and Procgen (Ma- zoure et al., 2020; Stooke et al., 2021; Raileanu and Fergus, 2021; Mazoure et al., 2021). Nottewryhh instances of paaretically-motivttee representation learning methods for RL include heuristic-guidee guided learning (Sun et al., 2018; Cheng et al., 1, and random Fourier features (Nachum and Yang, 2021).
Offiine reinforcement learning When learning from a static dataset, agents should balance inter- polation 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 (Fujimoto et al., 2019), MBS (Liu et al., 2020c), and CQL (Kumar et al., 2020) rely on behavior regularization loss (Wu et al., 2019) as tool to control the extrapolation error. Some methods, such as F-BRC (Kostrikov et al., 2021) are defined only for continuous action spaces while others, such as MOReL (Kidambi et al., 2020) estimate a pessimistic transition model. The major issue with current offline RL al- gorithms such as CQL is that they are perhaps overly pessimistic for generalization purposes, i.e. CQL and MBS ensure that the policy improvement well-supported by the batch of data. As we will show in our empirical comparisons, using overly conservative policy updates can prevent the representation from fully leveraging the information of the training dataset.
3 PROBLEM SETTING
3.1 PARTIALLY-ORSERVABLE MARKOV DECISION PROCESSES
A (infinite-horizon) partially-observabee Markov decision process (POMDP, Murphy, 2000) M is defined by the tuple M (S, po: A. ps. O, po, r, 7), where S is state space. Pa = P[ is the starting state distribution, .A is an action space, ps Pf|||s, a.] : S x A ASS) is transition function, O is an observation space, Po = |[ S + (()) is an observation function, x A - [rimin mmax] 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 on (- At every timestep t = 1, the agent, parameterized by a policy т. O + ((A), samples an action a, ~ T T-loa)) The environment transitions into next state 8141 ps(- a.) and emits a reward r, = r(se, az) along with next observation 0++1 ~ (-(
The goal of an RL agent is to maximize the cumulative rewards Et=p 7'r obtained over the entire episode. Value-basee off-policy RL algorithms achieve this by estimating the state-action value function under target policy
Q" (5 az) = Epp
(r( (ctk)sss,a), ar), Vs, E S, E
(I)
A=
where P. denotes the joint distribution of (8t4k) d obtained by executing in the environ- ment.
(AY') denotes the entire of distributions over the space X.
3

---

Under review as a conference paper at ICLR 2022
An important distinction from online RL is that, instead of sampling access to the environment, assume access to a historical dataset D collected by logging experience the policy, μ, in the form {oi,s, di,t, Fi,e i=,,1= i=N.t=T where, for practical purposes, the episode is truncated at timesteps. Furthermore, we assume that the agent can only be trained on a limited collection of POMDPs Mimin = {M; [M. and its performance is evaluated on the set of test POMDPs Miest- We assume that both Musin and Must were sampled from common task distribution and that every POMDP M2 E M = -Mainn U MMst shares the same transition dynamics and reward function with M but has different observation function Pı..- Importantly, since we perform control from pixels, we are in the POMDP setting (see Yarats et al., 2019) and therefore emphasize the difference between observations or and corresponding states 8, throughout the paper.
3.2 REPRESENTATION LEARNING
Previous works in the RL literature have studied the use of auxiliary signals improve generaliza- tion performance. Among others, Liu et al. (2020a); Agarwal et al. (2021) define the similarity of two observations depend on the distance between action sequences rolled out from that observa- tion under their respective optimal policies. They achieve this by finding latent space 2 S S in which the distance (z( z') for 2, 2' E Z equivalent to distance between true latent states (s(8,8) s') for all s' E S; the aforementionned works learn Z by optimizing action-based similari- ties between observations.
In practice, latent space is decoded from observation usingg latent state decoder f O + 2 from observation O₁. Through the paper. we assume that all value functions have a linear form in the latent decoded state, i.c. (2(o,a) a) 0_ (o) 0- 2g, which agrees with practical implementation of all algorithms. Within this model family, the ability of an RL agent to correctly decode latent states from unseen observations directly affects its policy, and therefore, its generalization capabilities. In the next section, we 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.
4 MOTIVATING EXAMPLE
Level 1
Level
Actions
8
11
11
13
8
0
13
Figure 1: Two levels of the Climber game from the Procgen benchmark (Cobbe et al., 2020) with near-identical true latent states and near-identical value functions but drastically different action sequences.
Multiple recently proposed self-supervised objectives (Liu et al., 2020a; Agarwal al., 2021) con- jecture that observations 0₃ E M2, 02 E M2 that emit similar future action sequences under optimal policies x mī should be decoded into nearby latent states 21, 22. While this heuristic can correctly group observations with respect to their true latent state in simple action spaces, it fails to identify similar pairs of trajectories POMDPs with multiple optimal policies. For instance, two trajectories might visit an identical set of latent states, but have drastically different actions.
Fig. shows one such example: two levels of the Climber game have near-identical true latent state (see Appendix) and value function (average normalized mcan-squared eπror of 0.0398 across episode), while having very different action sequences from a same PPO policy (average total vari- ation 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. 4), iee. various levels can have multiple dras- tically 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. In POMDP with
4

---

Under review as a conference paper at ICLR 2022
different structure, grouping representations by action sequences can be optimal. So how do unify these similarity metrics under a single framework?
In the next section, this insight to design general way of improve representation learning through self-supervised learning of discounted future behavior.
METHOD
We propose measure generalized notion of future behavior similarity using generalized value functions, as defined by the corresponding cumulant function. The choice of cumulant determines which components of the future trajectory are most relevant for generalization.
5.1
QUANTIFYING FUTURE BEHAVIOR WITH GVFs
An RL agent's 's future discounted behavior can be quantified not only by its the value function, but other auxiliary signals, for example, by its observation occupancy measure, known as successor fea- tures (Dayan, 1993; Barreto et al., 2016). The choice of the examined signal quantifies the properties the agent will exhibit the future, such as accumulated returns, or observation visitation density. See Thm. 2 in the Appendix for the connection between successor features and interpolation error our method.
Following the work of Sutton et al. (2011), we can broaden the class of value functions to any kind of cumulative discounted signal, as defined by a bounded cumulant function c : x A + Rd s.t. le(o, a)] S Cmax for Cmax aupo,aexx.. c(o, a). While typically cumulants scalar-valued functions (e.g. reward), we also make use of the vector-valued case for learning the successor features (Barreto et al., 2016), in which case the norm of c(o, a) is bounded.
Definition (Generalized value function) Let e be any bounded function over Id,, let y E [] and μ any policy. The generalized value function defined as
G" G"(o,) = Eggl
(2)
1=1
for any timestep 2 1 and or E O.
Since, in our case, we can learn G" for each distinct POMDP M, for the dataset D"r we index the GVF using the POMDP index, i.e. G = LearnGVF(c, D",i) (in practice, the learning is parallelized).
Algorithm 1: LearnGvF((, D", D,, (()) a,,)) 7): Offline estimation of GVF ČH
Input : CCmuulat function dataset DM, POMDP label i, initial parameters ((0), target
parameters 9, latent state decoder f. iterations J, learning rate discount
for j = L.... do
a. a D[i]; Sample transition from POMDP i
c t a):
o - random crop(o):
(l3) + عاو-ا a pts in (Ggts- ()) -c cc -G (,())))
Update target parameters á with of online parameters 0:
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 on Mi and 02 E M2 by first estimating the latent state with = f(o) using (learedd) latent state decoder J. and then evaluating the distance
dy & oz') (f(o)) - (f(2)) i 1,2,
(3)

---

Under review as a conference paper at ICLR 2022
measure of dissimilarity that can then be used in contrastive loss.
However, the distance between GVFs from two different POMDPs can have drastically different scales: (o1) - (02) I S.nn + making point-wise comparison meaningless. The issue 1-y
less acute for cumulants which induce a unnormalized density estimate (e.g. indicator functions for successor representation), and more problematic when the cumulant incorporates the extrinsic reward function. To avoid this problem, we suggest performing a comparison based on order statis-
tics.
robust distance estimate between GVF signals across POMDPs can be obtained by looking at
the cumulative distribution function of G; denoted (g(g) (o:) S g] for all or O. G deterministic GVF with the set of discontinuity points of mcasure O, and as such F& can be understood through the induced state distribution P. (using continuous mapping theorem from Mann and Wald (1943)). It can be estimated from n independent and identically distributed samples of D
- = Σ 1, G. LearnGVF(e, D", 1), 9
Ci,max
max
(4)
1=1
and its inverse, the empirical quantile function (van der Vaart, 1998)
(p) in E [-
G.max
Ci,max
P R(g)), p [0.1]
(5)
We use the empirical quantile function to partition the range of all GVFs into K quantile bins, i.e. disjoint sets with identical size where the set corresponding to quantile defined as (k)) = (o
M, ( (o) S (**))*)))
) and it's aggregated version as T(k) = (k)..
Importantly, we augment the dataset DD with observation-specife labels, which correspond to the index of the quantile bin into which the GVF of observation E M,, falls into:
(o) = max (.))(()
(6)
k
These self- supervieed labels are then used in multiclass InfoNCE loss (Oord et al., 2018), which variation of metric learning with respect to the quantile distance defined above (Khosla al., 2020; Song and Ermon, 2020).
5.3 SELF-SUPERVISED LEARNING OF GSFS
After augmenting the offline dataset with observation labels, we use simple self-supervised learn- ing procedure to minimize distance in the latent representation space between observations with identical labels.
First, the observation is encoded using non-linear encoder fo O - Z with parameters vb into a latent state representation f (o) The representation is then passed into two separate trunks: 1) linear matrix 0, which recovers the state-action value function (o a) = 0g z, and 2) a non-linear projection network he : Z - Z with parameters 9, to obtain a new embedding, used for contrastive learning.
The projection h(z) is then used used a multiclass InfoNCE loss (Oord et al., 2018; Song and Ermon, 2020) where a linear classifier W E x!ž|* aims to correctly predict the observation labels (i.e. quantile bins & 1,2,.., K) from h(2):
K
(NCE(Op, b., W) = -E
where 7 >> is temperature parameter.
Our empirical findings suggest that this version of the loss is more stable than other multi-class contrastive losses (see Appendix 7.3).
This encoder different from the used evaluate the GVFs.

---

Under review as a conference paper at ICLR 2022
Mj
Control
a
0
Q(z
LTD
Representation learning
1. Estimate G Gj from D
h(z)
ENCE (xCE(1(2)) (h(z),
h
labal())
h
Minibatch B (o. ((o))
2. Label obscrvations in minibatch B using P-1
3. Train agent
Figure 2:
Schematic view of GSF : the offline dataset D is used to estimate POMDP-specific GVFs wrt some cumulant function c, whose quantiles are then used label each observation in the dataset. These labels are then used in multi-class contrastive learning procedure along with offline RL learning.
5.4 ALGORITHM
Our method relies the approximation oracle LearnGVE, to produce GVF estimates, later used in the contrastive learning phase.
Since concerned with the offline RL setting, add auxiliary loss top of Conservative Q-learning (CQL, Kumar et al., 2020), a strong baseline. CQL trained using a linear combination of Q-learning (Watkins and Dayan, 1992; Ernst al., 2005) and behavior-regullaizztionn
lcor. (0) Eo (( max (.. d'EA
for X 2 0, 0 target network parameters and L.S.F being the log-sum-exp operator
Algorithm 2: GSF Offline RL with future behavior observation matching
Input : Dataset D ~ fs, initialized Q-function Qo with encoder Jo, and action weights 0z. per-POMDP set of GVFs G = (G lie, state projection network hg, epoch number J, number of POMDPs moving average parameter 3
number of quantiles K, temperature parameter т, exponential
for epoch j = j 1,2, do for minibatch B ~ do
Data augmentation en observation
+ random crop(o) for all e B;
2 (o) for all o B;
/* Update CQL agent
Update 0z, va using (o.
/* Compute G quantiiiiii
for POMDP M, = 1, m do
Estimate P:11 G/ from B
for observation o E B M, do
I(o) + f ( ( S (o)
/* Update encoder and projection network
Update 0n. 16 W using Vel.v. w v. W) computed with 2, (()()) and ;
10
11
Update CQL agent's target network with of online parameters vb,
Alg. 2 summarizes the learning procedure for GSF as implemented on top of a CQL agent for discrete action space. In experiments, all baselines random crops data augmentation.
A copy of 0 updated solely using an exponential moving average (sce Appendix). wrg/wiiggiioooooo https:: //en, wiki wiw.p oordia.or//ie wwi wik http://wk//l/g/s//ees// Kxg
7

---

Under review as a conference paper at ICLR 2022
Connection to existing methods Our framework is able recover objectives similar to those of prior works by carefully designing the cumulant function.
Cross-State Self-Constrain (CSSC, Liu et al., 2020a): In CSSC, observations or 02 are considered similar they have identical future action sequences of length under some fixed policy; total of |A| distinct classes are possible. This approach can be approximated in our framework by picking c(or, at) = lo (a), Va E A. The problem reduces to wwy classification problem for observations of timestep t, which GSF approximates using quantiles.
Policy similarity embedding (PSE, Agarwal et al., 2021): PSEs balance the distance between local optimal behaviors and long-term dependencies in the transitions, notably using drv. If we consider the space of Boltzmann policies Boltzman with respect to an POMDP-specific value function Q, then choosing c(or, at) r(st, 0z) in GSF will effec- tively compute the distance between unnormalized policies.
5.5 CHOICE OF NUMBER OF QUANTILES
How should the number of quantiles K be set, and what is the effect of smaller/ larger values of K on the observation distance? Thm. highlights a trade-off when choosing the number of quantiles bins empirically.
Theorem Let G1, G2 be generalized value functions with cumulants c1, c, from respective POMDPs M1, M2. K be the number of quantile bins, Th1 02 the number of sample transitions from each POMDP. Suppose that IP (1ss. les (o1,,, (((11,)) - C2 cz(02.t p(02.)) > (1-7)=/1) S 6.
Then, for any k = 1, K and & > 0 the following holds without loss of generality:
P
sup (G1( - G2((22) 38 (2
S 2e
22 22
+ p(ni, K,E) 8
(9)
where
p(n, K2 =
sup
AT (K+2/K)) (4/K) >E]
(10)
k= 1=
The proof can be found in the Appendix Sec. 7.2. For POMDP M1, the error decreases monoton- ically with increasing bin number K (second term) but the variance of bin labels depends on the number of sample transitions n1 (first term). The inter-POMDP error (third term) does not affect the bin assignment. Hence, choosing a large K will amount to pairing states by rankings, but results in high variance, as orderings are estimated from data and each bin will have 1. Setting K too small will group together unrelated observations, inducing high bias.
6 EXPERIMENTS
Unlike for single task offline RL (Fu et al., 2020), most works on zero-shot generalization from of- fline data either come up with ad hoc solution suiting their needs, e.g. (Ball et al., 2021), assess performance on benchmarks that do not evalute generalization across observation functions (e.g., Yu al,, 2020). To accelerate progress in this field, we devised the offline Procgen benchmark, an of- fline RL. dataset to directly test for generalization of offline RI. agents across observation functions
Offline Procgen benchmark We evaluate the proposed approach on offline version of the Procgen benchmark (Cobbe et al., 2020), which is widely used to evaluate zero-shot generalization across complex visual perturbations. Given random seed, Procgen allows to sample procedurally generated level configurations for 16 games under various complexity modes: "casy". "hard" and "exploration". The dataset is obtained as follows: we first pre-train PPO (Schulman et al., 2017) agent for 25M timesteps on 200 levels of "easy" distribution for each environment ("easy" mode is widely used to test generalization capabilities (Cobbe al., 2020; Raileanu and Fergus, 2021; Ma- zoure et al., 202.)). All agents use the IMPALA encoder architecture (Espeholt et al., 2018), which has enough parameters to allow better generalization performance, compared to other models (e.g.. Mnih et al., 2015).
'The benchmark will be open-sourced. We use the TFAgents' implementation (Guadarrama et al., 2018)
8

---

Under review as a conference paper at ICLR 2022
We compare the zero-shot performance the entire distribution of "easy" POMDPs for GSF against that strong RL and representation learning baselines: behavioral cloning (BC) to assess the quality of the PPO policy, CQL (Kumar et al., 2020) the current otate-of-the-art multiple offlin benchmarks which balances RL and BC objectives, CURL (Srinivas et al., 2020), CTRL (Mazoure al., 2021), DeepMDDP (Gelada et al., 2019) which learns metric closely lated to bisimulation across the MDP, Value Prediction Network (VPN, Oh et al., 2017) which combines model-free and model-based learning of values, observations, next observations, rewards and discounts, Cross-State Self- Self-Constraint (CSSC, Liu et al., 2020a) which boosts similarity of ob- servations with identical action sequences, as well as Policy Similarity Embeddings (Agarwal al., 2020), which groups observation representations based on distance in optimal policy space.
Results
CoL-nomalized
1.0
0.0
1.0-
BC
DeepMDP
VPN
CSSC
PSE
GSF reward
GSF SF norm
GS
CURL
CTRL
Algorithm
action count
Figure 3: Returns on the offline Procgen benchmark (Cobbe et al., 2020) after 1M training steps. Boxplots are constructed over 5 random seeds and all 16 games; each method is normalized by the per-game median CQL performance. White dots represent average distribution.
Fig. 3 shows the performance of all methods over 5 random seeds and all 16 games the offline Procgen benchmark after 1 million training steps. Per-game average scores for all methods can be found in 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.
Using different cumulants functions can lead to different label assigmments and hence different sim- ilarity groups. Fig. 3 examines the performance GSF with respect cumulants: r(s,, at), rewards s.t. GSF learns the policy's Qu-vvlue, 2) 1., (o). the successor representation" (Dayan, 1993; Barreto et al., 2016) s.t. GSF learns induced distribution over D (Machado al., 2020) and 3) 1o, (a), action counts, s.t. GSF learns discounted policy. While rewards and successor feature cumulant choices leads to similar performance, using action-based distance leads to larger variance.
7 DISCUSSION
In this work we proposed GSF, a novel algorithm which combines reinforcement learning with rep- resentation learning improve zero-shot generalization performance on challenging, pixel-based control tasks. GSF relies on computing the similarity between observation pairs with respect to any instantancous accumulated signal, which leads to improved empirical performance on the newly in- troduced offline Procgen benchmark. Theoretical results suggest that GSF hyperparameter choice depends trade-off between finite sample approximation and extrapolation error.
While our work answered some questions regarding zero-shot generalization in offline RL, some questions persist: can Gaseb distances be included in contrastive objective without the need for quantile discretization (perhaps through re-scaling or order statistics)? Can the cumulant function be chosen priori for specific task structure other than Procgen, and shown to lead to optimal representations?
In the continuous observation space, we learn d-dimensional successor feature vector Zi via TD and computing the quantiles over || |/1-

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
To ensure reproducibility of our work, we attach the source code to the submission, provide all proofs of both theorems in the appendix (with assumptions) and will open-source the offline Procgen benchmark.
REFERENCES
A. Agarwal, M. Henaff, S. Kakade, and W. Sun. Pe-pg: Policy cover directed exploration for provable policy gradient learning. Neural Information Processing Systems, 2020.
R. Agarwal, M. C. Machado, P. S. Castro, and M. G. Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. arXiv preprint arX1v:2101.05266, 2021.
P. J. Ball, C. Lu, J. Parker-Holder, and S. Roberts. Augmented world models facilitate zero-shot dynamics generalization from single offline environment. International Conference on Machine Learning. 2021.
A. Barreto, W. Dabney, R. Munos, J. J. Hunt, T. Schaul, H. Van Hasselt, and D. Silver. Successor features for transfer in reinforcement learning. arXiv preprint arXiv:1606.053: 2016.
P. L. Bartlett. The sample complexity of pattern classification with neural networks: the size of the weights is more important than the size the network. IEEE transactions on Information Theory, 44(2):525-536, 1998.
M. Bellemare, Candido, P. S. Castro, Gong, M. C. Machado, Moitra, S. Ponda, and Z. Wang. Autonomous navigation of stratospheric balloons using reinforcement learning, Nature, 588(7836):77-82, 2020.
J. Boyan and A. W. Moore. Generalization in reinforcement learning: Safely approximating the value function. Advances neural information processing systems, pages 369-376, 1995.
M. Caron, I. Misra, J. Mairal, P. Goyal, P. Bojanowski, and A. Joulin. Unsupervised learning of visual features by contrasting cluster assignments. arXiv preprint 2rXiv:2006.008882 2020.
S. 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.
P. S. Castro and D. Precup. Using bisimulation for policy transfer in mdps. In Twenty- AAAI Conference on Artificial Intelligence, 2010.
A. S. Chen, S. Nair, and C. Finn. Learning generalizable robotic reward functions from" in-the-wild" human videos. arXiv preprint arXiv:2103.16817, 2021.
C.A.. Cheng, A. Kolobov, and A. Swaminathan. Hcuristic-guide -uidee reinforcement learning. arXiv preprint arXiv: 102.02000.2777 2021.
K. Cobbe, O. Klimov, C. Hesse, T. Kim, and J. Schulman. Quantifying generalization reinforce- ment learning. In International Conference Machine Learning, pages 1282-1289. PMLR, 2019.
K. Cobbe, C. Hesse, J. Hilton, and J. Schulman. Leveraging procedural generation to benchmark reinforcement learning. International conference on machine learning, pages 2048-2056. PMLR, 2020.
Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Computation, 5(4):613-624, 1993.
A. Dvoretzky, J. Kiefer, and J. Wolfowitz. Asymptotic minimax character of the sample distribution function and of the classical multinomial estimator, The Annals of Mathematical Statistics, pages 642-669, 1956.
10

---

Under review as a conference paper at ICLR 2022
D. Ernst, P. Geurts, and L. Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6:503-556, 2005.
L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International Conference Muchine Learning, pages 1407-1416. PMLR, 2018.
N. Ferns, P. Panangaden, and D. Precup. Metrics for finite markov decision processes. In UAI, volume 4, pages 162-169, 2004.
J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven rein- forcement learning, 2020.
S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062. PMLR, 2019.
C. Gelada, S. Kumar, J. Buckman, O. Nachum, and M. G. Bellemare. Deepmdp: Learning con- tinuous latent space models for representation learning, International Conference on Machine Learning. pages 2170-21777 PMLR, 2019.
J.-B. Grill, F. Strub, F. Altché, C. Tallec, P. H. Richemond, E. Buchatskaya, C. Doersch, B. A. Pires, Z. D. Guo, M. G. Azar, et al. Bootstrap your own latent: new approach to self-supervised learning, arXiv preprint arXiv:2006.001 077733,, 2020.
Guadarrama, A. Korattikara, O. Ramirez, P. Castro, E. Holly, Fishman, K. Wang. E. Go- nina, N. Wu, E. Kokiopoulou, L. Sbaiz, J. Smith, G. Bartók, J. Berent, C. Harris, V. Van- houcke, and E. Brevdo. TF-Agents: A library for reinforcement learning in tensorflow. https://githin
wom/tensooflow httssoorllo//ggntst... agents, 2018. URL forsorfloo/ageettt.. [Online; accessed 4-October-2021J.
com/
K. He, H. Fan, Y. Wu, S. Xie, and R. Girshick. Momentum contrast for unsupervised visual repre- sentation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729-9738, 2020.
M. Hessel, H. Soyer, L. Espeholt, W. Czarnecki, Schmitt, and H. van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 33, pages 3796-3803, 2019.
I. Higgins, A. Pal, A. Rusu, L. Matthey, C. Burgess, A. Pritzel, M. Botvinick, C. Blundell, and A. Lerchner. Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning, pages 1480-1490. PMLR, 2017.
R. D. Hjelm, A. Fedorov, S. LLvoiie-Marchlllee Marchildon, K. Grewal, P. Bachman, A. Trischler, and Y. Ben- gio. Learning deep representations by mutual information estimation and maximization. arXiv preprint arXiv: 066700, 2018.
N. K. Jong and Stone. State abstraction discovery from irrelevant state variables. In IJCAI, volume 8, pages 752-757. Citeseer, 2005,
Khosla, P. Teterwak, C. Wang, A. Sarna, Y. Tian, P. Isola, A. Maschinot, C. Liu, and D. Krishnan. Supervised contrastive learning. Neural Information Processing Systems, 2020.
R. Kidambi, A. Rajeswaran, P. Netrapalli, and T. Joachims. Morel: Model-based offline reinforce- ment learning. arXiv preprint arXiv:2005.05951, 2020.
I. Kostrikov, D. Yarats, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.
I. Kostrikov, J. Tompson, R. Fergus, and O. Nachum. Offline reinforcement learning with fisher divergence critic regularization. arXiv preprint arXiv:2103.08050, 2021.
A. Kumar, A. Zhou, G. Tucker, and S. Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint 07X17 2020.

---

Under review as a conference paper at ICLR 2022
S. Lange, T. Gabel, and M. Riedmiller. Batch reinforcement learning. In Reinforcement learning, pages 45-73. Springer, 2012.
L. Li, J. Walsh, and M. L. Littman. Towards a unified theory of state abstraction for mdps. ISAIM, 4:5, 2006.
Z. Lin. Yang, Zhao, I. Qin, G. Yang, and T.-Y. Liu. Rd 2: Reward decomposition with representation decomposition. Advances in Neural Information Processing Systems, 33, 2020.
G. T. Liu, P.-J Cheng, and G. Lin. Cross-state self-constraint for feature generalization in deep reinforcement learning. 2020a.
L. Liu, W. Hamilton, G. Long. J. Jiang, and H. Larochelle. universal representation transformer layer for few-shot image classification. arXiv preprini urXiv:2006. 1702, 2020b.
Y. Liu, A. Swaminathan, A. Agarwal, and E. Brunskill. Provably good batch reinforcement learning without great exploration. NeurlPS, 2020c.
M. C. Machado, M. G. Bellemare, and M. Bowling. Count-based exploration with the successor representation. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pages 5125-5133, 2020.
H. B. Mann and A. Wald. On stochastic limit and order relationships. The Annals of Mathematical Statistics, 14(3):217-226, 1943.
B. Mazoure, R. T. d. Combes, T. Doan, P. Bachman, and R. D. Hjelm. Deep reinforcement and infomax learning. Neural Information Processing Systems, 2020.
B. Mazoure, A. M. Ahmed, P. MacAlpine, R. D. Hjelm, and A. Kolobov. Cross-trajectory represen- tation learning for zero-shot generalization in rl. arXiv preprint arXiv:2106.02193, 2021.
D. Misra, M. Henaff, A. Krishnamurthy, and J. Langford. Kinematic state abstraction and prov- ably efficient rich-observation reinforcement learning. In International conference machine learning, pages 6961-6971. PMLR, 2020.
V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Ried- miller, A. K. Fidjeland. G. Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533, 2015.
K. P. Murphy. A survey of pomdp solution techniques. environment, 2:X3, 2000.
O. Nachum and M. Yang. Provable representation learning for imitation with contrastive fourier features. Neural Information Processing Systems, 2021.
Oh, S. Singh, and H. Lee. Value prediction network. arXiv preprint arXiv:1707.03497, 2017.
A. v. d. Oord, Y. Li, and O. Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXhv:1807.0377 2018.
M. L. Puterman. Markov decision processes. Handbooks in operations research and management science, 2:331-434, 1990.
R. Raileanu and R. Fergus. Decoupling value and policy for generalization in reinforcement learn- ing. International Conference Machine Learning, 2021.
Schulman, Wolski, Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707.07347, 2017.
M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman. Data-efficient re- inforcement learning with self preficrree representations. International Conference on Learning Representations, 2020.
M. Schwarzer, N. Rajkumar, M. Noukhovitch, A. Anand, L. Charlin, D. Hjelm, P. Bachman, and A. Courville, Pretraining representations for data-efficient reinforcement learning. arXiv preprint (xiv:2106.04470 2021.
12

---

Under review as a conference paper at ICLR 2022
S. Sinha and A. Garg. S4rl: Surprisingly simple sulf--uup for offline reinforcement learning. arXiv preprint arXiv:2103.06326, 2021.
J. Song and S. Ermon. Multi-label contrastive predictive coding. Neural Information Processing Systems, 2020.
X. Song, Y. Jiang, S. Tu, Y. Du, and B. Neyshabur. Observational overfitting in reinforcement learning. International Conference on Learning Representations, 2019.
A. Srinivas, M. Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforce- ment learning. International Conference on Machine Learning, 2020.
A. Stone, O. Ramirez, K. Konolige, and R. Jonschkowski. The distracting control suite-a challeng- ing benchmark for reinforcement learning from pixels. arXiv preprint arXi::2101.02722 2021.
A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforce- ment learning. In International Conference on Machine Learning, pages 9870-9879, PMLR, 2021.
W. Sun, J. A. Bagnell, and B. Boots. Truncated horizon policy search: Combining reinforcement learning & imitation learning. arXiv preprint arXiv: /240, 2018.
R. S. Sutton, J. Modayil, M. Delp, T. Degris, P. M. Pilarski, A. White, and D. Precup. Horde: A scal- able real-time architecture for learning knowledge from unsupervised sensorimotor interaction. In The 10th International Conference on Autonomous Agents and Multiagent Destems-Volume 2, pages 761-768, 2011.
A. Swaminathan and T. Joachims. Batch learning from logged bandit feedback through counterfac- tual risk minimization. The Journal of Machine Learning Research, 16(1):1331-1755 2015.
A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. arXiv preprint arXiv:2103.00 07945, 2021.
E. Triantafillou, T. Zhu, V. Dumoulin, Lamblin, U. Evei, K. Xu, R. Goroshin, C. Gelada, K. Swer- sky, P.-A. Manzagol, et al. Meta-dataset: A dataset of datasets for learning to learn from few examples. International Conference on Learning Representations, 2019.
G. Valle-Pérez and A. Louis. Generalization bounds for deep learning. arXiv preprint arXiv:2012.04115, 2020.
A. W. van der Vaart. Asymptotic statistics. cambridge series in statistical and probabilistic mathe- matics, 1998.
R. Wang, Y. Wu, R. Salakhutdinov, and S. M. Kakade. Instabilities of offline rl with pre-trained neural representation. arXiv preprint arXiv:2103. 04947, 2021a.
Y. Wang, R. Wang, and S. M. Kakade. An exponential lower bound for linearly-realizable mdps with constant suboptimality gap. arXiv preprint arXiv:2 122.103....00 /2690, 2021b.
C. J. Watkins and P. Dayan. Q-learning. Machine learning, 8(3-4):279-292. 1992.
Y. Wu, G. Tucker, and O. Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv: 11361, 2019.
D. Yarats, A. Zhang, I. Kostrikov, B. Amos, J. Pineau, and R. Fergus. Improving sample efficiency in model-free reinforcement learning from images. arXiv preprint 1rxiv:1910.0/741 2019.
Yu, Quillen, Z. He, R. Julian, K. Hausman, C. Finn, and S. Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, pages 1094-1100. PMLR, 2020.
A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. International Conference Learning Represen- tations, 2020.
13


Reviews:
[{"review_id": "xTRt1Ii2Umf", "reviewer": "Reviewer_vyyv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed method is an interesting method for defining 'similarity' which can then be used for contrastive learning. The paper is well written and does not appear to have any technical or theoretical flaws. The topic of generalization in RL is relevant and the proposed method is a valuable contribution to this area. \n\nOne disadvantage of this method is that it requires the ability to pre-traing the generalized value function, hence the authors' decision to only apply it to offline RL. This leads to the, imo, main weakness of the paper: While the authors evaluate the method against a range of relevant baselines, they only do so in one domain, namely their newly proposed offline procgen benchmark. Hence the results will depend heavily on the authors effort to not only tune the hyperparameters of their own method, but also those of the baselines. Unfortunately, there is no information in the paper how the hyperparameters were chosen (they do provide the final hyperparameters of their own method only), so it's hard to evaluate how much one should trust the experimental evidence.\n\nHence, I think the paper could be strengthend significantly by evaluating on additional domains and/or providing detailed information about they applied hyperparameter search.\n\nQuestion/Nitpick: \n* Can you still call it contrastive loss? Eq. (7) just seems to be classification? ", "summary_of_the_paper": "The authors build on previous work which regularized states with similar value-function or future action sequences to have similar representations, by instead looking at successor features, i.e. the accumulated value of future cumulants. To deal with continuous value function whose scale can vary accross environments, they apply a quantile-binning technique. They also propose an offline version of Procgen as benchmark.", "summary_of_the_review": "Overall interesting method, but more experimental detail needed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "h5oinElEZz", "reviewer": "Reviewer_f74c", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strength:\n\n- RL generalization in an offline setting is an important and interesting topic for our community. \n- The approach is well-motivated by the example in Section 4. \n- The proposed GVF is a general framework that can recover objectives of existing works. \n- The proposed approach outperforms baselines such as CQL, CCSC and PSE. \n\nWeakness: \n\n- The reviewer found the approach section is not easy to follow. Some technical details seem missing. Please see detailed comments below:  \n    - The latent state decoder f seems to be a critical part of the proposed pipeline. Could you elaborate on how you learn this decoder?\n    - Definition 1 states that the cumulant function c can be any bounded function over R^d. It is unclear to the reviewer how to select the cumulant function c. What cumulant function is used in the experiments. Does each task require a different cumulant function?   \n\n\n- In addition, the reviewer found the notations confusing. Please see detailed comments below:\n   - In line 2 of Algorithm 1, what does D[i] represent? What is the difference between D[i] and D^{\\mu}\n   - In line 6 of algorithm 1, does G represent GVF? If yes, why does it take a latent representation as input? According to Definition 1, shouldn’t GVF take an observation o as input?\n   - \"Target parameter\" and \"online parameter\" are not formally defined\n   - $\\beta$ in line 7 of algorithm 1 is undefined?     \n\n\n- Experiments  \n  - The authors discussed F-BRC (Kostrikov et al., 2021) and MOReL (Kidambi et al., 2020), which are state-of-the-art offline RL approaches, in Section 2. However, there is no comparison with the aforementioned existing works provided. The experimental section could be more convincing if the comparison with F-BRC and MOReL is provided.   \n\n  - Recent works [a] show that data augmentation only could significantly improve the generalization capability of RL agents. The proposed approach also uses data augmentation during training (Line 4 of Algorithm 1).  It would be interesting to see an ablation study on the data augmentation and the proposed GVF.  \n\n[a] Reinforcement Learning with Augmented Data, Laskin et al. NeurIPS 2020\n", "summary_of_the_paper": "This paper studied  zero-shot generalization in an offline reinforcement learning setting. The authors proposed to improve the generalization via a better representation learning. Specifically, the authors hypothesized that observation with similar future behaviors should be assigned to similar representations. To this end, the generalized similarity function (GSF) that aggregates latent representations with the future behavior is proposed. The proposed approach is evaluated on an offline version of Procgen. The proposed approach outperforms baselines across different tasks in Procgen. \n", "summary_of_the_review": "The reviewer thinks the topic of this paper is significant. The approach seems interesting. However, the reviewer has some concerns about the experimental results and the clarity of the presentation. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "d6uDEXgrHDQ", "reviewer": "Reviewer_nnNd", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strenght:\n* The paper is clearly written and the various choices are reasonably justified\n* The proposition of an offline version of Procgen is a valuable asset for the offline RL research community\n* the approach of quantization in this context is novel to the best of my knowledge\n* The experimental results seem significant for the considered benchmark dataset\n* the introduction of generalization value function seem novel and pertinent\n\nWeakness:\n* The comparison to cURL (https://arxiv.org/abs/2004.04136) is missing, especially considering the fact that this approach is also leveraging metric learning in latent state space.\n* The D4RL dataset would have been useful to compare too, especially to evaluate over the latent recent corpus of approaches of offline RL.\n* The approach relies on PPO, which is an online RL algorithm, and CQL which is a conservative approach of Q-Learning. Maybe it would have been interesting to evaluate how the method behaves with SAC also.", "summary_of_the_paper": "The paper addresses the problem of generalization and representation learning in reinforcement learning.\nConcretely, the authors propose an approach of self-supervision for improving offline-RL based on similarity learning over the considered state space.\nIn this context, the concept of generalized value function is introduced.\nThe similarity measure is defined as a classification problem over quantized latent space.\nThe paper is evaluated in an offline setting of Procgen which is a dataset for multi-task reinforcement learning generalization evaluation.\nMore particularly, the proposed approach claims to be particularly efficient in improving zero-shot generalization performance on one offline RL benchmark, offline Procgen.\nThe offline Procgen dataset is a second contribution of the paper.", "summary_of_the_review": "The paper is addressing an important problem of sequential decision learning research.\nThe proposing is interesting while novelty to existing approach could be improved.\nThe proposition of a new dataset for zero-shot generalization of sequential decision-making in an offline context is valuable.\nThe comparison on another popular benchmark like D4RL and against the cURL approach would have been appreciated.\nThe introduction of generalized value function seems novel and pertinent in the context of self-supervised reinforcement learning and is clearly illustrated in the proposal of this work and could be valuable to the state of the art.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "ajVZZdrPb7j", "reviewer": "Reviewer_8hYV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The idea of aggregating observations based on behavior similarity has already been explored in some papers, but the approach proposed in this paper seem to be new.\nWhile the paper discusses its relation to CSSC and PSE and compares with these two methods in the experiments too, it misses a discussion or empirical comparison with CTRL. In addition, it is not clear whether the proposed algorithm's advantage over CSSC and PSE in the experiments is due to the use of a better aggregation method, which is the main claimed novelty. Can other aggregation methods be combined with CQL?\n\nThe problem setup and the motivation are not clear and often confusing. \n* There is a consistent confusion of MDP and POMDP throughout the paper. For example, in the introduction, the paper claims improving zero-shot generalization across a family of POMDPs, but the cited reference is Puterman's article on MDP. It then sometimes called the models POMDPs or MDPs. \n* In the historical dataset ${\\cal D}^{\\mu}$, what is required for $\\mu$?\n* What is a concrete example of the problem setting described in Section 3.1? In the motivating example, it seems the first images for Level 1 and Level 2 should be considered similar, but it is not clear why. The main paper states that they have near identical true latent states and value functions, but didn't explain why.\n* It's not clear what the different MDP/POMDP models are in the experiments.\n\nThe technical writing for the proposed approach requires a careful revision too. Here are some issues:\n- The term cumulant function is confusing as cumulant function is a well-known concept in statistics.\n- Section 3.1: for POMDP, a policy is usually maps a belief to an action, not an observation to an action. Also note that the Q-function in Eq. (1) is a function of the state and the action, and thus not directly computable because the state is not observed.\n- $c$'s output may be a vector, and $|c(o, a)|$ and $\\sup_{o, a} c(o, a)$ need to be defined.\n- Alg. 1: D[i] in line 1 is undefined; line 6 works only for scalar c; line 7 is not clear.\n- Eq. (3): RHS depend on $i$ and $j$, but LHS does not. Where is the $d_{\\mu}$ used?\n- $c_{1,\\max}^{\\mu}$ and $c^{\\mu}_{2, max}$ are undefined.\n- In $F_{i}(g) = \\mathbb{P}[G_{i}(o_{t}) \\le g]$, is $o_{t}$ treated as a random variable? What is the underlying distribution used? Is $G_{i}$ the same as $G_{i}^{\\mu}$?\n- Eq. (4): what does $G_{i} \\sim {\\cal D}^{\\mu}$ mean? Note that ${\\cal D}^{\\mu}$ is defined as a set of $(o, a, r)$ triplets.\n- Eq. (6): $i$ appears on the RHS but not on the LHS.\n- Eq. (7): $l(z)$ and $o \\sim {\\cal D}$ undefined. Does ${\\cal D}$ refer to ${\\cal D}^{\\mu}$?\n- Eq. (8): $(s, a, r) \\sim {\\cal D}^{\\mu}$? \n- Alg. 1: should line 10 be inside the for loop? Line 11 is not clear.\n\nThe representation learned by DeepMDP is closely connected to bisimulation, but it doesn't learn bisimulation relations.\n\n**Post-rebutall**\n\nI appreciate the clarification and the additional results. I've raised my score in view of these changes. It'd be helpful to further stengthen the experiments and improve the writing. The writing is generally clear now, but there are still quite a few issues (e.g., $o_{1}$ and $o_{2}$ in Eq. (3) now have superscripts, but this now is inconsistent with other ocurrences of $o_{1}$ and $o_{2}$ around Eq. (3)), and a careful revision will help readers to better understand the idea. For the experiments, it seems the baselines are not fine-tuned, and it's also not clear how the online algorithms are set-up for the offline setting. Providing details on how CSSC/PSE/CTRL is combined with CQL will also be helpful. In addition, comparison to existing offline algorithms will be helpful (e.g. refer to Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems by Levine et al.).\n\n", "summary_of_the_paper": "This paper proposes improving the generalization performance offline RL algorithms by using a new approach to aggregate observations based on the similarity of their expected future behavior.", "summary_of_the_review": "The proposed approach seems novel, but the motivation and presentation of the idea is confusing. It is not clerar whether the reported performance improvement is due to the proposed aggregation method contributes, and further discussion or experiments with related works may be needed.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a new offline RL technique to generalize across domains. The paper was initially confusing (i.e., MDP vs POMDP) and weak empirically.  The authors greatly improved the paper.  However, a the end of the day, it is still not clear why the proposed approach performs better than existing techniques.  We can think of the cumulant function with the discrete labels as essentially computing some statistics of future actions, observations and rewards.  This is what every self supervised technique does.  They differ in terms of their particular choice of statistics and architecture.  The paper does not sufficiently motivate the particular architecture.  Interestingly, in the experiments, the best statistics are cumulative rewards, which are closely related to the Q-values. In that case, it is even less clear why the approach should be beneficial since RL techniques that generalize across domains by learning state representations to predict Q-values seem very closely related.  

Despite the updates to the paper, the POMDP references are still confusing.  The issue is that the paper embeds observations as if they were sufficient to predict future observations and rewards.  This corresponds to the memoryless approach where a policy is optimized based on the last observation instead of the history of past actions and observations.  Memoryless strategies are effective only when the last observation is a sufficient statistic, meaning that we really have a (near) fully observable MDP.  The paper should discuss this and acknowledge that the approach will suffer in domains where memory of past actions and observations is critical.
Decision: Reject