Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion

We thank the reviewer for their feedback. We are glad that the reviewer finds our approach to be interesting and the presentation of methodology to be clear and helpful. We will be addressing your questions and concerns below.

Q: “The proposed method adopts a VQVAE-like model to capture the complex 3D world, as mentioned in the introduction challenge (i). Classic BEV (the method in BEVFusion, BEVFormer, etc.,) can also realize this ability IMO.”

A: Methods like BEVFusion and BEVFormer are image encoders and tackle supervised object detection; they do not have the ability to generate Lidar point clouds.

Q: “How about the masked image modeling? What will the results or inference time be like if no masked image modeling (or even MaskGIT) is applied?”

A: Masked image modeling is mostly used for representation learning, and often cannot serve as a generative model on its own. The task of world modeling is a generative task. Regarding the connection between masked image modeling and MaskGIT, it is true that after tokenization, MaskGIT is very similar to masked image modeling, but MaskGIT can serve as a generative model.

Q: “The intuition of using different masking strategies for world model training comes from the robotics field. It would be valuable if ablations on this could be presented as well.”

A: We only use two types of attention masks during training, and both are necessary. The masks are either a causal mask, or an identity mask that only allows each frame to attend to itself. The causal attention mask is standard in GPT language models (the original Transformer paper [1] says that they “modify the self-attention sub-layer … to prevent positions from attending to subsequent positions”). The identity (attention) mask is necessary for classifier-free diffusion guidance (CFG). Without it, CFG cannot be applied. To better illustrate how CFG works in our world model and why the second type of mask is necessary, we have included an additional diagram Figure 10 in the Appendix.

Q: “A naive baseline should be simple diffusion modeling with simple BEV features.”

A: Simple BEV features from supervised object detectors such as BEVFusion and BEVFormer cannot generate point clouds. In addition, assuming that our proposed tokenizer is used, other diffusion models such as Gaussian diffusion models are not necessarily simpler, and often require many more diffusion steps than our discrete diffusion model (which only requires 10 diffusion steps) in the absence of additional distillation techniques. Besides, our proposed method for classifier-free diffusion guidance in the world modeling context is likely important regardless of the continuous/discrete nature of the diffusion world model.

Q: “Taking the current action as input is reasonable for a world model, but much information about history could make the prediction rely heavily on the past.”

A: Since a self-driving environment is a partially observable rather than a fully observable environment, conditioning on both past observations and past actions is a reasonable thing to do. The more information from the past that is provided to the world model, the better it can estimate the current state of the world and make predictions about the future. The Transformer can learn to decide which part of information is most relevant or irrelevant, but it can only learn to do so when the necessary information is available.

Q: “As this task implicitly involves ego-planning, this could lead to causal confusion though impressive results are obtained under the open-loop scenario.” “In which paper, the future ego vehicle pose is provided?”

A: For a world model in self-driving scenes, the actions are where the vehicle plans to drive, represented by future ego vehicle poses. The previous state-of-the-art method 4D Occupancy [2] also uses future ego vehicle poses. Since the future ego poses are provided to the world model, the world modeling task alone does not involve ego-planning. We acknowledge that transferring this to a closed-loop setting while mitigating causal confusion is an interesting and open research problem, but the goal of our paper is to provide a more accurate world model that can be learned from unlabeled data, as we believe this is the major current bottleneck.

Q: “Why not report full results under all metrics in ablation study tables?”

A: In our updated draft, we have updated the ablation study tables to include all metrics within the ROI.

Thanks for the clarifications. I have also detailedly gone over other reviewers' comments and authors' feedback. I have some remaining concerns below.

• The motivation of discrete tokenization, in my original review. Maybe I didn't clearly state my thoughts. This point is somehow similar to the concern of Reviewer 2kSt. The method mainly includes three main components besides the neural rendering decoder, VQVAE-like BEV tokenizer, MaskGIT-based discrete diffusion model, and the spatio-temporal world model. My original concern is why use VQVAE-like BEV tokenizer (discrete BEV tokens) instead of continuous BEV features like normal BEV methods (BEVFormer, etc) or even voxel-based methods? The point is continuous vs discrete. I cannot fully get the idea of why they do not have the ability to generate Lidar point clouds. They should behave similarly to the tokenized one for the reconstruction decoder.

  • If thus, an ideal ablation table should be (the last two rows are current Tab.4). They are likely inferior and inefficient, but the comparisons as well as proper motivation in the introduction could make the overall project more solid and convincing: | Method | Metrics | FPS | Others | | -------------- | :-------: | :--------: | :------: | | continuous BEV + naive diffusion | - | - | - | | tokenized BEV + naive diffusion | - | - | - | | tokenized BEV + MaskGIT | - | - | - | | tokenized BEV + MaskGIT-based discreate diffusion | - | - | - |

• I agree with the authors that conditioning on past observations and actions is reasonable for a world model. My concern is that this could lead to an unfair comparison with previous methods, as suggested by Reviewer amws as well. The gain of point cloud forecasting mostly comes from CFG. Without this, the results are close to 4D-Occ.

We thank the reviewer for their feedback. We are glad that the reviewer finds the paper’s contribution to be clear and important, and the writing to be easy-to-follow. We will be addressing your questions and concerns below.

Q: “The metrics for evaluating the performance of the world model may not be reasonable enough.”

A: The metrics we use follow the evaluation protocol from previous point cloud forecasting literature [1]. In addition, our qualitative results clearly show that our model drastically outperforms the previous SOTA on dynamic objects. See the second scene in Figure 5, the third scene in Figure 11, the second and third scene in Figure 15; those examples clearly show that our method can predict the motion of dynamic objects well even in complex scenes, while prior methods are unable to. In fact, our approach seems to be the first point-cloud forecasting method that can make reasonable predictions for dynamic objects without any labels.

Q: “It is unclear to me what the advantage of using such a world model is for MBRL, for example, compared to object segmentation and tracking pipelines.”

A: The main advantage is that it is an unsupervised learning algorithm, so that it can learn from any unlabeled data. The traditional object segmentation and tracking pipelines require annotations and labels to learn, which are very expensive.

Q: “There are many complex modules in the pipeline, including VQ-VAE, diffusion model, neural feature grid, and transformer. Not sure if it is easy to reproduce the results and extend it to other datasets or tasks.”

A: We have included all the technical details in the appendix for reproducibility. We have also added two additional diagrams, Figure 6 and Figure 7, about the architecture of VQ-VAE and spatio-temporal Transformer. It is important to recognize that our approach is conceptually simple: turning sensor observations into discrete tokens, and then applying discrete diffusion. It only has two models: the tokenizer and the world model. The complex designs are about specific architecture choices inside those two models: the neural feature grid is an architecture choice for the tokenizer decoder, and the transformer is an architecture choice for the world model. To extend the approach to other datasets and tasks (beyond self-driving), the specific architectural choices might need to change, but the general approach of tokenization and discrete diffusion are broadly applicable.

Q: “Could the authors elaborate more on ‘using past agent history as CFG conditioning improves world modeling’. What agent history is used here and how is it used? Does it introduce additional knowledge and make the comparison unfair?”

A: To better illustrate how CFG is applied in our world model, we have included an additional diagram Figure 8 in the Appendix. Section 4.3 also covers this part in detail. The agent history means the past observations and action history of the ego vehicle. We assigned a specific symbol to denote agent history: $c^{t-1}$. It is used as classifier-free diffusion guidance. The past observations and action history are already inputs in the world model; they do not require any additional knowledge.

Q: “Since the proposed method introduces a world model, could the authors demonstrate some qualitative examples of different future predictions with different actions?”

A: In our updated draft, we have included visualizations of predicted futures under counterfactual actions in Figure 10 (in the appendix). Here we modify the future trajectories of the ego vehicle (which are action inputs to the world model), and we demonstrate that the world model is able to imagine alternative futures given different actions. The imagined futures are consistent with the counterfactual action inputs.

References:

[1] Khurana, Tarasha, et al. "Point Cloud Forecasting as a Proxy for 4D Occupancy Forecasting." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

We thank the reviewer for their feedback. We are glad that the reviewer finds our proposed method to be interesting and novel, and that the strong empirical performance of our method is recognized. We will be addressing your questions and concerns below.

Q: “The proposed architecture is highly specific to point-cloud occupancy prediction (the novelty lies largely in the encoder and decoder, which are task-specific architectures)”

A: It is true that our tokenizer is highly specific to point clouds. After tokenization, however, the specific 3D representations are abstracted away from the world model. The architecture of the world model is very general-purpose: it is a spatio-temporal Transformer, and both its inputs and outputs are discrete tokens.

Additionally, the concept of Bird-Eye View (BEV) is quite general and goes beyond point clouds, since many popular approaches for camera-based perception such as Lift-Splat-Shoot [1] and BEVFormer [2] also lift camera data to BEV. Therefore, camera data can potentially be tokenized using BEV tokens as well. In this case, the spatio-temporal transformer acting as the world model can remain unchanged.

Q: “It's not entirely clear from the introduction/related work how the proposed method relates to MaskGIT, and although MaskGIT is heavily referenced it is never clearly described.”

A: In our updated draft, we have included a clear and concise description of MaskGIT in the background section. To include it here: “MaskGIT (Chang et al., 2022) has significantly simpler training and sampling procedures based on BERT alone: for training, it masks a part of the input tokens (with an aggressive masking schedule) and then predicts the masked tokens from the rest; for sampling, it iteratively decodes tokens in parallel based on predicted confidence.”

Q: “It would be useful to also present ablations of the other task-specific model components (encoder/decoder, and maybe also the BEV token grouping?) as compared to general-purpose versions”

A: Due to the cost and compute time required to train a tokenizer from scratch, it is too costly to ablate all components in the tokenizer. That being said, the individual components in the tokenizer are mostly off-the-shelf: the encoder follows standard designs in point-cloud based object detection; the vector quantization layer follows VQVAE; the decoder follows NeRF-like differentiable depth rendering already known to work well. Each off-the-shelf component was designed for general-purpose use; the contribution of our tokenizer is to demonstrate that combining the three components together can allow us to tokenize an entire self-driving scene with detailed reconstructions.

Q: “I'm not sure how they could be improved but it's quite difficult to analyze results/see what's changed between two different images, both in Fig. 1 and Fig. 5.”

A: In our updated draft, we have modified Figure 1 to provide a more zoomed-in view. We have also updated Figure 5 to provide zoomed-in orange boxes, which highlight that our method is significantly better in terms of not just the overall structure of the scene but also geometric details, especially at 3 seconds into the future. We will continue to find better ways to visualize the point clouds for the camera-ready version.

Q: “It would be interesting to conduct a more thorough analysis of the modifications to MaskGIT, perhaps on a simpler discrete diffusion problem.”

A: We agree with the reviewer that this would be interesting, but the focus of this work is unsupervised world modeling in large-scale self-driving scenes. Given that our experiment section is centered around point cloud forecasting, we think that perhaps it makes the most sense to leave such investigation to future work.

Q: “Why is L1 median reported inside the ROI but L1 mean reported for the full scene?”

A: As explained in the paper, L1 median within the ROI is a much better metric because (1) the median is more robust to outliers than the mean, and (2) in most self-driving applications predictions within the ROI is what matters for downstream tasks such as planning. For a more direct comparison with the baselines we still computed L1 mean since that is what was reported in their original work.

We thank the reviewer for their feedback. We are glad that the novelty of our method and the significance of our results are recognized, and that the reviewer finds our method section to be clear. We will be addressing your questions and concerns below.

Q: “The unnumbered first equation in Section 3 should be explained better”

A: We have added additional explanations in the updated draft. This equation is the commonly used diffusion objective (Equation 13 in the original diffusion paper [1]).

Q: “It is not fully clear to me what ‘SE(3) ego poses’ mean”

A: We express the ego vehicle actions as ego pose transformations of the ego vehicle during two consecutive time steps, where each action is a 3-D homogeneous transformation matrix consisting of a translation and rotation (i.e., SE(3)). Since the world model is action-conditioned, we need to tell the world model where the agent plans to drive. As the ego vehicle can be assumed to have a rigid body, we do so in the form of rigid transformations. As mentioned in Section 4.4, we simply flatten those 4x4 matrices and feed the 16-dim vector as additional inputs to the spatio-temporal Transformer.

Q: “Figure 1 and 5 might be easier to read if you zoom in on the circled areas.”

A: In our updated draft, we have modified Figure 1 to provide a more zoomed-in view and better viewing angles. We have also updated Figure 5 to provide zoomed-in orange boxes, which highlight that our method is significantly better in terms of not just the overall structure of the scene but also geometric details.

Q: “Are you planning to release your code?”

A: We are unable to release the code, but we have included all the technical details of our model in the appendix. We have also added two additional diagrams, Figure 6 and Figure 7, about our model architectures to improve reproducibility.

Q: “What hardware is required to train your model?”

A: The tokenizer is trained on 16 T4 GPUs; and the world model is trained on 8 A10 GPUs.

References:

[1] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." International conference on machine learning. PMLR, 2015.

Author Response:

We thank the reviewer for their feedback. We are glad that the novelty of our approach and the rigor of our experimental evaluation are recognized, and that the reviewer finds our exposition to be clear and well-structured. We address the questions and concerns below.

Q: “The paper mostly addresses the prediction of near term future states but it is not clear if we go much further, how would the accuracy be?” “How is the result if we predict much further, like 9s?”

A: We answer this question from several perspectives:

• While we agree with the reviewer that long horizon prediction is an interesting problem, we believe finding a good metric to evaluate such capability will require extensive research. Evaluating sensor predictions 9s ahead via reconstruction metrics is not very meaningful. In autonomous driving scenarios, vehicles often cover a considerable distance within 9 seconds, leading to substantial transformations in the scene. Many foreground actors (like vehicles) and background elements (like trees) present in the new scene might not have been visible just 9 seconds earlier. Consequently, observations from 9 seconds prior often become insufficient, if not almost irrelevant, for making accurate future predictions. As such, unsupervised world modeling has a significantly greater degree of uncertainty than supervised object-level trajectory prediction, especially for long-horizon prediction. Note that supervised object-level trajectory prediction systems circumvent this issue by focusing solely on the trajectories of existing foreground actors that are visible at the current scene.
• Generating coherent long-horizon futures will likely require training the model to predict longer than 3s. While our framework allows us to train a model to predict 9s ahead, this will require significantly more compute resources, and is beyond the scope of our work.
• We note that our approach already significantly outperforms existing SOTA in 1s and 3s predictions, which have been the common benchmarks for unsupervised point cloud forecasting. Predicting 9s is very challenging, and there are no baselines to compare to, but we will consider it for future work.

Q: “with a diffusion model, the inference could be slow”

A: It is true that our current model requires 10 diffusion steps for generating each frame. However, we also point out the following:

• Inference speed is not our current focus in this work. The goal of our work is to showcase what results can be obtained with a discrete diffusion world model. As a pioneering work studying how to apply diffusion to unsupervised world modeling on real-world data, our model can already achieve good results using a reasonably small number of diffusion steps (10 steps).
• The field is actively working on improving the speed of diffusion models. Until recently, the field has primarily focused on Gaussian diffusion models such as DDPM [1], which initially required 1000 diffusion steps per sample; later on, DDIM [2] required 50 steps, progressive distillation [3] required around 8 steps, and Consistency Models [4] only required 1-2 diffusion steps. Even though discrete diffusion has not yet received as much attention from the field, we expect that similar distillation techniques can likely be applied in the future.

Q: “could you also share some insights on how this model could exactly be integrated within modern self-driving systems and work with other modules such as planning”

A: With a world model, the problem of autonomy can be seen as a model-based reinforcement learning (MBRL) problem. The autonomy system is an RL agent taking actions in the world. Previously, to apply MBRL, the bottleneck was about whether such a world model can be learnt on real-world observations in the self-driving scenes. Our work is a contribution towards removing this bottleneck. With an effective world model at hand, research can start to apply MBRL approaches such as Dreamer-v2 [5] to both learn a driving policy during training and plan with model-predictive control (MPC) at test time. While Dreamer-v2 assumes a given reward function, we can learn such a reward function from real-world driving data via techniques such as GAIL [6]. This is just one example of how this model can be integrated with modern autonomy systems.

Thanks for clarifying some of my questions. I think the notion of using a discrete diffusion process is indeed novel and improves the prediction accuracy nontrivially, so I keep my rating. But predicting 9s is still meaningful, though the author claimed they don't have the computing resources to do that which is reasonable. However, that does not mean that predicting 9s or some horizon longer than 3s isn't important. Since the vehicle speed could vary, in some cases, horizon longer than 3s is still important to predict to get a sense how far the vehicle is able to predict without making too much errors. There are indeed industry models that can predict long horizon trajectories, though some are not open sourced, which may make it hard to compare with.