I MPROVING Z ERO - SHOT G ENERALIZATION IN O FFLINE R EINFORCEMENT L EARNING USING G ENERALIZED S IMILARITY F UNCTIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
Reinforcement learning (RL) agents are widely used for solving complex sequential decision making tasks, but still exhibit difﬁculty in generalizing to scenarios not seen during training.
While prior online approaches demonstrated that using additional signals beyond the reward function can lead to better generalization capabilities in RL agents, i.e. using self-supervised learning (SSL), they struggle in the ofﬂine RL setting, i.e. learning from a static dataset.
We show that performance of online algorithms for generalization in RL can be hindered in the ofﬂine setting due to poor estimation of similarity between observations.
We propose a new theoretically-motivated framework called Generalized Similarity Functions (GSF), which uses contrastive learning to train an ofﬂine RL agent to aggregate observations based on the similarity of their expected future behavior, where we quantify this similarity using generalized value functions .
We show that GSF is general enough to recover existing SSL objectives while also improving zero-shot generalization performance on a complex ofﬂine RL benchmark, ofﬂine Procgen.
I NTRODUCTION
Reinforcement learning (RL) is a powerful framework for solving complex tasks that require a sequence of decisions.
The RL paradigm has allowed for major breakthroughs in various ﬁelds, e.g. outperforming humans on video games (Mnih et al., 2015; Schwarzer et al., 2020), controlling stratospheric 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, i.e. in which training involves a ﬁnite number of related tasks sampled from some distribution, with a potentially distinct sampling distribution during test time (Cobbe et 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: it 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 to as zero-shot generalization.
In our work, we are interested in the problem of zero-shot generalization where the difference between tasks is predominantly due to perceptually distinct observations.
An example of this setting is any environment with distractor features (Cobbe et al., 2020; Stone et al., 2021), i.e. features with no dependence on the reward signal nor the agent’s decisions.
This generalization setting has recently received much attention (Liu et al., 2020a; Agarwal et al., 2021; Mazoure et al., 2021), due to its particular relevance to real-world scenarios, for example deploying a single autonomous driving agent at day or at night.
Generalization capabilities of an agent can be analyzed through the prism of representation learning , under which the agent’s current belief about a rich and high-dimensional environment are summarized in a low-dimensional entity, called a representation.
Recent work in online RL has shown that learning state representations with speciﬁc properties such as disentanglement (Higgins et al., 2017) or linear separability (Lin et al., 2020) can improve zero-shot generalization performance.
Achieving this with limited data (i.e. ofﬂine RL) is challenging, since the representation will have a large estimation error over regions of low data coverage.
A common solution to mitigate this task-speciﬁc overﬁtting and extracting the most information out of the data consists in introducing auxiliary learning signals other than instantaneous reward (Raileanu and Fergus, 2021).
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 Procgen remains limited even when training on 200M frames, while generalization-oriented agents (Raileanu and Fergus, 2021; Mazoure et al., 2021) 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 an ofﬂine, ﬁxed dataset of experience (Ernst et al., 2005; Lange et al., 2012).
A natural limitation for generalization from ofﬂine data is that policy improvement is dependent on dataset quality.
Speciﬁcally, highdimensional 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 ofﬂine RL research.
Improving zero-shot generalization performance from an ofﬂine dataset of high-dimensional observations is therefore a hard problem due to limitations on dataset size and quality.
In this work, we are interested in improving zero-shot generalization across a family of PartiallyObservable Markov decision processes (MDPs, Puterman, 1990) in an ofﬂine RL setting, i.e. by training agents on a ﬁxed dataset.
We hypothesize that in order for an RL agent to be able to generalize 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 et al., 2011) to capture future behavior with respect to an 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., 2021).
Our main contributions are as follows:
We propose Generalized Similarity Functions (GSF), a novel self-supervised learning algorithm for reinforcement learning, that aggregates latent representations by the future behavior (or generalized value function) under their respective observations.
We devise a new benchmark constructed to test zero-shot generalization of ofﬂine RL algorithms: ofﬂine Procgen.
It consists of 5M transitions from 200 related levels of 16 distinct games.
We evaluate performance of GSF and other baseline methods on ofﬂine Procgen, and show that GSF outperforms both previous state-of-the-art ofﬂine RL and representation learning baselines on the entire distribution of levels.
We analyze the theoretical properties of GSF and describe the impact of hyperparameters and cumulant functions on empirical behavior.
R ELATED W ORKS
Generalization in reinforcement learning Generalizing a model’s predictions across a variety of unseen, high-dimensional inputs has been extensively studied in the static supervised learning setting (Bartlett, 1998;
Triantaﬁllou et al., 2019; Valle-P´erez and Louis, 2020; Liu et al., 2020b).
Generalization in RL has received a lot of attention: extrapolation to unseen rewards (Barreto et al., 2016;
Misra et al., 2020), observations (Zhang et al., 2020; Raileanu and Fergus, 2021; Liu et al., 2020a;
Agarwal et al., 2021; Mazoure et al., 2021) and transition dynamics (Ball et al., 2021).
Each generalization scenario is best solved by their respective set of methods: sufﬁcient 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 augmentation is a promising technique, but typically relies on handcrafted domain information, which might not be available a priori .
In fact, we will show in our experiments that generalization in the ofﬂine RL setting is poor even when using such handcrafted data augmentations, without additional representation learning mechanisms.
In this work, we posit that representation learning should use instantaneous auxiliary signals in order to prevent overﬁtting onto a unique signal (e.g. reward across tasks) and improve generalization performance.
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 optimizing for the reward alone.
However, more complex settings may require additional auxiliary signals in order to ﬁnd state abstractions better suited for control.
The problem of learning meaningful state representations (or abstractions) for planning and control has been extensively studied previously (Jong and Stone, 2005; Li et al., 2006), but saw real breakthroughs only recently, in particular due to advances in self-supervised learning (SSL).
Outside of RL, SSL has achieved spectacular results by closing the gap between unsupervised and supervised learning on certain datasets (Hjelm et al., 2018;
Oord et al., 2018; Caron et al., 2020;
Grill et al., 2020).
Representation learning, and speciﬁcally self-supervised learning, has also been used to achieve state-of-the-art generalization and sample efﬁciency results in RL on challenging control problems such as data efﬁcient Atari (Schwarzer et al., 2020; 2021), DeepMind Control (Agarwal et al., 2021) and Procgen (Mazoure et al., 2020; Stooke et al., 2021; Raileanu and Fergus, 2021; Mazoure et al., 2021).Noteworthy instances of theoretically-motivated representation learning methods for RL include heuristic-guided learning (Sun et al., 2018; Cheng et al., 2021), and random Fourier features (Nachum and Yang, 2021).
Ofﬂine reinforcement learning When learning from a static dataset, agents should balance interpolation and extrapolation errors, while ensuring proper diversity of actions (i.e. prevent collapse to most frequent action in the data).
Popular ofﬂine RL algorithms such as BCQ (Fujimoto et al., 2019),
MBS (Liu et al., 2020c), and CQL (Kumar et al., 2020) rely on a behavior regularization loss (Wu et al., 2019) as a tool to control the extrapolation error.
Some methods, such as F-BRC (Kostrikov et al., 2021) are deﬁned only for continuous action spaces while others, such as MOReL (Kidambi et al., 2020) estimate a pessimistic transition model.
The major issue with current ofﬂine 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.
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.
P ROBLEM SETTING 3. P ARTIALLY - OBSERVABLE M ARKOV DECISION PROCESSES
A (inﬁnite-horizon) partially-observable Markov decision process (POMDP, Murphy, 2000)
M is deﬁned by the tuple M = h S , p 0 , A , p S , O , p O , r,  i , where S is a state space, p 0 =
P [ s 0 ] is the starting state distribution,
A is an action space, p S = P
[ ·| s t , a t ] : S ⇥ A !
( S ) is a transition function, O is an observation space, p O = P [ ·| s t ] : S !
( O ) 1 is an observation function, r :
S ⇥ A !
[ r min , r max ] is a reward function and  2 [0 , 1) is a discount factor.
The system starts in one of the initial states s
⇠ p 0 with observation o
⇠ p O ( ·| s 0 ) .
At every timestep t = 1 , 2 , 3 , .. , the agent, parameterized by a policy ⇡ : O !
( A ) , samples an action a t ⇠ ⇡ ( ·| o t ) .
The environment transitions into a next state s t +1 ⇠ p S ( ·| s t , a t ) and emits a reward r t = r
( s t , a t ) along with a next observation o t
[Equation]
The goal of an RL agent is to maximize the cumulative rewards
P 1 t =
t r t obtained over the entire episode.
Value-based off-policy RL algorithms achieve this by estimating the state-action value function under a target policy ⇡ :
[Equation]
where P ⇡ t denotes the joint distribution of { s t + k , a t + k } 1 k =1 obtained by executing ⇡ in the environ- ment.
An important distinction from online RL is that, instead of sampling access to the environment, we assume access to a historical dataset D µ collected by logging experience of the policy, µ , in the form { o i,t , a i,t , r i,t } i = N,
t = T i =1 ,t =1 where, for practical purposes, the episode is truncated at T timesteps.
Furthermore, we assume that the agent can only be trained on a limited collection of POMDPs
M train =
{ M i } mi =1 , and its performance is evaluated on the set of test POMDPs M test .
We assume that both
M train and M test were sampled from a common task distribution and that every POMDP M i 2 M = M train [ M test shares the same transition dynamics and reward function with M but has a different observation function p i,
O .
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 o t and corresponding states s t throughout the paper.
3. R EPRESENTATION LEARNING
Previous works in the RL literature have studied the use of auxiliary signals to improve generalization performance.
Among others, Liu et al.
(2020a)
; Agarwal et al.
(2021) deﬁne the similarity of two observations to depend on the distance between action sequences rolled out from that observation under their respective optimal policies.
They achieve this by ﬁnding a latent space Z ✓ S in which the distance d Z ( z, z 0 ) for all z, z 0 2 Z is equivalent to distance between true latent states d
S ( s, s 0 ) for all s, s 0 2 S ; the aforementionned works learn Z by optimizing action-based similari- ties between observations.
In practice, latent space z is decoded from observation o using a latent state decoder f : O !
Z from observation o t .
Through the paper, we assume that all value functions have a linear form in the latent decoded state, i.e. Q ✓ ( o, a ) =
✓ > a f ( o ) =
✓ > a z , which agrees with our 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 ofﬂine RL, and how existing action-based similarity metrics fail to recover the true latent states for important families of
POMDPs.
M OTIVATING EXAMPLE
Multiple recently proposed self-supervised objectives (Liu et al., 2020a; Agarwal et al., 2021) conjecture that observations
o 1
M 1 , o 2 2 M 2 that emit similar future action sequences under optimal policies ⇡ ⇤ 1 , ⇡ ⇤ 2 should be decoded into nearby latent states z 1 , z 2 .
While this heuristic can correctly group observations with respect to their true latent state in simple action spaces, it fails to identify similar pairs of trajectories in POMDPs with multiple optimal policies.
For instance, two trajectories might visit an identical set of latent states, but have drastically different actions.
Fig.
1 shows one such example: two levels of the Climber game have a near-identical true latent state (see Appendix) and value function (average normalized mean-squared error of 0.0398 across episode), while having very different action sequences from a same PPO policy (average total variation distance of 0.4423 across episode).
The problem is especially acute in Procgen, since the PPO policy is high-entropy for some environments (see Fig. 4), 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 a better state representation than reasoning about action similarities.
In a POMDP with a different structure, grouping representations by action sequences can be optimal.
So how do we unify these similarity metrics under a single framework?
In the next section, we use this insight to design a general way of improve representation learning through self-supervised learning of discounted future behavior.
M ETHOD
We propose to measure a generalized notion of future behavior similarity using generalized value functions, as deﬁned by the corresponding cumulant function.
The choice of cumulant determines which components of the future trajectory are most relevant for generalization.
5. Q UANTIFYING FUTURE BEHAVIOR WITH GVF S
An RL agent’s future discounted behavior can be quantiﬁed not only by its the value function, but other auxiliary signals, for example, by its observation occupancy measure, known as successor features (Dayan, 1993; Barreto et al., 2016).
The choice of the examined signal quantiﬁes the properties the agent will exhibit in 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 in 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 deﬁned by a bounded cumulant function c : O ⇥ A !
R d , s.t.
| c
( o, a ) |  c max for c max = sup o,a 2 O ⇥
A c ( 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 (Barreto et al., 2016), in which case the norm of c ( o, a ) is bounded.
Deﬁnition 1 (Generalized value function) Let c be any bounded function over R d , let  2 [0 , 1] and µ any policy.
The generalized value function is deﬁned as
[Equation]
for any timestep t  1 and o t 2 O .
Since, in our case, we can learn G µ for each distinct POMDP M i for the dataset D µ , we index the GVF using the POMDP index, i.e. G µi =
LearnGVF ( c, D µ , i ) (in practice, the learning is parallelized).
5. M EASURING DISTANCES BETWEEN GVF S OF DIFFERENT POMDP S
Examining the difference between future behaviors of two observations quantiﬁes the exact amount of expected behavior change between these two observations.
Using the GVF framework, we could compute the distance between o 1 2 M 1 and o 2 2 M 2 by ﬁrst estimating the latent state with z = f ( o ) using a (learned) latent state decoder f , and then evaluating the distance
[Equation]
a measure of dissimilarity that can then be used in a contrastive loss.
However, the distance between GVFs from two different POMDPs can have drastically different scales: |
[Equation]
, making point-wise comparison meaningless.
The issue is 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 statistics.
A robust distance estimate between GVF signals across POMDPs can be obtained by looking at the cumulative distribution function of G i denoted F i ( g ) =
P [ G i ( o t )  g ] for all o t 2 O .
G i is a deterministic GVF with the set of discontinuity points of measure 0, and as such F i can be understood through the induced state distribution P µt (using continuous mapping theorem from Mann and Wald (1943)).
It can be estimated from n independent and identically distributed samples of D µ as
[Equation]
and its inverse, the empirical quantile function (van der Vaart, 1998)
[Equation]
We use the empirical quantile function to partition the range of all GVFs into K quantile bins, i.e. disjoint sets with identical size where the set corresponding to quantile k is deﬁned as I i ( k ) =
{ o
M i : F  1 i ( kK )  G µi ( o )  F  1 i ( k +1 K ) }
and it’s aggregated version as I ( k ) =
[ mi =1 I i ( k ) .
Importantly, we augment the dataset D µ with observation-speciﬁc labels, which correspond to the index of the quantile bin into which the GVF G of an observation o 2 M i falls into:
[Equation]
These self-supervised labels are then used in a multiclass InfoNCE loss (Oord et al., 2018), which is a variation of metric learning with respect to the quantile distance deﬁned above (Khosla et al., 2020; Song and Ermon, 2020).
5. S ELF - SUPERVISED LEARNING OF GSF S
After augmenting the ofﬂine dataset with observation labels, we use a simple self-supervised learning procedure to minimize distance in the latent representation space between observations with identical labels.
First, the observation o is encoded using a non-linear encoder f : O !
Z with parameters into a latent state representation
z
= f ( o ) 2 .
The representation z is then passed into two separate trunks: 1) a linear matrix ✓ a which recovers the state-action value function Q ✓ ( o, a )
=
✓ > a z , and 2) a non-linear projection network h ✓ : Z !
Z with parameters ✓ h to obtain a new embedding, used for contrastive learning.
The projection h ( z ) is then used used in a multiclass InfoNCE loss (Oord et al., 2018; Song and
Ermon, 2020) where a linear classiﬁer W 2 R |Z| ⇥ K aims to correctly predict the observation labels (i.e. quantile bins k = 1 , 2 , .., K ) from h ( z ) :
[Equation]
where ⌧ > 0 is a temperature parameter.
Our empirical ﬁndings suggest that this version of the loss is more stable than other multi-class contrastive losses (see Appendix 7.3).
5. A LGORITHM
Our method relies on the approximation oracle LearnGVF , to produce GVF estimates, later used in the contrastive learning phase.
Since we are concerned with the ofﬂine RL setting we add our auxiliary loss on top of Conservative
Q-learning (CQL, Kumar et al., 2020) a strong baseline.
CQL is trained using a linear combination of Q-learning (Watkins and Dayan, 1992; Ernst et al., 2005) and behavior-regularization:
[Equation]
[Equation]
and being
Input : Dataset D ⇠ µ , initialized Q-function Q ✓ with encoder f ✓ f and action weights ✓ a , per-POMDP set of GVFs G = { G µi } i 2 , state projection network h , epoch number J , number of POMDPs m , number of quantiles K , temperature parameter ⌧ , exponential moving average parameter  1 for epoch
j = 1 , 2 , .., J do for minibatch B ⇠ D do /
* Data augmentation on observation */ o
[Table]
crop ( o for all o 2 B ; z f ( o for all o 2 B ; / *
[Equation]
CQL
[Equation]
Update ✓ a , using r ✓ a , CQL ✓ )
for POMDP M i = 1 , 2 , .., m do Estimate ˆ F  1 i of G µi from B ; for observation o 2 B \ M i do
k Update encoder and projection network
Update ✓ h , , W using r ✓ h , , W ` NCE ( ✓ h , , W ) computed with z, l ( o ) and ⌧ ; Update CQL agent’s target network with  of online parameters , ✓ ;
Alg. 2 summarizes the learning procedure for GSF as implemented on top of a CQL agent for a discrete action space.
In our experiments, all baselines use random crops as data augmentation
Connection to existing methods Our framework is able to recover objectives similar to those of prior works by carefully designing the cumulant function.
• Cross-State Self-Constraint (CSSC, Liu et al., 2020a) :
In CSSC, observations o 1 , oare considered similar if they have identical future action sequences of length K under some ﬁxed policy; a total of |A| K distinct classes are possible.
This approach can be approximated in our framework by picking c ( o t , a t ) =
a t ( a ) , 8 a 2 A .
The problem reduces to a |A| T  t -way classiﬁcation problem for observations of timestep t , which GSF approximates using K 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 d TV .
If we consider the space of Boltzmann policies ⇡ Boltzmann with respect to an POMDP-speciﬁc value function Q , then choosing c ( o t , a t ) =
r ( s t , a t ) in GSF will effectively compute the distance between unnormalized policies.
5. C HOICE OF NUMBER OF QUANTILES K
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.
1 highlights a trade-off when choosing the number of quantiles bins empirically.
Theorem 1 Let G 1 , G 2 be generalized value functions with cumulants c 1 , c 2 from respective
POMDPs M 1 , M 2 , K be the number of quantile bins, n 1 , n 2 the number of sample transitions from each POMDP.
Suppose that P [sup t =1 , 2 ,.. | c 1 ( o 1 ,t , µ ( o 1 ,t ))
c 2 ( o 2 ,t , µ ( o 2 ,t )) | > (1   ) " /  ]   .
Then, for any k = 1 , 2 , .., K and " > 0 the following holds without loss of generality:
[Equation]
where p
[Equation]
The proof can be found in the Appendix Sec. 7.2.
For POMDP M 1 , the error decreases monotonically with increasing bin number K (second term) but the variance of bin labels depends on the number of sample transitions n 1 (ﬁrst term).
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 n = 1 .
Setting K too small will group together unrelated observations, inducing high bias.
E XPERIMENTS
Unlike for single task ofﬂine RL (Fu et al., 2020), most works on zero-shot generalization from ofﬂine data either come up with an ad hoc solution suiting their needs, e.g. (Ball et al., 2021), or assess performance on benchmarks that do not evalute generalization across observation functions (e.g., Yu et al., 2020).
To accelerate progress in this ﬁeld, we devised the ofﬂine Procgen benchmark, an ofﬂine RL dataset to directly test for generalization of ofﬂine RL agents across observation functions 5 .
Ofﬂine Procgen benchmark We evaluate the proposed approach on an ofﬂine version of the
Procgen benchmark (Cobbe et al., 2020), which is widely used to evaluate zero-shot generalization across complex visual perturbations.
Given a random seed, Procgen allows to sample procedurally generated level conﬁgurations for 16 games under various complexity modes: “easy”, “hard” and “exploration”.
The dataset is obtained as follows: we ﬁrst pre-train a PPO (Schulman et al., 2017) agent for 25M timesteps on 200 levels of “easy” distribution for each environment 6 (“easy” mode is widely used to test generalization capabilities (Cobbe et al., 2020; Raileanu and Fergus, 2021; Mazoure et al., 2021)).
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).
Results We compare the zero-shot performance on the entire distribution of ”easy” POMDPs for GSF against that of strong RL and representation learning baselines: behavioral cloning (BC) - to assess the quality of the PPO policy, CQL (Kumar et al., 2020) - the current state-of-the-art on multiple ofﬂine benchmarks which balances RL and BC objectives, CURL (Srinivas et al., 2020), DeepMDP (Gelada et al., 2019) - which learns a metric closely related to bisimulation across the MDP, Value Prediction Network (VPN, Oh et al., 2017) - which combines model-free and modelbased learning of values, observations, next observations, rewards and discounts, Cross-State SelfConstraint (CSSC, Liu et al., 2020a) - which boosts similarity of observations with identical action sequences, as well as Policy Similarity Embeddings (Agarwal et al., 2020), which groups observation representations based on distance in optimal policy space.
are
Fig.
3 shows the performance of all methods over 5 random seeds and all 16 games on the ofﬂine
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 (ofﬂine RL) performance, in this case implemented by CQL.
It can be seen that GSF performs better than other ofﬂine RL and representation learning baselines.
Using different cumulants functions can lead to different label assignments and hence different similarity groups.
Fig.
3 examines the performance of GSF with respect to 3 cumulants: 1) r ( s t , a t ) , rewards s.t.
GSF learns the policy’s Q µ -value, 2) o t ( o ) , the successor representation 7 (Dayan, 1993; Barreto et al., 2016)
s.t.
GSF learns induced distribution over D µ (Machado et al., 2020) and 3) a t ( 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.
D ISCUSSION
In this work we proposed GSF, a novel algorithm which combines reinforcement learning with representation learning to improve zero-shot generalization performance on challenging, pixel-based control tasks.
GSF relies on computing the similarity between observation pairs with respect to any instantaneous accumulated signal, which leads to improved empirical performance on the newly introduced ofﬂine Procgen benchmark.
Theoretical results suggest that GSF ’s hyperparameter choice depends on a trade-off between ﬁnite sample approximation and extrapolation error.
While our work answered some questions regarding zero-shot generalization in ofﬂine RL, some questions persist: can GVF-based distances be included in a contrastive objective without the need for quantile discretization (perhaps through re-scaling or order statistics)?
Can the cumulant function be chosen a priori for a speciﬁc task structure other than Procgen, and shown to lead to optimal representations
R EPRODUCIBILITY S TATEMENT
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 ofﬂine Procgen benchmark.