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

摘要：

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. To verify the effectiveness of TLDA, we conduct extensive experiments on DeepMind Control suite, CARLA and DeepMind Manipulation tasks, showing that TLDA improves both sample efficiency in training time and generalization in test time. It outperforms previous state-of-the-art methods across the 3 different visual control benchmarks1.

视觉强化学习(RL)的关键挑战之一是学习能够推广到看不见的环境的策略。近年来，旨在增强数据多样性的数据增强技术在提高学习策略的泛化能力方面表现出了良好的性能。然而，由于RL训练的敏感性，幼稚地应用数据增强以任务无关的方式对每个像素进行变换，可能会出现不稳定，损害样本效率，从而进一步加剧泛化性能。这种现象的核心是面对放大的图像时动作分布的分歧和高方差值估计。为了缓解这一问题，我们提出了一种基于任务的Lipschitz数据增强算法(TLDA)，它显式地识别出具有较大Lipschitz常数的与任务相关的像素，并且只增加与任务无关的像素。为了验证TLDA的有效性，我们在DeepMind Control Suite、Carla和DeepMind操作任务上进行了大量的实验，结果表明TLDA在训练时间和测试时间上都提高了样本效率和泛化能力。在3个不同的视觉控制基准中，它的表现优于以前最先进的方法1。

问题

在视觉RL中，存在一个两难境地：大量的数据增强对于更好的泛化是至关重要的，但它会导致样本效率和训练稳定性的显著降低。
原因
- 数据增强通常执行像素级图像转换，其中每个像素以与任务无关的方式进行转换。
- 然而，在视觉RL中，观察中的每个像素与任务和奖励函数具有不同的相关性。
- 如上图：task-agnostic的数据增强会影响agent的决策，损害agent的学习性能

本文工作

提出了一种视觉RL中的任务感知数据增强方法，该方法学习增加与任务相关性较小的像素
Task-aware Lipschitz Data Augmentation (TLDA) 任务感知的利普希兹数据增强
该算法无需增加辅助目标或额外的可学习参数，即可在任何下游视觉RL算法上实现。
理论分析和实验结果表明，TLDA可以有效地缓解动作分布漂移和高方差Q-估计问题
在3种不同的基准测试中，TLDA的样本效率和泛化能力与以前的最先进方法相当或更好

方法简述

通过Lipschitz常数的思想引入策略的Lipschitz常数公式
数据增强是对像素的变化，因此引入衡量像素和任务之间相关性的Lipschitz常数，并指导了增强策略。
对一像素施加扰动，通过扰动前后的策略变化计算像素对应的Lipschitz常数
将Lipschitz常数较大的像素点视为与任务相关的像素点，并避免对其进行增强。（避免策略的剧烈变化）

方法介绍

最优不变状态变换

数据增强的一个理想效果是满足最优不变状态变换的形式，同时向观测添加失真或干扰噪声

Lipschitz 常量

Lipschitz常数经常被用来衡量模型的稳健性
对于函数f，如果存在非负常数K满足下列公式，则函数f是Lipschitz连续的，同时最小的常数K就称为f的Lipschitz常数
从公式中可知，将||x-y||除过来之后，k就可以用来表示基于x和y的f函数的一个变化程度，变化越剧烈则K值越大，相反则K值越小

policy层面的Lipschitz 常量

d(·，·)为状态空间的距离度量，特定数据增强方法Φ，策略π的Lipschitz常数如下表示
认为较小的Lipschitz常数表示对输入方差的稳定性较高
这一命题表明，如果在一个特定的增广下获得较小的Lipschitz常数，则在实现数据增强时，我们将以较小的方差获得较紧的Q值估计的界。
sup为上确界，即因为K需要取到大于等于右侧公式的最小值
证明在附录中

Task-aware Lipschitz Data Augmentation (TLDA)

目标：保持训练的稳定性和提高泛化能力
K-matrix 构造
- 使用kernel对原始图像o进行扰动得到扰动图像A(o)
- 选择以A(o)的位置(i,j)为中心的pixels，使用mask进行对扰动像素进行选择
- 对于给定的观测值o，计算像素(i，j)的Lipschitz常数（其中分子可以解释为两个动作分布之间的距离，分母是原始观测和扰动观测之间的距离。）
- 利用每个像素的Lipschitz常数，构造能够反映任务相关性信息的矩阵，并将其应用于整个观测。
- 目标是捕捉与任务相关的位置，这些位置具有大的Lipschitz常数
  - （即在相同扰动水平下，这些位置往往会导致策略/值输出的高度变化）
基于K矩阵实现任务感知的Lipschitz增强
- 判断像素修改的标准，使用K矩阵的平均值作为阈值（对K矩阵进行二值化，N为像素数）
- 仅在与任务相关性较低的区域中修改了输出
- 方法倾向于保留K值较大的像素，只修改K值较小的像素。这增加了一个隐式约束，以保持政策和价值网络的稳定输出。（因此，它与定义1中的最优不变状态转换相呼应。）
将TLDA加入RL
- 基础强化学习算法使用了SAC，在其critic loss 中增加了正则化项用于处理增强数据
- 实例化算法如下所示

实验

前情设置

对距离度量方式的选择d(·|·)
- 比较了不同度量下的K矩阵
  - L2范式效果更尖锐更集中
  - 全方差距离具有模糊效果
  - KL散度则结果更暗
- 最终在所有的实验中选择了L2范式的方式进行距离度量的计算。
其他细节
- 实验中并没有对每个pixel计算其Lipschitz常数，而是选择了每5个像素计算一次
  - 减少了复杂度和训练时间
  - 同样有较好的效果
- 使用高斯模糊作为kernel进行图像o扰动得到图像A（o）
- 以（i,j）为中心的二维高斯作为掩码M（i,j）

Evaluation on DeepMind Control Suite

使用SAC作为基本算法
图像输入采用卷积神经网络。
要考虑在现有技术方法中应用的两种增强方法：随机卷积(通过随机卷积层传递输入)和随机覆盖(将观测o与额外图像i线性组合）。
Baselines：
- (1)DrQ(Kostrikov等人，2020)：具有弱扩充(随机移动)的SAC；
- (2)PAD(Hansen等人，2020)：增加一个适应不可见环境的辅助任务；
- (3)Soda(Hansen和Wang，2021)：通过采用类似BYOL的架构来最大化潜在表征之间的相互信息(Grill等人，2020)；
- (4)Svea(Hansen等人，2021)：修改Q-Target的形式
强数据增强下的样本效率
- 与SVEA和random patches
- svea是使用强数据增强，对比证明TLDA的有效性
- random patches(将图片中的随机块替换为另一张图片的中随机块)则是没有任何像素到任务的相关信息，证明对原始图像的盲目修改或保留并不能提高agent的性能。
泛化性能
- 从DMControl-GB的两个设置上评估了该代理的泛化能力：
  - (I)背景和代理的随机颜色；(II)动态视频背景。
  - 用TLDA训练的代理能够在不同的不可见环境中获得良好的健壮策略。
  - 已有方法对不同数据增强方法的差异很敏感，即在不同的方法下测试性能差别较大。
TLDA定性结果
- 对K矩阵进行可视化
- 从测试环境上的K矩阵来看，经过TLDA训练的代理将在机器人身上提供更大的Lipschitz常数，而Svea代理倾向于专注于照明视觉背景。
- 我们的方法能够学习影响性能的主要因素，而忽略阻碍泛化的无关领域。

对动作分布和Q值估计的影响
- 这两个结果表明，TLDA有可能在训练中实现更高的样本效率，并学习更稳健的策略，以便在看不见的环境中表现良好。

Evaluation on Autonomous Driving in CARLA

目的：进一步评估TLDA的性能，将该方法应用于具有更真实观察的任务
任务设置：
- 使用一个摄像头作为驾驶任务的输入观察
- 智能体的目标是在1000个时间步长内尽可能地沿着弯曲的道路行驶，而不与移动的车辆、行人和障碍物相撞
- 在与以往工作相同的设置下对agent进行训练，本方法达到了最好的训练样本效率
- 选择达到100米距离的成功率作为驾驶评估指标
- 泛化测试：CARLA用内置的参数提供了不同的天气条件。在4种不同的天气条件下对我们的方法进行了评估，分别是不同的光照条件、真实的降雨和湿滑。
- TLDA在样本效率和泛化能力上都优于所有基本算法，具有更稳定的驱动策略。
- 我们评估了不同天气下的碰撞强度，并在表5中报告了训练环境和看不见环境中的平均值。

Evaluation on DMC Manipulation Tasks

机器人操作是视觉RL的另一组具有挑战性和有意义的任务
DMC提供了一套可配置的操作任务，带有机械臂和卡扣式砖块。
考虑了两个实验任务：reach push
- reach：agent需要通过操纵手臂到达所示的红砖
- push：agent目标是将红砖推到白色标记点的位置
所有agent都在默认背景上进行培训，并在不同颜色的手臂和平台上进行评估。
样本效率和泛化性能是该算法的优势所在
所有方法都在操作任务上进行了500k步的训练，训练结果如图11所示。尽管DrQ显示了最佳的训练性能，但表3中针对泛化的比较结果表明，我们的方法的性能显著优于基线，这意味着DrQ可能倾向于过度适应环境，而Svea在这些更具挑战性的基准下可能很容易出现分歧。