---
title: "Don’t Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning"
authors:
  - "Zhecheng Yuan"
  - "Guozheng Ma"
  - "Yao Mu"
  - "Bo Xia"
  - "Bo Yuan"
  - "Xueqian Wang"
  - "Ping Luo"
  - "Huazhe Xu"
doi: "10.24963/ijcai.2022/514"
url: "https://ijcai.org/proceedings/2022/514"
pdf: "https://ijcai.org/proceedings/2022/0514.pdf"
venue: "Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence"
year: 2022
pages: "3702-3708"
track: "Main Track"
keywords:
  - "Machine Learning: Deep Reinforcement Learning"
  - "Robotics: Learning in Robotics"
  - "Machine Learning: Reinforcement Learning"
---

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

Zhecheng Yuan, Guozheng Ma, Yao Mu, Bo Xia, Bo Yuan, Xueqian Wang, Ping Luo, Huazhe Xu

*Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence*, pages 3702-3708. DOI: [10.24963/ijcai.2022/514](https://doi.org/10.24963/ijcai.2022/514)

## Abstract

One of the key challenges in visual Reinforcement Learning (RL) is to learn policies that can generalize to unseen environments. Recently, data augmentation techniques aiming at enhancing data diversity have demonstrated proven performance in improving the generalization ability of learned policies. However, due to the sensitivity of RL training, naively applying data augmentation, which transforms each pixel in a task-agnostic manner, may suffer from instability and damage the sample efficiency, thus further exacerbating the generalization performance. At the heart of this phenomenon is the diverged action distribution and high-variance value estimation in the face of augmented images. To alleviate this issue, we propose Task-aware Lipschitz Data Augmentation (TLDA) for visual RL, which  explicitly identifies the task-correlated pixels with large Lipschitz constants, and only augments the task-irrelevant pixels for stability. We verify the effectiveness of our approach on DeepMind Control suite, CARLA and DeepMind Manipulation tasks. The extensive empirical results show that TLDA improves both sample efficiency and generalization; it outperforms previous state-of-the-art methods across 3 different visual control benchmarks.

## Full text

## 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

Figure 1: Augmenting observation in a task-agnostic manner (in the middle) distracts the agent's decision, hence it will damage agent's asymptotic performance. This problem can be alleviated by taskaware data augmentation (in the bottom).

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000000_5c34d3cd16fd1f402781fa013b0c80641e12b1abad6a68a21b0c110fd8678b5e.png)

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.

Figure 2: Action Distribution. We use t-SNE to show high-dimensional actions employed by the same agent. The grey dots are the actions given the observations without augmentation (w/o aug); the blue dots (a) and orange dots (c) are the actions given the same observations under strong (random conv) and weak augmentation (random shift), respectively. The visualized results demonstrate that there is a significant action distribution shift under strong augmentation, while under weak augmentation, the policy is closer to the original distribution. The red dots (b) are TLDA under strong augmentation, which comes up with an action distribution that remains similar to the grey dots.

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000001_270219cc4cba2da0429881486843e0c0f87c0e25c7ff6c91c2ee9be09655de5d.png)

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.

Figure 3: Q-estimation error. Top: We measure the Q-estimation mean square error of the different augmented observation vs.the nonaugmented observation. The blue bar and the red bar are the error between strong augmented data and TLDA-augmented data vs nonaugmented data, separately. It shows that TLDA can significantly reduce the Q-estimation error to alleviate the high-variance estimation problems. Bottom: The distribution of Q-estimation. TLDA comes up with a closer Q-estimation distribution with the original one.

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000002_dd0d34f23218d7b4e8afa776a6428e775cd4406991c90976adc4394435f72454.png)

## 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:

<!-- formula-not-decoded -->

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

<!-- formula-not-decoded -->

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:

<!-- formula-not-decoded -->

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) :

<!-- formula-not-decoded -->

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 :

Figure 4: Overview of TLDA. This figure shows two examples and the pipeline of TLDA. The agent generates the K-matrix in a frame, and then preserves the larger Lipschitz constant areas under strong augmentation. The preserved areas are highlighted in the K-matrix.

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000003_74e0db2a6d8c91ea6e64d4bf06b4ec7879143b4f7d7e4b729bea510f88ee0d44.png)

<!-- formula-not-decoded -->

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:

<!-- formula-not-decoded -->

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:

<!-- formula-not-decoded -->

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 :

<!-- formula-not-decoded -->

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:

<!-- formula-not-decoded -->

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)]:

<!-- formula-not-decoded -->

with

<!-- formula-not-decoded -->

Figure 5: Sample efficiency in training environment. We compare TLDA, SVEA, and random patch under two kinds of augmentations. Top row and Bottom Row are corresponding to Random Conv and Random Overlay training curves of the episode return respectively. TLDA (red line) shows better sample efficiency on the training period. Mean and standard deviation of 5 runs.

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000004_aff1a07f0bc5e1e18a1dd43f402546372c60a71e8ca25846c403e9f03bc00c01.png)

## 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.

![Image](https://ijcai.org/proceedings/2022/fulltext/514/image_000005_641586773c7fbd996088b1fd11a7aedda0d1c97be577032f546e1107e52c732f.png)

Table 1: DMC-GB Generalization Performance. Episode return in test environments. The agents are trained on a fixed environment and evaluated on random colors (Bottom) and video backgrounds (Top) these two unseen test environments. Mean and std.deviation of 5 runs.

| Setting   | DMControl                                                                                                                                       | DrQ                                                                      | PAD                                                                   | SVEA (conv)                                                             | SVEA (overlay)                                                       | SODA (conv)                                                             | SODA (overlay)                                                        | TLDA (conv)                                                           | TLDA (overlay)                                                       |
|-----------|-------------------------------------------------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------|-----------------------------------------------------------------------|-------------------------------------------------------------------------|----------------------------------------------------------------------|-------------------------------------------------------------------------|-----------------------------------------------------------------------|-----------------------------------------------------------------------|----------------------------------------------------------------------|
|           | Cartpole-Swingup Walker-Stand Walker-Walk Ball in cup-Catch Cheetah-Run Cartpole-Swingup Walker-Stand Walker-Walk Ball in cup-Catch Cheetah-Run | 485±105 873±83 682±89 318±157 102±30 586±52 770±71 520±91 365±210 100±27 | 521±76 935±20 717±79 436±55 206±34 630±63 797±46 468±47 563±50 159±28 | 606±85 795±70 612±144 659±110 292±32 837±23 942±26 760±145 961±7 264±51 | 782±27 961±8 819±71 871±106 249±20 832±23 933±24 749±61 959±5 273±23 | 474±143 903±56 635±48 539±111 229±29 831±21 930±12 697±66 892±37 294±34 | 758±62 955±13 768±38 875±56 223±32 805±28 893±12 692±68 949±19 238±28 | 607±74 962±15 873±34 892±68 356±52 748±40 919±24 743±83 932±32 371±51 | 671±57 973±6 868±63 863±56 366±57 760±60 947±26 823±58 930±40 358±25 |

## 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.

Table 2: CARLA Driving. We report the success rate for reaching 100m distance under the unseen weather during 250 episodes across 5 seeds for each weather. (50 episodes for each seed)

| Setting        | DrQ   | SVEA   | Ours   |
|----------------|-------|--------|--------|
| Training       | 24%   | 49%    | 52%    |
| Wet Noon       | 0.8%  | 8.8%   | 18%    |
| SoftRain noon  | 0.4%  | 1.2%   | 7.6%   |
| Wet Sunset     | 0.8%  | 1.6%   | 9.2%   |
| MidRain Sunset | 0.0%  | 5.2%   | 12%    |

## 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.

Table 3: DMC Manipulation Tasks. M in the Setting column means: Modified. TLDA can better focus on the aim objects in the noisy and colorful visual backgrounds.

| Task   | Setting   | DrQ      | SVEA   | Ours    |
|--------|-----------|----------|--------|---------|
| Reach  | Training  | 136 ±20  | 49 ±48 | 124 ±32 |
| Reach  | MArm      | 68 ±20   | 21 ±25 | 55 ±21  |
| Reach  | MPlatform | 0.8 ±1.3 | 24 ±25 | 89 ±40  |
| Reach  | M Both    | 1 ±2     | 13 ±14 | 36 ±25  |
| Push   | Training  | 141 ±47  | 42 ±40 | 109 ±27 |
| Push   | MArm      | 88 ±52   | 21 ±16 | 60 ±43  |
| Push   | MPlatform | 4 ±1     | 34 ±28 | 95 ±33  |
| Push   | M Both    | 5 ±1     | 32 ±20 | 56 ±42  |

## 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.

## References

1. [Cobbe et al., 2019] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying general- ization in reinforcement learning. In International Con- ference on Machine Learning, pages 1282-1289. PMLR, 2019.
2. [Cobbe et al., 2020] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International con- ference on machine learning, pages 2048-2056. PMLR, 2020.
3. [Dosovitskiy et al., 2017] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference on robot learning, pages 1-16. PMLR, 2017.
4. [Fan et al., 2021] Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Anima Anandku- mar. Secant: Self-expert cloning for zero-shot generaliza- tion of visual policies. arXiv preprint arXiv:2106.09678, 2021.
5. [Ferns et al., 2011] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1662-1714, 2011.
6. [Finlay et al., 2018] Chris Finlay, Jeff Calder, Bilal Abbasi, and Adam Oberman. Lipschitz regularized deep neural networks generalize and are adversarially robust. arXiv preprint arXiv:1808.09540, 2018.
7. [Grill et al., 2020] Jean-Bastien Grill, Florian Strub, Flo- rent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhao- han Daniel Guo, Mohammad Gheshlaghi Azar, et al. Boot- strap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733, 2020.
8. [Hansen and Wang, 2021] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Confer- ence on Robotics and Automation (ICRA), pages 13611- 13617. IEEE, 2021.
9. [Hansen et al., 2020] Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. arXiv preprint arXiv:2007.04309, 2020.
10. [Hansen et al., 2021] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vi- sion transformers under data augmentation. arXiv e-prints, pages arXiv-2107, 2021.
11. [Kalashnikov et al., 2018] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vin- cent Vanhoucke, et al. Scalable deep reinforcement learn- ing for vision-based robotic manipulation. In Conference on Robot Learning, pages 651-673. PMLR, 2018.
12. [Kostrikov et al., 2020] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Reg- ularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.
13. [Laskin et al., 2020] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems, 33, 2020.
14. [Mnih et al., 2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidje- land, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533, 2015.
15. [Packer et al., 2018] Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. arXiv preprint arXiv:1810.12282, 2018.
16. [Pauli et al., 2021] Patricia Pauli, Anne Koch, Julian Berberich, Paul Kohler, and Frank Allgower. Training robust neural networks using lipschitz bounds. IEEE Control Systems Letters, 2021.
17. [Raileanu et al., 2021] Roberta Raileanu, Maxwell Gold- stein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Au- tomatic data augmentation for generalization in reinforce- ment learning. Advances in Neural Information Process- ing Systems, 34, 2021.
18. [Shorten and Khoshgoftaar, 2019] Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmen- tation for deep learning. Journal of Big Data, 6(1):1-48, 2019.
19. [Tobin et al., 2017] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural net- works from simulation to the real world. In IROS, pages 23-30. IEEE, 2017.
20. [Tunyasuvunakool et al., 2020] Saran Tunyasuvunakool, Al- istair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continu- ous control. Software Impacts, 6:100022, 2020.
21. [Zhang et al., 2018] Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study on overfitting in deep reinforcement learning. arXiv preprint arXiv:1804.06893, 2018.
22. [Zhang et al., 2020] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742, 2020.
23. [Zhu et al., 2017] Yuke Zhu, Roozbeh Mottaghi, Eric Kolve, Joseph J Lim, Abhinav Gupta, Li Fei-Fei, and Ali Farhadi. Target-driven visual navigation in indoor scenes using deep reinforcement learning. In ICRA, pages 3357-3364. IEEE, 2017.
