Autonomous Driving | GAIA-1: A Generative World Model for Autonomous Driving

GAIA-1 （‘Generative AI for Autonomy’，生成式自主人工智能），使用 video, text, and action为输入，生成真实的驾驶场景，同时提供对自车行为和场景特征的细粒度（fine-grained）控制
该算法将输入映射到离散的token上，将世界模型变为无监督序列建模的问题，然后预测下一个token

1 Introduction

GAIA-1致力于世界模型与视频生成；

使用 large language models 做物体的检测和识别，再用 video diffusion models 做场景生成；

2 Model

2.1 Encoding Video, Text and Action

使用了三个模态的输入：video, text, action，将其编码在维的空间中；

Image tokens

每一帧的video表示一个离散的token，使用预训练的图像token化模块（use a pre-trained image tokenizer）
考虑序列长度为的图像，每张图片再离散化为的离散 token；到这一步，就获得了一个序列，且对于有 $z_t = (z_{t,1},...,z_{t,n})R^{n} $，$ n=, D=16 $，其中$ H $和$ W $表示输入图像的高和宽，$ D$表示图像token化的下采样参数；
之后将这些tokens放进 embedding layer ，映射到维空间中；

Text tokens

在每个时刻都会从文本和行为中整合信息，对于文本信息，使用预训练的 T5-large model，然后在每个时间步产生长度为的文本tokens
- Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 2020.
这些tokens也通过线性层映射到维的空间中，表示为

Action tokens

对行为使用个参数进行建模，分别表示速度和曲率；每个变量单独通过线性层映射到维的空间中
因此在时间步时，行为token表示为
分别计算出image、text、action的tokens之后，将三者按照text - image - action的顺序堆叠，可得
然后使用分解的时空位置嵌入
1. 可学习的时间嵌入在给定时间步的所有token之间共享，例如，有T个时间嵌入。
2. 可学习的空间嵌入表示时间步长内的token位置，例如，有个空间嵌入（ text tokens, image tokens, and action tokens），维度

2.2 Image Tokenizer

序列长度是指描述数据所需的离散token的数量。词汇表的大小对应于单个token可能取值的数量。两者之间存在取舍关系；
类似于嵌入矩阵每一行是one-hot的，相同语义信息情况下，列数越多，行数越短，反之亦然；
为了提升效率使用了 discrete image autoencoder，利用VAD的思想，将原始图像转换为抽象的特征，然后再作为输入给到世界模型，主要有两个目标：
1. 压缩来自原始像素的信息以使序列建模问题易于处理。图像包含大量冗余和噪声信息。我们希望减少描述输入数据所需的序列长度。
2. 引导产生有意义的表征，例如语义，而不是高频信号。由此产生的世界模型的输入空间将更容易构建，并且不受可能大大减慢学习过程的高频信号的支配。
3. 参考：Neural discrete representation learning. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
对于输入图像进行了下采样处理，使用 DINO model 进行图像语义信息提取

Increasing semantic content of image tokens through DINO distillation

离散的自动编码器是一个全卷积的2D U-Net结构，获得的图像tokens表示为
训练损失为：
- Image reconstruction loss
  - 该项损失为图像重建损失（image reconstruction loss）的损失、损失、perceptual loss 、GAN loss 的加权和
- Quantization loss
  - 为了更新embedding vectors，使用了embedding loss 和 commitment loss
- Inductive bias loss
  - 图像特征的量化被训练来匹配预训练的 DINO 模型，使用cosine相似度损失，将来自DINO的信息蒸馏到已学习的token中非常重要，因为这使它们能够从该模型的归纳偏差中受益。

2.3 World Model

根据2.1所示，世界模型的输入是，世界模型是一个autoregressive transformer network架构；其训练目标是使用Transformer块的注意力矩阵中的因果mask，根据所有历史的token来预测序列中的下一个图像token

（对于每个时刻，先预测该时刻所有的image tokens，然后再根据历史信息预测下一个时刻的 image token，每一个token是原始图像长宽的1/16）

在训练过程中随机地dropout条件tokens，因此世界模型的生成能力得以加强：
- 实现无条件的生成（unconditional generation）
- 实现行为条件下的生成（action-conditioned generation）
- 实现文本条件下的生成（text-conditioned generation）
为了进一步缩短世界模型的序列长度，对视频进行了时域下采样，从25Hz调整为6.25Hz，这使得世界模型能够进行更长期地解释现象，减少了难以处理的序列长度
为了实现全帧率的视频数据输入，使用了video decoder，将在2.4节中介绍

2.4 Video Decoder

GAIA-1使用了denoising video diffusion models，将问题建模为在扩散过程中对一系列帧进行去噪，模型可以沿时间访问信息，这对于输出视频的时间一致性有很大帮助
参考以下研究并使用一个具有分解的空间和时间注意力层的3D U-Net
- Video diffusion models. In arXiv preprint, 2022.
在训练中，视频扩散模型以从预训练的图像表征器 token化输入图像获得的图像表征为条件；
在推断期间，扩散模型以来自世界模型的预测图像表征为输入条件。
The video decoder loss is

其中，是去噪视频模型；上去噪目标，使用v-parameterization 以提高时间一致性；是采样的离散扩散时间；是长度为的视频序列；表示有噪声的视频，和是关于时间的noise schedule；表示conditioning image tokens序列；是由训练任务指定的图像mask序列

3 Data

4700 hours at 25Hz的驾驶数据，训练过程中筛选了一部分数据以控制数据分布情况

对于 tokenizer，平衡了latitude, longitude, weather category几类数据的量级
对于世界模型和视频扩散模型，平衡了latitude, longitude, weather category, steering behavior category, speed behavior category的量级关系

总结

GAIA-1实现了自动驾驶场景的可控生成，所谓可控包括几个方面：场景、障碍物、车辆行为、天气等，在4700小时的数据集上取得了不错的效果，但模型只输出场景图像，没有BEV和轨迹信息，是个小小的遗憾；此外，文本数据的标注/生成也是个问题；本文的图像、文本等数据的编码和处理均由之前的成熟模型完成，本文更多的是一个数据构造和模型集成的工作；