[论文解读] Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation

Emilio Parisotto, Ruslan Salakhutdinov|arXiv (Cornell University)|Apr 4, 2021

Reinforcement Learning in Robotics被引用 9

一句话总结

本文提出演员-学习者蒸馏（Actor-Learner Distillation, ALD），一种持续蒸馏方法，可在演员-延迟约束的强化学习设置中，将大型高容量Transformer学习者模型的知识迁移至小型快速LSTM演员模型。通过利用Transformer的样本效率，同时保持LSTM的推理速度，ALD在训练时间与小型模型相当的情况下，实现了更优的样本效率。

ABSTRACT

Many real-world applications such as robotics provide hard constraints on power and compute that limit the viable model complexity of Reinforcement Learning (RL) agents. Similarly, in many distributed RL settings, acting is done on un-accelerated hardware such as CPUs, which likewise restricts model size to prevent intractable experiment run times. These "actor-latency" constrained settings present a major obstruction to the scaling up of model complexity that has recently been extremely successful in supervised learning. To be able to utilize large model capacity while still operating within the limits imposed by the system during acting, we develop an "Actor-Learner Distillation" (ALD) procedure that leverages a continual form of distillation that transfers learning progress from a large capacity learner model to a small capacity actor model. As a case study, we develop this procedure in the context of partially-observable environments, where transformer models have had large improvements over LSTMs recently, at the cost of significantly higher computational complexity. With transformer models as the learner and LSTMs as the actor, we demonstrate in several challenging memory environments that using Actor-Learner Distillation recovers the clear sample-efficiency gains of the transformer learner model while maintaining the fast inference and reduced total training time of the LSTM actor model.

研究动机与目标

• 解决在推理速度严格受限的演员-延迟约束强化学习环境中部署高容量模型的挑战。
• 实现在机器人等现实应用中使用样本高效的大模型（如Transformer），其中计算和功耗限制了推理阶段的模型大小。
• 开发一种持续蒸馏程序，将高容量学习者模型的知识压缩至低容量演员模型，同时不损失样本效率。
• 通过使用快速、小型的演员模型，在分布式同步策略强化学习中保持低实验运行时间，同时仍能受益于大型学习者模型的学习进展。
• 证明在部分可观察环境中，不同架构之间（特别是Transformer与LSTM）的蒸馏可有效传递归纳偏置和性能提升。

提出的方法

• ALD框架采用一个中心化学习者，其使用大型Transformer模型，通过交互学习并提升样本效率。
• 一个独立的轻量化LSTM基演员模型在推理速度约束下执行推理，确保低延迟决策。
• 通过持续策略蒸馏，将知识从Transformer学习者传递至LSTM演员，使演员模型被训练以模仿学习者的策略行为。
• 蒸馏过程在训练期间在线应用，使用学习者策略输出作为软目标，监督演员策略的更新。
• 该方法运行于分布式同步策略强化学习设置中，演员收集轨迹并将其发送至中心化学习者以执行策略更新。
• 该方法在保持LSTM演员计算效率的同时，捕获了Transformer学习者在内存密集型环境中的优越样本效率。

实验结果

研究问题

• RQ1 能否在不牺牲推理速度的前提下，使用大型样本高效Transformer模型来训练小型快速LSTM演员模型？
• RQ2 从高容量学习者到低容量演员的持续蒸馏是否能在部分可观察环境中保持学习者的样本效率优势？
• RQ3 ALD能否在训练运行时间与小型模型相当的同时，达到大型模型的性能水平？
• RQ4 在演员-延迟约束的强化学习环境中，不同架构之间（特别是Transformer与LSTM）的知识迁移效果如何？
• RQ5 从Transformer到LSTM的蒸馏是否能提升内存密集型强化学习任务中的泛化能力和鲁棒性？

主要发现

• ALD使LSTM演员在具有挑战性的记忆环境中实现了与Transformer学习者相当的样本效率，显著优于从零开始训练的基线。
• 尽管使用了大型Transformer进行学习，该方法仍保持了与LSTM单模型基线相似的实验运行时间。
• 在部分可观察环境中，ALD在保留LSTM低延迟推理的同时，恢复了Transformer的样本效率优势。
• 持续蒸馏过程有效将Transformer的归纳偏置传递至LSTM，提升了学习稳定性和性能。
• ALD表明，不同架构之间的知识蒸馏在演员-延迟约束的强化学习中可极为有效，即使学生与教师具有不同的归纳偏置。
• 该方法与架构效率改进正交，意味着可与高效Transformer变体结合，进一步提升性能。