基于关键状态的扩散模型轨迹规划方法

Key-State-Conditioned Diffusion Models for Trajectory Planning

摘要: 在离线强化学习的轨迹规划任务中，传统基于自回归的规划方法因误差逐级累积效应而限制了模型性能。近年来，扩散模型凭借其出色的分布建模能力被引入该领域，以缓解误差累积问题。然而，现有方法在高维动作空间生成长时序轨迹时仍面临性能不足的挑战。为此，本文提出了一种基于关键状态的扩散模型轨迹规划方法，通过提取原始轨迹中的关键状态特征，并结合条件扩散生成模型进行轨迹规划，将传统的自回归式轨迹规划范式转化为基于关键状态的条件生成问题。在确保生成轨迹时序连续性的同时，提升了模型轨迹规划的性能。在D4RL基准测试的Gym-Mujoco、Maze2d、AntMaze和Adroit等多个环境中进行的实验表明，本文方法在轨迹规划性能和算法鲁棒性方面均优于现有方法。

Abstract: In trajectory planning for offline reinforcement learning, conventional autoregressive planning methods suffer from performance limitations due to error accumulation effects. While diffusion models have recently been introduced to this domain to mitigate error accumulation through their exceptional distribution modeling capabilities, existing approaches still face performance challenges when generating long-horizon trajectories in high-dimensional action spaces. To address this, we propose a Key-State-Conditioned Diffusion Models for Trajectory Planning method that integrates key states with diffusion models. Our approach extracts critical state features from original trajectories and combines them with conditional diffusion generative models for trajectory planning, effectively transforming the traditional autoregressive planning paradigm into a key state-conditioned generation problem. This method not only maintains temporal continuity in generated trajectories but also significantly enhances planning performance. Extensive experiments conducted on multiple D4RL benchmark environments, including Gym-Mujoco, Maze2d, AntMaze, and Adroit, demonstrate that our method outperforms existing approaches in both trajectory planning performance and algorithmic robustness.

1. 引言

强化学习(Reinforcement Learning, RL)在机器人控制[1] [2]、自动驾驶[3]等序列决策任务中展现出巨大潜力。然而，传统在线强化学习面临安全风险高、试错成本大等固有缺陷。在此背景下，离线强化学习(Offline Reinforcement Learning)通过利用静态数据集进行策略优化，有效规避了实时交互需求，逐渐发展成为该领域的主流范式[4]。学界已提出多种创新方法：Fujimoto等人提出批次约束Q学习(Batch-Constrained Q-Learning, BCQ) [5]算法，通过行为克隆与Q学习的协同机制显式约束策略空间；Kumar等人提出保守Q学习算法(Conservative Q learning, CQL) [6]，创新性地在价值函数中引入正则化项以逼近真实Q值的下界；Kidambi等人提出的基于模型的离线强化学习算法(Model-based Offline Reinforcement Learning, MOReL) [7]则通过构建环境模型来补偿离线交互的缺失。

然而，传统离线强化学习方法在复杂轨迹规划任务中仍面临诸多挑战：自回归式规划方法受限于单步预测误差在时序决策中的累积效应，其复合误差会导致生成轨迹逐渐偏离真实环境动态[8]。此外，基于单步Q值估计的方法难以有效应对稀疏奖励场景下的信用分配问题(Credit Assignment Problem, CAP)，同时受限于离线数据集的次优轨迹分布，传统方法易受分布偏移的影响[9]。

生成式模型的技术突破为离线强化学习开辟了新路径。研究者通过将序列建模与传统规划方法相结合，突破了传统策略迭代的框架，将序贯决策转化为直接的序列建模问题，展现出独特的优势。轨迹Transformer算法(Trajectory Transformer, TT) [10]率先引入Transformer架构[11]，利用其自注意力机制有效捕捉状态–动作序列的长期依赖关系。决策Transformer (Decision Transformer, DT) [12]进一步将强化学习重构为序列预测任务，通过嵌入网络将历史轨迹编码为连续表征。

近年来，扩散模型(Diffusion Model)在图像生成领域获得显著成就[13]，也为轨迹规划提供了新工具[14]。Janner等人提出了基于扩散模型结构的条件生成建模算法Diffuser [15]，通过反向扩散过程生成满足约束条件的轨迹序列，开创了基于扩散模型的规划范式。相较于传统方法，这类方法通过联合建模完整轨迹的概率分布，有效缓解了自回归方法的累积偏差和分布偏移问题。扩散模型在建模高维连续动作空间时展现出的多模态捕捉能力，使其成为复杂轨迹生成的有效工具。

然而，现有基于扩散模型的轨迹规划方法仍存在关键瓶颈：

1) 时序特征学习的不足：现有方法缺乏对原始轨迹时序特征的有效建模，导致难以平衡轨迹规划长度与质量。短视距规划易受稀疏奖励和局部不确定性的影响，造成轨迹规划出现偏差；而长视距规划则要求模型具备更强的跨步依赖建模能力。

2) 轨迹生成的局限性：部分方法因缺乏价值引导，难以生成高价值轨迹，且环境约束不充分会导致生成的轨迹可能违反物理限制。

3) 生成式模型噪声问题：无分类器引导机制可能导致模型倾向于生成低概率区域轨迹，这些偏离主分布的样本会显著影响策略的可靠性。

针对上述挑战，本研究提出了一种基于关键状态的扩散模型轨迹规划方法(Key-State-Conditioned Diffusion Models for Trajectory Planning, KSDP)，本文的主要贡献如下：

• 构建了一个融合Transformer时序建模与β-VAE特征压缩的混合架构，从原始轨迹中提取包含状态转移规律和语义上下文的关键特征。

• 构建了基于关键状态条件的扩散模型，通过价值函数引导生成高回报状态序列，并耦合逆动力学模型确保动作轨迹的物理可行性。

• 提出基于核密度估计的自适应筛选策略，通过动态阈值调整有效抑制生成模型的分布外噪声。

2. 基于关键状态的扩散模型轨迹生成

2.1. 关键状态提取

在本研究方法中，关键状态被定义为能够有效表征原始轨迹中跨T个时间步长状态序列的显著性特征。这些关键状态不仅需要捕捉轨迹的核心动态信息，还需编码其上下文时序关联特征，从而为后续扩散模型提供结构化条件输入，以生成高保真度的状态轨迹。在连续控制任务中，轨迹数据的复杂性和时序依赖性对特征提取提出了较高要求。关键状态的提取旨在减少冗余信息的影响，突出对轨迹规划具有决定性意义的状态特征。近年来，Transformer架构因其在自然语言处理领域的显著成功而备受关注，其核心在于通过自注意力机制(Self-Attention)有效建模序列元素的长距离依赖关系，这一特性使其在处理长时序数据时表现出色，尤其适用于捕捉轨迹数据中的长期动态模式。因此，我们首先使用Transformer模型初步提取关键状态的上下文时序特征。

数据集中的轨迹样本的形式为$\tau =\left(s_1,a_1,r_1,s_2,a_2,r_2,\cdots ,s_T\right)$ ，其中$s\in {ℝ}^{d_s}$ 表示状态，$a\in {ℝ}^{d_s}$ 表示动作，$r\in {ℝ}^{d_s}$ 表示奖励。首先，我们需要定义模型的输入形式及其编码目标。轨迹τ是一个长度为T的序列，轨迹的每个时间步由状态、动作和奖励组成，每个时间步的输入定义为$x_t=\left[s_t,a_t,r_t\right]$ 。然而动作空间、状态空间和奖励空间的维度大小往往存在差异。为了统一模型的输入维度，我们通过独立的全连接层将各模态映射到统一的隐空间$d_h$ ，再融合为统一的时间步输入：

$e_s^{\left(t\right)}=W_s\cdot s_t+b_s,e_a^{\left(t\right)}=W_a\cdot a_t+b_a,e_s^{\left(t\right)}=W_r\cdot r_t+b_r$ (1)

其中$W_s\in {ℝ}^{d_h\times d_s}$ ，$W_a\in {ℝ}^{d_h\times d_a}$ ，$W_r\in {ℝ}^{d_h\times d_r}$ 为嵌入权重矩阵，$b$ 为偏置向量。通过上述嵌入层，状态、动作和奖励都被映射到相同的维度$d_h$ ，输入的嵌入层如图1所示。原始数据轨迹经过嵌入后的序列可以表示为$E\left[e_1,e_2,\cdots ,e_T\right]\in {ℝ}^{T\times d_h}$ ，$e_t$ 表征了时间步t的完整信息。

Transformer模型摒弃了传统循环神经网络(RNN)的递归结构，完全依赖自注意力机制处理序列数据。然而，输入元素的顺序变化不会影响注意力权重的计算结果，导致模型无法区分序列中元素的时序关系，所以必须显式引入位置信息以捕捉轨迹的马尔可夫性及长期依赖关系。

Figure 1. Embedding layer schematic diagram

图1. 嵌入层示意图

传统的绝对位置编码(Absolute Positional Encoding)通过为每个时间步添加固定的正弦–余弦函数或可学习向量来表示位置。然而，这种方法在长序列中可能难以充分捕捉动态的时序关系，尤其当状态、动作和奖励之间的依赖随相对位置变化时。为此，我们采用Shaw等人提出的相对位置编码(Relative Positional Encoding, RPE) [16]。与传统的绝对位置编码不同，该方法通过建模序列中任意两位置i和j的相对距离j − i，增强模型对长序列结构的感知能力。为控制计算复杂度，设定最大相对距离$k_{\max }$ ，超出范围的距离被截断：

$\text{clipped\_distance}=\text{clip}\left(j-i,-k_{\max },k_{\max }\right)$ (2)

其中，$k_{\max }$ 表示考虑的最大相对距离，超出此范围的距离被截断，相对距离被限制在$\left[-k_{\max },k_{\max }\right]$ 范围内，减少了模型参数量，并假设远距离依赖的影响较小。在标准的自注意力机制中，注意力权重由查询(Query)和键(Key)的点积计算：

$\text{Attention}\left(Q,K,V\right)=\mathrm{softmax}\left(\frac{Q{K}^{\text{T}}}{\sqrt{d_h}}\right)V$ (3)

其中$Q=X{W}_Q$ ，$K=X{W}_K$ ，$V=X{W}_V$ 分别是查询、键和值矩阵，$W_Q$ $W_K$ $W_V$ 是可学习的权重。这种计算仅依赖于Q和K的内容信息，忽略了时间步之间的相对位置关系。相对位置编码通过在注意力分数中引入位置偏置项，显式建模i和j之间的相对距离。具体而言，在计算$Q{K}^{\text{T}}$ 时引入一个相对位置表示矩阵RPE，从而在注意力分数中显式建模位置关系：

$\text{Attention}\left(Q,K,V\right)=\mathrm{softmax}\left(\frac{Q{K}^{\text{T}}+RPE}{\sqrt{d_h}}\right)V$ (4)

其中，$RPE\in {ℝ}^{T\times T}$ 是一个相对位置偏置矩阵，矩阵中的元素$RP{E}_{i,j}$ 表示位置i和j之间的相对位置信息，通过神经网络学习生成。此机制赋予模型动态选择关键时间步的能力，而非简单地对所有时间步进行均等处理。嵌入序列通过多头自注意力层处理，生成全局特征表示M。多头机制通过不同子空间并行计算注意力，提升特征表达能力：

${\text{head}}_m=\text{Attention}\left(Q{W}_m^Q,K{W}_m^K,V{W}_m^V\right)$ (5)

$M=\text{Concat}\left({\text{head}}_1{\text{,head}}_2\text{,}\cdots {\text{,head}}_m\right)W_O$ (6)

其中，$W_i^Q$ ，$W_i^K$ ，$W_i^V$ 为查询、键和值的投影矩阵，m为注意力头数，$W_O\in {ℝ}^{d_h\times d_h}$ 是输出的线性变换矩阵。在多头自注意力机制完成序列的全局特征交互后，输出M通过残差连接(Residual Connection)与层归一化(Layer Normalization)接入由全连接层组成前馈神经网络，通过ReLU激活函数引入非线性表达能力。自注意力层和前馈神经网络层组成了一个基本的上下文特征提取模块，通过重复堆叠构成了完整的编码器结构，如图2所示。

Figure 2. Transformer encoder architecture diagram

图2. Transformer编码器结构图

Transformer编码器输出的$M\in {ℝ}^{T\times d_h}$ 蕴含丰富的上下文信息，但由于其保留了输入序列的维度，仍可能包含冗余或不关键的信息。为进一步提炼关键状态特征，我们将M输入至β-变分自编码器(β-Variational Autoencoder, β-VAE) [17]模型中，以生成隐变量$z~q\left(z|M\right)$ 。其目的是压缩轨迹特征，同时保留足够的信息用于后续生成任务。解码器部分被建模为逆动力模型，解码器的输入为相邻状态和关键状态特征z，输出为模型预测产生这个状态转移的动作a，即$q\left(a_t|s_t,s_{t+1},z\right)$ 。逆动力模型的设计不仅作为解码器辅助编码器训练，还作为轨迹规划中的动作轨迹生成器，这部分将在下一小节详细讲解。β-VAE模型的训练目标结合了重构损失和KL散度，确保重构精度和隐变量分布的正则化：

$L_{\beta }=E_{a_t~D}\left[{\displaystyle \sum _{t=1}^T{\Vert a_t-D_{\psi }\left(a_t|s_t,s_{t+1},z\right)\Vert }_2^2-\beta \cdot D_{KL}\left(E_{\varphi }\left(z|M\right)\parallel p\left(z\right)\right)}\right]$ (7)

其中第一项为解码器预测动作和原始轨迹中动作的重构损失，表示逆动力模型在给定潜在变量z和状态序列s的情况下重构动作序列a的能力。第二项为KL (Kullback-Leibler)散度项，约束了关键状态特征分布$E_{\varphi }\left(z|M\right)$ 与先验分布$p\left(z\right)$ 之间的接近程度。$\beta \in \left[0,1\right]$ 为超参数，用于平滑重构损失和KL散度之间的权重，当β = 1时，β-VAE退化为传统的VAE。通过调整β的值，β-VAE模型能够在重构质量和拟合潜在变量分布之间取得更好的平衡。通过Transformer和β-VAE编码器，我们从原始轨迹$\tau$ 中提取到了关键状态的紧凑表示z，其捕获了轨迹的高级语义信息，如策略模式与环境动态等，且有效滤除了噪声，为生成任务提供紧凑且信息丰富的关键状态表示。

2.2. 扩散轨迹生成

现有方法大多直接利用扩散模型生成状态–动作轨迹，这种方式虽然直观，但在强化学习场景中存在局限性。在强化学习的任务环境中，状态通常具有连续性，而动作则呈现出更大的多样性，且本质上往往是离散的。在机器人控制场景中，以关节力矩表示的动作序列通常表现出较高的频率和平滑性，这显著增加了预测和建模的难度[18]。为应对这一问题，我们首先使用扩散模型生成状态序列，再使用上节提到的逆动力模型(Inverse-Dynamics) [19]，根据状态序列推断出动作序列，如图3所示。

Figure 3. Training flowchart of KSDP algorithm

图3. KSDP算法的训练流程图

具体而言，扩散模型$G_{\theta }$ 用于生成T个时间步长的状态轨迹序列，其中k表示去噪过程中的时间步长：

$x^k\left(\tau \right)={\left(s_t,s_{t+1},\cdots ,s_{t+T-1}\right)}^k$ (8)

在生成状态轨迹后使用逆动力学模型预测动作：

$a_t=D_{\psi }\left(a_t|s_t,s_{t+1},z\right)$ (9)

我们将强化学习中的轨迹规划问题定义为条件扩散模型的序列生成问题，其中上标表示扩散过程的时间步，下标表示强化学习中的时间步：

${\max }_{\theta }{E}_{\tau ~D}\left[\log p_{\theta }\left(x^0\left(\tau \right)|y\left(\tau \right)\right)\right]$ (10)

我们使用关键状态z作为轨迹生成模型的条件，通过条件信息$y\left(\tau \right)$ 生成轨迹，将扩散过程限制在轨迹状态的上下文信息范围内。扩散模型的前向过程是一个马尔可夫过程，通过逐步向原始轨迹数据$x^0$ 添加高斯噪声，使其分布逐渐趋向各向同性的高斯分布。该过程定义为马尔可夫链：

$q\left(x^{1:K}|x^0\right)={\displaystyle \prod _{k=1}^{K}q\left(x^k|x^{k-1}\right)}$ (11)

$q\left(x^k|x^{k-1}\right)=N\left(x^k;\sqrt{1-{\beta }^k}{x}^{k-1},{\beta }^{k}I\right)$ (12)

其中，${\beta }^k$ 是时间步k的噪声调度参数，控制每一步添加的噪声强度。通过重参数化技巧，可直接从原始状态轨迹${\tau }^0$ 采样任意时刻K的噪声轨迹，前向过程的闭式解为：

$q\left(x^K|x^0\right)=N\left(x^k;\sqrt{{\overline{\alpha }}^k}{x}^0,\left(1-{\overline{\alpha }}^k\right)I\right)$ (13)

其中${\alpha }^k=1-{\beta }^k$ ，方差调度采用余弦退火策略保证训练稳定性。原始状态轨迹数据在$k\to K$ 时逐渐趋向于标准高斯分布$x^K~N\left(0,I\right)$ 。反向过程的目标是从纯噪声$x^K~N\left(0,I\right)$ 开始，逐步去噪以恢复原始轨迹$x^0$ 。该过程同样被建模为马尔可夫过程，其条件分布为：

$p_{\theta }\left(x^{0:K}|z\right)=p\left(x^K\right){\displaystyle \prod _{k=1}^K{p}_{\theta }\left(x^{k-1}|x^k,z\right)}$ (14)

$p_{\theta }\left(x^{k-1}|x^k,z\right)=N\left(x^{k-1};{\mu }_{\theta }\left(x^k,k,z\right),{\Sigma }^{k}I\right)$ (15)

其中，${\mu }_{\theta }$ 是参数化的均值函数，由神经网络预测。${\Sigma }^k$ 通常设为固定的协方差矩阵${\beta }^k$ 以简化计算。反向去噪的均值可以通过噪声预测网络建模：

${\mu }_{\theta }\left(x^k,k,z\right)=\frac{1}{\sqrt{{\alpha }^k}}\left(x^k-\frac{{\beta }^k}{\sqrt{1-{\overline{\alpha }}^k}}{\epsilon }_{\theta }\left(x^k,k,z\right)\right)$ (16)

为平衡生成轨迹的多样性与条件约束，我们采用无分类器扩散(Classifier-free Diffusion) [20]框架。在每轮训练中，我们从均匀分布$k~\mathcal{U}\left(1,K\right)$ 中随机采样时间步k，根据前向过程$q\left(x^K|x^0\right)$ 将采样的原始状态轨迹$x^0$ 转换为加噪样本$x^K$ 。随后，神经网络${\epsilon }_{\theta }$ 以$x^K$ 、k和z为输入，预测前向过程中添加的噪声$\epsilon$ 。同时引入条件丢弃机制：在每次前向传播中，以概率β将条件z替换为无效条件$\varnothing$ ，从而使模型同时学习条件和无条件分布。模型的损失函数基于DDPM [21]简化后的扩散模型目标，即预测噪声的均方误差。给定真实噪声$\epsilon ~N\left(0,I\right)$ ，损失函数定义为：

$L\left(\theta \right)=E_{k,\tau \in D}\left[{\Vert \epsilon -{\epsilon }_{\theta }\left(x^k\left(\tau \right),\left(1-\beta \right)z+\beta \varnothing ,k\right)\Vert }^2\right]$ (17)

通过最小化该损失，模型能够在生成符合条件z的轨迹的同时，在无条件情况下保持生成结果的多样性。通过调整指导强度$\beta$ ，可以灵活控制条件对生成结果的影响，允许在多样性(无条件生成)和条件一致性(有条件生成)之间进行权衡。然而，这种方式无法显式优化任务目标，无法直接引导模型生成更高回报的轨迹。

为解决这一问题，我们额外训练了一个轨迹价值预测器$V_{\varsigma }$ ，通过关键状态特征z来预测扩散生成轨迹的累计回报，以此作为轨迹生成中的价值梯度引导。特别地，我们在累计回报预测中隐式引入折扣因子，确保其与价值函数的定义一致，从而提升引导精度。在训练时使用离线数据中长度为T的轨迹的累计回报作为监督信号：

$V\left(\tau \right)={\displaystyle \sum _{t=1}^T{\gamma }^{t-1}R\left(s_t,a_t\right)}$ (18)

轨迹价值预测器$V_{\varsigma }$ 训练目标为最小化预测值与真实折扣回报之间的重构损失：

$L_{\varsigma }=E_{\tau ~D}\left[{\Vert V_{\varsigma }\left(z\right)-V\left(\tau \right)\Vert }_2^2\right]$ (19)

其中D表示离线数据集，z是编码器提取的关键状态特征。最终采样过程，结合无分类器方法和轨迹价值梯度引导方式实现：

$\widehat{\epsilon }=\omega {\epsilon }_{\theta }\left(x^k,z,k\right)+\left(1-\omega \right){\epsilon }_{\theta }\left(x^k,\varnothing ,k\right)-\sqrt{1-{\overline{\alpha }}^k}\nabla V_{\varsigma }\left(z\right)$ (20)

其中$\omega$ 表示无分类器方法在条件扩散模型中的引导权重。设置$\omega =1$ 时，模型等效于标准条件生成模型；设置$\omega >1$ 时，增强条件信息的影响，生成结果更贴近条件z。${\overline{\alpha }}^k$ 表示累积的噪声方差，随时间步增加而减小，则$\sqrt{1-{\overline{\alpha }}^k}$ 随时间步增大而增大。价值梯度项的引导强度自动适应扩散过程中的时间步k，无需手动设置超参数，从而在扩散后期强化价值引导效果。在扩散后期，带噪轨迹${\tau }^k$ 已接近真实分布，此时价值梯度引导更可靠，能够更精确地优化任务目标，避免早期过强引导引入的偏差。通过在扩散生成过程中融合无分类器引导和轨迹价值梯度引导，模型既能生成符合关键状态约束的轨迹，又能产出高价值的样本。

2.3. 轨迹规划

在模型训练结束后的应用阶段，KSDP算法的轨迹规划流程如图4所示。

Figure 4. Flowchart of trajectory planning in KSDP algorithm

图4. KSDP算法的轨迹规划流程图

我们从初始状态$s_t$ 开始进行轨迹规划，为了获取关键状态特征，我们需要对先验分布$p\left(z|s_t\right)$ 建模，即需要训练一个模型将当前状态映射到关键状态特征。关键状态特征z是一个低维的抽象表示，考虑到扩散模型对多模态分布的强大建模能力，我们使用扩散模型来实现这一过程，即$P_{\sigma }\left(z|s_t\right)$ 。

我们将编码器生成的关键状态特征z作为标签，通过直接预测原始的关键状态特征的方式来计算重构损失，从而训练扩散模型$P_{\sigma }$ ，Jun [22]等人的研究表明，这种方式在隐空间中比预测噪声ε训练模型效果更好。

$L\left(\sigma \right)=E_{z^0\sim E_{\varphi }\left(z|M\right),z^K~q\left(z^K|z^0\right)}\left({\Vert z^0-{\mu }_{\sigma }\left(z^K,s_t,k\right)\Vert }^2\right)$ (21)

值得注意的是，在轨迹规划的流程中，需要用到两次扩散模型：用于生成关键状态的$P_{\sigma }$ 和用于生成轨迹的$G_{\theta }$ ，这是两个不同参数的网络。

对于$P_{\sigma }$ 我们用无分类器引导的方式来简单高效的训练模型，然而一些研究表明[23]，无分类器引导方式可能促使生成模型倾向于生成分布中低概率区域的轨迹。这些样本往往偏离训练数据的主流分布，从而为生成模型引入噪声。

为解决这一问题，我们设计了一种筛选机制，以确保生成的关键状态特征z更贴近数据分布的特性。Pearce等人[23]提出了一种基于核密度估计(Kernel Density Estimation, KDE)的动作采样方法。我们在此基础上对其进行了改进，并将其应用于筛选扩散模型生成的关键状态特征。核密度估计是一种非参数估计方法，与参数估计不同，它不预先假设数据服从特定的分布形式，而是直接通过数据本身来估计分布。KDE通过计算每个候选样本z在特征空间中的密度值来评估其代表性。密度较高的样本通常位于数据分布的高密度区域，因此更可能被视为高质量的样本，适合用于后续的任务。

传统的核密度估计通常基于欧氏距离计算样本间的相似性，这种方法在面对简单一维或低维数据时效果良好，但其局限在于忽略了数据维度之间的相关性。特别是在关键状态特征空间这种相对高维的情景中，欧氏距离可能导致密度估计不够精确。为此，我们引入马氏距离(Mahalanobis Distance)替代欧氏距离，马氏距离通过融入数据的协方差矩阵，能够更准确地反映数据点之间的真实距离，从而提升密度估计的精确性和筛选效果。改进后的核密度估计公式定义为：

$\widehat{p}\left(z\right)=\frac{1}{n}{\displaystyle \sum _{i=1}^n{K}_M\left(\frac{z-z_i}{h}\right)}$ (22)

其中z是待评估的候选样本，$z_i$ 是n个候选样本中的第i个样本，h是带宽(bandwidth)，用于控制核函数的覆盖范围，h越大覆盖范围越广，密度估计越平滑。核函数$K_M(\cdot )$ 是KDE的核心，我们使用高斯核函数，其定义为：

$K_M\left(\frac{z-z_i}{h}\right)=\frac{1}{{\left(2\pi h^2\right)}^{d/2}{|\sum |}^{1/2}}\exp \left(-\frac{1}{2h^2}{\left(z-z_i\right)}^{\text{T}}{\sum }^{-1}\left(z-z_i\right)\right)$ (23)

核函数的作用是基于马氏距离计算$z_i$ 和z的相似度。d是z的维度。${\left(z-z_i\right)}^{\text{T}}{\sum }^{-1}\left(z-z_i\right)$ 是通过协方差矩阵Σ计算的马氏距离，能够捕捉维度间的相关性。指数部分能使高斯核在z接近$z_i$ 时赋予较高的权重，随着距离增加权重呈指数衰减，能够有效捕捉数据分布的局部结构。公式的系数为归一化因子，确保核函数的输出满足概率密度的性质。协方差矩阵计算公式为：

$\sum =\frac{1}{n-1}{\displaystyle \sum _{i=1}^n\left(z_i-\mu \right)}{\left(z_i-\mu \right)}^{\text{T}}$ (24)

$\mu =\frac{1}{n}{\displaystyle \sum _{i=1}^n{z}_i}$ (25)

通过改进的KDE概率密度估计公式，我们对每个生成的z样本计算其概率密度$\widehat{p}\left(z\right)$ 。最终，最优样本$z^{\ast }$ 通过以下方式筛选得出：

$z^{\ast }=\underset{z\in \left\{z_1,z_2,\cdots ,z_n\right\}}{\arg \max }p\left(z\right)$ (26)

在获得筛选后的关键状态特征后，其作为扩散模型$G_{\theta }$ 的条件输入，指导模型生成状态序列，并且利用逆动力模型生成动作序列，其中第一个动作被智能体执行与环境交互。当发生状态转移后，根据当前状态$s_{t+1}$ 重复轨迹规划流程。

3. 实验结果与分析

3.1. 实验环境和数据集

本研究在D4RL [24]数据集上进行实验，该数据集是离线强化学习领域的标准基准，提供多样化的任务环境以评估算法性能。实验涉及Gym-Mujoco任务(如HalfCheetah、Hopper、Walker2d)，聚焦智能体的基础运动控制，如奔跑和跳跃；Adroit任务(如Pen、Door、Hammer)，挑战高维动作空间中的复杂机械臂操作，如抓取和开门；Maze2D任务，测试智能体在二维迷宫中的路径规划能力；以及AntMaze任务，控制四足机器人Ant在复杂迷宫中完成长期导航，结合路径规划与机器人控制。实验严格遵循D4RL的标准化评估协议，通过在多场景、多难度级别上的组合测试，系统地验证了算法在不同数据分布下的泛化能力。通过上述任务和数据集的设置，本研究旨在全面评估所提算法在离线强化学习中的性能表现，涵盖从基础运动控制到复杂操作和长期导航的多种场景。

3.2. 对比实验

在本节中，我们通过对比实验来评估KSDP方法在离线强化学习数据集中的表现。我们选择了多种基线方法，包括模仿学习方法行为克隆(BC)；基于值的方法批量约Q学习(BCQ) [5]，保守Q学习(CQL) [6]，隐式Q学习(IQL) [9]；传统轨迹规划方法MPPI [25]，MoReL [7]，HiGoC [26]，MBOP [27]；基于Transformer的生成模型方法轨迹变换器(TT) [10]，决策变换器(DT) [12]；基于扩散模型的方法Diffuser [15]，DD [28]。这些基线方法涵盖了离线强化学习轨迹规划方法的主流范式，以确保对KSDP的全面评估，实验结果如下。

Gym-Mujoco是一个经典的连续控制任务数据集，其中高维的连续动作空间为模型决策提出了挑战。尤其是replay和medium数据中包含大量次优轨迹，对模型的鲁棒性要求较高。如表1所示，KSDP在Gym-Mujoco数据集的平均得分达到82.7，显著高于所有基线方法，包括Diffuser (77.5)、DT (74.7)和TT (78.9)。特别是在medium-replay中，KSDP的平均得分72相比于Diffuser的67.3约提升了7%的性能。KSDP的关键状态特征提取机制能够有效过滤次优数据中的噪声，结合逆动力解码器生成高质量的动作序列。这使其在面对次优数据和动态变化时表现出较强的鲁棒性。

Table 1. Comparative experimental results on Gym-Mujoco datasets

表1. Gym-Mujoco数据集中的对比试验结果

Datasets	BC	MBOP	MoReL	TT	DT	Diffuser	KSDP
halfcheetah-medium-expert-v2	55.2	105.9	53.3	95.0	86.8	88.9	89.3
walker2d-medium-expert-v2	107.5	70.2	95.6	101.9	108.1	106.9	107.4
hopper-medium-expert-v2	52.5	55.1	108.7	110.0	107.6	103.3	110.5
halfcheetah-medium-v2	42.6	44.6	42.1	46.9	42.6	42.8	47.1
walker2d-medium-v2	75.3	41.0	77.8	79.0	74.0	79.6	80.7
hopper-medium-v2	52.9	48.8	95.4	61.1	67.6	74.3	92.4
halfcheetah-medium-replay-v2	36.6	42.3	40.2	41.9	36.6	37.7	42.5
walker2d-medium-replay-v2	26.0	9.7	49.8	82.6	66.6	70.6	79.2
hopper-medium-replay-v2	18.1	12.4	93.6	91.5	82.7	93.6	94.3
Average	51.9	47.8	72.9	78.9	74.7	77.5	82.7

Maze2D是一个二维导航任务数据集，该任务环境是一个典型了稀疏奖励环境，只有当智能体到达目标时才会获得奖励反馈，这显著增加了模型的规划难度。相比其他数据集，虽然动作–状态空间的维度较低，但路径规划复杂。此外，在Multi2D设置下的每个episode开始时目标位置是随机初始化的，由于目标位置的随机性，环境的复杂度和不确定性更高，智能体需要具备更强的适应性和泛化能力。如表2所示，KSDP在Maze2D数据集上平均得分达到92.4，为所有基线方法中最高。AntMaze数据集使用MuJoCo的Ant机器人进行导航，迷宫环境和Maze2D相同，机器人运动涉及多关节协调，增加了控制难度。如表3所示，KSDP在AntMaze数据集上的平均得分达到82.7，证明了在稀疏奖励任务中，KSDP借助逆动力解码器同样可以适应高维的动作空间控制任务，展现了对高维状态和复杂动态的处理能力。KSDP在稀疏奖励环境下的优异表现归功于其条件扩散模型和关键状态引导机制。扩散模型能够生成覆盖长时程的高质量轨迹，而关键状态提取帮助模型聚焦于导航中的重要决策点。

Table 2. Comparative experimental results on AntMaze datasets

表2. AntMaze数据集中的对比试验结果

Datasets	MPPI	IQL	HiGoC	Diffuser	DD	KSDP
Maze2D-U-Maze-3	14.4	23.2	61.2	82.6	83.9	89.8
Maze2D-U-Maze-3	5.7	19.8	59.8	87.8	85.8	92.2
Maze2D-Large-2	3.9	31.1	45.4	87.9	87.3	93.4
Multi2D-U-Maze-3	17.8	16.5	67.9	85.4	86.9	95.1
Multi2D-Medium-2	8.1	8.9	52.4	85.6	88.2	92.6
Multi2D-Large-2	4.5	10.3	42.1	89.3	91.7	91.2
Average	9.1	18.3	54.8	86.4	87.3	92.4

Table 3. Comparative experimental results on Maze2D datasets

表3. Maze2D数据集中的对比试验结果

Datasets	IQL	HiGoC	Diffuser	KSDP
AntMaze-U-Maze	62.2	91.2	76.0	85.2
AntMaze-Medium	70.0	79.3	45.5	82.6
AntMaze-Large	47.5	67.3	22.0	80.3
Average	59.9	79.3	47.8	82.7

Adroit数据集是离线强化学习中最具挑战性的任务之一，涉及超高维手部操作任务，同时需要模拟手的物理交互，动态建模难度大。如表4所示，KSDP在pen-cloned任务中获得最高得分47.7，优于CQL和IQL方法，但在hammer和door任务中得分较低(分别为2.8和1.8)，所有方法均表现不佳。KSDP在pen-cloned中的表现显示其在高难度任务中仍有潜力，关键状态提取和扩散模型能够在一定程度上捕捉任务的关键特征。手部操作任务需要极高的动作精度和复杂的动态建模，KSDP当前的设计在这方面的优化不足，导致性能受限。我们分析，Adroit的超高维动作空间和物理交互复杂性，超出了KSDP对于当前特征提取和压缩表征的能力，模型难以生成满足任务要求的精确动作，未来可进一步优化KSDP在此类任务中的表现。

Table 4. Comparative experimental results on Adroi datasets

表4. Adroi数据集中的对比试验结果

Datasets	BC	BCQ	IQL	CQL	KSDP
pen-cloned	37.0	44.0	37.3	39.2	47.7
hammer-cloned	0.6	0.4	2.1	2.1	2.8
door-cloned	0.0	0.0	1.6	0.4	1.8

3.3. 消融实验

3.3.1. 关键状态提取模块消融实验和可视化分析

为了评估关键状态特征提取模块中各组成部分对轨迹生成性能的贡献，我们设计并开展了消融实验。本实验旨在分析Transformer编码器和β-VAE编码器在KSDP模型中的作用及其对生成轨迹质量的影响。通过控制变量的方式，所有模型在相同的训练和评估条件下进行，确保结果的公平性和可比性。我们设计了以下三种版本的关键状态特征提取模块进行对比实验：

β-VAE编码器提取关键状态：移除Transformer编码器，直接将原始轨迹输入β-VAE编码器提取关键状态特征。为适配序列输入，我们使用RNN网络在β-VAE前添加了一个简单的序列编码层。

随机关键状态：不使用任何编码器，直接将Transformer嵌入层的轨迹向量通过全连接层映射到关键状态特征维度，将其输入条件扩散模型作为条件。

KSDP：输入轨迹首先通过Transformer编码器处理，生成上下文增强的特征序列，随后输入β-VAE编码器，压缩为关键状态特征。这些特征作为条件输入条件扩散模型，最终生成完整轨迹。

实验基于D4RL数据集中的Gym-Mujoco任务开展，实验结果如表5所示。

Table 5. Ablation study results of Key-State extraction module

表5. 关键状态提取模块消融实验结果

Datasets	β-VAE提取关键状态	随机关键状态	KSDP
halfcheetah-medium-v2	42.3	34.7	47.1
walker2d-medium-v2	76.5	68.4	80.7
hopper-medium-v2	88.4	74.2	92.4
halfcheetah-medium-replay-v2	37.1	31.6	42.5
walker2d-medium-replay-v2	76.6	64.3	79.2
hopper-medium-replay-v2	88.9	78.5	94.3

实验结果表明，KSDP取得了最佳表现，显著优于随机关键状态特征的模型。这一结果验证了Transformer编码器和β-VAE编码器在关键状态提取中的重要性和协同效应。KSDP充分利用了Transformer的自注意力机制捕捉轨迹中重要状态信息和长期依赖关系，以及β-VAE的潜空间压缩能力提取关键信息。扩散模型通过关键状态的引导显著提升了生成轨迹的质量。

与仅使用β-VAE编码器的模型相比，KSDP的性能提高了约6%。这表明Transformer编码器在理解轨迹序列的上下文和动态关系方面至关重要，仅靠简单的RNN网络无法替代其功能。Transformer通过自注意力机制能够直接建模序列中任意两个元素之间的关系，克服了RNN在处理长序列时的梯度消失问题，能够更好地提取上下文特征。

为深入分析关键状态特征的分布特性，我们先将关键状态特征映射到对应任务的状态空间，然后使用t-SNE (t-distributed Stochastic Neighbor Embedding)将其高维特征降维至二维空间进行可视化。如图5所示，图左为原始数据集中状态空间，图右为提取到的关键状态空间。

如图所示，关键状态特征能够很好地表征原始数据分布，同时呈现出清晰的聚类结构。不同轨迹类别的特征分布明显分离，表明其成功捕获了轨迹的语义信息和动态模式，凸显了Transformer编码器和β-VAE编码器协同作用的显著效果。

Figure 5. Original trajectory spatial distribution map (left), key state spatial distribution map (right)

图5. 原始轨迹空间分布图(左)，关键状态空间分布图(右)

3.3.2. 轨迹价值引导消融实验

在KSDP方法中，轨迹价值预测模块通过预测关键状态所表征轨迹的价值V，为扩散模型的采样过程提供价值梯度引导，从而改进了传统无分类器引导的轨迹生成方式。为了评估该模块对生成轨迹质量的贡献，我们设计了消融实验，对比KSDP基准模型与移除价值引导的变体模型，分析不同引导策略对性能的影响。具体而言，我们设置了以下两种版本的模型进行对比：

无价值引导模型：移除轨迹价值预测器，仅依赖无分类器引导进行轨迹采样。这种配置依赖扩散模型的原始生成能力，未引入额外的价值信息。

KSDP：同时使用无分类器引导和轨迹价值引导。

实验基于D4RL数据集中的Maze2D任务进行评估。该任务具有稀疏奖励和长时程依赖的特点，因此适合检验轨迹生成质量的优劣。实验结果如表6所示。

Table 6. Ablation study results of trajectory value guidance

表6. 轨迹价值引导消融实验结果

Datasets	无价值引导模型	KSDP
Maze2D-U-Maze-3	83.1	89.8
Maze2D-Medium-2	86.5	92.2
Maze2D-Large-2	84.7	93.4

实验结果表明，KSDP在所有Maze2D子任务中的表现均优于无价值引导模型，平均分数提升约8.3%。在复杂度最高的Maze2D-Large-2中，KSDP得分比无价值引导模型提高10%。这种性能提升可以归因于轨迹价值引导的引入。Maze2D任务的稀疏奖励特性要求模型生成能够有效连接起点和目标的长时程轨迹，而无分类器引导仅依赖于数据分布的先验知识，容易生成偏离最优路径的轨迹。相比之下，KSDP通过价值预测模块为扩散模型提供了额外的优化方向，使得采样过程更倾向于生成高价值的轨迹。这种引导机制在高复杂度迷宫环境尤为显著，因为更复杂的迷宫需要更强的目标导向能力，而价值梯度恰好弥补了无分类器引导在这方面的不足。

本消融实验展示了轨迹价值引导在扩散模型轨迹生成中的重要性。通过结合无分类器引导和轨迹价值梯度引导的方式，为生成高价值轨迹提供了强有力的支持。

3.3.3. 筛选机制消融实验

在KSDP算法的轨迹规划过程中，关键状态筛选机制是其核心组成部分。该机制利用基于核密度估计的方法，从扩散模型生成的多个关键状态候选样本中筛选出最优样本，以减少噪声并确保关键状态符合数据分布特性。为了评估这一机制的作用，我们设置了以下两种版本的模型进行对比：

无筛选机制模型：移除筛选机制，直接使用扩散模型生成的一个关键状态样本作为轨迹生成的条件输入。

KSDP：保留完整的筛选机制。

为确保实验的可控性，两种模型在状态轨迹生成和动作执行阶段保持一致，仅在关键状态筛选机制的有无上有所差异。实验在D4RL数据集中的AntMaze任务上进行，该任务环境状态空间更加复杂，在关键状态特征生成时容易产生噪声，实验结果如表7所示。

Table 7. Ablation study results of the selection mechanism

表7. 筛选机制消融实验结果

Datasets	无筛选机制模型	KSDP
AntMaze-U-Maze	84.7	85.2
AntMaze-Medium	80.4	82.6
AntMaze-Large	76.1	80.3

实验结果表明，KSDP在所有AntMaze子任务中的成功率均高于无筛选机制模型。在U-Maze和Medium环境下，无筛选机制模型与KSDP分数相差不大，我们分析在相对简单的任务环境，即使没有筛选机制，模型也能生成较为合理的轨迹，因此分数下降幅度较小。在Large环境下，KSDP分数相较于无筛选机制模型有较大提升。这表明在高维复杂任务中，筛选机制的作用尤为重要，能够有效减少噪声、优化关键状态的选择，从而提升轨迹规划的成功率。

4. 结论

本文提出了一种基于关键状态的扩散模型轨迹规划方法，该方法利用Transformer和β-VAE模型从原始轨迹中提取出关键状态特征，并以此作为条件驱动扩散模型生成状态轨迹，随后通过逆动力模型生成相应的动作序列。在轨迹规划过程中，引入了基于核密度估计的关键状态筛选机制，有效降低了模型生成过程中的噪声干扰。在离线强化学习标准数据集D4RL中的多个任务环境进行的实验表明，KSDP的模型性能优于目前的先进方法，且在次优数据集和复杂环境中的鲁棒性更高。最后，通过对比实验和消融实验，在多个数据集上验证了KSDP的性能优势，充分展示了其在轨迹规划任务中的有效性和鲁棒性。未来可探索更高效的扩散模型架构，以解决其计算复杂性并缩短模型采样时间。