Don’t Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning

1 Introduction

Deep Reinforcement Learning (DRL) from visual observations has carved out brilliant paths in many domains such as video games [Mnih et al., 2015], robotics manipulation [Kalashnikov et al., 2018], and visual navigation [Zhu et al., 2017]. However, it remains challenging to obtain generalizable policies across different environments with visual variations due to overfitting [Zhang et al., 2018].

Data Augmentation [Shorten and Khoshgoftaar, 2019] and Domain Randomization [Tobin et al., 2017] based approaches are widely used to learn generalizable visual representations. However, recent work [Hansen et al., 2021] find that in visual RL, there is a dilemma: heavy data augmentations are vital for better generalization, but it will cause a significant decrease in both the sample efficiency and the training stability. One of the main reasons is that data augmentation conventionally perform pixel-level transformation, where each pixel is transformed in a task-agnostic manner. However, each pixel in the observation has different relevance to the task and the reward function. Hence, it is worth rethinking data augmentation in the new context of visual RL.

∗Corresponding Author

To better understand the effect of data augmentation in visual RL, we visualize the action distribution output from policies trained with various data augmentation choices in Figure 2. We find that the agent's actions vary dramatically when faced with different augmentation methods. Specifically, when weak augmentation such as shifting is applied, the action distribution remains closer to the original distribution that has no augmentation (Figure 2(c)); however, when strong augmentation e.g., random convolution is applied, the action distribution drastically changes (Figure 2(a)) and the Q-estimation yields the discrepancy with the un-augmented data, as shown in Figure 3. This comparison reveals the severe problem that causes instability when data augmentation is applied blindly without knowing the task information.

In this work, we propose a task-aware data augmentation method in visual RL that learns to augment the pixels less correlated to the task, namely Task-aware Lipschitz Data Augmentation (TLDA), as shown in Figure 1. A desirable quality for such a method is that it maintains a stable policy output even on augmented observations. Following this insight, we introduce the Lipschitz constant that measures the relevance between the pixel and the task, then guides the augmentation strategy. Specifically, we first impose a perturbation on a certain pixel, and calculate the corresponding Lipschitz constant for the pixel via the policy change before and after the perturbation. Then, to avoid the occurrence of drastic policy changes, we treat the pixels with larger Lipschitz constant as the task-relevant ones and avoid augmenting them. Therefore, the output could be more stable while keeping the diversity of augmented data.

We conduct experiments on 3 benchmarks: DMControl Generalization Benchmark (DMC-GB) [Hansen and Wang, 2021], CARLA [Dosovitskiy et al., 2017], and DMControl manipulation tasks [Tunyasuvunakool et al., 2020]. We train agents in a fixed environment and evaluate on the unseen environments. Extensive experiments show that TLDA outperforms the prior state-of-the-art methods due to more stable and efficient training and robust generalization performance.

Our main contributions are summarized as follows: (i) We propose Task-aware Lipschitz Data Augmentation (TLDA), which can be implemented on downstream visual RL algorithm easily without adding auxiliary objectives or additional learnable parameters; (ii) We provide theoretical understanding and experiments to show TLDA can alleviate the action distribution shift and high variance Q-estimation problems effectively; (iii) TLDA achieves competitive or better sample efficiency and generalization ability than previous state-ofthe-art methods in 3 different kinds of benchmarks.

2 RelatedWork

Generalization in RL. Researchers have been investigated RL generalization from various perspectives, such as visual appearances [Cobbe et al., 2019], dynamics [Packer et al., 2018] and environment structures [Cobbe et al., 2020]. In this paper, we focus on generalization over different visual appearances. Two popular paradigms are proposed to address the overfitting issue in current visual RL research. The first is to regard generalization as a representation learning problem. Bi-simulation metric [Ferns et al., 2011] is implemented to learn robust representation features [Zhang et al., 2020]. The other paradigm is to design auxiliary tasks. SODA [Hansen and Wang, 2021] adds a BYOL-like [Grill et al., 2020] architecture and introduces an auxiliary loss which encourages the representation to be invariant to task-irrelevant properties of the environment. In contrast to the previous efforts, our method does not require to employ a specific metric to learn representation, nor to introduce additional modules.

Data Augmentation for RL. Data Augmentation is an efficient method to improve the generalization of visual RL. RAD [Laskin et al., 2020] compares different data augmentation methods and reveals that the benefits of different augmentation methods to RL tasks are not the same. SECANT [Fan et al., 2021] mentions that weak augmentation can improve sample efficiency but not generalization ability. It also shows that the simple use of strong augmentation is likely to cause training divergence, though generalization ability is improved. Automatic data augmentation is proposed in [Raileanu et al., 2021] to make better use of data augmentation. We advocate this paradigm and believe that one crucial factor for improving sample efficiency and generalization lies in the design of data augmentation, namely, how we can diversify the input as much as possible while maintaining the invariance of output. We show that how strong augmentation affects action distribution shifts and causes high variance of Q estimation, and illustrate that our approach is effective in alleviating these two problems.

3 Preliminaries

We consider learning in a Markov Decision Process (MDP) formulated by the tuple ⟨S,A, r,P, γ⟩ where S is the state space, A is the action space, r : S × A ↦→ R is a reward function, P (st+1 | st, at) is the state transition function, γ ∈ [0, 1) is the discount factor. The goal is to learn a policy π∗ to maximize the expected cumulative return π∗ = argmaxπ Eat∼π(·|st),st∼P [∑Tt=1 γ t r (st, at)], starting from an initial state s0 ∈ S and following the policy πθ (· | st) which is parameterized by a set of learnable parameters θ. Meanwhile, we expect the learned policy π∗ θ can be well generalized to new environments, which have the same structure and definition of the original MDP, but with different observation space O constructed from the same state space S.

3.1 Data Augmentation

Definition 1 (Optimality-Invariant State Transformation) Given an MDP M, we define an augmentation method ϕ : O → O′ as an optimality-invariant transformation if ∀o ∈ O, a ∈ A, ϕ(o) ∈ O′ , where O′ is a new set of observation satisfies:

A desirable quality for data augmentation is to satisfy the form of Optimality-Invariant State Transformation while distortion or distracting noise is added to the observation.

3.2 Lipschitz Constant

The Lipschitz constant is frequently utilized to measure the robustness of a model, we introduce Lipschitz continuity of the policy here. A function f: Rn → Rm is Lipschitz continuous on X ⊆ Rn if there exists a non-negative constant K ≥ 0 such that

The smallest such K is called the Lipschitz constant of f [Pauli et al., 2021].

Definition 2 (Lipschitz constant of the policy) Assume the state space is equipped with a distance metric d(·, ·). Under a certain augmentation method ϕ, the Lipschitz constant of a policy π is defined as follows:

where DTV (P||Q) = 12 ∑a∈A |P(a) - Q(a)| is the total variation distance between distributions. If Kπ is finite, the policy π is Lipschitz continuous.

For a certain model, a smaller Lipschitz constant represents higher stability against the variance of input [Finlay et al., 2018]. The following proposition illustrates that the estimation error of Q-value can be bounded by Lipschitz constant:

Proposition 1 We consider an MDP M, a policy π and an augmentation method ϕ. Suppose the rewards are bounded by rmax and state space is equipped with a distance metric d(·, ·), such that ∀a ∈ A, ∀s ∈ S, |r(s, a)| ≤ rmax, the following inequality holds, where ∥d(ϕ)∥∞ = sups∈S d(ϕ(s), s) :

This proposition indicates that if a smaller Lipschitz constant under one specific augmentation is acquired, we will have a tighter bound of the Q-value estimation with a lower variance while implementing data augmentation.

4 Method

To maintain the training stability and improve the generalization ability, we propose: Task-aware Lipschitz Data Augmentation (TLDA), an efficient and general task-aware data augmentation method for visual RL.

4.1 Construct the K-matrix

We first calculate the Lipschitz constant from perturbed input images. By using a kernel to perturb the original image o ∈ RH×W , we obtain the perturbed image denoted as A(o). Next, we choose the pixels centered with the location (i, j) of A(o) as in the Eq (5), denoted as Φ(o, i, j). Specifically, we use the Hadamard product ⊙ to choose the perturbed pixels around location(i, j) by an image mask M(i, j) ∈ (0, 1)H×W :

To derive the Lipschitz constant, we use the notation d(Φ(o, i, j), o) to represent the distance between input o and Φ(o, i, j) under metric d(·, ·). As in Definition 2, for a given observation o, the Lipschitz constant of the pixel (i, j) can be computed as follows:

where the numerator can be interpreted as distance between two action distributions: π(· | Φ(o, i, j)), π(· | o), and the denominator is the distance between the original observation and the perturbed one. With the per-pixel Lipschitz constant in hand, we then construct the matrix that can reflect the taskrelevance information and be applied on the whole observation. By arranging Kπ ij into a matrix which have the same size as o following Eq (7), we denote this matrix as the K-matrix:

We aim to capture the task-related locations with large Lipschitz constants which tend to cause high variance in the policy/value output during the same level of perturbation.

4.2 Task-Aware Lipschitz Augmentation (TLDA) with the K-matrix

Intuitively, data augmentation operations should not modify the task-related pixels indicated by large Lipschitz constants. We follow this intuition and propose a simple yet effective way to decide which areas can be modified. We use the mean value of the K-matrix as a threshold, and binarize the K-matrix by the following way, where N is the number of pixels (H ×W), Kmean = 1N × ∑ij Kπ ij :

The obtained mask M K is used to decide which pixels can be augmented. For any data augmentation method o′ = Aug(o), we apply the following operation:

We note that the output ˜ o is only modified in the areas that have low relevance to the task.

As mentioned above, TLDA tends to preserve the pixels with large Kπ ij and augment only the pixels associated with the small ones, which adds an implicit constraint to maintain the stable output of the policy and value network. Hence, it echoes with the Optimality-Invariant State Transformation as in Definition 1. Figure 4 demonstrates the overall framework of TLDA. During the training process, the K-matrix is calculated on the fly against every training step on augmented observations. Take cutout (adding a black patch to the image) in Figure 4 as an example, since the corresponding K-matrix shows that the upper part of the robot's body features large Lipschitz constants, therefore, blindly augmenting the image might touch the pixels in this area and cause catastrophic action/value changes. In contrast, TLDA preserves the critical parts of the original observations indicated by K-matrix.

4.3 Reinforcement Learning with TLDA

We use soft-actor-critic (SAC) as the base reinforcement learning algorithm for TLDA. Similar to previous work, we also include a regularization term RQ(θ) to the SAC critic loss JQ (θ) to handle augmented data. Our critic loss LQ(θ) is as follows, where saug t is calculated by Eq (9), and ˆQ (st, at) = r (st, at) + γEst+1∼P [V (st+1)]:

with

5 Experiment

In this section, we explore how TLDA can affect the agent's sample efficiency and generalization performance. We compare our method with other baselines on a wide spectrum of tasks including DeepMind control suite, CARLA simulator, as well as DeepMind Manipulation tasks.

5.1 Evaluation on DeepMind Control Suite

Setup. For comparison, we mainly consider two augmentation ways applied in the prior state-of-the-art methods: random convolution (passing input through a random convolutional layer) and random overlay (linearly combining the observation o with the extra image I, ϕ(o) = αo + (1 - α)I). Baselines. We benchmark TLDA against the following state-of-the-art methods: (1) DrQ [Kostrikov et al., 2020]: SAC with weak augmentation (random shift); (2) PAD [Hansen et al., 2020]: adding an auxiliary task for adapting to the unseen environment; (3) SODA [Hansen and Wang, 2021]: maximizing the mutual information between latent representation by employing a BYOL-like [Grill et al., 2020] architecture; (4) SVEA [Hansen et al., 2021]: modifying the form of Q-target.

Sample efficiency under strong augmentations. We compare the sample efficiency with SVEA to exhibit the effectiveness of TLDA. We also include another baseline that preserves random patches from the un-augmented observation as opposed to TLDA that preserves task-related parts. We call this baseline random patch. By contrast, SVEA only uses the strong augmentation method but retains no raw pixel. Figure 5 demonstrates that TLDA achieves better or comparable asymptotic performance in the training environment than baselines on DM-control suite while having better sample efficiency. The results also indicate that random patch will hinder the performance in some tasks. We reckon that since random patch does not have any pixel-to-task relevance knowledge, it inevitably destroys the image's integrity and even leads to further distortion to the observations after data augmentation. Therefore, blindly keeping the original observation's information cannot improve the agent's training performance. It is the retention of areas with larger Lipschitz constants, instead of random original areas, that boosts the sample efficiency of training agents.

Generalization Performance. We evaluate the agent's generalization ability on two settings from DMControlGB [Hansen and Wang, 2021]. Results are shown in Table 1. TLDA outperforms prior state-of-the-art methods in 7 out of 10 instances. The agent trained with TLDA is able to acquire a good robust policy in different unseen environments. Meanwhile, we notice that prior methods are sensitive to augmentation methods, which makes their testing performance varies dramatically. On the contrary, our method with task-aware observations is more stable and not susceptible to this issue.

Effect on Action Distribution and Q-estimation. In this section, we analyze how TLDA influences the output of the policy and value networks. Given a DrQ agent trained in the original environment, we assess the Q-value estimation and the action distribution under different augmentation. To get a better understanding of this issue, we visualize the action distribution of the agent under different augmentation methods, as shown in Figure 2. For weak augmentation, although its action distribution is closest to the un-augmented one (Figure 2(c)), it cannot improve generalization, as shown in Table 1(DrQ). Strong augmentation, on the other hand, will cause an obvious distribution shift (Figure 2(a)), thus significantly hindering the training process. TLDA has a closer action distribution than simply applying strong augmentation (Figure 2 (b)) by using the Lipschitz constant to identify and preserve the task-aware areas. Furthermore, as shown in Figure 3, we find that the Q-estimation of TLDA has a lower variance than that of naively applying strong augmentation. These two results illustrate that TLDA has the potential to achieve higher sample efficiency in training and learn a more robust policy to perform well in unseen environments.

5.2 Evaluation on Autonomous Driving in CARLA

To further evaluate the TLDA's performance, we apply this method in the tasks with more realistic observations: autonomous driving in the CARLA simulator. In our experiment, we use one camera as our input observation for driving tasks, where the goal of the agent is to drive along a curvy road as far as possible in 1000 time-steps without colliding with the moving vehicles, pedestrians and barriers. We adapt the reward function and train an agent under the weather with the same setting from previous work [Zhang et al., 2020]. We find that our method achieves the best training sample efficiency. For generalization, CARLA provides different weather conditions with built-in parameters. We evaluate our method in 4 kinds of weather with different lighting conditions, realistic raining and slipperiness. Results are in Table 2, where we choose the success rate for reach 100m distance as the driving evaluation metric. TLDA outperforms all base algorithms in both sample efficiency and generalization ability with a more stable driving policy.

5.3 Evaluation on DMCManipulation Tasks

Robot manipulation is another set of challenging and meaningful tasks for visual RL. DM control [Tunyasuvunakool et al., 2020] provides a set of configurable manipulation tasks with a robotic Jaco arm and snap-together bricks. We consider two tasks for experiments: reach and push. All the agents are trained on the default background and evaluated on different colors of arms and platforms. The generalization performance are shown in Table 3. The results show that our method can be adapted to the unseen environments more appropriately. The Modified Platform and Modified Both setting are challenging for agents to discern the target objects from the noisy backgrounds. SVEA under strong data augmentation suffers from instability and divergence for training, while TLDA can augment the pixel in a task-aware manner, thus further maintaining the training stability. Despite that DrQ shows better training performance, it barely generalizes to the environments with different visual layouts. In summary, sample efficiency and generalization performance contribute to exhibiting the superiority of the proposed algorithm.

6 Conclusion

In this paper, we propose Task-aware Lipschitz Data Augmentation (TLDA) for visual RL, which can reliably identify and augment pixels that are not strongly correlated with the learning task while keeping task-related pixels untouched. This technique aims to provide a principled mechanism for boosting the generalization ability of RL agents and can be seamlessly incorporated into various existing visual RL frameworks. Experimental results on 3 challenging benchmarks confirm that, compared with the baselines, TLDA not only features higher sample efficiency but also helps the agents generalize well to the unseen environments.