强化学习中自适应的辅助任务加权(Adaptive Auxiliary Task Weighting)

RL Researcher

Adaptive Auxiliary Task Weighting for Reinforcement Learning. NeurIPS 2019

在强化学习中，使用辅助任务来帮助主任务的学习是重要的一类方法。主任务在学习中，由于奖励的稀疏性和任务的困难，可以使用辅助任务帮助学习特征表示。常见的辅助任务如下表所示。在学习RL主任务的同时，可以利用收集到的样本，用监督学习的方法来学习环境模型（forward dynamics），逆环境模型(inverse dynamics)，利用图像状态的特殊性来学习Egomotion, Autoencoder的特征，利用光流来学习特征。由于学习这些辅助任务的时并不需要依赖环境提供的奖励信号，而只需要利用智能体在经验池中收集到的样本，因此这些任务被称为 self-supervised 学习方法。辅助任务将和主任务共享某些特征表示，在训练辅助任务的同时可以训练良好的特征表示，该特征表示将有利于加快主任务的学习。

以往的辅助任务学习方法，直接将主任务和多个辅助任务联合进行训练，辅助任务的权重需要进行调参。本文在此基础上提出了一种辅助任务的自适应加权方法。如下图，多个辅助任务将和主任务共享卷积特征表示，随后主任务和辅助任务被映射成低维特征表示进行连接，最后输出每个辅助任务的损失。如果使用固定的权重来加权各个辅助任务的损失，会存在以下问题：

有些辅助任务在训练的前期比较重要，在后期慢慢变得不重要，权重应该是变化的
有些辅助任务甚至会拖慢训练的进程，对主任务的学习没有帮助，应该给予较小的权重

本文的创新点在于各个损失到最终损失的过程进行了自适应的加权。每个辅助任务的权重 w_i 是可以自适应的调整的。

自适应的辅助任务加权（1-step update）

本文对辅助任务学习目标的假设是辅助任务能够使得主任务的损失更快的下降。如果一个辅助任务和主任务的损失的梯度方向是一致或接近的（比如夹角小于90度），那么该任务被认为是重要的，有利于主任务学习的。如果二者梯度方向相反，则认为该辅助任务不利于主任务的学习，应该给予较小的权重。本文对加权过程提供了比较理论化的方法。

考虑有 K 个辅助任务和一个主任务，辅助任务的损失记为 L_{aux,i}(\theta_t) , i\in\{0,1,...,k\} . 主任务的损失记为 L_{main} ，则总体任务的损失为：

L(\theta_t)=L_{main}(\theta_t)+\sum_{i=1}^{K}w_i L_{aux,i}(\theta_t)

其中 \theta_t 是 t 时间步的网络权重。在该损失的定义下，权重 \theta_t 的迭代公式为：

\theta_{t+1}=\theta_t+\alpha\nabla_{\theta_t}L(\theta_t)

我们希望在辅助任务的作用下，主任务迭代的更快。我们希望最大化 main 网络的损失变化：

V_t(w)=L_{main}(\theta_{t+1})-L_{main}(\theta_t)

由于 \theta_{t+1} 是 \theta_t 随着总体的损失函数 L(\theta_t) 来迭代的，因此

L_{main}(\theta_{t+1})=L_{main}(\theta_t+\alpha\nabla_{\theta_t}L(\theta_t))

进而，根据一阶泰勒公式 f(x+a)=f(a)+\nabla f(a) \cdot x ，则

L_{main}(\theta_t+\alpha\nabla_{\theta_t}L(\theta_t))\approx L_{main}(\theta_t)+\alpha \nabla_{\theta_t}L_{main}(\theta_t)^{\top} \nabla_{\theta_t} L(\theta_t)

带入 V_t(w) 的表达式，得

V_t(w)=L_{main}(\theta_t)+\alpha \nabla_{\theta_t}L_{main}(\theta_t)^{\top} \nabla_{\theta_t} L(\theta_t)-L_{main}(\theta_t)

消元得

V_t(w)=\alpha \nabla_{\theta_t}L_{main}(\theta_t)^{\top} \nabla_{\theta_t} L(\theta_t)

由于辅助任务的权重 w 的目的是最大化 V_t(w) ，因此：

w:=w+\frac{\partial V_t(w_i)}{\partial w_i}

其中，偏导数项中，根据 V_t(w) 的定义，只有在包含辅助任务的损失 L(\theta_t) 中才含有 w 项，因此可以求解如下：

\frac{\partial V_t(w_i)}{\partial w_i}=\alpha\nabla_{\theta_t}L_{main}(\theta_t)^{\top}\nabla_{\theta_t}L_{aux,i}(\theta_t)

由此我们得到了辅助任务权重 w_i,i\in\{0,1,...,k\} 的训练梯度。 w 对应的更新梯度正比于

“ L_{main} 对参数的梯度和第 i 个辅助任务 L_{aux,i} 对参数的梯度的数量积”

我们知道， a\cdot b 可以写成 \|a\|\cdot \|b\|\cdot \cos<a,b> . 因此，辅助任务 i 对应的权重 w_i 的梯度与主任务的梯度，辅助任务 i 的梯度，以及二者梯度的夹角有关。

如果辅助任务对主任务的学习是有利的，二者梯度的夹角应该很小，此时cos值接近于1
如果辅助任务与主任务不相关，二者的夹角接近于正交，此时cos值接近于0
如果辅助任务对主任务的学习是不利的，二者的梯度方向相反，此时 cos值小于0

因此，当辅助任务和主任务对 \theta_t 产生的梯度方向基本一致时，会有较大的权重，二者的梯度会产生叠加的效果。其余情况下，会给予辅助任务较小的权重来避免影响主任务的学习。

自适应的辅助任务加权（N-step update）

以上 w_i 最大化的是 L_{main}(\theta_{t+1})-L_{main}(\theta_{t}) ，因此是一种 1-step update 目标。本文进而提出了一个扩展，使用 L_{main}(\theta_{t+N})-L_{main}(\theta_{t}) 作为学习的目标。

目标函数可以写成：

V_t(w)=L_{main}(\theta_{t+N})-L_{main}(\theta_t)

将 \theta_{t+N} 展开为总体损失 L 的单步的更新结果：

V_t(w)=L_{main}(\theta_{t+N-1}+\alpha\nabla_{\theta_{t+N-1}}L(\theta_{t+N-1}))-L_{main}(\theta_t)

利用泰勒公式展开第一项 L_{main}(\theta_{t+N-1}+\alpha\nabla_{\theta_{t+N-1}}L(\theta_{t+N-1}))=L_{main}(\theta_{t+N-1})+\alpha\nabla_{\theta_{t+N-1}} L_{main}(\theta_{t+N-1})^{\top}\nabla_{\theta_{t+N-1}}L(\theta_{t+N-1})

类似的，展开上式中的第一项 L_{main}(\theta_{t+N-1}) ，得到

\begin{split} L_{main}(\theta_{t+N-1})&=L_{main}(\theta_{t+N-2}+\alpha\nabla_{\theta_{t+N-2}}L(\theta_{t+N-2}))\\ &=L_{main}(\theta_{t+N-2})+\alpha\nabla_{\theta}L_{main}(\theta_{t+N-2})^{\top}\nabla_{\theta_{t+N-2}}L(\theta_{t+N-2})) \end{split}

再次展开上式中的第一项 L_{main}(\theta_{t+N-2}) ,...以此类推直到展开可得

\begin{split} V_t(w)&=L_{main}(\theta_{t})+\alpha\sum_{j=0}^{N-1}\nabla_{\theta_{t+j}}L_{main}(\theta_{t+j})^{\top}\nabla_{\theta_{t+j}}L(\theta_{t+j})-L_{main}(\theta_t)\\ &=\alpha\sum_{j=0}^{N-1}\nabla_{\theta_{t+j}}L_{main}(\theta_{t+j})^{\top}\nabla_{\theta_{t+j}}L(\theta_{t+j}) \end{split}

进而

\frac{\partial V_t(w_i)}{\partial w_i}=\alpha\sum_{j=0}^{N-1}\nabla_{\theta_{t+j}}L_{main}(\theta_{t+j})^{\top}\nabla_{\theta_{t+j}}L_{aux,i}(\theta_{t+j})

使用 N-step 的目标来更新相比于 1-step 有如下优势：

使用多步的目标使得算法能够考虑到较长时间段内的总体损失，相比于单步损失更合理
使用多步的目标可以平滑梯度，相比于单步损失产生的梯度更加稳定。此外，一般我们可以认为，辅助任务的权重应该在训练过程中有一个稳定的变化而不是跳变。

总结 N-step 的方法来更新的算法如下：

实验

1. 简单的监督学习实验

主任务： L_0(x,y)=x^2+y^2

辅助任务1： L_1(x,y)=(x-0.5)^2+(y-0.5)^2

辅助任务2： L_2(x,y)=-L_0(x,y)

这三个任务中，明显的，辅助任务2的损失和主任务是完全相反的。总体损失定义为：

L=L_0+w_1L_1+w_2L_2

如果将辅助任务权重设置为固定的值，如 w_1=w_2=1 ，那么 L_0 会和 L_2 抵消，那么最终收敛到的结果是 L_1 的最优解 (0.5,0.5) ，与主任务的最优解不一致。如下图左所示。

下图中显示了自适应的权重调节结果。从不同方位初始化的点会先收敛到 (0.5,0.5) ，然后再收敛到 (0,0) . 右侧显示了辅助任务的权重变化，任务2刚开始权重就明显降低，原因是与主任务之间有抵触。任务1则先升高后下降，符合预期。

2. 其他实验

本文使用 Atari 中 Breakout, Pong, Seaquest 和 Gym Robotics中 FetchReach, HandReach等使用HER算法在多目标环境下进行实验。在这些环境的主任务中，添加本文刚开始所示的Table 1所示的辅助任务，包括前向环境模型，逆环境模型，egomotion, auto-encoder, 光流等。

文章包括以下Baseline，其中与本文直接相关的是 Cosine Similarity. 但该baseline是一种硬约束，当二者梯度夹角小于 90 度的时候可以使用辅助任务，否则将其权重置0.