世界模型再進(jìn)化!DrivingWorld:通過視頻GPT構(gòu)建自動(dòng)駕駛世界模型(港科技&地平線)
寫在前面 & 筆者的個(gè)人理解
最近自回歸(AR)生成模型的成功,如自然語言處理中的GPT系列,促使人們努力在視覺任務(wù)中復(fù)制這一成功。一些工作試圖通過構(gòu)建能夠生成逼真的未來視頻序列和預(yù)測(cè)自車狀態(tài)的基于視頻的世界模型,將這種方法擴(kuò)展到自動(dòng)駕駛。然而,先前的工作往往產(chǎn)生不令人滿意的結(jié)果,因?yàn)榻?jīng)典的GPT框架旨在處理1D上下文信息,如文本,并且缺乏對(duì)視頻生成所必需的空間和時(shí)間動(dòng)態(tài)進(jìn)行建模的固有能力。本文介紹了DrivingWorld,這是一個(gè)GPT風(fēng)格的自動(dòng)駕駛世界模型,具有多種時(shí)空融合機(jī)制。這種設(shè)計(jì)能夠有效地對(duì)空間和時(shí)間動(dòng)態(tài)進(jìn)行建模,從而促進(jìn)高保真、長(zhǎng)持續(xù)時(shí)間的視頻生成。具體來說,我們提出了一種下一狀態(tài)預(yù)測(cè)策略來模擬連續(xù)幀之間的時(shí)間一致性,并應(yīng)用下一token預(yù)測(cè)策略來捕獲每個(gè)幀內(nèi)的空間信息。為了進(jìn)一步提高泛化能力,我們提出了一種新的掩碼策略和重新加權(quán)策略用于token預(yù)測(cè),以緩解長(zhǎng)期漂移問題并實(shí)現(xiàn)精確控制。我們的工作展示了制作高保真、持續(xù)時(shí)間超過40秒的一致視頻片段的能力,這比最先進(jìn)的駕駛世界模型長(zhǎng)2倍多。實(shí)驗(yàn)表明,與先前的工作相比,我們的方法實(shí)現(xiàn)了卓越的視覺質(zhì)量和更精確的可控未來視頻生成。
- 開源鏈接:https://github.com/YvanYin/DrivingWorld
總結(jié)來說,本文介紹了DrivingWorld,這是一個(gè)基于GPT風(fēng)格視頻生成框架的駕駛世界模型。我們的主要目標(biāo)是在自回歸框架中增強(qiáng)時(shí)間一致性的建模,以創(chuàng)建更準(zhǔn)確可靠的世界模型。為了實(shí)現(xiàn)這一目標(biāo),我們的模型結(jié)合了三個(gè)關(guān)鍵創(chuàng)新:1)時(shí)間感知標(biāo)記化:我們提出了一種時(shí)間感知標(biāo)記器,將視頻幀轉(zhuǎn)換為時(shí)間相干標(biāo)記,將未來視頻預(yù)測(cè)的任務(wù)重新表述為預(yù)測(cè)序列中的未來標(biāo)記。2)混合token預(yù)測(cè):我們引入了一種下一狀態(tài)預(yù)測(cè)策略來預(yù)測(cè)連續(xù)狀態(tài)之間的時(shí)間一致性,而不是僅僅依賴于下一個(gè)token預(yù)測(cè)策略。之后,應(yīng)用下一個(gè)token預(yù)測(cè)策略來捕獲每個(gè)狀態(tài)內(nèi)的空間信息。3)長(zhǎng)時(shí)間可控策略:為了提高魯棒性,我們?cè)谧曰貧w訓(xùn)練過程中實(shí)施了隨機(jī)標(biāo)記丟棄和平衡注意力策略,從而能夠生成具有更精確控制的持續(xù)時(shí)間更長(zhǎng)的視頻。DrivingWorld使用AR框架增強(qiáng)了視頻生成中的時(shí)間連貫性,學(xué)習(xí)了未來進(jìn)化的有意義表示。實(shí)驗(yàn)表明,所提出的模型具有良好的泛化性能,能夠生成超過40秒的視頻序列,并提供準(zhǔn)確的下一步軌跡預(yù)測(cè),保持合理的可控性。
相關(guān)工作回顧
世界模型。世界模型捕捉了環(huán)境的全面表示,并根據(jù)一系列行動(dòng)預(yù)測(cè)了未來的狀態(tài)。世界模型在游戲和實(shí)驗(yàn)室環(huán)境中都得到了廣泛的探索。Dreamer利用過去的經(jīng)驗(yàn)訓(xùn)練了一個(gè)潛在動(dòng)力學(xué)模型,以預(yù)測(cè)潛在空間內(nèi)的狀態(tài)值和行為。DreamerV2基于最初的Dreamer模型構(gòu)建,在雅達(dá)利游戲中達(dá)到了人類水平的性能。DreamerV3使用了更大的網(wǎng)絡(luò),并成功地學(xué)會(huì)了從零開始在Minecraft中獲取鉆石。DayDreamer擴(kuò)展了Dreamer,在現(xiàn)實(shí)世界中訓(xùn)練了四個(gè)機(jī)器人,成功地完成了運(yùn)動(dòng)和操縱任務(wù)。
最近駕駛場(chǎng)景的世界模型在學(xué)術(shù)界和工業(yè)界都引起了極大的關(guān)注。之前的大多數(shù)工作僅限于模擬器或控制良好的實(shí)驗(yàn)室環(huán)境。Drive WM使用擴(kuò)散模型探索了現(xiàn)實(shí)世界中的駕駛規(guī)劃者。GAIA-1基于自回歸模型研究了現(xiàn)實(shí)世界的駕駛規(guī)劃者,但GAIA-1具有較大的參數(shù)和計(jì)算需求,隨著條件框架數(shù)量的增加而增加。在本文中,我們提出了一個(gè)自回歸框架下的自動(dòng)駕駛場(chǎng)景的有效世界模型。
VQVAE。VQVAE通過矢量量化學(xué)習(xí)離散碼本表示,以對(duì)圖像分布進(jìn)行建模。VQGAN通過結(jié)合LPIPS損失和對(duì)抗性PatchGAN損失提高了真實(shí)感。MoVQ通過將空間變異信息嵌入量化向量中,解決了VQGAN的空間條件歸一化問題。LlamaGen進(jìn)一步微調(diào)了VQGAN,表明較小的碼本矢量維數(shù)和較大的碼本大小可以提高重建性能。雖然基于VQGAN的結(jié)構(gòu)被廣泛使用,但一些方法探索了更高效的架構(gòu)。ViT VQGAN用視覺變換器取代了卷積編碼器-解碼器,提高了模型捕獲長(zhǎng)距離依賴關(guān)系的能力。VAR采用多尺度結(jié)構(gòu)來預(yù)測(cè)先前尺度的后續(xù)尺度,從而提高了發(fā)電質(zhì)量和速度。然而,這些方法側(cè)重于單一圖像處理,阻礙了它們捕獲時(shí)間一致性。為了解決這個(gè)問題,我們提出了一種時(shí)間感知標(biāo)記器和解碼器。
視頻生成。目前有三種主流的視頻生成模型:基于GAN、基于擴(kuò)散和基于GPT的方法?;贕AN的方法經(jīng)常面臨幾個(gè)挑戰(zhàn),例如模式崩潰,生成器生成的視頻的多樣性受到限制。此外,生成器和鑒別器之間的對(duì)抗性學(xué)習(xí)可能會(huì)導(dǎo)致訓(xùn)練過程中的不穩(wěn)定?;跀U(kuò)散的方法的一個(gè)主要問題是它們無法生成精確控制的視頻。擴(kuò)散過程的隨機(jī)性在每一步都引入了隨機(jī)性,使得難以對(duì)生成內(nèi)容中的特定屬性進(jìn)行嚴(yán)格控制。另一方面,傳統(tǒng)的基于GPT的方法允許一定程度的控制,但它們的計(jì)算成本隨序列長(zhǎng)度呈二次增長(zhǎng),顯著影響了模型效率。本文提出了一種解耦的時(shí)空世界模型框架,該框架在確保精確控制的同時(shí),顯著降低了計(jì)算成本,提高了模型效率。
DrivingWorld方法詳解
我們提出的世界模型DrivingWorld利用GPT風(fēng)格的架構(gòu)高效預(yù)測(cè)未來狀態(tài),能夠以10Hz的頻率將預(yù)測(cè)時(shí)間延長(zhǎng)到40秒以上。該模型旨在理解過去的現(xiàn)實(shí)世界狀態(tài),并預(yù)測(cè)未來的視頻內(nèi)容和車輛運(yùn)動(dòng)。DrivingWorld專門專注于根據(jù)時(shí)間1到T的歷史狀態(tài)預(yù)測(cè)時(shí)間T+1的下一個(gè)狀態(tài),我們可以通過逐一順序預(yù)測(cè)未來狀態(tài)來生成長(zhǎng)視頻。
如圖2所示,我們提出的DrivingWorld不僅可以根據(jù)過去的觀測(cè)結(jié)果生成未來狀態(tài),還可以通過操縱車輛的位置和方向來支持復(fù)雜駕駛場(chǎng)景的可控模擬。
Tokenizer
標(biāo)記化將連續(xù)數(shù)據(jù)轉(zhuǎn)換為離散標(biāo)記,從而能夠與語言模型和增強(qiáng)的多模態(tài)序列建模集成。在我們的方法中,標(biāo)記器將多模態(tài)狀態(tài)映射到統(tǒng)一的離散空間中,從而實(shí)現(xiàn)了精確可控的多模態(tài)生成。為了為圖像生成時(shí)間一致的嵌入,我們提出了一種時(shí)間感知的矢量量化標(biāo)記器。我們提出的車輛姿態(tài)標(biāo)記器將姿態(tài)軌跡離散化,并將其整合到我們的DrivingWorld中。
前言:Single Image Vector Quantized Tokenizer。單圖像矢量量化(VQ)標(biāo)記器旨在將圖像特征圖轉(zhuǎn)換為離散標(biāo)記q。量化器利用包含K個(gè)矢量的學(xué)習(xí)離散碼本,將每個(gè)特征f(i,j)映射到Z中最接近代碼的索引。這種方法能夠?qū)⑦B續(xù)圖像數(shù)據(jù)轉(zhuǎn)換為離散token。
時(shí)間感知矢量量化標(biāo)記器。單圖像VQ標(biāo)記器通常難以產(chǎn)生時(shí)間一致的嵌入,導(dǎo)致不連續(xù)的視頻預(yù)測(cè),阻礙了世界模型的訓(xùn)練。
為了解決這個(gè)問題,我們提出了一種時(shí)間感知的矢量量化標(biāo)記器,旨在確保隨時(shí)間推移的一致嵌入。具體來說,為了捕捉時(shí)間依賴性,我們?cè)赩QGAN量化之前和之后都插入了一個(gè)self-att,其中注意力沿著時(shí)間維度進(jìn)行操作。這使得我們的模型能夠捕捉幀之間的長(zhǎng)期時(shí)間關(guān)系,提高生成序列的連貫性和一致性。我們的模型基于LlammaGen的開源VQGAN實(shí)現(xiàn)。我們直接而有效的時(shí)間self-att的集成可以無縫地整合到原始框架中,然后進(jìn)行微調(diào),以開發(fā)一個(gè)健壯且通用的時(shí)間感知VQ標(biāo)記器。
車輛位姿標(biāo)記器。為了準(zhǔn)確表示車輛的自車狀態(tài),包括其方向θ和位置(x,y),我們采用以自車輛為中心的坐標(biāo)系,如圖2所示。我們采用相鄰時(shí)間步長(zhǎng)之間的相對(duì)姿態(tài),而不是全局姿態(tài)。這是因?yàn)樵陂L(zhǎng)期序列中,由于絕對(duì)姿態(tài)值的增加,全球姿態(tài)帶來了重大挑戰(zhàn)。這種增長(zhǎng)使得歸一化變得困難,并降低了模型的魯棒性。隨著序列變長(zhǎng),管理這些大的姿勢(shì)值變得越來越困難,阻礙了有效的長(zhǎng)期視頻生成。
World Model
世界模型旨在理解過去的狀態(tài)輸入,模擬現(xiàn)實(shí)世界的動(dòng)態(tài),并預(yù)測(cè)未來的狀態(tài)。在我們的背景下,它預(yù)測(cè)了即將到來的駕駛場(chǎng)景,并規(guī)劃了可行的未來軌跡。為此,世界模型將歷史狀態(tài)標(biāo)記連接成一個(gè)長(zhǎng)序列,其中2D圖像標(biāo)記以鋸齒形順序展開為1D形式。因此,目標(biāo)是預(yù)測(cè)下一個(gè)狀態(tài)?;谶^去的觀測(cè)序列,捕捉時(shí)間和多模態(tài)依賴關(guān)系。請(qǐng)注意,來自不同模態(tài)的所有離散token在被饋送到世界模型之前,都由其各自的可學(xué)習(xí)嵌入層映射到共享的潛在空間中。所有后續(xù)過程都在這個(gè)潛在空間內(nèi)進(jìn)行。
前言:下一個(gè)token預(yù)測(cè)。一種直接的方法是使用GPT-2結(jié)構(gòu)進(jìn)行1D順序下一個(gè)token預(yù)測(cè)。圖3(a)顯示了一個(gè)簡(jiǎn)化示例。因果注意被應(yīng)用于下一個(gè)token預(yù)測(cè),T+1中的第i個(gè)token被建模為:
因此我們提出了一種下一狀態(tài)預(yù)測(cè)管道,它由兩個(gè)模塊組成:一個(gè)集成時(shí)間和多模態(tài)信息以生成下一狀態(tài)特征(即時(shí)間多模態(tài)融合模塊),另一個(gè)是自回歸模塊(即內(nèi)部狀態(tài)自回歸模塊)以生成高質(zhì)量的內(nèi)部狀態(tài)token。時(shí)間多模態(tài)融合模塊。我們的時(shí)間多模態(tài)模塊由一個(gè)單獨(dú)的時(shí)間層和一個(gè)多模態(tài)層組成。這將時(shí)間和多模態(tài)信息的處理解耦,從而提高了訓(xùn)練和推理速度,同時(shí)也降低了GPU內(nèi)存消耗。如圖3(b)所示,我們建議在時(shí)間轉(zhuǎn)換層Fa(·)中使用因果注意力掩碼,其中每個(gè)token只關(guān)注自身和所有先前幀中相同順序位置的token,充分利用時(shí)間信息。
在多模態(tài)信息融合層Fb(·)中,我們?cè)谕粠胁捎秒p向掩碼,旨在充分整合內(nèi)部狀態(tài)多模態(tài)信息,并促進(jìn)模態(tài)之間的交互。每個(gè)token處理來自同一時(shí)間步的其他token:
內(nèi)部狀態(tài)自回歸模塊。在時(shí)間多模態(tài)模塊之后,我們獲得了用于未來幀狀態(tài)預(yù)測(cè)的特征。一種天真的方法是同時(shí)預(yù)測(cè)下一個(gè)狀態(tài)tokenht。最近,多圖像生成工作提出,用于下一個(gè)token預(yù)測(cè)的自回歸流水線可以生成更好的圖像,甚至優(yōu)于擴(kuò)散方法。受此啟發(fā),我們提出了一個(gè)內(nèi)部狀態(tài)自回歸模塊來生成下一時(shí)間步的姿勢(shì)和圖像(見圖3(b))。
然后,它們被輸入到內(nèi)部狀態(tài)自回歸Transformer層Fc(·)。因果掩碼在這些層中使用,因此每個(gè)token只能出席自己并前綴內(nèi)部狀態(tài)token。自回歸過程如方程式6所示。由于我們的管道同時(shí)包含了下一個(gè)狀態(tài)預(yù)測(cè)和下一個(gè)內(nèi)部狀態(tài)token預(yù)測(cè),我們?cè)谟?xùn)練中實(shí)施了兩種教師強(qiáng)制策略,即一種用于幀級(jí)別,另一種用于內(nèi)部狀態(tài)級(jí)別。
訓(xùn)練損失交叉熵:
Decoder
使用世界模型預(yù)測(cè)下一個(gè)狀態(tài)標(biāo)記,然后我們可以利用解碼器為該狀態(tài)生成相應(yīng)的相對(duì)方向、相對(duì)位置和重建圖像。這個(gè)過程使我們能夠?qū)㈩A(yù)測(cè)的潛在表示映射回物理輸出,包括空間和視覺數(shù)據(jù)。
Vehicle Pose Decoder:
Temporal-aware Decoder:
Long-term Controllable Generation
Token Dropout實(shí)現(xiàn)無漂移自動(dòng)回歸。在訓(xùn)練過程中,世界模型使用過去的地面真實(shí)token作為條件來預(yù)測(cè)下一個(gè)token。然而,在推理過程中,模型必須依賴于先前生成的表征進(jìn)行調(diào)節(jié),這可能包含缺陷。僅使用完美的GT圖像進(jìn)行訓(xùn)練可能會(huì)在推理過程中導(dǎo)致內(nèi)容漂移問題,導(dǎo)致生成的輸出迅速退化并最終失敗。為了解決這個(gè)問題,我們提出了一種隨機(jī)掩蔽策略(RMS),其中一些來自地面真實(shí)token的token被隨機(jī)丟棄。每個(gè)標(biāo)記有50%的機(jī)會(huì)被該幀中的另一個(gè)隨機(jī)標(biāo)記替換,并且這種丟失以30%的概率應(yīng)用于整個(gè)調(diào)節(jié)圖像序列。如圖4所示,這種dropout策略顯著緩解了推理過程中的漂移問題。
平衡注意力實(shí)現(xiàn)精確控制。世界模型利用廣泛的注意力操作在代幣之間交換和融合信息。然而,每個(gè)前視圖圖像被離散化為512個(gè)標(biāo)記,而只有2個(gè)標(biāo)記表示姿勢(shì)(方向和位置)。這種不平衡會(huì)導(dǎo)致模型忽略姿態(tài)信號(hào),導(dǎo)致可控生成不令人滿意。為了解決這個(gè)問題,我們提出了一種平衡的注意力操作,通過在注意力機(jī)制中優(yōu)先考慮自車狀態(tài)標(biāo)記,而不是平等地關(guān)注所有標(biāo)記,來實(shí)現(xiàn)更精確的控制。具體來說,我們手動(dòng)增加注意力圖中方向和位置標(biāo)記的權(quán)重(在softmax層之前),分別為這些標(biāo)記添加0.4和0.2的恒定權(quán)重。此外,我們結(jié)合了QK范數(shù)和2D旋轉(zhuǎn)位置編碼,以進(jìn)一步穩(wěn)定訓(xùn)練并提高性能。
實(shí)驗(yàn)結(jié)果
結(jié)論和未來工作
總之,DrivingWorld通過利用GPT風(fēng)格的框架來生成更長(zhǎng)、高保真的視頻預(yù)測(cè),并提高了泛化能力,從而解決了以前自動(dòng)駕駛視頻生成模型的局限性。與在長(zhǎng)序列中難以保持連貫性或嚴(yán)重依賴標(biāo)記數(shù)據(jù)的傳統(tǒng)方法不同,DrivingWorld生成了逼真、結(jié)構(gòu)化的視頻序列,同時(shí)實(shí)現(xiàn)了精確的動(dòng)作控制。與經(jīng)典的GPT結(jié)構(gòu)相比,我們提出的時(shí)空GPT結(jié)構(gòu)采用了下一狀態(tài)預(yù)測(cè)策略來模擬連續(xù)幀之間的時(shí)間一致性,然后應(yīng)用下一token預(yù)測(cè)策略來捕獲每個(gè)幀內(nèi)的空間信息。展望未來,我們計(jì)劃整合更多的多模態(tài)信息,并整合多視圖輸入。通過融合來自不同模態(tài)和視角的數(shù)據(jù),我們的目標(biāo)是提高動(dòng)作控制和視頻生成的準(zhǔn)確性,增強(qiáng)模型理解復(fù)雜駕駛環(huán)境的能力,并進(jìn)一步提高自動(dòng)駕駛系統(tǒng)的整體性能和可靠性。