Paper: Large Trajectory Models are Scalable Motion Predictors and Planners (ICLR-2024)

Reviewer comment (Experiments):
The story is well told, but the experimental support is weak, some key ablations are missing.

Author response:
Thank you for the time to read and review our paper. We find your review very constructive!

We have rewritten the abstract, introduction, and some other parts of the paper to focus on scalability for better clarity. We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. If you have any further concerns, we would be keen to address them. We will incorporate all of these changes in the final revision.

# General comments
1. **Re Ablations on the decoders:** We added a new section 5.5, and Table 3 for the ablation study on the Key Points and the diffusion decoders.
2. **Re NuPlan CLS:** We measured the STR(CKS)-16m without diffusion decoders on two CLS metrics. The scores are higher than PDM-Open and lower than PDM-Hybird. We are exploring different flavors of post-procession on the Key Points for better and more stable performance across different test sets. We will include these scores (both CLS-Interactive and CLS-NonInteractive) in the Appendix in the final version.
3. **Re Not significantly better than MTR-e2e on all metrics:** As you might notice, we used the STR(CPS)-16m version to compare. We are working on larger models and will include a better performance (on minADE and minFDE) in the final version.

Reviewer comment (Experiments):
There is a lack of ablation studies for the model design, such as quantifying the improvement brought about by the proposed keypoints module, or comparing the GMM decoder versus the diffusion decoder.

Author response:
Thank you for the time to read and review our paper. We find your review very constructive!

We have rewritten the abstract, introduction, and some other parts of the paper to focus on scalability for better clarity. We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. If you have any further concerns, we would be keen to address them. We will incorporate all of these changes in the final revision.

# General comments
1. **Re Ablations on the decoders:** We added a new section 5.5, and Table 3 for the ablation study on the Key Points and the diffusion decoders.
2. **Re NuPlan CLS:** We measured the STR(CKS)-16m without diffusion decoders on two CLS metrics. The scores are higher than PDM-Open and lower than PDM-Hybird. We are exploring different flavors of post-procession on the Key Points for better and more stable performance across different test sets. We will include these scores (both CLS-Interactive and CLS-NonInteractive) in the Appendix in the final version.
3. **Re Not significantly better than MTR-e2e on all metrics:** As you might notice, we used the STR(CPS)-16m version to compare. We are working on larger models and will include a better performance (on minADE and minFDE) in the final version.

Reviewer comment (Experiments):
For the planning module, there is a lack of closed-loop evaluation, which is crucial for the planning component.

Author response:
We implement quick additional scaling experiments on 20% of the training set on WOMD with similar model sizes as we did for the NuPlan dataset. This additional figure indicates similar strong scalability on the evaluation loss during training. 

Figure link: http://180.167.251.46:880/ICLR_Submission5908/WOMD_quick_scale.jpg

As for CLS on NuPlan, we made clear arguments in **Section 5.1, paragraph Metrics for Planning** that CLS is a notable misalignment metric as discussed in the PDM [1] paper. Could you please help us understand why this misaligned metric is so important for us to prove our model's scalability? We think Closed-Loop Scores are clear "better-to-have" results but not compulsory in the scope of this paper.

[1] Daniel Dauner, Marcel Hallgarten, Andreas Geiger, and Kashyap Chitta. Parting with misconceptions about learning-based vehicle motion planning. In Conference on Robot Learning (CoRL), 2023.

Reviewer comment (Experiments):
Why haven't the closed-loop results on NuPlan been included? Since you've already integrated the open-loop evaluation with NuPlan, conducting a closed-loop test should be quite straightforward. Is it because the results were not satisfactory?

Author response:
Thank you for the constructive advice and clarification. 

A full scaling analysis like Table 1 on WOMD with other prediction models is interesting but expensive considering the number of experiments. We humbly question how much the value of this analysis would have because:
1. We demonstrate the scalability of STR on WOMD in the previous eval loss figure. 
2. We pointed out that MotionLM (the only previous method that did scalability tests) does not scale like STR.
We can at least provide the converged loss and mAP/minADE/minFDE results of some other open-sourced methods on some scales in the final version for comparison. 

Other potential relevant questions:
1. Does eval loss represent scalability or do we need to evaluate on mAP/minADE/minFDE? As discussed previously with reviewer 2D3R, we hope reviewers understand why we chose evaluation loss to measure the scalability rather than mAP/minADE/minFDE in Figure 1. We want to point the reviewers to revisit the scaling law analysis on LLM [1]. They did all measurements on evaluation and test loss for a reason. 
2. Scalablity automatically brings SOTA performance (like better than MTR++?) on all metrics on WOMD? As long as our results indicate, this is not the case. Selections of priors (like intention points of MTR) do matter. Post-processions (like NMS) do matter. Ensembling models does matter. All these tricks affect how well our model scales and the final mAP/minADE/minFDE results. For now, there is still a gap and we are exploring which trick is still necessary after we scale. And that is beyond the scope of this paper. 

For a fair comparison on CLS, we need to adopt our method into the NuPlan planner API and run expensive closed-loop simulations and those take time. We will make sure to include these results and experiment settings with details in our final version.

[1] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.

Reviewer comment (Presentation):
The method description is not well organized and some details are missing, e.g., the proposal, key points and $O_t$ in Fig. 1. Besides, the introduction of the conditional diffusion (Sec 3.2) has no tight logic connection with the method (Sec 4). The only place that mentioned diffusion model in Sec 4 is "STR is compatible with diverse decoders like Gaussian Mixture Model or diffusion decoders."

Author response:
Thank you for the time to read and review our paper. 

We simplified the structure and focused on our core insight, which is sequence modeling brings scalability and solves many previous challenges. The design choices of the encoders and decoders are not important in this work. We have rewritten most of the abstract and the introduction to focus on this insight. Please see our uploaded revision. 

We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. You can easily check all implementation details about data preprocessing, encoders, and decoders in our code.

We added another section to the related works about previous diffusion decoders on motion planning and motion prediction as an introduction before Section 3.2 and Section 4 for a better continuity flow on writing.

Reviewer comment (Novelty):
The novelty of this work. Though it is interesting to apply transformer and diffusion on motion planning/prediction tasks, it remains ambiguous as to the unique contributions or innovations introduced by this work in adapting the diffusion model to this specific domain.

Author response:
Thank you for the time to read and review our paper. 

We simplified the structure and focused on our core insight, which is sequence modeling brings scalability and solves many previous challenges. The design choices of the encoders and decoders are not important in this work. We have rewritten most of the abstract and the introduction to focus on this insight. Please see our uploaded revision. 

We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. You can easily check all implementation details about data preprocessing, encoders, and decoders in our code.

We added another section to the related works about previous diffusion decoders on motion planning and motion prediction as an introduction before Section 3.2 and Section 4 for a better continuity flow on writing.

Reviewer comment (Novelty):
It is not new to adopt an encoder-decoder architecture to take the context as input, and predict the future states, which has been studied by MPNet [R1]. STR uses a different but existing neural architecture which seemingly lacks new contributions.

Author response:
**Re novelty compared with MPNet:** As we pointed out in the abstract and introduction, learning from large-scale real-world demonstrations brings non-trivial additional challenges compared to [1] which seems to be a more relevant work than MPNet. Specifically, MPNet learns from a small set of RRT* planning results. Additionally, MPNet does not consider other agents (or dynamic environments) at all. Specifically, MPNet does not provide the environment's history, including the past trajectories of other agents and the traffic lights for the motion prediction and planning tasks.

Reviewer comment (Reproducibility):
What is the data format of the context? Do you need to convert it into a set of vectors?

Author response:
We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. You can easily check all implementation details about data preprocessing, encoders, and decoders in our code.

Reviewer comment (Experiments):
As in Section 5.2, you have different choices of encoders for different datasets, what is the concern here? Is there any guidance on how to select the proper encoder when facing a different dataset?

Author response:
We simplified the structure and focused on our core insight, which is sequence modeling brings scalability and solves many previous challenges. The design choices of the encoders and decoders are not important in this work. We have rewritten most of the abstract and the introduction to focus on this insight. Please see our uploaded revision. 

We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. You can easily check all implementation details about data preprocessing, encoders, and decoders in our code.

Reviewer comment (Experiments):
Similar to above questions, Sec 4 mentioned that "STR is compatible with diverse decoders like Gaussian Mixture Model or diffusion decoders." Do you have any thoughts on how to select the proper decoders and what are their pros/cons?

Author response:
We simplified the structure and focused on our core insight, which is sequence modeling brings scalability and solves many previous challenges. The design choices of the encoders and decoders are not important in this work. We have rewritten most of the abstract and the introduction to focus on this insight. Please see our uploaded revision. 

We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproduction. You can easily check all implementation details about data preprocessing, encoders, and decoders in our code.

Reviewer comment (Writing):
The described model is too general. Instead of describing what was actually implemented, the authors write about what *could* have been implemented. This is partially because the authors are trying to describe their work as one model on two different datasets; in reality, they have implemented two *different* albeit related models, customized towards the two datasets they work on. The description should be reworked to be as clear as possible about what was implemented rather than describing possibilities.

Author response:
1. **Re The described model is too general:** We understand your concern. As previously addressed, the study of scalability is the core of this paper. We intentionally test different encoders on different datasets to test if the scalability holds. For better clarity, we deleted other potentials (goal points or heat map) for the Proposal from Section 4 - paragraph Proposal. We also deleted the GMM decoder description for the Key Points from Section 4 - paragraph Key Points. Changes are marked in red in the revised version.

Reviewer comment (Reproducibility):
It is not important what the proposals *could* be; what is important is what was actually implemented in this work. For example, statements like "These Proposals can be intention points (Shi et al., 2022), goal points (Gu et al., 2021), or endpoints heat maps (Gilles et al., 2021)" should instead clearly state what was implemented in this paper.

Author response:
1. **Re The described model is too general:** We understand your concern. As previously addressed, the study of scalability is the core of this paper. We intentionally test different encoders on different datasets to test if the scalability holds. For better clarity, we deleted other potentials (goal points or heat map) for the Proposal from Section 4 - paragraph Proposal. We also deleted the GMM decoder description for the Key Points from Section 4 - paragraph Key Points. Changes are marked in red in the revised version.

Reviewer comment (Experiments):
Statements such as "For each Key Point, STR is compatible with diverse decoders like Gaussian Mixture Model (GMM) decoders (Chai et al., 2019), or diffusion decoders to produce K different predictions" are not relevant if those components were not actually implemented. The description should focus only on what was implemented and evaluated.

Author response:
1. **Re The described model is too general:** We understand your concern. As previously addressed, the study of scalability is the core of this paper. We intentionally test different encoders on different datasets to test if the scalability holds. For better clarity, we deleted other potentials (goal points or heat map) for the Proposal from Section 4 - paragraph Proposal. We also deleted the GMM decoder description for the Key Points from Section 4 - paragraph Key Points. Changes are marked in red in the revised version.

Reviewer comment (Reproducibility):
Remove diffusion models from the paper. The diffusion models are tangential to the main point about scaling, and they add significant complexity that will make the work difficult to reproduce and evaluate.

Author response:
We simplified the structure and focused on our core insight, which is sequence modeling brings scalability, solving many previous challenges. We have rewritten most of the abstract and the introduction to focus on this insight. Please see our uploaded revision. 

We also attached our full code (with readme and instructions), which will be open-sourced with pre-trained checkpoints for easy reproducing. Since the scalability comes from the formulation (Note the results in Fig. 2 do not involve the diffusion decoders), readers can easily train a fairly good model without the diffusion decoder’s two-stage training (See new ablation section 5.5).

Reviewer comment (Experiments):
Consider separating the model used for NuPlan and the model used for WOMD, and describe those models separately.

Author response:
2. **Why do we choose different encoders and decoders for each Dataset?** Scalability is the core of this paper not the design choice of the encoders or decoders for each dataset. To clarify this insight, we rewrote the abstract and the introduction. We observe scalability on both NuPlan and WOMD with both raster and vector encoders. Note adapting additional encoders to a large-scale dataset is a non-trivial task. We will add a vectorized encoder (adapted from the PDM) for NuPlan and a rasterized encoder for WOMD along with their experiment results in the final version.

Reviewer comment (Reproducibility):
Describe the split of parameters between the encoders and the transformer.

Author response:
3. **Re split of parameters:** As introduced in section 5.3, the number of trainable parameters is only from the backbone transformers. Hence, the number of parameters of the encoders is less relevant to the study of the scaling. For your reference, the detailed number of the trainable parameters of each part in our model is listed below.

|      | Mini | Small | Medium | Large |
| ----------- | ----------- | ----------- | ----------- | ----------- |
| Waymo Vector Encoder      | 738,736 | 6,337,072 | 52,624,944 | 80,278,960 |
| NuPlan Raster Encoder      | 11,290,464 | 11,340,672 | 11,474,560 | 11,692,128 |
| Backbone Transformer.      |  3,332,096 | 16,287,488 | 124,439,808 | 1,557,611,200 |
| KP MLP Decoder.              | 18,436 | 270,340 | 2,383,876 | 10,291,204 |
| Key Point Diffusion Decoder | 12,793,700 | 13,410,404 | 16,496,740 | 25,984,868 |
| Trajectory MLP Decoder | 18,436 | 270,340 | 2,383,876 | 10,291,204 |

Reviewer comment (Experiments):
The first weakness I want to mention is that I'm missing the motivation for some of the design decisions of the model. A central claim of the paper is that they reduce motion prediction to a "sequence modeling task". However, for the historical trajectories and map information that the model conditions on, I don't see the motivation for encoding this information causally in time. Full attention across time should work just as well while being more parameter-efficient (for instance, as was done in wayformer), especially when some of the history is partially occluded as in motion prediction data. Additionally, since the authors use diffusion to model future keypoints, I also don't see the motivation to force the diffusion model to be autoregressive. In my mind, the diffusion model should be able to model all future points at once without generating each one-by-one. If performance was amazing with these choices, I'd understand it, but Table 2 shows that motion prediction performance is well below SOTA. If the authors want to claim that these choices improve the scalability of the model, they should compare against a model with full attention in the encoder and all-at-once diffusion in the scaling plot Figure 2.

Author response:
1. **Re Missing the value of sequence modeling formulation:** We find that sequence modeling formulation is the key to scalability. Our experiments indicate great scalability, whether with or without the diffusion decoders on both NuPlan and WOMD. So, the choices for neither encoders nor decoders are relevant to the scalability. We also added a comparison with the recent MotionLM [1] paper, indicating fundamental scalability differences with previous non-sequence modeling models. (see the revised version, Section 1 - Intro, paragraph 3 marked in red). MotionLM also used autoregressive down-sampled (2Hz) key points for prediction but they cannot scale over 10M.

2. **Re Full attention without masks?:** We did implement a mask for invalid data of the WOMD dataset. We skip the details because this is a common implementation for the WOMD dataset only. We believe more shapes of attention for better parameter efficiency are not relevant to scalability and we decided to leave further experiments to future works.

3. **Re Why Autoregressively Diffusion (1by1) for the Key Points?:** We tried diffusion decoders in both 1-by-1 and 5-at-once styles. The results are almost the same across all metrics, as shown in the table below. Since the 1-by-1 generation provides more options for post-processing, we chose this setting as default.

|      | 8sADE  | 3sFDE | 5sFDE | 8sFDE |
| ----------- | ----------- | ----------- | ----------- | ----------- |
| STR(CKS) 16m Bwd w Diffusion (1by1) | 2.095 | 1.129 | 2.519 | 5.300 |
| STR(CKS) 16m Bwd w Diffusion (AllAtOnce) | 2.107 | 1.140 | 2.535 | 5.327 |

Reviewer comment (Writing):
The second weakness I see is that the way in which the model is adapted from MTR or other motion prediction models is not stated as clearly as it should be. My sense is that the method that the authors propose involves training a GPT-like model with temporally-causal mask to regress (x,y) locations at (8,4,2,1,0.5) seconds, then freezing the encoder and training the keypoint decoder with diffusion instead of regression, then keeping the encoder frozen and training an output head that upsamples the keypoints to a 10hz 8-second trajectory. The authors also reference "classification scores" though, so I don't think I'm getting the full picture. My sense from Appendix A.1 is that the reason for these choices is that they stabilize the training of the diffusion model. If that's the case, I think most of the motivation of the paper should emphasize that the goal is to improve stability of diffusion for motion prediction, not to turn motion prediction to a sequence prediction task which is currently stated as the main motivation.

Author response:
4. **Re Diffusion Training in 3 stages?:** There are only two stages for training the diffusion decoders. It is a common training pipeline like training text-to-image diffusion models. The diffusion decoder is trained together with the trajectory decoder in the second stage. We revise and elaborate on this detail in the new Appendix A5 section.

Reviewer comment (Experiments):
Abstract statement "learn to make complex reasonings for long-term planning, extending beyond the horizon of 8 seconds" - where do the authors show that the model generalizes beyond 8 seconds in the paper?

Author response:
1. **Re Long-term reasonings in the abstract:** This is illustrated with the lower right example in Fig. 3. The vehicle changed to a left lane to prepare a left turn coming next and the turning happened beyond 8 seconds in the future. We will add more examples and discussion to the Appendix in the final version.

Reviewer comment (Novelty):
Introduction statement "We arrange the past and future states of the target into one sequence for learning and generation" - I feel the main novelty of the paper is in the design of the decoder. So maybe it's only important to mention that future states are encoded in a sequence.

Author response:
2. **Re Writings in the Intro:** We rephrased the paragraph to avoid the ambiguities you pointed out.

Reviewer comment (Theory):
Section 3.1 first paragraph - there's a mismatch in notation between $s^F$ and $s^T$. Additionally, it should probably be acknowledged in this paragraph that this paper approximates the multi-agent future trajectory distribution as independent, e.g. $p(s^T | c, s_{ego}, s) \approx p(s^T_{ego} | c, s_{ego}, s) \cdot p(s^T_0 | c, s_{ego}, s) \cdot ... \cdot p(s^T_N | c, s_{ego}, s)$.

Author response:
3. **Re 3.1 S^F:** Revised.

Reviewer comment (Experiments):
Section 4 statement "Finally, the model learns the pattern to generate all future states given the contexts and conditions with direct regression losses" - Would it make sense to represent the full trajectory as a sequence of key points? Regression can lead to out-of-distribution trajectories, even when key points at t=(0.5,1,2,4,8) seconds are used for conditioning.

Author response:
4. **Re Section 4:** Revised for better clarity. See the paragraph on “​​Future States” in the revised version.

Reviewer comment (Presentation):
Figure 2 - for the plot on the right, what do the circular dots and the square/triangle dots represent?

Author response:
5. **Re Figure 2 Legends:** Legends added. Thank you.

Reviewer comment (Writing):
Section 5.2 "decoders" - I found this section hard to understand. What are "stacked MLPs"? What is the "additional MSE loss" in addition to? Is "distance regression loss" different from MSE loss? For the "classification scores", what is the network classifying between?

Author response:
7. **Re Section 5.2 - Writing issues about the Decoder paragraph:** We rewrote the whole paragraph. See the paragraph on “Decoders” in the revised version.

Reviewer comment (Experiments):
Appendix A.1 "a plain MLP decoder is used when training the Transformer backbone from scratch, then the backbone is frozen and we train a diffusion-based key points decoder" - are there issues when the model is trained with diffusion from scratch? Is there a reason to make the decoder autoregressive if diffusion is being used?

Author response:
8. **Re Training details of the diffusion decoders:** Addressed in General Comments 3.

Reviewer comment (Experiments):
Table 2 - Can the authors expand on why they only compare to MTR-e2e? MTR and MTR++ get mAP above 0.40 on WOMD.

Author response:
9. **Re Compare with MTR (not e2e) and MTR++:** The MTR (not e2e) adds NMS and other non-learning tricks to boost the performance, making it not a fair comparison. However, we are working on comparing the larger STRs with them on the test set of WOMD (on the leaderboard). The results will be included in the final version.