【论文笔记】CARFF: Conditional Auto-encoded Radiance Field for 3D Scene Forecasting

原文链接：https://browse.arxiv.org/abs/2401.18075

1. 引言

人类可以从部分视觉上下文中想象不能看到的部分（物体的存在与位置，以及场景与物体的形状、颜色、纹理等），这对安全决策至关重要。而自动驾驶系统的传统方法是将传感器输入转化为被跟踪物体的位置、速度等信息，通常会基于最坏的假设（可视区域边界存在高速运动的车辆）进行规划。

NeRF可以从无遮挡视角生成新视图，从而包含目标检测器可能丢失的重要视觉信息；此外，其使用隐式的密度图表达显式几何，可以无需渲染就进行运动规划。但使用NeRF进行遮挡处的概率预测是很困难的。本文提出CRAFF，第一个使用场景的部分观测进行3D表达随机预测，并有效整合视觉感知与显式几何的方法。

CARFF提出PC-VAE（姿态条件变分自编码器），一种基于卷积和ViT的图像编码器，将部分观测的图像映射为潜在场景表达，包含带隐式概率的状态信念（belief）。此后会训练神经辐射场作为3D解码器，恢复任意视角下的3D场景。此外，还引入混合密度模型，在编码器的信念空间中随机预测3D场景随时间的变化。

3. 方法

本文的主要目标是在动态环境中建立模型进行3D场景的概率预测。但直接整合NeRF的3D场景表达与概率模型（如VAE）会导致非凸和相互影响的优化，从而导致训练不稳定。因此，本文将训练分为两个阶段（如下图所示），即首先训练PC-VAE得到视图不变的场景表达，再将解码器替换为NeRF从隐空间学习3D场景。这样，隐空间表达捕捉了环境的可能的状态和动态，而NeRF则在信念空间生成新视图，从而使未观测区域可见。预测时，不确定性可建模为从预测的高斯混合模型中自回归采样隐变量，从而进行高效决策。因此，本文将场景预测视为隐分布上部分可观测的马尔科夫决策过程（POMDP），在不确定性中捕捉多模态信念。
在这里插入图片描述

3.1 NeRF姿态条件VAE（PC-VAE）

结构：给定时间 $t$ 下的场景 $S_t$ ，从相机姿态 $c$ 捕获以自车为中心的观测图像 $I_c^t$ 。目标是形成图像的3D表达，从而进行未来场景预测。本文使用以隐变量 $z$ 为条件的神经辐射场，隐变量从后验分布 $q_\phi(z|I_c^t)$ 中采样。本文使用PC-VAE学习后验分布。使用卷积与在ImageNet上预训练的ViT建立编码器，学习将图像空间映射为高斯分布的隐空间 $q_\phi(z|I_c^t)=\mathcal N(\mu,\sigma^2)$ ，其中 $\mu$ 为均值， $\sigma^2$ 为方差。解码器 $p (I ∣ z, c)$ 以相机姿态为条件，将隐变量 $z\sim\mathcal N(\mu,\sigma^2)$ 映射到图像空间得到 $I$ 。这使编码器生成与相机姿态无关的隐变量。

为进行3D场景建模，使用Instant-NGP，使用哈希编码和占用网格增强计算效率。使用小型MLP $F_\theta(z)$ 建模密度与外观：
$F_\theta(z):(x,d,z)\rightarrow((r,g,b),\sigma)$

其中 $x\in\mathbb{R}^3$ 和 $d=(\theta,\phi)$ 分别表示位置向量和视角。

训练方法：本文使用两阶段训练流程以稳定训练。首先在语义丰富的2D像素空间优化基于卷积ViT的编码器和姿态条件卷积解码器，以处理复杂而真实的场景。通过将解码器限制在相机姿态上，可以实现相机视角与场景上下文的解耦，使得表达是视角不变的、编码器是3D感知的。然后，将解码器替换为以隐变量为条件的NeRF，并冻结编码器，在3D空间中重建编码器的信念，并进行新视图合成。

损失：使用标准VAE损失，即均方误差（MSE）和给定证据下界（ELBO）的KL散度（KLD）：
$\mathcal{L}_{PC-VAE}=\mathcal L_{MSE,PC-VAE}+\mathcal L_{KLD}=\|p(I|z,c'')-I_{c''}^t\|^2+\mathbb{E}_{q(z|I_c^t)}[\log p(I|z)]-w_{KL}D_{KL}(q_{\phi}(z|I_c^t)||p(I|z))$

其中 $w_{KL}$ 为KL散度损失权重， $z\sim q_\phi(z|I_c^t)$ 。为使编码器是3D感知的，后验使用相机姿态 $c$ 编码，而解码器以随机采样的姿态 $c^{''}$ 为条件。

KL散度正则化隐空间，以平衡条件重建和遮挡下的随机性。增大 $w_{KL}$ 可以使隐变量趋向正态分布 $\mathcal{N}(0,1)$ ，以保证部分观测下的概率采样。但过度的正则化会使隐变量不可分离，从而使模式崩溃。本文使用延迟线性KL散度损失权重调度，以得到平衡的权重。

然后在VAE的后验上使用基于NeRF的解码器以建模场景。在每个时刻 $t$ ，使用逐像素的MSE损失训练NeRF：
$\mathcal L_{MSE,NeRF}=\|I_c^t-render(F_\theta(\cdot|q_\phi(z|I_c^t)))\|^2$

使用标准的渲染算法。最后，在姿态条件编码器学习到的隐空间上使用预测模型。

3.2 场景预测

问题表达：为预测未来的场景，需要预测未来的隐变量分布。本文将预测视为在PC-VAE隐空间中后验分布 $q_\phi(z|I_c^t)$ 上部分可观测的马尔科夫决策过程（POMDP）。

推断时，学习代表不同场景可能性的混合高斯分布。使用混合密度网络（MDN）建模POMDP，用多头MLP预测 $K$ 个高斯的混合。每个时刻 $t$ 下的分布为：
$q'_\phi(z_t|I_c^{t-1})=MDN(q_\phi(z_{t-1}|I_c^{t-1}))$

模型以后验分布 $q_\phi(z_{t-1})$ 为条件，预测每个时刻的后验分布 $q_\phi(z_t|I_c^{t-1})$ 。预测的后验分布由高斯的混合给出：
$q'_\phi(z_t)=\sum_{i=1}^K\pi_i\mathcal{N}(\mu_i,\sigma_i^2)$

其中 $\pi_i,\mu_i,\sigma_i^2$ 为第 $i$ 个高斯的权重、均值和方差，均以后验 $q_\phi(z_{t-1})$ 和采样隐变量 $z_{t-1}$ 为条件。从混合高斯 $q'_\phi(z_t)$ 中采样 $z_t$ ，其结果最可能落在某个高斯模式内。通过NeRF渲染，可以反映该模式对应的场景。

损失：通过最小化负对数似然函数优化MDN：
$\mathcal L_{MDN}=-\sum_{i=1}^N\log(\sum_{j=1}^K\pi_j\mathcal{N}(y_i;\mu_j,\sigma_j^2))$

其中 $y_i\sim q_\phi(z_t)$ 从编码器学到的 $z_t$ 的分布中采样， $N$ 为样本数。

在这里插入图片描述
推断：如上图所示。将图像通过编码器得到后验 $q_\phi(z_t)$ ，然后预测未来可能的后验分布 $q'_\phi(z_{t+1})$ ，采样场景隐变量并进行规划。这可通过(a)NeRF的密度探测或(b)使用现成的方法如YOLO分割渲染的新视图图像（如下图所示）实现。这样，可以得到编码器隐空间中的高斯分布 $q_\phi(z_{t+1})$ ，并自回归地送回MDN预测下一时刻。
在这里插入图片描述

4. 结果

本文使用CARLA驾驶模拟器，并设计基于CARFF的控制器，进行下游规划任务。

4.2 CARFF评估

PC-VAE性能和消融：实验表明，PC-VAE在给定部分观测的情况下能有效重建复杂环境。使用PSNR进行定量评估合成图像与真实图像的差异。

为衡量编码器生成的隐空间质量，还使用t-SNE可视化每张图像的隐变量样本分布，并使用基于SVM的指标定量衡量可视化簇，其值越高，关于时间的簇越可分，从而证明隐空间是视图不变的。错误分类样本或边界样本通常表达部分或完全遮挡的区域，可以对这些样本建模概率行为，并进行预测。结果表明，本文的PC-VAE比其余方法有优势，特别对于无条件VAE而言，因其只重建图像，而不捕捉3D几何。其余方法需要仔细微调超参数，且对相机剧烈运动的泛化能力较差。

因此，本文的方法可以通过隐变量表达保留随机特性，且同时能进行精确重建。

3D新视图合成：给定部分观测，本文方法可以保留所有可能的当前状态信念，并为每个信念生成任意视角下的新视图。

在完整和部分观测下的推断：完整观测下，CARFF可以根据参与者定位结果，生成稳定的信念，而避免不正确的预测；部分观测下，CARFF可以预测多种可能情况。