AIxiv专栏是AI在线发布学术、技术内容的栏目。过去数年，AI在线AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：[email protected]；[email protected]
论文一作为之江实验室研究专家、香港中文大学在职博士李蓝青，指导老师为香港中文大学计算机科学与工程系王平安 (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. 对于分布外任务的泛化性
图左为分布外任务的构造方式：以 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.