当前位置：首页 > ai >正文

【论文解读】Decision Transformer：用序列建模重新定义强化学习

ai 2025/6/27 8:02:54

1st author

Lili Chen
Kevin Lu

paper:Decision Transformer: Reinforcement Learning via Sequence Modeling

code: kzl/decision-transformer: Official codebase for Decision Transformer: Reinforcement Learning via Sequence Modeling.

5. 总结 (结果先行)

Decision Transformer (DT) 标志着强化学习领域一个转变：将传统的试错与价值迭代过程，重塑为一个序列建模问题。它摒弃了显式的策略学习或价值函数拟合，转而利用 Transformer 序列模式识别能力，直接从离线数据中学习“在何种期望回报下，给定历史状态与动作，未来应该采取什么动作”。

这种方法不仅在多个离线强化学习基准测试中取得了与顶尖模型相当甚至更优的性能，更重要的是，它简化了强化学习的算法设计，并为利用大规模预训练模型解决决策问题提供了新的思路。DT 的成功表明，通过恰当的序列化和条件化，复杂的决策过程或许能以一种更直接、更“数据驱动”的方式被学习。

Figure 1: Decision Transformer 的结构

1. 思想

大问题:
- 传统强化学习 (RL)，尤其是离线 RL，面临诸多挑战：如值函数估计的准确性、引导探索的复杂性、以及在固定数据集上避免分布外动作的过高估计 (value overestimation)。
- TD-learning 等方法依赖于 Bellman 方程的迭代更新，这在长时序任务和稀疏奖励环境下可能效率低下或不稳定（例如“致命三元素”问题）。
- 如何设计一种简单、可扩展且能有效利用离线数据进行策略学习的 RL 框架？
  - 能否将 RL 中的“策略”隐式地编码为一个条件序列生成模型？
小问题:
- 如何将“目标”（即期望达成的回报）作为模型的输入条件，以指导动作的生成？
- Transformer 架构能否有效地捕捉状态、动作和回报之间的长程依赖关系，从而实现有效的信用分配？
核心思想:
1. RL 即序列建模 (RL as Sequence Modeling):
  将强化学习问题转化为一个序列预测问题。具体而言，模型学习预测在给定期望的未来回报 (return-to-go)、过去的状态序列和过去的动作序列的条件下，当前应该采取的动作。
2. 条件自回归生成 (Conditional Autoregressive Generation):
  模型以自回归的方式生成动作。通过在输入中给定一个期望的未来回报总和 (target return-to-go)，模型被“提示” (prompted) 去生成能够达到该回报水平的动作序列。
3. Transformer 作为骨干 (Transformer Backbone):
  利用 Transformer 架构（特别是 GPT 类的 Causal Transformer）来处理状态、动作和回报-累计的序列。

2. 方法

2.1 轨迹表示 (Trajectory Representation)

RL 过程中的一条轨迹 $\tau$ 被表示为一个包含回报-累计 (returns-to-go)、状态 (states) 和动作 (actions) 的序列：
$\tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \dots, \hat{R}_T, s_T, a_T)$
其中：

$\hat{R}_t = \sum_{t'=t}^T r_{t'}$ 是回报-累计 (return-to-go)，表示从当前时刻 $t$ 到轨迹结束 $T$ 的累积奖励。这是 DT 的核心思想体现，用于指导模型生成期望回报的动作。 $r_{t'}$ 是在时间步 $t^{'}$ 获得的即时奖励。
$s_t \in \mathcal{S}$ 是在时间步 $t$ 的状态。
$a_t \in \mathcal{A}$ 是在时间步 $t$ 执行的动作。

2.2 模型架构 (Model Architecture)

采用 GPT 风格的因果 Transformer (causal transformer)。在预测时间步 $t$ 的动作 $a_t$ 时，模型只能关注到时间步 $t$ 及之前的信息（即 $\hat{R}_{\le t}, s_{\le t}, a_{<t}$ ）。
输入编码:
- 将最近 $K$ 个时间步的 $(\hat{R}_i, s_i, a_i)$ 三元组（或在预测 $a_t$ 时为 $\hat{R}_t, s_t$ 和之前的 $a_{<t}$ ）作为输入。
- 每个模态 (回报-累计、状态、动作) 都有其专属的线性嵌入层 (modality-specific linear embeddings)，将原始输入投影到高维向量空间。
- 对每个时间步 $i$ (对应三个 token: $\hat{R}_i, s_i, a_i$ ) 添加一个可学习的情景时间步编码 (episodic timestep encoding)，而非标准的 Transformer 位置编码。这帮助模型区分不同时间步的信息。
- 最终的 token 嵌入是原始嵌入和时间步嵌入的总和。
Transformer 处理: 这些嵌入后的 token 序列被送入因果 Transformer 模块进行处理。
动作预测: 对于每个时间步 $t$ ，模型基于其在 $s_t$ 位置的 Transformer 输出，通过一个线性解码器 (linear decoder) 来预测当前状态 $s_t$ 下应执行的动作 $a_t$ 。

2.3 训练 (Training)

从离线数据集中采样长度为 $K$ 的轨迹片段 $(\hat{R}_i, s_i, a_i)_{i=t-K+1}^{t}$ 。
训练目标是最小化模型预测的动作 $\hat{a}_i$ 与数据集中实际动作 $a_i$ 之间的差异。
- 对于离散动作空间，使用交叉熵损失 (cross-entropy loss)。
- 对于连续动作空间，使用均方误差损失 (mean-squared error loss)。
损失在所有 $K$ 个时间步上进行平均。
注意，DT 仅预测动作，而不预测未来的状态或回报-累计。

2.4 推理/评估 (Inference/Evaluation)

在评估阶段，DT 以自回归方式与环境交互：

设定初始目标回报: 设定一个期望的初始回报-累计 $\hat{R}_{target}$ (例如，专家级回报)。
获取初始状态: 从环境中获取初始状态 $s_1$ 。
构建输入序列: 用当前的 $\hat{R}_{target}$ 、 $s_1$ (以及可能为空的初始动作和历史) 构成模型的初始输入序列。
预测动作: 模型预测动作 $a_1$ 。
与环境交互: 在环境中执行 $a_1$ ，观察到新的状态 $s_2$ 和即时奖励 $r_1$ 。
更新目标回报: 将目标回报更新为 $\hat{R}_{target} \leftarrow \hat{R}_{target} - r_1$ 。
迭代: 将新的 $(\hat{R}_{target}, s_2, a_1)$ 加入到输入序列的末尾（并可能移除最旧的元素以保持上下文长度 $K$ ），然后重复步骤 4-7，直到结束。

这种方法巧妙地将策略学习问题转化为一个条件序列建模问题，通过控制 $\hat{R}_{target}$ 来引导模型的行为。

3. 实验

在 Atari 游戏、OpenAI Gym 中的连续控制任务以及 Key-to-Door 这种需要长程信用分配的任务上评估。

对比基准 (Baselines):
- TD-learning: 主要与 Conservative Q-Learning (CQL) 进行比较，是当时最先进的 model-free 离线 RL 算法之一。
- 行为克隆 (Behavior Cloning, BC): 包括在整个数据集上进行 BC，以及一种作者提出的变体——百分位行为克隆 (%BC)，即只在数据集中回报最高的 X% 的轨迹上进行 BC。
评价指标 (Metrics):
- 主要使用归一化得分，其中 0 分代表随机策略，100 分代表专家策略（或专业玩家得分，如 Atari）。
核心实验结论:
1. 性能具有竞争力:
  - 在 Atari 任务上 (使用 1% DQN Replay 数据集)，DT 的表现在 4 个游戏中，有 3 个与 CQL 相当，并且全面优于或持平于其他基线 (REM, QR-DQN, BC)。
  - 在 OpenAI Gym 的 D4RL 数据集上 (medium, medium-replay, medium-expert)，DT 在大多数任务上取得了最高分，并在其余任务上具有竞争力。平均而言，DT 显著优于 CQL 和其他传统 RL 算法。
2. 超越简单的行为克隆子集:
  - 通过与 %BC 对比，DT 通常能达到或超过在“最佳”数据子集上训练的 %BC 的性能，而无需预先知道哪个子集最优。这表明 DT 能够有效地从整个数据分布中学习，并根据期望回报条件进行泛化。
  - 在数据量较少的情况下（如 Atari），DT 比 %BC 表现更好，说明 DT 能更好地利用所有轨迹（即使是次优轨迹）来提升泛化能力。
3. 有效的回报条件化 (Return Conditioning):
  实验表明 (Figure 4)，通过调整输入给 DT 的目标回报值，实际获得的回报与目标回报高度相关。在某些任务上 (如 Pong, HalfCheetah, Walker)，DT 能够精确匹配期望的回报。在 Seaquest 等任务中，DT 甚至能够通过提示比数据集中最高回报更高的目标值，来实现超越数据集的性能 (外推能力)。
4. 长上下文的重要性 (Benefit of Longer Context Length):
  与仅使用当前状态 (K=1) 的模型相比，使用更长历史上下文 (K=30 或 K=50) 的 DT 性能显著更好 (表 5, Atari)。这表明历史信息对于 DT 理解当前策略上下文和提升性能至关重要。
5. 有效的长程信用分配 (Long-term Credit Assignment):
  在 Key-to-Door 任务中 (一个需要先拾取钥匙再开门的稀疏奖励任务)，DT 仅基于随机游走数据进行训练，就能学习到接近最优的策略 (表 6, 成功率 71.8%-94.6%)。相比之下，CQL 在此任务上表现很差 (13.1%-13.3%)，表明 DT 在处理长程依赖和信用分配方面具有优势。
6. 稀疏奖励下的鲁棒性 (Robustness in Sparse Reward Settings):
  将 D4RL 数据集中的密集奖励改为仅在轨迹结束时给予一次总奖励（延迟奖励），DT 的性能几乎不受影响。而 CQL 在这种设置下性能急剧下降。这进一步证明了 DT 对奖励信号的密度不敏感。
7. 无需显式悲观主义或行为正则化:
  与 CQL 等需要显式机制 (如值函数惩罚或策略约束) 来处理分布外动作的算法不同，DT 通过其序列建模的本质自然地避免了这些问题，因为它直接学习在数据支持的区域内行动以达到目标回报。

4. Q&A

Q1. 什么是“致命三元素” (The Deadly Triad)？

“致命三元素”是指在强化学习中，当以下三个元素同时出现时，可能导致学习过程不稳定甚至发散的情况：

函数逼近 (Function Approximation): 使用参数化函数（如神经网络）来估计价值函数或策略，而不是使用表格记录。
自举 (Bootstrapping): 更新当前估计值时，依赖于其他学习到的估计值。例如，在时序差分学习 (TD learning) 中，Q 值的更新会用到下一个状态的估计 Q 值： $\leftarrow r + \gamma \max_{a'} Q(s', a')$ 。
离策略学习 (Off-policy Learning): 智能体学习的策略 (目标策略 $\pi$ ) 与其在环境中交互并收集数据的策略 (行为策略 $\mu$ ) 不同。这允许智能体重用旧的经验。

当这三者结合时，一个估计误差可能会通过自举被放大，并通过离策略学习传播到策略未覆盖的区域，导致价值函数估计的持续恶化和发散。

DT 在一定程度上规避了这个问题：

它不依赖于显式的价值函数自举。DT 将回报-累计 $\hat{R}_t$ 作为输入条件，直接预测动作，而不是迭代地更新价值函数。
虽然它也是离策略的（从固定的离线数据集中学习），并且使用函数逼近（Transformer 本身），但由于缺乏传统意义上的 Bellman 更新和价值自举，它对“致命三元素”的敏感性大大降低。

Q2. 由于模型训练目标是最小化模型预测的动作 $\hat{a}_i$ 与数据集中实际动作 $a_i$ 之间的差异。如果某个 $s$ 的最优动作 $a^*$ 并不存在于数据集中, 如何保证最优策略?

这是一个非常关键的问题，触及了所有基于离线数据学习的方法的根本局限性。

DT 的能力边界:
决策 Transformer 无法保证找到数据集中从未以任何形式出现过的“全新”最优动作 $a^*$ 。 它的核心是学习数据集中存在的模式，即“在状态 $s_i$ 下，如果期望获得回报 $\hat{R}_i$ ，那么历史上采取的动作是 $a_i$ ”。
如何接近最优:
1. 条件化的力量: 通过将期望回报 $\hat{R}_t$ 作为条件，DT 可以在测试时被“引导”去模仿数据集中那些导致了高回报的轨迹片段。如果数据集中包含了足够多样和高质量的轨迹（即使没有一条是全局最优的），DT 可能会通过“缝合”(stitching) 不同轨迹中的优秀片段来组合出接近最优的策略。
  例如，对于某个状态 $s$ ，数据集中可能没有直接的 $a^*, \text{高回报})$ 的样本。但可能有：
  - $a_1, \text{中回报})$
  - $a_2, \text{高回报})$ ，其中 $s^{'}$ 与 $s$ 相似，或者 $a_2$ 的某种泛化形式在 $s$ 下也是好的。
    Transformer 的泛化能力可能会让它在被提示高回报时，倾向于生成一个“类似 $a_2$ ” 的动作，或者从 $s, a_1)$ 的经验中推断出如果想要更高回报，动作应该如何调整（基于数据集中其他类似情况的模式）。
2. 泛化而非凭空创造: 它更像是在学习一个“回报-条件下的行为先验”。如果数据集质量足够高，包含了接近最优的子序列，那么当以高目标回报进行提示时，模型可以生成高质量的动作序列。但如果最优动作的“信号”在数据集中完全缺失，DT 也无能为力。
3. “最优”的定义: 在离线 RL 的背景下，“最优”通常是指在给定数据集的约束下能达到的最好性能。DT 的目标是学习一个能够匹配或超越数据集中最佳轨迹性能的策略，而不是发现理论上的、数据中未体现的全局最优策略。

总结: DT 不是在“发明”新的最优动作，而是在数据驱动的模式下，根据期望回报，智能地“选择”或“组合”出数据集中暗示的最佳行为。其性能高度依赖于离线数据集的质量和覆盖度。

3. 将 $\hat{R}_t$ 作为输入有什么直觉/理论上的意义? 与直接将 $(s, a)$ 序列作为输入有什么不同?

将回报-累计 $\hat{R}_t$ (return-to-go) 作为输入是决策 Transformer 的核心创新之一，其意义深远：

直觉意义:
1. 目标导向 (Goal-Conditioning): $\hat{R}_t$ 充当了一个明确的“目标”或“指令”。模型被告知：“我希望从现在开始获得 $\hat{R}_t$ 这么多的回报，请告诉我当前状态 $s_t$ 下应该采取什么动作 $a_t$ ”。这使得模型可以学习不同“技能水平”的行为。
2. 行为解耦 (Disentangling Behaviors): 在一个混合数据集中，同一个状态 $s_t$ 可能对应着不同的动作 $a_t$ ，这些动作后续导向了不同的累积回报。
  - 如果只输入 $(s, a)$ 序列（如标准行为克隆），模型会学习这些行为的“平均”，可能导致次优策略。它无法区分哪些 $(s, a)$ 对是“好”的。
  - 引入 $\hat{R}_t$ 后，模型可以学习到：在状态 $s_t$ 下，如果期望高回报，则采取动作 $a_t^{(high)}$ ；如果期望低回报（或者数据集中只有低回报的经验），则采取动作 $a_t^{(low)}$ 。它将行为与该行为的“后果”（未来的回报）联系起来。
3. 利用次优数据: 即使数据集中充满了次优轨迹，DT 仍然可以学习。它能学到“为了获得 X 回报，你需要执行 Y 动作序列”，即使 X 不是最大的。这比简单地丢弃次优数据更有效。
理论意义 (或与现有框架的联系):
1. 序列建模替代动态规划: 传统 RL 依赖贝尔曼方程进行动态规划或价值迭代。DT 将此过程转化为一个序列预测问题。 $\hat{R}_t$ 作为输入，避免了显式地学习和迭代价值函数 $V (s)$ 或 $Q (s, a)$ 。从某种意义上说，价值信息被直接“注入”到模型输入中。
2. 与 Upside-Down RL (UDRL) 的关联: UDRL 的思想是学习一个策略，该策略以“命令”（如期望回报和剩余时间）为条件来生成动作。DT 可以看作是 UDRL 思想在更强大的 Transformer 架构和更灵活的序列表示上的实现和扩展。
3. 信用分配: Transformer 的自注意力机制可以直接在整个序列 $(\hat{R}, s, a)$ 上操作，理论上能够更好地捕捉动作与其远期回报之间的关系，从而实现更直接的信用分配，而不像 TD学习那样需要逐步反向传播奖励信号。
与直接将 $(s, a)$ 序列作为输入的区别:
- $(s, a)$ 序列 (例如，用于行为克隆的 Transformer):
  - 模型学习 $P(a_t | s_t, s_{t-1}, a_{t-1}, \dots)$ 。
  - 它会模仿数据集中出现过的行为模式，而不考虑这些行为的好坏。如果数据集主要是次优的，它会学到次优策略。
  - 无法在测试时主动追求高回报，只能复制训练数据中的平均行为。
- $(\hat{R}, s, a)$ 序列 (Decision Transformer):
  - 模型学习 $P(a_t | \hat{R}_t, s_t, \hat{R}_{t-1}, s_{t-1}, a_{t-1}, \dots)$ 。
  - 它学习的是一个回报条件下的策略。
  - 允许在测试时通过设定高的目标 $\hat{R}_1$ 来“查询”高回报策略。
  - 能够从包含多种质量轨迹的数据集中学习，并根据回报目标区分和生成不同质量的行为。