現如今，以 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 的分布完全由任務本身的獎勵函數及狀態轉移函數所決定?；诖耍覀兛梢詫ι鲜鲴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. 對于分布外任務的泛化性

圖左為分布外任務的構造方式：以 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.