Cross-State Self-Constraint for Feature Generalization in Deep Reinforcement Learning

1 INTRODUCTION . Deep Reinforcement learning has achieved tremendous success on mastering video games ( Mnih et al. , 2015 ) and the game of GO ( Silver et al. , 2017 ) . While training agent by using deep reinforcement learning algorithms , we usually assume that the agent could extract appropriate and effective features from different states and take actions accordingly . However , as more and more research works ( Zhang et al . ( 2018 ) , Song et al . ( 2019 ) , Dabney et al . ( 2020 ) ) have pointed out , even welltrained RL agents that learns from visualized input tend to memorizing spurious pattern rather than understanding the essential generic features of a given state . For example , an agent might pay more attention to the change of irrelevant background rather than noticing the obstacles or enemies ( Song et al. , 2019 ) . To improve generalization in the new environment , various kinds of regularization method like dropout ( Farebrother et al. , 2018 ) and data augmentation ( Laskin et al. , 2020 ) has been proposed and tested in combination with reinforcement learning . Conventional methods like dropout and batch-norm has been proven to be effective in supervised-learning , and for self-supervised learning like RL we see multiple related applications across various environments . Data augmentation like random crop ( Laskin et al. , 2020 ) or random convolution ( Lee et al. , 2019 ) have also been proposed recently and provide considerable generalization enhancement to the unseen levels of various tested environment ( Tassa et al . ( 2018 ) , Cobbe et al . ( 2018 ) , Cobbe et al . ( 2020 ) ) . The agent is acting on multiple augmented views of the same input and learn from these prior injected data . However , modifying state information ( injecting prior to the data ) may be risky or even detrimental for representation learning because vital features may be altered or lost ( ex : flipping state image might change the corresponding behavioral meaning , cropping the input image might lose critical features like the enemy position in the game ) . To avoid losing informative features of the visualized input , we choose a different approach . As a human learner , we rarely depend on multiple augmented views of the same input to discriminate important or fictitious features . Instead , human learners try to recognize general patterns across multiple states and act accordingly . In other words , if the same action ( or behavior ) has been conducted by a well-trained agent in two different states , we would infer that the agent has conceived similar feature patterns in these states . For example , if one car stops for ninety seconds at two different intersections , we would guess that the car might be stopped by the red light at both places ( Figure 1b ) . From here we get the intuition about the relation between action ( or behavior ) and representation feature space learned by the agent : for a RL agent acting rationally across sequence of states , its behavior would be similar if it perceives similar critical patterns while acting differently if it perceives different patterns . Based on this intuition , we designed a novel constraint that performs regularization directly in the representation feature space of the learning agent . Our hypothesis for this constraint is simple : states in which the rational agent behaves identically should share more representation resemblance than states the agent behave differently . We test this novel constraint in combination with Rainbow ( Hessel et al. , 2017 ) , the state-of-theart Q-learning method that combines numerous improvements with the original DQN model ( Mnih et al. , 2015 ) . Inspired by the pair-wise structure used in BPR-opt ( Rendle et al. , 2012 ) , we design the self-constraint based on the agent ’ s behavior and utilize implicit feedback between positive and negative state pairs in the replay buffer . One thing worth noticing is that no change is needed for the underlying RL algorithm or model to adapt the proposed method , and it can be easily applied to other models with minimal effort ( Figure 1a ) . To measure the improvement on generalization to unseen environment levels , We test CSSC across 16 games on the OpenAI ProcGen benchmark ( Cobbe et al. , 2020 ) . We see significant improvement on generalization capability for most ProcGen games in comparison with base Rainbow model . We highlight the main contributions of CSSC as follows : • Directly optimize feature generalization across various input states • Requires no additional modification on the input data and the base model • Enhance generalization and data efficiency on existing RL methods 2 RELATED WORK . 2.1 AUXILIARY TASK FOR REPRESENTATION LEARNING IN RL . Representation Learning has been a vital part for deep RL algorithms . Even though the main goal of RL is to find the optimal value function , Mccallum & Ballard ( 1996 ) and Li et al . ( 2006 ) show that a representation specialized to this function may not be suitable for the sequence of value functions leading to it . On the other hand , Dabney et al . ( 2020 ) argue that the path toward optimal value function might be hindered by overfitting the representation to any intermediate value function during training . One popular and effective method to address this challenge is by using auxiliary tasks ( Jaderberg et al . ( 2016 ) , Bellemare et al . ( 2019 ) ) . Dabney et al . ( 2020 ) also propose new auxiliary tasks by considering the value-improvement path holistically . 2.2 REGULARIZATION ON RL . 2.2.1 ADDING STOCHASTICITY TO RL . Conventional practices like stochastic policy ( Hausknecht & Stone , 2015 ) , random starts ( Mnih et al. , 2015 ) , sticky actions ( Hausknecht & Stone , 2015 ) and frame skipping are wildly used for popular tasks like Atari ( Machado et al. , 2017 ) . By adding stochasticity to the environments during training and testing , we would prevent simple algorithms like trajectory tree ( Kearns et al. , 1999 ) and brute ( Machado et al. , 2017 ) from optimizing over open-loop sequences of actions without even considering the input states . However , as Zhang et al . ( 2018 ) has pointed out , injecting stochasticity to the maze-like environments does not necessarily prevent the deep RL agents from overfitting . 2.2.2 CONVENTIONAL REGULATIZATION ON RL . For deep supervised learning , regularizations has long been a major issue as generalization from training to test set is the major concern . Common regularization practices like L2 regularization , dropout , batch-normalization from supervised learning has also been tested in the reinforcement learning scenario . Cobbe et al . ( 2018 ) and Farebrother et al . ( 2018 ) show that dropout and L2 regularization improve generalization to test environments in Coinrun ( Cobbe et al. , 2018 ) and Atari 2600 ( Bellemare et al. , 2013 ) respectively . 2.2.3 DATA AUGMENTATION ON RL . For supervised learning related to images or computer vision , data augmentation is a popular way to enhance performance by injecting useful priors on the training data . In recent years , we see numerous data augmentation methods for RL that improve generalization or data efficiency . Lee et al . ( 2019 ) introduce a randomized convolution network that randomly perturbs input features . Srinivas et al . ( 2020 ) utilize contrastive learning and show considerable performance gains in terms of sample-efficiency . Kostrikov et al . ( 2020 ) uses augmented data and weighted Q-functions to achieve state-of-the-art data-efficiency on the DMControl ( Tassa et al. , 2018 ) . Laskin et al . ( 2020 ) investigate ten different data augmentation on RL and point out that random crop is the most effective on DMControl ( Tassa et al. , 2018 ) 3 BACKGROUND . CSSC is a general framework of cross-state representation regularization for RL . In principle , one can apply CSSC to other variation of DQN or policy-based models for discrete action-space environments . In this work we pick Rainbow ( Hessel et al. , 2017 ) and PPO ( Schulman et al. , 2017 ) as our base model to show that CSSC is compatible with the original Nature DQN ( Mnih et al. , 2015 ) , its multiple improvements and policy-based algorithms . In the following subsections we will review Rainbow DQN and introduce the concept of implicit feedback in Bayesian Personalized Ranking ( BPR ) proposed by Rendle et al . ( 2012 ) . 3.1 RAINBOW . In combination with a convolutional neural network as visualized input encoder , Deep Q network ( Mnih et al. , 2015 ) demonstrates that it is possible to use neural network as a function approximator that maps raw pixels to Q values . Since then , multiple improvements such as Double Q Learning ( van Hasselt et al. , 2015 ) , Dueling Network Architecture ( Wang et al. , 2015 ) , Prioritized Experience Replay ( Schaul et al. , 2015 ) and Noisy Networks ( Fortunato et al. , 2017 ) have been proposed . In addition , Bellemare et al . ( 2017 ) proposed the method of predicting a distribution over possible value support through the C51 algorithm . By combining all the above techniques into a single off-policy algorithm , Rainbow DQN showcases the state-of-the-art sample-efficiency on Atari Benchmarks . The resulting loss function for Rainbow DQN is as follows : LRainbow = DKL ( Φzd ( n ) t ||dt ) ( 1 ) where Φz is the projection onto fixed support z. d ( n ) t and dt are the n-step target return distribution and the model-predicted return distribution respectively . 3.2 BPR-OPT . As opposed to explicit feedback like user ratings , implicit feedback in recommendation systems focuses on interactions like click or view between users and items . In implicit feedback system only positive observations are available , and non-observed user-item pairs -e.g . a user has not view an item yet - are considered as negative observations . Instead of rating prediction , BPR-opt ( Rendle et al. , 2012 ) is designed for direct ranking optimization based on the implicit feedback between users and items . BPR extends the user preference from observed interaction pairs to non-observed data by ranking the preference of positive observation and negatived observation across the training data . The formulation of the maximum posterior estimator for the personalized ranking optimization ( BPR-opt ) can be written as follows : ln p ( Θ| > u ) = ∑ ( u , i , j ) ∈Ds lnσ ( x̂uij ) − λΘ||Θ||2 ( 2 ) where Θ represents the parameter vector of the base model . Ds represent the batch samples from training data . p ( Θ| > u ) is the posterior probability conditioned on latent preference structure > u for user u , λΘ is the model specific regularization parameter , σ is the sigmoid function , and x̂uij ( Θ ) is the real-valued function model Θ which captures the preference relationship between user u , item i and item j. Rendle et al . ( 2012 ) use x̂uij : = x̂ui − x̂uj to indicate that the user u prefers item i over item j . 4 METHOD . The main idea of CSSC is to extend the similarity ranking across various state representation based on our hypothesis : representations motivating identical behavior should share more similarity than those motivating different behavior . We name this method as ” self-constraint ” to emphasize that the state pairing is decided by the agent ’ s behavior rather than pre-defined fixed labels . In the following subsection below , we introduce the following concepts to provide further explanation and insight : ( i ) definition of the behavior of an agent at a given state ( ii ) implicit feedback between state and behavior ( iii ) cross-state self-constraint as an auxiliary loss for representation regularization

Summary: This paper proposes a new method for learning representations of images with the goal of improving generalization in RL. The key idea is to regularize the learned feature space by forcing embeddings of states followed by the same action (or sequence of actions) to be more similar than embeddings of states followed by different actions. The method is evaluated on the Procgen and achieves superior performance relative to Rainbow, a standard RL algorithm.