蛙跳扩散模型轨迹预测

扩散模型轨迹预测

参考论文《Leapfrog Diffusion Model for Stochastic Trajectory Prediction》
CVPR2024

1. 问题定义

目的是得到model $g_{\theta }(\cdot )$，参数$\theta$来生成分布${\mathcal{P}}_{\theta }=g_{\theta }(\mathbf{X},{\mathbb{X}}_{\mathcal{N}})$，基于分布${\mathcal{P}}_{\theta }$来画$K$个样本， Y ^ = { Y 1 ^ , Y 2 ^ , . . . , Y K ^ } \hat{\mathcal{\mathbf{Y}}}=\{\hat{\mathbf{Y_1}},\hat{\mathbf{Y_2}},...,\hat{\mathbf{Y_K}}\} Y^={ Y1​^​,Y2​^​,...,YK​^​}，这样至少有一个样本是接近真实的未来轨迹。总体问题定义：
${\theta }^{\ast }={\min }_{\theta }{\min }_{{\hat{\mathbf{Y}}}_i\in \hat{\mathcal{Y}}}D\left({\hat{\mathbf{Y}}}_i,\mathbf{Y}\right),\text{ s.t. }\hat{\mathcal{Y}}\sim {\mathcal{P}}_{\theta }$
$\mathbf{X}$和${\mathbb{X}}_{\mathcal{N}}$分别表示ego车辆的过去轨迹和neighboring车辆，$\mathbf{Y}$是ego车辆的未来轨迹。
通过一系列去噪步骤来学习轨迹分布，先执行前向diffusion的加噪到 未来轨迹的ground-truth上，然后，用条件去噪过程从过去轨迹的噪声中来恢复未来轨迹。
Diffusion 过程：

2. 方法论

2.1 前向扩散

$${\mathbf{Y}}^0=\mathbf{Y}$$
初始化扩散轨迹，
$${\mathbf{Y}}^{\gamma }=f_{\text{diffuse }}\left({\mathbf{Y}}^{\gamma -1}\right),\gamma =1,\cdots ,\Gamma$$
使用前向$f_{diffuse}(\cdot )$向${\mathbf{Y}}^{\gamma -1}$添加连续噪声来获取扩散后的${\mathbf{Y}}^{\gamma }$，其中${\mathbf{Y}}^{\gamma }$是第$\gamma$次diffusion步骤

2.2 逆过程

$${\hat{\mathbf{Y}}}_k^{\Gamma }\stackrel{i.i.d}{\sim }\mathcal{P}\left({\hat{\mathbf{Y}}}^{\Gamma }\right)=\mathcal{N}\left({\hat{\mathbf{Y}}}^{\Gamma };0,\mathbf{I}\right)\text{, sample }K\text{ times }$$
从正态分布中抽取$K$个独立同分布的样本初始化去噪轨迹$\widehat{{\mathbf{Y}}_k^{\Gamma }}$
$${\hat{\mathbf{Y}}}_k^{\gamma }=f_{\text{denoise }}\left({\hat{\mathbf{Y}}}_k^{\gamma +1},\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right),\gamma =\Gamma -1,\cdots ,0$$
迭代应用去噪操作$f_{denoise}(\cdot )$以过去轨迹$\mathbf{X},{\mathbb{X}}_{\mathcal{N}}$为条件获取去噪轨迹。
$\widehat{{\mathbf{Y}}_k^{\gamma }}$是第$\gamma$次去噪轨迹第$k$次采样，最终$K$个预测轨迹$$\hat{\mathcal{Y}=\left\{{\hat{\mathbf{Y}}}_1^0,{\hat{\mathbf{Y}}}_2^0,\dots ,{\hat{\mathbf{Y}}}_K^0\right\}}$$
Note：
前向扩散处理不会用于推理步骤，在训练期间，${\mathbf{Y}}^{\gamma }$是第$\gamma$步$\widehat{{\mathbf{Y}}_k^{\gamma }}$的监督
每个去噪步骤都是扩散步骤的逆过程，每个${\mathbf{Y}}^{\gamma }$和$\widehat{{\mathbf{Y}}_k^{\gamma }}$共享基础分布。
以上是问题定义建模，传统扩散模型受限于大量去噪的步骤的运算时间限制，但是轨迹预测需要实时推理，如果去噪的步骤很少会导致未来分布的表示能力很弱。
方法
[图片]

2.3 蛙跳扩散模型的步骤

$\mathbf{X}$和${\mathbb{X}}_{\mathcal{N}}$分别是ego和neighboring智能体过去的轨迹，$\mathbf{Y}$是ego的未来轨迹，$\tau$是leapfrog的步数。
$${\mathbf{Y}}^0=\mathbf{Y}$$
$${\mathbf{Y}}^{\gamma }=f_{\text{diffuse }}\left({\mathbf{Y}}^{\gamma -1}\right),\gamma =1,\cdots ,\Gamma$$
$${\hat{\mathcal{Y}}}^{\tau }\stackrel{K}{\sim }\mathcal{P}\left({\hat{\mathbf{Y}}}^{\tau }\right)=f_{\mathrm{LSG}}\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right)$$
和标准扩散模型的步骤最主要的区别在这里，初始化器$f_{LSG}(\cdot )$直接对第$\tau$个去噪分布$\mathcal{P}(\widehat{{\mathbf{Y}}^{\tau }})$建模，假设等价于执行$(\Gamma -\tau )$个去噪步骤，从分布$\mathcal{P}(\widehat{{\mathbf{Y}}^{\tau }})$中抽样并获取$K$个未来轨迹${\hat{\mathcal{Y}}}^{\tau }=\left\{{\hat{\mathbf{Y}}}_1^{\tau },{\hat{\mathbf{Y}}}_2^{\tau },\dots ,{\hat{\mathbf{Y}}}_K^{\tau }\right\}$
$${\hat{\mathbf{Y}}}_k^{\gamma }=f_{\text{denoise }}\left({\hat{\mathbf{Y}}}_k^{\gamma +1},\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right),\gamma =\tau -1,\cdots ,0$$
在这一步，只需要对每个轨迹$\widehat{{\mathbf{Y}}_k^{\gamma }}$应用剩余$\tau$个去噪步骤来获取最终的预测$\hat{\mathcal{Y}}=\left\{{\hat{\mathbf{Y}}}_1^0,{\hat{\mathbf{Y}}}_2^0,\dots ,{\hat{\mathbf{Y}}}_K^0\right\}$
Note：
去噪步骤由$\Gamma$减少到了$\tau$，其远远小于$\Gamma$，模型初始化器对$\tau$去噪步骤直接提供了轨迹，加快了推理。和标准扩散模型相比，这里的抽样并非来自独立同分布的结果。
新模型和蛙跳模型使用相同的前向扩散过程，保证了表达能力。

2.4 蛙跳初始化器

通过学习的方式建模第$\tau$个去噪分布$\mathcal{P}(\hat{\mathbf{Y}})$，将分布拆解为三个部分：均值、全局方差和样本预测部分。过程如下：
$${\mu }_{\theta }=f_{\mu }\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right)\in {\mathbb{R}}^{T_{\mathrm{f}}\times 2}$$
$${\sigma }_{\theta }=f_{\sigma }\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right)\in \mathbb{R}$$
$${\hat{\mathbb{S}}}_{\theta }=\left[{\hat{\mathbf{S}}}_{\theta ,1},\cdots ,{\hat{\mathbf{S}}}_{\theta ,K}\right]=f_{\hat{\mathbf{S}}}\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}},{\sigma }_{\theta }\right)\in {\mathbb{R}}^{T_f\times 2\times K}$$
$${\hat{\mathbf{Y}}}_k^{\tau }={\mu }_{\theta }+{\sigma }_{\theta }\cdot {\hat{\mathbf{S}}}_{\theta ,k}\in {\mathbb{R}}^{T_{\mathrm{f}}\times 2}{\hat{\mathbf{Y}}}_k^{\tau }={\mu }_{\theta }+{\sigma }_{\theta }\cdot {\hat{\mathbf{S}}}_{\theta ,k}\in {\mathbb{R}}^{T_{\mathrm{f}}\times 2}$$
其中，$f_{\mu }(\cdot ),f_{\sigma }(\cdot ),f_{\hat{\mathbb{S}}}(\cdot )$是三个可训练模块，${\mu }_{\theta },{\sigma }_{\theta }$是$\mathcal{P}\left({\hat{\mathbf{Y}}}^{\tau }\right)$的均值和方差，${\hat{\mathbf{S}}}_{\theta ,k}$是第$k$次采样的正则化位置。
$${\mathbf{e}}_{{\mu }_{\theta }}^{\text{social }}=\mathrm{softmax}\left(\frac{f_{\mathrm{q}}(\mathbf{X})f_{\mathrm{k}}{\left({\mathbb{X}}_{\mathcal{N}}\right)}^{\top }}{\sqrt{d}}\right)f_{\mathrm{v}}\left({\mathbb{X}}_{\mathcal{N}}\right)$$注意力模块社交编码
${\mathbf{e}}_{{\mu }_{\theta }}^{\mathrm{temp}}=f_{\mathrm{GRU}}\left(f_{\mathrm{conv}1\mathrm{D}}(\mathbf{X})\right)$GRU时序编码
${\mu }_{\theta }=f_{\text{fusion }}\left(\left[{\mathbf{e}}_{{\mu }_{\theta }}^{\text{social }}:{\mathbf{e}}_{{\mu }_{\theta }}^{\text{temp }}\right]\right)$融合均值估计，得到轨迹均值（MLP）
采样预测模块$f_{\hat{\mathbb{S}}}(\cdot )$将标准差的估计作为输入，计算过程如下：
${\mathbf{e}}_{{\hat{\mathbb{S}}}_{\theta }}^{\sigma }=f_{\text{encode }}\left({\sigma }_{\theta }\right)$将标准差的估计经过编码生成高维的embedding${\mathbf{e}}_{{\hat{\mathbb{S}}}_{\theta }}^{\sigma }$，这样标准差的估计也在样本的预测过程中涉及了。
${\hat{\mathbb{S}}}_{\theta }=f_{\text{fusion }}\left(\left[{\mathbf{e}}_{{\mathbb{S}}_{\theta }}^{\text{social }}:{\mathbf{e}}_{{\hat{\mathbb{S}}}_{\theta }}^{\text{temp }}:{\mathbf{e}}_{{\hat{\mathbb{S}}}_{\theta }}^{\sigma }\right]\right)$
从蛙跳初始化器获取到K个样本${\hat{\mathcal{Y}}}^{\tau }=\left\{{\hat{\mathbf{Y}}}_1^{\tau },{\hat{\mathbf{Y}}}_2^{\tau },\dots ,{\hat{\mathbf{Y}}}_K^{\tau }\right\}$，然后执行$\tau$个去噪步骤来迭代精炼预测轨迹。
问题：剩余的$\tau$步指的是哪$\tau$步？采样K次是为了什么？

2.5 去噪模块

去噪模块$f_{denoise}(\cdot )$从过去的轨迹$(\mathbf{X},{\mathbb{X}}_{\mathcal{N}})$的条件下对轨迹，有两个训练模块，一个是基于Transformer的上下文编码模块学习社交-时序embedding和一个噪声估计模块$f_{ϵ}(\cdot )$用以估计需要减少的噪声。第$\gamma$步去噪的流程如下：
$\mathbf{C}=f_{\text{context }}\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right)$使用上下文编码模块从过去的轨迹中获取上下文条件$\mathbf{C}$,$f_{context}(\mathbf{X},{\mathbb{X}}_{\mathcal{N}})$和$f_{\mu }(\cdot )$是相同的结构
${\mathbf{ϵ}}_{\theta }^{\gamma }=f_{\mathbf{ϵ}}\left({\hat{\mathbf{Y}}}_k^{\gamma +1},\mathbf{C},\gamma +1\right)$通过上下文$\mathbf{C}$的多层感知机实现的噪声估计$f_{ϵ}(\cdot )$来估计带噪声轨迹${\hat{\mathbf{Y}}}_k^{\gamma +1}$的噪声${\mathbf{ϵ}}_{\theta }^{\gamma }$
$${\hat{\mathbf{Y}}}_k^{\gamma }=\frac{1}{\sqrt{{\alpha }_{\gamma }}}\left({\hat{\mathbf{Y}}}_k^{\gamma +1}-\frac{1-{\alpha }_{\gamma }}{\sqrt{1-{\bar{\alpha }}_{\gamma }}}{\mathbf{ϵ}}_{\theta }^{\gamma }\right)+\sqrt{1-{\alpha }_{\gamma }\mathbf{z}}$$ 标准去噪步骤
其中${\bar{\alpha }}_{\gamma }={\prod }_{i=1}^{\gamma }{\alpha }_i$和${\alpha }_{\gamma }$是扩散过程的参数，$\mathbf{z}\sim \mathcal{N}(\mathbf{z};0,\mathbf{I})$是噪声

3. 实践

3.1训练

分为两个阶段，第一阶段训练去噪模块、第二阶段聚焦于蛙跳初始化器。
（蛙跳初始化器的在给定分布$\mathcal{P}({\hat{\mathbf{Y}}}^{\tau })$的情况下，蛙跳初始化器的训练更加稳定）
第一阶段使用扩散模型的标准训练模式利用噪声估计的loss训练：
$${\mathcal{L}}_{\mathrm{NE}}={\Vert \mathbf{ϵ}-f_{\mathbf{ϵ}}\left({\mathbf{Y}}^{\gamma +1},f_{\text{context }}\left(\mathbf{X},{\mathbb{X}}_{\mathcal{N}}\right),\gamma +1\right)\Vert }_2$$
其中， γ ∼ U { 1 , 2 , ⋯   , Γ } , ϵ ∼ N ( ϵ ; 0 , I ) \gamma \sim U\{1,2, \cdots, \Gamma\},\boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}) γ∼U{ 1,2,⋯,Γ},ϵ∼N(ϵ;0,I)，扩散的轨迹${\mathbf{Y}}^{\gamma +1}=\sqrt{{\bar{\alpha }}_{\gamma }}{\mathbf{Y}}^0+\sqrt{1-{\bar{\alpha }}_{\gamma }}\mathbf{ϵ}$
反向传播loss和训练的参数在上下文编码模块和噪声估计模块。
第二阶段使用可训练的蛙跳初始化器优化蛙跳扩散模型并且冻结去噪模块。对于每一个样本，loss为
$$\begin{array}{rl}\mathcal{L}& ={\mathcal{L}}_{\text{distance }}+{\mathcal{L}}_{\text{uncertainty }}\\ & =w\cdot \underset{k}{\min }{\Vert \mathbf{Y}-{\hat{\mathbf{Y}}}_k\Vert }_2+\left(\frac{{\sum }_k{\Vert \mathbf{Y}-{\hat{\mathbf{Y}}}_k\Vert }_2}{{\sigma }_{\theta }^{2}K}+\log {\sigma }_{\theta }^2\right)\end{array}$$
其中$w\in \mathbb{R}$是超参数权重，第一项约束了K个预测的最小距离。直观上说，如果蛙跳初始化器生成分布$\mathcal{P}({\hat{\mathbf{Y}}}^{\tau })$的高质量估计，K个预测中之一一定与groud-truth非常接近。
第二项通过不确定性损失对重新参数化的方差估计${\sigma }_{\theta }$进行归一化，平衡预测多样性和平均精度。方差估计控制预测的分散性，弥合场景复杂性和预测多样性。
$\frac{{\sum }_k{\Vert \mathbf{Y}-{\hat{\mathbf{Y}}}_k\Vert }_2}{{\sigma }_{\theta }^{2}K}$将${\sigma }_{\theta }$的值与场景的复杂度呈正比关系。
$\log {\sigma }_{\theta }^2$使用正则化器为所有的预测生成高方差。
附：作者表达了技术手段，他们并不在第二阶段对蛙跳初始化器进行估计时使用显式监督的手段的原因如下：
显式监督的实现：在初始化器的估计过程中，分布$\mathcal{P}({\hat{\mathbf{Y}}}^{\Gamma })$可以从正态分布中去噪，为了完成显式监督，从正态分布下的$\mathcal{P}({\hat{\mathbf{Y}}}^{\Gamma })$中抽样M个（M>>K），然后通过$f_{denoise}$去噪直到获得去噪轨迹${\hat{\mathbf{Y}}}^{\tau }$，然后使用这M个样本计算统计量作为均值估计$f_{\mu }(\cdot )$和方差估计$f_{\sigma }(\cdot )$的显式监督。
然而，由于$\tau <<\Gamma$，那么对$M>>K$个样本运行$(\Gamma -\tau )\approx \Gamma$步去噪获得统计数据的过程会导致训练时间和存储消耗变得无法接受。（NBV数据集上一个epoch要6天）

3.2 推理阶段

在推理过程中，蛙跳扩散模型只需要$\tau$步，而不是$\Gamma$步去噪，从而加快了推理速度。具体来说，我们首先生成 K 个相关样本，以使用经过训练的蛙跳初始化器对分布$\mathcal{P}\left({\tilde{\mathbf{Y}}}^{\tau }\right)$进行建模。然后，这些样本将被输入去噪过程并迭代微调以产生最终的预测；参见算法 :

3.3 源码

实验复现：

conda create -n led python=3.7
conda activate led

注意spconv 1.x版本已经废弃，找不到了

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
git clone -b v1.2.1 --recursive https://github.com/traveller59/spconv.git
cd spconv
python setup.py bdist_wheel
cd dist
pip install spconv-1.2.1-cp37-cp37m-linux_x86_64.whl

https://drive.google.com/drive/folders/1Uy8-WvlCp7n3zJKiEX0uONlEcx2u3Nnx
数据下载地址
训练

python main_led_nba.py --cfg led_augment --gpu 0 --train 1 --info try1

3.4 问题记录

问题：

查看日志文件

看到了两处报错，一处是LIBC_PTHREAD找不到，一处是nvcc报错。
经过排查，需要完整克隆指令（https://github.com/traveller59/spconv/issues/264）

0.11.4这个不兼容torch 1.8.0