论文精读——Decision Transformer: Reinforcement Learning via Sequence Modeling

系列文章目录

---

提示：写完文章后，目录可以自动生成，如何生成可参考右边的帮助文档

---

Abstract

提示：这里可以添加本文要记录的大概内容：

We introduce a framework that abstracts Reinforcement Learning (RL) as a sequence modeling problem. This allows us to draw upon the simplicity and scalability of the Transformer architecture, and associated advances in language modeling such as GPT-x and BERT. In particular, we present Decision Transformer, an architecture that casts the problem of RL as conditional sequence modeling. Unlike prior approaches to RL that fit value functions or compute policy gradients, Decision Transformer simply outputs the optimal actions by leveraging a causally masked Transformer. By conditioning an autoregressive model on the desired return (reward), past states, and actions, our Decision Transformer model can generate future actions that achieve the desired return. Despite its simplicity, Decision Transformer matches or exceeds the performance of state-of-the-art model-free offline RL baselines on Atari, OpenAI Gym, and Key-to-Door tasks.
翻译
我们提出了一个将强化学习（RL）抽象为序列建模问题sequence modeling problem的框架。这使我们能够利用 Transformer 架构的简洁性和可扩展性scalability，以及诸如 GPT-x 和 BERT 等在语言建模方面的相关进展。具体而言，我们提出了决策 Transformer（Decision Transformer），这是一种将强化学习问题转化为条件序列建模的架构。与以往拟合价值函数或计算策略梯度的强化学习方法不同，Decision Transformer 通过leverage因果掩码(causally masked) Transformer 直接输出最优行动。通过根据期望的回报（奖励）、过去的状态和行动对自回归模型autoregressive model进行条件设定，我们的Decision Transformer 模型可以生成实现期望回报的未来行动。尽管决策 Transformer 架构简单，但在 Atari、OpenAI Gym 和开门（Key-to-Door）等任务上，它的性能与最先进的无模型离线强化学习基线相当，甚至超越了它们。

---

提示：以下是本篇文章正文内容，下面案例可供参考

Ⅰ. INTRODUCTION

  最近的研究表明，Transformer [1] 能够对大规模的语义概念分布进行建模，包括在语言方面实现出色的零样本泛化 [2] 以及在分布外图像生成方面取得令人瞩目的成果 [3]。这与强化学习（RL）中的许多工作形成了鲜明对比，强化学习通常学习单一策略来对特定的、范围较窄的行为分布进行建模。鉴于 Transformer 模型应用的多样性及其影响力，我们试图探究其在序列决策问题中的应用。具体而言，我们不是将 Transformer 作为传统强化学习算法的一种架构选择 [4, 5]，而是研究轨迹建模（类似于语言建模）是否可以替代传统的强化学习算法。

我们考虑以下范式paradigm转变：不再通过强化学习中占主导地位的传统强化学习算法（如时间差分（TD）学习 [6]）来训练策略，而是使用序列建模目标，在收集到的经验数据on collected experience上训练 Transformer 模型。这使我们能够绕过通过自举法bootstrapping传播回报的需求，从而避免了导致强化学习不稳定的 “致命三角” [6] 之一。它还避免了像 TD 学习中通常所做的那样对未来奖励进行折扣，对未来奖励进行折扣可能会导致不良的短视行为。此外，我们可以利用在语言和视觉领域广泛使用且易于扩展的现有 Transformer 框架，借助大量关于 Transformer 模型稳定训练的研究成果；这种方法仅依靠常见的监督学习系统，就消除了对专门的强化学习框架的需求。鉴于 Transformer 已被证明具有对**长序列和广泛数据分布(model long sequences and wide data distributions)**进行建模的能力，它还有其他优势。与缓慢传播奖励且容易受到 “干扰” 信号影响(are prone to “distractor” signals)的贝尔曼备份（Bellman backups）不同，Transformer 可以通过自注意力机制直接进行功劳分配 [7]。这使得 Transformer 在稀疏或有干扰的奖励情况下仍能有效工作。此外，Transformer 建模方法可以对广泛的行为分布进行建模，从而实现更好的泛化和迁移。虽然 “倒置” 强化学习（UDRL）[8, 9, 10] 也使用基于目标回报的监督损失，但我们的工作是受序列建模驱动，而非监督学习，并且旨在从对长行为序列的建模中获益。更多相关工作的讨论见第 6 节。

我们通过考虑离线强化学习(offline RL)来探索我们的假设，在离线强化学习中，我们让智能体从次优数据中学习策略，即从固定的、有限的经验中产生最有效的行为。由于误差传播(error propagation)和价值高估(value overestimation) [11]，这个任务在传统上具有挑战性。然而，在以序列建模为目标进行训练时，这是一项自然的任务。通过在状态、行动和回报的序列上训练自回归模型，我们将**策略采样(policy sampling)**简化为自回归生成建模。我们可以通过手动设置回报token来指定策略的专业性，即查询哪种 “技能”，这些tokens可作为生成的提示。

离线强化学习（offline reinforcement learning）：强化学习的一种设置，智能体仅使用预先收集的固定数据集进行学习，而不与环境进行实时交互。比如，利用之前记录的机器人在各种场景下的行动和奖励数据来训练机器人的行为策略，而不需要机器人在训练时再次与实际环境互动。
自回归生成建模（autoregressive generative modeling）：基于自回归模型的生成式建模方法，根据序列的历史数据生成新的数据。在强化学习中，根据之前的状态、行动和回报生成下一个行动，不断迭代生成完整的行动序列。

示例说明。为了直观理解我们的提议，考虑将在有向图上寻找最短路径的任务视为强化学习问题。到达目标节点时reward为 0，否则为 -1。我们训练一个 GPT [12] 模型，以预测returns-to-go（sum of future rewards）、states和actions的下一个token。仅在随机游走数据上进行训练，且没有专家演示，在测试时，我们可以通过添加一个先验来生成尽可能高的回报（附录中有更多详细信息和实证结果），从而生成最优轨迹，然后在此基础上生成行动。因此，通过将序列建模工具与事后回报信息相结合，我们无需动态规划就能实现策略改进。

受这一观察结果的启发，我们提出了Decision Transformer，我们使用 GPT 架构对轨迹进行自回归建模（如图 1 所示）。我们通过在 Atari [13]、OpenAI Gym [14] 和开门（Key-to-Door）[15] 环境的离线强化学习基准测试中评估Decision Transformer，来研究序列建模是否能够进行策略优化。我们发现，在不使用动态规划的情况下，Decision Transformer 在这些基准测试中的表现与最先进的无模型离线强化学习算法 [16, 17] 相当。此外，在需要长期功劳分配的任务中，决策 Transformer 的性能明显优于传统强化学习算法。通过这项工作，我们希望将 Transformer 模型的最新重大进展与强化学习问题联系起来。

2. Preliminaries

2.1 Offline reinforcement learning

我们考虑在由元组$(S,A,P,R)$描述的马尔可夫决策过程（MDP）中进行学习。这个MDP元组包含状态$s\in S$、行动$a\in A$、转移动力学$P(s^{\prime }|s,a)$以及奖励函数$r=R(s,a)$。我们分别用$s_t$、$a_t$和$r_t=R(s_t,a_t)$来表示在时间步(t)时的状态、行动和奖励。强化学习的目标是学习一种策略，在MDP中最大化期望回报$\mathbb{E}[{\sum }_{t=1}^T{r}_t]$。在离线强化学习中，我们不是通过与环境交互来获取数据，而是只能访问一些固定的有限数据集，这些数据集由来自环境的轨迹组成。这种设置更具挑战性，因为它剥夺了智能体探索环境和收集额外反馈的能力。

2.2 Transformers

Transformer由Vaswani等人 [1] 提出，是一种用于高效对序列进行建模的架构。它由堆叠的自注意力层和残差连接组成。每个自注意力层接收$n$个对应于唯一输入令牌的嵌入 { x i } i = 1 n \{x_{i}\}_{i = 1}^{n} {xi​}i=1n​，并输出$n$个嵌入 { z i } i = 1 n \{z_{i}\}_{i = 1}^{n} {zi​}i=1n​，保持输入维度不变。第(i)个令牌通过线性变换映射为一个键$k_i$、一个查询$q_i$和一个值$v_i$。自注意力层的第$i$个输出是通过用查询$q_i$与其他键$k_j$之间的归一化点积对值$v_j$进行加权得到的，公式为：
z i = ∑ j = 1 n softmax ( { < q i , k j ′ > } j ′ = 1 n ) j ⋅ v j z_{i}=\sum_{j = 1}^{n} \text{softmax}\left(\left\{\left< q_{i}, k_{j'}\right>\right\}_{j' = 1}^{n}\right)_{j} \cdot v_{j} zi​=j=1∑n​softmax({⟨qi​,kj′​⟩}j′=1n​)j​⋅vj​
这使得该层能够通过查询和键向量的相似性（最大化点积）隐式地形成状态 - 回报关联，从而进行 “功劳分配”。在这项工作中，我们使用GPT架构 [12]，它通过因果自注意力掩码对Transformer架构进行了修改，以实现自回归生成，将对$n$个令牌的求和/softmax操作改为仅对序列中前面的令牌$(j\in [1,i])$进行操作。其他架构细节请参考原始论文。

• 因果自注意力掩码（causal self-attention mask）：在自注意力机制的基础上，用于限制模型在生成过程中只能关注当前位置及之前的位置信息，而不能看到未来的信息。在文本生成任务中，当生成第(i)个单词时，模型只能根据已经生成的前(i - 1)个单词的信息来进行预测，因果自注意力掩码就起到了这样的限制作用。在GPT架构中，利用因果自注意力掩码实现了自回归生成，符合人类语言生成的顺序性特点。
• 自回归生成（autoregressive generation）：一种生成数据的方式，模型根据已生成的部分数据来预测下一个数据。在语言模型中，就是根据已经生成的前面的单词，不断预测下一个单词，逐步生成完整的文本。例如，在生成一个句子时，模型首先根据初始输入生成第一个单词，然后结合第一个单词生成第二个单词，依此类推，直到生成完整的句子。

Ⅲ. Method

在本节中，我们将介绍Decision Transformer，它对Transformer架构进行了最小限度的修改，以自回归方式对轨迹进行建模，总结如图1和算法1所示。

3.1Trajectory representation

• 英文原文：The key desiderata in our choice of trajectory representation are (a) it should enable transformers to learn meaningful patterns and (b) we should be able to conditionally generate actions at test time. It is nontrivial to model rewards since we would like the model to generate actions based on future desired returns, rather than past rewards. As a result, instead of modeling the rewards directly, we model the returns-to-go ${\hat{R}}_t={\sum }_{t^{\prime }=t}^T{r}_{t^{\prime }}$. This leads to the following trajectory representation which is amenable to autoregressive training and generation:
  [\tau = (\hat{R}{1}, s{1}, a_{1}, \hat{R}{2}, s{2}, a_{2}, \ldots, \hat{R}{T}, s{T}, a_{T}).]
• 翻译：我们选择轨迹表示trajectory representation的关键要求是：（a）它应使Transformer能够学习有意义的模式；（b）我们应能够在测试时根据条件生成行动。对奖励进行建模并非易事，因为我们希望模型基于未来期望回报生成行动，而非过去的奖励。因此，我们不直接对奖励进行建模，而是对未来回报returns-to-go ${\hat{R}}_t={\sum }_{t^{\prime }=t}^T{r}_{t^{\prime }}$ 进行建模注意是仅仅对当前时间步过后的回报进行建模。这就产生了以下适用于自回归训练和生成的轨迹表示：
  $$\tau =({\hat{R}}_1,s_1,a_1,{\hat{R}}_2,s_2,a_2,\dots ,{\hat{R}}_T,s_T,a_T).$$

• 未来回报（returns-to-go）：从当前时间步 $t$ 到整个决策过程结束时（时间步 $T$ ）的所有未来奖励的总和，用公式 ${\hat{R}}_t={\sum }_{t^{\prime }=t}^T{r}_{t^{\prime }}$ 表示。它反映了从当前状态出发，采取一系列行动后可能获得的总收益，是决策Transformer建模的关键变量。模型根据未来回报来决定当前应该采取的行动，以实现期望的累积收益。

3.2 Architecture

我们将最后 $K$ 个时间步的信息输入到Decision Transformer中，总共形成 $3K$ 个token（每个模态：return-to-go, state, or action）。为了获得token embeddings，我们为每个模态学习一个线性层，将原始输入投影到嵌入维度，然后进行层归一化 [18]居然还要把状态进行投影，没看懂。哦是为了将离散的状态/回报/行动投影到连续的空间中，至少是变成连续的序列。对于有视觉输入的环境，状态会被输入到卷积编码器中，而非线性层。此外，为每个时间步学习一个Position Embedding，并将其添加到每个token中 —— 注意，这与Transformer使用的标准位置嵌入不同，因为一个时间步对应三个token。然后，这些token由GPT [12] 模型进行处理，该模型通过自回归建模预测未来的action token。

3.3 Training

因为是offline RL，所以其实采样的各种轨迹都已经是预先准备好的经验了。果然内核还是机器学习
我们从数据集中采样长度为 $K$ 的序列小批量数据minibatches of sequence。对应于输入令牌 $s_t$ 的预测头被训练来预测 $a_t$，对于离散行动使用交叉熵损失，对于连续行动使用均方误差损失，并且对每个时间步的损失进行平均更多的不应该是平均，而应该对某些奖励相差交大的时间步增加损失。我们发现预测状态或未来回报对于良好的性能并非必要，尽管这是可行的（如5.3节所示），并且可能是未来工作中一个有趣的研究方向。仅仅根据状态输入来预测应该产生的行为a，而没有要求预测r'与s'

Evaluation

• 英文原文：During evaluation rollouts, we specify a target return based on our desired performance (e.g., specify maximum possible return to generate expert behavior) as well as the environment starting state, to initialize generation. After executing the generated action, we decrement the target return by the achieved reward and obtain the next state. We repeat this process of generating actions and applying them to obtain the next return-to-go and state until episode termination.
• 翻译：在评估过程中，我们根据期望的性能指定一个目标回报（例如，指定最大可能的回报以生成专家行为）以及环境的起始状态，来初始化生成过程。执行生成的行动后，我们将目标回报减去获得的奖励，得到下一个状态。我们重复这个生成行动并应用行动以获得下一个未来回报和状态的过程，直到episode结束。这个episode指的是RL中的每个episode都从初始状态开始，然后到结束的一个episode吗
• 核心专业术语解释：
  • 评估（Evaluation）：在训练完成后，对决策Transformer模型性能进行测试和分析的过程。通过在不同的环境状态下运行模型，观察模型的决策效果和获得的回报，来判断模型是否学习到了有效的策略，以及模型在不同场景下的泛化能力。
  • 目标回报（Target return）：在评估过程中预先设定的期望获得的回报值。它用于引导模型生成相应的行动策略，例如，如果希望模型表现出专家级别的行为，可以设置最大可能的目标回报，模型会尝试生成一系列行动以接近或达到这个目标回报。
  • 情节（Episode）：在强化学习中，一个情节指的是智能体从初始状态开始，与环境进行一系列交互，直到达到终止状态的整个过程。在评估决策Transformer时，每个情节都从指定的起始状态开始，模型根据环境反馈不断生成行动，直到满足终止条件（如达到目标、超出时间限制等），一个情节结束。

Algorithm1 Decision Transformer

# R, s, a, t：未来回报、状态、行动、时间步
# K：上下文长度（输入到Decision Transformer的每个序列的长度）
# transformer：带有因果掩码的Transformer（GPT）
# embed_s, embed_a, embed_R：线性嵌入层
# embed_t：学习得到的情节位置嵌入
# pred_a：线性行动预测层

# 主模型
def DecisionTransformer(R, s, a, t):
    # 计算令牌的嵌入
    pos_embedding = embed_t(t)  # 每个时间步（注意不是每个令牌）
    s_embedding = embed_s(s) + pos_embedding
    a_embedding = embed_a(a) + pos_embedding
    R_embedding = embed_R(R) + pos_embedding

    # 交叉排列令牌，如 (R_1, s_1, a_1,..., R_K, s_K)
    input_embeds = stack(R_embedding, s_embedding, a_embedding)

    # 使用Transformer获取隐藏状态
    hidden_states = transformer(input_embeds=input_embeds)

    # 选择用于行动预测令牌的隐藏状态
    a_hidden = unstack(hidden_states).actions

    # 预测行动
    return pred_a(a_hidden)

# 训练循环
for (R, s, a, t) in dataloader:  # 维度：(批量大小, K, 维度)
    a_preds = DecisionTransformer(R, s, a, t)
    loss = mean((a - a_preds)**2)  # 连续行动的L2损失
    optimizer.zero_grad(); loss.backward(); optimizer.step() 

# 评估循环
target_return = 1  # 例如，专家级回报
R, s, t, done = [target_return], [env.reset()], [1], False
while not done:  # 自回归生成/采样
    # 采样下一个行动
    action = DecisionTransformer(R, s, a, t)[-1]  # 对于连续行动
    new_s, r, done, _ = env.step(action)

    # 将新令牌添加到序列中
    R = R + [R[-1] - r]  # 用奖励递减未来回报
    s, a, t = s + [new_s], a + [action], t + [len(R)]
    R, s, a, t = R[-K:],...  # 仅保留长度为K的上下文

解释

1. 整体功能：这段伪代码展示了决策Transformer在连续行动场景下的工作流程，包括模型定义、训练和评估三个主要部分。
2. 模型定义（DecisionTransformer函数）
   • 嵌入计算：首先通过不同的线性嵌入层（embed_s、embed_a、embed_R ）分别对状态 $s$、行动 $a$、未来回报 $R$ 进行嵌入，并且都加上针对每个时间步 $t$ 学习到的情节位置嵌入（embed_t ）。这样做是为了让模型能融合不同模态信息并感知时间步信息。
   • 令牌排列与输入：将处理后的三种模态的嵌入进行交叉排列组合（stack操作），形成输入给Transformer的嵌入表示（input_embeds ）。
   • 隐藏状态获取与行动预测：把组合后的嵌入输入到带有因果掩码的Transformer（这里是GPT架构）中，得到隐藏状态（hidden_states ）。然后从隐藏状态中提取出用于行动预测的部分（a_hidden ），最后通过线性行动预测层（pred_a ）输出预测的行动。
3. 训练循环：从数据加载器（dataloader ）中按批次取出未来回报 $R$、状态 $s$、行动 $a$ 和时间步 $t$ 数据。将这些数据输入到决策Transformer模型中得到预测行动（a_preds ），使用均方误差（L2损失）衡量预测行动和真实行动 $a$ 的差异。通过反向传播（loss.backward() ）计算梯度，并利用优化器（optimizer ）更新模型参数（optimizer.step() ）。
4. 评估循环：设定一个目标回报（target_return ），例如想要达到专家级别的回报水平。初始化未来回报 $R$、状态 $s$、时间步 $t$ 和结束标志 done 。在循环中，不断通过决策Transformer模型生成行动（action ），并在环境中执行（env.step(action) ），根据环境反馈更新未来回报（考虑减去获得的奖励）、状态等信息。同时，为了保持输入到模型的上下文长度为 $K$，对相关数据进行截断处理（R[-K:] 等操作） ，直到满足结束条件（done 为True）。