论文理解【Offline RL】——【TT】Offline Reinforcement Learning as One Big Sequence Modeling Problem
- 标题：Offline Reinforcement Learning as One Big Sequence Modeling Problem
- 文章链接：Offline Reinforcement Learning as One Big Sequence Modeling Problem
- 代码：jannerm/trajectory-transformer
- 官方主页：Offline Reinforcement Learning as One Big Sequence Modeling Problem
- 官方博客：Sequence Modeling Solutions for Reinforcement Learning Problems
- 发表：NIPS 2021
- 领域：离线强化学习（offline/batch RL）—— Transformer-Based / Model-Based
- 摘要：强化学习（RL）通常涉及到估计平稳（model-free）策略或单步模型（model-based），利用马尔可夫特性在时间上分解问题。然而，我们也可以将 RL 视为一个通用序列建模问题，其目标是产生一系列的动作，来获取一系列高奖励。从这个角度看，很容易想到高容量序列预测模型是否也可以为 RL 问题提供有效的解决方案，这些模型在 NLP 等其他领域工作良好。为此，我们使用 Transformer 架构来建模轨迹分布，将 beam search 作为一种规划算法，探索了如何通过序列建模工具来解决 RL 问题。这种从 RL 问题到序列建模问题的转换简化了一系列设计决策，使我们能够去掉 Offline RL 算法中常见的许多组件。我们在长期动态预测、IL、Goal conditioned RL 和 Offline RL 上展示了该方法的灵活性。此外，我们还证明了该方法可以与现有的 model-free 算法相结合，从而得到稀疏奖励且 long-horizon 任务中最先进的规划器。
文章目录
- 1.背景
- 2. 本文方法
-
- 2.1 思想
- 2.2 方法
-
- 2.2.1 模型设计
- 2.2.2 Beam Search (波束搜索)
- 3. 实验
-
- 3.1 建模能力
- 3.2 Attention 模式
- 3.3 各种任务上的性能
-
- 3.3.1 Offline RL
- 3.3.2 Imitation Learning & goal-reaching
- 4. 总结
1.背景
Offline RL是这样一种问题设定：Learner 可以获取由一批 episodes 或 transitions 构成的固定交互数据集，要求 Learner 直接利用它训练得到一个好的策略，而且禁止 Learner 和环境进行任何交互，示意图如下
关于 Offline RL 的详细介绍，请参考 Offline/Batch RL简介- Offline RL 是近年来很火的一个方向，下图显示了 2019 年以来该领域的重要工作，本文出现在 21 年，和同期的 DT 一样是最先纯粹使用 Transformer 模型解 Offline RL 问题的文章，可能也是最先把 Offline RL 当做序列建模问题来解的文章
- 本文使用的 Transformer 模型是 GPT，这是 Transformer 的 Decoder 部分，可以作为标准语言模型使用。所谓 “标准语言模型”，就是它会吃进去句尾的若干个词（token）然后预测下一个词是什么，这个过程可以反复进行从而实现文本生成的效果。GPT 相对传统语言模型的优势在于，GPT 模型结构中的 masked self attention 层利用 Q , K , V Q,K,V Q,K,V 矩阵在各个 token 间建立起了显式连接，从而解决了 RNN/LSTM 等传统序列模型的长跨度信息的遗忘问题。详细说明可以参考 快速串联 RNN / LSTM / Attention / transformer / BERT / GPT
2. 本文方法
2.1 思想
- 和之前介绍的非常类似的 【DT】Decision Transformer: Reinforcement Learning via Sequence Modeling 方法对比一下，二者都使用具有 AutoRegress 生成能力的 GPT 结构
- DT 对序列 τ = ( R ^ 1 , s 1 , a 1 , R ^ 2 , s 2 , a 2 , . . . , R ^ T , s T , a T ) \tau = \big(\hat{R}_1,s_1,a_1,\hat{R}_2,s_2,a_2,...,\hat{R}_T,s_T,a_T \big) τ=(R^1,s1,a1,R^2,s2,a2,...,R^T,sT,aT) 进行建模，其中 R ^ t = ∑ t ′ = t T r t ′ \widehat{R}_{t}=\sum_{t^{\prime}=t}^{T} r_{t^{\prime}} R t=∑t′=tTrt′ 代表未来获取的 return，其核心思想是 “基于历史序列，在当前状态下，以 return-to-go 作为目标来预测合适的 action”
其中对 R ^ , s , a \hat{R},s,a R^,s,a 三种 token 均使用线性层进行嵌入（如果 state 是图像形式则用卷积层嵌入），不像 NLP 里那样设置 vocab table 然后使用 nn.Embedding 进行嵌入 - TT（本文方法）直接对 Offline 轨迹 τ = ( s 1 , a 1 , r 1 , s 2 , a 2 , r 2 , … , s T , a T , r T ) \boldsymbol{\tau}=\left(\mathbf{s}_{1}, \mathbf{a}_{1}, r_{1}, \mathbf{s}_{2}, \mathbf{a}_{2}, r_{2}, \ldots, \mathbf{s}_{T}, \mathbf{a}_{T}, r_{T}\right) τ=(s1,a1,r1,s2,a2,r2,…,sT,aT,rT) 进行建模，然后使用 beam search 进行规划。这类似 model-based 的思路，但 TT 通过建模轨迹，实质上同时对环境转移、奖励函数和 behavior policy 进行了建模，这使其具有更多灵活性（比如可以做 IL）
上图可见 TT 使用了 NLP 里 GPT 那样的传统方法进行嵌入，先把 s , a s,a s,a 的每个维度离散为 vocab table，再使用 nn.Embedding 进行嵌入
- DT 对序列 τ = ( R ^ 1 , s 1 , a 1 , R ^ 2 , s 2 , a 2 , . . . , R ^ T , s T , a T ) \tau = \big(\hat{R}_1,s_1,a_1,\hat{R}_2,s_2,a_2,...,\hat{R}_T,s_T,a_T \big) τ=(R^1,s1,a1,R^2,s2,a2,...,R^T,sT,aT) 进行建模，其中 R ^ t = ∑ t ′ = t T r t ′ \widehat{R}_{t}=\sum_{t^{\prime}=t}^{T} r_{t^{\prime}} R t=∑t′=tTrt′ 代表未来获取的 return，其核心思想是 “基于历史序列，在当前状态下，以 return-to-go 作为目标来预测合适的 action”
2.2 方法
-
强化学习某种程度上可以视为序列生成问题，要生成能在环境交互中产生一系列高 reward 收益的一系列 action。具体而言，轨迹可以视为 state、action 和 reward 三类 token 组成的非结构化序列。过去已经有许多文章研究了应用 LSTM 等序列模型来表示标准 RL 算法中的组件，如策略、值函数和模型，但这些工作仍然依赖于标准 RL 组件，本文的目标是使用序列建模来取代尽可能多的 RL pipeline，从而产生一种更简单的方法，其有效性取决于序列模型的表征能力而非算法的复杂性。具体而言，作者训练了一个单一的高容量序列模型来表示 state、action 和 reward 序列上的联合分布，这可以
- 作为 model-based 模型，进而通过规划得到策略
- 作为 behavior policy，用于模仿学习
- 作为 policy constraint，用于 Offline RL
第 3 点多解释一句，传统基于 TD-Learning 的 Offline RL 需要 “策略约束” 或者说 “对不确定性的悲观假设”，以避免 agent 经历 Offline 数据集中少见或不存在的 transition，因为这些 transition 对应的 Q ( s , a ) Q(s,a) Q(s,a) 无法良好估计，具体可参考 论文理解【Offline RL】——【BCQ】Off-Policy Deep Reinforcement Learning without Exploration 。本文中 TT 直接对 Offline 轨迹建模，它同时生成数据集分布内的 state 和 action，因此不用担心生成 OOD 的 ( s , a ) (s,a) (s,a)
某种程度上，Trajectory Transformer 可以看做 model-based RL 和（隐式）policy constraints 的结合
2.2.1 模型设计
-
在实现层面上，TT 这种 “GPT 建模 + Beam Search 规划” 和策略几乎与 NLP 中常见的做法相同。因此本文主要考虑的不是体系结构设计，而是如何表示轨迹数据。由于轨迹可能由连续的 state 和 action 组成，而传统 NLP 中 token 都是离散的，需要进行适当的离散化。具体而言，作者对动作空间和状态空间的每一个维度都分别进行离散，即把原始轨迹
τ = ( s 1 , a 1 , r 1 , s 2 , a 2 , r 2 , … , s T , a T , r T ) \boldsymbol{\tau}=\left(\mathbf{s}_{1}, \mathbf{a}_{1}, r_{1}, \mathbf{s}_{2}, \mathbf{a}_{2}, r_{2}, \ldots, \mathbf{s}_{T}, \mathbf{a}_{T}, r_{T}\right) τ=(s1,a1,r1,s2,a2,r2,…,sT,aT,rT) 离散为（这里假设状态空间有 N N N 维，动作空间有 M M M 维）
τ = ( . . . , s t 1 , s t 2 , . . . , s t N , a t 1 , a t 2 , . . . , a t N , r t , . . . ) t = 1 , . . . , T \tau = \left(...,s_t^1, s_t^2,...,s_t^N,a_t^1, a_t^2,...,a_t^N,r_t,...\right) \quad t=1,...,T τ=(...,st1,st2,...,stN,at1,at2,...,atN,rt,...)t=1,...,T 这样处理后，轨迹长度从 3 T 3T 3T 增长为 T ( N + M + 1 ) T(N+M+1) T(N+M+1)根据代码来看，作者具体实现的 GPT 堆叠了 4 个 transformer decoder block，使用 4 个 attention 头。state、action、reward 和 reward-to-go (见下文) 都离散为 100 个取值，这样 token 总数为 100 ∗ ( o b s _ d i m + a c t i o n _ d i m + 2 ) + 1 100*(obs\_dim+action\_dim+2)+1 100∗(obs_dim+action_dim+2)+1，最后的 + 1 +1 +1 是 stop token
虽然这种各个维度独立离散化的操作效率较低，但尽量保持了轨迹分布模型的容量，和简化为对角多维高斯分布相比，其能描述的轨迹分布更多样。作者比较了以下两种离散化方法
Uniform：就是对每个连续维度的取值范围进行均匀划分，若某 state 维的 vocab size 为 V V V，则该维离散后的第 i i i 个离散值覆盖的取值空间宽度为 ( max s i − min s i ) / V (\max s^i−\min s^i)/V (maxsi−minsi)/V。这样的好处是保留了原始连续空间中关于欧氏距离的信息，这可能比训练数据的分布更能反映问题的结构，但易受到离群点影响，导致很多离散取值没有对应的 Offline 样本Quantile：划分后的每个维度的每个离散值都占有等量的概率质量。这样破坏了欧氏距离信息，但是所有离散取值标记都在数据中存在对应的样本，避免分布不均匀问题干扰监督学习
-
训练时使用 teacher-forcing 方式，设 TT 参数为 θ \theta θ，最大化以下联合对数似然函数
L ( τ ) = ∑ t = 1 T ( ∑ i = 1 N log P θ ( s t i ∣ s t < i , τ < t ) + ∑ j = 1 M log P θ ( a t j ∣ a t < j , s t , τ < t ) + log P θ ( r t ∣ a t , s t , τ < t ) ) \mathcal{L}(\tau)=\sum_{t=1}^{T}\left(\sum_{i=1}^{N} \log P_{\theta}\left(s_{t}^{i} \mid \mathbf{s}_{t}^{
2.2.2 Beam Search (波束搜索)
- 本文使用的规划方法
Beam Search（波束搜索）是使用序列模型时的常见组件。对于序列生成任务，模型输出是一个时间步一个时间步依次生成的，而且前面时间步的结果还会影响后面时间步的结果（AutoRegress）。由于每步输出都是基于历史序列的条件概率，需要一个称为解码的额外动作来融合模型多个时间步的输出，使得最终得到的序列的每一步条件概率连乘起来最大。一个简单的想法是遍历所有可能计算联合概率，从而找出最大概率序列，但这种方法的复杂度是指数上升的，假设 vocab size 为 v v v，那么长 T T T 的序列种类就有 v T v^T vT 种之多，通常无法实现。要降低复杂度，就要控制展开的程度，一个极端情况是贪心搜索，即每一个时间步都取出一个条件概率最大的输出，再将从开始到当前步的结果作为输入去获得下一个时间步的输出。如下文本生成例子中，用贪心策略生成的序列为 “ABC”
这样做将原来指数级别的求解空间直接压缩到了线性级别，但由于丢弃了绝大多数的可能解，这种用局部最优组成的序列质量不高。Beam Search 就是一个可以通过调节展开程度来调节 “轨迹质量” 和 “计算复杂度” 间平衡位置/倾向性的算法，它的想法也很简单，就是每次展开保留 top N N N 个最可能的输出，可见贪心搜索就是 N = 1 N=1 N=1 情况下的特殊 Beam Search。下图是一个 N = 2 N=2 N=2 的例子
每步我们保留 top 2 最可能的序列，如图可见分别是 A/C、AB/CE 和 ABD/CED，假如我们的目标输出长度就是 3，那么只需在 ABD/CED 中选择概率较高的那个输出即可。可见 Beam Search 的计算复杂度是贪心搜索的 N N N 倍，是一种牺牲时间换性能的方法。本文给出 Beam Search 伪代码如下
其中 ( ) () () 代表空序列， ∘ \circ ∘ 代表序列连接 - 在用 Beam Search 规划时使用不同的标准，可以做 IL、Goal conditioned RL、Offline RL 等等任务。具体而言
Imitation learning：模仿学习的目标是再现训练数据中的轨迹分布，这时直接优化 Offline 轨迹的出现概率即可，因此无需修改上述伪代码，模型输入和上面一样设计为 τ = ( . . . , s t 1 , s t 2 , . . . , s t N , a t 1 , a t 2 , . . . , a t N , r t , . . . ) t = 1 , . . . , T \tau = \left(...,s_t^1, s_t^2,...,s_t^N,a_t^1, a_t^2,...,a_t^N,r_t,...\right) \quad t=1,...,T τ=(...,st1,st2,...,stN,at1,at2,...,atN,rt,...)t=1,...,T Beam Search 时概率条件 x \pmb{x} x 设置为 s t s_t st（和可选的历史轨迹 τ < t \tau_{Goal conditioned RL：作者将 goal 设定为将在未来轨迹末尾达到的状态 s T \mathbf{s}_T sT，这样需要优化的状态条件概率为
P θ ( s t i ∣ s t < i , τ < t , s T ) P_{\theta}\left(s_{t}^{i} \mid \mathbf{s}_{t}^{代码里似乎也没有 GCRL 这部分的相关代码，不过我看作者的
ConditionalGPT类是直接继承GPT类实现的，可能模型输入还是带有 action 和 reward 的Offline RL：Offline RL 需要生成 action 来最大化累计折扣 reward，只需将 Beam Search 中的优化目标从预测概率 P θ ( Y ∣ x ) P_\theta(Y|\pmb{x}) Pθ(Y∣x) 改成预测 reward 即可，这可以理解为 “用最优 token 的对数概率代替最相似 token 的对数概率”。但这种类似 greedy 策略的做法会导致短视问题，为此作者引入了额外的 reward-to-go 信息
R t = ∑ t ′ = t T γ t ′ − t r t ′ R_t = \sum_{t'=t}^T \gamma^{t'-t}r_{t'} Rt=t′=t∑Tγt′−trt′ 加入这个 token 后，轨迹将由四元组 ( s t , a t , r t , R t ) (\pmb{s}_t,\pmb{a}_t,r_t,R_t) (st,at,rt,Rt) 组成，将 s , a , r , R \pmb{s,a},r,R s,a,r,R 分别离散化后得到模型输入，每 N + M + 1 N+M+1 N+M+1 个 token 预测一个 reward-to-go，Beam Search 规划时最大化 r t + R t r_t+R_t rt+Rt 即可。总体上看，用 TT 做 Offline RL 的过程可以概括为 “预训练阶段对 Offline 轨迹建模，规划阶段先抽样候选动作，再使用模型评估它们的效果，选择奖励最大化的轨迹”。另外有一点值得强调：通过联合建模 state 和 action 并使用相同的过程对它们进行采样，可以避免查询 OOD 的 ( s , a ) (s,a) (s,a) pair
3. 实验
3.1 建模能力
- TT 本身是一个 long-horizon policy-conditioned predictive model，作者首先检验其建模能力。在 humanoid 环境中使用一个训练好的策略收集交互数据，然后分别训练 TT 和来自 PETS（一个 SOTA model-based 方法）的 feedforward Gaussian dynamics model，再分别预测 100 步轨迹，可见 Transformer 在长跨度轨迹建模上具有巨大优势
事实上，过去的很多 Model-based 方法都会有意控制 horizon 在较短的范围内以避免上面这样因级联错误而导致崩溃，还没有任何 Model-based 方法能在类似维度的任务上得到类似长度的高精度预测 - 作者进一步进行定量分析来评估模型的误差累积与预测范围。使用不同模型从一个固定起点采样 1000 条轨迹来估计每个模型在每个时间步预测的状态边际，然后考察 behavior policy 访问这些预测边缘内状态的概率，如下
这里 Transformer 就是 condition on previous trajectory 的 TT，而 Markovian Transformer 是 TT 的一个消融，只 condition on 最近的一个 MDP 步（由于各个维度进行了离散表示，所以即使一个时间步也是有一定长度的序列）。另外右边图是一个 humanoid 环境的 POMDP 变体，每个状态的每个维度都以 50% 的概率被掩盖。这里得到三个结论- Transformer-based 模型的轨迹建模能力远强于过去的 Model-based 方法
- 全观测情况下，只 condition on 最近的一个时间步（和过去 Model-based 方法相同）的性能和 condition on 更长历史序列差不多，说明 “模型结构差异” 和 “状态离散化增加的表达能力” 在 TT 的长范围精度中发挥了很大的作用
- 部分观测情况下，考察更长的历史序列可以提高建模准确性
3.2 Attention 模式
- 作者可视化了 TT 预测过程中的 attention 向量，发现了两种主要的注意模式。下图每行都是一个轨迹，从上到下显示了轨迹自回归生成的过程，用颜色表示关注程度，即每一列都是同一个 token 在不同时刻的关注程度。这里的两张图分别是第一层和第三层注意力头生成的，省略了 reward 维度
- 如左图发现了 Markovian 关注模式，即生成时只关注最近的 transition
- 如右图发现了 action 的各个维度强烈地依赖于之前多个时间步长前的相同维度，而 state 没有呈现这种关系
3.3 各种任务上的性能
3.3.1 Offline RL
-
这是 TT 最难适用的一个任务，因为奖励最大化行为通常和混合质量 Offline 数据集中的行为分布差异最大。作者使用 D4RL benchmark，对比方法包括 BC、CQL（RL-Based 路线 SOTA）、MBOP（Model-Based 路线 SOTA）和 DT。性能对比如下
TT 的性能与之前的所有方法媲美或更好，两种离散化 TT 性能差不多，除了 HalfCheetah 的 Med-Expert，这个数据集中速度变化范围太大，导致均匀离散化不佳
-
另外，作者尝试了将 TT 和 Q 函数结合。TT 本质还是一种 MC，其中的 Reward-to-go 是 MC 得到的，这对于许多标准 Offline RL 任务是足够的，但在稀疏奖励和长 horizon 设置中，Offline 数据过于缺乏信息以至无法指导基于 Beam Search 的规划过程。作者使用 AntMaze 这一稀疏奖励环境，将 TT 中的 Reward-to-go 替换为另外训练的 RL-Based Model-free 方法 IQL 中所得的 Q 函数，结果如下
可见这种结合方式性能最佳，甚至超过了提供 Q 函数的 IQL，这说明启发式搜索方法更不容易受到 Q 函数中错误的影响，反过来讲，启发式搜索方法对 Q 函数的质量要求相比传统方法低
3.3.2 Imitation Learning & goal-reaching
- 作者测试 TT 做 IL 的效果，在 Hopper 和 Walker2d 环境上直接用 TT 对专家轨迹建模并规划，得到了 104% 和 109% 的 return，说明了 TT 做 IL 的有效性。虽然这个结果不出乎意料，但它说明了 NLP 建模的解码算法也可有效用于控制
- 测试用 TT 做 GCRL 时，作者使用了经典的 four rooms 环境，先训练一个 goal-reaching agent，通过在环境中均匀采样起点和终点，收集预训练 agent 的轨迹作为 Offline 数据集。TT 的效果不错
4. 总结
- 本文证明了大规模序列模型可以直接解决 Offline RL 问题，且在 IL 和 goal-reaching 等变形问题上均有效，而无需任何传统方法中常见的价值估计等组件。但其问题在于 Transformer 的计算量非常高，当上下文窗口太大时，需要数秒钟进行动作选择，无法用于大多数动态系统的实时控制
- 和之前介绍的 DT 一样，这两篇文章都完全抛弃了传统 RL 中的 MDP 因果图结构。之前的 PETS, world model, dramer 等 Model-based 方法都会遵循标准 MDP 中对策略、环境转移和奖励函数等的定义（也就是状态分布的条件是上一步的状态，而动作、奖励、价值都由当前的状态决定）。整个强化学习社区一般相信这样能提高样本效率，不过这样的图结构其实也可能是一种制约。自然语言领域从 RNN 到 Transformer 以及计算机视觉领域 CNN 到 Transformer 的转变其实都体现了：随着数据增加，让网络自己学习图结构更有利于获得表现更好的模型
- 另外，TT 同时对环境和 behavior policy 建模，这带来诸多优势，比如
- 不用担心 Offline RL 中常见的外推误差问题，无需策略约束或悲观假设等组件
- 具有轨迹生成能力，有能力做数据增强，实时上今年 NIPS 的一篇论文就进一步挖掘了这个优势