【论文阅读】Decision Transformer: Reinforcement Learning via Sequence Modeling

1 本文解决了什么问题？

本文将强化学习抽象为一个序列建模问题，并利用 Transformer 架构进行建模，提出 Decision Transformer。该架构与之前的强化学习方法（值函数和策略梯度方法）均不同，Decision Transformer 只通过一个因果掩码 Transformer 来输出最有动作。通过将期望收益（奖励）、过去的状态和动作进行自回归模型建模，Decision Transformer 可以生成实现期望收益的未来动作。

本文并非是使用 Transformer 作为传统强化学习方法的框架，而是 研究轨迹建模（trajectory modeling，类似于语言建模）是否可以替代传统的强化学习算法。

2 本文的解决方法是什么？

本文提出的轨迹建模模型 Decision Transformer 训练流程如下：

Trajectory representation.

本文选择轨迹表示（trajectory representation）的关键要求是：

(a) 它应该使 Transformer 能够学习有意义的模式；

(b) 我们应该能够在测试时有条件地生成动作。

对奖励进行建模是非常重要的，因为我们希望模型能够根据未来的期望收益生成动作，而非过去的收益。因此，本文并不直接使用瞬时收益作为奖励，而是将奖励建模为 return-to-go（未来收益之和），即 R^t=∑t′=tTrt′\widehat{R}_{t}=\sum_{t^{\prime}=t}^{T} r_{t^{\prime}}Rt=∑t′=tTrt′。

于是，轨迹表示（trajectory representation）可写成如下形式：
τ=(R^1,s1,a1,R^2,s2,a2,…,R^T,sT,aT)\tau=\left(\widehat{R}_{1}, s_{1}, a_{1}, \widehat{R}_{2}, s_{2}, a_{2}, \ldots, \widehat{R}_{T}, s_{T}, a_{T}\right) τ=(R1,s1,a1,R2,s2,a2,…,RT,sT,aT)

Architecture.

作者将最后 KKK 个时间步输入 Decision Transformer，因此总共有 3K3K3K 个 token（return-to-go、状态和动作）。为了对 token 进行编码（嵌入），作者训练每个模态的线性层，它将原始输入投影到嵌入维度，并进行层归一化。

值得注意的是，这与 Transformer 使用的标准位置嵌入不同，此处一个时间步对应了 3 个 token。

随后，token 将由 GPT 模型进行处理，该模型通过自回归模型预测未来的动作 token。

Training.

本文从数据集中采样小批量序列长度为 KKK 的样本。与输入 token sts_tst 对应的预测头经过训练，以预测 ata_tat（离散动作使用交叉熵损失函数，连续动作使用均方误差），并对每个时间步的损失取平均值。

Evaluation.

在评估 rollout 期间，作者根据期望的性能（例如，指定生成专家行为的最大可能回报）以及环境启动状态指定目标回报，以初始化动作生成。在执行完生成的动作后，作者用所获得的奖励减去目标返回值，并获得下一个状态。重复这个生成动作并执行的过程，以获得下一个 return-to-go 和状态，直到 episode 终止。

Decision Transformer 的伪代码如下：

3 有什么关键的实验结果？

本文主要与 TD Learning 和模仿学习做对比实验。TD Learning 代表着传统强化学习的先进技术，而模仿学习有类似于 Decision Transformer 的公式。

offline-RL：https://zhuanlan.zhihu.com/p/414497708

TD Learning：这些方法大多使用了 action-space constraint 或 value pessimism，代表标准的 RL 方法。最先进的无模型方法是 Conservative Q-Learning （CQL），这是本文的主要比较方法。此外，本文还比较了其他无模型 RL 算法，如 BEAR 和 BRAC。
Imitation learning：该方法同样使用监督损失进行训练，而不是使用 Bellman backups。本文在这里使用行为克隆（BC）。

作者评估了 离散的（Atari） 和 连续的（OpenAI Gym） 控制任务。前者需要长期的信用分配，而后者需要细粒度的连续控制，代表了一系列不同的任务。下图总结了本文的主要结果，其中显示了每个环境的平均专家标准化性能。

Atari

由于 Atari 的高维视觉输入，以及由于动作和结果奖励之间的延迟而导致的信用分配困难，使它具有挑战性。作者在 DQN-replay 数据集中所有样本的 1% 上评估我们的方法，代表在线 DQN 代理在训练期间观察到的 5000万个 transitions 中的 50万个。作者实验了 3 个种子的均值和标准差，将职业玩家的分数归一化，其中 100 代表职业玩家分数，0 代表随机策略。

作者比较了 4 个Atari 游戏，包括 Breakout，Qbert，Pong 和 Seaquest。本文的方法在 4 个游戏中有 3 个可以与 CQL 竞争，并且在所有 4 个游戏中都超过或匹敌 REM，QR-DQN 和 BC。

OpenAI Gym

本文将考虑来自 D4RL 基准的连续控制任务，还考虑了不属于基准测试的 2D reacher 环境，并使用类似于 D4RL 基准测试的方法生成数据集。“猎豹”是一种目标条件任务，奖励很少，所以它代表了不同于标准运动环境（HalfCheetah、Hopper 和 Walker）的设置。下面描述了不同的数据集设置。

Medium：Medium 策略生成的 100万个时间步，其得分约为专家策略的三分之一；
Medium-Replay：智能体的回放缓冲区经过 Medium 策略的性能训练（在我们的环境中大约为 25k-400k 时间步）；
Medium-Expert：Medium 策略生成的 100万个时间步与专家策略生成的 100万个时间步相连接。

本文的结果如下表所示。Decision Transformer 在大多数任务中都取得了最高的分数，并且在剩下的任务中可以与最先进的技术相媲美。

4 相关讨论

4.1 Does Decision Transformer perform behavior cloning on a subset of the data?

本文试图深入了解是否可以将 Decision Transformer 视为对数据子集进行模仿学习，并获得一定的回报。为了研究这一点，作者提出了一种新的方法，百分位行为克隆（percent Behavior Cloning， %BC），其中作者只在数据集中前 X% 的时间步上运行行为克隆，按 episode 的收益排序。百分比 X% 在标准 BC（X = 100%）之间插入，该 BC 在整个数据集上训练，只克隆观察到的最佳数据轨迹。

当数据充足时 —— 就像在 D4RL 模式中一样 —— 作者发现 %BC 可以匹配或击败其他离线 RL 方法。在大多数环境中， Decision Transformer 的性能与最好的 %BC 不相上下，这表明它可以在整个数据集分布上训练后专注于特定子集。

相比之下，当研究低数据机制时 —— 比如 Atari，我们使用 1% 的重放缓冲区作为数据集 —— %BC 的表现比较弱。这表明，在数据量相对较低的情况下，Decision Transformer 可以通过使用数据集中的所有轨迹来提高泛化，从而超过 %BC，即使这些轨迹与返回条件反射目标不同。

结果表明，Decision Transformer 可以比简单地在数据集的子集上进行模仿学习更有效。在本文考虑的任务中，Decision Transformer 要么优于 %BC，要么与之性能相当，而不需要选择最优子集。

4.2 How well does Decision Transformer model the distribution of returns?

作者通过在大范围内改变期望的目标回报来评估 Decision Transformer 理解 return-to-go tokens 的能力——评估 transformers 的多任务分布建模能力。下图显示了 agent 在不同目标收益值的评估过程中累计的平均采样收益。在每一项任务中，期望的目标回报和实际观察到的回报是高度相关的。

在 Pong、HalfCheetah 和 Walker 等任务中，Decision Transformer 生成的轨迹几乎与预期回报完美匹配。此外，在Atari 的一些任务中，如 Seaquest，作者可以用比数据集中最大 episode 回报更高的回报来提示 Decision Transformer，这表明 Decision Transformer 有时能够进行推断。

4.3 Does Decision Transformer perform effective long-term credit assignment?

为了评估模型的长期信用分配能力，作者考虑了 Mesnard 等人中提出的 Key-to-Door 环境的一个变种。这是一个基于网格的环境，有三个阶段：

第一阶段，智能体被放置在有钥匙的房间里；
然后，智能体被放置在一个空房间里；
最后，智能体被放置在一个有门的房间里。

当智能体在第三阶段到达门时，它会得到一个二元奖励，但前提是它在第一阶段拿起钥匙。这个问题对于信用分配来说很困难，因为信用必须从 episode 的开始传播到最后，而忽略了中间采取的行动。

本文在应用随机动作生成的轨迹数据集上进行训练，并在下表中展示成功率。此外，在 Key-to-Door 环境中，作者使用整个 episode 长度作为上下文，而不是像其他环境中那样拥有固定的上下文窗口。

Decision Transformer 模型和 %BC（只在成功的 episode 上训练）能够学习有效的策略——产生接近最优的路径，尽管只在随机漫步上训练。TD 学习（CQL）不能在长期范围内有效地传播 q 值，得到较差的性能。

4.4 Can transformers be accurate critics in sparse reward settings?

前面已经证实了 Decision Transformer 可以产生有效的策略（actors）。现在评估 Decision Transformer 是否也可以成为有效的 critics。本文修改 Decision Transformer，使其除了输出“钥匙到门”环境中的动作 token 外，还输出收益 token。本文发现 Transformer 会根据事件不断更新奖励概率，如下图（左）所示。此外，本文发现 Transformer 关注情节中的关键事件（捡起钥匙或到达门），如下图（右）所示，表明形成了 Raposo 等人讨论的状态-奖励关联，并使准确的值预测成为可能。

4.5 Does Decision Transformer perform well in sparse reward settings?

TD 学习算法的一个已知弱点是，为了表现出色，它们需要密集的奖励，这是不现实的并且代价昂贵的。相比之下，Decision Transformer 可以在这些设置中提高鲁棒性，因为它对奖励的密度做出最小的假设。为了评估这一点，我们考虑 D4RL 基准的延迟返回版本，其中代理在轨迹上没有收到任何奖励，而是在最后的时间步中收到轨迹的累计奖励。

我们对延迟回报的结果如上表所示。延迟收益对 Decision Transformer 影响最小；由于训练过程的性质，而模仿学习方法是奖励不可知论的。虽然 TD 学习崩溃了，但 Decision Transformer 和 %BC 仍然表现良好，这表明 Decision Transformer 对延迟奖励更稳健。