NeurIPS 2022

摘要

The purpose of multi-task reinforcement learning (MTRL) is to train a single policy that can be applied to a set of different tasks. Sharing parameters allows us to take advantage of the similarities among tasks. However, the gaps between contents and difficulties of different tasks bring us challenges on both which tasks should share the parameters and what parameters should be shared, as well as the optimization challenges due to parameter sharing. In this work, we introduce a parameter-compositional approach (PaCo) as an attempt to address these challenges. In this framework, a policy subspace represented by a set of parameters is learned. Policies for all the single tasks lie in this subspace and can be composed by interpolating with the learned set. It allows not only flexible parameter sharing but also a natural way to improve training. We demonstrate the state-of-the-art performance on Meta-World benchmarks, verifying the effectiveness of the proposed approach.

多任务强化学习（MTRL）的目的是训练一个可以应用于一组不同任务的单一策略。共享参数使我们能够利用任务之间的相似性。然而，不同任务的内容和难度之间的差距给我们带来了挑战，即哪些任务应该共享参数，哪些参数应该共享，以及参数共享带来的优化挑战。在这项工作中，我们引入了一个参数组合方法（PaCo），试图解决这些挑战。在这个框架中，由一组参数代表的策略子空间被学习。所有单一任务的策略都在这个子空间中，可以通过对所学参数集的插值来组成。它不仅允许灵活的参数共享，而且也是改善训练的一种自然方式。我们在Meta-World基准上展示了最先进的性能，验证了所提方法的有效性。

研究动机

• 多任务强化学习旨在训练一种有效的策略，可以应用于同一机器人解决不同的任务
• MTRL的挑战主要在于任务间应该共享什么以及怎样共享

主要贡献

• 提出一个通用参数组成（PACO）MTRL训练框架，该框架可以学习用于制定不同任务策略的代表性参数集
• 提出一个方法利用PACO的分解结构来稳定MTRL的训练
• 在Meta-World环境中验证了方法的性能

方法描述

• Categories of MTRL：

  • 建模共享结构（本文属于这一类）
  • 解决梯度冲突的算法
  • 选择或学习更好的任务表征
• Multi-Goal MTRL：任务由不同的"goal"定义，state和transition相同，reward不同；

• Multi-Skill MTRL：除了相同环境的目标变化，还包括不同环境共享相似动态过程（transition function），要求智能体掌握不同的skill。不同skill的任务state不同。

  • 个人理解：例如操作同一个机械臂拿不同的物体，由于物体不同state和transition存在差异，但也有类似的时候，例如用于表示机械臂相关的state是相同的，机械别各种操作的transition是类似的

• MTRL Challenges：

  • Parameter-Sharing：参数共享可以充分利用任务间的相似性。但是任务难度和内容的不同带来的问题是：哪些任务应该参数共享？哪部分参数应该共享？如果错误的共享可能导致训练效率不如各任务独立训练
  • Multi-Task Training Stability：不同技能同时训练存在的冲突，例如部分任务的损失爆炸可能影响其他的训练任务

• Formulation

  • parameter-compositional model参数组合模型：$\theta_{\tau} = \Phi \mathbf{w}_\tau$

    • $\Theta = \Phi \mathbf{W} \quad \Leftrightarrow \quad [\theta_1,\cdots,\theta_T] = \Phi[\mathbf{w}_1,\cdots,\mathbf{w}_T]$
    • $\theta_\tau \in \mathbb{R}^n$：任务 $\tau$ 的所有可训练参数
    • $\Phi=[\phi_1,\phi_2,\cdots,\phi_K] \in \mathbb{R}^{n \times K}$：$K$ 个参数向量 $\{\phi_i\}_{i=1}^K$ 组成的矩阵
    • $\mathbf{w}_\tau \in \mathbb{R}^K$：组合向量，是任务 $\tau$ 的可训练嵌入向量
    • 对于单任务而言，该方法等同于标准的参数化模型方法
    • 对于多任务而言，模型被分为两部分：任务无关task-agnostic $\Phi$ 和 任务相关task-aware $\mathbf{w}_\tau$

  • multi-task loss训练损失：$J_\Theta \triangleq \sum_\tau J_\tau(\theta)$

    • $J_\tau(\theta)$：任务 $\tau$ 的actor和critic损失（SAC实现）
    • $\frac{\partial J_\Theta}{\partial \Phi} = \sum_\tau \frac{\partial J_\tau}{\partial \Phi}$：所有 $T$ 个任务对学习共享参数集合 $\Phi$ 的贡献
    • $\frac{\partial J_\Theta}{\partial \mathbf{W}} = \sum_\tau \frac{\partial J_\tau}{\partial \mathbf{w}_\tau}$：单个任务损失 $J_\tau$ 只影响其自身任务的特定组合向量 $\mathbf{w}_\tau$

  • addressing the MTRL challenges：

    • 通过学习每个任务的特定组合向量 $\mathbf{w}_\tau$ 进行灵活的参数共享
    • 由于task-specific和task-agnostic参数的完全独立，提供了一种自然解决方案改善MTRL训练的稳定性

• Stable Multi-Task Reinforcement Learning

  • 对共享参数更新的共享可能来源于当下主导的任务，将影响其他任务的学习
  • loss maskout：当某个任务损失 $J_\eta$ 超过阈值 $\epsilon$，将不对其进行优化

  • $\mathbf{w}$-reset：$\mathbf{w}_\eta = \sum_{j \in \mathcal{V}} \beta_j \mathbf{w}_j, \quad \beta = [\beta_1,\beta_2,\cdots] \sim \triangle^{|\mathcal{V}|-1}$

    • 重新初始化任务 $\eta$ 的task-specific参数 $\mathbf{w}_\eta$
    • 其中 $\mathcal{V} \triangleq \{j|J_j \le \epsilon\}$
    • ==没看到参数 $\beta \sim \triangle^{|\mathcal{V}-1|}$ 的表示有什么意义==

• Unified Perspective on Some Existing Methods

  • Single-Task Model

    • $\Phi=[\phi_1,\phi_2,\cdots]$
    • $\mathbf{w}_\tau$ 为 one-hot向量
    • 即每个任务都有自己的专有参数 $\theta_\tau = \phi_\tau$

  • Multi-Task Model

    • $\Phi=[\phi] \in \mathbb{R}^{n \times 1}$
    • $\mathbf{w}_1 = \mathbf{w}_2 = \cdots = 1$
    • 即所有任务共享所有参数 $\theta_\tau = \phi$

  • Multi-Head Multi-Task Model

    • $\Phi=\left[\begin{array}{cccccc}\phi^{\prime} & \phi^{\prime} & \cdots & \phi^{\prime} & \cdots & \phi^{\prime} \\ \psi_1 & \psi_2 & \cdots & \psi_\tau & \cdots & \psi_K\end{array}\right] \in \mathbb{R}^{n \times K}$

      • $\psi_\tau$：输出层子参数
    • $\mathbf{w}_\tau$ 为 one-hot向量
    • 即每个任务既有共享上层参数，也有独立下层参数 $\theta_\tau = \left[\begin{array}{c} \phi' \\ \psi_\tau \ \end{array}\right]$

  • Soft-Modularization

    • $f_{\boldsymbol{\theta}_\tau}=\left[\begin{array}{cccc}{\left[f_{\boldsymbol{\phi}_1^1}\right.} & f_{\boldsymbol{\phi}_2^1} & \cdots & \left.f_{\boldsymbol{\phi}_K^1}\right] z^1(s, \tau) \\ &&\vdots \\{\left[f_{\boldsymbol{\phi}_1^m}\right.} & f_{\boldsymbol{\phi}_2^m} & \cdots & \left.f_{\boldsymbol{\phi}_K^m}\right] z^m(s, \tau)\end{array}\right]$
    • 其中可以理解为 $\Phi$ 为每层 $f_{\phi_j^i}$ 的组合；$\mathbf{w}_\tau$ 为 $z^i(s,\tau)$
    • 但区别是 $\mathbf{w}$ 只与 $\tau$ 相关；$z$ 与 $\tau$ 和 $s$ 都相关

理论分析

无

实验验证

• 实验环境：Meta-World

• Baselines：

  • Multi-task SAC：one-hot task作为输入的SAC
  • Multi-Head SAC：Multi-Head的SAC
  • SAC+FiLM：FiLM实现task-conditional policy
  • PCGrad：http://darkdawn.top/index.php/archives/18/
  • Soft-Module：http://darkdawn.top/index.php/archives/22/
  • CARE：http://darkdawn.top/index.php/archives/19/

• 实验设置

  • 10个并行环境
  • 20 million环境步数（2 million per task）
  • 10个随机种子
• 实验结果

• 消融实验

  • 稳定训练：Maskout、w-reset

• Parameter Set Size：参数集合数量K

• Compositional Structure Variations：不同网络部分使用参数组合

  • Actor-only：只有Actor网络参数组合；Output-only：只有输出层参数组合（类似Multi-Head）

• 可视化

  • 策略子空间（normalize到一个圆上）

    • 图a展示不同任务学到的参数比确实不一样
    • 图b展示稍微更改参数比，智能体的策略有一定可解释性

• 将10个任务的参数比可视化（PCA到二维）

  • 有的任务很独立（Reach）；
  • 有的任务很相似（window-open v.s door-open, window-open v.s. window-close）；
  • 有趣的发现：w很接近peg-insert v.s. window-open/window-close/door-open/drawer-open；可能的原因：他们虽然语义上差距很大，但都使用类似的轨迹完成任务（解释CARE直接引入自然语言元数据可能存在不合理）

其他思考

• 对比Soft-Module：Soft-Module是以网络层作为最小单位进行参数选择，本文是以单个参数作为最小单位进行参数选择。而且选择方式不太一样，Soft-Module是根据状态和任务id生成选择routing，而PaCo每个任务独立存在一个参数w，用于选择参数集合。
• 其实本文的这种参数组合方式和CARE中的Mixture Encoder比较类似，区别在于没有利用元数据进行mixture的生成，而是直接独立学习。再回溯该方法其实和MMOE的Gate模块就很类似了，不过这里是对参数进行加权和，而MMOE是对输出进行加权和
• 其实本文还可以进行一个验证，就是测试一下泛化性，或者是预训练模型的few-shot，具体而言，我们认为学习的到 $\Phi$ 是shared knowledge，那么针对新的任务则只需要学习对应的 $w$ 权重即可，但文中只在附录中简单说明，并没有完整的实验演示。
• 还有一个未解决的问题是，没办法自动的进行参数K的选择，该超参数的选择也需要消耗大量的计算资源去实验，当任务数量进一步扩大时，该参数将更难确定。
• 对于Maskout和w-reset对训练的稳定性有帮助，个人持保留意见。当网络学不会的时候（loss爆炸），则停止训练以免影响其他任务或者重新初始化再试一次，这种做法其实也是比较trick的方法。因为没有研究CARE的源码，所以作者提到的”stop training“的trick可以关注一下，确实在其正文中完全没有提及。

原文链接：https://arxiv.org/abs/2210.11653

OpenReview：https://openreview.net/forum?id=LYXTPNWJLr

Code：https://github.com/TToTMooN/paco-mtrl （暂时还没有）