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