現(xiàn)如今，以 GPT 為代表的大語(yǔ)言模型正深刻影響人們的生產(chǎn)與生活，但在處理很多專業(yè)性和復(fù)雜程度較高的問題時(shí)仍然面臨挑戰(zhàn)。在諸如藥物發(fā)現(xiàn)、自動(dòng)駕駛等復(fù)雜場(chǎng)景中，AI 的自主決策能力是解決問題的關(guān)鍵，而如何進(jìn)行決策大模型的高效訓(xùn)練目前仍然是開放性的難題。

強(qiáng)化學(xué)習(xí)（RL）作為一種經(jīng)典的時(shí)序決策模型的訓(xùn)練方法，勢(shì)必成為決策大模型訓(xùn)練及微調(diào)的核心技術(shù)之一。而由于任務(wù)和數(shù)據(jù)的復(fù)雜性，我們希望模型在訓(xùn)練時(shí)能擺脫傳統(tǒng)強(qiáng)化學(xué)習(xí)與環(huán)境在線交互的方式，實(shí)現(xiàn)在海量歷史數(shù)據(jù)中進(jìn)行離線、多任務(wù)的高效學(xué)習(xí)，這一新范式被稱為「離線元強(qiáng)化學(xué)習(xí) 」（Offline Meta-RL）。

近期，圍繞離線元強(qiáng)化學(xué)習(xí)，來自之江實(shí)驗(yàn)室、香港中文大學(xué)、同濟(jì)大學(xué)等單位的研究團(tuán)隊(duì)提出了全新算法 UNICORN。該方法基于信息論，首次系統(tǒng)性地提出了一套關(guān)于強(qiáng)化學(xué)習(xí)中任務(wù)表示學(xué)習(xí)（task representation learning）的理論框架 UNICORN（UNIfied Information Theoretic Framework of Context-Based Offline Meta-ReiNforcement Learning），將現(xiàn)有主流方法利用一個(gè)基于任務(wù)表征的通用互信息優(yōu)化目標(biāo)進(jìn)行了統(tǒng)一，并憑借理論創(chuàng)新和全面的實(shí)驗(yàn)驗(yàn)證，成為離線及元強(qiáng)化學(xué)習(xí)領(lǐng)域的重要里程碑，被人工智能三大頂級(jí)會(huì)議 NeurIPS 2024 接收為 Spotlight 文章（中稿率 2.08%）。

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

問題背景

在經(jīng)典強(qiáng)化學(xué)習(xí)中，智能體（agent）通過與外部環(huán)境實(shí)時(shí)交互來收集反饋，在不斷試錯(cuò)（trial-and-error）中積累經(jīng)驗(yàn)進(jìn)行學(xué)習(xí)。然而在諸如自動(dòng)駕駛、疾病治療等現(xiàn)實(shí)場(chǎng)景中，試錯(cuò)帶來的風(fēng)險(xiǎn)往往是無(wú)法承受的，從而使人們開始關(guān)注如何擺脫與環(huán)境的在線交互，僅從歷史數(shù)據(jù)中進(jìn)行學(xué)習(xí)，這一新范式被稱為「離線強(qiáng)化學(xué)習(xí)」（offline RL）。

另一方面，復(fù)雜多變的真實(shí)場(chǎng)景使得智能體處理多任務(wù)能力的必要性與日俱增，這種使智能體像人類一樣同時(shí)學(xué)習(xí)多種技能并進(jìn)行舉一反三的范式被稱作「元強(qiáng)化學(xué)習(xí)」（meta-RL）。

離線強(qiáng)化學(xué)習(xí)和元強(qiáng)化學(xué)習(xí)作為強(qiáng)化學(xué)習(xí)的兩個(gè)分支，有著各自獨(dú)特的優(yōu)勢(shì)。前者由于擺脫了與環(huán)境的在線交互，可以重復(fù)利用歷史數(shù)據(jù)進(jìn)行訓(xùn)練，具有高安全性、高樣本效率的特點(diǎn)；而后者聚焦多任務(wù)及遷移學(xué)習(xí)，在泛化能力方面表現(xiàn)突出，兩者優(yōu)勢(shì)互補(bǔ)。

于是在 2021 年前后，人們開始嘗試結(jié)合兩種范式來訓(xùn)練更加強(qiáng)大的智能體，其中主流的一類方法被稱為「基于語(yǔ)境的離線元強(qiáng)化學(xué)習(xí)」（Context-Based Offline Meta-RL，COMRL)，其核心思想是將當(dāng)前任務(wù)的表征作為額外的狀態(tài)信息，訓(xùn)練一個(gè)適用于任意任務(wù) / 環(huán)境的通用策略（universal policy）:

在該框架下，如何學(xué)習(xí)魯棒、有效的任務(wù)表征 Z 成為核心問題，而其中最重要的挑戰(zhàn)是語(yǔ)境偏移（context shift）。由于智能體的訓(xùn)練數(shù)據(jù)是離線也就是固定分布的，但在測(cè)試時(shí)面臨的任務(wù)語(yǔ)境未知且多變，導(dǎo)致訓(xùn)練和測(cè)試集間可能在狀態(tài) - 動(dòng)作（state-action）維度或者任務(wù)維度上存在巨大分布偏移，這對(duì)于模型的魯棒性、泛化性提出了極高要求。

針對(duì)上述問題，現(xiàn)有主流方法例如 FOCAL［1］、CORRO［2］和 CSRO［3］陸續(xù)提出了多種優(yōu)化目標(biāo)，利用度量學(xué)習(xí)（metric learning）、對(duì)比學(xué)習(xí)（contrastive learning）等思想進(jìn)行任務(wù)表征學(xué)習(xí)：

• FOCAL

• CORRO

• CSRO

然而，現(xiàn)有方法主要聚焦于對(duì)損失函數(shù)的經(jīng)驗(yàn)性改進(jìn)，缺乏針對(duì)任務(wù)表示學(xué)習(xí)尤其是語(yǔ)境偏移的系統(tǒng)性理論支持和設(shè)計(jì)指導(dǎo)。

基于信息論的統(tǒng)一理論框架 UNICORN

UNICORN 的核心創(chuàng)新在于借助信息論，從數(shù)學(xué)定義、因果關(guān)系分解、中心定理三個(gè)層面依次遞進(jìn)，首次系統(tǒng)性地定義和解構(gòu)了 COMRL 中的任務(wù)表示學(xué)習(xí)這一問題，并通過嚴(yán)格理論證明將現(xiàn)有方法的優(yōu)化目標(biāo)進(jìn)行了統(tǒng)一，由此提出并驗(yàn)證了兩種新的算法實(shí)現(xiàn)，以啟迪未來更多新方法的設(shè)計(jì)。

1. 任務(wù)表示學(xué)習(xí)的數(shù)學(xué)定義

由于在 COMRL 中，數(shù)據(jù)覆蓋多個(gè)任務(wù)，假設(shè)這些任務(wù)采樣于一個(gè)特定的任務(wù)分布，我們將遵從該分布的任務(wù)變量定義為隨機(jī)變量，與之對(duì)應(yīng)的數(shù)據(jù)樣本和任務(wù)表示變量分別定義為、，則三者構(gòu)成如下馬爾可夫鏈：

任務(wù)表示學(xué)習(xí)的數(shù)學(xué)定義：COMRL 中的任務(wù)表示學(xué)習(xí)意在找到一個(gè)數(shù)據(jù)相對(duì)于任務(wù)變量的充分統(tǒng)計(jì)量（sufficient statistics） 。

2. 因果關(guān)系分解

在 COMRL 中，原則上，任務(wù)變量 與進(jìn)行采樣的行為策略（behavior policy）應(yīng)該相互獨(dú)立。我們發(fā)現(xiàn)如果將數(shù)據(jù)樣本看作包含狀態(tài)、動(dòng)作、獎(jiǎng)勵(lì)的四元組 (s, a, s’, r)，則其中前兩項(xiàng)（s, a）的分布主要與行為策略相關(guān)，而在給定 s、a 的情況下，s’、r 的分布完全由任務(wù)本身的獎(jiǎng)勵(lì)函數(shù)及狀態(tài)轉(zhuǎn)移函數(shù)所決定。基于此，我們可以對(duì)上述馬爾可夫鏈進(jìn)行如下拆解：

該分解反映了由 s’, r 構(gòu)成的 X_t 包含了與任務(wù)變量、的絕大部分因果關(guān)系，而由 s, a 構(gòu)成的 X_b 由于與行為策略高度相關(guān)，含有大量虛假相關(guān)性（spurious correlation）。同時(shí)在數(shù)學(xué)上，我們發(fā)現(xiàn)該因果關(guān)系分解正好可以對(duì)應(yīng)如下的互信息分解：

基于上述分析，我們將右邊第一項(xiàng)命名為主因果關(guān)系（primary causality），第二項(xiàng)稱為次因果關(guān)系（lesser causality）。

3. 中心定理

本文從數(shù)學(xué)上嚴(yán)格證明了如下不等式（中心定理）：

該中心定理引申出 2 個(gè)重要結(jié)論，為未來 COMRL 領(lǐng)域的新方法設(shè)計(jì)指明了道路：

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

基于上述洞察，為了展示 UNICORN 框架的指導(dǎo)意義，通過對(duì) I (Z; M) 的近似，我們提出了兩種新的算法實(shí)現(xiàn)：

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

• 自監(jiān)督 UNICORN：將 I (Z; M) 的求解近似為數(shù)據(jù)重建（生成式）+ 度量學(xué)習(xí)（對(duì)比式）

實(shí)驗(yàn)結(jié)果：UNICORN 的廣泛適用性和魯棒性

為了證明 UNICORN 理論框架的普適性，我們?cè)诙喾N機(jī)器人連續(xù)控制任務(wù)的相關(guān)設(shè)定下對(duì)新提出的兩種方法進(jìn)行了廣泛的實(shí)驗(yàn)驗(yàn)證：

1. Behavior IID/OOD (訓(xùn)練集與測(cè)試集的行為策略采樣于相同分布 / 不同分布)

結(jié)論：UNICORN 算法在同分布測(cè)試集上性能媲美 SoTA，在分布外測(cè)試集上性能顯著優(yōu)于現(xiàn)有其他方法。

2. 不同質(zhì)量的數(shù)據(jù)集表現(xiàn)

結(jié)論：UNICORN 算法（尤其無(wú)監(jiān)督版本）在不同質(zhì)量的數(shù)據(jù)集上的性能均達(dá)到 SoTA。

3. 不同模型架構(gòu)的可遷移性（應(yīng)用于 Decision Transformer（DT）的測(cè)試結(jié)果）

結(jié)論：UNICORN 算法在 MLP/Decision Transformer 架構(gòu)上相比現(xiàn)有方法均呈現(xiàn)明顯優(yōu)勢(shì)，可以作為即插即用的模塊廣泛應(yīng)用于其他 RL 算法中。

4. 對(duì)于分布外任務(wù)的泛化性

圖左為分布外任務(wù)的構(gòu)造方式：以 Ant-Dir 為例，訓(xùn)練任務(wù)的目標(biāo)方向采樣自第二、三象限，測(cè)試任務(wù)分布于第一、四象限，兩者完全不重疊。圖右為測(cè)試結(jié)果：自監(jiān)督 UNICORN 為唯一取得正向小樣本遷移（positive few-shot transfer）的算法。

結(jié)論：利用無(wú)監(jiān)督 UNICORN 中的自編碼器進(jìn)行 domain randomization 和 model-based RL，可以將智能體的能力外推至分布外的任務(wù)，這一點(diǎn)是現(xiàn)有其他方法都無(wú)法做到的。

UNICORN 的未來展望：為拓展決策大模型的能力邊界提供理論基礎(chǔ)

UNICORN 為離線元強(qiáng)化學(xué)習(xí)提供了統(tǒng)一理論基礎(chǔ)和算法設(shè)計(jì)準(zhǔn)則，對(duì)于決策大模型的大規(guī)模離線、多任務(wù)預(yù)訓(xùn)練及微調(diào)，從而進(jìn)一步拓展決策大模型的能力邊界具有指導(dǎo)意義。該技術(shù)有助于解決藥物設(shè)計(jì)、精準(zhǔn)醫(yī)療、具身智能等前沿領(lǐng)域面臨的 AI 模型的泛化性、多目標(biāo)優(yōu)化、樣本利用率等挑戰(zhàn)，同時(shí)，團(tuán)隊(duì)也在探索將 UNICORN 框架進(jìn)一步推廣到在線強(qiáng)化學(xué)習(xí)等更多場(chǎng)景中。

參考文獻(xiàn)：

［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.