HDGT:Heterogeneous Driving Graph Transformer for Multi-Agent Trajectory Prediction via Scene Encoding
今天分享的这篇自动驾驶运动预测的工作发布于2023年的TPAMI,名为HDGT(Heterogeneous Driving Graph Transformer),它在2021年于Waymo Open Motion Dataset(WOMD)上取得了SOTA水准的结果。
与上一篇介绍的GoRela类似,HDGT同样将驾驶场景建模为异构空间图,其中包含多种类型的结点(主要是代理结点【交通参与者】和道路结点)和多种类型的边连接关系(agent2agent、agent2map、map2agent、map2map)。同时,在对场景数据进行预处理时,HDGT同样采用了相对位置编码,以使得输入到模型当中的数据具有视角不变性。
摘要
将感知模块得到的驾驶场景数据编码为向量表示(vector representations)有助于下游的运动预测任务。驾驶场景中通常包含多种异构的元素,比如不同类型的交通实体(代理、道路以及交通信号灯),并且交通实体之间的语义关联是丰富且多样的。
同时,交通参与者之间同样存在相关性,这意味着交通场景中的空间关联(spatial relation)是一种相对性的关系,它们应该被编码在以交通参与实体为中心(ego-centric)的坐标系当中,而不是全局坐标系当中。
基于上述观察,本文提出异构驾驶图Transformer(Heterogeneous Driving Graph Transformer),它是一种用于对异构场景图进行特征提取的backbone。在对异构图进行构建时,本文方法会根据不同类型的语义关联,来连接不同类型的结点。在对空间相关性进行编码时,结点的坐标及其入边将会在以结点为中心(node-centric)的坐标系统当中进行编码。
针对图神经网络的消息聚合模块(aggregation module),本文以层级的方式使用Transformer结构来对异构的输入进行聚合。
实验结果表明,HDGT在Waymo Open Motion Challenge和INTERACTION Prediction Challenge两项自动驾驶运动预测任务上取得了SOTA水准的结果。
Introduction(精简版)
自动驾驶运动预测任务的一个关键性质是,它的数据本质上是异质的(heterogeneity),意味着数据当中的实体以及实体之间的关联可能具有不同的类型。
输入的数据可能包含代理的运动状态,交通信号灯的位置,带有可行驶区域的车道折线以及人行横道等多种类型。
将上述信息进行全面的编码并不简单,因为与CV当中的图像数据或是NLP当中的文本数据相比,驾驶场景数据的结构性很差,无法直接采用CNNs或Transformers等常用的神经网络来对驾驶场景所包含的非欧式数据进行处理。
为了对驾驶场景的数据进行处理,需要显式地考虑输入当中所包含的不同类别并处理输入当中实体之间的语义关联。
本文提出Heterogeneous Driving Graph Transformer,即HDGT,它将驾驶场景编码为异构图,其中包含两种类型的结点(代理与道路)以及用于编码这两种结点之间语义关系的边。具体来说,HDGT中基于Transformer设计了用于对图中结点和边的特征进行聚合【使用多头注意力实现聚合函数】和更新【使用FFN实现更新函数】的函数。为了充分地利用异构图当中的异质性,类别的异质将会作为模型的先验知识,不同类型的结点和边将会采用参数不同的Transformer模块进行处理。
与LaneGCN当中显式地设计四种Agent2Map、Map2Agent、Agent2Agent和Map2Map交互来对场景当中的元素进行更新不同,HDGT与GoRela类似,会在异构图当中同时对图中结点和边的特征进行更新。
除了异质性之外,驾驶场景数据的另一个特点是相关性,即:每一个代理都应该在各自的局部坐标系统当中处理局部邻域内与场景中其他元素之间的交互。如果在全局坐标系统当中对这种相关性进行处理,模型很容易对训练场景当中所使用的绝对坐标产生过拟合,从而无法处理在训练期间没有见过的场景。例如,在对车辆的未来轨迹进行预测时,模型应该在车辆的局部坐标系中考虑它是会直行,还是左转或右转,不应该在全局坐标系统当中对这种相对关系进行处理。
综上所述,相较于直接使用绝对坐标系统对场景数据进行处理,更合理的做法是在局部坐标系统当中对元素之间的相对关系进行处理。
过去的大多数方法通常会选取一个参考点(reference point),比如自车(ego agent,即自动驾驶场景数据的采样车辆,或是要预测的目标代理),之后,将包括车辆和道路在内元素的坐标全部变换到以参考点为坐标系统中心的坐标系当中。尽管上述方式显式地利用了场景中元素与自车的相对关系,但是非自车元素的坐标表示仍然不是在它们各自的局部坐标系统当中的,而是在以自车为坐标中心的系统当中。因此,正如Waymo Open Motion Dataset当中所建议的那样,如果采用自车作为参考点,为了在多代理预测任务上获得最好的表现,应该轮流以每一个要预测的目标代理作为参考点对整个场景的坐标进行多次标准化,显然这种方法会产生极大的计算开销。
本文设计了一种模型,它会在所有代理各自的局部坐标系统当中对称地对代理的信息进行处理(all agents’ are processing information in their local coordinate system and they are treated symmetrically),从而使得针对场景中所有代理的未来轨迹预测可以在单次前向传播中完成。
本文提出了HDGT,在本文方法中,结点的空间特征以及入边属性都将会被表示在以结点自身为中心的局部坐标系当中(in each node’s local, ego-centric system)。因此,结点特征的聚合以及更新都将会在它们各自的参考帧当中完成。
本文方法的主要贡献如下:
- 本文方法设计了一种统一的异构GNN框架,用于进行自动驾驶运动预测,并层级地使用Transformer对异质的输入进行处理。显式地对场景中所有的语义信息以及空间关联进行建模可以极大程度地提升预测的精度。统一的Transformer架构以及同时对所有元素的特征进行更新的方式使得异构GNN非常容易在模型中以基础模块的形式进行堆叠和拓展(stack and scale)。
- 与过去大多数方法中所采用的选取一个代理作为参考点的方法不同,本文方法提出了一种以对称的方式对场景中元素之间的空间关联进行处理的方法,该方法极大程度地提高了模型的泛化能力,并使得模型可以在不降低预测精度的前提下,在单次前向传播中完成多个代理未来多模态轨迹的预测。
- 本文方法在Waymo Open Motion Dataset以及INTERACTION数据集上进行了实验。实验结果表明,本文提出的HDGT可以在这两项数据集上达到SOTA水准的结果。
Related Work(精简版)
原文当中的相关工作对Heterogeneous Graph Neural Networks、Transformer以及Trajectory Prediction for Autonomous Driving进行了介绍,此处只对第三部分“Trajectory Prediction for Autonomous Driving”的阅读进行记录,有关前两部分的内容感兴趣的读者可以参考原文。
Scene Encoder for Trajectory Prediction
在驾驶场景中,代理的运动不仅仅会受其他代理的影响,包括道路、交通信号灯、交通标志在内的交通场景静态元素同样会对代理的未来轨迹产生影响。因此,在对自动驾驶运动预测模型进行设计时,应该同时将地图元素考虑在内。由于高精地图当中的地图元素是高度非结构化的(highly unstructured),处理地图元素的一种方法是将场景转换为栅格图,再使用视觉模型对栅格图进行特征提取。但栅格化方法会产生大量的信息丢失,并且这种特征表示的方法效率较低【实际上以今天的视角来看,已经鲜有新的工作使用栅格化方法了,基于图结构或序列化结构的表示是自动驾驶运动预测特征表示的主流方法】。
近期大多数方法采用向量化表示对驾驶场景当中的代理历史运动轨迹已经高精地图当中的静态元素进行表示,这类方法通常使用1D CNN或LSTM等序列模型来对时序数据进行特征提取【此处的时序数据指的就是代理的历史运动信息,在不使用分解注意力(factorized attention)的方法中,通常使用多尺度1D CNN、LSTM/GRU或是Transformer等处理序列的模型来对经过预处理的代理历史运动信息进行处理】,使用PointNet来对折线数据进行处理,并使用GNN来对代理以及道路元素之间的交互进行处理。
Output Decoder for Trajectory Prediction
自动驾驶运动预测的pipeline可以被总结为两部分,分别是场景编码(scene encoding)和轨迹解码(trajectory decoding)。一类工作聚焦于对场景编码的方式进行探索,比如本文提出的HDGT,而还有一类工作主要关注于对轨迹解码的方法进行探索,较为经典的方法包括TNT、DenseTNT和LaneRCNN等。
综上所述,本文方法主要关注于场景编码器的设计,与现有的方法相比,本文方法是第一个显式地使用一张异构图对驾驶场景中所有的结点类型以及边的语义以及几何连接关系进行表示的方法,并且本文提出的模型还能够以异构图为基础,同时对图中所有结点和边的特征进行更新,并且针对不同类型的结点和边,本文方法将会使用不同的权重对它们分别进行处理。
此外,本文方法还将在每一个结点各自的局部坐标系统当中对它们进行处理,这种具有视角不变性的表示方法可以提升模型的泛化能力,并使得模型可以在单次前向传播中完成多个代理的多模态轨迹预测。
值得注意的是,与HDGT同期的Multipath++和HiVT均显式地对自动驾驶运动预测当中的视角不变性表示进行了探索。基于类似的使用局部坐标系统来对结点特征进行处理的思想,Multipath++以及HiVT均在公开数据集上取得了较好的结果,进一步佐证了这篇文章在前文所提到的相对性(relativity)的重要性。
Methodology
Problem Formulation and Approach Overview
运动预测任务的输入包含两部分,分别是代理的历史运动信息以及高精地图当中的道路元素。运动预测任务的目标是对代理未来
T
T
T个时刻的轨迹【通常是2D坐标】进行预测。
Construction of the Heterogeneous Graph
HDGT将各个包含多个代理及道路元素的驾驶场景表示为有向的异构图(directed heterogeneous graph)。一张有向图可以被表示为 G = { V , E } G=\{V, E\} G={V,E},其中 V V V代表结点集合而 E E E代表边集合。一张异构图包含一个结点类型映射函数 τ ( v ) : V → T y p e v \tau(v): V\rightarrow{Type_v} τ(v):V→Typev和一个边类型映射函数 ϕ ( e ) : E → T y p e e \phi(e): E\rightarrow{Type_e} ϕ(e):E→Typee。在HDGT中,边的类型 e : u → v e: u\rightarrow v e:u→v是由源结点 u u u的类型以及目标结点 v v v的类型以及二者之间的关联类型所决定的。当构建场景的异构图时,会考虑场景中所有可能的组合以及语义关联。
针对代理结点,HDGT与VectorNet或是LaneGCN当中的做法(仅仅将距离小于某个阈值的结点进行连接)不同。这项工作观察到,在一些场景中,车辆会在未来时刻行驶一段很长的距离(in some scenarios, vehicles traversed a long way in the prediction horizon)。例如,在WOMD(Waymo Open Motion Dataset)中,一些车辆会在8秒的预测时长内行驶超过300米。综上,为所有代理设置一个固定的邻域半径并连接邻域半径内的结点是不必要的,并且还会花费额外的计算开销。例如,针对行人,行驶距离非常远的车辆不会对行人的未来行为产生影响。
对于所有代理结点,本文方法将会在最后一个观测时刻(即当前时刻)计算代理与其它代理之间的距离。对于所有折线形结点(nodes with the polyline shape),以折线的中心作为参考点来计算折线形结点与其它结点之间的距离。对于所有四边形结点(nodes with a polygon shape),以四边形的形心作为参考点。
HDGT采取与LaneRCNN当中类似的方法,对于每一个代理,将它们的距离阈值设置为它们当前行驶的速度乘以一个与类型相关的未来预测时长值(to be its speed times the prediction horizon with a type-specific buffer value)。
对于车道或是交通信号结点,本文方法将它们与代理结点根据上文提到的阈值进行反向连接(即,根据每一个代理 S p e e d × T h r e s h o l d Speed \times Threshold Speed×Threshold得到的阈值,对阈值内的车道与交通信号结点进行检索,将检索到的车道与交通信号结点与这个代理结点进行连接)。此外,对于车道结点,本文方法将会根据四种类型的边连接关系:左邻、右邻、进入(entry)、退出(exit),来对车道结点进行连接。由于车道的长度可能变化较大,本文方法进一步将长度超过某个阈值的车道分割为一个个小的车道段,使得车道结点的输入都是相似的。
综上所述,异构场景图中所包含的结点类型包括代理、车道以及交通信号灯。边类型包括Agent-Lane、Agent-Trafficsign、Lane-Agent、Lane-NextLane、Lane-PreviousLane、Lane-LeftLane、Lane-RightLane、Trafficsign-Agent。
Relative Spatial Relation Encoding
在将驾驶场景转换为异构场景图之后,需要对异构场景图当中的结点和边的特征进行初始化,正如标题中所说的相对空间关联编码,本文方法将会使用局部坐标系统对场景中结点的特征进行表示,本节也是这篇文章的精华所在。
正如上文所提到的,直接将原始的数据输入到模型当中并不能得到最优解,因为原始的空间数据是处于由数据集本身所定义的全局坐标系统当中的。考虑到同一种空间关联可能会由于处于不同的坐标系统中而产生不同的表示,空间数据应该被统一标准化为相同的形式。
此外,没有经过标准化的输入将会使得模型针对训练数据中出现的极端值产生过拟合,从而使得模型的泛化性较差。
因此,为了对数据进行标准化,一种常用的方法是选择场景当中的一个代理(通常是目标代理)作为参考(它的坐标点为原点,它的行驶方向为正方向),来对整个场景当中的数据进行标准化。尽管上述标准化的方式可以实现自车的相对关系,但使用自车结点作为参考中心的标准化方式会使得非自车结点不处于它们自身的局部坐标系统当中。【这一段阐述了以代理为中心(agent-centric)的场景编码方式的缺陷】
WOMD数据集当中提到,为了在多代理联合预测任务上取得最好的结果,采用上述以代理为中心的标准化方法时,应该在每一个预测之前均以当前要预测的代理为中心,重新对场景坐标进行标准化,这将会带来额外的计算开销。
综上所述,HDGT设计了一种编码方法,使得所有代理的空间位置都可以在它们各自的局部坐标系统当中进行处理,使得每一个代理都可以进行统一且对称的处理。基于局部坐标系统的表示方法,整个场景的代理未来轨迹预测可以在单次前向传播的过程中完成,而不需要显式地进行多次以代理为中心的标准化。
【下面对HDGT中使用局部坐标系统对场景当中的元素进行处理的方法进行记录:】
对于某个表示代理或地图元素的结点,本文方法将会在它们各自的局部坐标系统当中对它们的结点特征进行初始化。具体来说,对于每一个结点
v
v
v,本文方法将它们的参考系定义为
(
x
r
e
f
v
,
y
r
e
f
v
,
θ
r
e
f
v
)
(x^v_{ref}, y^v_{ref}, \theta^v_{ref})
(xrefv,yrefv,θrefv),其中
x
r
e
f
v
,
y
r
e
f
v
x^v_{ref}, y^v_{ref}
xrefv,yrefv表示的是局部坐标系统的原点,而
θ
r
e
f
v
\theta^v_{ref}
θrefv表示的是局部坐标系统的方向。对于结点
v
v
v的所有带有空间坐标的特征,使用
[
x
,
y
]
[x, y]
[x,y]减去
[
x
r
e
f
v
,
y
r
e
f
v
]
[x^v_{ref}, y^v_{ref}]
[xrefv,yrefv],使得结点
v
v
v当中的所有坐标均相对于其自身的参考系。之后,对于坐标及速度特征,使用
−
θ
r
e
f
v
-\theta^v_{ref}
−θrefv对它们进行旋转,即首先通过
−
θ
r
e
f
v
-\theta^v_{ref}
−θrefv得到旋转矩阵,再与
[
x
,
y
]
,
[
v
e
l
o
c
i
t
y
x
,
v
e
l
o
c
i
t
y
y
]
[x, y], [velocity_x, velocity_y]
[x,y],[velocityx,velocityy]进行相乘。此外,对于代理结点当中带有行驶方向的特征,将它们的行驶方向与
θ
r
e
f
v
\theta^v_{ref}
θrefv进行相减。
通过上述方式,所有结点的特征均相对于其自身的局部坐标系统。对于代理结点,将它们的参考坐标 ( x r e f v , y r e f v ) (x^v_{ref}, y^v_{ref}) (xrefv,yrefv)设置为它们最后一个观测时刻(即当前时刻)的坐标,而参考方向 θ r e f v \theta^v_{ref} θrefv则设置为它们在最后一个观测时刻的行驶方向,这意味着代理的历史运动轨迹及速度都是相对于其当前状态的。在对HDGT进行训练时,代理的训练标签,即未来轨迹,也会进行相同的坐标变换,使得代理的未来轨迹也是相对于当前状态的。
对于折线形的道路结点,本文方法选取折线的中间位置及折线起点至终点向量的方向作为折线形结点的参考坐标
(
x
r
e
f
v
,
y
r
e
f
v
,
θ
r
e
f
v
)
(x^v_{ref}, y^v_{ref}, \theta^v_{ref})
(xrefv,yrefv,θrefv)。对于四边形结点,选取四边形的形心作为参考坐标,选取四边形最长边的方向作为参考方向。通过上述方式,所有的道路元素均以一种统一的方式进行了预处理。
至此,所有的结点特征都被表示在了各个结点的局部坐标系统当中。然而,它们仍然是非结构化的,即:代理结点的特征由代理的历史运动状态序列组成,道路结点由折线和四边形组成。为了能够使用图神经网络并行地处理上述异构的信息,首先要做的是将上述结点的特征初始化为相同长度的张量。
针对代理结点,目前它的特征由标准化后的历史运动状态组成,首先使用1D-CNN来对代理历史运动信息进行编码。对于所有的地图元素结点,使用PointNet对折线特征和四边形特征进行编码。
此外,针对类型特定的信息,比如车道的类型(单车道、双车道等)以及道路信息,使用可学习的embeddings对它们进行表示,并将它们与由1D CNN和PointNet处理过后得到的几何特征拼接起来。
对于有向图当中的边,本文方法将会在目标结点的局部坐标系统当中对边的特征进行标准化。具体来说,对于边 u → v u\rightarrow v u→v,将 u u u视为源结点,将 v v v视为目标结点。边特征表示的是在结点 v v v视角下结点 u u u的信息。例如,边特征 e u → v e_{u\rightarrow v} eu→v表示的可能是在代理 v v v视角下车道 u u u的信息。回忆一下,上文已经得到了结点 u u u的向量表示,它是在结点 u u u自身的局部坐标系统当中进行过标准化的。但需要注意的是,结点 u u u的向量表示为 u − u 0 u-u^0 u−u0,它被表示在结点 u u u的视角下,而非目标结点 v v v的视角下。因此,需要先将参考坐标 u 0 u^0 u0变换到目标结点 v v v的局部坐标系统当中。
将结点 u u u和 v v v的参考坐标定义为 ( x r e f u , y r e f u , θ r e f u ) (x^u_{ref}, y^u_{ref}, \theta^u_{ref}) (xrefu,yrefu,θrefu)和 ( x r e f v , y r e f v , θ r e f v ) (x^v_{ref}, y^v_{ref}, \theta^v_{ref}) (xrefv,yrefv,θrefv)。为了将源结点 u u u变换到目标结点 v v v的视角下,需要的信息包括: ( x r e f v − x r e f u , y r e f v − y r e f u , cos ( θ r e f v − θ r e f u ) , sin ( θ r e f v − θ r e f u ) ) = ( Δ x u → v , Δ y u → v , cos Δ θ u → v , sin Δ θ u → v ) (x^v_{ref} - x^u_{ref}, y^v_{ref} - y^u_{ref}, \cos(\theta^v_{ref} - \theta^u_{ref}), \sin(\theta^v_{ref} - \theta^u_{ref})) = (\Delta x_{u\rightarrow v}, \Delta y_{u\rightarrow v}, \cos{\Delta \theta_{u\rightarrow v}}, \sin{\Delta \theta_{u \rightarrow v}}) (xrefv−xrefu,yrefv−yrefu,cos(θrefv−θrefu),sin(θrefv−θrefu))=(Δxu→v,Δyu→v,cosΔθu→v,sinΔθu→v),将上述信息记作 Δ P o s e u → v \Delta Pose_{u \rightarrow v} ΔPoseu→v。为了简便,使用MLP对上述信息进行处理,得到边的初始特征 e u → v 0 e^0_{u \rightarrow v} eu→v0:
e u → v 0 = M L P τ ( u ) ( c o n c a t [ u 0 , Δ P o s e u → v ] ) e^0_{u \rightarrow v} = MLP_{\tau(u)}(concat[u^0, \Delta Pose_{u \rightarrow v}]) eu→v0=MLPτ(u)(concat[u0,ΔPoseu→v])
其中 τ ( u ) \tau(u) τ(u)的含义是,用于对边特征进行处理的MLP权重取决于源结点 u u u的类型,异构的边将会使用不同的权重来完成。通过上述方式,使得结点 v v v的入边特征与结点 v v v的局部坐标系统相关,并且该操作是对所有结点对称的。将上述用于对边特征进行初始化的MLP称为ViewShift Unit,即视角变换单元。
综上所述,基于本节所提出的空间关系编码方式,本文方法将空间编码的过程分解为两步:
首先,对场景中单个元素的运动或几何属性进行编码,针对代理结点,使用代理在当前时刻的位置和行驶方向作为单个代理的局部参考坐标,对代理的历史运动状态进行标准化,再使用1D CNN对经过标准化的代理历史运动状态进行编码,得到代理结点相应的特征向量,作为代理结点的初始化特征。而针对折线形以及四边形的道路结点,将折线的中点以及四边形的形心作为局部坐标系的坐标原点,将折线起点至终点的方向以及四边形最长边的方向作为局部坐标系的参考方向,对道路元素进行标准化,使用PointNet对标准化后的坐标信息进行特征提取,得到道路元素初始化的结点特征。
然后,对场景中两两物体之间的关联进行编码,即对边特征进行编码。以边 e u → v e_{u\rightarrow v} eu→v为例,结点 u u u为源结点而结点 v v v为目标结点,HDGT中重点强调边 e u → v e_{u\rightarrow v} eu→v为目标结点 v v v的入边,其包含的边特征为将结点 u u u变换到结点 v v v视角下的相对姿态。即 e u → v 0 = M L P τ ( u ) ( c o n c a t [ u 0 , Δ P o s e u → v ] ) , Δ P o s e u → v = ( x r e f v − x r e f u , y r e f v − y r e f u , cos ( θ r e f v − θ r e f u ) , sin ( θ r e f v − θ r e f u ) ) = ( Δ x u → v , Δ y u → v , cos Δ θ u → v , sin Δ θ u → v ) e^0_{u \rightarrow v} = MLP_{\tau(u)}(concat[u^0, \Delta Pose_{u \rightarrow v}]), \Delta Pose_{u \rightarrow v}=(x^v_{ref} - x^u_{ref}, y^v_{ref} - y^u_{ref}, \cos(\theta^v_{ref} - \theta^u_{ref}), \sin(\theta^v_{ref} - \theta^u_{ref})) = (\Delta x_{u\rightarrow v}, \Delta y_{u\rightarrow v}, \cos{\Delta \theta_{u\rightarrow v}}, \sin{\Delta \theta_{u \rightarrow v}}) eu→v0=MLPτ(u)(concat[u0,ΔPoseu→v]),ΔPoseu→v=(xrefv−xrefu,yrefv−yrefu,cos(θrefv−θrefu),sin(θrefv−θrefu))=(Δxu→v,Δyu→v,cosΔθu→v,sinΔθu→v)。需要注意的是,初始的边特征没有显式地对源结点 u u u的特征进行变换,来将它变换到目标结点 v v v的视角当中,而是保存了将 u u u变换到 v v v视角下所需的信息,这部分信息被保存在 Δ P o s e u → v \Delta Pose_{u \rightarrow v} ΔPoseu→v当中,即 u u u的局部坐标系统与 v v v的局部坐标系统之间的变换。
基于上述双阶段的编码方式,场景中的所有结点均被表示在其各自的局部坐标系统当中,而结点的入边特征包含着源结点与目标结点局部坐标系统之间的相对关系。
Heterogeneous Driving Graph Transformer Layer
在对异构场景图进行构建并对结点以及边特征进行初始化之后,下一步要做的就是在结点和边之间进行消息传递,并对结点和边的特征进行更新。
Transformer
Transformer可以被视为一种具有排列不变性的集合函数,它可以通过点积注意力来以成对(pair-wise)的方式对集合当中元素的信息进行融合。Transformer主要由两部分组成,分别是多头注意力模块(multi-head attention module)和前馈模块(feedforward module),多头注意力模块可以将来自于集合中其它元素的信息与当前元素进行融合,而前馈模块可以对集合当中元素的信息进行更新。
具体来说,假设我们有一个包含 m m m个特征的集合,特征的维度为 h h h,记作: Z = { z i ∣ i = 1 , 2 , . . . , m } ∈ R m × h Z = \{z_i | i=1,2,...,m\}\in R^{m\times h} Z={zi∣i=1,2,...,m}∈Rm×h。在注意力机制中,我们要进行三次线性变换( W q ∈ R h × h q , W k ∈ R h × h k , W v ∈ R h × h v , h q = h k W_q\in R^{h\times h_q}, W_k\in R^{h\times h_k}, W_v\in R^{h\times h_v}, h_q=h_k Wq∈Rh×hq,Wk∈Rh×hk,Wv∈Rh×hv,hq=hk),得到Q、K和V:
Q , K , V = Z W q , Z W k , Z W v Q, K, V = ZW_q, ZW_k, ZW_v Q,K,V=ZWq,ZWk,ZWv
之后,执行点积注意力,得到权重:
A = S o f t m a x ( Q K T / h q ) A = Softmax(QK^T/\sqrt{h_q}) A=Softmax(QKT/hq)
最后,将权重 A A A与值 V V V相乘,得到:
Z ′ = A V Z' = AV Z′=AV
综上,本文将注意力函数定义为 Z ′ = A t t n ( Z ) Z' = Attn(Z) Z′=Attn(Z)。多头注意力对注意力进行了进一步拓展,将相同的输入使用多个权重不同的线性层进行处理,记作 M H A ( z ) MHA(z) MHA(z),即 Z ′ ′ = M H A ( Z ) = [ A t t n 1 ( Z ) , . . . , A t t n n h ( Z ) ] W o Z'' = MHA(Z) = [Attn_1(Z), ..., Attn_{n_h}(Z)]W_o Z′′=MHA(Z)=[Attn1(Z),...,Attnnh(Z)]Wo。
前馈模块是使用带有激活函数 δ ( ⋅ ) \delta(\cdot) δ(⋅)的双层MLP来实现:
Z ′ ′ ′ = F F N ( Z ′ ′ ) = δ ( Z W 1 + b 1 ) W 2 + b 2 Z''' = FFN(Z'') = \delta(ZW_1 + b_1)W_2 + b_2 Z′′′=FFN(Z′′)=δ(ZW1+b1)W2+b2
Z ′ ′ ′ Z''' Z′′′是Transformer layer的输出,它的维度通常与 Z Z Z相同,从而可以加入跳层连接。MHA和FFA当中均加入了跳层连接和Layer Norm(LN)用于稳定训练过程:
Z ′ ′ = M H A ( L N ( Z ) ) + Z Z'' = MHA(LN(Z)) + Z Z′′=MHA(LN(Z))+Z
Z ′ ′ ′ = F F N ( L N ( Z ′ ′ ) ) + Z ′ ′ Z''' = FFN(LN(Z'')) + Z'' Z′′′=FFN(LN(Z′′))+Z′′
Graph Neural Network
在HDGT中,信息的交换以及特征的更新均使用GNN来实现。
Aggregation Function
聚合函数 ρ \rho ρ以集合作为输入,将集合中的元素进行降维,得到单个元素,即聚合函数进行“聚合”的结果。聚合函数必须对输入具有排列不变性,并且它的输入数量是可变的(即,集合当中所含元素的数量是不固定的,对应到图中即为:每一个结点均有数量不固定的相邻结点)。
Update Function
更新函数 ϕ \phi ϕ将会以结点当前的特征以及从邻域结点得到的经过聚合函数 ρ \rho ρ处理的特征作为输入,来对当前结点的特征进行更新。
HDGT的输入是异构图,其中包含两类聚合函数和两类更新函数。聚合函数 ρ V \rho^V ρV作用于结点,并从每个结点的入边对入边的特征进行聚合;而 ρ E \rho^E ρE作用于边,并从每条边的源结点聚合特征。更新函数 ϕ V \phi^V ϕV和 ϕ E \phi^E ϕE基于聚合到的信息和当前特征分别对结点和边特征进行更新。
在HDGT中,包含 K K K层GNN,意味着结点和边的特征将会被更新 K K K次。
Node Aggregation and Update
针对用于结点的聚合函数 ρ V \rho^V ρV,本文方法使用Transformer的MHA来从每一个结点的入边聚合信息,正如上文中所提到的,目标结点入边的特征将会使用源结点特征 u 0 u^0 u0以及从源结点局部坐标系到目标结点局部坐标系的相对姿态 Δ P o s e u → v \Delta Pose_{u\rightarrow v} ΔPoseu→v进行初始化。
然而,Transformer在最初是被设计用于文本数据的,它假设集合当中所有的元素都是同质的,即单词(words)。这一假设对于驾驶场景数据来说是不成立的,因为驾驶场景当中包含多种类型的结点。
因此,针对所有的入边使用相同的权重是不合理的。此外,由于每一个代理结点在场景中的行驶速度是不同的,这将会导致它们的邻域半径不同,作者发现一些代理结点可能具有上千条入边。由于注意力机制是稀疏的,在入边过多时,很难对那些真正起到关键作用的信息进行聚合,反而会导致模型性能的下降。
为了缓解上述问题,本文层级地使用Transformer,首先单独地对每一种语义类型的关联进行聚合,再使用MLP来聚合不同类型的信息。在对入边特征进行聚合时,每一种类型的入边共享同一组MHA的权重,并执行cross-attention的计算,将结点特征视为 Q Q Q,将入边特征视为 K , V K, V K,V。通过这种层级的方式,将会得到一系列聚合特征,每一种边类型对应一个聚合特征集合。最后,将上述聚合特征拼接,并使用MLP进行处理,得到最终的每一个结点的聚合特征。
具体来说,对于目标结点 u u u及其 m m m条入边的集合 { e i } i = 0 i = m \{e_i\}^{i=m}_{i=0} {ei}i=0i=m,在第 k k k层时,结点的特征记作 v k − 1 v^{k-1} vk−1,其入边特征记作 { e i k − 1 } i = 0 i = m \{e^{k-1}_i\}^{i=m}_{i=0} {eik−1}i=0i=m。需要注意的是,模型第一层的输入就是上文所提到的结点和边的初始化特征。将目标结点 v v v所有的来自 k − 1 k-1 k−1层的类型为 j j j的入边集合记作 E j k − 1 ∈ R m j × h E^{k-1}_j \in R^{m_j \times h} Ejk−1∈Rmj×h,其中 m j m_j mj指的是所有类型为 j j j的入边的数目,而 h h h为边特征的维度。根据下述公式计算针对目标结点的入边类型特定(即首先计算类型为 j j j的入边)的注意力权重:
a j = S o f t m a x ( q j K j ) T / h ∈ R m j a_j = Softmax(q_j K_j)^T / \sqrt{h} \in R^{m_j} aj=Softmax(qjKj)T/h∈Rmj
其中 q j = v k − 1 W j , q u e r y ∈ R h q_j = v^{k-1}W_{j, query} \in R^h qj=vk−1Wj,query∈Rh,而 K j = E j W j , k e y ∈ R m j × h K_j = E_j W_{j, key} \in R^{m_j \times h} Kj=EjWj,key∈Rmj×h, W j , q u e r y W_{j, query} Wj,query和 W j , k e y W_{j, key} Wj,key是针对入边类型 j j j的线性变换矩阵。
通过上述方式,可以实现类型特定的注意力分数计算,它将场景图的异质性作为先验知识。例如,为了确保一辆车可以安全地行驶,它的注意力应该放在其它代理身上,而对于车道结点则需要区分对待。
对于所有类型的关联,使用上述的方式执行MHA(多头注意力),得到的输出是一个聚合特征的集合 { v j k − 1 , ′ ′ ∈ R h ∣ j = 1 , . . . , n r } \{v^{k-1,''}_j \in R^h | j=1, ..., n_r\} {vjk−1,′′∈Rh∣j=1,...,nr},其中 n r n_r nr是目标结点 v v v所有可能的入边的类型。使用MLP来各种类型的聚合特征进行融合,得到:
v k − 1 , ′ ′ = M L P ( c o n c a t ( { v j k − 1 , ′ ′ ∣ j = 1 , . . . , n r } ) ) v^{k-1, ''} = MLP(concat(\{v^{k-1, ''}_j | j=1, ..., n_r\})) vk−1,′′=MLP(concat({vjk−1,′′∣j=1,...,nr}))
至此,便实现了针对结点的聚合函数 ρ V \rho^V ρV。
而对于结点特征的更新函数 ϕ V \phi^V ϕV,使用FFN来对所有的结点特征进行更新。考虑到结点的特征本质上是异质的,针对不同类型的结点,使用不同权重的FFN对结点进行更新。
v k + 1 = F F N T y p e ( v ) ( v k − 1 , ′ ′ ) v^{k+1} = FFN_{Type(v)}(v^{k-1, ''}) vk+1=FFNType(v)(vk−1,′′)
其中 v k + 1 v^{k+1} vk+1是结点 v v v更新之后得到的特征。
为了描述简便,原文在对此处进行描述时,省略了跳层连接以及LayerNorm。
Edge Aggregation and Update
【HDGT的每一个GNN layer不仅会对目标结点的特征进行更新,还会同时对边特征进行更新,这也是HDGT的一大亮点,常见的图神经网络仅仅对结点的特征进行更新,而边特征不会变化】
边特征 e u → v e_{u \rightarrow v} eu→v表示的是在目标结点 v v v的视角下源结点 u u u的特征表示,它主要由源结点在自身局部坐标系统下的特征 u 0 u^0 u0以及源结点 u u u到目标结点 v v v局部坐标系统的姿态变换 Δ P o s e u → v \Delta Pose_{u \rightarrow v} ΔPoseu→v组成。因此,如果想要对边特征 e u → v e_{u \rightarrow v} eu→v进行更新,首先需要结点 u u u在当前层的特征,并将它变换到结点 v v v的视角下。
具体来说,假定在第 k k k层,我们有 u k − 1 , v k − 1 u^{k-1}, v^{k-1} uk−1,vk−1和 e u → v k − 1 e^{k-1}_{u \rightarrow v} eu→vk−1,分别表示的是来自上一层的源结点 u u u、目标结点 v v v以及连接二者的边 e u → v e_{u \rightarrow v} eu→v的特征。如果此时正处于第一层,则这些特征是上文所提到的初始化特征。
对于边的特征聚合函数 ρ E \rho^E ρE,首先对结点 u u u的特征进行检索,得到 u k − 1 u^{k-1} uk−1,并计算从结点 u u u的局部坐标系变换到结点 v v v的局部坐标系所需的相对姿态 Δ P o s e u → v \Delta Pose_{u \rightarrow v} ΔPoseu→v。
对于边的特征更新函数 ϕ E \phi^E ϕE。将聚合到的 u k − 1 u^{k-1} uk−1和 Δ P o s e u → v \Delta Pose_{u \rightarrow v} ΔPoseu→v以及上一层得到的边特征 e u → v k − 1 e^{k-1}_{u \rightarrow v} eu→vk−1进行拼接,并输入到MLP当中,得到边特征 e u → v k e^k_{u \rightarrow v} eu→vk。对于不同类型的边,此处仍然使用不同权重的MLP,以确保异质性:
e u → v k = M L P t y p e ( u → v ) ( c o n c a t [ u k − 1 , Δ P o s e u → v , e u → v k − 1 ] ) e^k_{u \rightarrow v} = MLP_{type(u \rightarrow v)}(concat[u^{k-1}, \Delta Pose_{u \rightarrow v}, e^{k-1}_{u \rightarrow v}]) eu→vk=MLPtype(u→v)(concat[uk−1,ΔPoseu→v,eu→vk−1])
在原文中,此处仍然省略了跳层连接和Layer Norm。
综上所述,通过上文所提到的针对结点和边的特征聚合与更新函数,可以完成对异构场景图的特征提取。所有的聚合与更新函数均与Transformer的操作类似,聚合函数使用MHA进行实现,而更新函数使用FFN进行实现,恰好对应于Transformer中最关键的两个模块。
Output Head and Training Objective
HDGT方法强调的是场景编码器的重要性,在轨迹解码阶段,使用一个regression-MLP和一个classification-MLP的组合来进行完成。
具体来说,针对每一个目标结点,使用HDGT最后一层输出的结点特征作为用于轨迹解码的特征表示。其中一个MLP模块用于对轨迹进行回归,输出的张量维度为 N × K × T × 2 N \times K \times T \times 2 N×K×T×2,其中 N N N是目标代理的数量, K K K是所需的模式, T T T是需要预测的长度, 2 2 2对应的是二维 ( x , y ) (x,y) (x,y)坐标。另一个MLP模块用于输出每一个代理多模态的未来轨迹中每一个模式对应的分数,其维度为 N × K N \times K N×K。不同类型的代理(车辆/行人/非机动车)将会使用不同的参数。
假定针对某个代理,输出为:
O r e g = { ( p 1 k , p 2 k , . . . , p T k ) } k ∈ [ 1 , K ] O_{reg} = \{(p^k_1, p^k_2, ..., p^k_T)\}_{k \in [1, K]} Oreg={(p1k,p2k,...,pTk)}k∈[1,K]
O c l s = { c k } k ∈ [ 1 , K ] O_{cls} = \{c^k\}_{k \in [1, K]} Ocls={ck}k∈[1,K]
其中 p i k = ( x i k , y i k ) p^k_i = (x^k_i, y^k_i) pik=(xik,yik)是代理在时刻 i i i所对应第 k k k个模式的坐标, c k c_k ck为置信分数。
对于每一个代理,HDGT在训练时仅仅对regression head中回归损失最低的那条轨迹的损失进行反向传播,以避免模式崩塌,并鼓励模型输出更多样的轨迹。对于classification head,使用交叉熵损失作为分类损失,回归损失最低的模式将会被视为正样本,而其他的模式均会被视为负样本。
Conclusion
本文提出了一种原则性的方法,用于对自动驾驶轨迹预测任务中的场景进行表示,从而更好地完成轨迹预测任务。该方法将驾驶场景建模为异构图,并使用Transformer来对图神经网络当中的聚合以及更新函数进行实现。
驾驶场景图当中结点所对应的空间特征均被表示在结点各自的局部坐标系统当中,从而实现了模型的视角不变性,使得模型不会受全局坐标系统平移或旋转所带来的场景变换的影响。
本文方法在Waymo Open Motion Dataset以及INTERACTION两项motion forecasting benchmarks上达到了SOTA水准的结果。