【论文阅读笔记】Traj-MAE: Masked Autoencoders for Trajectory Prediction

Abstract

通过预测可能的危险，轨迹预测一直是构建可靠的自动驾驶系统的关键任务。一个关键问题是在不发生碰撞的情况下生成一致的轨迹预测。为了克服这一挑战，我们提出了一种有效的用于轨迹预测的掩蔽自编码器(Traj-MAE)，它能更好地代表驾驶环境中智能体的复杂行为。

具体来说，我们的Traj-MAE采用了多种掩蔽策略来预训练轨迹编码器和地图编码器，允许捕获智能体之间的社会和时间信息，同时利用来自多个粒度的环境影响。为了解决使用多种屏蔽策略对网络进行预训练时出现的灾难性遗忘问题，我们引入了一个连续预训练框架，该框架可以帮助Traj-MAE有效地从各种策略中学习有价值的、多样化的信息。我们在多智能体和单智能体设置下的实验结果表明，Traj-MAE使用最先进的方法取得了具有竞争力的结果，并且显著优于我们的基线模型。该准则一经发布，将向公众开放。

1. Introduction

轨迹预测的目标是预测移动主体(例如行人和车辆)的未来轨迹，这是构建安全、舒适、可靠的自动驾驶系统的关键问题[30,61,35,12,49]。许多有前途的工作[21,7,47,25,65,56]已经被学术界和工业界提出了极大的兴趣和需求。研究表明，对智能体之间复杂的相互作用进行建模[45,42,44,6,26]在轨迹预测中非常重要。在此基础上，为了解决碰撞预测问题并生成一致的轨迹预测，有必要对智能体之间的社会和时间关系进行建模，并对地图[2]有一个全局的理解。在本文中，我们使用自监督学习来研究这个问题。

自监督学习旨在从未标记的数据中学习潜在语义，而不是基于人类注释构建表示。近年来，自监督学习在自然语言处理[14,57]和计算机视觉[55,36,4]中的应用取得了显著进展。其中最有前途的自监督方法之一是掩码自编码器(MAE)[22]，它在各种任务中都取得了成功[37,51]。此外，在相同的小规模数据集上进行预训练和微调对于学习良好的表示b[15]也是必不可少的。受这些作品的启发，我们的目标是使用掩码自编码器探索智能体与地图的多粒度之间的复杂相互作用。

如何设计一个有效的掩码自编码器来生成一致的轨迹预测?我们试图从以下几个方面来回答这个问题:(i)轨迹和高清地图的信息密度与图像的信息密度有很大的不同。

图像是具有高度空间冗余的自然信号，而轨迹则是具有复杂社会交互作用的连续时间序列信号，而高清地图则包含高度结构化的信息。考虑到这些差异，针对轨迹预测的模型需要进行相应的调整以捕获信息特征。因此，我们研究了不同的掩蔽策略和合适的掩蔽比的轨迹和高清地图。

我们开发了社会和时间掩蔽，使轨迹编码器能够从不同的角度捕获信息。我们还研究了多粒度掩蔽，以强制地图编码器从高清地图捕获结构信息。此外，我们发现无论采用何种掩蔽策略，高掩蔽率(50% ~ 60%)都能产生良好的效果，这就要求编码器对历史轨迹和高清地图有全面的了解。

(ii)缺乏有效的多策略预训练框架对有效的多模态轨迹预测提出了挑战。由于任务的复杂性，传统的从头开始的多任务学习[67]可能难以收敛，而传统的持续学习方法[11,38]由于无法在不忘记先前学习的知识的情况下用多个任务训练网络而受到限制。为了解决这个问题，我们提出了一种新的方法，利用先前学习的参数初始化网络，与原始策略同时训练新策略。因此，我们确保我们的网络可以在保留以前获得的知识的同时获得新的知识。

在此基础上，我们提出了一种高效实用的自监督轨迹预测框架——掩蔽轨迹自编码器(Traj-MAE)。如图1所示，Traj-MAE利用输入轨迹和高清地图的部分屏蔽，分别利用轨迹编码器和地图编码器重建被屏蔽的片段。通过采用多种掩蔽策略对输入轨迹和高清地图的缺失部分进行重构，轨迹编码器和地图编码器可以从多个角度全面理解输入的潜在语义。此外，我们还引入了一种新的连续预训练框架，这是一种高效的学习方法，可以同时使用多种策略训练模型，从而减轻灾难性遗忘的问题。

我们的核心贡献如下:•据我们所知，我们是第一个提出一个简洁高效的自监督轨迹预测的掩蔽轨迹自编码器。

•我们探索了不同的掩蔽策略，充分利用MAE来挖掘历史轨迹和高清地图的潜在语义。同时，提出了一种连续预训练框架，利用多策略对模型进行高效训练。

•我们对Argoverse和INTERACTION进行了广泛的实验，用于自动驾驶轨迹预测，并对TrajNet++的合成分区进行了行人轨迹预测。我们的TrajMAE在这些基准测试中取得了具有竞争力的结果，并且显著优于我们的基线模型。

2. Related Works

轨迹预测被广泛认为是一项序列建模任务，许多基于RNN的方法[1,66,32,8]被提出来建模智能体未来位置的轨迹模式，因为RNN(例如LSTM[23])在序列建模方面取得了显著的成功。由于transforms[53]捕获远程依赖关系的强大能力，许多基于变压器的方法已经出现并蓬勃发展。提出了STAR[60]，通过空间和时间转换器之间的交错来捕获复杂的时空相互作用。mmTransformer[30]设计用于分层地聚合过去的轨迹、道路信息和社会互动。为了预测多智能体的未来轨迹，AgentFormer[61]和AutoBots[20]给出了同时对时间维度和社会维度建模的解决方案。增强编码器在两个维度上对信息建模的能力是这项工作的一个有趣的中心焦点。

近年来，自监督学习在自然语言处理和计算机视觉领域取得了显著的成功，尤其是自编码方法。去噪自编码器(DAE)[54]是一种从损坏的输入中重建原始信号的学习表示方法。

BERT[14]可以看作是DAE的发展，它屏蔽输入令牌并训练模型来预测缺失的内容。随着BERT中提出的掩模语言建模(MLM)任务，提出了许多MLM变体[57,5]来提高变压器预训练的性能。同样，在计算机视觉中，自动编码方法经常关注不同的借口任务进行预训练[36,4,22]。其中最流行的一种方法是MAE[22]，它随机屏蔽输入的patch，并训练模型在像素空间中恢复被屏蔽的patch。基于MAE的不断进步[16,37,3]验证了其有效性。遵循MAE的概念，我们的方法专注于利用MAE作为一种工具来预训练具有强大特征提取能力的模型编码器。

持续学习是一种解决灾难性遗忘问题的方法，这种问题发生在连续学习不同输入模式的样本中。这些方法大致可分为重播、基于正则化和参数隔离方法[13]。回放方法[40,41,24,10,50]是在学习新任务的同时回放之前的任务样本，以减轻遗忘。相反，在学习新数据时，基于正则化的方法[48,39,63,28]经常在损失函数中引入正则化项来巩固之前的知识。参数隔离方法[31,46]为每个任务分配不同的模型参数，以防止任何可能的遗忘。

在这项工作中，我们提出了一个持续的预训练框架来解决遗忘问题，通过利用相关掩蔽策略的训练样本中包含的特定信息，我们能够提高模型编码器的泛化。

3. Approach

我们的Traj-MAE是一种复杂而有效的自监督方法。图2提供了Traj-MAE框架的概述。在本节中，我们首先介绍我们的骨干网。然后，我们深入分析了轨迹和高清地图重建的掩蔽策略。最后，我们讨论了如何将TrajMAE纳入我们的持续预训练框架。

3.1. Network Backbone

在这项工作中，我们使用具有变压器编码器-解码器架构(详见补充材料)的Autobots[20]作为基线模型来验证所提出方法的有效性。我们的Traj-MAE对输入轨迹和高清地图中的随机部分进行掩码，然后分别重建缺失部分。继MAE[22]和VideoMAE[51]之后，我们采用非对称编解码器设计来减少计算量。

Traj-MAE Encoder。在汽车人中，历史轨迹被编码成上下文张量，连同可学习的种子参数和地图上下文，被传递给解码器来预测未来的轨迹。受此设计启发，我们采用汽车人编码器作为轨迹编码器。然而，在《汽车人》中，高清地图是直接输入到解码器的，这使得模型很难捕捉到高清地图的固有信息。为了解决这一限制，我们引入了一个与轨迹编码器结构相似的地图编码器，以更好地重建被屏蔽的高清地图。然而，我们观察到，直接将地图编码器添加到汽车人的结果几乎没有改善(见补充材料)。然而，我们发现用我们提出的掩蔽和重建策略对地图编码器进行预训练可以进一步提高准确率，验证了我们预训练策略的有效性。

Traj-MAE Encoder。Traj-MAE中的编码器仅处理输入的未屏蔽部分，而解码器则从潜在表示和掩码令牌中重建缺失部分。掩码令牌是共享向量，指示需要预测的缺失部分的存在。此外，将位置嵌入添加到所有令牌中以提供位置信息。Traj-MAE解码器采用比编码器浅的Transformer块设计，仅在预训练期间用于执行轨迹和地图重建策略。这使得解码器架构灵活且独立于编码器架构。使用轻量级解码器进行预训练可以显著减少预训练时间。