henry 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI

世界模型，也能持續(xù)學(xué)習(xí)了！

剛剛，紐約大學(xué)聯(lián)合LeCun初創(chuàng)AMI帶來JEPA系列的最新成果——

AdaJEPA

與過去在預(yù)訓(xùn)練結(jié)束后就凍結(jié)參數(shù)的世界模型不同，AdaJEPA能夠在與環(huán)境交互中，基于測(cè)試時(shí)自適應(yīng)（Test-Time Adaptation, TTA），實(shí)時(shí)調(diào)整世界模型的編碼器和預(yù)測(cè)器參數(shù)，從而實(shí)現(xiàn)持續(xù)學(xué)習(xí)。

具體而言，AdaJEPA通過計(jì)劃、執(zhí)行、觀測(cè)、更新、再規(guī)劃的閉環(huán)，在每次交互中只執(zhí)行MPC規(guī)劃出的第一段動(dòng)作，然后把真實(shí)觀察到的下一幀狀態(tài)，當(dāng)成自監(jiān)督信號(hào)來更新世界模型。

由此，在下一輪規(guī)劃時(shí)用的就不再是剛部署時(shí)那個(gè)凍結(jié)模型，而是已經(jīng)被當(dāng)前環(huán)境“校準(zhǔn)”過的模型。

這個(gè)思路有點(diǎn)像經(jīng)典強(qiáng)化學(xué)習(xí)里的Dyna

模型不是一次訓(xùn)練完就結(jié)束，而是在真實(shí)交互中不斷修正自己對(duì)世界的理解。

實(shí)驗(yàn)結(jié)果也表明，無論是在分布內(nèi)環(huán)境，還是面對(duì)各類分布外偏移，AdaJEPA的規(guī)劃成功率都明顯優(yōu)于固定世界模型。

這是怎么做到的？

計(jì)劃、執(zhí)行、觀測(cè)、更新、再規(guī)劃

一直以來，基于JEPA路線的隱空間世界模型，都有一個(gè)默認(rèn)前提，就是模型訓(xùn)練完，就凍結(jié)參數(shù)。

整個(gè)流程大致是這樣：

模型先在離線軌跡上學(xué)習(xí)把高維圖像壓進(jìn)latent space，然后再在這個(gè)隱空間里預(yù)測(cè)未來。

到了測(cè)試階段，MPC（Model Predictive Control，模型預(yù)測(cè)控制）就會(huì)調(diào)用這個(gè)凍結(jié)的世界模型，在隱空間里向前滾動(dòng)“想象”未來，優(yōu)化出一串動(dòng)作，再把第一步動(dòng)作拿到真實(shí)環(huán)境里執(zhí)行。

（注：MPC的核心思想是每次只往前預(yù)測(cè)一小段，算出一串動(dòng)作，但先只執(zhí)行第一步。等真實(shí)環(huán)境反饋回來，再重新預(yù)測(cè)、重新規(guī)劃）

但問題在于，環(huán)境一變，凍結(jié)世界模型就容易失準(zhǔn)。

當(dāng)系統(tǒng)面臨測(cè)試時(shí)分布偏移（Test-time Distribution Shift）時(shí)，在latent space里看起來能到達(dá)目標(biāo)的動(dòng)作，落到真實(shí)環(huán)境里，可能一步都不對(duì)。

更麻煩的是，MPC本來就靠短時(shí)域滾動(dòng)規(guī)劃，單步誤差往后滾幾步，就會(huì)被放大。

為解決這一問題，論文提出AdaJEPA框架。

它的核心判斷是：世界模型不該訓(xùn)練完就固定在那里。它應(yīng)該像真正部署中的智能體一樣，一邊行動(dòng)，一邊用新經(jīng)驗(yàn)校準(zhǔn)自己。

具體來說，AdaJEPA的循環(huán)可以分為四步：

• ：模型先把當(dāng)前觀測(cè)編碼成latent state，然后用當(dāng)前世界模型進(jìn)行MPC，在隱空間里向前滾動(dòng)預(yù)測(cè)，找出一串最接近目標(biāo)狀態(tài)的動(dòng)作。
• 執(zhí)行：模型不會(huì)一次性執(zhí)行完整動(dòng)作序列，而是只執(zhí)行第一段動(dòng)作。隨后，真實(shí)環(huán)境返回下一幀觀測(cè)。
• 更新：把這次真實(shí)狀態(tài)轉(zhuǎn)移存進(jìn)在線緩存區(qū)。AdaJEPA再讓模型根據(jù)觀測(cè)和動(dòng)作預(yù)測(cè)下一步latent state，并和真實(shí)狀態(tài)預(yù)測(cè)編碼出的latent state對(duì)齊。預(yù)測(cè)錯(cuò)在哪里，梯度就從哪里回來。
• 再規(guī)劃：更新后的世界模型立刻進(jìn)入下一輪MPC。論文里默認(rèn)只更新視覺編碼器和預(yù)測(cè)器的最后幾層，每次重規(guī)劃只做1步梯度下降。

由此，AdaJEPA的循環(huán)不再只是傳統(tǒng) MPC 的：規(guī)劃，執(zhí)行，再規(guī)劃。

而是變成了：規(guī)劃，執(zhí)行，觀測(cè)，更新，再規(guī)劃。

世界模型也因此不再只是一個(gè)被動(dòng)調(diào)用的“想象器”，而變成了一個(gè)會(huì)在部署過程中持續(xù)校準(zhǔn)自己的模塊。

實(shí)現(xiàn)細(xì)節(jié)

在實(shí)現(xiàn)上，AdaJEPA的底座依然是JEPA（Joint-Embedding Predictive Architectures），聯(lián)合嵌入預(yù)測(cè)架構(gòu)。

和傳統(tǒng)像素級(jí)預(yù)測(cè)世界模型不同，JEPA并不直接預(yù)測(cè)未來圖像，而是先把圖像壓進(jìn)一個(gè)更緊湊的隱空間，只在latent space里預(yù)測(cè)未來狀態(tài)。

具體來說，整個(gè)模型由三個(gè)核心組件組成：

• 狀態(tài)編碼器，把當(dāng)前觀測(cè)編碼成隱狀態(tài)。
• 動(dòng)作編碼器 ，把動(dòng)作編碼成動(dòng)作嵌入。
• 預(yù)測(cè)器，根據(jù)當(dāng)前隱狀態(tài)和動(dòng)作嵌入，預(yù)測(cè)下一步隱狀態(tài)。

AdaJEPA的在線更新，就發(fā)生在這個(gè)隱空間里。

每次執(zhí)行動(dòng)作后，系統(tǒng)都會(huì)把真實(shí)狀態(tài)轉(zhuǎn)移存入在線緩存區(qū)。這個(gè)緩存區(qū)不會(huì)無限增長，論文里默認(rèn)只保留最近N條轉(zhuǎn)移。

更新時(shí)，AdaJEPA會(huì)讓模型根據(jù)當(dāng)前觀測(cè)和動(dòng)作預(yù)測(cè)下一時(shí)刻的隱狀態(tài)，再和真實(shí)下一幀觀測(cè)編碼出的隱狀態(tài)對(duì)齊。

為了防止在線更新把原本的表征空間拉崩，論文做了兩個(gè)限制：

一是對(duì)目標(biāo)表征使用stop-gradient；二是只更新少量參數(shù)。

實(shí)驗(yàn)?zāi)J(rèn)只更新視覺編碼器和預(yù)測(cè)器的最后幾層，并且每次MPC重規(guī)劃只進(jìn)行1步梯度下降。

所以，這并不是把整個(gè)世界模型在線重訓(xùn)一遍。

它更像是每走一步，就用剛剛從真實(shí)環(huán)境里獲得的新反饋，把世界模型往當(dāng)前環(huán)境上輕輕校準(zhǔn)一下。

實(shí)驗(yàn)測(cè)試

為了驗(yàn)證測(cè)試時(shí)自適應(yīng)能否讓凍結(jié)世界模型在環(huán)境變化中重新校準(zhǔn)自己，論文在PushT/PushObjPointMaze兩個(gè)基準(zhǔn)上做了實(shí)驗(yàn)。

實(shí)驗(yàn)表明，在 PushObj未見過形狀上，AdaJEPA幾乎把規(guī)劃成功率翻倍。

而在PointMaze的未見過布局中，提升同樣明顯。

其中，GD規(guī)劃成功率從53.3%提到 78.7%，CEM從49.3%提到70.7%。

更關(guān)鍵的是，這種在線更新帶來的額外延遲并不高。

論文默認(rèn)只更新視覺編碼器和預(yù)測(cè)器的最后幾層，每次MPC重規(guī)劃只做1步梯度下降，額外延遲只有 0.01到0.03秒。

也就是說，AdaJEPA并不是用很重的在線訓(xùn)練，強(qiáng)行換來更高成功率。

它更像是在原有世界模型的基礎(chǔ)上，加入了一個(gè)輕量的“部署時(shí)自我校準(zhǔn)”機(jī)制。

總體來看，這篇論文想說明的是：

世界模型不必訓(xùn)練完就被凍結(jié)，只要讓它在部署時(shí)利用真實(shí)交互結(jié)果做輕量更新，就已經(jīng)能顯著提升面對(duì)環(huán)境變化時(shí)的魯棒性。

作者介紹

最后，讓我們來簡單介紹一下這篇論文的作者們。

作者Ying Wang目前是紐約大學(xué)數(shù)據(jù)科學(xué)中心CILVR Lab 的博士生，研究方向是世界模型，導(dǎo)師是Mengye RenYann LeCun

另一位作者Oumayma Bounou目前是紐約大學(xué)博士后研究員，她的研究興趣集中在世界模型、控制和優(yōu)化。

目前，她正與LeCun合作研究世界模型。

此外，紐約大學(xué)計(jì)算機(jī)科學(xué)與數(shù)據(jù)科學(xué)助理教授Mengye Ren和圖靈獎(jiǎng)獲得者Yann LeCun共同擔(dān)任指導(dǎo)作者。

參考鏈接[1]https://arxiv.org/pdf/2606.32026