D2-World: An Efficient World Model through Decoupled Dynamic Flow (2024)

2.1 Stage I: Vision-based Occupancy Prediction

In this section, we introduce the architecture of the occupancy network, which takes visual images as input and predicts the occupancy state and semantics for a single frame.

Image Encoder. The image encoder is designed to encode the input multi-camera 2D images into high-level features. Our image encoder comprises a backbone for high-level feature extraction and a neck for multi-resolution feature aggregation. By default, we use the classical ImageNet pre-trained ResNet-50 as the backbone in ablation studies, and Swin-Transformer-B[8] as the backbone for submission. Although employing a stronger image backbone can enhance prediction performance, we considered the trade-offs between resource usage and training time, and ultimately decided against using huge backbones such as InternImage-XL[11].

View Transformation.We utilize LSS[4] for view transformation, which densely predicts the depth of each pixel through a classification method, allowing us to project the image features into 3D space. Moreover, to introduce the temporal information in our model, we adopt the technique proposed in [6], dynamically warping and fusing one historical volume feature to produce a fused feature.

Ocupancy Head. We adopt the semantic scene completion module proposed in [12] as our occupancy head, which contains several 3D convolutional blocks to learn a local geometric representation. The features from different blocks are concatenated to aggregate information. Finally, a linear projection is utilized to map the features into $C_{0}$ dimensions, where $C_{0}$ is the number of classes.

Losses.To alleviate the class-imbalance issue in occupancy prediction, we utilize class-weighted cross-entropy and Lovasz losses. Our multi-task training losses are a combination of occupancy prediction loss and depth loss.

2.2 Stage II: 4D Occupancy Forecasting

In this section, we introduce the process of future point cloud forecasting. The framework consists of an occupancy encoder, a flow decoder, flow guided warping and refine, and a rendering process.

Initially, the 3D occupancy data is preprocessed into spacetime tokens. The spatial-temporal transformer effectively captures the spatial structures and local spatiotemporal dependencies within these tokens.Following the encoding of historical tokens, the flow decoder is employed to predict future flow in each voxel grid. Then, warping and refinement generate the final occupancy density.To fully leverage the temporal information across the entire sequence, we utilize a non-autoregressive approach for decoding, which achieves impressive forecasting performance alongside high efficiency.Finally, a differentiable volume rendering process is used to generate the point cloud from the predicted occupancy.

3D Occupancy Encoding.Given a sequence of historically observed $N_{h}$ frames 3D occupancy $\mathcal{O}_{T-N_{h}:T}$, where each occupancy $\mathcal{O}_{i}\in\mathbb{R}^{H_{0}\times W_{0}\times D_{0}}$, we first encode the occupancy sequence into spacetime tokens.Here, $H_{0}$, $W_{0}$ and $D_{0}$ represent the resolution of the surrounding space centered on the ego car. Each voxel is assigned as one of $C_{0}$ classes, denoting whether it is occupied and which semantic category it is occupied with.

To reduce the computational burden, we transform the 3D occupancy in the BEV representation.Take a single-frame occupancy as an example, it first uses a learnable class embedding to map the 3D occupancy into occupancy embedding $\hat{\mathbf{y}}\in\mathbb{R}^{H_{0}\times W_{0}\times D_{0}\times C}$.Then, it reshapes the 3D occupancy embedding along the height dimension to obtain a BEV representation $\mathbf{\tilde{y}}\in\mathbb{R}^{H_{0}\times W_{0}\times DC}$.The BEV embedding then is decomposed into non-overlapping 2D patches $\mathbf{y}_{p}\in\mathbb{R}^{H\times W\times C^{\prime}}$, where $H=H_{0}/P$, $W=W_{0}/P$, $C^{\prime}=P^{2}\cdot C_{0}$, and $P$ is the resolution of each image patch.After that, a lightweight encoder composed of several 2D convolution layers, i.e., Conv2d-GroupNorm-SiLU, is followed to extract the patch embeddings. After considering the sequence of patch embeddings, we obtain the historical occupancy spacetime tokens $\mathbf{{y}}\in\mathbb{R}^{N_{h}\times H\times W\times C}$.

Ego Pose Encoding. We represent the ego pose as relative displacements between adjacent frames in the 2D ground plane. Given the historical ego poses, we employ multiple linear layers followed by a ReLU activation function to obtain the ego tokens$\mathbf{e}\in\mathbb{R}^{N_{h}\times C}$.

Spatial-Temporal Transformer.The spatial-temporal transformer jointly models the evolution of the surrounding scene and plans the future trajectory of the ego vehicle. Inspired by previous works on video prediction[10, 9, 3], we incorporate several spatial-aware local-temporal (SALT) attention blocks within the Spatial-Temporal Transformer. As shown in Fig.2(a), in each SALT block, 2D convolution layers are first utilized to generate the query map and paired key-value embeddings for the spacetime tokens, effectively preserving structural information through this spatial-aware CNN operation.Subsequently, the standard multi-head attention mechanism is employed to capture the temporal correlations between tokens. This approach allows for the learning of temporal correlations while preserving the spatial information of the sequence.Furthermore, we replace the traditional feed-forward network (FFN) layer with a 3D convolutional neural network (3DCNN) to introduce local temporal clues for enhanced sequential modeling.

Decoupled Dynamic Flow.As illustrated in Fig.1 and Fig.2(b), we design a decoupled dynamic flow to simplify the occupancy forecasting problem.Specifically, the flow decoder—which comprises multiple stacked SALT blocks—processes the encoded historical BEV features and forecasts the absolute future flows with respect to the current ego coordinate.Utilizing the occupancy semantics, we decouple the dynamic and static grids, forecasting the future voxel features via the warping operation.For the dynamic voxels, we transform the absolute flow for each future timestamp using the future ego poses, ensuring alignment with the current frame.For the static ones, we directly transform them through future ego poses.Finally, we apply a refinement module composed of several simple CNNs to enhance the coarse warped features.

Rendering & Losses.We utilize the same rendering process and losses as ViDAR[13] for optimizing the point cloud forecasting, which is a ray-wise cross-entropy loss to maximize the response of points along its corresponding ray. For pose regression, we use L1 loss during the training.

3.1 Experimental Setups

Dataset.We conduct our experiments on the OpenScene dataset[2], which is derived from the nuPlan dataset[1].Due to some scenes in OpenScene lacking corresponding occupancy labels, we ignore these scenes during our experiments.For submission, the challenge utilizes an online server that provides historical images along with normalized ray directions for point forecasting.

Metric.For this challenge, model evaluation is conducted using the Chamfer Distance (CD)[5]. The Chamfer Distance quantifies the similarity between predicted and ground-truth point clouds by computing the average nearest-neighbor distance from points in one set to those in the other set, in both directions.

Training Strategies.During the training process, both stages are trained with AdamW optimizer with gradient clipping and a cyclic learning rate policy. The initial learning rates a $2e^{-4}$ and $1e^{-3}$ for stage I and stage II, respectively.In stage I, we utilize a total batch size of 24, distributed across 24 NVIDIA V100 GPUs. In stage II, the total batch size is reduced to 16, leveraging 16 NVIDIA V100 GPUs.For the ablation studies, stage II is trained using 8 NVIDIA V100 GPUs with a total batch size of 8.Both stages are trained for 24 epochs.

Network Details. For stage I, the input image resolution is $512\times 1408$ incorporating common data augmentation techniques such as flipping and rotation, applied to both the images and the 3D space.The resolution of the generated 3D voxel grid is $200\times 200\times 16$.Prior to feeding the predicted occupancy into stage II, we apply grid sampling operations to align the occupancy annotations from the range of [-50m, -50m, -4m, 50m, 50m, 4m] to the LiDAR point cloud range of [-51.2m, -51.2m, -5.0m, 51.2m, 51.2m, 3.0m].

3.2 Quantitative Results

Main Results & Ablation Study. The main results are presented in Tab.1.In addition to showing the overall performance of our model (D²-World), we also demonstrate the performance of our model without decoupled dynamic flow (D²-World vanilla).Our method demonstrates superior performance across all timestamps when compared to the baseline model, with further performance enhancements observed upon the introduction of the decoupled dynamic flow. Our best submission ranks 2nd on the leaderboard, achieving a Chamfer Distance (CD) of 0.79, with both stages trained on the full dataset.

Training Efficiency. To further validate the efficiency of our approach, we compare training hours and GPU memory usage across different models, as shown in Tab.2. The baseline method, ViDAR, requires up to 63 GB of GPU memory and 23.50 hours for training. Even its efficient version[13], which does not supervise all future frames, still demands high GPU memory (38 GB) and considerable training time (18.5 hours).In contrast, although our method necessitates pre-training an occupancy prediction model, our world model can be trained in approximately 3 hours with only 28 GB of GPU memory under the same conditions. Additionally, our model, even with the decoupled dynamic flow, maintains reasonable training hours and GPU memory.

The Effects of Occupancy Performance.The results using different occupancy performances are presented in Tab.3, where only 1/8 mini dataset are used to train.We first train our world model with binary occupancy prediction (empty and occupied) as inputs. The results from Version A to Version E denote the performance of the world model when the occupancy performance changes.We find that the world model performs better when the occupancy performance is improved.

Furthermore, introducing decoupled dynamic flow with semantic occupancy inputs yields additional performance enhancements, as shown in Versions F to H.Interestingly, the performance does not significantly improve even when ground truth occupancy with 100% mIoU and IoU is used as input. Our analysis indicates that this is due to the inherently sparse nature of point cloud forecasting, which primarily requires predicting the foremost visible surfaces of objects in the 3D space, whereas IoU evaluation for occupancy encompasses the entire dense space.