ICLR 2023 Spotlight | 节省95%训练开销，清华黄隆波团队提出强化学习专用稀疏训练框架RLx2

大模型时代，模型压缩和加速显得尤为重要。传统监督学习可通过稀疏神经网络实现模型压缩和加速，那么同样需要大量计算开销的强化学习任务可以基于稀疏网络进行训练吗？本文提出了一种强化学习专用稀疏训练框架，可以节省至多 95% 的训练开销。

深度强化学习模型的训练通常需要很高的计算成本，因此对深度强化学习模型进行稀疏化处理具有加快训练速度和拓展模型部署的巨大潜力。然而现有的生成小型模型的方法主要基于知识蒸馏，即通过迭代训练稠密网络，训练过程仍需要大量的计算资源。另外，由于强化学习自举训练的复杂性，训练过程中全程进行稀疏训练在深度强化学习领域尚未得到充分的研究。

清华大学黄隆波团队提出了一种强化学习专用的动态稀疏训练框架，“Rigged Reinforcement Learning Lottery”（RLx2），可适用于多种离策略强化学习算法。它采用基于梯度的拓扑演化原则，能够完全基于稀疏网络训练稀疏深度强化学习模型。RLx2 引入了一种延迟多步差分目标机制，配合动态容量的回放缓冲区，实现了在稀疏模型中的稳健值学习和高效拓扑探索。在多个 MuJoCo 基准任务中，RLx2 达到了最先进的稀疏训练性能，显示出 7.5 倍至 20 倍的模型压缩，而仅有不到 3% 的性能降低，并且在训练和推理中分别减少了高达 20 倍和 50 倍的浮点运算数。

一、背景

在游戏、机器人技术等领域，深度强化学习（DRL）已经取得了重要的应用。然而，深度强化学习模型的训练需要巨大的计算资源。例如，DeepMind 开发的 AlphaGo-Zero 在围棋游戏中击败了已有的围棋 AI 和人类专家，但需要在四个 TPU 上进行 40 多天的训练。OpenAI-Five 是 OpenAI 开发的 Dota2 AI，同样击败了人类半职业 Dota 高手，但是需要高达 256 个 GPU 进行 180 天的训练。实际上，即使是简单的 Rainbow DQN [Hessel et al. 2018] 算法，也需要在单个 GPU 上训练约一周时间才能达到较好的性能。

高昂的资源消耗限制了深度强化学习在资源受限设备上的训练和部署。为了解决这一问题，作者引入了稀疏神经网络。稀疏神经网络最初在深度监督学习中提出，展示出了对深度强化学习模型压缩和训练加速的巨大潜力。在深度监督学习中，SET [Mocanu et al. 2018] 和 RigL [Evci et al. 2020] 等常用的基于网络结构演化的动态稀疏训练（Dynamic sparse training - DST）框架可以从头开始训练一个 90% 稀疏的神经网络，而不会出现性能下降。

在深度强化学习领域，已有的工作已经成功生成了极度稀疏的深度强化学习网络。然而，他们的方法仍然需要迭代地训练稠密网络，往往需要预训练的稠密模型作为基础，导致深度强化学习的训练成本仍然过高，无法直接应用于资源有限设备。

从头开始训练一个稀疏深度强化学习模型，如果能够完美实现，将极大地减少计算开销，并实现在资源受限设备上的高效部署，具备优秀的模型适应性。然而，在深度强化学习中从头开始训练一个超稀疏网络（例如 90% 的稀疏度）具有挑战性，原因在于自举训练（Bootstrap training）的非稳定性。在深度强化学习中，学习目标不是固定的，而是以自举方式给出，训练数据的分布也可能是非稳定的。此外，使用稀疏网络结构意味着在一个较小的假设空间中搜索，这进一步降低了学习目标的置信度。因此，不当的稀疏化可能对学习路径造成不可逆的伤害，导致性能较差。最近的研究 [Sokar et al. 2021] 表明，在深度强化学习中直接采用动态稀疏训练框架仍然无法在不同环境中实现模型的良好压缩。因此，这一重要的开放问题仍然悬而未决：

能否通过全程使用超稀疏网络从头训练出高效的深度强化学习智能体？

二、方法

清华大学黄隆波团队对这一问题给出了肯定的答案，并提出了一种强化学习专用的动态稀疏训练框架，“Rigged Reinforcement Learning Lottery”（RLx2），用于离策略强化学习（Off-policy RL）。这是第一个在深度强化学习领域以 90% 以上稀疏度进行全程稀疏训练，并且仅有微小性能损失的算法框架。RLx2 受到了在监督学习中基于梯度的拓扑演化的动态稀疏训练方法 RigL [Evci et al. 2020] 的启发。然而，直接应用 RigL 无法实现高稀疏度，因为稀疏的深度强化学习模型由于假设空间有限而导致价值估计不可靠，进而干扰了网络结构的拓扑演化。

因此，RLx2 引入了延迟多步差分目标（Delayed multi-step TD target）机制和动态容量回放缓冲区（Dynamic capacity buffer），以实现稳健的价值学习（Value learning）。这两个新组件解决了稀疏拓扑下的价值估计问题，并与基于 RigL 的拓扑演化准则一起实现了出色的稀疏训练性能。为了阐明设计 RLx2 的动机，作者以一个简单的 MuJoCo 控制任务 InvertedPendulum-v2 为例，对四种使用不同价值学习和网络拓扑更新方案的稀疏训练方法进行了比较。

下图展示了 RLx2 算法的主要部分，包括基于梯度的拓扑演化、延迟多步差分目标和动态容量回放缓冲区。

2.1 基于梯度的拓扑演化

在 RLx2 中，作者采用了与 RigL [Evci et al. 2020] 相同的方法来进行拓扑结构的演化。作者计算了损失函数对网络权重的梯度值。然后，周期性地增加稀疏网络中具有较大梯度的连接，并移除权重绝对值最小的现有连接。通过周期性的结构演化，获得了一个结构合理的稀疏神经网络。

2.2 延迟多步差分目标

RLx2 框架还引入了多步差分目标：

这一概念在现有研究中 [Munos et al. 2016] 已被证实能够改善差分学习（TD learning）。作者还发现，引入多步目标可以通过折扣因子减少稀疏网络的拟合误差，如下式所示：

然而，训练伊始立即采用多步差分目标可能会导致更大的策略不一致误差。因此，作者额外采用了延迟方案来抑制策略不一致性并进一步提高值函数的学习效果。

2.3 动态容量回放缓冲区

离策略（Off-policy）算法使用回放缓冲区（Replay buffer）来存储收集到的数据，并使用从缓冲区中抽样的批次数据来训练网络。研究表明 [Fedus et al. 2020]，当使用更大的回放容量时，算法的性能通常会提高。然而，无限大小的回放缓冲区会因为多步目标的不一致性和训练数据的不匹配导致策略不一致性。动态容量回放缓冲区是一种通过调整缓冲区容量控制缓冲区中数据的不一致性，以实现稳健值函数学习的方法。作者引入了以下策略距离度量来评估缓冲区中数据的不一致性：

随着训练的进行，当回放缓存中的策略距离度量大于阈值时，则停止增加缓冲区容量，使得策略距离度量始终小于设定的阈值。

三、实验

作者在四个 MuJoCo 环境（HalfCheetah-v3、Hopper-v3、Walker2d-v3 和 Ant-v3），和两个常见的深度强化学习算法 TD3 和 SAC 中进行了实验。作者定义了一个终极压缩比率，即在该比率下，RLx2 的性能下降在原始稠密模型的 ±3% 之内。这也可以理解为具有与原始稠密模型完全相同性能的稀疏模型的最小大小。根据终极压缩比率，作者在下表中呈现了不同算法在不同环境采用相同参数量的神经网络的性能。

性能在所有四个环境中，RLx2 的性能在很大程度上优于所有基准算法（除了 Hopper 环境中与 RigL 和 SAC 的性能相近）。此外，小型稠密网络（Tiny）和随机静态稀疏网络（SS）的性能平均最差。SET 和 RigL 的性能较好，但在 Walker2d-v3 和 Ant-v3 环境中无法保持性能，这意味着在稀疏训练下稳健的价值学习是必要的。

模型压缩RLx2 实现了优秀的压缩比，并且仅有轻微的性能下降（不到 3%）。具体而言，使用 TD3 算法的 RLx2 实现了 7.5 倍至 25 倍的模型压缩，在 Hopper-v3 环境中获得了最佳的 25 倍压缩比。在每个环境中，演员网络（Actor network）可以压缩超过 96% 的参数，评论家网络（Critic network）可以压缩 85% 至 95% 的参数。SAC 算法的结果类似。另外，使用 SAC 算法的 RLx2 实现了 5 倍至 20 倍的模型压缩。

节省训练开销与基于知识蒸馏或行为克隆的方法 [Vischer et al. 2021] 不同，RLx2 在整个训练过程中使用了稀疏网络。因此，它具有加速训练并节省计算资源的额外潜力。四个环境的平均结果表明，表格中 RLx2-TD3 分别减少了 12 倍和 20 倍的训练和推理浮点运算数，RLx2-SAC 分别减少了 7 倍和 12 倍的训练和推理浮点运算数。

四、总结

作者提出了一种用于离策略强化学习的稀疏训练框架 RLx2，能够适用于各种离策略强化学习算法。这一框架利用基于梯度的结构演化方法实现了高效的拓扑探索，并通过延迟多步差分目标和动态容量回放缓冲区建立了稳健的值函数学习。RLx2 不需要像传统剪枝方法一样预训练稠密网络，却能够在训练过程中使用超稀疏网络来训练高效的深度强化学习智能体，并且几乎没有性能损失。

作者在使用 TD3 和 SAC 的 RLx2 上进行了实验，结果表明其稀疏训练性能非常出色：模型压缩比例为 7.5 倍至 20 倍，性能下降不到 3%，训练和推理的浮点运算数分别减少高达 20 倍和 50 倍。作者认为未来有趣的工作包括将 RLx2 框架扩展到更复杂的 RL 场景，这些场景对计算资源的需求更高，例如多智能体、离线强化学习等场景，也包括真实世界的复杂决策问题而非标准的 MuJoCo 环境。

参考文献

[1] Hessel, Matteo, et al. "Rainbow: Combining improvements in deep reinforcement learning." Proceedings of the AAAI conference on artificial intelligence. Vol. 32. No. 1. 2018.

[2] Mocanu, Decebal Constantin, et al. "Scalable training of artificial neural networks with adaptive sparse connectivity inspired by network science." Nature communications 9.1 (2018): 2383.

[3]Evci, Utku, et al. "Rigging the lottery: Making all tickets winners." International Conference on Machine Learning. PMLR, 2020.

[4]Sokar, Ghada, et al. "Dynamic sparse training for deep reinforcement learning." arXiv preprint arXiv:2106.04217 (2021).

[5]Munos, Rémi, et al. "Safe and efficient off-policy reinforcement learning." Advances in neural information processing systems 29 (2016).

[6]Fedus, William, et al. "Revisiting fundamentals of experience replay." International Conference on Machine Learning. PMLR, 2020.

[7]Vischer, Marc Aurel, Robert Tjarko Lange, and Henning Sprekeler. "On lottery tickets and minimal task representations in deep reinforcement learning." arXiv preprint arXiv:2105.01648 (2021).

作者：谭亦钦 文章来源：【机器之心】

Illustration by unDraw

-The End-

本周上新！

扫码观看！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线480+期talk视频，2400+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。

投稿内容

// 最新技术解读/系统性知识分享 //

// 前沿资讯解说/心得经历讲述 //

投稿须知

稿件需要为原创文章，并标明作者信息。

我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励

投稿方式

发送邮件到

chenhongyuan@thejiangmen.com

或添加工作人员微信（chemn493）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。

关于我“门”

将门是一家以专注于发掘、加速及投资技术驱动型创业公司的新型创投机构，旗下涵盖将门创新服务、将门技术社群以及TechBeat人工智能社区。公司致力干通过连接技术与商 业，发掘和培育具有全球影响力的科技创新企业，推动企业创新发展与产业升级。

将门成立于2015年底，创始团队由微软创投在中国的创始团队原班人马构建而成，曾为微软优选和深度孵化了126家创新的技术型创业公司。

如果您是技术领域的初创企业，不仅想获得投资，还希望获得一系列持续性、有价值的投后服务，欢迎发送或者推荐项目给我“门”:

bp@thejiangmen.com

点击右上角，把文章分享到朋友圈

⤵一键送你进入TechBeat快乐星球