LWAIL:基于隐空间 Wasserstein 距离的对抗式模仿学习
模仿学习(Imitation Learning, IL)作为强化学习的重要分支,通过模仿专家演示实现智能体决策能力的学习,但传统方法普遍依赖大量中高质量的专家演示数据(含动作信息),在实际场景中难以满足。 ICLR 2026 论文《Latent Wasserstein Adversarial Imitation Learning》提出了一种全新的对抗式模仿学习框架LWAIL(Latent Wasserstein Adversarial Imitation Learning),聚焦于仅状态分布匹配,通过在动力学感知的隐空间中计算 Wasserstein 距离,大幅降低了对专家数据的依赖。
LWAIL 通过预训练阶段利用少量随机生成的仅状态数据训练意图条件值函数(ICVF, Intention Conditioned Value Function),学习到捕捉状态空间动力学结构的隐空间;在模仿阶段,基于该隐空间的欧氏距离作为代价函数,结合 Wasserstein 对抗式模仿学习框架,实现仅用 1 条或少数条仅状态专家轨迹即可达到专家级性能。在 MuJoCo 多个环境的实验表明,LWAIL 显著优于现有基于 Wasserstein 距离和对抗式的模仿学习方法,为低数据量、仅状态的模仿学习提供了新的有效解决方案。
原文链接:Latent Wasserstein Adversarial Imitation Learning
代码链接:https://github.com/JackyYang258/LWAIL
沐小含持续分享前沿算法论文,欢迎关注...
一、研究背景与动机
1.1 强化学习与模仿学习的痛点
强化学习(RL)在游戏、大语言模型训练等领域取得了显著成功,但核心依赖精确定义的奖励信号—— 在机器人控制等实际场景中,奖励信号难以设计,且需针对不同环境定制,限制了其落地。
模仿学习作为替代方案,直接从专家演示中学习策略,无需显式奖励设计,但其面临两大核心挑战:
- 专家动作缺失:实际场景中常仅能获取专家的状态序列(如视频观测),无法得到动作信息,因此从观测中模仿(LfO, Learning from Observations) 成为研究重点;
- 专家数据稀缺:即使是仅状态的专家演示,获取成本也往往极高,亟需能从极少量专家数据中学习的方法。
1.2 现有对抗式模仿学习的局限性
对抗式模仿学习(AIL)是 LfO 的主流方法,通过匹配智能体与专家的状态分布实现策略学习,核心是度量分布间的差异。现有方法主要存在两类问题:
(1)
- 散度的分布覆盖约束
多数 AIL 方法采用 、
、
等
- 散度度量分布差异,但
- 散度要求分布具有相同的支撑集(即非专家分布需完全覆盖专家分布),否则会产生数值误差。实际中低质量非专家数据(如随机数据)难以满足该约束,限制了方法的适用性。
(2)Wasserstein 距离的度量缺陷
为解决 - 散度的问题,Wasserstein 距离(地球移动距离 EMD)被引入模仿学习,其无需分布覆盖约束,且能捕捉状态空间的几何特性。但现有基于 Wasserstein 距离的 AIL 方法(利用 Kantorovich-Rubinstein (KR) 对偶)存在关键缺陷:
- 计算 Wasserstein 距离需先定义状态间的基础度量(ground metric),现有方法普遍使用欧氏距离,但欧氏距离无法捕捉环境的动力学特性(如两个状态欧氏距离近,但因环境动力学限制无法相互到达),会严重误导学习过程;
- 虽有方法尝试使用原始 Wasserstein 距离规避该问题,但引入了新的复杂度,且需要替代函数,实用性不足。
1.3 核心研究问题
基于上述痛点,论文提出核心问题:能否从极少量低质量的仅状态数据中,学习到编码环境动力学的距离度量,使得智能体仅用少数条仅状态专家轨迹就能高效模仿?
为回答该问题,LWAIL 提出两阶段框架:先通过少量随机仅状态数据训练 ICVF 得到动力学感知的隐空间,再在该隐空间中用欧氏距离作为 Wasserstein 距离的基础度量,结合对抗式模仿学习实现高效模仿。
二、预备知识
本节梳理论文涉及的核心基础概念,为 LWAIL 的理解奠定理论基础。
2.1 马尔可夫决策过程(MDP)
MDP 是强化学习的基础框架,定义为元组
,其中:
:状态空间,
为单个状态;
:动作空间,
为单个动作;
- P(s′∣s,a):状态转移概率,描述从状态s执行动作a到s′的概率;
- R(s,a)∈R:即时奖励函数;
∈[0,1):折扣因子,衡量未来奖励的重要性。
智能体的目标是学习策略
,最大化累积折扣奖励
。
论文重点关注状态占用度和状态对占用度,即状态 / 状态对的访问频率:
- 状态占用度:

- 状态对占用度:

2.2 Wasserstein 距离
Wasserstein 距离用于度量两个概率分布间的 “运输成本”,相比 - 散度更鲁棒,且能捕捉空间几何特性。
对于度量空间
(c:M×M→R为状态间的基础度量),两个分布p(x)和q(x)的1-Wasserstein 距离定义为:W1(p,q)=infΠ(p,q)∫M×Mc(x,y)dΠ(x,y)其中Π(p,q)是p和q的联合分布,满足边缘分布为p和q,直观上表示将分布p的 “质量” 移动到q的最小总成本。
实际中常用KR 对偶形式计算 Wasserstein 距离,避免直接求解复杂的优化问题:

其中
表示f是1-Lipschitz 函数(即
)。
关键缺陷:为保证 Lipschitz 性,现有方法普遍通过梯度正则化实现,这使得基础度量 被限制为欧氏距离,无法捕捉环境动力学。
2.3 意图条件值函数(ICVF)
ICVF 是一种嵌入学习算法,通过无动作、无奖励的仅状态数据学习能捕捉环境动力学的状态表示,核心是学习增广值函数
,定义为:从状态 出发,采用针对目标意图(目标状态)
的最优策略
,未来访问状态
的未归一化状态占用度,形式化表示为:

其中
为指示函数(满足条件为 1,否则为 0),作为伪奖励标签。
ICVF 的价值函数被设计为分解形式,以提取动力学感知的状态嵌入:
![]()
:状态嵌入,将原始状态映射到 维隐空间,是 LWAIL 的核心;
:意图矩阵,编码目标意图 ;
:结果表示,编码未来状态 。
ICVF 可通过离线 RL 算法(如 IQL) 从少量随机仅状态数据中训练,无需专家动作和显式奖励,这是 LWAIL 能利用低质量数据的关键。
三、LWAIL 核心框架
LWAIL 的核心是将动力学感知的 ICVF 隐空间与 Wasserstein 对抗式模仿学习结合,解决传统 Wasserstein AIL 的基础度量缺陷,实现仅状态、少数据的高效模仿。框架分为预训练阶段和模仿阶段,整体流程如图 1 所示。

3.1 核心目标:Wasserstein 对抗式状态对占用度匹配
LWAIL 的核心目标是最小化智能体与专家的状态对占用度分布间的 1-Wasserstein 距离,因为状态对
能更好地捕捉环境的动力学转移特性,相比单独的状态占用度更适合模仿学习。
形式化目标为:
![]()
其中 是智能体策略
的状态对占用度,
是专家的经验状态对占用度。
结合 KR 对偶形式和策略优化,将上述目标转化为双层优化问题(对抗式学习的核心):
![]()
其中:
:兼具三重角色 ——KR 对偶函数、对抗式学习的判别器(区分智能体和专家的状态对)、强化学习的奖励函数;
- 对智能体策略
,目标简化为
,即通过 RL 算法最大化 ,对应奖励
。
3.2 动力学感知的基础度量:ICVF 隐空间
传统方法中,KR 对偶的 Lipschitz 约束迫使基础度量 为欧氏距离,无法捕捉动力学。LWAIL 的核心创新是:将原始状态映射到 ICVF 训练的隐空间,用隐空间的欧氏距离作为 Wasserstein 距离的基础度量,即:
![]()
其中
是 ICVF 学习的状态嵌入。
理论支撑:定理 3.1
论文证明了 ICVF 嵌入与状态对占用度的线性关系,为该度量的有效性提供理论依据:
定理 3.1:在近确定性 MDP 且
<1的条件下,对于 ICVF 中针对目标
训练收敛的任意策略
,存在向量 η,使得对任意相邻状态对 (s,s′),有
),即状态对占用度近似为 ICVF 嵌入的线性组合。
直观解释:ICVF 的价值函数
、智能体的状态对占用度
、专家的状态对占用度
均建模了 “从 到
的访问概率”,三者共享由
编码的动力学线性结构,因此在 ICVF 隐空间中学习能对齐 Wasserstein 优化与环境动力学,降低学习难度。
实证支撑:t-SNE 可视化
论文通过 t-SNE 将原始状态空间和 ICVF 隐空间的轨迹可视化(图 2),颜色表示状态的真实奖励(绿色越高),连线表示轨迹中的相邻状态。结果表明:ICVF 隐空间能更好地捕捉状态间的动力学关系,相同轨迹在隐空间中呈现出与奖励更相关的连续结构,而原始状态空间的分布杂乱,无法反映动力学。

3.3 LWAIL 两阶段算法细节
LWAIL 严格分为预训练和模仿两个阶段,预训练阶段学习动力学感知的隐空间,模仿阶段在隐空间中进行对抗式模仿学习,且预训练的 ICVF 嵌入在模仿阶段完全冻结,避免过拟合专家数据。
3.3.1 预训练阶段
输入:少量随机仅状态转移数据 (可由随机策略采集,仅需 1% 的在线交互数据)、空的重放缓冲区
;
输出:冻结的 ICVF 状态嵌入 、预训练的判别器
。
步骤 1:采集随机数据:用随机初始化的策略与环境交互,采集状态对转移数据 ,存入重放缓冲区
(若已有随机数据集可跳过);
步骤 2:训练 ICVF 获取嵌入:用 IQL 算法通过式 (11) 的损失训练 ICVF 的价值函数
,从分解形式中提取状态嵌入 并冻结;
![]()
步骤 3:预训练判别器 :以隐空间的状态对
为输入,用随机策略的状态对占用度 和专家的
训练判别器f,目标为:
![]()
采用 WGAN-GP 的梯度惩罚保证f的 1-Lipschitz 性。
3.3.2 模仿阶段
输入:冻结的 ICVF 嵌入、预训练的判别器
、仅状态专家数据集
、重放缓冲区
、初始策略
;
输出:模仿专家的策略 。
核心是将原始状态替换为隐空间嵌入,在隐空间中求解 3.1 节的双层优化问题,目标为:
![]()
具体步骤(对应算法 1):
- 交互采集数据:用当前策略
与环境交互,采集轨迹
; - 生成伪奖励:计算隐空间状态对的伪奖励
,其中 为 sigmoid 函数,将奖励归一化到 [0,1],提升 RL 训练的稳定性;
- 存储数据:将
存入重放缓冲区 ;
- 更新判别器:每隔固定步长,用B中的智能体数据和专家数据更新判别器f(保持 Lipschitz 性);
- 更新策略:从
中采样小批量数据,用TD3 算法更新策略
(TD3 为离线 RL 算法,稳定性好,适合对抗式学习的奖励波动);
- 迭代:重复步骤 1-5,直到策略收敛。
3.3.3 算法 1:LWAIL 伪代码

四、实验设计与结果分析
论文通过多类环境、多维度对比、大量消融实验验证 LWAIL 的有效性,核心回答三个问题:
- LWAIL 学习的奖励与真实奖励的差异?
- LWAIL 在复杂连续控制环境中的性能如何?
- ICVF 嵌入对 LWAIL 的贡献有多大?LWAIL 对环境噪声的鲁棒性如何?
实验代码开源于:https://github.com/JackyYang258/LWAIL
4.1 实验设置
4.1.1 实验环境
- 简单导航环境:Maze2D-umaze-v0(2D 点质量迷宫导航,验证奖励学习的有效性);
- 连续控制环境:MuJoCo 的 4 个经典环境(Hopper、HalfCheetah、Walker2D、Ant),均为高维连续动作的机器人运动任务,验证复杂环境的性能;
- 复杂导航环境:Maze2d-medium/large、Antmaze-umaze-v2,引入初始状态高斯噪声,验证泛化能力;
- 其他环境:DeepMind Control 的 Ball-in-Cup,验证跨环境的适用性。
4.1.2 数据集设置
- 专家数据:仅使用1 条仅状态专家轨迹(D4RL 专家数据集),模拟数据稀缺场景;
- 预训练数据:10K 随机状态对转移数据(仅为在线交互数据的 1%),模拟低质量数据场景;
- 所有方法均使用相同的专家数据和预训练数据,保证对比公平。
4.1.3 对比基线
将基线分为 4 类,覆盖经典 IL、Wasserstein-based IL、LfO、离线到在线 IL,部分基线需额外专家动作(标注于结果中):
- 经典 IL:BC(行为克隆)、GAIL、AIRL;
- Wasserstein-based IL:PWIL、WDAIL、IQlearn;
- LfO(仅状态):BCO、GAIfO、DIFO、LS-IQ、DiffAIL、OPOLO、DACfO;
- 离线到在线 IL:OLLIE。
4.1.4 评价指标
采用D4RL 标准化奖励(越高越好),专家的标准化奖励为 100 左右,结果为 5 次独立实验的均值 ± 标准差。
4.2 简单环境实验:Maze2D-umaze-v0
验证 LWAIL 学习的奖励是否具有动力学感知特性,结果如图 3 所示。

关键结论:ICVF 学习的隐空间使得 LWAIL 的奖励能捕捉迷宫的动力学约束,避免智能体向墙壁等不可达区域探索,而真实稀疏奖励因延迟性导致收敛慢,证明了 LWAIL 奖励学习的有效性。
4.3 连续控制实验:MuJoCo
这是论文的核心对比实验,验证 LWAIL 在复杂连续控制任务中对基线的超越性,学习曲线如图 4,最终结果如表 1。

表 1 MuJoCo 环境性能对比

关键结论:
- LWAIL 的平均性能达到 99.07,显著优于所有基线,甚至超越了需要额外专家动作的 WDAIL;
- 最优的仅状态基线 DIFO 平均性能为 92.51,LWAIL 提升约 7%,且方差更小,训练更稳定;
- 经典 IL 方法(如 BC、GAIL)因仅依赖少量专家数据,性能极差,证明了 LWAIL 对少数据场景的适应性。
4.4 复杂导航环境实验:Maze2d/Antmaze
验证 LWAIL 在有噪声、高难度导航任务中的泛化能力,引入初始状态高斯噪声(标准差 0.2/0.5),迫使智能体从专家演示中未出现的状态出发,结果如表 2、表 3。
表 2 Maze2d-medium/large 带初始噪声的性能

表 3 Antmaze-umaze-v2 性能

关键结论:
- ICVF 嵌入是泛化的核心:无 ICVF 的方法在噪声下性能骤降为负数,而 LWAIL 即使在噪声 0.5 时仍保持高性能,证明 ICVF 隐空间捕捉的动力学能让智能体从陌生状态恢复并找到最优路径;
- LWAIL 在 Antmaze 中性能与最优基线 SMODICE 接近,且无需调整超参数,而 SMODICE 需针对不同散度调整超参数,证明 LWAIL 的鲁棒性。
4.5 消融实验:验证 LWAIL 的核心组件
论文设计了大量消融实验,验证 ICVF 嵌入、下游 RL 算法、噪声鲁棒性等核心组件的重要性,核心结果如下。
4.5.1 ICVF 嵌入的有效性
对比 ICVF 与其他对比学习嵌入(CURL、PW-DICE),以及无嵌入的情况:表 4 不同嵌入方法的消融

结论:ICVF 嵌入显著优于其他对比学习嵌入,无嵌入时性能大幅下降(尤其是 Walker2D 和 Ant),证明 ICVF 是为 Wasserstein AIL 量身设计的动力学感知嵌入,更适合少数据、仅状态的模仿学习。
4.5.2 环境噪声的鲁棒性
在 MuJoCo 环境中为动作引入高斯噪声(标准差 0.1/0.2/0.5),模拟真实场景的随机动力学:表 5 LWAIL 在带动作噪声的 MuJoCo 中的性能

结论:LWAIL 在噪声 0.2 以内保持高鲁棒性,平均性能仅下降约 3.5%;即使噪声达到 0.5,Hopper 仍保持近 108 的性能,证明 LWAIL 对适度的环境随机 ness 具有良好的适应性。
4.5.3 ICVF 嵌入对其他基线的提升
将 ICVF 嵌入应用到其他基线(WDAIL、IQlearn、GAIL),验证 ICVF 是否为 LWAIL 专属提升:表 6 其他基线 + ICVF 的性能

结论:ICVF 嵌入对其他基线的提升有限,甚至部分基线性能下降,证明LWAIL 的框架是 ICVF 嵌入的最佳载体——LWAIL 将 ICVF 与 Wasserstein AIL、TD3 深度结合,而其他基线的优化目标与 ICVF 的动力学感知特性不匹配。
4.5.4 下游 RL 算法的选择

对比 TD3 与其他 RL 算法(PPO、DDPG),验证 TD3 的适配性:结论:TD3 的平均性能为 99.07,远优于 PPO(24.87)和 DDPG(29.01),因为 TD3 是离线、抗干扰的 RL 算法,适合对抗式学习中奖励信号的波动,而 PPO(在线)、DDPG(易过估计)在该场景中表现极差。
4.6 其他消融实验
论文还进行了数据集大小 / 质量、专家轨迹数量、子采样专家轨迹、超参数敏感性等消融实验,核心结论如下:
- 预训练数据量:仅需 10K 随机数据(1% 在线数据)即可达到最优性能,增加数据量性能无明显提升,证明 LWAIL 的预训练阶段数据效率极高;
- 预训练数据质量:随机、中等、专家质量的预训练数据性能接近,证明 LWAIL 无需高质量预训练数据,适配各种低质量数据场景;
- 专家轨迹数量:1 条、5 条、10 条甚至全部专家数据的性能基本一致,证明 LWAIL 是真正的少数据模仿学习方法;
- 子采样专家轨迹:即使仅取 1/10 的专家状态对,LWAIL 仍保持 96.27 的平均性能,远超同条件下的 WDAIL 和 IQlearn;
- 超参数敏感性:LWAIL 在学习率、更新频率、判别器训练轮数的合理范围内性能稳定,证明工程实现的鲁棒性。
五、相关工作
论文从三个核心方向梳理相关工作,明确 LWAIL 的创新点:
5.1 基于 Wasserstein 距离的模仿学习
现有方法分为原始 Wasserstein和KR 对偶 Wasserstein两类:
- 原始 Wasserstein:无基础度量限制,但需正则化、替代函数,计算复杂(如 PWIL);
- KR 对偶 Wasserstein:计算简单,但基础度量被限制为欧氏距离,无法捕捉动力学(如 WDAIL、IQlearn)。
LWAIL 的创新:首次将 ICVF 嵌入与 KR 对偶 Wasserstein 结合,解决了基础度量的动力学缺失问题,同时保持计算简单。
5.2 离线到在线模仿学习
现有方法需中等质量的离线数据和专家动作,无法利用随机仅状态数据(如 OLLIE)。LWAIL 的创新:仅用少量随机仅状态数据预训练,无需专家动作,且离线预训练与在线模仿的解耦设计提升了灵活性。
5.3 状态嵌入学习
现有方法分为理论型(双模拟度量)和实证型(对比学习、自编码器):
- 双模拟度量:需动作信息,聚焦于奖励差异的匹配;
- 对比学习(如 CURL):为视觉 RL 设计,未考虑模仿学习的分布匹配目标。
LWAIL 的创新:首次将 ICVF 嵌入应用于模仿学习,ICVF 无需动作和奖励,能从仅状态数据中学习动力学感知嵌入,且与 Wasserstein 距离的分布匹配目标天然对齐。
六、结论与展望
6.1 核心贡献
论文的三大核心贡献可总结为:
- 理论创新:证明 ICVF 隐空间能在少量低质量仅状态数据下捕捉动力学感知的基础度量,首次解决了传统 Wasserstein AIL 的欧氏距离几何缺陷;
- 算法创新:提出两阶段的 LWAIL 框架,将 ICVF 预训练与 Wasserstein 对抗式模仿学习解耦,实现仅用 1 条仅状态专家轨迹达到专家级性能;
- 实证创新:在 MuJoCo、Maze2d、Antmaze 等多类环境中验证 LWAIL 的优越性,消融实验充分证明了 ICVF 嵌入、TD3、隐空间度量等核心组件的重要性。
6.2 局限性
- LWAIL 的对抗式训练需要平衡判别器和策略的更新频率,这是所有 AIL 方法的共性问题;
- 目前仅在状态观测的环境中验证,尚未拓展到视觉观测(如像素级)的场景,这是未来的重要研究方向。
6.3 未来研究方向
- 视觉观测的 LWAIL:将 ICVF 与视觉表示学习结合,处理像素级的仅观测模仿学习;
- 多任务 / 跨环境模仿:利用 ICVF 的动力学感知特性,实现跨环境的少数据模仿学习;
- 模型基 LWAIL:将 LWAIL 与环境动力学模型结合,进一步提升数据效率;
- 更复杂的噪声场景:研究在强噪声、非平稳环境中的 LWAIL 改进方法。
七、总结
LWAIL 作为 ICLR 2026 的最新研究成果,针对模仿学习中专家数据稀缺、专家动作缺失、传统距离度量无动力学感知三大痛点,提出了一种简洁而高效的解决方案。其核心创新是将 ICVF 的动力学感知隐空间与 Wasserstein 对抗式模仿学习结合,通过两阶段框架实现了从极少量低质量仅状态数据到专家级策略的高效学习。
在实验中,LWAIL 在多个经典环境中显著超越现有基线,且展现出良好的鲁棒性和泛化能力,为实际场景中的模仿学习落地提供了新的思路。同时,LWAIL 的设计也为后续研究指明了方向:将表示学习与模仿学习的目标深度对齐,是提升少数据、仅观测模仿学习性能的关键。
论文核心亮点:用最简单的框架解决最核心的问题,无需复杂的模块设计,仅通过隐空间的改造就实现了性能的质的飞跃,兼具理论深度和工程实用性。
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。
更多推荐


所有评论(0)