Paper: Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion (ICLR-2024)

Reviewer comment (Experiments):
Table 4 presents the ablation of the discrete diffusion algorithm. The motivation to use MaskGIT seems its parallel decoding strategy. 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?

Author response:
*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.

Reviewer comment (Experiments):
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.

Author response:
*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.

Reviewer comment (Experiments):
In light of the above points, a naive baseline should be simple diffusion modeling with simple BEV features.

Author response:
*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.

Reviewer comment (Theory):
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. If the prediction horizon is longer, such a long history is also weird as poses at the very beginning are intuitively unhelpful.

Author response:
*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: "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."*

A: We have explained to Reviewer amws that the past observations and action history are already inputs to the world model; they do not require any additional knowledge, and therefore do not lead to any unfair comparison.

Reviewer comment (Experiments):
As this task implicitly involves ego-planning, this could lead to causal confusion though impressive results are obtained under the open-loop scenario. The authors have stated that combining the world modeling approach with model-based RL is a future direction, yet, it is important to demonstrate its effectiveness for the decision-making task.

Author response:
*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: “The authors have stated that combining the world modeling approach with model-based RL is a future direction, yet, it is important to demonstrate its effectiveness for the decision-making task”*

A: To apply model-based RL (MBRL) approaches to a domain, a prerequisite is that a reasonably good world model can be learned in this domain. Our proposed model represents a leap in the fidelity of unsupervised world models in the self-driving domain, which unlocks the possibility of doing model-based RL on raw driving data. In other words, our work is a necessary step before MBRL can be applied. Adding MBRL introduces additional complexities and is beyond the scope of one paper; as such, we leave this for future work.

Reviewer comment (Writing):
In the introduction, the task definition of point cloud forecasting is 'to predict future point cloud observations given past observations and future ego vehicle poses'. In which paper, the future ego vehicle pose is provided?

Author response:
*Q: “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.

Reviewer comment (Evaluation):
Why not report full results under all metrics in ablation study tables?

Author response:
*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.

Reviewer comment (Novelty):
The conclusion says “One particularly exciting aspect of our approach is that it is broadly applicable to many domains. We hope that future work will combine our world modeling approach with model-based reinforcement learning to improve the decision-making capabilities of autonomous agents.” 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.

Author response:
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.

Reviewer comment (Reproducibility):
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.

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Presentation):
The introduction/related work are somewhat intermixed, which is fine, but leads to confusing presentation. 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. Background of discrete diffusion could also be better described.

Author response:
*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.”

Reviewer comment (Experiments):
Ablations of differences to MaskGIT are good, but 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.

Author response:
*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.

Reviewer comment (Presentation):
Minor note: point cloud prediction visualizations are a little difficult to parse – 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.

Author response:
*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.

Reviewer comment (Experiments):
It would be interesting to conduct a more thorough analysis of the modifications to MaskGIT, perhaps on a simpler discrete diffusion problem.

Author response:
*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.

Reviewer comment (Evaluation):
Why is L1 median reported inside the ROI but L1 mean reported for the full scene?

Author response:
*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.

Reviewer comment (Experiments):
It's mentioned that the model is relatively small; how much do the results change when scaling the model (up/down)?

Author response:
*Q: “It's mentioned that the model is relatively small; how much do the results change when scaling the model (up/down)?”*

A: While there are ways to further scale our Transformer world model, studying the scaling properties and scaling laws are beyond the scope of this paper, as it requires additional dedicated time and resources than what is feasible during this rebuttal period.

Reviewer comment (Writing):
The unnumbered first equation in Section 3 should be explained better.

Author response:
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]).

Reviewer comment (Presentation):
Figure 1 and 5 might be easier to read if you zoom in on the circled areas.

Author response:
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.

Reviewer comment (Reproducibility):
What hardware is required to train your model?

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