- 标题: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.背景
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}^{<i}, \boldsymbol{\tau}_{<t}\right)+\sum_{j=1}^{M} \log P_{\theta}\left(a_{t}^{j} \mid \mathbf{a}_{t}^{<j}, \mathbf{s}_{t}, \boldsymbol{\tau}_{<t}\right)+\log P_{\theta}\left(r_{t} \mid \mathbf{a}_{t}, \mathbf{s}_{t}, \boldsymbol{\tau}_{<t}\right)\right) L(τ)=t=1∑T(i=1∑NlogPθ(sti∣st<i,τ<t)+j=1∑MlogPθ(atj∣at<j,st,τ<t)+logPθ(rt∣at,st,τ<t))这里 τ < t \boldsymbol{\tau}_{<t} τ<t 表示从时间步 0 到 t − 1 t−1 t−1 的轨迹, s t < i \boldsymbol{s}_t^{<i} st<i 表示状态 s t \boldsymbol{s}_t st 的 0 到 i − 1 i−1 i−1 维, a t < j \boldsymbol{a}_t^{<j} at<j 也类似。具体实现上,作者使用 Adam 优化器来优化交叉熵损失,学习率从 0 线性 warmup 到 2.5 × 1 0 − 4 2.5×10^{−4} 2.5×10−4
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_{<t} τ<t)即可。这有点类似轨迹级的模仿学习,我们优化整个轨迹的分布,而非一般 BC 那样优化 ( s , a ) (s,a) (s,a) pair 的分布。如果这里将预测的序列长度设置为动作维度,那么该方法恰好对应于 Autoregress 的最简单的行为克隆形式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}^{<i}, \boldsymbol{\tau}_{<t}, \mathbf{s}_{T}\right) Pθ(sti∣st<i,τ<t,sT) GPT 的 mask self attention 机制本身就是以序列中的历史部分为条件进行生成的,借助该机制,作者使用来自 upside down RL 的 relabel 技巧,直接将 Offline 轨迹的最后一个状态 s T \mathbf{s}_T sT 作为目标状态放到序列最前面,这样就能让后面每个 token 生成时都关注到它,不用改模型结构。文章中这里写得不怎么清晰,如果模型直接输入调整过顺序的状态序列 s T , s 1 , s 2 , . . . , s T − 1 \mathbf{s}_{T},\mathbf{s}_{1},\mathbf{s}_{2},...,\mathbf{s}_{T-1} sT,s1,s2,...,sT−1 那么该模型只能生成 “为了实现 goal 应去向的状态”,但是不知道如何去向这里代码里似乎也没有 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 的一篇论文就进一步挖掘了这个优势