【Motion Forecasting】SEPT：自监督学习与运动预测任务相结合达到新SOTA

SEPT: Towards Efficent Scene Representation Learning For Motion Prediction

今天分享一篇自动驾驶运动预测方向的文章《SEPT: Towards Efficent Scene Representation Learning For Motion Prediction》，发布于2023年，来自于清华大学。

原论文链接：https://arxiv.org/abs/2309.15289

摘要

运动预测对于自动驾驶汽车在复杂的交通场景中安全地行驶来说是一项关键的任务。有效地提取驾驶场景中交通参与者之间的时空关联（spatiotemporal relationships）对于进行准确的运动预测来说是至关重要的。受预训练的大语言模型的启发，本文提出了SEPT，它是一种利用自监督学习任务来加强模型对复杂交通场景时空信息理解的建模框架。预训练的编码器会通过微调的方式进一步应用于下游的运动预测任务。大量的实验结果表明，SEPT不需要复杂的模型设计或人为的特征工程，即可在Argoverse 1和Argoverse 2的运动预测数据集上取得SOTA水准的结果。

结论

本文提出了SEPT（Scene Encoding Predictive Transformer，SEPT），它是一种用于运动预测任务的准确且有效的建模框架。通过利用三种针对场景理解的掩码重建自监督学习任务，SEPT以轻量且统一的模型架构在两个运动预测的benchmark数据集上达到了SOTA水准的结果。消融实验的结果进一步证明了本文方法所引入的预训练任务可以以一致且加性的方式大幅度地提升运动预测模型的性能。

同时，作者表明，SEPT仍然存在局限性，由于SEPT采用的是以代理为中心的场景表示（agent-centric scene representation），SEPT是一种单代理预测框架（single-agent prediction framework），将它拓展到多代理联合预测任务上并不简单。在未来的工作中，作者将会进行进一步的探索，来将SEPT拓展为效果更好的多代理联合预测框架。

Introduction（精简版）

针对运动预测任务，过去的SOTA方法通常采用复杂的模型架构设计，并通常依赖于基于锚点的建模（anchor-based modeling）或是基于提议的改进策略（proposal-refinement scheme）来进一步提升模型的预测性能。上述方法通常会导致模型的设计过于复杂。

尽管过去针对场景编码的方法聚焦于特征工程以及模型设计的创新，但作者认为，建立在统一架构之上的场景编码器可以通过正确的训练策略来进一步提高模型对于交通场景的理解能力。

受大语言模型当中自监督学习的启发，作者认为，通过利用有效的自监督学习目标，运动预测模型可以隐式地学习到环境动态信息以及交通参与者之间交互等有用的交通场景信息。

基于此，本文方法提出了Scene Encoding Predictive Transformer（SEPT），它是一种简洁并且强大的运动预测框架，它利用了自监督学习任务来提高模型对交通场景中时空信息的理解。

作者为SEPT设置了三项自监督预训练学习任务，来捕捉三种场景上下文的关键信息：

1. 代理历史运动信息的时序依赖；
2. 道路网络的空间结构；
3. 道路和代理之间的交互；

上述三种信息分别对应于三种预训练任务：

1. Marked Trajectory Modeling（MTM）；
2. Masked Road Modeling（MRM）；
3. Tail Prediction（TP）；

正如上图中所示，MTM会随机掩住并重建代理历史运动轨迹当中的一些轨迹点，以编码由运动约束而产生的时序依赖；类似地，MRM会随机地掩住与道路信息相关的特征的一部分，并通过未被掩住的部分对特征进行重建，使得编码器可以有效地对道路的拓扑结构以及连接关系进行捕捉；上述两种任务均是针对单一模态的（分别针对代理和道路），第三个任务TP聚焦于捕捉模态之间的交互关系，通过进行短程运动预测任务来完成。在TP任务中，本文方法将代理的轨迹分为两部分，分别是头（head）和尾（tail），任务的目标是基于头部分（head section）以及道路的上下文信息对尾部分（tail section）进行预测。

在通过上述三种自监督学习任务对编码器进行预训练之后，编码器会进一步进行微调，用于下游的运动预测任务。

Related Works（精简版）

Scene encoding with Transformers

目前，注意力机制以及Transformer架构已经被广泛地应用于运动预测任务，来对交通场景中的长程依赖以及复杂交互进行更好的建模。在早期，研究者们在不同的子模块中利用注意力机制，来对代理交互或是代理-地图交互等空间关系进行编码。也有一些工作使用图结构对道路元素进行编码，并使用Transformer架构来捕捉图结构输入之间的关联。

近期的SOTA方法倾向于同时对输入的所有模态进行处理，包括代理的历史运动轨迹以及道路网络，这类方法使用层级堆叠的（hierarchically stacked）transformer blocks对输入进行处理。这类方法通常会将整个交通场景表示为3-D tensors（space-time-feature），并交替地在时间和空间维度上使用分解注意力（factorized attention）来捕捉场景当中的时空依赖关系。

而对于SEPT，本文方法通过使用同一种模型架构设计实现了更加紧致的信息处理。在SEPT中，空间和时间信息会被顺序地进行编码，使得模型的结构更加简洁，并且参数量更少。

Self-supervised learning in motion prediction

近期的一些方法同样探索了自监督学习在运动预测任务当中使用的可能性，比如Traj-MAE和Forecast-MAE。从方法的命名不难看出，这两种方法都是受MAE的启发而提出的。其实早在2020年，自动驾驶运动预测方向向量化方法的代表VectorNet就已经引入了自监督学习任务，通过随机掩住图当中的某个结点的特征，再通过邻域结点对掩住结点的特征进行重建，从而鼓励模型学习更全面的场景交互信息。

SEPT与上述方法均有一定程度的不同，基于简洁且统一的模型架构设计，SEPT通过三种预训练任务使得模型在两项运动预测任务上达到了SOTA水准的精度，并且其消融实验表明，每一种预训练任务均是有效的。

Approach

Input Representation

Agent Trajectories

本文方法将代理的历史运动轨迹表示为$[A,T,D_h]$，它捕捉了与目标代理相距最近的$A$个邻域代理的轨迹，包括目标代理自身的轨迹。每一条代理的轨迹被表示为时间长度为$T$的状态向量，包括代理的位置坐标、当前位置的时间戳、代理的类型以及由数据集得到的代理的其它属性。所有的位置坐标都将会基于目标代理在当前时刻$T$的位置和方向进行标准化（agent-centric）。此外，并非所有的代理都具有完整的$T$时刻的历史观测，那些少于$T$时刻的观测将会使用掩码进行填充，以进行进一步的注意力计算。

Road Network

驾驶场景的道路网络将会被表示为$R$个道路向量$[R,D_r]$。每一个道路向量都是包含道路段起点和终点、道路转向以及来自于数据集当中的其它道路语义信息的有向道路中心线段（each raod vector is a directed lane centerline segment with features）。为了捕捉道路网路细粒度的特征，道路中心线将会被分割为长度不超过五米的向量。

通过上述方式得到的道路网络表示将会使得输入向量的集合规模过大，使得在计算注意力时带来额外的计算量和存储开销。因此，SEPT引入了道路剪枝模块（road pruning module），来识别目标代理在预测时长内可能的未来道路，从而构造一个更加紧致的道路向量子集$[S,D_r]$。

Model Architecture

SEPT采用了有效且简洁的模型架构设计，由场景编码器和未来轨迹解码器组成。

Projection Layer

Projection Layer会将不同模态的输入映射到一个共享的高维向量空间$R^D$当中，来执行后续的注意力计算。

TempoNet

TempoNet由$K_T$各堆叠的Transformer Encoder Blocks组成，用于代理历史轨迹的编码。TempoNet将会以代理的历史运动轨迹编码$[A,T,D]$作为输入，并沿时间维度$T$执行self-attention的计算。之后，使用max-pooling来对时间维度的特征进行聚合，得到代理的嵌入表示$[A,D]$。

TempoNet中使用T5当中提出的简化版的Positional Encoding，来对时间维度上的相对位置进行编码。

SpaNet

SpaNet与TempoNet的架构类似，采用了$K_S$个Transformer Encoder Blocks的堆叠，它的目标是对交通场景中的空间关联进行捕捉。SpaNet的输入由代理和道路的嵌入表示$[A+S,D]$组成。

与过去的SOTA方法相比，SEPT避免使用复杂的模块来分别对代理之间的交互、道路图交互以及代理和道路之间的交互进行编码。相反，SpaNet直接在空间维度$A+S$利用self-attention，来以一种统一的方式对代理-代理、道路-道路、代理-道路之间的交互进行编码（与ProphNet中AcSR采用的方法类似）。

Cross Attender

Cross Attender由$K_C$个堆叠的Cross Attention Layers构成。该模块会计算$N$个Learnable Queries的集合$[N,D]$与场景编码$[A+S,D]$之间的cross-attention，生成$N$个嵌入表示$[N,D]$，对应于$N$种输出的模态。每一个嵌入表示向量将会通过两个MLPs从语义空间映射到物理空间当中，产生2D的代理未来轨迹以及轨迹相应的置信分数。

Scene Understanding Training

在该阶段，将会对SEPT的编码器进行预训练，使得模型更好地捕捉交通场景当中的空间和时间关联，并为下游的轨迹预测任务提供高质量的语义表示。编码器的预训练是通过三项自监督学习任务完成的，它们分别针对TempoNet、SpaNet以及二者的组合。

在该阶段，编码器将会基于三项任务训练目标的加和来进行训练。

Task 1：Masked Trajectory Modeling（MTM）

MTM的目标是提升TempoNet对代理历史运动信息时序依赖性的理解。正如图3(a)所示，本文方法会随机使用可学习的mask token来以一定比例掩住可用的（eligible）输入轨迹。此处“可用的（eligible）”指的是比某个阈值更长的轨迹，过短的观测轨迹质量较低，其所包含的信息不足以使得模型学习到有用的表示。

对于MTM以及后续的两项场景理解任务，隐向量将会被输入到较浅的MLP解码器种来生成重建向量，MTM及后续的两项任务的目标函数均被定义为重建向量和原始向量之间的掩码$L_2$损失（mask $L_2$ loss）。

Task 2：Masked Road Modeling（MRM）

如图3(b)所示，MRM与MTM的思路类似，它的目标是使SpaNet学习到更多与道路向量输入相关的空间结构。与MTM当中的轨迹不同，来自于道路图当中的道路段并非序列，之间使用位置编码是不可取的。如果不施加与相对位置相关的归纳偏置，模型很难以token-level的mask进行学习。

相反，针对MRM，本文方法使用属性级别的掩码，通过将所选的一部分道路段除起点之外的其他属性设置为零来完成。被掩住的道路段的起点会被视为SpaNet的提示，用于重建其原始向量。通过上述方式，SpaNet必须学习对无序的道路向量的连通性和连续性有所感知，才能够恢复出道路段终点等被掩住的属性。

Task 3：Tail Prediction（TP）

在TP任务的设置种，嗲里的历史运动轨迹会被划分为两部分，分别是头部（head）和尾部（tail）。模型会基于头部轨迹，使用TempoNet和SpaNet来对尾部轨迹进行预测，它可以被视为完整运动预测任务的简化版，如图3c所示。

具体来说，代理的历史运动轨迹在被输入到TempoNet之前，会在$T_h$处被截断。短于$k{T}_h$的轨迹不会被使用。之后，截断轨迹的嵌入表示以及道路向量会被拼接，并输入到SpaNet当中。最后，对应代理的特征将会被输入到MLP种对尾部轨迹进行解码预测。

TP任务被用于将MTM学习到的轨迹表示和MRM学习到的道路上下文表示进行对齐。

Motion Prediction Training

在训练阶段，将预训练的编码器和解码器进行拼接，并以端到端的方式进行微调。运动预测任务的损失函数$L$包含两部分，分别是轨迹的回归损失$L_{reg}$和分类损失$L_{cls}$。

Experiments

SEPT在Argoverse 1和Argoverse 2上达到了SOTA水准的结果。

总结

• SEPT采用了非常简洁和标准的模型架构设计，在TempoNet种使用Transformer Encoder对代理的历史运动轨迹进行编码，并在时间维度上使用max-pooling对特征进行聚合，从而得到代理的历史运动编码$[A,D]$，其中$A$是包括目标代理在内的代理数量，$D$是特征维度。在SpaNet中将代理编码和道路编码进行拼接，得到$[A+S,D]$维度的特征，计算self-attention实现简洁的代理-代理、道路-道路、代理-道路交互编码。最后使用DETR-like architecture，使用learnable embedding与场景编码进行cross-attention，并进一步对未来轨迹进行解码。
• 基于上述简单的模型架构设计，SEPT使用了三种自监督预训练任务，来对模型的编码器进行预训练，从而对场景编码器进行性能提升。首先使用MTM，提升TempoNet对代理历史运动轨迹的时序依赖的理解；之后使用MRM，提升SpaNet道路结构的理解；最后使用TP任务来进一步对MTM学习到的轨迹表示以及MRM学习到的道路上下文结构信息进行对齐。经过预训练的编码器将会和解码器拼接，以在下游的轨迹预测任务上进行端到端的微调，从而进行最后的预测。
• SEPT在Argoverse 1/2上均达到了SOTA，它的成功不仅证明了自监督学习与运动预测任务相结合的可行性，在近期的工作大多数聚焦于轨迹解码改进的研究背景下，我认为SEPT的成功进一步证明了场景编码在整个运动预测框架当中的重要性。