Style-Agnostic Reinforcement Learning

The main target of the domain generalization in RL can be summarized as training an agent to learn a robust policy that can be generalized to unseen environments. This allows RL algorithms to be applied in more realistic situations because agents are often tested in different environments from the training stage. One example is deploying a policy learned from the simulation to the real world in the robot manipulation task.

Data randomization is a promising technique for such domain generalization in many cases [Robot2, Robot3]. However, it is difficult to build an accurate and practical simulator that enables using data randomization. Visual augmentation, on the other hand, is much easier to apply as it is based on simple image transformations. Laskin & Lee et al. [RAD], for example, demonstrated that simply using data augmentation, such as random cropping or gray scaling, is indeed helpful in improving the generalization capacity of RL agents. Also, Yarats & Kostrikov et al. [DrQ] suggested using regularization terms for stabilizing the model training when using data augmentation.

However, although data augmentation is potentially effective, it has several limitations. For example, a naïve choice of the augmentation type may degrade the generalization performance [RAD]. Applying cropping to an essential part of the image may confuse the agent, or training the model to produce the same action from a rotated image may be unreasonable. Here, we present a method for domain generalization by diversifying the training examples without requiring a complex strategy for data augmentation. The generator in the SAR model generates new feature examples having different styles and helps the agents with learning style-agnostic representations.

Adversarial feature learning has become popular for domain generalization in computer vision tasks [SONG, ADV, MMDAAE, featureAttack, SagNet]. Li et al. [MMDAAE] showed that adversarial objectives help a model learn universal feature representations across different domains. Furthermore, Nam & Lee et al. [SagNet] proposed a method of reducing the style gap for domain generalization in the image classification task. Inspired by this work, we investigate the adversarial feature learning for RL agents, but with a simpler training procedure, i.e., without dividing training phases or considering the environment style’s classes.

We note that adopting adversarial training for RL is not new [RARL, DARL]. To our best knowledge, however, exploiting adversarial learning to the latent features in RL framework and the min-max game scheme is not presented before. Especially, our method can be interpreted as domain randomization beyond pixel space. Mixing styles with linear interpolation for representation learning in RL setting has been proposed in the earlier work [MixStyle]. However, unlike in the previous study, the style perturbation generator in SAR produces new synthetic styles that will not be seen with a simple interpolation. The adversarial examples help the actor extract style-agnostic embeddings without any label of styles and, finally, learn a robust policy for unseen environments.

RL agents interact and get trained with the world environment within a Markov decision process, which is defined as a tuple of $($state space $S$, action space $A$, transition probability $P$, reward space $R$, and discount factor $\gamma )$; at every timestep $t$, the agent observes a state $s_t\in S$ and takes an action $a_t\in A$ from its policy $\pi \left(a_t|s_t\right)$ [MDP]. Then, the agent is rewarded with $r_t\in R$, and moves to the next state $s_{t+1}$ sampled from the transition probability $P(s_{t+1}|s_t,a_t)$.

The policy of the agent is optimized to maximize the discounted sum of rewards $G_t={\sum }_{k=t}^{\infty }{\gamma }^k{r}_k$. With given state $s_t$, the value of the state $V\left(s_t\right)$ is estimated as $E_{\tau \sim \pi }\left[G_t|s_t\right]$ and the value of the state-action $Q(s_t,a_t)$ is computed as $E_{\tau \sim \pi }[G_t|s_t,a_t]$, with trajectory $\tau$ sampled from the policy $\pi$.

With deep RL algorithms, the policy $\pi$ gets parameterized by a set of learnable parameters $\psi$, and value function $V$ or $Q$ is optimized with network parameter $\varphi$. Also, especially for visual-based RL, since the images only offer partial observations, Mnih et al. [stackObs] has proposed that defining the state $s_t$ as a stacked consecutive image frames $(o_{t-k},o_{t-k+1},\dots ,o_t)$, where $O$ is a high-dimensional image space and $o\in O$, is effective.

Our method is based on the concept of style transfer, which was proven to be successful in generating images with new styles [CIN, StyleGAN]. The style perturbation layer shifts the style of observations $z$ with the generated perturbation mean ${\beta }_{adv}\left(z\right)$ and variance ${\gamma }_{adv}\left(z\right)$, to build style-perturbed feature map $z_{adv}$, or StylePerturb(z), with the following equation:

Then, the SAR agent should take the same action from $z_t$ and $z_{adv,t}$ to be robust among different environments, as the perturbed feature indicates an observation with different styles but the same semantics, e.g., in Procgen, the same player, enemies, and items, but shifted texture of the background image, the colors of projectiles, and the shapes of obstacles. We will further explain the objectives to achieve this generalization.

Primarily, the policy network is updated via PPO or SAC objectives. Thus, the actor loss of SAR is adopted from Equation 2 with PPO baseline or from Equation 4 when using SAC. We will denote this loss be $L_{actor}^{\circ }$. Also, for the critic loss, as suggested in RAD [RAD], we adopt the critic objective of PPO or SAC interchangeably, denoted as $L_{critic}^{\circ }$.

Another big goal of SAR is to be robust to different environments. Therefore, the agent should learn its policy by minimizing the difference between the distributions of actions from the style-perturbed features $z_{adv,t}$ and the original ones $z_t$. By leveraging KL-divergence, we can calculate the objective as $L_{div}=KL\left[\pi \right(\cdot |z_t\left)||\pi \right(\cdot |z_{adv,t}\left)\right]$. Integrating this with a weight coefficient $\lambda$, the objective for the SAR actor module can be written as:

On the other hand, the generator participates in the min-max game in another manner: to maximize the differences between the action distributions. This module is trained with the objective of the same $L_{div}$ but with a converted sign. Unlike the previous works using class label information of the environment style [DARL] or additional heavy background images [SVEA], the objectives for the robust policy (i.e., adversarial loss) do not demand any secondary labors. Hence, the overall goals for the generator can be formalized as:

where ${\lambda }^{\prime }$ can be different coefficient from that of actor objective.

Finally, the critic gets updated to guide the actor to optimize its policy to maximize the value function. Meanwhile, we observed that the sharing critic network, for predicting the value for both style-perturbed features and the original ones, does not bring a huge difference in the performance from decoupling the critic network but lighter training computation. Instead, we add a regularization term $G_{critic}$ for the value function, to minimize the difference between the value predicted from the adversarial example, i.e., $\left(V_{\varphi }\right(z_t)-V_{\varphi }(z_{adv,t}){)}^2$, which helps stabilization. Thus, the critic’s objectives can be computed as follows:

with hyperparameter $\kappa$¹¹1The values used for each hyperparameters $\lambda$, ${\lambda }^{\prime }$, $\kappa$ in the experiment are described in the supplementary material..

Here, we present the pseudo-code of the SAR algorithm. As depicted in Figure 1, to maximize the effect of style transfer, we design the $z_t$ to pass a Style Mixing module and a Style Perturbation module with two divided branches. In the Style Mixing module, the styles of observations in the mini-batch get interpolated with Equation 9. In Style Perturbation module, on the other hand, the styles of observations are shifted with new styles generated from the generator network with Equation 10.

With two different features $z_t$ and $z_{adv,t}$, the SAR agent predicts two different action distributions ${\pi }_t$ and ${\pi }_{adv,t}$. The difference between these predictions $L_{div}$ is computed, and it gets interpreted in two different ways: by the generator to produce more unfamiliar styles and by the actor to make its policy more robust.

In this section, we exhibit the experiment results for the generalization performance of our SAR model on Procgen [ProcGen] and Distracting Control Suite [DistCS] benchmarks. Recently, these benchmarks have become a standard for measuring the generalization performance of visual-based RL algorithms [RAD, DrQ, DrAC, SODA, SVEA]. These contain reasonably challenging and diverse tasks, which are highly relevant to real-world robot learning.

While the Procgen benchmark is with a discrete action space, the Distracting Control Suite presents continuous control tasks. We exploited PPO as the basic baseline on the Procgen, and SAC as the basic baseline on the Distracting Control Suite, showing that the SAR algorithm can be applied to both on-policy and off-policy algorithms. Figure 2 visualizes some examples of training and test environments in the two different benchmarks.

The baselines are six visual-learning RL algorithms showing state-of-the-art results on Procgen. PPO [PPO] is the vanilla on-policy RL baseline, and RAD [RAD] uses data augmentation on top of PPO. We performed random translation (denoted as ‘trans’) and random color cutout (denoted as ‘color’) for RAD, as they are reporting the best performance. Among many advanced algorithms on RAD, UCB DrAC[DrAC], and Meta DrAC [DrAC] are chosen to be compared with our method among three variants of DrAC; the former one presents the best performance among the variants. Mixstyle [MixStyle] exploits the style mixing, and DARL [DARL] uses an adversarial objective for regularization with style $label$s. ²²2We reproduced all the results of the baselines. The results showed better than the reported performance in several tasks as more training steps [RAD, DrAC, MixStyle].

In this experiment, we purposely tested different baselines from Procgen to compare SAR with various algorithms. SAC [SAC] is the vanilla off-policy RL algorithm, and CURL [CURL] uses a contrastive objective for representation learning on top of SAC. DrQ [DrQ] was chosen as the representative baseline using the data augmentation with additional regularization terms. PAD [PAD] adapts to a new test environment using self-supervision.³³3We reproduced all the results of the baselines and applied ‘trans’ to DrQ. The results with zero noise well match the reported performances in most cases [SAC, CURL, DrQ].⁴⁴4The performance of PAD differs from the reported value because of the simultaneous application of natural video backgrounds, color noise, and camera angle noise.

This ablation study answers two questions regarding (1) whether the generator module helps generalization performance and (2) generalization term $G_V$ is important for stabilization. Comparing SAR to PPO baseline, MixStyle using only style mixing, and SAR ($\lambda =0$), Table 4 shows that the adversarial objective helps improve the mean of the performances. Comparing SAR to SAR ($\kappa =0$), Table 4 shows that $G_V$ helps stabilize the variance of the performances. ⁵⁵5Note that the results in Table 4 are slightly different from Table 1, as we applied exponential moving average smoothing before averaging with coefficient value 0.98.

Furthermore, we qualitatively examine the features extracted from the encoder learned with the SAR objectives. The feature $z$ we analyzed is from encoded features before entering the AdaIN layer to exclude the effect of explicit style mixing.

We demonstrate three analyses on the embedding:

• GradCAM [GradCAM] visualization for the high-level understanding interpretation.

• Reconstruction images from the feature maps.

• t-SNE [t-SNE] for analyzing the latent representations.