模型、專家智慧體和資料集都已開源。
隨著 Llama 3 釋出,未來大模型的引數量已飆升至驚人的 4000 億。儘管每週幾乎都有一個聲稱效能超強的大模型出來炸場,但 AI 應用還在等待屬於它們的「ChatGPT 時刻」。其中,AI 智慧體無疑是最被看好的賽道。
就連吳恩達都說,GPT-4 加上 AI 智慧體,可能提前達到 GPT-5 的效果。
不過,我們熟知的智慧體往往有點「偏科」。例如,第一個 AI 軟體工程師 Devin,專精於程式碼。會打遊戲的智慧體往往也只能在某一個遊戲裡秀操作。尋找一個能夠同時擅長多個領域,並能在其中無縫切換的通用模型仍是機器學習研究中的一個關鍵目標。
為了解決這個問題,研究者們對於智慧體如何結合計算機視覺(CV)和自然語言處理(NLP)任務進行了廣泛探索,但將強化學習(RL)任務整合進來的研究相對較少。這是由於 RL 任務本質上是異質的,這使得將 RL 任務與對話和影像識別等其他任務結合起來更加困難。這要求智慧體能融會貫通不同領域任務中的不同模態、任務複雜性和資料型別。要達到全能型智慧體,主要需要解決以下問題:(1)如何設計一個能夠處理多種資料型別和模態的統一模型結構?(2)如何有效地平衡不同任務的學習進度和優先順序?(3)如何確保智慧體制定合適的學習目標,以避免不同任務之間的干擾和負向遷移?
來自 Hugging Face、法國國家資訊與自動化研究所(INRIA)和波爾多大學的四位研究者提出了智慧體中的「六邊形戰士」——Jack of All Trades (JAT)。JAT 是一個基於 Transformer 的多模態通用強化學習智慧體框架。在此框架下,智慧體能夠透過同一套引數應對不同複雜度的多種任務,化身既會打遊戲,又能控制機器人的全能高手。論文同時釋出了大量 RL 智慧體與 JAT 資料集。這是首個用於通用智慧體訓練的資料集 JAT 資料集,包含了由專家智慧體收集的數十萬條軌跡。
論文名稱:《Jack of All Trades, Master of Some, a Multi-Purpose Transformer Agent》
論文連結:https://huggingface.co/papers/2402.09844
程式碼連結:https://github.com/huggingface/jat
專案連結:https://huggingface.co/jat-project/jat
資料集:https://huggingface.co/datasets/jat-project/jat-dataset
模型架構
JAT 的核心結構基於 Transformer,使用了 EleutherAI 的 GPT-Neo 實現。JAT 最大的創新點在於其嵌入機制,從本質上解決了資料型別不同的問題。JAT 模型將觀察嵌入與其對應的獎勵值和動作嵌入交錯排列,形成一個序列。
圖 1.JAT 網路架構。對於序列中的決策任務,一方面輸入觀察嵌入與獎勵值,另一方面行動嵌入被編碼並被交錯放置。模型使用因果掩碼自迴歸地生成下一個嵌入,並根據預期的模態進行解碼。
因此,每個嵌入要麼對應一個與獎勵相關聯的觀察嵌入,要麼對應一個動作嵌入。JAT 如何進一步對這些資訊進行編碼呢?這要取決於資料的型別。如果觀察嵌入或動作嵌入的資料型別是影像,那麼 JAT 將使用 CNN。如果是連續向量,則使用線性層。如果是離散值,則使用線性投影層。模型的輸出也遵循相同的邏輯,具體取決於預測目標的資料型別。預測基於因果推理進行,將觀察嵌入向後移動一個時間步,確保智慧體可以根據所有先前的觀察和動作嵌入來預測下一個動作嵌入。
這種嵌入設計讓研究團隊在訓練智慧體執行 NLP 和 CV 任務時興致盎然。對於和文字相關的任務,作者讓 JAT 模型採用 GPT-2 的分詞策略,將文字轉換為一個整數序列,然後透過一個查詢表對映到一個嵌入向量序列。對於和影像有關的任務,JAT 模型將選擇 ViT 方法,將影像切割成小塊後,透過線性層轉換為嵌入向量序列。JAT 模型再將影像和文字的向量序列拼接在一起,形成一個統一的序列,輸入到 Transformer 中。
考慮到資料的模態變來變去,JAT 如何計算損失函式呢?它將針對每種模態分別計算 loss。對於影像和連續值,它使用均方誤差(MSE)損失。對於離散值,它使用交叉熵損失。最終的損失是序列中每種元素損失的平均值。那麼,這是否意味著 JAT 在預測動作嵌入和觀察嵌入時的權重是相同的呢?實際上不是,在此後的章節中將一步探討這個問題。
實驗結果
研究團隊共採用了 157 個訓練任務來 JAT 評估。他們將這些任務分為 10 類,並記錄了 JAT 的總獎勵值。
JAT 模型在最終的檢查點上達到了 65.8% 的專家得分,說明 JAT 能夠在非常廣泛的任務上達到專家水平。以下具體列出了 JAT 在四個常見的智慧體訓練環境中的得分:
對於 Atari 57,應用 JAT 模型的智慧體實現了專家分數的 14.1%,這相當於人類表現的 37.6%。Atari 影片遊戲廣泛被用作評估和開發強化學習演算法的基準環境,其中《吃豆人》是一款標誌性遊戲。在這一系列的 21 款遊戲中,JAT 智慧體的表現已經超越了人類玩家。值得注意的是, JAT 只用了單一網路就在所有 Atari 影片遊戲中達到了這種水平;
對於 BabyAI,應用 JAT 模型的智慧體達到了專家分數的 99.0%,只有一個任務的表現未能超過專家水平的 50%;
對於 Meta-World,應用 JAT 模型的智慧體達到了專家分數的 65.5%;
對於 MuJoCo,應用 JAT 模型的智慧體達到了專家分數的 84.8%。
JAT 智慧體在 Atari 57 基線上和人類表現的對比
JAT 智慧體在小遊戲中的表現
這些 JAT 智慧體都可以透過專案主頁下載,進一步測試和體驗。更多細節請參閱論文原文。
專家智慧體和 JAT 資料集
專家策略
傳統的強化學習往往在單一環境中尋找專家策略,即在一個特定任務中尋找讓模型表現最優的方法。構建跨領域的多功能智慧體,也離不開這種方法。論文作者選擇了 Atari、BabyAI、Meta-World 和 MuJoCo 一系列性質不同,難度各異的訓練環境,直到訓練出表現最好的智慧體。這一系列採用 JAT 框架的專家智慧體已經在專案主頁上釋出。
JAT 資料集
論文作者隨論文同步釋出了 JAT 資料集,這是首個針對通用智慧體訓練的專項資料集。其中包含了數十萬條由上述專家智慧體收集的軌跡資料。使用起來也很方便,可以像載入 Hugging Face 平臺上的其他資料集一樣簡單。以下是呼叫程式碼示例:
JAT 資料集不僅包含強化學習的資料,還整合了來自維基百科等文字資料集,以及 Oscar、OK-VQA、Conceptual Captions 等針對視覺任務的資料集,提供了更豐富的資料型別選擇。
增加模型預測觀察嵌入的能力
智慧體學得更好更快了
在訓練強化學習智慧體時,主要目標是使其在未曾遇到的任務中實現獎勵最大化。然而,如果要求智慧體預測未來可能遇到的情境,這一額外任務會促進還是阻礙其學習過程呢?
關於這個問題存在兩種相反的觀點。一方面,學會預判可能會讓智慧體對環境有更深入的理解,從而學得更好更快。另一方面,這可能會分散智慧體對其主要目標的注意力,導致在預測觀察嵌入和行動嵌入時都表現平庸。
為了得到問題的答案,論文作者進行了一個實驗,使用了一個結合了觀察損失和行動損失的損失函式,並透過權重引數 k 來平衡這兩種損失。
研究團隊在 95% 的置信區間內,針對選定任務,測量了預判將如何影響模型學習。每項任務進行了 100 次評估,基於這些評估得到了 k 值的範圍。結果表明,適當選擇 k 值可以顯著提升智慧體的表現。
當 k 值過高(高於 0.5)時,預測觀察嵌入的額外任務阻礙了學習過程。但當 k 值較低時,對學習的影響可以忽略不計,且智慧體的表現與沒有額外預判任務時的表現相似。
研究團隊發現,當 k=0.005 時,存在一個最佳臨界點。這意味著,只要平衡得當,為智慧體增加預測觀察嵌入的任務,實際上可以提高智慧體的學習效率。這一發現對於設計類似的智慧體具有重要意義,突顯了輔助目標在提升智慧體學習效率方面的潛在價值。
未來展望
JAT 專案為通用智慧體研究領域開闢了全新的方向。研究團隊表示目前只是初步探索,以下幾點思路可供未來研究者深入挖掘:
改進資料的質量:儘管填補了之前少有通用智慧體訓練資料集的空缺,JAT 資料集仍處於初級階段。其中的專家軌跡僅來自每個環境中的一名專家智慧體,這可能導致一些誤差。雖然研究團隊已盡力讓智慧體達到最優表現,但某些環境仍具挑戰性。在這些環境中,智慧體仍有很大進步空間。收集到更多資料,訓練更多的專家智慧體,將在很大程度上解決這些問題。
使用離線強化學習:JAT 智慧體是仿照基線一比一地訓練出來的。這意味著,其一,智慧體無法利用次優的軌跡;其二,JAT 智慧體無法超越專家。論文選擇了這種方法是因為它比較簡單,但研究團隊相信,使用離線強化學習可以提高智慧體的效能,同時,實現起來也不會過於複雜。
發揮更智慧的多工取樣策略的全部潛力:目前,JAT 智慧體均勻地從所有任務中取樣資料,但這種方法可能限制了它的全部潛力。透過動態調整取樣率,專注於最具挑戰性的任務,或許也可以加速智慧體的學習過程,並解鎖顯著的效能提升。
參考連結:
https://huggingface.co/blog/jat
https://twitter.com/QGallouedec/status/1782430246957994422