Adversarial Style Transfer for Robust Policy Optimization in Reinforcement Learning

1 INTRODUCTION . The reinforcement learning ( RL ) environments often provide observation , which is a highdimensional projection of the true state , complicating policy learning as the deep neural network model might mistakenly correlate reward with irrelevant information . Thus deep neural networks might overfit irrelevant features in the training data due to their high flexibility and the long-stretched time duration of the RL training process , which leads to poor generalization ( Hardt et al. , 2016 ; Zhang et al. , 2018a ; Cobbe et al. , 2019 ; 2020 ; Zhang et al. , 2018b ; Machado et al. , 2018 ; Gamrian & Goldberg , 2019 ) . These irrelevant features ( e.g. , background color ) usually do not impact the reward ; thus , an optimal agent should avoid focusing on them during policy learning . Even worse , this might lead to incorrect state representations , which prevents deep RL agents from performing well even in slightly different environments . Thus , to learn a correct state representation , the agent needs to avoid those features that are not important for optimal policy learning . In recent times , generalization in RL has been explored extensively . In zero-shot generalization framework ( Zhang et al. , 2018b ; Song et al. , 2020 ; Wang et al. , 2019 ; Packer et al. , 2018 ) the agent is trained on a finite training set of MDP levels and then evaluated on the full distribution of the MDP levels . The distribution of the MDP levels can be generated in various ways , including changing dynamics , exploit determinism , and adding irrelevant confounders to high-dimensional observation . An agent trained in these setups can overfit any of these factors , resulting in poor generalization performance . Many approaches have been proposed to address this challenges including various data augmentation approaches such as random cropping , adding jitter in image-based observation ( Cobbe et al. , 2019 ; Laskin et al. , 2020a ; Raileanu et al. , 2020 ; Kostrikov et al. , 2020 ; Laskin et al. , 2020b ) , random noise injection ( Igl et al. , 2019 ) , network randomization ( Osband et al. , 2018 ; Burda et al. , 2018 ; Lee et al. , 2020 ) , and regularization ( Cobbe et al. , 2019 ; Kostrikov et al. , 2020 ; Igl et al. , 2019 ; Wang et al. , 2020 ) have shown to improve generalization . The common theme of these approaches is to increase diversity in the training data so as the learned policy would better generalize . However , this perturbation is primarily done in isolation of the RL reward function , which might change an essential aspect of the observation , resulting in sub-optimal policy learning . Moreover , the random perturbation in various manipulations of observations such as cropping , blocking , or combining two random images from different environment levels might result in unrealistic observations that the agent will less likely observe during testing . Thus these techniques might work poorly in the setup where agents depend on realistic observation for policy learning . It is also de- sirable to train the agent with realistic observations , which helps it understand the environments ’ semantics . Otherwise , the agent might learn unexpected and unsafe behaviors while entirely focusing on maximizing rewards even by exploiting flaws in environments such as imperfect reward design . This paper focuses on a particular form of generalization where the agent gets to observe a highdimensional projection of the true state . In contrast , the latent state dynamics remain the same across all the MDP levels . This assumption has been shown to affect generalization ( Song et al. , 2020 ) in well studied benchmarks ( Nichol et al. , 2018 ) . It has been observed that the agent overfit the scoreboard or timer and sometimes achieves the best training performance even without looking at the other part of the observation . Consequently , this policy completely fails to generalize in test environments . The high-capacity model even can memorize the whole training environment , thus severely affecting the generalization ( Zhang et al. , 2018b ) . This paper proposes a style transfer-based observation translation method that considers the observation ’ s content and reward signal . In particular , a generator performs the style transfer while the policy tries to be robust toward such style changes . Note that , style transfer-based approach has been used to target generalization in Anonymous ( 2021 ) where the authors propose to train style transfer and reinforcement learning in isolation . In contrast , in this paper , we propose a unified learning approach where the style transfer is integrated into the reinforcement learning training . Thus the generator produces more effective translation targeted toward RL objective , making the RL policy more robust . The trajectory data from the agent ’ s replay buffer is first clustered into different categories , and then observation is translated from one cluster ’ s style to another cluster ’ s style . Here the style is determined by the attribute commonality of observation features in a cluster . The agent should be robust toward changes of such features . Moreover , the translated trajectories correspond to those that possibly appear in testing environments , assisting the agent in adapting to unseen scenarios . Figure 1 shows an overview of our method . Our approach consists of a max-min game theoretic objective where a generator network modifies the observation by style transferring it to maximize the agent ’ s probability of taking a different action . In contrast , the policy network ( agent ) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward . Based on this objective , we propose a practical deep reinforcement learning algorithm , Adversarial Robust Policy Optimization ( ARPO ) , that aims to find a robust policy by mitigating the observational overfitting and eventually generalizes to test environments . We empirically show the effectiveness of our ARPO agent in generalization and sample efficiency on challenging Procgen ( Cobbe et al. , 2020 ) and Distracting Control Suite ( Stone et al. , 2021 ) benchmark with image-based observation . The Procgen leverages procedural generation to create diverse environments to test the RL agents ’ generalization capacity . The distracting control adds various visual distractions on top of the Deepmind Control suite ( Tassa et al. , 2020 ) to facilitate the evaluation of generalization performance . Empirically , we observed that our agent ARPO performs better in generalization and sample efficiency than the baseline PPO ( Schulman et al. , 2017 ) and a data augmentation-based RAD ( Laskin et al. , 2020a ) algorithm on the many Procgen environments . Furthermore , APRO performs better compared to PPO and SAC ( Haarnoja et al. , 2018 ) on the evaluated two Distracting Control Suite environments in various setups . In summary , our contributions are listed as follows : • We propose Adversarial Robust Policy Optimization ( ARPO ) , a deep reinforcement learning algorithm to find a robust policy that generalizes to test environments in a zero-shot generalization framework . • We evaluate our approach on challenging Procgen ( Cobbe et al. , 2020 ) and Distracting Control Suite ( Stone et al. , 2021 ) benchmark in generalization ans sample efficiency settings . • Empirically , we observed that our agent ARPO performs better generalization and sample efficiency performance compared to standard PPO ( Schulman et al. , 2017 ) , SAC ( Haarnoja et al. , 2018 ) , and a data augmentation-based RAD ( Laskin et al. , 2020a ) approach in various settings . 2 PRELIMINARIES AND PROBLEM SETTINGS . Markov Decision Process ( MDP ) An MDP is denoted by M = ( S , A , P , r ) where at every timestep t , from an state st ∈ S , the agent takes an action at from a set of actionsA . The agent then receives a reward rt and the environment and move to a new state st+1 ∈ S based on the transition probability P ( st+1|st , at ) . Reinforcement Learning . Reinforcement learning aims to learn a policy π ∈ Π that maps state to actions . The policy ’ s objective is to maximize cumulative reward in an MDP , where Π is the set of all possible policies . Thus , the policy which achieves the highest possible cumulative reward is the optimal policy π∗ ∈ Π. Generalization in Reinforcement Learning . In the zero-shot generalization framework ( Song et al. , 2020 ) , we assume the existence of a distribution of levels D over an MDP and a fixed optimal policy π∗ which can achieve maximal return over levels generated from the distribution D. The levels may differ in observational variety , such as different background colors . In this setup , the agent has access to a fixed set of MDP levels during training the policy . The trained agent is then tested on the unseen levels to measure the generalization performance of the agent . Style Transfer . The task of image-to-image translation is to change a particular aspect of a given image to another , such as red background to green background . Generative adversarial network ( GAN ) -based method has achieved tremendous success in this task ( Kim et al. , 2017 ; Isola et al. , 2017 ; Zhu et al. , 2017 ; Choi et al. , 2018 ) . Given training data from two domains , these models learn to style-translate images from one domain to another . The domain is defined as a set of images sharing the same attribute value , such as similar background color and texture . The shared attributes in a domain are considered as the “ style ” of that domain . Many translation methods require paired images from two domains , which is not feasible in the reinforcement learning setup . The agent collects the data during policy learning , and annotating them in pairs is not feasible . Thus , we leverage unpaired image to image translation , which does not require a one-to-one mapping of annotated images from two domains . In particular , we build upon the work of StarGAN ( Choi et al. , 2018 ) which is capable of learning mappings among multiple domains efficiently . The model takes in training data of multiple domains and learns the mappings between each pair of available domains using only a single generator . This method automatically captures special characteristics of one image domain and figures out how these characteristics could be translated into the other image collection , making it appealing in the reinforcement learning setup . However , the RL agent generates experience trajectory data which is not separated into domains . Thus , we further apply a clustering approach which first clusters the trajectory observation into domains , and then we train the generator , which style-translates among those domains . 3 ADVERSARIAL ROBUST POLICY OPTIMIZATION ( ARPO ) . Our approach consists of a max-min game theoretic objective where a generator network modifies the observation by style transferring it to maximize the agent ’ s probability of taking a different action . In contrast , the policy network ( agent ) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward . Based on this objective , we propose a practical deep reinforcement learning algorithm , Adversarial Robust Policy Optimization ( ARPO ) , to find a robust policy that generalizes to test environments . We now discuss the details of the policy network and generator .

Summary: The paper propose a method for improving generalisation and robustness across visual factors of variation in reinforcement learning, using a style transfer network to adversarially perturb the input to the policy, while the policy is trained to be invariant to this visual perturbation. The style transfer network is trained using StarGAN, based on domains created through a Gaussian Mixture Model clustering of the observations from the policy's experience. The method is demonstrated on OpenAI Procgen easy mode, and compared against PPO and RAD. The method does better than both of these baselines on 8 of the 16 games. Ablations on a single game show that the method is fairly robust to different choices of hyperparameters (number of clusters for the GMM, and weight of the adversarial loss for the style transfer network and policy).