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

2023-09-26 05:09

©作者 | 何浩然

单位 | 上海AI Lab&上海交大APEX实验室

研究方向 | Embodied AI、强化学习

扩散模型（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”，旨在解决这个问题并希望启发后续通用决策智能的研究：

论文题目：

Diffusion Model is an Effective Planner and Data Synthesizer for Multi-Task Reinforcement Learning

论文链接：

https://arxiv.org/abs/2305.18459

背景

数据驱动的大模型在 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 来生成数据进行数据增强提升强化学习策略的性能，具体框架如图所示：

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

其中

是轨迹的标准化累积回报，是 Demonstration Prompt，可以表示为：

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

对做 classifi-free guidance，并在测试推理过程中生成未来步长的动作序列，在实验中我们选取第一个动作与环境交互。

对于，它的损失函数可以表示为：

其中

值得注意的是这里我们需要扩散模型建模整个（s,a,r）的完整轨迹，并且由于目标是完成数据生成任务，需要尽可能多样化的数据，不需要模型 condition 在上，自然也不需要 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 训练，从表中我们可以观察到在见过的任务上，我们的方法均取得了最好的性能。甚至给定一段 demonstration prompt，能泛化到没见过的任务上并取得较好的表现。我们选取四个任务对原数据和生成的数据做 T-SNE 可视化分析，发现我们生成的数据的分布基本匹配原数据分布，并且在不偏离的基础上扩展了分布，使数据覆盖更加全面。

同样也具有泛化性，更多实验分析可以参见论文。

总结

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

更多阅读