NeurIPS Spotlight | 決策模型有了全新預訓練範式統一框架

論文一作爲之江實驗室研究專家、香港中文大學在職博士李藍青，指導老師爲香港中文大學計算機科學與工程系王平安 (Pheng Ann Heng) 教授。同濟大學碩士生張海爲共同第一作者，指導老師趙君嶠教授爲論文通訊作者。

現如今，以 GPT 爲代表的大語言模型正深刻影響人們的生產與生活，但在處理很多專業性和複雜程度較高的問題時仍然面臨挑戰。在諸如藥物發現、自動駕駛等複雜場景中，AI 的自主決策能力是解決問題的關鍵，而如何進行決策大模型的高效訓練目前仍然是開放性的難題。

強化學習（RL）作爲一種經典的時序決策模型的訓練方法，勢必成爲決策大模型訓練及微調的核心技術之一。而由於任務和數據的複雜性，我們希望模型在訓練時能擺脫傳統強化學習與環境在線交互的方式，實現在海量歷史數據中進行離線、多任務的高效學習，這一新範式被稱爲「離線元強化學習 」（Offline Meta-RL）。

近期，圍繞離線元強化學習，來自之江實驗室、香港中文大學、同濟大學等單位的研究團隊提出了全新算法 UNICORN。該方法基於信息論，首次系統性地提出了一套關於強化學習中任務表示學習（task representation learning）的理論框架UNICORN（UNIfied Information Theoretic Framework of Context-Based Offline Meta-ReiNforcement Learning），將現有主流方法利用一個基於任務表徵的通用互信息優化目標進行了統一，並憑藉理論創新和全面的實驗驗證，成爲離線及元強化學習領域的重要里程碑，被人工智能三大頂級會議 NeurIPS 2024 接收爲 Spotlight 文章（中稿率 2.08%）。

問題背景

在經典強化學習中，智能體（agent）通過與外部環境實時交互來收集反饋，在不斷試錯（trial-and-error）中積累經驗進行學習。然而在諸如自動駕駛、疾病治療等現實場景中，試錯帶來的風險往往是無法承受的，從而使人們開始關注如何擺脫與環境的在線交互，僅從歷史數據中進行學習，這一新範式被稱爲「離線強化學習」（offline RL）。

另一方面，複雜多變的真實場景使得智能體處理多任務能力的必要性與日俱增，這種使智能體像人類一樣同時學習多種技能並進行舉一反三的範式被稱作「元強化學習」（meta-RL）。

離線強化學習和元強化學習作爲強化學習的兩個分支，有着各自獨特的優勢。前者由於擺脫了與環境的在線交互，可以重複利用歷史數據進行訓練，具有高安全性、高樣本效率的特點；而後者聚焦多任務及遷移學習，在泛化能力方面表現突出，兩者優勢互補。

於是在 2021 年前後，人們開始嘗試結合兩種範式來訓練更加強大的智能體，其中主流的一類方法被稱爲「基於語境的離線元強化學習」（Context-Based Offline Meta-RL，COMRL)，其核心思想是將當前任務的表徵作爲額外的狀態信息，訓練一個適用於任意任務 / 環境的通用策略（universal policy）:

在該框架下，如何學習魯棒、有效的任務表徵 Z 成爲核心問題，而其中最重要的挑戰是語境偏移（context shift）。由於智能體的訓練數據是離線也就是固定分佈的，但在測試時面臨的任務語境未知且多變，導致訓練和測試集間可能在狀態 - 動作（state-action）維度或者任務維度上存在巨大分佈偏移，這對於模型的魯棒性、泛化性提出了極高要求。

針對上述問題，現有主流方法例如 FOCAL［1］、CORRO［2］和 CSRO［3］陸續提出了多種優化目標，利用度量學習（metric learning）、對比學習（contrastive learning）等思想進行任務表徵學習：

然而，現有方法主要聚焦於對損失函數的經驗性改進，缺乏針對任務表示學習尤其是語境偏移的系統性理論支持和設計指導。

基於信息論的統一理論框架 UNICORN

UNICORN 的核心創新在於藉助信息論，從數學定義、因果關係分解、中心定理三個層面依次遞進，首次系統性地定義和解構了 COMRL 中的任務表示學習這一問題，並通過嚴格理論證明將現有方法的優化目標進行了統一，由此提出並驗證了兩種新的算法實現，以啓迪未來更多新方法的設計。

1. 任務表示學習的數學定義

2. 因果關係分解

3. 中心定理

該中心定理引申出 2 個重要結論，爲未來 COMRL 領域的新方法設計指明瞭道路：

基於上述洞察，爲了展示 UNICORN 框架的指導意義，通過對 I (Z; M) 的近似，我們提出了兩種新的算法實現：

實驗結果

UNICORN 的廣泛適用性和魯棒性

爲了證明 UNICORN 理論框架的普適性，我們在多種機器人連續控制任務的相關設定下對新提出的兩種方法進行了廣泛的實驗驗證：

1. Behavior IID/OOD (訓練集與測試集的行爲策略採樣於相同分佈 / 不同分佈)

結論：UNICORN 算法在同分布測試集上性能媲美 SoTA，在分佈外測試集上性能顯著優於現有其他方法。

2. 不同質量的數據集表現

結論：UNICORN 算法（尤其無監督版本）在不同質量的數據集上的性能均達到 SoTA。

3. 不同模型架構的可遷移性（應用於 Decision Transformer（DT）的測試結果）

結論：UNICORN 算法在 MLP/Decision Transformer 架構上相比現有方法均呈現明顯優勢，可以作爲即插即用的模塊廣泛應用於其他 RL 算法中。

4. 對於分佈外任務的泛化性

圖左爲分佈外任務的構造方式：以 Ant-Dir 爲例，訓練任務的目標方向採樣自第二、三象限，測試任務分佈於第一、四象限，兩者完全不重疊。圖右爲測試結果：自監督 UNICORN 爲唯一取得正向小樣本遷移（positive few-shot transfer）的算法。

結論：利用無監督 UNICORN 中的自編碼器進行 domain randomization 和 model-based RL，可以將智能體的能力外推至分佈外的任務，這一點是現有其他方法都無法做到的。

UNICORN 的未來展望

爲拓展決策大模型的能力邊界提供理論基礎

UNICORN 爲離線元強化學習提供了統一理論基礎和算法設計準則，對於決策大模型的大規模離線、多任務預訓練及微調，從而進一步拓展決策大模型的能力邊界具有指導意義。該技術有助於解決藥物設計、精準醫療、具身智能等前沿領域面臨的 AI 模型的泛化性、多目標優化、樣本利用率等挑戰，同時，團隊也在探索將 UNICORN 框架進一步推廣到在線強化學習等更多場景中。

參考文獻：

［1］. Lanqing Li, Rui Yang, and Dijun Luo. Focal: Efficient fully-offline meta-reinforcement learning via distance metric learning and behavior regularization. ICLR 2021.

［2］. Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reinforcement learning via contrastive learning. ICML 2022.

［3］.Yunkai Gao, et al. Context shift reduction for offline meta-reinforcement learning. NeurIPS 2023.