Improving zero-shot generalization in offline reinforcement learning using generalized similarity functions

1 INTRODUCTION . 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 fields , 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 finite 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 specific 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 . offline RL ) is challenging , since the representation will have a large estimation error over regions of low data coverage . A common solution to mitigate this task-specific overfitting and extracting the most information out of the data consists in introducing auxiliary learning signals other than instantaneous reward ( 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 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 , 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 offline RL research . Improving zero-shot generalization performance from an offline 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 offline RL setting , i.e . by training agents on a fixed dataset . We hypothesize that in order for an RL agent to be able to generalize across 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 : 1 . 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 . 2 . We devise a new benchmark constructed to test zero-shot generalization of offline RL algorithms : 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 on offline Procgen , and show that GSF outperforms both previous state-of-the-art offline RL and representation learning baselines on 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 . 2 RELATED WORKS . 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 ; Triantafillou et al. , 2019 ; Valle-Pérez 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 : 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 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 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 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 find state abstractions better suited for control . The problem of learning meaningful state representations ( or abstractions ) for planning and control has been extensively studied previously ( 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 specifically self-supervised learning , has also been used to achieve state-of-the-art generalization and sample efficiency results in RL on challenging control problems such as data efficient 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 ) . Offline reinforcement learning When learning from a static dataset , agents should balance interpolation and extrapolation errors , while ensuring proper diversity of actions ( i.e . prevent collapse to most frequent action in the data ) . Popular offline RL algorithms such as BCQ ( 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 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 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 . 3 PROBLEM SETTING . 3.1 PARTIALLY-OBSERVABLE MARKOV DECISION PROCESSES . A ( infinite-horizon ) partially-observable Markov decision process ( POMDP , Murphy , 2000 ) M is defined by the tuple M = hS , p0 , A , pS , O , pO , r , i , where S is a state space , p0 = P [ s0 ] is the starting state distribution , A is an action space , pS = P [ ·|st , at ] : S ⇥ A ! ( S ) is a transition function , O is an observation space , pO = P [ ·|st ] : S ! ( O ) 1 is an observation function , r : S ⇥ A ! [ rmin , rmax ] is a reward function and 2 [ 0 , 1 ) is a discount factor . The system starts in one of the initial states s0 ⇠ p0 with observation o0 ⇠ pO ( ·|s0 ) . At every timestep t = 1 , 2 , 3 , .. , the agent , parameterized by a policy ⇡ : O ! ( A ) , samples an action at ⇠ ⇡ ( ·|ot ) . The environment transitions into a next state st+1 ⇠ pS ( ·|st , at ) and emits a reward rt = r ( st , at ) along with a next observation ot+1 ⇠ pO ( ·|st+1 ) . The goal of an RL agent is to maximize the cumulative rewards P 1 t=0 trt obtained over the entire episode . Value-based off-policy RL algorithms achieve this by estimating the state-action value function under a target policy ⇡ : Q⇡ ( st , at ) = EP⇡t [ 1X k=1 kr ( st+k , at+k ) |st , at ] , 8st 2 S , at 2 A ( 1 ) where P⇡t denotes the joint distribution of { st+k , at+k } 1k=1 obtained by executing ⇡ in the environment . 1 ( X ) denotes the entire set of distributions over the space X . 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 { oi , t , ai , t , ri , t } i=N , t=Ti=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 Mtrain = { Mi } mi=1 , and its performance is evaluated on the set of test POMDPs Mtest . We assume that both Mtrain and Mtest were sampled from a common task distribution and that every POMDP Mi 2 M = Mtrain [ Mtest shares the same transition dynamics and reward function with M but has a different observation function pi , 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 ot and corresponding states st throughout the paper .

Summary: The paper addresses the problem of generalization and representation learning in reinforcement learning. Concretely, the authors propose an approach of self-supervision for improving offline-RL based on similarity learning over the considered state space. In this context, the concept of generalized value function is introduced. The similarity measure is defined as a classification problem over quantized latent space. The paper is evaluated in an offline setting of Procgen which is a dataset for multi-task reinforcement learning generalization evaluation. More particularly, the proposed approach claims to be particularly efficient in improving zero-shot generalization performance on one offline RL benchmark, offline Procgen. The offline Procgen dataset is a second contribution of the paper.