本文經自動駕駛之心公眾號授權轉載，轉載請聯系出處。

論文鏈接：https://arxiv.org/pdf/2410.15774
項目網頁：https://tsinghua-mars-lab.github.io/StateTransformer/
代碼開源：https://github.com/Tsinghua-MARS-Lab/StateTransformer

主要內容：

大型實際駕駛數據集推動了有關自動駕駛數據驅動運動規劃器的各個方面的研究，包括數據增強、模型架構、獎勵設計、訓練策略和規劃器架構。在處理復雜和少樣本情況下，這些方法有較好的表現。但是由于設計過于復雜或訓練范式的問題，這些方法在規劃性能上的泛化能力有限。在本文中，我們回顧并比較了以前的方法，重點關注泛化能力。實驗結果顯示，隨著模型的適當擴展，許多設計元素變得冗余。我們介紹了StateTransformer-2 (STR2)，這是一種可擴展的、僅使用解碼器的運動規劃器，它結合了Vision Transformer (ViT) 編碼器和混合專家(MoE) 的Transformer架構。MoE骨干通過訓練期間的專家路由解決了模態崩潰和獎勵平衡問題。在NuPlan數據集上的大量實驗表明，我們的方法在不同測試集和閉環模擬中比以前的方法具有更好的泛化能力。此外，我們評估了其在真實城市駕駛場景中的可擴展性，顯示出隨著數據和模型規模的增長其一致的準確性提升。

下面我們詳細解釋STR2的模型設計。我們選擇的不是向量化的輸入，而是柵格化圖片的輸入，可以方便進行規模化訓練。輸入還加入了聚類的軌跡作為引導線，結果可以看到在大曲率的場景表現會更合理一些。輸入經過一個使用MoE架構增強的Transformer模型，自回歸的方式輸出引導線類別，關鍵點以及軌跡。

ViT編碼器。我們采用僅解碼的ViT圖像編碼器，以實現更好的可擴展性和性能，它由堆疊的12層Transformer組成。柵格化的圖像被切分成16個小塊。我們選擇GeLU作為ViT編碼器的激活函數。

Mixture-of-Expert。語言建模任務要求模型從復雜且通常具有統計爭議的專家數據獎勵中學習和實現平衡。受MoE模型在語言建模任務上泛化結果的啟發，我們將GPT-2骨干網絡替換為MoE骨干網絡用于序列建模。MoE層通過專用內核和專家并行（EP）提供了更好的內存效率。我們還利用了Flash Attention2 和數據并行（DP）以提高訓練效率。

自回歸。在生成序列中我們添加了聚類軌跡作為嵌入特征用于模態分類，并使用交叉熵損失。我們使用K-Means聚類，從0.7百萬個隨機選擇的動態可行軌跡中按其時空距離提取了512個候選軌跡。每條歸一化的軌跡包括未來8秒的80個軌跡點（x, y和偏航角）。

liauto數據集上的規模化實驗。我們采用了liauto數據集進行了scaling law的探索，liauto數據集是一個工業級的超大規模現實世界駕駛數據集。該數據集包括車道級導航地圖和來自7個RGB攝像頭、1個LiDAR和1個毫米波雷達（MMWR）的傳感器設置的跟蹤結果。我們選擇了過去6個月內收集的城市駕駛場景，其中沒有任何人工標注。我們篩選出錯誤的導航路線，因為這些路線與實際的未來駕駛軌跡不匹配。最終，我們將所有駕駛日志重新整理為長達10秒的訓練和測試樣本，其中包括過去的2秒和未來的8秒。最終的訓練數據集擁有超過1b訓練樣本。實驗結果（如圖2）可以看到隨著數據規模的增加以及模型參數的增加，test loss都有下降的趨勢。均衡考慮訓練成本和收益的關系，我們最后采用的為800m的模型。

從圖3 nuplan閉環仿真的結果可以看到STR2取得了全面SOTA。結果來看專家軌跡的NR分數較高，R的分數卻很低，說明專家軌跡沒有一個適應環境變化的能力，仿真環境和實車有一定的區別。結果可以看出PDM-Hybrid相比于其他的方法的R得分明顯高于NR，說明生成-評估的范式在模型泛化性能上表現優異。我們借鑒了PDM的生成-評估范式，基于STR2模型的輸出結果進行了候選軌跡生成，再經過PDM打分器輸出得分最高的軌跡。

圖4可視化結果可以看出PDM-Hybrid軌跡因為基于當前車道中心線，沒有主動變道和繞障的能力，我們的模型由于泛化性能強，所以很多場景下表現比PDM-Hybrid合理很多。本工作在nuplan testhard 數據集上的實驗結果證明了生成-評估范式的合理性。對于生成-評估范式，在后續工作中我們還可以有更多嘗試，例如如何在模型層面輸出更合理的候選軌跡，如何將未來的不確定性考慮進打分器中，同時可以考慮將本文的方法作為真值標注的一種方式，幫助模型朝著更合理的方向迭代。為了感受模型的泛化性能我們將nuplan數據訓練得到的模型直接應用于liauto數據集上推理，結果表現良好，對于動態障礙物的避讓交互等都有合理的輸出。