Original: A DVERSARIAL S TYLE T RANSFER FOR R OBUST P OLICY O PTIMIZATION IN R EINFORCEMENT L EARNING
Revised: A DVERSARIAL S TYLE T RANSFER FOR R OBUST P OLICY O PTIMIZATION IN R EINFORCEMENT L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: This paper proposes an algorithm that aims to improve generalization for reinforcement learning agents by removing overﬁtting to confounding features.
Revised: This paper proposes an algorithm that aims to improve generalization for reinforcement learning agents by removing overﬁtting to confounding features.

Original: Our approach consists of a max-min game theoretic objective.
Revised: Our approach consists of a max-min game theoretic objective.

Original: A generator transfers the style of observation during reinforcement learning.
Revised: A generator transfers the style of observation during reinforcement learning.

Original: An additional goal of the generator is to perturb the observation, which maximizes the agent’s probability of taking a different action.
Revised: An additional goal of the generator is to perturb the observation, which maximizes the agent’s probability of taking a different action.

Original: In contrast, a policy network updates its parameters to minimize the effect of such perturbations, thus staying robust while maximizing the expected future reward.
Revised: In contrast, a policy network updates its parameters to minimize the effect of such perturbations, thus staying robust while maximizing the expected future reward.

Original: Based on this setup, we propose a practical deep reinforcement learning algorithm, Adversarial Robust Policy Optimization (ARPO), to ﬁnd an optimal policy that generalizes to unseen environments.
Revised: Based on this setup, we propose a practical deep reinforcement learning algorithm, Adversarial Robust Policy Optimization (ARPO), to ﬁnd a robust policy that generalizes to unseen environments.

Original: We evaluate our approach on Procgen and Distracting Control Suite for generalization and sample efﬁciency.
Revised: We evaluate our approach on Procgen and Distracting Control Suite for generalization and sample efﬁciency.

Original: Empirically, ARPO shows improved performance compared to a few baseline algorithms, including data augmentation.
Revised: Empirically, ARPO shows improved performance compared to a few baseline algorithms, including data augmentation.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

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

Original: Thus deep neural networks might overﬁt irrelevant features in the training data due to their high ﬂexibility and the long-stretched time duration of the RL training process, which leads to poor generalization (Hardt et al., 2016;
Revised: Thus deep neural networks might overﬁt irrelevant features in the training data due to their high ﬂexibility and the long-stretched time duration of the RL training process, which leads to poor generalization (Hardt et al., 2016;

Original: Zhang et al., 2018a; Cobbe et al., 2019; 2020; Zhang et al., 2018b; Machado et al., 2018; Gamrian & Goldberg, 2019).
Revised: Zhang et al., 2018a; Cobbe et al., 2019; 2020; Zhang et al., 2018b; Machado et al., 2018; Gamrian & Goldberg, 2019).

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

Original: Even worse, this might lead to incorrect state representations, which prevents deep RL agents from performing well even in slightly different environments.
Revised: Even worse, this might lead to incorrect state representations, which prevents deep RL agents from performing well even in slightly different environments.

Original: Thus, to learn a correct state representation, the agent needs to avoid those features that are not important for optimal policy learning.
Revised: Thus, to learn a correct state representation, the agent needs to avoid those features that are not important for optimal policy learning.

Original: In recent times, generalization in RL has been explored extensively.
Revised: In recent times, generalization in RL has been explored extensively.

Original: In zero-shot generalization framework (Zhang et al., 2018b; Song et al., 2020;
Revised: In zero-shot generalization framework (Zhang et al., 2018b; Song et al., 2020;

Original: Wang et al., 2019; Packer et al., 2018)
Revised: Wang et al., 2019; Packer et al., 2018)

Original: the agent is trained on a ﬁnite training set of MDP levels and then evaluated on the full distribution of the MDP levels.
Revised: the agent is trained on a ﬁnite training set of MDP levels and then evaluated on the full distribution of the MDP levels.

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

Original: An agent trained in these setups can overﬁt any of these factors, resulting in poor generalization performance.
Revised: An agent trained in these setups can overﬁt any of these factors, resulting in poor generalization performance.

Original: 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;
Revised: 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;

Original: Laskin et al., 2020a; Raileanu et al., 2020;
Revised: Laskin et al., 2020a; Raileanu et al., 2020;

Original: 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;
Revised: 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;

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

Original: The common theme of these approaches is to increase diversity in the training data so as the learned policy would better generalize.
Revised: The common theme of these approaches is to increase diversity in the training data so as the learned policy would better generalize.

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

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

Original: Thus these techniques might work poorly in the setup where agents depend on realistic observation for policy learning.
Revised: Thus these techniques might work poorly in the setup where agents depend on realistic observation for policy learning.

Original: It is also de- sirable to train the agent with realistic observations, which helps it understand the environments’ semantics.
Revised: It is also de- sirable to train the agent with realistic observations, which helps it understand the environments’ semantics.

Original: Otherwise, the agent might learn unexpected and unsafe behaviors while entirely focusing on maximizing rewards even by exploiting ﬂaws in environments such as imperfect reward design.
Revised: Otherwise, the agent might learn unexpected and unsafe behaviors while entirely focusing on maximizing rewards even by exploiting ﬂaws in environments such as imperfect reward design.

Original: This paper focuses on a particular form of generalization where the agent gets to observe a highdimensional projection of the true state.
Revised: This paper focuses on a particular form of generalization where the agent gets to observe a highdimensional projection of the true state.

Original: In contrast, the latent state dynamics remain the same across all the MDP levels.
Revised: In contrast, the latent state dynamics remain the same across all the MDP levels.

Original: This assumption has been shown to affect generalization (Song et al., 2020) in well studied benchmarks (Nichol et al., 2018) .
Revised: This assumption has been shown to affect generalization (Song et al., 2020) in well studied benchmarks (Nichol et al., 2018) .

Original: It has been observed that the agent overﬁt the scoreboard or timer and sometimes achieves the best training performance even without looking at the other part of the observation.
Revised: It has been observed that the agent overﬁt the scoreboard or timer and sometimes achieves the best training performance even without looking at the other part of the observation.

Original: Consequently, this policy completely fails to generalize in test environments.
Revised: Consequently, this policy completely fails to generalize in test environments.

Original: The high-capacity model even can memorize the whole training environment, thus severely affecting the generalization (Zhang et al., 2018b).
Revised: The high-capacity model even can memorize the whole training environment, thus severely affecting the generalization (Zhang et al., 2018b).

Original: This paper proposes a style transfer-based observation translation method that considers the observation’s content and reward signal.
Revised: This paper proposes a style transfer-based observation translation method that considers the observation’s content and reward signal.

Original: In particular, a generator performs the style transfer while the policy tries to be robust toward such style changes.
Revised: In particular, a generator performs the style transfer while the policy tries to be robust toward such style changes.

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

Original: In contrast, in this paper, we propose a uniﬁed learning approach where the style transfer is integrated into the reinforcement learning training.
Revised: In contrast, in this paper, we propose a uniﬁed learning approach where the style transfer is integrated into the reinforcement learning training.

Original: Thus the generator produces more effective translation targeted toward RL objective, making the RL policy more robust.
Revised: Thus the generator produces more effective translation targeted toward RL objective, making the RL policy more robust.

Original: The trajectory data from the agent’s replay buffer is ﬁrst clustered into different categories, and then observation is translated from one cluster’s style to another cluster’s style.
Revised: The trajectory data from the agent’s replay buffer is ﬁrst clustered into different categories, and then observation is translated from one cluster’s style to another cluster’s style.

Original: Here the style is determined by the attribute commonality of observation features in a cluster.
Revised: Here the style is determined by the attribute commonality of observation features in a cluster.

Original: The agent should be robust toward changes of such features.
Revised: The agent should be robust toward changes of such features.

Original: Moreover, the translated trajectories correspond to those that possibly appear in testing environments, assisting the agent in adapting to unseen scenarios.
Revised: Moreover, the translated trajectories correspond to those that possibly appear in testing environments, assisting the agent in adapting to unseen scenarios.

Original: Figure 1 shows an overview of our method.
Revised: Figure 1 shows an overview of our method.

Original: Our approach consists of a max-min game theoretic objective where a generator network modiﬁes the observation by style transferring it to maximize the agent’s probability of taking a different action.
Revised: Our approach consists of a max-min game theoretic objective where a generator network modiﬁes the observation by style transferring it to maximize the agent’s probability of taking a different action.

Original: In contrast, the policy network (agent) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward.
Revised: In contrast, the policy network (agent) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward.

Original: Based on this objective, we propose a practical deep reinforcement learning algorithm, AdversarialRobust
Revised: Based on this objective, we propose a practical deep reinforcement learning algorithm, Adversarial

Original:  Policy Optimization (ARPO), that aims to ﬁnd an optimal policy by mitigating the observational overﬁtting and eventually generalizes to test environments.
Revised: Robust Policy Optimization (ARPO), that aims to ﬁnd a robust policy by mitigating the observational overﬁtting and eventually generalizes to test environments.

Original: We empirically show the effectiveness of our ARPO agent in generalization and sample efﬁciency on challenging Procgen (Cobbe et al., 2020) and Distracting Control Suite (Stone et al., 2021) benchmark with image-based observation.
Revised: We empirically show the effectiveness of our ARPO agent in generalization and sample efﬁciency on challenging Procgen (Cobbe et al., 2020) and Distracting Control Suite (Stone et al., 2021) benchmark with image-based observation.

Original: The Procgen leverages procedural generation to create diverse environments to test the RL agents’ generalization capacity.
Revised: The Procgen leverages procedural generation to create diverse environments to test the RL agents’ generalization capacity.

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

Original: Empirically, we observed that our agent ARPO performs better in generalization and sample efﬁciency than the baseline PPO (Schulman et al., 2017) and a data augmentation-based RAD (Laskin et al., 2020a)
Revised: Empirically, we observed that our agent ARPO performs better in generalization and sample efﬁciency than the baseline PPO (Schulman et al., 2017) and a data augmentation-based RAD (Laskin et al., 2020a)

Original: algorithm on the many Procgen environments.
Revised: algorithm on the many Procgen environments.

Original: Furthermore, APRO performs better compared to PPO and SAC (Haarnoja et al., 2018) on the evaluated two Distracting Control Suite environments in various setups.
Revised: Furthermore, APRO performs better compared to PPO and SAC (Haarnoja et al., 2018) on the evaluated two Distracting Control Suite environments in various setups.

Original: In summary, our contributions are listed as follows:
Revised: In summary, our contributions are listed as follows:

Original: • We propose Adversarial Robust Policy Optimization (ARPO), a deep reinforcement learning algorithm to ﬁnd an optimal policy that generalizes to test environments in a zero-shot generalization framework.
Revised: • We propose Adversarial Robust Policy Optimization (ARPO), a deep reinforcement learning algorithm to ﬁnd a robust policy that generalizes to test environments in a zero-shot generalization framework.

Original: • We evaluate our approach on challenging Procgen (Cobbe et al., 2020) and Distracting Control Suite (Stone et al., 2021) benchmark in generalization ans sample efﬁciency settings.
Revised: • We evaluate our approach on challenging Procgen (Cobbe et al., 2020) and Distracting Control Suite (Stone et al., 2021) benchmark in generalization ans sample efﬁciency settings.

Original: • Empirically, we observed that our agent ARPO performs better generalization and sample efﬁciency 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.
Revised: • Empirically, we observed that our agent ARPO performs better generalization and sample efﬁciency 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.

Original:  P RELIMINARIES AND P ROBLEM S ETTINGS
Revised:  P RELIMINARIES AND P ROBLEM S ETTINGS

Original: Markov Decision Process (MDP)
Revised: Markov Decision Process (MDP)

Original: An MDP is denoted by M = ( S , A , P , r ) where at every timestep t , from an state s t ∈ S , the agent takes an action a t from a set of actions A .
Revised: An MDP is denoted by M = ( S , A , P , r ) where at every timestep t , from an state s t ∈ S , the agent takes an action a t from a set of actions A .

Original: The agent then receives a reward r t and the environment and move to a new state s t +1 ∈ S based on the transition probability P ( s t +1 | s t , a t ) .
Revised: The agent then receives a reward r t and the environment and move to a new state s t +1 ∈ S based on the transition probability P ( s t +1 | s t , a t ) .

Original: Reinforcement Learning .
Revised: Reinforcement Learning .

Original: Reinforcement learning aims to learn a policy π ∈ Π that maps state to actions.
Revised: Reinforcement learning aims to learn a policy π ∈ Π that maps state to actions.

Original: The policy’s objective is to maximize cumulative reward in an MDP, where Π is the set of all possible policies.
Revised: The policy’s objective is to maximize cumulative reward in an MDP, where Π is the set of all possible policies.

Original: Thus, the policy which achieves the highest possible cumulative reward is the optimal policy π ∗ ∈ Π .
Revised: Thus, the policy which achieves the highest possible cumulative reward is the optimal policy π ∗ ∈ Π .

Original: Generalization in Reinforcement Learning .
Revised: Generalization in Reinforcement Learning .

Original: 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 ﬁxed optimal policy π ∗ which can achieve maximal return over levels generated from the distribution D .
Revised: 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 ﬁxed optimal policy π ∗ which can achieve maximal return over levels generated from the distribution D .

Original: The levels may differ in observational variety, such as different background colors.
Revised: The levels may differ in observational variety, such as different background colors.

Original: In this setup, the agent has access to a ﬁxed set of MDP levels during training the policy.
Revised: In this setup, the agent has access to a ﬁxed set of MDP levels during training the policy.

Original: The trained agent is then tested on the unseen levels to measure the generalization performance of the agent.
Revised: The trained agent is then tested on the unseen levels to measure the generalization performance of the agent.

Original: Style Transfer .
Revised: Style Transfer .

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

Original: 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).
Revised: 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).

Original: Given training data from two domains, these models learn to style-translate images from one domain to another.
Revised: Given training data from two domains, these models learn to style-translate images from one domain to another.

Original: The domain is deﬁned as a set of images sharing the same attribute value, such as similar background color and texture.
Revised: The domain is deﬁned as a set of images sharing the same attribute value, such as similar background color and texture.

Original: The shared attributes in a domain are considered as the “style” of that domain.
Revised: The shared attributes in a domain are considered as the “style” of that domain.

Original: Many translation methods require paired images from two domains, which is not feasible in the reinforcement learning setup.
Revised: Many translation methods require paired images from two domains, which is not feasible in the reinforcement learning setup.

Original: The agent collects the data during policy learning, and annotating them in pairs is not feasible.
Revised: The agent collects the data during policy learning, and annotating them in pairs is not feasible.

Original: Thus, we leverage unpaired image to image translation, which does not require a one-to-one mapping of annotated images from two domains.
Revised: Thus, we leverage unpaired image to image translation, which does not require a one-to-one mapping of annotated images from two domains.

Original: In particular, we build upon the work of StarGAN (Choi et al., 2018) which is capable of learning mappings among multiple domains efﬁciently.
Revised: In particular, we build upon the work of StarGAN (Choi et al., 2018) which is capable of learning mappings among multiple domains efﬁciently.

Original: The model takes in training data of multiple domains and learns the mappings between each pair of available domains using only a single generator.
Revised: The model takes in training data of multiple domains and learns the mappings between each pair of available domains using only a single generator.

Original: This method automatically captures special characteristics of one image domain and ﬁgures out how these characteristics could be translated into the other image collection, making it appealing in the reinforcement learning setup.
Revised: This method automatically captures special characteristics of one image domain and ﬁgures out how these characteristics could be translated into the other image collection, making it appealing in the reinforcement learning setup.

Original: However, the RL agent generates experience trajectory data which is not separated into domains.
Revised: However, the RL agent generates experience trajectory data which is not separated into domains.

Original: Thus, we further apply a clustering approach which ﬁrst clusters the trajectory observation into domains, and then we train the generator, which style-translates among those domains.
Revised: Thus, we further apply a clustering approach which ﬁrst clusters the trajectory observation into domains, and then we train the generator, which style-translates among those domains.

Original:  A DVERSARIAL R OBUST P OLICY O PTIMIZATION (ARPO)
Revised:  A DVERSARIAL R OBUST P OLICY O PTIMIZATION (ARPO)

Original: Our approach consists of a max-min game theoretic objective where a generator network modiﬁes the observation by style transferring it to maximize the agent’s probability of taking a different action.
Revised: Our approach consists of a max-min game theoretic objective where a generator network modiﬁes the observation by style transferring it to maximize the agent’s probability of taking a different action.

Original: In contrast, the policy network (agent) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward.
Revised: In contrast, the policy network (agent) updates its parameters to minimize the effect of perturbation while maximizing the expected future reward.

Original: Based on this objective, we propose a practical deep reinforcement learning algorithm, Adversarial Robust Policy Optimization (ARPO), to ﬁnd an optimal policy that generalizes to test environments.
Revised: Based on this objective, we propose a practical deep reinforcement learning algorithm, Adversarial Robust Policy Optimization (ARPO), to ﬁnd a robust policy that generalizes to test environments.

Original: We now discuss the details of the policy network and generator.
Revised: We now discuss the details of the policy network and generator.

Original: 3. P OLICY N ETWORK
Revised: 3. P OLICY N ETWORK

Original: The objective of the policy network is two-fold: maximize the cumulative RL return during training and be robust to the perturbation on the observation imposed by the generator.
Revised: The objective of the policy network is two-fold: maximize the cumulative RL return during training and be robust to the perturbation on the observation imposed by the generator.

Original: Ideally, the action probability of the translated image should not change from the original image, as the translated observation is assumed to represent the same semantic as the original observation.
Revised: Ideally, the action probability of the translated image should not change from the original image, as the translated observation is assumed to represent the same semantic as the original observation.

Original: The reinforcement learning loss L π can be optimized with various reinforcement learning mechanisms, such as policy gradient (optimization).
Revised: The reinforcement learning loss L π can be optimized with various reinforcement learning mechanisms, such as policy gradient (optimization).

Original: However, we do not make any assumption the type of
Revised: However, we do not make any assumption the type of

Original: RL algorithms to use for optimizing L π ; thus, any algorithms can be incorporated with our method.
Revised: RL algorithms to use for optimizing L π ; thus, any algorithms can be incorporated with our method.

Original: Here we discuss the policy optimization formulation, which is based on the proximal policy optimization (PPO) (Schulman et al., 2017) objective:
Revised: Here we discuss the policy optimization formulation, which is based on the proximal policy optimization (PPO) (Schulman et al., 2017) objective:

Original: [Equation]
Revised: [Equation]

Original: where, π θ is the current policy and π θ old is the old policy; A t is estimated from the sampled trajectory by policy π θ old leveraging a value function estimator denoted as V .
Revised: where, π θ is the current policy and π θ old is the old policy; A t is estimated from the sampled trajectory by policy π θ old leveraging a value function estimator denoted as V .

Original: Here, both the policy π and value network V are represented as neural networks.
Revised: Here, both the policy π and value network V are represented as neural networks.

Original: An additional objective of the policy is to participate in the adversarial game with the style transfer generator.
Revised: An additional objective of the policy is to participate in the adversarial game with the style transfer generator.

Original: Thus the ﬁnal policy loss L θ is deﬁned as follows:
Revised: Thus the ﬁnal policy loss L θ is deﬁned as follows:

Original: [Equation]
Revised: [Equation]

Original: The policy parameter θ is updated to minimize equation 3.
Revised: The policy parameter θ is updated to minimize equation 3.

Original: For the observation x t at timestep t , the KL-component measure the distortion in the action distribution of the policy π θ due to the perturbation ( x (cid:48) t ).
Revised: For the observation x t at timestep t , the KL-component measure the distortion in the action distribution of the policy π θ due to the perturbation ( x (cid:48) t ).

Original: The hyperparameter β 1 determines the participation in an adversarial objective with the generator.
Revised: The hyperparameter β 1 determines the participation in an adversarial objective with the generator.

Original: By minimizing the KL-part, the policy becomes robust to the change in perturbation proposed by the generator.
Revised: By minimizing the KL-part, the policy becomes robust to the change in perturbation proposed by the generator.

Original: 3. G ENERATOR N ETWORK
Revised: 3. G ENERATOR N ETWORK

Original: It takes the observation x t as input and outputs the style-translated observation x (cid:48) t .
Revised: It takes the observation x t as input and outputs the style-translated observation x (cid:48) t .

Original: The objective of this network is to change the agent’s (policy network’s) probability of taking a different action for the corresponding input observation.
Revised: The objective of this network is to change the agent’s (policy network’s) probability of taking a different action for the corresponding input observation.

Original: However, the content or semantic of the observations needs to be intact, and the only style of the images will differ.
Revised: However, the content or semantic of the observations needs to be intact, and the only style of the images will differ.

Original: The ﬁrst step of this process is to curate a dataset with different style features such as background color and texture of objects present in the observation.
Revised: The ﬁrst step of this process is to curate a dataset with different style features such as background color and texture of objects present in the observation.

Original: We ﬁrst use experience trajectory observation data from the environment and separate them into different classes based on visual features (Figure 2).
Revised: We ﬁrst use experience trajectory observation data from the environment and separate them into different classes based on visual features (Figure 2).

Original: The task of the generator is then to translate the images from one class images to another class images.
Revised: The task of the generator is then to translate the images from one class images to another class images.

Original: In this case, the “style” is deﬁned as the commonality among images in a single class.
Revised: In this case, the “style” is deﬁned as the commonality among images in a single class.

Original: We now discuss details about how we separate the trajectory observation to get different clusters.
Revised: We now discuss details about how we separate the trajectory observation to get different clusters.

Original: Clustering Trajectory Observation .
Revised: Clustering Trajectory Observation .

Original: The trajectory data is ﬁrst clustered into n clusters.
Revised: The trajectory data is ﬁrst clustered into n clusters.

Original: Then, we use the Gaussian Mixture Model (GMM) for clustering and ResNet (He et al., 2016) for feature extraction and dimension reduction.
Revised: Then, we use the Gaussian Mixture Model (GMM) for clustering and ResNet (He et al., 2016) for feature extraction and dimension reduction.

Original: Firstly, a pre-trained (on ImageNet)
Revised: Firstly, a pre-trained (on ImageNet)

Original: ResNet model reduces an
Revised: ResNet model reduces an

Original: RGB image into a 1000 dimensional vector.
Revised: RGB image into a 1000 dimensional vector.

Original: This step extracts useful information from the raw image and reduces the dimension, allowing faster cluster training.
Revised: This step extracts useful information from the raw image and reduces the dimension, allowing faster cluster training.

Original: Note that this process focuses entirely on the visual aspect of the observation.
Revised: Note that this process focuses entirely on the visual aspect of the observation.

Original: After passing through the ResNet, the observation dataset is clustered using GMM.
Revised: After passing through the ResNet, the observation dataset is clustered using GMM.

Original: Generator Training .
Revised: Generator Training .

Original: Generator G tries to fool discriminator D in an adversarial setting by generating a realistic image represented by the true image distribution.
Revised: Generator G tries to fool discriminator D in an adversarial setting by generating a realistic image represented by the true image distribution.

Original: We extend the generator’s ability, and an additional objective of the generator is to fool the RL agent (policy network) by transforming the style of the original observation.
Revised: We extend the generator’s ability, and an additional objective of the generator is to fool the RL agent (policy network) by transforming the style of the original observation.

Original: For example, changing the background color to a different one does not change the semantic of the observation, but the image style would be different.
Revised: For example, changing the background color to a different one does not change the semantic of the observation, but the image style would be different.

Original: Ideally, the action mapping of the policy needs to be the same on both the original and translated observation.
Revised: Ideally, the action mapping of the policy needs to be the same on both the original and translated observation.

Original: However, a policy that focuses on the irrelevant background color will suffer from this translation, and thus its probability distribution deviates because of the style translation.
Revised: However, a policy that focuses on the irrelevant background color will suffer from this translation, and thus its probability distribution deviates because of the style translation.

Original: This is deﬁned as the
Revised: This is deﬁned as the

Original: KL-divergence KL ( π θ ( .
Revised: KL-divergence KL ( π θ ( .

Original: | x t ) , π θ ( .
Revised: | x t ) , π θ ( .

Original: | x (cid:48) t )) , where x t is the given observation and x (cid:48) t is the translated observation by the generator.
Revised: | x (cid:48) t )) , where x t is the given observation and x (cid:48) t is the translated observation by the generator.

Original: The generator tries to generate a style that eventually increases the KL quantity.
Revised: The generator tries to generate a style that eventually increases the KL quantity.

Original: On the other hand, the policy should learn to be robust toward such changes.
Revised: On the other hand, the policy should learn to be robust toward such changes.

Original: The policy training described above addresses this and tries to minimize the KL-divergence due to the style transfer, which eventually allows robust policy learning and ultimately prevents the agent from overﬁtting irrelevant features from the high-dimensional observation.
Revised: The policy training described above addresses this and tries to minimize the KL-divergence due to the style transfer, which eventually allows robust policy learning and ultimately prevents the agent from overﬁtting irrelevant features from the high-dimensional observation.

Original: We build on the generator on previous works (Choi et al., 2018; Zhu et al., 2017) which is a uniﬁed framework for a multi-domain image to image translation.
Revised: We build on the generator on previous works (Choi et al., 2018; Zhu et al., 2017) which is a uniﬁed framework for a multi-domain image to image translation.

Original: The discriminator loss is
Revised: The discriminator loss is

Original: [Equation]
Revised: [Equation]

Original: which consist of the adversarial loss L adv and domain classiﬁcation loss L rcls .
Revised: which consist of the adversarial loss L adv and domain classiﬁcation loss L rcls .

Original: The discriminator detects a fake image generated by the generator G from the real image in the given class data.
Revised: The discriminator detects a fake image generated by the generator G from the real image in the given class data.

Original: On the other hand, the generator loss is
Revised: On the other hand, the generator loss is

Original: [Equation]
Revised: [Equation]

Original: which consists of image translation components and policy component which is the KL-part.
Revised: which consists of image translation components and policy component which is the KL-part.

Original: From the translation side, the L adv is adversarial loss with discriminator L fcls is the loss of detecting fake image.
Revised: From the translation side, the L adv is adversarial loss with discriminator L fcls is the loss of detecting fake image.

Original: Finally, a cycle consistency loss (Choi et al., 2018; Zhu et al., 2017)
Revised: Finally, a cycle consistency loss (Choi et al., 2018; Zhu et al., 2017)

Original: L rec is used which makes sure the translated input can be translated back to the original input, thus only changing the domain related part and not the semantic.
Revised: L rec is used which makes sure the translated input can be translated back to the original input, thus only changing the domain related part and not the semantic.

Original: Some more details of these losses can be found in the supplementary materials.
Revised: Some more details of these losses can be found in the supplementary materials.

Original: Our addition to the generator objective is the KL-part in equation 5 which tries to challenge the RL agent and fool its policy by providing a more challenging translation.
Revised: Our addition to the generator objective is the KL-part in equation 5 which tries to challenge the RL agent and fool its policy by providing a more challenging translation.

Original: However, the eventual goal of this adversarial game is to help the agent to be robust to any irrelevant features in the observations.
Revised: However, the eventual goal of this adversarial game is to help the agent to be robust to any irrelevant features in the observations.

Original: The hyperparameter β 2 control the participation in the adversarial objective of the generator network.
Revised: The hyperparameter β 2 control the participation in the adversarial objective of the generator network.

Original: Now we discuss a practical RL algorithm based on the above objectives.
Revised: Now we discuss a practical RL algorithm based on the above objectives.

Original: 3. ARPO A LGORITHM
Revised: 3. ARPO A LGORITHM

Original: The training of the policy network and the generator network happen in an alternating way.
Revised: The training of the policy network and the generator network happen in an alternating way.

Original: First, the policy is trained on a batch of replay buffer trajectory observation, and then the generator parameters are updated based on the objectives discussed above.
Revised: First, the policy is trained on a batch of replay buffer trajectory observation, and then the generator parameters are updated based on the objectives discussed above.

Original: The more overall training progresses, the better the generator at translating observation, and the policy also becomes better and becomes robust to irrelevant style changes.
Revised: The more overall training progresses, the better the generator at translating observation, and the policy also becomes better and becomes robust to irrelevant style changes.

Original: For stable training, the input images data for the generator is kept ﬁxed initially, and we apply the GMM cluster only at the beginning to generate different class images (distribution).
Revised: For stable training, the input images data for the generator is kept ﬁxed initially, and we apply the GMM cluster only at the beginning to generate different class images (distribution).

Original: The pseudocode of the algorithm is given in Algorithm 1.
Revised: The pseudocode of the algorithm is given in Algorithm 1.

Original: do 3: do 4: a t ∼ π
Revised: do 3: do 4: a t ∼ π

Original: ( a t | x t ) 5: x t +1 ∼ P ( x t +1 | x t , a t ) 6: r t ∼ R ( x t , a t ) 7:
Revised: ( a t | x t ) 5: x t +1 ∼ P ( x t +1 | x t , a t ) 6: r t ∼ R ( x t , a t ) 7:

Original: D ←− D ∪ { ( x t , a t , r t , x t +1 ) } 8: end for 9: for each observation x t in D do 10: x (cid:48)
Revised: D ←− D ∪ { ( x t , a t , r t , x t +1 ) } 8: end for 9: for each observation x t in D do 10: x (cid:48)

Original: t ←− Generator ( x t ) // generate translated observation 11:
Revised: t ←− Generator ( x t ) // generate translated observation 11:

Original: Compute L π from data D using Equation 1, and 2 // update policy for RL objective 12:
Revised: Compute L π from data D using Equation 1, and 2 // update policy for RL objective 12:

Original: Compute L θ using Equation 3 // update policy with adversarial KL part 13:
Revised: Compute L θ using Equation 3 // update policy with adversarial KL part 13:

Original: Update generator by computing Equation 6, and14: end for 15: end for
Revised: Update generator by computing Equation 6, and14: end for 15: end for

Original: Discussion on convergence .
Revised: Discussion on convergence .

Original: The adversarial min-max KL component KL ( π θ ( .
Revised: The adversarial min-max KL component KL ( π θ ( .

Original: | x t ) , π θ ( .
Revised: | x t ) , π θ ( .

Original: | x (cid:48) t )) become zero when the RL policy is optimal ( π ∗ ) and the generators only perturbs (translates) the irrelevant part of all observations.
Revised: | x (cid:48) t )) become zero when the RL policy is optimal ( π ∗ ) and the generators only perturbs (translates) the irrelevant part of all observations.

Original: In that case, the optimal policy only focuses on the relevant part of the observations, that is, the true state; thus, any changes in irrelevant part due to style transfer will be ignored.
Revised: In that case, the optimal policy only focuses on the relevant part of the observations, that is, the true state; thus, any changes in irrelevant part due to style transfer will be ignored.

Original: At that point the KL component become zero as the π ∗ θ ( .
Revised: At that point the KL component become zero as the π ∗ θ ( .

Original: | x t ) =
Revised: | x t ) =

Original: π ∗ θ ( .
Revised: π ∗ θ ( .

Original: | x (cid:48) t ) .
Revised: | x (cid:48) t ) .

Original: Note that in practice, the algorithm is usually trained for limited timesteps, and thus the adversarial training might not converge to optimal.
Revised: Note that in practice, the algorithm is usually trained for limited timesteps, and thus the adversarial training might not converge to optimal.

Original: However, empirically we observe that our algorithm ARPO achieves improved performance in many challenging Procgen and Distracting control suite tasks.
Revised: However, empirically we observe that our algorithm ARPO achieves improved performance in many challenging Procgen and Distracting control suite tasks.

Original:  E XPERIMENTS 4. S ETUP
Revised:  E XPERIMENTS 4. S ETUP

Original: Our experiments cover both discrete (Procgen) and continuous (Distracting control) action spaces with image-based observation.
Revised: Our experiments cover both discrete (Procgen) and continuous (Distracting control) action spaces with image-based observation.

Original: Procgen Environment .
Revised: Procgen Environment .

Original: We evaluate our proposed agent ARPO on the challenging RL generalization benchmark Procgen (Cobbe et al., 2020).
Revised: We evaluate our proposed agent ARPO on the challenging RL generalization benchmark Procgen (Cobbe et al., 2020).

Original: We use the standard evaluation protocol from Cobbe et al.
Revised: We use the standard evaluation protocol from Cobbe et al.

Original: (2020) where the agent is allowed to train on only 200 levels of each environment, with the difﬁculty level set as easy .
Revised: (2020) where the agent is allowed to train on only 200 levels of each environment, with the difﬁculty level set as easy .

Original: Then the agents are tested over more than 100K levels, full distribution.
Revised: Then the agents are tested over more than 100K levels, full distribution.

Original: Thus, to better generalize to the test environment, the agent must master the skill and avoid overﬁtting to spurious features in the training environment.
Revised: Thus, to better generalize to the test environment, the agent must master the skill and avoid overﬁtting to spurious features in the training environment.

Original: The observations are raw pixel images, and the environment varies drastically between levels; some snippets are given in Figure 3.
Revised: The observations are raw pixel images, and the environment varies drastically between levels; some snippets are given in Figure 3.

Original: Distracting Control Suite Environment .
Revised: Distracting Control Suite Environment .

Original: The distracting control adds various visual distractions on top of the Deepmind Control suite (Tassa et al., 2020) to facilitate the evaluation of general- ization performance.
Revised: 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.

Original: We experimented on two distracting settings, background  and color for easy difﬁculty mode.
Revised: We experimented on two distracting settings, background (dynamic) and color for easy difﬁculty mode.

Original: In the distracting background, videos played in the background, and for the color distraction, the body of the robot color changes in different agent interactions.
Revised: In the distracting background, videos played in the background, and for the color distraction, the body of the robot color changes in different agent interactions.

Original: The train and test environments differ by different distractions, such as separate videos for train and test.
Revised: The train and test environments differ by different distractions, such as separate videos for train and test.

Original: As the agent has to use image-based observation, the distractions make the task challenging.
Revised: As the agent has to use image-based observation, the distractions make the task challenging.

Original: Some snippets of the environments can be found in Figure 4.
Revised: Some snippets of the environments can be found in Figure 4.

Original: Baselines .
Revised: Baselines .

Original: We compare our algorithm with Proximal Policy Optimization ( PPO ), an on-policy algorithm motivated by how we can quickly optimize our policy while staying close to the current policy.
Revised: We compare our algorithm with Proximal Policy Optimization ( PPO ), an on-policy algorithm motivated by how we can quickly optimize our policy while staying close to the current policy.

Original: PPO performs effectively in many RL benchmarks, which includes the Procgen environments (Cobbe et al., 2020).
Revised: PPO performs effectively in many RL benchmarks, which includes the Procgen environments (Cobbe et al., 2020).

Original: Furthermore, we evaluated our method with a data augmentation technique, RAD (Laskin et al., 2020a) on Procgen environments.
Revised: Furthermore, we evaluated our method with a data augmentation technique, RAD (Laskin et al., 2020a) on Procgen environments.

Original: RAD is a plug-and-play module that can be used with existing RL algorithms and shows effective empirical performance in complex
Revised: RAD is a plug-and-play module that can be used with existing RL algorithms and shows effective empirical performance in complex

Original: RL benchmarks, including some Procgen environments.
Revised: RL benchmarks, including some Procgen environments.

Original: In particular, we evaluated Cutout Color augmentation technique which has shown better results in many Procgen environments compared to other augmentation techniques evaluated in Laskin et al.
Revised: In particular, we evaluated Cutout Color augmentation technique which has shown better results in many Procgen environments compared to other augmentation techniques evaluated in Laskin et al.

Original: (2020a).
Revised: (2020a).

Original: This augmentation is applied on top of standard PPO.
Revised: This augmentation is applied on top of standard PPO.

Original: Note that our ARPO algorithm can be easily plugged in with these data augmentation-based methods.
Revised: Note that our ARPO algorithm can be easily plugged in with these data augmentation-based methods.

Original: However, in this paper, we evaluated how our ARPO agent performs even when no augmentation is used.
Revised: However, in this paper, we evaluated how our ARPO agent performs even when no augmentation is used.

Original: Furthermore, in the Distracting control suite, in addition to
Revised: Furthermore, in the Distracting control suite, in addition to

Original: PPO, we compare ARPO with Soft Actor-Critic ( SAC ) (Haarnoja et al., 2018) which is an off-policy algorithm that optimizes a stochastic policy.
Revised: PPO, we compare ARPO with Soft Actor-Critic ( SAC ) (Haarnoja et al., 2018) which is an off-policy algorithm that optimizes a stochastic policy.

Original: Implementation .
Revised: Implementation .

Original: We implemented our ARPO on the Ray RLlib framework (Liang et al., 2018) and used a CNN-based model for the policy network for these experiments.
Revised: We implemented our ARPO on the Ray RLlib framework (Liang et al., 2018) and used a CNN-based model for the policy network for these experiments.

Original: For all the experiments with ARPO, we set the β 1 = β 2 = 20 (in equation 3, and 5).
Revised: For all the experiments with ARPO, we set the β 1 = β 2 = 20 (in equation 3, and 5).

Original: The generator’s hyperparameters are set to λ cls = 1 , and λ rec = 10 .
Revised: The generator’s hyperparameters are set to λ cls = 1 , and λ rec = 10 .

Original: Number of cluster is set to 3.
Revised: Number of cluster is set to 3.

Original: Details of model conﬁgurations and parameters are given in the Appendix.
Revised: Details of model conﬁgurations and parameters are given in the Appendix.

Original: We will open-source our code and data for reproducibility and facilitating further research.
Revised: We will open-source our code and data for reproducibility and facilitating further research.

Original: 4. P ROCGEN R ESULTS
Revised: 4. P ROCGEN R ESULTS

Original: Generalization .
Revised: Generalization .

Original: The results in Figure 5 show the test results of ARPO and baselines PPO and
Revised: The results in Figure 5 show the test results of ARPO and baselines PPO and

Original: RAD on 16 Procgen environments with image observation.
Revised: RAD on 16 Procgen environments with image observation.

Original: Overall, ARPO performs better in many environments.
Revised: Overall, ARPO performs better in many environments.

Original: We observe that, almost in half of the environments, our agent ARPO outperforms both baselines.
Revised: We observe that, almost in half of the environments, our agent ARPO outperforms both baselines.

Original: It also performs comparably in a few other environments.
Revised: It also performs comparably in a few other environments.

Original: Note that in some environments such as Ninja, Plunder, Chaser, Heist, all the agents fail to obtain a reasonable training accuracy (Cobbe et al., 2020) and thus perform poorly during testing (generalization).
Revised: Note that in some environments such as Ninja, Plunder, Chaser, Heist, all the agents fail to obtain a reasonable training accuracy (Cobbe et al., 2020) and thus perform poorly during testing (generalization).

Original: [Table]
Revised: [Table]

Original: This result shows that our proposed agent ARPO achieves a strong performance and reduces the generalization gap compared to the tested baselines on generalization.
Revised: This result shows that our proposed agent ARPO achieves a strong performance and reduces the generalization gap compared to the tested baselines on generalization.

Original: We further show how each agent achieves generalization during the entire timesteps.
Revised: We further show how each agent achieves generalization during the entire timesteps.

Original: Consequently, we compute the maximum score achieved by each agent in the entire duration and take the average best score.
Revised: Consequently, we compute the maximum score achieved by each agent in the entire duration and take the average best score.

Original: Finally we report mean and standard deviation of these best scores across 3 seed runs.
Revised: Finally we report mean and standard deviation of these best scores across 3 seed runs.

Original: The generalization results are computed by evaluating the trained agents on test levels (full distribution).
Revised: The generalization results are computed by evaluating the trained agents on test levels (full distribution).

Original: The performance is computed by evaluating each agent for 10 random episode trials after a training interval of 10.
Revised: The performance is computed by evaluating each agent for 10 random episode trials after a training interval of 10.

Original: All agents ran for 22 million timesteps for training.
Revised: All agents ran for 22 million timesteps for training.

Original: Table 1 shows the results comparison for all the agents on Procgen environments.
Revised: Table 1 shows the results comparison for all the agents on Procgen environments.

Original: Similar to the previous discussion, we see that in many of the environments, our ARPO agents perform better than the baselines.
Revised: Similar to the previous discussion, we see that in many of the environments, our ARPO agents perform better than the baselines.

Original: ARPO performs better than both the baselines in half of the environment, and in 12 out 16 environments,
Revised: ARPO performs better than both the baselines in half of the environment, and in 12 out 16 environments,

Original: ARPO performs better than at least on baselines.
Revised: ARPO performs better than at least on baselines.

Original: This result shows that our agent outperforms the baselines in the generalization results on Procgen benchmarks.
Revised: This result shows that our agent outperforms the baselines in the generalization results on Procgen benchmarks.

Original: Sample Efﬁciency .
Revised: Sample Efﬁciency .

Original: We evaluate and compare our agent on the sample efﬁciency during training.
Revised: We evaluate and compare our agent on the sample efﬁciency during training.

Original: This result shows how quickly our agent achieves a better score during training.
Revised: This result shows how quickly our agent achieves a better score during training.

Original: Learning curve can be found in Appendix B (Figure 7).
Revised: Learning curve can be found in Appendix B (Figure 7).

Original: In many environments, we see our agent ARPO perform better than the baselines PPO and RAD.
Revised: In many environments, we see our agent ARPO perform better than the baselines PPO and RAD.

Original: These results show that despite optimizing for the generalization, our adversarial training helps the
Revised: These results show that despite optimizing for the generalization, our adversarial training helps the

Original: ARPO agent learn a robust policy which is also better during training.
Revised: ARPO agent learn a robust policy which is also better during training.

Original: 4. D ISTRACTING C ONTROL R ESULTS
Revised: 4. D ISTRACTING C ONTROL R ESULTS

Original: Figure 6 shows the results of ARPO and PPO on Walker walk environment from distracting control suite.
Revised: Figure 6 shows the results of ARPO and PPO on Walker walk environment from distracting control suite.

Original: We observe that ARPO performs better compared to the baseline algorithm PPO in both sample efﬁciency (train) and generalization (test) for background and color distraction.
Revised: We observe that ARPO performs better compared to the baseline algorithm PPO in both sample efﬁciency (train) and generalization (test) for background and color distraction.

Original: Note that we observe large variance (standard deviation) in the results across the run, particularly in test time.
Revised: Note that we observe large variance (standard deviation) in the results across the run, particularly in test time.

Original: These results correspond to the benchmark results where Stone et al.
Revised: These results correspond to the benchmark results where Stone et al.

Original: (2021) also ﬁnd comparatively larger variance in reported results.
Revised: (2021) also ﬁnd comparatively larger variance in reported results.

Original: In some cases (Figure 6b), the test performance surpasses the corresponding timestep’s train reward.
Revised: In some cases (Figure 6b), the test performance surpasses the corresponding timestep’s train reward.

Original: This scenario might happen because the videos for train and test environments have different difﬁculties.
Revised: This scenario might happen because the videos for train and test environments have different difﬁculties.

Original: Thus in some cases, the agent ﬁnds better test performance than the training.
Revised: Thus in some cases, the agent ﬁnds better test performance than the training.

Original: However, our ARPO’s performance remains better during testing compared to PPO.
Revised: However, our ARPO’s performance remains better during testing compared to PPO.

Original: We also observe that in these settings ARPO (and also PPO) performs better compared to SAC.
Revised: We also observe that in these settings ARPO (and also PPO) performs better compared to SAC.

Original: The results are in the Appendix D (Figure 10).
Revised: The results are in the Appendix D (Figure 10).

Original: Furthermore, results analysis for
Revised: Furthermore, results analysis for

Original: Cheetah run environment is in the Appendix C.
Revised: Cheetah run environment is in the Appendix C.

Original: Please see the Appendix for ablation study (Section E), and related work (Section F).
Revised: Please see the Appendix for ablation study (Section E), and related work (Section F).

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: Lack of visual diversity in the observation of the environment might result in poor clustering, which eventually leads to less challenging style translation by the adversarial generator.
Revised: Lack of visual diversity in the observation of the environment might result in poor clustering, which eventually leads to less challenging style translation by the adversarial generator.

Original: However, in those scenarios, the reinforcement learning algorithms often perform well as the train and test environment remain consistent (e.g., same background).
Revised: However, in those scenarios, the reinforcement learning algorithms often perform well as the train and test environment remain consistent (e.g., same background).

Original: Furthermore, this setup might lead the agent to overﬁt the training environment, which results in the agent performing poorly in a slightly different environment (Song et al., 2020; Zhang et al., 2018b).
Revised: Furthermore, this setup might lead the agent to overﬁt the training environment, which results in the agent performing poorly in a slightly different environment (Song et al., 2020; Zhang et al., 2018b).

Original: In this paper, we are interested in the scenario where the train and test environment are visually different, which is a more practical setup (Cobbe et al., 2020; Stone et al., 2021).
Revised: In this paper, we are interested in the scenario where the train and test environment are visually different, which is a more practical setup (Cobbe et al., 2020; Stone et al., 2021).

Original: In summary, we proposed an algorithm, Adversarial Robust Policy Optimization (ARPO), to improve generalization for reinforcement learning agents by removing the effect of overﬁtting in highdimensional observation space.
Revised: In summary, we proposed an algorithm, Adversarial Robust Policy Optimization (ARPO), to improve generalization for reinforcement learning agents by removing the effect of overﬁtting in highdimensional observation space.

Original: Our method consists of a max-min game theoretic objective where a generator is used to transfer the style of high-dimensional observation (e.g., image), thus perturb original observation while maximizing the agent’s probability of taking a different action for the corresponding input observation.
Revised: Our method consists of a max-min game theoretic objective where a generator is used to transfer the style of high-dimensional observation (e.g., image), thus perturb original observation while maximizing the agent’s probability of taking a different action for the corresponding input observation.

Original: In contrast, a policy network updates its parameters to minimize the effect of such translation, thus being robust to the perturbation while maximizing the expected future reward.
Revised: In contrast, a policy network updates its parameters to minimize the effect of such translation, thus being robust to the perturbation while maximizing the expected future reward.

Original: We tested our approaches on Procgen and distracting control benchmarks in generalization setup and conﬁrmed the effectiveness of our proposed algorithm.
Revised: We tested our approaches on Procgen and distracting control benchmarks in generalization setup and conﬁrmed the effectiveness of our proposed algorithm.

Original: Furthermore, empirical results show that our method generalizes better in unseen test environments than the tested baseline algorithms.
Revised: Furthermore, empirical results show that our method generalizes better in unseen test environments than the tested baseline algorithms.

Original: A A PPENDIX B A DDITIONAL P ROCGEN R ESULTS
Revised: A A PPENDIX B A DDITIONAL P ROCGEN R ESULTS

Original: Figure 7 shows the sample efﬁciency results for ARPO, PPO and RAD.
Revised: Figure 7 shows the sample efﬁciency results for ARPO, PPO, and RAD.

Original: Procgen aggregated results .
Revised: Procgen aggregated results .

Original: We further show the aggregated results on all the environments in the
Revised: We further show the aggregated results on all the environments in the

Original: Procgen environment.
Revised: Procgen environment.

Original: Figure 8 shows the aggregated results comparison on all 16 environments of
Revised: Figure 8 shows the aggregated results comparison on all 16 environments of

Original: Procgen which is a probabilistic measure that highlights how likely an improvement of the algorithm over other proposed in Agarwal et al. (2021b).
Revised: Procgen, which is a probabilistic measure that highlights how likely an improvement of the algorithm over other proposed in Agarwal et al. (2021b).

Original: C A DDITIONAL D ISTRACTING C ONTROL R ESULTS
Revised: C A DDITIONAL D ISTRACTING C ONTROL R ESULTS

Original: Results for ARPO and PPO on Cheetah run with background distraction are in Figure 9.
Revised: Results for ARPO and PPO on Cheetah run with background distraction are in Figure 9.

Original: We see that both agents perform comparably in both background and color distraction.
Revised: We see that both agents perform comparably in both background and color distraction.

Original: We observe similar results for both the agents in this setup of Cheetah run environments.
Revised: We observe similar results for both the agents in this setup of Cheetah run environments.

Original: Note that in this training timestep the PPO agents perform poorly in the training and testing.
Revised: Note that in this training timestep the PPO agents perform poorly in the training and testing.

Original: The highest achievable reward for this environment is much higher (see Stone et al.
Revised: The highest achievable reward for this environment is much higher (see Stone et al.

Original: Consequently the the ARPO agent could not improve over the base algorithm PPO.
Revised: Consequently the the ARPO agent could not improve over the base algorithm PPO.

Original: This suggest that the base algorithm needs to trained sufﬁciently to achieve any reasonable reward before we beneﬁt of the adversarial training.
Revised: This suggest that the base algorithm needs to trained sufﬁciently to achieve any reasonable reward before we beneﬁt of the adversarial training.

Original: Note that the resulting score is low compared to the reported results in the benchmark paper Stone et al.
Revised: Note that the resulting score is low compared to the reported results in the benchmark paper Stone et al.

Original: (2021) which suggests the PPO itself could not learn helpful behavior in this environment.
Revised: (2021) which suggests the PPO itself could not learn helpful behavior in this environment.

Original: D C OMPARISON WITH SAC
Revised: D C OMPARISON WITH SAC

Original: Results comparison on Distracting Control for ARPO, PPO, and SAC are in Figure 10 (Walker walk), and 11 (Cheetah run).
Revised: Results comparison on Distracting Control for ARPO, PPO, and SAC are in Figure 10 (Walker walk), and 11 (Cheetah run).

Original: We see that ARPO (and PPO) perform better compared to the SAC in both sample efﬁciency (train) and generalization (test) in these settings.
Revised: We see that ARPO (and PPO) perform better compared to the SAC in both sample efﬁciency (train) and generalization (test) in these settings.

Original: E A BLATION S TUDY E. H YPERPARAMETERS OF ARPO
Revised: E A BLATION S TUDY E. H YPERPARAMETERS OF ARPO

Original: We conduct a study on how ARPO agent’s performance varies due to its hyperparameters: (a) the number of clusters generated by the GMM model for generator training, and (b) β 1 , β 2 which indicate the amount of participation of policy and generator in the adversarial objective.
Revised: We conduct a study on how ARPO agent’s performance varies due to its hyperparameters: (a) the number of clusters generated by the GMM model for generator training, and (b) β 1 , β 2 which indicate the amount of participation of policy and generator in the adversarial objective.

Original: Figure 12 shows the ablation results.
Revised: Figure 12 shows the ablation results.

Original: We observe that our ARPO agent shows improvement in generalization (test) performance with the increase in number of cluster on Climber environment (Figure 12a).
Revised: We observe that our ARPO agent shows improvement in generalization (test) performance with the increase in number of cluster on Climber environment (Figure 12a).

Original: As the number of cluster increase, the generator has more option to choose for translation, and thus the policy might face a hard challenge and thus learn a more robust policy.
Revised: As the number of cluster increase, the generator has more option to choose for translation, and thus the policy might face a hard challenge and thus learn a more robust policy.

Original: In this paper, we report the ARPO agent’s results on cluster n cluster = 3 .
Revised: In this paper, we report the ARPO agent’s results on cluster n cluster = 3 .

Original: On the other hand, and ablation results on different β values are shown in Figure 12b.
Revised: On the other hand, and ablation results on different β values are shown in Figure 12b.

Original: When β values are large the test performance goes up; however, training performance suffer a bit.
Revised: When β values are large the test performance goes up; however, training performance suffer a bit.

Original: In this paper, we report the results on β 1 = β 2 = 20 , which shows a balance in train and test performance.
Revised: In this paper, we report the results on β 1 = β 2 = 20 , which shows a balance in train and test performance.

Original: E. Q UALITATIVE A NALYSIS OF A DVERSARIAL S TYLE T RANSFER
Revised: E. Q UALITATIVE A NALYSIS OF A DVERSARIAL S TYLE T RANSFER

Original: In this section we discuss how the generator performs in generating adversarially translated observation during agent training.
Revised: In this section we discuss how the generator performs in generating adversarially translated observation during agent training.

Original: Figure 13,

Original: [Figure]
Revised: [Figure]

Revised: Figure 13, 14, 13, 15, and 16 show sample translation from the generator.

Revised: Please see the corresponding caption for detailed discussion.

Original: E. KL R EGULARIZATION
Revised: E. KL R EGULARIZATION

Original: Figure 17 shows the KL regularization results in different policy iteration.
Revised: Figure 17 shows the KL regularization results in different policy iteration.

Original: We for each environment we run the experiment and collect KL value ( KL ( π θ ( .
Revised: We for each environment we run the experiment and collect KL value ( KL ( π θ ( .

Original: | x t ) , π θ ( .
Revised: | x t ) , π θ ( .

Original: | x (cid:48) t )) ) for each policy iteration.
Revised: | x (cid:48) t )) ) for each policy iteration.

Original: To calculate moving average, we select a window size of 500 and the standard deviation is calculated on these 500 values, and showed in the shaded area in Figure 17.
Revised: To calculate moving average, we select a window size of 500 and the standard deviation is calculated on these 500 values, and showed in the shaded area in Figure 17.

Original: F R ELATED W ORK
Revised: F R ELATED W ORK

Original: Data augmentation in RL .
Revised: Data augmentation in RL .

Original: In recent time, many approaches have been proposed to address generalization challenges including various data augmentation approaches (Cobbe et al., 2019; Laskin et al., 2020a; Kostrikov et al., 2020; Raileanu et al., 2020;
Revised: In recent time, many approaches have been proposed to address generalization challenges including various data augmentation approaches (Cobbe et al., 2019; Laskin et al., 2020a; Kostrikov et al., 2020; Raileanu et al., 2020;

Original: Laskin et al., 2020b; Igl et al., 2019;
Revised: Laskin et al., 2020b; Igl et al., 2019;

Original: Osband et al., 2018; Lee et al., 2020; Wang et al., 2020; Zhang et al., 2021) have shown to improve generalization.
Revised: Osband et al., 2018; Lee et al., 2020; Wang et al., 2020; Zhang et al., 2021) have shown to improve generalization.

Original: The common theme of these approaches is to increase diversity in the training data; however, these perturbations are primarily performed in isolation of the RL reward function, which might change crucial aspects of the observation, resulting in sub-optimal policy learning.
Revised: The common theme of these approaches is to increase diversity in the training data; however, these perturbations are primarily performed in isolation of the RL reward function, which might change crucial aspects of the observation, resulting in sub-optimal policy learning.

Original: In contrast, we propose a style transfer that tries to generate semantically similar but different visual style observation to train a robust policy.
Revised: In contrast, we propose a style transfer that tries to generate semantically similar but different visual style observation to train a robust policy.

Original: Furthermore, our style perturbation takes into account the reward signal from the policy.
Revised: Furthermore, our style perturbation takes into account the reward signal from the policy.

Revised: Note that the base RL algorithm (e.g., PPO) in our ARPO method still uses the original observation from the environment while participating in adversarial objectives (see Figure 1).

Revised: Thus, data augmentation can be applied to the original observations before feeding them to the base RL algorithm to potentially improve the base RL agent, as observed in many data augmentation techniques discussed above.

Original: Style Transfer in RL .
Revised: Style Transfer in RL .

Original: Many approaches have been proposed which use style transfer to improve reinforcement learning algorithms.
Revised: Many approaches have been proposed which use style transfer to improve reinforcement learning algorithms.

Original: Gamrian & Goldberg (2019) use image-to-image translation between domains to learn the policy in zero-shot learning and imitation learning setup.
Revised: Gamrian & Goldberg (2019) use image-to-image translation between domains to learn the policy in zero-shot learning and imitation learning setup.

Original: They use different levels’ data and train unpaired GANs between two levels (domains), which requires access to the annotated agent’s trajectory data in both source and target domains.
Revised: They use different levels’ data and train unpaired GANs between two levels (domains), which requires access to the annotated agent’s trajectory data in both source and target domains.

Original: Furthermore, Smith et al.
Revised: Furthermore, Smith et al.

Original: (2019) performs pixel-level image translation via CycleGAN to convert the human demonstration to a video of a robot.
Revised: (2019) performs pixel-level image translation via CycleGAN to convert the human demonstration to a video of a robot.

Original: These translated data can then be used to convert the human demonstration into a video of a robot used to construct a reward function that helps learn complex tasks.
Revised: These translated data can then be used to convert the human demonstration into a video of a robot used to construct a reward function that helps learn complex tasks.

Original: However, their
Revised: However, their

Original: CycleGAN operates on two predeﬁned domains, human demonstration and robot video.
Revised: CycleGAN operates on two predeﬁned domains, human demonstration and robot video.

Original: In contrast, we ﬁrst automatically generate the domains using a clustering algorithm and then train the style-translation generator.
Revised: In contrast, we ﬁrst automatically generate the domains using a clustering algorithm and then train the style-translation generator.

Original: In our case, we do not need the information of levels; instead, we automatically cluster the trajectory data based on visual features of the observations.
Revised: In our case, we do not need the information of levels; instead, we automatically cluster the trajectory data based on visual features of the observations.

Original: Additionally, our visual-based clustering may put observations from multiple levels into a cluster, potentially preventing GAN from overﬁtting Gamrian & Goldberg (2019) to any particular environment levels.
Revised: Additionally, our visual-based clustering may put observations from multiple levels into a cluster, potentially preventing GAN from overﬁtting Gamrian & Goldberg (2019) to any particular environment levels.

Original: Adversarial Approach in RL .
Revised: Adversarial Approach in RL .

Original: Adversarial methods have been used in the context of reinforcement learning to improve policy training.
Revised: Adversarial methods have been used in the context of reinforcement learning to improve policy training.

Original: Li et al.
Revised: Li et al.

Original: (2021) proposes to combine data augmentation with auxiliary adversarial loss to improve generalization.
Revised: (2021) proposes to combine data augmentation with auxiliary adversarial loss to improve generalization.

Original: In particular, they use an adversarial discriminator to predict the label of the observation.
Revised: In particular, they use an adversarial discriminator to predict the label of the observation.

Original: However, they need to know the domain label, limiting the applicability as it might not be available, and the number of labels might be large, for example, in procedural generation.
Revised: However, they need to know the domain label, limiting the applicability as it might not be available, and the number of labels might be large, for example, in procedural generation.

Original: In contrast, we use an unsupervised clustering that automatically ﬁnds out similar visual features later used for generator training.
Revised: In contrast, we use an unsupervised clustering that automatically ﬁnds out similar visual features later used for generator training.

Original: (Pinto et al., 2017b) use adversarial min-max training to learn robust policy against forces/disturbances in the system (Pinto et al., 2017b).
Revised: (Pinto et al., 2017b) use adversarial min-max training to learn robust policy against forces/disturbances in the system (Pinto et al., 2017b).

Original: Adversarial training has been used between two robots to improve object grasping (Pinto et al., 2017a). Adversarial training in  multi-agent self-play scenariohas been applied (Heinrich & Silver, 2016).
Revised: In addition, adversarial training has been used between two robots to improve object grasping (Pinto et al., 2017a) and in the context of multi-agent self-play scenario (Heinrich & Silver, 2016).

Revised: Zhang & Guo (2021) generate adversarial examples for observations in a trajectory by minimizing the cumulative reward objective of the reinforcement learning agent.

Original: In contrast, we propose a visual-based style-transfer which eventually guide to learn a robust policy.
Revised: In contrast to these methods, we use adversarial visual-based style-transfer of the observation to guide robust policy learning by enforcing the policy to produce similar output action distributions.

Original: Generalization in RL .
Revised: Generalization in RL .

Original: Many methods have been proposed to reduce the generalization gap between training and unseen test environment in the context of reinforcement learning.
Revised: Many methods have been proposed to reduce the generalization gap between training and unseen test environment in the context of reinforcement learning.

Original: State learning approaches (Higgins et al., 2017; Agarwal et al., 2021a; Zhang et al., 2020) and auto-encoder based latent state learning (Lange et al., 2012; Lange & Riedmiller, 2010; Hafner et al., 2019) with reconstruction loss have been proposed.
Revised: State learning approaches (Higgins et al., 2017; Agarwal et al., 2021a; Zhang et al., 2020) and auto-encoder based latent state learning (Lange et al., 2012; Lange & Riedmiller, 2010; Hafner et al., 2019) with reconstruction loss have been proposed.

Original: The sequential structure in RL (Agarwal et al., 2021a) and behavioral similarity between states (Zhang et al., 2020) to learn invariant states have been leveraged to improve RL robustness and generalization.
Revised: The sequential structure in RL (Agarwal et al., 2021a) and behavioral similarity between states (Zhang et al., 2020) to learn invariant states have been leveraged to improve RL robustness and generalization.

Original: In contrast, we learn the invariant state by providing the RL agent with an additional variant of the same observation while retaining the reward structure using an adversarial max-min objective.
Revised: In contrast, we learn the invariant state by providing the RL agent with an additional variant of the same observation while retaining the reward structure using an adversarial max-min objective.

Original: Our primary focus is to demonstrate how the adversarial style transfer helps learn a robust policy and improves the robustness of the base RL policy used.
Revised: Our primary focus is to demonstrate how the adversarial style transfer helps learn a robust policy and improves the robustness of the base RL policy used.

Original: Thus our method is orthogonal to many existing algorithms, including strong baseline IDAAC (Raileanu & Fergus, 2021).
Revised: Thus our method is orthogonal to many existing algorithms, including strong baseline IDAAC (Raileanu & Fergus, 2021).

Original: Combining our method with these algorithms could improve performance; investigating them can be an interesting avenue for future work.
Revised: Combining our method with these algorithms could improve performance; investigating them can be an interesting avenue for future work.

Original: G T RAINING L OSS OF S TAR GAN G ENERATOR
Revised: G T RAINING L OSS OF S TAR GAN G ENERATOR

Original: The adversarial loss is calculated as the Wasserstein GAN (Arjovsky et al., 2017) objective with gradient penalty which is deﬁned as
Revised: The adversarial loss is calculated as the Wasserstein GAN (Arjovsky et al., 2017) objective with gradient penalty which is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where ˆ x is sampled uniformly along a straight line between a pair of real and generated fake images.
Revised: where ˆ x is sampled uniformly along a straight line between a pair of real and generated fake images.

Original: Here, the gradient penalty helps to stabilize the GAN training process and has been shown to perform better than traditional GAN training (Choi et al., 2018;
Revised: Here, the gradient penalty helps to stabilize the GAN training process and has been shown to perform better than traditional GAN training (Choi et al., 2018;

Original: Arjovsky et al., 2017; Gulrajani et al., 2017).
Revised: Arjovsky et al., 2017; Gulrajani et al., 2017).

Original: We set the hyperparameter λ gp = 10 .
Revised: We set the hyperparameter λ gp = 10 .

Original: The classiﬁcation loss of real image is deﬁned as
Revised: The classiﬁcation loss of real image is deﬁned as

Original: [Equation]
Revised: [Equation]

Original:  where D cls ( c (cid:48) | x ) is the probability distribution over all clusteer labels.
Revised: (7) where D cls ( c (cid:48) | x ) is the probability distribution over all clusteer labels.

Original: Similarly, the classiﬁcation loss of fake generated image is deﬁned as
Revised: Similarly, the classiﬁcation loss of fake generated image is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: The reconstruction loss with the generator objective deﬁned as
Revised: The reconstruction loss with the generator objective deﬁned as

Original: [Equation]
Revised: [Equation]

Original: where we use the L 1 norm.
Revised: [Table]

Original: H E NVIRONMENT D ETAILS
Revised: H E NVIRONMENT D ETAILS

Original: Procgen We conducted experiments on OpenAI Procgen Cobbe et al.
Revised: Procgen We conducted experiments on OpenAI Procgen Cobbe et al.

Original: (2020) benchmark, consisting of diverse procedurally-generated environments with different action sets.
Revised: (2020) benchmark, consisting of diverse procedurally-generated environments with different action sets.

Original: This environment has been used to measure how quickly (sample efﬁciency) a reinforcement learning agent learns generalizable skills.
Revised: This environment has been used to measure how quickly (sample efﬁciency) a reinforcement learning agent learns generalizable skills.

Original: These environments greatly beneﬁt from the use of procedural content generation, the algorithmic creation of a near-inﬁnite supply of highly randomized content.
Revised: These environments greatly beneﬁt from the use of procedural content generation, the algorithmic creation of a near-inﬁnite supply of highly randomized content.

Original: The design principles consist of high diversity, tunable difﬁculty, shared action, shared observation space, and tunable dependence on exploration.
Revised: The design principles consist of high diversity, tunable difﬁculty, shared action, shared observation space [Table] and tunable dependence on exploration.

Original: Procedural generation logic directs the level layout and other game-speciﬁc details.
Revised: Procedural generation logic directs the level layout and other game-speciﬁc details.

Original: Thus, to master any of these environments, agents must learn an effective policy across all environment variations.
Revised: Thus, to master any of these environments, agents must learn an effective policy across all environment variations.

Original: We use all 16 environments available in this benchmarks.
Revised: We use all 16 environments available in this benchmarks.

Original: All environments use a discrete 15 dimensional action space which generates 64 × 64 × 3 RGB image observations.
Revised: All environments use a discrete 15 dimensional action space which generates 64 × 64 × 3 RGB image observations.

Original: Note that some environments may use no-op actions to accommodate a smaller subset of actions.[Table] [Table]
Revised: Note that some environments may use no-op actions to accommodate a smaller subset of actions.

Original: I I MPLEMENTATION D ETAILS
Revised: I I MPLEMENTATION D ETAILS

Original: Procgen Experiments For experimenting on Procgen environments, we used RLlib Liang et al.
Revised: Procgen Experiments For experimenting on Procgen environments, we used RLlib Liang et al.

Original: (2018) to implement our ARPO [Table] and baselines PPO and RAD cutour color algorithms.
Revised: (2018) to implement our ARPO, and baselines PPO and RAD cutour color algorithms.

Original: For all the agents’ policy network (model), we use a CNN architecture used in IMPALA Espeholt et al.
Revised: For all the agents’ policy network (model), we use a CNN architecture used in IMPALA Espeholt et al.

Original: (2018) which is the best performing model in Procgen benchmark (Cobbe et al., 2020).
Revised: (2018) which is the best performing model in Procgen benchmark (Cobbe et al., 2020).

Original: We use the same policy parameters for all agents for a fair comparison.
Revised: We use the same policy parameters for all agents for a fair comparison.

Original: Policy learning hyperparameter settings for all the agents (ARPO, PPO, and RAD) are set same for fair comparison and detect the effect of our proposed method (max-min adversarial objective with perturbation network).
Revised: Policy learning hyperparameter settings for all the agents (ARPO, PPO, and RAD) are set same for fair comparison and detect the effect of our proposed method (max-min adversarial objective with perturbation network).

Original: These hyperparameters are given in Table 2.
Revised: These hyperparameters are given in Table 2.

Original: Note that only the custom parameters are given here, other defaults parameter values can be found in the RLlib library Liang et al.
Revised: Note that only the custom parameters are given here, other defaults parameter values can be found in the RLlib library Liang et al.

Original: Distracting Control Experiments For experimenting on Distracting Control Suite (Stone et al., 2021) environments, we used RLlib Liang et al.
Revised: Distracting Control Experiments For experimenting on Distracting Control Suite (Stone et al., 2021) environments, we used RLlib Liang et al.

Original: (2018) to implement our ARPO, and baselines PPO and SAC algorithms.
Revised: (2018) to implement our ARPO, and baselines PPO and SAC algorithms.

Original: We use the PPO’s CNN-based policy network (model) from RLlib for ARPO, and PPO.
Revised: We use the PPO’s CNN-based policy network (model) from RLlib for ARPO, and PPO.

Original: The policy speciﬁc parameters for ARPO, and PPO are the same which are given in Table 3.
Revised: The policy speciﬁc parameters for ARPO, and PPO are the same which are given in Table 3.

Original: For SAC, we use it’s CNN-based model available in RLlib.
Revised: For SAC, we use it’s CNN-based model available in RLlib.

Original: The policy speciﬁc parameters are given in Table 4.
Revised: The policy speciﬁc parameters are given in Table 4.

Original: Note that only the custom parameters for the RLlib implementation are given here, other defaults parameter values can be found in the RLlib library Liang et al. 
Revised: Note that only the custom parameters for the RLlib implementation are given here, other defaults parameter values can be found in the RLlib library Liang et al. [Table]

Original: [Table]
Revised: [Table]

Original: Computing details .
Revised: Computing details .

Original: We used the following machine conﬁguration to run our experiments: 20 coreCPU with 256 GB of RAM, CPU Model Name: Intel(R)
Revised: We used the following machine conﬁguration to run our experiments: 20 coreCPU with 256 GB of RAM, CPU Model Name: Intel(R)

Original: Xeon(R)
Revised: Xeon(R)

Original: Silver 4114 CPU @
Revised: Silver 4114 CPU @

Original: 2.20GHz, and a Nvidia A100 GPU.
Revised: 2.20GHz, and a Nvidia A100 GPU.