世界模型(World Model)作為學習到的動力學模擬器,正變得越來越強大。它們能夠在高維視覺空間中預測長序列的未來觀測,並以過去難以想象的方式跨任務泛化。然而,擁有強大的預測模型並不等同於能夠有效地用於控制或規劃。長時域規劃在實踐中仍然脆弱:最佳化變得病態,非貪心結構導致糟糕的區域性極小值,高維潛在空間引入微妙的失敗模式。
針對這些問題,來自伯克利AI研究團隊(BAIR)的Michael Psenka、Mike Rabbat、Aditi Krishnapriyan、Yann LeCun和Amir Bar提出了GRASP(Gradient RelAxed Stochastic Planner),一種基於梯度的隨機規劃器。GRASP透過三種創新設計使長時域規劃變得實用:提升軌跡到虛擬狀態、向狀態迭代新增隨機性以及重塑梯度。
世界模型與規劃挑戰
世界模型是一種學習模型,給定當前狀態和未來動作序列,預測接下來會發生什麼。在規劃中,最簡單的做法是透過模型展開動作序列並最小化終端誤差。但長時域展開會建立深而病態的計算圖,導致梯度爆炸或消失。此外,長時域任務通常需要非貪心行為(如繞過障礙物),這使得損失景觀充滿區域性陷阱。
一種自然的解決方案是採用搭配法(Collocation),將動力學約束作為軟約束,同時最佳化動作和狀態序列。搭配法允許平行計算所有時間步,且梯度不再透過深層展開回傳,避免了長期依賴問題。然而,對於基於深度學習的世界模型,直接最佳化狀態會遭遇對抗魯棒性問題:狀態輸入梯度極其脆弱,因為訓練過程只正則化資料流形切向方向,而法向方向的行為不受約束,導致模型對未見過狀態高度敏感。
GRASP的核心創新
GRASP建立在搭配法框架之上,但針對深度學習模型的梯度特性進行了專門設計。其關鍵洞察是:雖然狀態梯度不可靠,但動作梯度通常是良好行為的,因為動作空間維數低且訓練充分。因此,GRASP僅依賴動作梯度進行規劃。
GRASP包含兩個主要成分:
- 狀態噪聲探索:在最佳化過程中向虛擬狀態更新注入高斯噪聲,幫助跳躍到不同的損失盆地,同時動作仍由梯度引導。
- 梯度重塑:停止狀態到世界模型的梯度傳播,避免對抗脆弱性。同時新增密集目標項(鼓勵每一步向目標靠近),與停止梯度的動力學損失平衡,防止陷入瑣碎解。
此外,GRASP每K次迭代進行一次短期的“同步”修正:從初始狀態使用當前動作展開,對原始序列損失進行少量梯度步,使狀態和動作保持與真實軌跡一致。
實驗結果
在Push-T任務上,GRASP與CEM、標準梯度下降(GD)和LatCo等方法進行了比較。結果顯示,隨著規劃時域增加,GRASP在成功率和規劃速度上均顯著優於其他方法。例如,當時域H=60時,GRASP成功率達到26.2%,而CEM僅為7.2%,GD為16.4%;同時GRASP的中位成功時間僅為49.1秒,遠快於其他方法。GRASP在長時域下不僅成功率更高,而且規劃更快。
未來展望
研究者認為,目前是利用世界模型進行規劃的有趣時期:經典規劃與控制文獻非常成熟,但現代大規模世界模型上的純規劃最佳化仍待充分探索。未來方向包括擴充套件到擴散世界模型、採用更先進的最佳化器和噪聲策略,以及將GRASP整合到閉環系統或強化學習策略中。