NeurIPS Spotlight | 基於信息論，決策模型有了全新預訓練範式統一框架

AIxiv專欄是機器之心發佈學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯繫報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

論文一作為之江實驗室研究專家、香港中文大學在職博士李藍青，指導老師為香港中文大學計算機科學與工程系王平安 (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%）。

論文標題：Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning
論文鏈接：https://openreview.net/pdf?id=QFUsZvw9mx
項目地址：https://github.com/betray12138/UNICORN

問題背景

在經典強化學習中，智能體（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）等思想進行任務表徵學習：

FOCAL

CORRO

CSRO

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

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

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

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

由於在 COMRL 中，數據覆蓋多個任務，假設這些任務采樣於一個特定的任務分佈，我們將遵從該分佈的任務變量定義為隨機變量

、

，則三者構成如下馬爾可夫鏈：

，與之對應的數據樣本和任務表示變量分別定義為

任務表示學習的數學定義：COMRL 中的任務表示學習意在找到一個數據

相對於任務變量

的充分統計量（sufficientstatistics）

。

2.因果關係分解

在 COMRL 中，原則上，任務變量與進行采樣的行為策略（behavior policy）應該相互獨立。我們發現如果將數據樣本

看作包含狀態、動作、獎勵的四元組(s,a,s’,r)，則其中前兩項（s,a）的分佈主要與行為策略相關，而在給定s、a的情況下，s’、r的分佈完全由任務本身的獎勵函數及狀態轉移函數所決定。基於此，我們可以對上述馬爾可夫鏈進行如下拆解：

該分解反映了由 s’, r 構成的 X_t 包含了與任務變量

、

的絕大部分因果關係，而由s,a構成的X_b由於與行為策略高度相關，含有大量虛假相關性（spuriouscorrelation）。同時在數學上，我們發現該因果關係分解正好可以對應如下的互信息分解：

基於上述分析，我們將右邊第一項命名為主因果關係（primary causality），第二項稱為次因果關係（lesser causality）。

3.中心定理

本文從數學上嚴格證明了如下不等式（中心定理）：

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

主因果關係作為不等式下界，雖不包含虛假相關性，但缺失部分因果相關性；而主因果關係與次因果關係之和作為上界，在囊括全部因果相關性的同時引入了虛假相關性。因此，理論上，一個有且僅包含全部因果相關性的「最優優化目標」應該介於兩者之間，而 I (Z; M) 剛好滿足這一要求。基於上述觀察，我們提出 I (Z; M) 應該作為任務表示學習優化目標的金標準（ground truth），其天然具有對於語境偏移的魯棒性。
現有主流方法本質都是在優化 I (Z; M) 的一個近似，例如 FOCAL、CORRO、CSRO 分別優化的是其上界、下界和兩者的線性插值。因此，找到更好的 I (Z; M) 近似方法將成為 COMRL 領域未來發展的關鍵方向。

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

有監督 UNICORN：將 I (Z; M) 的求解近似為離散化的分類問題

自監督 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.