在线模仿学习的耦合分布随机专家蒸馏

世界模型在线模仿学习的耦合分布随机专家蒸馏

Coupled Distributional Random Expert Distillation for World Model Online Imitation Learning

https://arxiv.org/pdf/2505.02228v1

摘要

模仿学习（IL）通过使智能体能够从专家演示中学习复杂行为，已在包括机器人技术、自动驾驶和医疗保健等各个领域取得了显著成功。然而，现有的模仿学习方法通常面临稳定性挑战，尤其是在世界模型框架中依赖对抗性奖励或价值公式时。在这项工作中，我们提出了一种新的在线模仿学习方法，该方法通过基于随机网络蒸馏（RND）的奖励模型来解决这些限制，用于密度估计。我们的奖励模型建立在世界模型潜在空间中专家分布和行为分布的联合估计之上。我们在各种基准测试中评估了我们的方法，包括 DMControl、Meta-World 和 ManiSkill2，展示了其在运动和操作任务中提供稳定性能并达到专家水平结果的能力。我们的方法在保持专家水平性能的同时，比对抗性方法表现出更高的稳定性。

1 引言

模仿学习（IL）最近在广泛的领域中显示出显著的有效性，特别是在解决复杂的现实世界挑战方面。在机器人技术中，模仿学习显著推进了操作任务的最新技术 [55,44,42,6]，使机器人能够精确且适应性地执行复杂操作。同样，模仿学习在运动任务中也取得了令人印象深刻的结果 [7,41,29]，为各种机器人平台开发了强大而灵活的运动控制器。除了机器人技术，模仿学习还在自动驾驶等领域展示了其多功能性 [36,3,5]，用于建模复杂的决策过程并确保车辆安全高效导航。此外，模仿学习已开始在医疗保健领域做出有意义的贡献 [10]，为医疗决策提供支持并增强复杂诊断过程的可解释性。这些成就凸显了模仿学习的广泛适用性及其在不同领域推动变革性进展的潜力。

模仿学习的最简单方法是将行为克隆直接应用于提供的专家数据集，如之前的 IBC [14] 和 Diffusion Policy [6] 等工作所示。然而，这种方法不了解动态，在遇到分布外状态时可能导致泛化能力不足。为了解决这些缺点，GAIL [27]、SQIL [38]、IQ-Learn [16]、MAIL [2] 和 CFIL [15] 等方法引入了价值或奖励估计，以促进对环境的更深入理解，同时利用在线交互来增强探索。具体而言，GAIL、MAIL 和 IQ-Learn 将模仿学习问题构建为对抗性训练过程，区分专家和学习者的状态 - 动作分布。

潜在世界模型在模仿学习中的最新进展取得了重大突破。包括 V-MAIL [37]、CMIL [32]、Ditto [9]、EfficientImitate [51] 和 IQMPC [34] 在内的几项先前工作已将对抗性模仿学习框架与世界模型集成，以解决模仿学习任务。然而，正如第 D.7 节所讨论的，我们发现即使使用世界模型，对抗性目标在某些情况下仍然可能存在不稳定性。为了克服这个问题，我们提出用一种基于随机网络蒸馏（RND）[4] 的新型密度估计方法来取代对抗性奖励或价值公式，以减轻不稳定性。具体而言，我们在世界模型的潜在空间中进行密度估计，利用潜在表示的优越特性及其增强的动态感知能力，因为潜在动态模型直接在该空间内训练。与现有的将 RND 用于模仿学习的方法 [45] 不同，我们的方法联合学习奖励模型和世界模型的其他组件，同时估计专家分布和行为分布。相比之下，现有的随机专家蒸馏 [45] 在原始观察和动作空间中估计分布，并将奖励模型学习与下游强化学习过程解耦，这使得在具有高维观察和动作空间的复杂任务中难以解决。我们在 DMControl [43]、Meta-World [52] 和 ManiSkill2 [18] 的一系列任务中评估了我们的方法，展示了稳定的性能并取得了专家水平的结果。

总之，我们工作的贡献总结如下：

• 我们为基于世界模型的在线模仿学习提出了一种新的奖励模型公式，该公式基于用于密度估计的随机网络蒸馏。

• 我们证明，与以前具有对抗性公式的方法相比，我们的方法表现出卓越的稳定性，并在一系列模仿学习任务中取得了专家水平的性能，包括运动和操作任务。

2 预备知识

我们将决策问题建模为马尔可夫决策过程（MDP）。MDP 可通过元组 ⟨S, A, p0, P, r, γ⟩ 定义。具体而言，S 和 A 分别表示状态空间和动作空间，p0 是初始状态分布，P : S × A → ∆S 描述转移概率，r (s, a) 是奖励函数，γ ∈ (0, 1) 为折扣因子。设 Z 为世界模型的潜在状态空间，Z × A 上由专家策略诱导的潜在状态 - 动作分布和由行为策略 π 诱导的潜在状态 - 动作分布分别记为 ρE 和 ρπ。

2.1 随机网络蒸馏

随机网络蒸馏（RND）[4] 是一种促进探索的技术。具体来说，它利用一个参数固定的随机网络 f̄θ(x) 和一个可学习的预测网络 fθ(x)。在训练过程中，RND 针对特定数据分布 ρ 的数据集 D 最小化以下均方误差损失：

2.2 世界模型

3 方法论

在本节中，我们将阐述我们提出的方法 —— 耦合分布随机专家蒸馏（Coupled Distributional Random Expert Distillation，简称 CDRED）的设计动机与详细方法论。我们将证明，相较于直接在基于世界模型的模仿学习中应用随机专家蒸馏（RED）[45]，我们的方法具有更高的稳定性与合理性。

3.1 动机

3 方法论

在本节中，我们将阐述我们提出的方法 —— 耦合分布随机专家蒸馏（Coupled Distributional Random Expert Distillation，简称 CDRED）的设计动机与详细方法论。我们将证明，相较于直接在基于世界模型的模仿学习中应用随机专家蒸馏（RED）[45]，我们的方法具有更高的稳定性与合理性。

3.1 动机

为解决这些困难，我们引入了一种耦合方法。该方法联合估计专家分布和行为分布；它在训练的早期阶段鼓励策略探索。我们在图 2 中提供了直观的图示，并在 3.3 节中描述了详细的方法论。在这种耦合方法中，我们需要在在线训练期间估计行为分布，正如 [49] 所指出的，这自然会引发最终奖励不一致的问题。因此，我们采用了他们跟踪数据出现频率的方法，我们将在 3.2 节中对此进行描述。

3.2 缓解奖励估计不一致问题

在 RND 模型训练的不同阶段可能会出现不一致性问题 [49]。在初始阶段，这些不一致性源于单个网络中的极端值，可通过使用多个目标网络（记为 K 个目标网络）来缓解。在最终阶段，当生成的奖励分布与实际状态 - 动作分布不匹配时，会产生不一致性。为解决这一问题，需要对状态 - 动作出现次数 n 构建无偏估计器。我们需要跟踪状态 - 动作的出现频率，以确保在在线训练分布式 RND 模型时保持一致性。

3.3 耦合分布随机专家蒸馏

公式 7 中的第一项衡量当前分布与专家分布之间的距离，而第二项通过惩罚策略的利用行为来鼓励探索。缩放因子 ζ 用于平衡这两项：在训练早期策略次优时，第二项占主导地位以促进探索；随着策略接近最优，第一项逐渐主导，使策略稳定在专家分布附近。通常 ζ 接近 1，确保初始探索阶段后第一项占据主导。

系数 σ 控制奖励函数的衰减率，其中第一项基于专家分布，第二项基于行为分布。为保证稳定性，奖励通过 K 个随机目标网络的输出均值计算。函数 g (x) 为单调递增函数，g (x)=exp (x) 和 g (x)=x 均适用（具体行为差异见附录 D.2）。

公式 8 中的标量系数 α 用于平衡第一项（L2 范数）和第二项（p1/n 估计量）的贡献。遵循文献 [49]，我们设定初始阶段第一项主导，随训练推进逐步转向第二项 —— 这可通过固定 α 实现，无需动态调整系数。该改进支持状态 - 动作分布的在线一致性估计，直接为在线模仿学习的奖励建模提供支撑。

3.4 与世界模型集成的模仿学习方法

世界模型通过将观测编码到潜在空间并学习潜在空间的转移模型，实现对策略和环境动态的建模。无解码器世界模型（如 TD-MPC 系列 [25,26]）已被证明是解决复杂强化学习任务的有效工具。我们采用的无解码器世界模型包含以下组件：

3.4 与世界模型集成的模仿学习方法

世界模型通过将观测编码到潜在空间并学习潜在空间的转移模型，实现对策略和环境动态的建模。无解码器世界模型（如 TD-MPC 系列 [25,26]）已被证明是解决复杂强化学习任务的有效工具。我们采用的无解码器世界模型包含以下组件：

奖励模型（CDRED）：由两个预测器和 K 个目标网络组成，通过估计专家分布和行为分布实现奖励近似；
编码器 h: S→Z：将观测（基于状态或视觉）映射为潜在表征；
潜在动态模型 d: Z×A→Z：学习潜在表征的转移动态，隐式建模环境动态；
价值函数：通过 CDRED 模型估计的奖励辅助训练时序差分目标，实现未来回报估计；
策略先验：学习随机策略以指导世界模型的规划过程。
训练流程见算法 1，规划过程详见算法 2。

模型训练

4 实验

我们在多种任务上开展实验，包括运动任务、操作任务，以及基于视觉和状态观测的任务。使用 DMControl [43]、Meta-World [52] 和 ManiSkill2 [18] 环境评估我们的方法。对于基线方法，我们将其与集成世界模型架构的 IQ-MPC [34]，以及无模型方法（具体为 IQ-Learn+SAC [16]（图中记为 IQL+SAC）、CFIL+SAC [15]、HyPE [39]（附录 D.3）和 SAIL [46]（附录 D.4））进行比较。此外，我们还将行为克隆（BC）作为基线方法纳入评估。所有实验中，专家轨迹均从训练好的 TD-MPC2 [26] 中采样。我们还对专家轨迹数量和函数 g 的选择进行了消融研究（详见附录 D.2），进一步评估了算法在噪声环境动态中的鲁棒性（附录 D.5），检验了在潜在空间中构建奖励模型的优势（附录 D.8），并突出了其相较于现有对抗训练方法的优越性（附录 D.7）。附录 D.6 还提供了衡量训练稳定性的定量结果。所有实验均在单张 RTX3090 显卡上进行。

4.1 Meta-World 实验

我们在 Meta-World 环境的 6 个任务上开展实验。每个任务使用 100 条专家轨迹，确保专家数据在所有算法中保持一致，以在各任务内进行公平比较。由于采用对抗训练方法，即使应用梯度惩罚，IQ-MPC 在这些任务中仍受限于过强的判别器。CFIL+SAC [15] 由于训练流模型的固有挑战，在训练过程中遇到不稳定性。我们的方法表现出稳定且达到专家水平的性能，在这些任务中优于这些基线方法。图 4 展示了回合奖励结果，表 1 列出了成功率结果。

4.2 DMControl 实验

我们在 DMControl [43] 环境的 6 个任务上开展实验。对于低维任务，我们使用 100 条专家轨迹；对于高维任务，使用 500 条专家轨迹。环境维度详情见附录 C。在 Hopper Hop、Walker Run 和 Humanoid Walk 任务上，我们的 CDRED 模型表现与 IQ-MPC 相当。然而，在 Cheetah Run 和 Dog Stand 任务中，IQ-MPC 出现长期不稳定性，导致智能体在长时间在线训练后失败。在 Reacher Hard 任务上，IQ-MPC 因判别器过强而难以学习到专家水平的策略。基线算法中的无模型方法在这些任务上未能实现稳定的专家水平性能。回合奖励结果如图 5 所示。

4.3 基于视觉的实验

除了使用基于状态观测的实验外，我们还在基于视觉观测的任务上对方法进行基准测试。具体而言，我们从 DMControl [43] 中选择了三个基于视觉观测的任务。为创建这些视觉数据集，我们基于基于状态的专家轨迹渲染视觉观测，以替换专家数据中原始的基于状态的观测。每个任务使用由训练好的 TD-MPC2 模型 [26] 生成的 100 条专家轨迹。结果如图 6 所示。有趣的是，我们观察到，当使用由训练好的基于状态的 TD-MPC2 策略生成的轨迹（其中状态观测被从这些状态渲染的 RGB 图像替换）时，视觉版 IQ-MPC 在 Cheetah Run 任务中遇到判别器过强的问题。然而，当使用直接在视觉观测上训练的 TD-MPC2 策略生成的专家轨迹时，IQ-MPC 表现良好。

5 结论

我们提出了一种基于世界模型的在线模仿学习新方法，其核心是创新的奖励模型公式。与传统对抗性方法可能在训练中引入不稳定性不同，我们的奖励模型基于专家和行为状态 - 动作分布的密度估计。这种公式在保持高性能的同时增强了稳定性。我们的模型在多个基准测试的各种任务中（包括 DMControl、Meta-World 和 ManiSkill2）均表现出专家级水平，并且在长期在线训练中始终保持稳定性能。凭借其强大的奖励建模和稳定性，我们的方法有望解决复杂的现实世界机器人控制任务，这些任务对可靠性和适应性有着关键要求。

F 相关工作

我们的研究建立在模仿学习与基于模型的强化学习的前期进展之上。

模仿学习

模仿学习（IL）的最新进展已借助深度神经网络和多样化方法提升性能。生成对抗模仿学习（GAIL）[27] 通过受生成对抗网络（GANs）[17] 启发的极小极大优化问题，为对抗性奖励学习奠定了基础。后续多项研究基于 GAIL 展开：基于模型的对抗模仿学习（MAIL）[2] 通过数据驱动方法训练前向模型，扩展了 GAIL 框架；ValueDICE [33] 将对抗框架转化为基于分布比率估计的离策略学习。离线模仿学习领域也涌现出重要进展，例如 Diffusion Policy [6] 将扩散模型应用于行为克隆，Ditto [9] 则将 Dreamer V2 [23] 与对抗技术结合。Implicit BC [14] 表明，基于隐式模型的监督策略学习可提升机器人任务的实际表现；DMIL [54] 在离线模仿学习中利用判别器评估动态模型准确性及模型滚动预测相对专家演示的次优性。

其他创新聚焦于融合先进强化学习技术：逆软 Q 学习（IQ-Learn）[16] 重构了 GAIL 的学习目标，并将其应用于软演员 - 评论家（soft actor-critic）[20] 和软 Q 学习智能体；SQIL [38] 提出基于软 Q 函数的在线模仿学习算法；CFIL [15] 引入耦合流方法，实现从专家演示中同步生成奖励与学习策略；随机专家蒸馏（RED）[45] 提出了一种替代方案，通过估计专家策略分布的支撑集构建奖励模型。

基于模型的模仿学习方法

基于模型的方法在推动模仿学习发展中同样至关重要：V-MAIL [37] 采用变分模型促进模仿学习，CMIL [32] 将保守世界模型应用于基于图像的操作任务；前期研究 [11,28,30] 强调了基于模型的模仿学习在现实机器人控制与自动驾驶中的潜力；Das 等人 [8] 提出的基于模型的逆强化学习方法通过关键点预测提升模仿任务表现；混合逆强化学习（Hybrid Inverse Reinforcement Learning）[39] 提出融合在线数据与专家演示的新策略，增强智能体在随机环境中的鲁棒性；EfficientImitate [51] 将 EfficientZero [50] 与对抗模仿学习结合，在 DMControl 任务 [43] 中取得优异性能。

基于模型的强化学习（MBRL）

基于模型的强化学习（MBRL）的最新进展利用通过数据驱动方法构建的动态模型，提升智能体的学习与决策能力。MBPO [31] 提出基于模型的策略优化算法，确保策略逐步单调改进；将其扩展至离线强化学习时，MOPO [53] 在奖励函数中引入基于动态模型不确定性的惩罚项，有效处理分布偏移问题；MBVE [13] 通过基于模型的轨迹滚动（model-based rollouts）增强无模型智能体，改善价值估计精度。

许多方法聚焦于在潜在空间中构建动态模型：PlaNet [22] 开创性地提出带证据下界（ELBO）训练目标的循环状态空间模型（RSSM），解决部分可观测马尔可夫决策过程（POMDP）中的挑战；基于 PlaNet，Dreamer 系列算法 [21,23,24] 利用学习的世界模型在潜在空间中模拟未来轨迹，实现高效学习与规划；TD-MPC 系列 [25,26] 进一步优化潜在空间建模，通过开发适用于模型预测控制（MPC）的可扩展世界模型，结合时序差分学习目标提升性能；类似地，MuZero [40] 将潜在动态模型与树搜索结合，在离散控制任务中实现规划与策略学习的无缝融合；EfficientZero 系列 [50,47] 对 MuZero 进行改进，在视觉强化学习任务中实现卓越的采样效率。

https://arxiv.org/pdf/2505.02228v1