再掀强化学习变革!DeepMind提出「算法蒸馏」:可探索的预训练强化学习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) ,通过建立因果序列模型将强化学习算法提取到神经网络中。


论文链接:https://arxiv.org/pdf/2210.14215.pdf

算法蒸馏将学习强化学习视为一个跨episode的序列预测问题,通过源RL算法生成一个学习历史数据集,然后根据学习历史作为上下文,通过自回归预测行为来训练因果Transformer。与蒸馏后学习(post-learning)或专家序列的序列策略预测结构不同,AD能够在不更新其网络参数的情况下完全在上下文中改进其策略。

[if !supportLists]· [endif]Transfomer收集自己的数据,并在新任务上最大化奖励;

[if !supportLists]· [endif]无需prompting或微调;

[if !supportLists]· [endif]在权重冻结的情况下,Transformer可探索、利用和最大化上下文的返回(return)!诸如Gato类的专家蒸馏(Expert Distillation)方法无法探索,也无法最大化返回。

实验结果证明了AD可以在稀疏奖励、组合任务结构和基于像素观察的各种环境中进行强化学习,并且AD学习的数据效率(data-efficient)比生成源数据的RL算法更高。AD也是第一个通过对具有模仿损失(imitation loss)的离线数据进行序列建模来展示in-context强化学习的方法。

1

『算法蒸馏』

2021年,有研究人员首先发现Transformer可以通过模仿学习从离线RL数据中学习单任务策略,随后又被扩展为可以在同域和跨域设置中提取多任务策略。这些工作为提取通用的多任务策略提出了一个很有前景的范式:首先收集大量不同的环境互动数据集,然后通过序列建模从数据中提取一个策略。把通过模仿学习从离线RL数据中学习策略的方法也称之为离线策略蒸馏,或者简称为策略蒸馏(Policy Distillation, PD)。尽管PD的思路非常简单,并且十分易于扩展,但PD有一个重大的缺陷:生成的策略并没有从与环境的额外互动中得到提升。例如,MultiGame Decision Transformer(MGDT)学习了一个可以玩大量Atari游戏的返回条件策略,而Gato通过上下文推断任务,学习了一个在不同环境中解决任务的策略,但这两种方法都不能通过试错来改进其策略。MGDT通过微调模型的权重使变压器适应新的任务,而Gato则需要专家的示范提示才能适应新的任务。简而言之,Policy Distillation方法学习政策而非强化学习算法。研究人员假设Policy Distillation不能通过试错来改进的原因是,它在没有显示学习进展的数据上进行训练。算法蒸馏(AD)通过优化一个RL算法的学习历史上的因果序列预测损失来学习内涵式策略改进算子的方法。


AD包括两个组成部分:1、通过保存一个RL算法在许多单独任务上的训练历史,生成一个大型的多任务数据集;2、将Transformer使用前面的学习历史作为其背景对行动进行因果建模。由于策略在源RL算法的整个训练过程中不断改进,AD必须得学习如何改进算子,才能准确模拟训练历史中任何给定点的行动。最重要的是,Transformer的上下文大小必须足够大(即跨周期),以捕捉训练数据的改进。


在实验部分,为了探索AD在in-context RL能力上的优势,研究人员把重点放在预训练后不能通过zero-shot 泛化解决的环境上,即要求每个环境支持多种任务,且模型无法轻易地从观察中推断出任务的解决方案。同时episodes需要足够短以便可以训练跨episode的因果Transformer。


在四个环境Adversarial Bandit、Dark Room、Dark Key-to-Door、DMLab Watermaze的实验结果中可以看到,通过模仿基于梯度的RL算法,使用具有足够大上下文的因果Transformer,AD可以完全在上下文中强化学习新任务。


AD能够进行in-context中的探索、时间上的信用分配和泛化,AD学习的算法比产生Transformer训练的源数据的算法更有数据效率。

2

『PPT讲解』

为了方便论文理解,论文的一作Michael Laskin在推特上发表了一份ppt讲解。


算法蒸馏的实验表明,Transformer可以通过试错自主改善模型,并且不用更新权重,无需提示、也无需微调。单个Transformer可以收集自己的数据,并在新任务上将奖励最大化。尽管目前已经有很多成功的模型展示了Transformer如何在上下文中学习,但Transformer还没有被证明可以在上下文中强化学习。为了适应新的任务,开发者要么需要手动指定一个提示,要么需要调整模型。如果Transformer可以适应强化学习,做到开箱即用岂不美哉?但Decision Transformers或者Gato只能从离线数据中学习策略,无法通过反复实验自动改进。


使用算法蒸馏(AD)的预训练方法生成的Transformer可以在上下文中强化学习。


首先训练一个强化学习算法的多个副本来解决不同的任务和保存学习历史。


一旦收集完学习历史的数据集,就可以训练一个Transformer来预测之前的学习历史的行动。由于策略在历史上有所改进,因此准确地预测行动将会迫使Transformer对策略提升进行建模。


整个过程就是这么简单,Transformer只是通过模仿动作来训练,没有像常见的强化学习模型所用的Q值,没有长的操作-动作-奖励序列,也没有像 DTs 那样的返回条件。在上下文中,强化学习没有额外开销,然后通过观察 AD 是否能最大化新任务的奖励来评估模型。Transformer探索、利用、并最大化返回在上下文时,它的权重是冻结的!另一方面,专家蒸馏(最类似于Gato)不能探索,也不能最大化回报。


AD 可以提取任何 RL 算法,研究人员尝试了 UCB、DQNA2C,一个有趣的发现是,在上下文 RL 算法学习中,AD更有数据效率。


用户还可以输入prompt和次优的demo,模型会自动进行策略提升,直到获得最优解!而专家蒸馏ED只能维持次优的demo表现。


只有当Transformer的上下文足够长,跨越多个episode时,上下文RL才会出现。AD需要一个足够长的历史,以进行有效的模型改进和identify任务。


通过实验,研究人员得出以下结论:

[if !supportLists]· [endif]Transformer可以在上下文中进行 RL

[if !supportLists]· [endif]带AD的上下文 RL 算法比基于梯度的源 RL 算法更有效

[if !supportLists]· [endif]AD提升了次优策略

[if !supportLists]· [endif]in-context强化学习产生于长上下文的模仿学习

参考资料:https://www.reddit.com/r/MachineLearning/comments/ye578r/r_incontext_reinforcement_learning_with_algorithm/https://twitter.com/MishaLaskin/status/1585265485314129926

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 212,454评论 6 493
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 90,553评论 3 385
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 157,921评论 0 348
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 56,648评论 1 284
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 65,770评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,950评论 1 291
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,090评论 3 410
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 37,817评论 0 268
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,275评论 1 303
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 36,592评论 2 327
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,724评论 1 341
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,409评论 4 333
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,052评论 3 316
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,815评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,043评论 1 266
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 46,503评论 2 361
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 43,627评论 2 350

推荐阅读更多精彩内容