该怎么把预训练Transformer范式用到强化学习里?
在当下的序列建模任务上,Transformer可谓是最强大的神经网络架构,并且经过预训练的Transformer模型可以将prompt作为条件或上下文学习(in-context learning)适应不同的下游任务。
大型预训练Transformer模型的泛化能力已经在多个领域得到验证,如文本补全、语言理解、图像生成等等。
从去年开始,已经有相关工作证明,通过将离线强化学习(offline RL)视为一个序列预测问题,那么模型就可以从离线数据中学习策略。
但目前的方法要么是从不包含学习的数据中学习策略(如通过蒸馏固定的专家策略),要么是从包含学习的数据(如智能体的重放缓冲区)中学习,但由于其context太小,以至于无法捕捉到策略提升。
DeepMind的研究人员通过观察发现,原则上强化学习算法训练中学习的顺序性(sequential nature)可以将强化学习过程本身建模为一个「因果序列预测问题」。
具体来说,如果一个Transformer的上下文足够长到可以包含由于学习更新而产生的策略改进,那它应该不仅能够表示一个固定的策略,而且能够通过关注之前episodes的状态、行动和奖励表示为一个策略提升算子(policy improvement operator)
这也提供了一种技术上的可行性,即任何RL算法都可以通过模仿学习蒸馏成一个足够强大的序列模型,并将其转化为一个in-context RL算法。
基于此,DeepMind提出了算法蒸馏(Algorithm Distillation, AD) ,通过建立因果序列模型将强化学习算法提取到神经网络中。
算法蒸馏将学习强化学习视为一个跨episode的序列预测问题,通过源RL算法生成一个学习历史数据集,然后根据学习历史作为上下文,通过自回归预测行为来训练因果Transformer。
与蒸馏后学习(post-learning)或专家序列的序列策略预测结构不同,AD能够在不更新其网络参数的情况下完全在上下文中改进其策略。
实验结果证明了AD可以在稀疏奖励、组合任务结构和基于像素观察的各种环境中进行强化学习,并且AD学习的数据效率(data-efficient)比生成源数据的RL算法更高。
AD也是第一个通过对具有模仿损失(imitation loss)的离线数据进行序列建模来展示in-context强化学习的方法。
算法蒸馏
2021年,有研究人员首先发现Transformer可以通过模仿学习从离线RL数据中学习单任务策略,随后又被扩展为可以在同域和跨域设置中提取多任务策略
这些工作为提取通用的多任务策略提出了一个很有前景的范式:首先收集大量不同的环境互动数据集,然后通过序列建模从数据中提取一个策略
把通过模仿学习从离线RL数据中学习策略的方法也称之为离线策略蒸馏,或者简称为策略蒸馏(Policy Distillation, PD)。
尽管PD的思路非常简单,并且十分易于扩展,但PD有一个重大的缺陷:生成的策略并没有从与环境的额外互动中得到提升。
例如,MultiGame Decision Transformer(MGDT)学习了一个可以玩大量Atari游戏的返回条件策略,而Gato通过上下文推断任务,学习了一个在不同环境中解决任务的策略,但这两种方法都不能通过试错来改进其策略。
MGDT通过微调模型的权重使变压器适应新的任务,而Gato则需要专家的示范提示才能适应新的任务。
简而言之,Policy Distillation方法学习政策而非强化学习算法。
研究人员假设Policy Distillation不能通过试错来改进的原因是,它在没有显示学习进展的数据上进行训练。
算法蒸馏(AD)通过优化一个RL算法的学习历史上的因果序列预测损失来学习内涵式策略改进算子的方法。
参考资料:
https://www.reddit.com/r/MachineLearning/comments/ye578r/r_incontext_reinforcement_learning_with_algorithm/
https://twitter.com/MishaLaskin/status/1585265485314129926