扩散模型解决多任务强化学习问题

大家好，扩散模型 (diffusion model) 在CV领域甚至NLP领域都已经有了令人印象深刻的表现。最近的一些工作开始将diffusion model用于强化学习 (RL) 中来解决序列决策问题，它们主要利用diffusion model来建模分布复杂的轨迹或提高策略的表达性。但是， 这些工作仍然局限于单一任务单一数据集，无法得到能同时解决多种任务的通用智能体。那么，diffusion model能否解决多任务强化学习问题呢？我们最近提出的一篇新工作——"Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning"，旨在解决这个问题并希望启发后续通用决策智能的研究：

论文目前已挂在arxiv上:

背景：

• 数据驱动的大模型在CV和NLP领域已经获得巨大成功，我们认为这背后源于模型的强表达性和数据集的多样性和广泛性。基于此，我们将最近出圈的生成式扩散模型 (diffusion model) 扩展到多任务强化学习领域 (multi-task reinforcement learning)，利用large-scale的离线多任务数据集训练得到通用智能体。
• 目前解决多任务强化学习的工作大多基于Transformer架构，它们通常对模型的规模，数据集的质量都有很高的要求，这对于实际训练来说是代价高昂的。基于TD-learning的强化学习方法则常常面临distribution-shift的挑战，在多任务数据集下这个问题尤甚，而我们将序列决策过程建模成条件式生成问题 (conditional generative process)，通过最大化likelihood来学习，有效避免了distribution shift的问题。

方法：

具体来说，我们发现diffusion model不仅能很好地输出action进行实时决策，同样能够建模完整的(s,a,r,s')的transition来生成数据进行数据增强提升强化学习策略的性能，具体框架如图所示：

对于决策规划任务，我们将模型称为 \text{MTDiff-p} ；对于数据生成（增强）任务，我们将模型称为 \text{MTDiff-s} 。我们的扩散模型基于DDPM的训练方式，利用 L_2 损失训练 \epsilon_\theta ，它能够预测逆向过程每一步的噪声 \epsilon 。对于 \text{MTDiff-p} ，它的损失函数可以表示为：

\begin{equation} \mathcal{L}^{\color{red}{p}}(\theta):=\mathbb{E}_{k\sim \mathcal{U}(1, K), \bm{x}_0(\tau)\sim q,\epsilon\sim\mathcal{N(\boldsymbol{\mathrm{0}},\boldsymbol{{I}})},\beta\sim \rm{Bern}(p)}\big[\big\|\epsilon-\epsilon_\theta\big(\bm{x}_k^{\color{red}{p}}(\tau),\bm{y}'(\tau),(1-\beta)R(\tau)+\beta\varnothing, k\big)\big\|^2\big]. \end{equation}\\

其中 \begin{equation} \bm{x}^{\color{red}{p}}_k(\tau):=(a_t, a_{t+1},...,a_{t+H-1})_k, \end{equation} \bm{y}'(\tau):=(Z,s_{t-L+1},...,s_{t}) , R(\tau) 是轨迹的标准化累积回报， Z 是 Demonstration Prompt ，可以表示为：

\begin{equation} Z:=\begin{bmatrix} s^*_{i}&s^*_{i+1}& \cdots& s^*_{i+J- 1}\\ a^*_{i}&a^*_{i+1}&\cdots&a^*_{i+J-1} \end{bmatrix}, \end{equation}\\

*表示专家轨迹， Z 即是从每个任务下的一条或几条专家轨迹采样得到的 J 时间步长的一段包含状态和动作的轨迹。

\text{MTDiff-p} 对 R(\tau) 做 classifi-free guidance ，并在测试推理过程中生成未来 H 步长的动作序列，在实验中我们选取第一个动作与环境交互。

对于 \text{MTDiff-s} ，它的损失函数可以表示为：

\begin{equation} \mathcal{L}^{\color{red}{s}}(\theta):=\mathbb{E}_{k\sim \mathcal{U}(1, K), \bm{x}_0(\tau)\sim q,\epsilon\sim\mathcal{N(\boldsymbol{\mathrm{0}},\boldsymbol{{I}})}}\big[\big\|\epsilon-\epsilon_\theta(\bm{x}_k^{\color{red}{s}}(\tau),\bm{y}^{\color{red}{s}}(\tau), k)\big\|^2\big]. \end{equation}\\

其中 \begin{equation} \bm{x}^{\color{red}{s}}_k(\tau):=\begin{bmatrix} s_t &s_{t+1}& \cdots & s_{t+H-1} \\ a_t &a_{t+1}&\cdots&a_{t+H-1} \\ r_t &r_{t+1}& \cdots & r_{t+H-1} \end{bmatrix}, \end{equation} \bm{y}^{\color{red}{s}}(\tau):=[Z]。 值得注意的是这里我们需要扩散模型建模整个 (s,a,r) 的完整轨迹，并且由于目标是完成数据生成任务，需要尽可能多样化的数据，不需要模型condition在 R(\tau) 上，自然也不需要classifier-free guidance。

模型结构：

为了更好地建模多任务数据，并且统一多样化的输入数据，我们用transformer架构替换了传统的U-Net网络，网络结构图如下：

实验：

我们首先在Meta-World MT50上开展实验并与baselines进行比较，我们在两种数据集上进行实验，分别是包含大量专家数据，从SAC-single-agent中的replay buffer中收集到的Near-optimal data（100M）；以及从Near-optimal data中降采样得到基本不包含专家数据的Sub-optimal data（50M）。实验结果如下：

可以看到我们的方法在同样大小的数据上不仅超越了大多数离线强化学习方法，并且也超过了两种多任务在线强化学习方法，其中PaCo是目前的SOTA方法。我们的方法通过 classifier-free guidance 得到最优行为，对数据集质量的依赖性不强，在Sub-optimal data上的表现相比其他sequence modeling的方法有很大提升。在我们的实验中也观察到Behavior Cloning（BC）在Near-optimal data上和我们的方法相比具有competitive performance，但BC严重依赖数据集质量，在Sub-optimal data上性能急剧下降。

在数据增强方面，我们的方法也同样有效，明显提升了数据集的质量以增强策略性能，实验结果如下：

我们选取45个任务的Near-optimal data训练 \text{MTDiff-s} ，从表中我们可以观察到在 \text{MTDiff-s} 见过的任务上，我们的方法均取得了最好的性能。甚至给定一段demonstration prompt， \text{MTDiff-s} 能泛化到没见过的任务上并取得较好的表现。 我们选取四个任务对原数据和 \text{MTDiff-s} 生成的数据做T-SNE可视化分析，发现我们生成的数据的分布基本匹配原数据分布，并且在不偏离的基础上扩展了分布，使数据覆盖更加全面。

\text{MTDiff-p} 同样也具有泛化性，更多实验分析可以参见论文。

总结：

我们提出了一种基于扩散模型（diffusion model）的一种新的、通用性强的多任务强化学习解决方案，它不仅可以通过单个模型高效完成多任务决策，而且可以对原数据集进行增强，从而提升各种离线算法的性能。我们未来将把 \text{MTDiff} 迁移到更加多样、更加通用的场景，旨在深入挖掘其出色的生成能力和数据建模能力，解决更加困难的任务。同时，我们会将 \text{MTDiff} 迁移到真实控制场景，并尝试优化其推理速度以适应某些需要高频控制的任务。