Disentangling Generalization in Reinforcement Learning

Generalization in Reinforcement Learning ( RL ) is usually measured according to concepts from supervised learning . Unlike a supervised learning model however , an RL agent must generalize across states , actions and observations from limited reward-based feedback . We propose to measure an RL agent ’ s capacity to generalize by evaluating it in a contextual decision process that combines a tabular environment with observations from a supervised learning dataset . The resulting environment , while simple , necessitates function approximation for state abstraction and provides ground-truth labels for optimal policies and value functions . The ground truth labels provided by our environment enable us to characterize generalization in RL across different axes : state-space , observation-space and action-space . Putting this method to work , we combine the MNIST dataset with various gridworld environments to rigorously evaluate generalization of DQN and QR-DQN in state , observation and action spaces for both online and offline learning . Contrary to previous reports about common regularization methods , we find that dropout does not improve observation generalization . We find , however , that dropout improves action generalization . Our results also corroborate recent findings that QR-DQN is able to generalize to new observations better than DQN in the offline setting . This success does not extend to state generalization , where DQN is able to generalize better than QR-DQN . These findings demonstrate the need for careful consideration of generalization in RL , and we hope that this line of research will continue to shed light on generalization claims in the literature . 1 Generalization in Reinforcement Learning . A Reinforcement Learning ( RL ) agent perpetually finds itself in novel states of an environment . To act intelligently , the agent must generalize its previous experience to new situations . Function approximation helps distill this previous experience into the agent ’ s learnable parameters , which allows previous experience to be leveraged with new state inputs ( Boyan and Moore , 1994 ; Sutton , 1995 ) . While there is a growing literature of new methods for improving generalization of deep RL algorithms , principled and quantitative methods for evaluating generalization remain lacking . This is due in part to the complexity of the MDP problem formulation and the difficulty of disentangling generalization from the performance of RL algorithms in terms of achieving higher expected cumulative reward , i.e . return . One common notion of generalization in RL evaluates an agent ’ s capabilities by checking if it can achieve similar performance in an environment that is similar to , but not exactly the same as , the environment in which it was trained ( Whiteson et al. , 2011 ) . This can be accomplished through randomization , i.e . by randomizing the parameters underlying the environment , such as wind velocity in helicopter hovering or the mass of objects in a simulator , thereby changing the transition dynamics ( Peng et al. , 2017 ) . Generalization in this sense draws parallels to supervised learning , where classifiers are often trained on a fixed dataset , and evaluated on a separate testing dataset under the I.I.D . assumption . The agent is said to generalize well if the difference between training error and testing error is small . While these supervised learning concepts are relevant to RL , there are two problems with taking a cross-environment approach to evaluating RL generalization . First , this paradigm does not disentangle the various aspects of generalization that are required for an RL agent to succeed in its task , whether it be the value estimates or the policy , and whether they are robust to variations in states , observations or actions . This is the question of “ what ” performance criterion should be measured when we discuss generalization . In RL , we have function approximators for many quantities : state-transition , reward , state-value , actionvalue and policy . While generalization of state-value is similar to regression , generalization with quantities related to action , such as policy or action-value , do not have supervised learning analoguesand hence require : . : : : : This : : is : : : : : : : because : : : : : : : policies : : : : and : : : : : : : : : : : : action-values : : : : have : : as : : : : : : many : : : : : : : outputs : : as : : : : : : : : actions , : : : : and : : : : only : : : : the : : : : : : action : : : : : taken : : : by : : : : the : : : : : agent : : is : : : : : : : : : updated , : : : : : : which : : : : : : : : : : necessitates separate consideration . Second , the paradigm follows practice in supervised learning to impose a strict separation of test and train environments . This practice effectively focuses on transfer performance across environments but fails to evaluate the agent ’ s ability to generalize within a single environment . I.e . it fails to answer the question of “ how ” to measure generalization under a performance criterion . Unlike supervised learning , the agent ’ s state distribution changes during learning because the policy changes . Even where we have randomly generated training and testing sets of environments , the environments ’ complex dynamics do not admit ground-truth labels to determine optimal actions or values for comparison and evaluation . These environments only allow an agent ’ s generalization capabilities to be measured in terms of Monte-Carlo rollouts of the learned policy , leaving us restricted to the single performance criterion of return and preventing us from specifically measuring important nuanced differences in the agent ’ s ability to generalize . In addition , because the upper bound on return is unknown to the researcher , informed judgments about the quality of the policy is very difficult to make . Within the RL-oriented literature on generalization , there are two distinct categories of research . The first proposes environments and methodologies for measuring RL generalization . Examples of this approach are ProcGen/CoinRun ( Cobbe et al. , 2020 ; ? ) : : : : : : : : : : : : : : : : : : : : : : : ( Cobbe et al. , 2020 ; 2019 ) , randomized-reward CartPole ( Zhang et al. , 2018a ) , the GridWorld maze ( Zhang et al. , 2018b ) , observation projection ( ? ) : : : : : : : : : : : : : : : : : ( Song et al. , 2020 ) and the hierarchy of state generalization ( Witty et al. , 2018 ) . These efforts are aimed at the second issue raised above , i.e . “ how ” generalization should be measured in RL . Underlying these approaches is specifying how to split the environment into testing and training scenarios . Early work by Zhang et al . ( 2018a ) proposes using separate seeds . This , however , does not ensure that the states encountered by the agent are truly separate . Other works , such as those by Cobbe et al . ( 2020 ) ; ? ) ; Elsayed et al . ( 2020 ) : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : Cobbe et al . ( 2020 ; 2019 ) ; Elsayed et al . ( 2020 ) , procedurally generate separate environments for testing and training . When generalization is measured on truly separate testing and training environments , we are able to determine whether an agent ’ s policy is generalizing from one environment to another . Again , this formulation does not allow us to study generalization within a single environment , nor does it allow us to measure generalization of the value functions . The second category proposes or investigates RL methods that improve generalization . These include regularization experiments in Atari ( Farebrother et al. , 2018 ) and continuous control ( ? ) : : : : : : : : : : : : : : : ( Liu et al. , 2021 ) , contrastive similarity embeddings ( Agarwal et al. , 2021 ) and bisimulation metrics ( ? ) : : : : : : : : : : : : : : : : : ( Zhang et al. , 2021 ) . There is also the hypothesis that better , and hence more generalizable , representations arise from auxiliary tasks ( ? ) : : : : : : : : : : : : : : : : : : : : ( Jaderberg et al. , 2017 ) , which is also suspected to be the reason for the success of distributional RL ( ? Bellemare et al. , 2017 ) : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : : ( Dabney et al. , 2018 ; Bellemare et al. , 2017 ) and has recently been investigated in the offline setting ( ? ) . : : : : : : : : : : : : : : : : : : : ( Agarwal et al. , 2020 ) : . : Previous work closest to ours is using Contextual Decision Processes ( CDPs ) as : : To : : : : : : : : : : understand : : : : : : : : : : : : generalization : : : in : : : RL , : : : we : : : : use : : : the : : : : : : : : : : Contextual : : : : : : : : Decision : : : : : : : Process : : : : : : : ( CDP ) : : : : : : : : : : framework , : : : : : which : : is : a problem class for theoretical analysis of RL algorithms that use function approximation ( Du et al. , 2019 ; Jiang et al. , 2017 ; Dann et al. , 2018 ) . The CDP problem formulation renders the states unobservable , but allows the agent to view observations that contain enough information to recover the state . This formulation enables the study of function approximation as applied to RL and has significant implications for RL in general . For example , it helps to : : : : : : : : : : : : : : : Du et al . ( 2019 ) show that there exist algorithms with exponentially more efficient explo- ration than Q-learning ( Du et al. , 2019 ) , a result suggesting that RL algorithms should be designed to take advantage of function approximation and its generalization abilities , rather than naively extending tabular algorithms with function approximation . However , no existing work has leveraged CDPs empirically , despite that it connects supervised learning with RL . The MNIST gridworld : : 2D : : : : : : : MNIST : : : : : : maze environment used by ? : : : : : : : : : : : : : : : Lee et al . ( 2019 ) may be considered a CDP , : : : : : simple : : : : : : CDP , : : : : : where : : : : the : : : : : : : : : : : observations : : : : are : : : : : : : : : : : : deterministic : : : : and : : : : : : : : : equivalent : : to : : : : : state , but is not recognized as such . The work by ? : : : : : : : : : : : : : : : : Song et al . ( 2020 ) proposed projecting the state of simple control environment and varying this projection between training and testing sets . These two works have similar goals to ours . However , the first work does not investigate generalization and neither makes use of the ground-truth values to probe generalization rigorously , such as in the label-corruption experiment by ( Zhang et al. , 2017 ) that we will extend to RL . : : : : : Our : : : : work : : is : : : : : : : : uniquely : : : : : : : : : : analogous : : to : : : : : : : : : : : : : generalization : : in : : : : : : : : : : supervised : : : : : : : : learning , : : : : : : : : : : : : : simultaneously : : : : : : : : : answering : : : : : : “ how ” : : : : : : : : : : : : generalization : : : : : : should : : : be : : : : : : : : : measured : : in : : : RL : : : : and : : : : : : “ what ” : : : : : : should : : be : : : : : : : : : measured . : : : In : : : : : : : : : answering : : : : : : theses : : : : two : : : : : : : : : questions , : : : we : : : : : : : : : : : disentangle : : : : : : : : : : : : generalization : : : : : : across : : : : : three : : : : axes : : : : : : : : states , : : : : : : : : : : : : observations : : : : and : : : : : : : : actions . : : : : : : : : Finally , : : : : this : : : : : : : : : : : : disentangled : : : : : : : : : : perspective : : : : : : shows : : : : : how : : : : : : : different : : : : : : : : : : : : : generalization : : : : : : : : : : : mechanisms : : : : : : : benefit : : : the : : : : : : : : different : : : : axes : : : of : : : : : : : : : : : : generalization .

Summary: This paper proposes an empirical evaluation method to measure of the generalization capacity of an RL agent. It relies on CDPs combining a tabular environment with a supervised learning dataset. Generalization is measured across three axis: state space, observation space and action space. The empirical evaluation is led on DQN and QR-DQN on the four room domain and corridor domain combined with the MNIST dataset in the online and offline settings. The authors find that dropout improves action generalization but not observation generalization while regularisation improves observation generalization. They also find that QR-DQN generalise better than DQN in the offline setting on the observation axis and action axis but not on the state space axis.