【源头活水】RL Transformer之Trajectory Transformers

“问渠那得清如许，为有源头活水来”，通过前沿领域知识的学习，从其他研究领域得到启发，对研究问题的本质有更清晰的认识和理解，是自我提高的不竭源泉。为此，我们特别精选论文阅读笔记，开辟“源头活水”专栏，帮助你广泛而深入的阅读科研文献，敬请关注。

来源：知乎—鱼子酱

地址：https://zhuanlan.zhihu.com/p/384458362

最近RL的大佬们又开始整活了，前有Pieter Abbeel放出Decision Transformers（DT）[1]，后有Sergey Levine放出Trajectory Transformers（TT）[2]，两者的motivation都是一样的：能不能尝试将RL问题（序列决策问题）看成一个序列生成任务，再借助modern architecture（e.g. Transformer）去完成此序列生成任务。

这样做其实就完全抛弃了传统RL用Bellman方程TD求解或者policy gradient求解的方式，但也有潜在好处，比如直接通过transformer里面的self-attention完成RL信度分配问题而不需要用Bellman回溯的方式缓慢地传播reward，同时对稀疏奖励更加effective（这是Pieter的statement，但我觉得没有足够的证据，更多的是实验证明），也不再需要像Offline RL那样加各种约束（这是Sergey的statement）。

Pieter和Sergey的motivation虽然一样，但是做法却迥然。可解释性上，个人觉得Sergey的TT做得更好，故本篇先从TT说起。

01

Background Requirements

Offline RL, GPT (不懂GPT的可以先去看一下，简单说就是GPT只用了Transformer decoder来做序列任务) and Beam search.

02

Brief Introduction

在Trajectory Transformers（TT）中，核心是如何表示trajectory。表示完了之后就用GPT transformer来做sequence modeling的训练。在实际测试时，给定一个初始状态，transformer自回归地生成trajectory，结合beam search找到最可能、reward最大的trajectory（自然就有了下一个动作，相当于做了决策了）。

03

Trajectory Representation

A trajectory  consists of N-dimensional states, M-dimensional actions, and scalar rewards:

假设state的维度是N，action的维度是M，reward是一个标量，T个transition。这里要注意的是state或者action的每个维度按理说应该是标量，但这里表示成矢量是因为实际输入的时候时该维度value对应字典里面的index，所以是一个one-hot向量（（别问什么是字典，问就是先去看最原始的Transformer做机器翻译，或者去参考GPT））。

构建字典

如果state和action离散的，则可以正常构建字典；但如果是连续的，则需要先做离散化处理。先设置每个维度一个bin的数量V，假设  , 则离散化为

作者在离散时设置了offset，为了保证所有维度不重叠（方便output，后续会说）。这样，action tokens  的offset就是Vx(N+j), reward tokens  的offset就是Vx(N+M)。所以，input vocabulary里面有Vx(N+M+1)个词（作者在offline RL setting中引入了一个reward-to-go，所以在此setting下，有Vx(N+M+2)个词）。而对于output vocabulary，因为我们设置了offset，所以output vocabulary可以缩减到含有V个词。

04

Training

GPT中的attention是causal的，即当前token只跟在它之前的token做attention。训练时因为有grountruth trajectory，所以我们可以直接并行输入trajectory tokens，然后输出所有的预测节点，每个节点是一个V维的向量，经过softmax，得到概率分布，再和groundtruth做交叉熵，得到log似然值（最小化交叉熵，相当于在最大化log似然），于是training loss可以写成

这个式子看上去很长，不好理解，但其实很好理解。先不看t维度的求和，先看括号里面。括号里面有三部分，即为三种case（预测state, action和reward）下的log似然，然后可以仔细观察下这三个条件概率，其实就是给定在当前时刻之前的trajectory和当前时刻当前维度之前的token，当前时刻的token的条件概率（如果把input tokens按照时间顺序从左到右展开，那么就是给定当前token所有左边的token，当前时刻的token的概率）。

05

Testing/Inference

如果transformer训练好了，该怎么测试？因为GPT结构只用transformer decoder，所以来了一个state，transformer decoder可以自回归向后预测trajectory，这自然include下一个动作，相当于做了决策。具体地，先用beam search搜索似然概率最大的B个预测的trajectory，再从中选取reward+reward-to-go最大的那个trajectory作为最终的预测结果。

Beam search算法：

讲道理这算法的伪代码写得... 感觉没有对符号、表示做充分的说明，本来很简单的东西，被搞得就很迷幻。这里， 是size为V的output的字典，L其实就是input要预测的序列长度（比如M=3，N=2，预测reward和reward-to-go，如果要预测10个transition的trajectory，则L=T*(M+N+2)=70），B是beam search的宽度， 表示{(  , q), ... }序列， 为trajectory，q是这个trajectory对应的概率。

06

Experiment

作者在3种setting（imitation learning，goal-conditioned RL和offline RL）上验证TT的效果。

Trajectory Prediction Accuracy

Reference指的是groundtruth，中间的TT预测的结果，最后一行是SOTA planning算法预测的结果，可以看出TT对于这种长时空的预测效果还是很好的，作者后面也分析了，这得益于transformer architecture powerful capacity。

此为量化比较，从log似然（左边curves）可以观察到Transformer和Markovian Transformer（将attention的感受野限制在有限的窗口种，不再是之前所有的token）的长时空预测效果很好，同时error propagation没有那么严重。也就是在full-observed的环境中其实缩短感受野并不是特别影响性能，那么可以推出可能transformer的architecture起了重要作用。右边curves是在ablation当我们把input state以50%的概率在每个维度上随机mask时的效果，这种setting下，全感受野（Transformer）的效果要好于有限窗口感受野（Markovian Transformer）。

Attention Patterns

这里作者想研究transformer到底在干些什么，于是可视化token之间的响应图（颜色越深，对应的两个token响应越大）。左边是根据第一层attention layer得到的，右边是第三层得到的。作者观察到这两种不同形式的响应图。从左边这种可以发现，transformer发现了Markovian strategy（因为任意一个state或者action的维度只跟它之前的token有较大响应）。在右边这种中，当前state的各维度跟之前state的相同维度响应较大，当前action跟之前的action响应较大（作者说让人联想起trajectory optimization algorithms中的action smoothing，不明觉厉，我只觉得像action sticking，哈哈哈）。

Offline RL

对比了Offline RL的SOTA method CQL，behavior cloning（BC），以及model-based方法MOPO和MBOP，可以发现Trajectory Transformers Optimization （TTO）的效果不亚于/好于这些方法。这里我也对比了TTO和Decision Transformer（DT）的效果，发现DT的效果似乎好一点。如果从理论上分析，TTO一定比BC好，因为BC其实可以认为假设了一阶马尔可夫性质，就是拿当前state就能推断出action，但实际环境中很难完全满足这种性质，TTO相当于建模了之前的trajectory，是一种Markovian-orthogonal的方法，不对马尔可夫性质有要求。

Imitation and Goal-reaching/conditioned RL

(1) imitation：作者尝试只用最大化似然配合beam search来control（也就是在标准TTO中不考虑reward-maximizing），然后发现这种情况下TTO做imitation learning的效果略好于专家。behavior cloning其实也能重现专家的效果。这暗示了用TTO这种decoding algorithm也能有效用于控制。

用一个pretrained goal-conditioned agent采集trajectory用作Offline data来训练TTO，然后做anti-causal goal-state conditioning的TTO training（这里不展开了，文章里有写，主要是说把goal state放到trajectory最前面，s_goal, s_t, a_t, r_t,..., 这种非因果的建模其实是在算给定之前的trajectory以及期望发生的状态的条件下，最大化下一个state/action/reward token的似然概率）。上图都是TTO产生的trajectory，也不知道作者放上去这个能说明啥，没有对比，啥也不是。。。

07

Discussion

Limitation: inference的时候速度慢，在连续环境上做离散化有点麻烦。

参考

1. Decision Transformer: Reinforcement Learning via Sequence Modeling https://arxiv.org/abs/2106.01345

2. Reinforcement Learning as One Big Sequence Modeling Problem https://arxiv.org/abs/2106.02039

本文目的在于学术交流，并不代表本公众号赞同其观点或对其内容真实性负责，版权归原作者所有，如有侵权请告知删除。

“源头活水”历史文章

更多源头活水专栏文章，

请点击文章底部“阅读原文”查看

分享、在看，给个三连击呗！