【論文筆記】World Models for Autonomous Driving: An Initial Survey

原文鏈接：https://ieeexplore.ieee.org/abstract/document/10522953

1. 世界模型的發展

A. 世界模型的結構基礎

在這里插入圖片描述

世界模型包含4個關鍵組件，以模擬人類連貫的思考和決策過程。

a）感知模塊使用如變分自編碼器（VAE）、掩膜自編碼器（MAE）或離散自編碼器（DAE）等編碼模塊，將環境輸入（如圖像、視頻、文本、控制指令等）處理并壓縮為更易于管理的格式。

b）記憶模塊用于記錄和管理過去、現在和預測的世界狀態及其相關代價或獎勵。

c）控制/行為模塊評估當前狀態和世界模型的預測，以決定實現目標（如最小化代價或最大化獎勵）的最優行為序列。該模塊與世界模型模塊區分開來，使用小參數集獨立訓練。

d）世界模型模塊估計當前世界狀態丟失的信息并預測未來的環境狀態。

模型預測的確定性和真實世界現象的不確定性之間的協調是一個挑戰，因此，可以使用溫度變量引入不確定性，或使用結構化框架（如循環狀態空間模型（RSSM）與聯合嵌入預測結構（JEPA）），以微調預測精度與靈活性的平衡。此外，使用Top-K采樣，以及從CNN模型轉換到Transformer模型（如Transformer狀態空間模型（TSSM）和時空Patchwise Transformer（STPT））通過將世界模型的預測與真實世界的可能發展對齊，能更好地近似真實世界的復雜程度和不確定性。

世界模型中最常用的核心結構是RSSM與JEPA。

1）循環狀態空間模型（RSSM） 在隱空間中實現純正向預測。
在這里插入圖片描述
上圖c展示了潛在動態模型的3次步進。隨機變量（圈）和確定性變量（方塊）在模型結構內交互——實線表示生成過程，虛線表示推斷過程。

與無法捕捉多種潛在未來的、確定性初始狀態的RNN（圖a）和完全隨機、無信息保留能力的狀態空間模型（SSM，圖b）不同，RSSM將狀態分解為隨機性狀態和確定性狀態，有效利用確定性元素的預測穩定性和隨機性元素的自適應潛力。

記觀測和行為序列為 $(x_0,a_1,x_1,a_2,x_2,\cdots,a_T,x_T)$ 。觀測到 $x_t$ 后，智能體進行行為 $a_{t+1}$ ，并接收下一步觀測 $x_{t+1}$ 。RSSM將觀測和狀態轉移建模為下列生成過程：
$p(x_{0:T}|a_{1:T})=\int\prod_{t=0}^{T}p(x_t|z_{\leq t},a_{\leq t})p(z_t|z_{<t},a_{\leq t})dz_{0:T}$

其中 $z_{0:T}$ 為隨機隱狀態。近似后驗定義為：
$q(z_{0:T}|x_{0:T},a_{1:T})=\prod_{t=0}^Tq(z_t|z_{<t},a_{\leq t},x_t)$

RSSM使用共享的GRU來將 $z_{<t}$ 和 $a_{\leq t}$ 壓縮為確定性編碼 $h_t$ ：
$h_t=GRU(h_{t-1},MLP(concat[z_{t-1},a_t]))$

隨后，用于計算先驗、似然和后驗的充分統計量：
$p(z_t|z_{<t},a_{\leq t})=MLP(h_t)\\ p(x_t|z_{\leq t},a_{\leq t})=\mathcal N(\hat x_t,1)\\ \hat x_t=Decoder(concat[h_t,z_t])\\ q(z_t|z_{<t},a_{\leq t},x_t)=MLP(concat[h_t,e_t])\\ e_t=Encoder(x_t)$

訓練目標是最大化證據下界（ELBO）：
$\log p(x_{0:T}|a_{1:T})\geq\mathbb E_q[\sum_{t=0}^T\log p(x_t|z_{\leq t},a_{\leq t})-\mathcal L_{KL}(q(z_t|z_{<t},a_{\leq t},x_t),p(z_t|z_{<t},a_{\leq t}))]$

2）聯合嵌入預測結構（JEPA） 關注表達空間而非直接且具體的預測。
在這里插入圖片描述
如圖所示，JEPA通過雙編碼器將輸入 $x$ 和目標 $y$ 抽象為表達 $s_x,s_y$ ，并使用隱變量 $z$ 進行預測。其預測過程更簡單，且能保證結果的相關性和可靠性。

JEPA的核心是能量函數 $E_w(x,y,z;\theta)$ ，其捕捉了模型的預測誤差。這里 $\theta$ 表示模型參數。能量函數定義為
$E_w(x,y,z;\theta)=\|s_y-Pred(s_x,z;\phi)\|_2^2+\lambda\|z\|_2^2$

其中，L2范數平方衡量了預測表達與目標表達 $s_y$ 的歐式距離，即模型預測誤差；預測函數 $P re d$ 將輸入表達 $s_x$ 和隱變量 $z$ 映射到目標空間，其參數為 $\phi$ 。正則化項 $\lambda\|z\|_2^2$ 則懲罰模型的復雜性，以避免過擬合。

優化的目標是通過尋找 $\theta,\phi,z$ 最小化能量函數 $E_w$ ，可表達為以數據分布為約束的復值拉格朗日優化問題：
$L(\theta,\phi,z;x,y,\alpha)=E_w(x,y,z;\theta)-\alpha(h(x,y,z;\theta,\phi)-c)$

其中 $L(\theta,\phi,z;x,y,\alpha)$ 是拉格朗日方程， $\alpha$ 為拉格朗日乘子，強制滿足約束 $h(x,y,z;\theta,\phi)=c$ ； $h(x,y,z;\theta,\phi)$ 為約束函數，其參數為 $\theta$ 和 $\phi$ ， $c$ 為函數 $h$ 的目標值（常數）。

JEPA的訓練包含高階優化方法，考慮二階導數以保證收斂：
$\theta_{t+1}=\theta_t-\eta\nabla_\theta^2L(\theta_t,\phi_t,z_t;x,y,\alpha_t)$

其中 $\theta_{t+1}$ 為更新后的參數向量， $\eta$ 為學習率； $\eta\nabla_\theta^2L(\theta_t,\phi_t,z_t;x,y,\alpha_t)$ 為拉格朗日函數 $L$ 對 $\theta$ 二階偏導數的Hessian矩陣。

考慮到 $z$ 的高維特性和多模態分布的可能性，JPEA可采用變分近似方法處理后驗 $p(z|x,y;\theta)$ ，得到變分下界：
$\log p(y|x;\theta,\phi)\geq\mathbb E_{q(z|x;\psi)}[\log p(y|x,z;\theta;\phi)]-KL[q(z|x;\psi)||p(z|x;\theta)]$

其中 $\log p(y|x;\theta,\phi)$ 為數據 $y$ 在給定 $x$ 和模型參數 $\theta,\phi$ 時的對數似然， $\mathbb E_{q(z|x;\psi)}[\cdot]$ 為關于分布 $q(z|x;\psi)$ 的期望， $KL[q(z|x;\phi)||p(z|x;\theta)]$ 為變分分布 $q(z|x;\psi)$ 和先驗分布 $p(z|x;\theta)$ 的KL散度。