上海AI Lab新研究:利用人類操作視頻訓練高效具身策略

投稿作者:白辰甲(上海人工智能實驗室青年研究員)、何浩然(上海人工智能實驗室實習生)

如何使機器人學習多任務通用具身策略是一項長期的挑戰。

從近期大語言模型發展的歷程看,獲得通用知識的關鍵是從互聯網中獲得大量數據,使用大規模網絡結構和無監督學習目標進行預訓練。

類似的,學習通用具身策略需要從大量機器人交互數據中獲得實體、任務、環境、動作的數據,從而更好的理解環境並作出決策。

然而,與視覺和自然語言處理不同,高質量的具身數據獲取是非常困難的,且不同機器人的數據往往難以通用。現有研究主要通過藉助基礎模型作爲基礎具身策略,但由於機器人和其他領域數據存在較大差異,基礎策略往往在具身場景中存在適應性和泛化難題。

近期,上海人工智能實驗室、香港科技大學、上海交通大學等聯合提出的大規模人類視頻預訓練和具身策略微調算法給出了一個合理的解決方案,提出了全新的基於視頻預測擴散模型的高效策略學習算法:Video-based Policy Learning via Discrete Diffusion(VPDD)來解決該問題。

通過從大規模人類操作數據 Ego4d 學習統一的視頻表徵,使用大量無動作視頻構建自監督視頻預測擴散模型預訓練任務,並在少量有動作標記的具身數據上進行高效策略微調,能夠使通用人類操作視頻中編碼的物理世界先驗知識適應於具身任務,僅利用少量機器人軌跡在 RLBench 等 3D 通用機械臂操作任務集合中獲得優異的性能。

通常,學習具身策略往往需要結構化的機器人數據集來進行強化學習或模仿學習訓練,數據集中包含機器人觀測、動作、獎勵或者專家狀態-動作。然而,針對特定場景的機器人數據往往非常有限,難以覆蓋完整的狀態-動作空間決策,在相似場景和真實世界的策略泛化中存在較大困難。一個直覺的解決方案是,能否利用在其他領域的大規模視頻數據,特別是人類操作視頻來幫助具身決策?人類在現實場景中第一視角的物體操作視頻和機器人操作任務具有高度的相似性,包含了物理世界的交互信息,並具有多元的任務場景和複雜的視覺背景,可以幫助具身策略學習物體操作的先驗知識。

近期部分工作開始利用人類操作數據去輔助策略學習,然而,現有研究主要集中於從人類視頻中提取圖像表徵或者Affordance區域,侷限在圖像的特徵表示而忽略了人類操作視頻中蘊含的豐富時序信息的行爲信息,不同於現有方法,本研究提出構建基於視頻預測(video prediction)來獲取智能體對未來軌跡的估計,同時通過機器人數據獲得可執行動作的智能體,挖掘在人類操作視頻和機器人數據上統一的行爲模式。爲了有效利用大量人類數據,設計了預訓練(pre-traiining)和微調(fine-tuning)的框架,前者可以 scaling up 到大規模的視頻數據集,後者可以利用少量數據快速遷移至下游任務。整體框架如下圖所示。

圖1 VPDD 總體思路

方法致力於從三個方面利用人類操作數據解決具身高效策略的學習問題:

1.統一視頻編碼

爲了從數據分佈極廣的各種類型的視頻數據中提取有效的信息輸入給神經網絡進行學習,設計視頻自編碼器 Video VQ-VAE 把視頻數據壓縮成離散的隱向量,隱向量從訓練得到的 VQ-VAE 的碼本中提取。這樣,對於人類視頻或機器人視頻,算法都可以用同一個碼本中的不同隱向量表徵,不僅統一了特徵空間去掉了冗餘信息,也減少了模型學習的難度。見下圖 Stage 1所示。

圖2 三階段學習框架

2.基於視頻的預訓練

在預訓練階段,想要從大量視頻中提取與物理交互有關的普適知識,設計了自監督學習實現該目標。給定一段歷史視頻和文本作爲 prompts,利用大規模擴散模型預測未來視頻 token 序列。當模型能很好地理解交互模式並預測到準確的未來軌跡時,智能體能夠對未來可能發生的行爲進行預估,從而用該信息去指導下游任務的決策過程。

爲了處理複雜和信息量豐富的離散視頻編碼,並且支持提出的預訓練及微調的兩階段訓練模式,我們採用表達力極強的離散擴散模型(Discrete Diffusion)進行數據建模和學習。不同於適用於連續狀態空間的 Gaussian 擴散模型,離散擴散模型通過 state masking 策略來進行加噪和去噪。VQ-VAE 編碼和擴散模型擴散過程可見下圖:

圖3 離散擴散模型,from paper “Unleashing Transformers: Parallel Token Prediction with Discrete Absorbing Diffusionfor Fast High-Resolution Image Generation from Vector-Quantized Codes”

在預訓練階段,爲了減小計算開銷,我們利用 Perceiver Transformer 作爲擴散模型的 backbone;在微調階段,由於只需要生成低維的 action,我們使用 GPT2 Transformer 作爲 backbone,以便於在小規模機器人數據集中進行策略學習。

3.機器人策略學習

通過從大規模人類數據集中學習到的普遍視頻預測模式,在下游機器人任務中僅需要依賴少量機器人數據就能夠快速的學習策略。具體的,在微調階段利用有限的機器人數據集,包括視頻和動作,可以輸出可執行動作的決策智能體。

方法在單視角視覺觀測的的 Meta-World 任務集合和使用多視角觀測的 3D 操作任務集合 RLBench 中評估有效性。結果發現,論文提出的方法方法可以成功預測比較準確的未來運動軌跡,無論是單視角還是多視角,這些都通過一個離散擴散模型生成。下面顯示了在關鍵幀附近的相鄰視頻預測結果。

在具體的決策任務上,本文方法也明顯優於以前的方法。重要的是,方法僅需要少量的數據集就可以在各種機械臂抓取任務上達到比較高的成功率,在 Meta-World 和 RLBench 上的實驗結果如下:

該論文創新性地提出了 VPDD,一種利用離散擴散模型生成未來運動軌跡(視頻)並將預訓練學習的知識快速遷移至決策中的方法。VPDD 可以靈活地處理各種視頻輸入的機械臂操作任務,包括單視角相機的 Meta-World(2D 操作)以及多視角相機的 RLBench(3D 操作)。受限於計算資源和模型規模,VPDD 在視頻生成上仍有瑕疵,對於某些樣本可能存在軌跡不連續或者視角不匹配的問題。未來的工作可以在這些方面繼續進行優化。