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 中的任務表示學習意在找到一個資料相對於任務變數
的充分統計量（sufficient statistics）
。
2. 因果關係分解
在 COMRL 中，原則上，任務變數 與進行取樣的行為策略（behavior policy）應該相互獨立。我們發現如果將資料樣本看作包含狀態、動作、獎勵的四元組 (s, a, s’, r)，則其中前兩項（s, a）的分佈主要與行為策略相關，而在給定 s、a 的情況下，s’、r 的分佈完全由任務本身的獎勵函式及狀態轉移函式所決定。基於此，我們可以對上述馬爾可夫鏈進行如下拆解：
該分解反映了由 s’, r 構成的 X_t 包含了與任務變數、
的絕大部分因果關係，而由 s, a 構成的 X_b 由於與行為策略高度相關，含有大量虛假相關性（spurious correlation）。同時在數學上，我們發現該因果關係分解正好可以對應如下的互資訊分解：
基於上述分析，我們將右邊第一項命名為主因果關係（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. 對於分佈外任務的泛化性