Diffusion Reward Learning Rewards via Conditional Video Diffusion
文章概括
引用:
@inproceedings{huang2025diffusion,
title={Diffusion reward: Learning rewards via conditional video diffusion},
author={Huang, Tao and Jiang, Guangqi and Ze, Yanjie and Xu, Huazhe},
booktitle={European Conference on Computer Vision},
pages={478--495},
year={2025},
organization={Springer}
}
Huang, T., Jiang, G., Ze, Y. and Xu, H., 2025. Diffusion reward: Learning rewards via conditional video diffusion. In European Conference on Computer Vision (pp. 478-495). Springer, Cham.
原文: https://link.springer.com/chapter/10.1007/978-3-031-72946-1_27
代码、数据和视频:diffusion-reward.github.io.
系列文章:
请在
《
《
《文章
》
》
》 专栏中查找
摘要
从专家视频中学习奖励是一种经济高效的方法,可用于指定强化学习(RL)任务的预期行为。在本工作中,我们提出了一个名为“Diffusion Reward”的新框架,该框架通过条件视频扩散模型从专家视频中学习奖励,用于解决复杂的视觉强化学习问题。我们的核心见解是,当扩散模型以专家轨迹为条件时,生成的多样性较低。因此,“Diffusion Reward”被形式化为条件熵的负值,以鼓励对专家行为的有效探索。我们在模拟平台和真实世界的机器人操作任务中使用视觉输入展示了该方法的有效性。此外,“Diffusion Reward”甚至能够成功且高效地解决未见过的任务,并在很大程度上超越了基线方法。项目页面和代码:diffusion-reward.github.io。
关键词: 从视频中学习奖励 · 机器人操作 · 视觉强化学习
1 引言
奖励的指定是强化学习(RL)中的一个基本挑战,它决定了智能体学习行为的有效性及其与预期目标的一致性。手动设计稠密的奖励函数既耗费大量劳动,有时在现实世界的任务中(如机器人操作 [40])甚至是不可行的,因为在这些任务中获取特权状态信息非常困难。作为一种替代方案,在这些场景中使用稀疏奖励通常更受青睐,因为它对人工干预的要求较低 [29]。然而,稀疏奖励提供的监督不足会阻碍强化学习的样本效率,可能给数据收集带来难以承受的负担。
从专家视频中学习奖励函数提供了一种有前景的解决方案,因为视频采集的工作量较小且视频中包含了丰富的任务执行信息 [6,54]。在这些方法中,生成模型自然地被研究者用于无监督地提取信息丰富的奖励,以供强化学习训练使用 [32,45]。一种经典的方法基于生成对抗学习,通过学习一个判别性奖励来区分智能体和专家的行为。尽管这种方法直接明了,但它未能充分利用时间信息,而时间信息在解决强化学习问题中已被证明非常重要 [56],并且该方法对对抗训练的鲁棒性较差。针对这些问题,最近的研究通过预训练 VideoGPT [49] 来编码时间信息,并直接使用预测的对数似然作为奖励 [10]。然而,这些方法在建模复杂的专家视频分布及相应的奖励时遇到了困难,尤其是在具有复杂动态的任务中。如图2所示,对于分布外的专家视频,即使这些视频在分布内视频中与最优行为一致,其学习到的回报也表现出明显下降的趋势。
图2. 视频模型和奖励的分析。VIPER (CE) 和 Diffusion Reward (LL) 分别用条件熵 (CE) 和对数似然 (LL) 替换了原始奖励。基于 CE 的奖励为未见过的专家视频分配了接近最优的奖励。这种提升得益于扩散模型强大的建模能力。结果在7个 MetaWorld 任务上取平均值。Suboptimal 表示包含25%随机操作的视频。
视频扩散模型 [17] 在计算机视觉领域表现出了卓越的性能,它们在捕捉视频的复杂分布方面展现了强大的能力,例如文本生成视频 [11,15] 和视频编辑 [5,26]。最近的研究还探讨了它们在建模专家视频方面的潜力,以作为机器人操作任务中的可泛化规划器 [9,19,21]。尽管取得了这些进展,从视频扩散模型中提取信息丰富的奖励仍是一个研究不足的领域,但其在指导强化学习从视频中获取专业技能并推广到未见任务方面展现了巨大的前景。
在本研究中,我们提出了“Diffusion Reward”,一种利用视频扩散模型捕捉专家视频分布并预训练视觉强化学习稠密奖励函数的奖励学习框架。我们的核心见解是,相较于以非专家轨迹为条件的扩散,以类似专家的轨迹为条件的扩散表现出更低的生成多样性。这一原理自然地通过寻找低生成多样性来指导强化学习对专家行为的探索。我们因此估计条件于历史帧的熵,形式化我们的见解,并结合寻求新奇性的奖励 [4] 和稀疏的环境奖励,形成用于高效强化学习的稠密奖励。此外,为了加速奖励推断,我们通过利用矢量量化编码 [13] 压缩高维观察,执行潜在的扩散过程。
我们通过在10个视觉机器人操作任务上的实验实证验证了我们框架的有效性,包括来自MetaWorld [51]的7个夹爪操作任务和来自Adroit [35]的3个灵巧操作任务,分别表现出比表现最佳的基线方法高出38%和35%的性能提升。此外,我们惊喜地观察到,所学习的奖励在未见任务上的零样本泛化性能表现良好。此外,在真实机器人任务上生成的合理奖励以及相关的离线强化学习性能,展示了Diffusion Reward在现实世界中的适用性。图1概述了本研究。本文的主要贡献可以总结如下:
- 我们引入了Diffusion Reward,这是一种新颖的奖励学习框架,利用视频扩散模型的生成建模能力,为高效强化学习提供信息丰富的稠密奖励信号。
- 我们证明了该框架在仿真和现实世界的机器人操作任务中,相比基线方法具有显著的性能优势。
- 我们观察到预训练的奖励模型能够指导强化学习在未见任务上进行学习,并在真实机器人任务中生成合理的奖励。
2 相关工作
从视频中学习奖励。 从专家视频中提取奖励为强化学习中的奖励指定提供了一种经济但有效的解决方案 [6,37,38]。若干研究尝试通过在潜在空间中测量当前观察与目标图像之间的距离来学习指示任务进展的稠密奖励 [3,25,54]。尽管这些方法前景广阔,但在仿真之外获取目标图像可能具有挑战性,从而限制了它们在开放世界任务中的适用性 [24]。相比之下,生成模型被广泛用于无需访问未来信息即可提取奖励。一种代表性的方法,属于仅基于状态的模仿学习 [34,44,58],通过生成对抗学习 [32,45] 的在线范式来区分类似专家和非专家的行为。这一思想在[10]中得到了进一步优化,其中从视频中预训练奖励函数,并预测对数似然作为奖励。然而,在复杂任务中准确建模专家视频分布时,该方法面临挑战,这些任务具有复杂的动态特性,从而使其容易生成无效的奖励。相比之下,我们的方法利用扩散模型的强大建模能力,并估计条件熵的负值作为更具辨别性的奖励,以加速强化学习的探索。
用于强化学习的扩散模型。 扩散模型已被广泛研究用于强化学习,例如提升策略的表达能力 [8,14,47],以及增强经验 [7,52]。除此之外,一些研究直接从离线数据中学习无条件规划的扩散模型 [19],或者根据任务回报指定条件规划的扩散模型 [2]。条件扩散的概念在[9,21]中进一步研究,使用文本作为任务规范,其中视频扩散模型作为与逆动力学相关的规划器。与这些方法不同,我们旨在通过条件于历史帧的视频扩散模型学习信息丰富的奖励,从而加速在线强化学习。这种历史条件化已在[18]中用于轨迹生成,但与我们的奖励学习重点不同。我们的工作也与[30]密切相关,该研究从两个拟合不同行为的扩散模型中提取奖励,针对的是基于状态的任务。而我们的区别在于,我们预训练的奖励模型无需动作标签和额外的任务无关数据。
强化学习的预训练。 许多方法已被探索用于直接从数据中预训练强化学习的有效策略。这些方法包括在在线微调之前采用离线强化学习技术的方法 [22,23,27],以及学习行为(技能)先验的方法 [33,36,39]。然而,这些方法需要额外收集奖励、动作标签或任务无关数据,这会带来一定的代价。相比之下,我们的方法专注于仅从无动作的纯专家视频中预训练奖励,从而减少了对这些额外数据源的需求。一些研究也探索了使用与我们数据集相似的数据集进行视觉表示的预训练 [28,31,48,53]。然而,这些方法并未提取用于指导强化学习的稠密任务执行信号。
3 前言
问题表述。 我们考虑一个与环境交互的强化学习智能体,该环境被建模为一个有限时域的马尔可夫决策过程(MDP),由一个元组 M = ( S , A , T , R , γ ) M = (S, A, T, R, \gamma) M=(S,A,T,R,γ)定义,其中 S S S是状态空间, A A A是动作空间, T T T是环境转移函数, R R R是奖励函数, γ \gamma γ是折扣因子。强化学习智能体的目标是学习一个最优策略 π \pi π,以最大化在 K K K个时间步内的期望回报 E [ ∑ k = 0 K − 1 γ k r k ] E[\sum_{k=0}^{K-1} \gamma^k r_k] E[∑k=0K−1γkrk]。
在本工作中,我们更具体地关注高维状态空间及二元稀疏奖励(即0/1奖励,用作任务成功的指示器)。特别地,我们将RGB图像 x ∈ R H × W × 3 x \in \mathbb{R}^{H \times W \times 3} x∈RH×W×3作为智能体观察到的状态。这种设置来源于强化学习的现实应用场景,例如机器人领域,其中基于视觉的传感数据更加易得,而指定复杂的奖励通常需要繁琐的手工设计,有时甚至不可行。然而,这对强化学习智能体提出了巨大挑战,因为通常需要大量与环境的交互,这被称为样本效率问题。
专家视频。 为了提高强化学习的样本效率,我们假设智能体可以访问一组由专家策略生成的无标签视频。这些视频显著特点是无动作信息,并从多个任务中收集而来,且未进行任务标识。我们将其表示为 D = { D 1 , D 2 , . . . , D N } D = \{D_1, D_2, ..., D_N\} D={D1,D2,...,DN},其中 D i D_i Di是任务 i i i的演示视频集。每个视频集 D i D_i Di由多个专家视频组成, τ = { x 0 , x 1 , . . . , x K − 1 } ∈ S K \tau = \{x_0, x_1, ..., x_{K-1}\} \in S^K τ={x0,x1,...,xK−1}∈SK。我们的目标是从专家视频中预训练有效的奖励函数,以加速强化学习的在线探索。接下来的部分将介绍我们的奖励学习框架。
马尔可夫决策过程(MDP)
定义: 一个有限时域的马尔可夫决策过程 (Markov Decision Process, MDP) 是一个数学模型,表示强化学习问题的基本结构,用来描述智能体与环境的交互。它由以下五元组定义: M = ( S , A , T , R , γ ) M = (S, A, T, R, \gamma) M=(S,A,T,R,γ)
- S S S(状态空间): 描述环境可能处于的所有状态集合。例如,在机器人控制中, S S S可能是机器人臂的所有可能位置和速度的组合。
- A A A(动作空间): 描述智能体可以采取的所有可能动作的集合。例如,机器人臂的所有可能运动方向。
- T T T(转移函数): 描述状态与动作之间的动态关系。给定当前状态 s ∈ S s \in S s∈S和动作 a ∈ A a \in A a∈A,转移函数 T ( s ′ ∣ s , a ) T(s'|s, a) T(s′∣s,a)定义了在执行动作 a a a后,转移到下一状态 s ′ s' s′的概率。
- R R R(奖励函数): 定义智能体在每个状态或采取某个动作后收到的奖励。例如,完成任务的机器人可能收到奖励1,而失败时奖励为0。
- γ \gamma γ(折扣因子): 表示未来奖励对当前决策的重要性,通常在 0 ≤ γ < 1 0 \leq \gamma < 1 0≤γ<1之间。如果 γ \gamma γ较小,则智能体更关注短期奖励。
目标: 智能体的目标是学习一个最优策略 π \pi π,即定义在状态空间上的概率分布 π ( a ∣ s ) \pi(a|s) π(a∣s),它能在有限时间步内 K K K最大化期望回报:
E [ ∑ k = 0 K − 1 γ k r k ] E[\sum_{k=0}^{K-1} \gamma^k r_k] E[k=0∑K−1γkrk]
例如,在一个自动驾驶场景中, S S S可能是车辆周围的环境状态, A A A可能是车辆的控制信号,智能体的目标是安全、高效地驾驶到目的地。
高维状态空间和稀疏奖励
高维状态空间
在一些实际问题中,状态是高维数据。例如,机器人看到的环境可以用RGB图像 x ∈ R H × W × 3 x \in \mathbb{R}^{H \times W \times 3} x∈RH×W×3表示,其中 H H H是图像高度, W W W是图像宽度, 3 3 3表示颜色通道。这种高维状态带来了以下挑战:
- 数据处理复杂: 从图像中提取有意义的特征需要更强的计算能力。
- 样本效率低: 在高维状态空间中探索需要更多的样本,训练变得昂贵。
稀疏奖励
在强化学习中,奖励函数通常用来指导智能体的学习。然而,在许多实际场景中,奖励非常稀疏,即只有在任务成功时(如机器人完成拼图任务),智能体才能获得奖励。这种设置的主要问题包括:
- 难以找到有效策略: 智能体可能在许多回合中都无法获得任何奖励,导致学习效率低。
- 探索难度大: 智能体需要尝试大量不同的动作组合才能触发奖励。
示例: 一个机器人被要求在房间中找到一个特定的物体。当机器人触碰到目标物体时,奖励为1;其他情况下,奖励为0。机器人可能需要花费大量时间搜索,甚至可能从未找到物体,导致学习困难。
专家视频
无标签视频
在本文设置中,智能体可以访问由专家策略生成的视频数据,这些视频的特点是:
- 无动作信息: 只包含状态序列 x 0 , x 1 , . . . , x K − 1 x_0, x_1, ..., x_{K-1} x0,x1,...,xK−1,而不包括相应的动作。
- 多任务混合: 视频来源于多个不同任务的执行,且未进行任务标识。
- 高质量演示: 视频中展示的是高效完成任务的专家行为。
这些无标签视频可以帮助智能体更快地学习任务。通过观察这些视频,智能体可以:
- 获取完成任务的示范样本。
- 通过模仿专家行为,减少随机探索的需求。
数据结构
专家视频集被表示为: D = { D 1 , D 2 , . . . , D N } D = \{D_1, D_2, ..., D_N\} D={D1,D2,...,DN} 其中, D i D_i Di是任务 i i i的演示视频集,每个视频 τ ∈ D i \tau \in D_i τ∈Di是一个状态序列: τ = { x 0 , x 1 , . . . , x K − 1 } ∈ S K \tau = \{x_0, x_1, ..., x_{K-1}\} \in S^K τ={x0,x1,...,xK−1}∈SK
例如,在自动驾驶场景中, D i D_i Di可以是不同城市的驾驶视频,视频中的每一帧 x t x_t xt表示一个驾驶场景。
奖励学习框架
本文的目标是从这些专家视频中预训练一个奖励函数,以帮助智能体更快地找到最优策略。关键点包括:
- 奖励函数学习: 使用无标签视频训练一个模型,将状态序列与任务成功概率相关联。
- 在线探索加速: 在强化学习阶段,利用预训练的奖励函数引导智能体更高效地探索环境。
总结 本文中提到的设置源于许多现实问题,例如机器人视觉控制和自动驾驶,这些问题普遍具有高维状态空间和稀疏奖励的特征。通过利用专家视频,强化学习可以显著提高样本效率,减少实际应用中的高额训练成本。
4 方法
我们提出了“Diffusion Reward”,这是一个通过条件视频扩散从专家视频中学习奖励的新框架,用于解决后续的视觉强化学习问题,如图1和算法1所示。在第4.1节中,我们首先在潜在空间中使用视频扩散模型对专家视频进行建模。在第4.2节中,我们通过估计条件于历史的熵,并将其负值用作强化学习的奖励,形式化了Diffusion Reward的核心见解。
图1. Diffusion Reward概述。(左)我们提出了一个使用视频扩散模型进行奖励学习的强化学习框架,该模型通过专家视频预训练。我们在条件于历史帧的情况下执行扩散过程,以估计条件熵作为奖励,从而鼓励强化学习探索类似专家的行为。(右)10个视觉机器人操作任务的平均成功率表明了我们提出的Diffusion Reward在5次运行中的有效性。阴影区域表示标准误差。
图1内容详解
左图:Diffusion Reward的概述
视频扩散模型(Video Diffusion Model)
- 概念: 视频扩散模型是一种基于扩散过程的生成模型,可以从噪声中逐步去噪,生成与输入条件匹配的视频帧。
- 工作原理:
- 输入历史帧 z c z_c zc(即RL智能体与环境交互的历史观测序列)作为条件。
- 在条件的指导下,通过扩散过程生成去噪后的样本(denoised samples)。
- 使用生成的样本估计条件熵 − H ( P θ ( ⋅ ∣ z c ) ) -\mathcal{H}(P_\theta(\cdot|z_c)) −H(Pθ(⋅∣zc)),其中 P θ P_\theta Pθ是视频扩散模型的概率分布。
- 目的: 利用条件熵作为奖励信号,从而引导强化学习智能体探索更接近专家行为的策略。
示例: 假设一个机械臂任务中,智能体需要从视频帧学习如何抓取目标物体:
- 历史帧 z c z_c zc可能包括机械臂接近目标物体的状态。
- 视频扩散模型基于这些历史帧生成可能的后续帧(?),例如机械臂正确抓住物体的状态。
- 如果生成的样本更清晰地表明目标行为(如抓取成功),则条件熵降低,智能体获得更高的奖励。
条件熵(Conditional Entropy)
- 定义: 条件熵衡量在给定条件(历史帧 z c z_c zc)下,模型预测的不确定性。
H ( P θ ( ⋅ ∣ z c ) ) = − ∑ x P θ ( x ∣ z c ) log P θ ( x ∣ z c ) \mathcal{H}(P_\theta(\cdot|z_c)) = -\sum_{x}P_\theta(x|z_c)\log P_\theta(x|z_c) H(Pθ(⋅∣zc))=−x∑Pθ(x∣zc)logPθ(x∣zc)- 含义: 较低的条件熵意味着模型生成的后续帧更加确定,表示智能体行为更接近专家轨迹。
条件熵与负条件熵的关系
- 条件熵( H ( P θ ( ⋅ ∣ z c ) ) \mathcal{H}(P_\theta(\cdot|z_c)) H(Pθ(⋅∣zc))): 它表示在给定历史帧 z c z_c zc的条件下,视频扩散模型生成后续帧的不确定性。条件熵越低,说明模型对未来状态的预测越确定,意味着智能体的行为更接近专家的轨迹。
- 负条件熵( − H ( P θ ( ⋅ ∣ z c ) ) -\mathcal{H}(P_\theta(\cdot|z_c)) −H(Pθ(⋅∣zc))): 负条件熵是奖励函数的核心,直接用于强化学习智能体的优化。因为条件熵越低,负条件熵越高,所以通过最大化负条件熵,智能体会倾向于降低条件熵,从而更接近专家行为。
为什么使用负条件熵?
强化学习的目标是最大化奖励。为了让智能体倾向于产生低条件熵的行为,需要将负条件熵作为奖励:
- 当条件熵较低时,负条件熵值较高,智能体获得更多奖励。
- 例如,假设 H ( P θ ( ⋅ ∣ z c ) ) = 0.1 \mathcal{H}(P_\theta(\cdot|z_c)) = 0.1 H(Pθ(⋅∣zc))=0.1(低条件熵),对应的奖励 − H = − 0.1 -\mathcal{H} = -0.1 −H=−0.1(较高)。
- 当条件熵较高时,负条件熵值较低,智能体获得较少奖励。
- 例如,假设 H ( P θ ( ⋅ ∣ z c ) ) = 1.0 \mathcal{H}(P_\theta(\cdot|z_c)) = 1.0 H(Pθ(⋅∣zc))=1.0(高条件熵),对应的奖励 − H = − 1.0 -\mathcal{H} = -1.0 −H=−1.0(较低)。
因此,使用负条件熵可以有效引导智能体朝着更接近专家轨迹的方向探索和学习。
总结
- 条件熵低,表示模型对未来帧的预测更加确定。
- 负条件熵高,对应智能体获得更高奖励。
- 最大化负条件熵的目标间接鼓励智能体降低条件熵,使其行为更接近专家。
奖励机制: Diffusion Reward使用负条件熵作为奖励: Reward = − H ( P θ ( ⋅ ∣ z c ) ) \text{Reward} = -\mathcal{H}(P_\theta(\cdot|z_c)) Reward=−H(Pθ(⋅∣zc))
- 如果智能体的行为导致视频扩散模型的预测更确定,则奖励更高。
- 这相当于鼓励智能体采取专家视频中类似的行为。
- 强化学习与环境交互
- 强化学习(RL)智能体与环境交互,收集观测数据 z c z_c zc。
- 智能体通过扩散奖励信号更新其策略,从而在未来的交互中实现更接近专家的表现。
右图:Diffusion Reward在任务中的表现
任务设置
- 图中展示了10个视觉机器人操作任务的平均成功率:7个来自Metaworld环境,3个来自Adroit环境。
- 每个任务代表不同类型的机器人操作,例如抓取、旋转或摆放物体。
比较方法
- Raw Sparse Reward(稀疏奖励): 直接使用原始的稀疏奖励信号,智能体需要在奖励稀疏的情况下学习。
- RND(Random Network Distillation): 利用随机网络输出的预测误差作为探索奖励。
- AMP(Adversarial Motion Priors): 基于对抗学习的方法,用于奖励专家行为。
- VIPER: 模拟模仿专家行为的框架。
- Diffusion Reward(本文提出的方法): 使用视频扩散模型计算条件熵奖励。
结果分析
- Metaworld环境: Diffusion Reward方法显著超越其他基线方法,成功率在训练结束时接近100%,表现出更强的任务完成能力。
- Adroit环境: Diffusion Reward依然优于其他方法,特别是在任务复杂的情况下,成功率差距更明显。
- 稳定性: Diffusion Reward的阴影区域(标准误差)较小,说明方法在不同运行中的表现更一致。
核心知识点总结
视频扩散模型的引入
- 通过预训练专家视频,学习如何根据历史状态预测未来轨迹。
- 使用扩散模型生成去噪样本,并通过条件熵量化不确定性。
Diffusion Reward的优势
- 有效利用专家视频作为指导信号。
- 避免了直接手工设计复杂奖励函数的需求。
- 提高了强化学习的样本效率和任务成功率。
实际应用场景
- 机器人操作: 从视觉数据学习复杂任务(如抓取、分类、装配)。
- 自动驾驶: 利用驾驶视频数据预训练奖励模型,指导车辆探索安全驾驶策略。
通过引入视频扩散模型和条件熵奖励,本文提出的方法显著提高了强化学习在高维稀疏奖励任务中的效率和效果。
算法框架
该算法分为两个主要阶段:
- 预训练奖励模型阶段(第1行-第2行):通过专家视频训练扩散模型,用于奖励函数的构建。
- 强化学习阶段(第3行-第11行):在学习的奖励函数指导下进行强化学习。
预训练奖励模型阶段:
第1行: 收集专家视频数据
- 目标: 从 K K K个任务中收集专家执行任务的视频数据集 D D D。
- 专家视频 D D D:
- 是无标签的轨迹数据 τ \tau τ,每个轨迹是状态序列 { x 0 , x 1 , … , x T − 1 } \{x_0, x_1, \dots, x_{T-1}\} {x0,x1,…,xT−1}。
- 这些视频展现了高效完成任务的行为。
- 作用: 提供模型训练所需的先验知识,用于帮助智能体模仿专家行为。
示例: 在一个机器人抓取任务中,专家视频可以包括机械臂从初始状态接近目标并成功抓取物体的连续帧序列。
第2行: 训练扩散模型 p θ p_\theta pθ
- 扩散模型(Diffusion Model):
- 一种生成模型,通过噪声添加和去噪过程学习数据分布。
- 在这里,它被用来建模专家视频的条件分布 p θ ( z k 0 : T ∣ z c ) p_\theta(z_k^{0:T}|z_c) pθ(zk0:T∣zc),其中 z c z_c zc是历史帧, z k 0 : T z_k^{0:T} zk0:T是后续帧。
- 目标: 让扩散模型能够根据历史帧生成符合专家行为的后续帧。
核心思想:
- 通过预训练,扩散模型学习了专家行为的特征分布,为后续奖励设计奠定基础。
示例: 训练模型学习从历史帧生成机械臂逐渐接近目标的轨迹。
强化学习阶段:
第3行: 主循环
- 目标: 智能体在环境中持续与环境交互,并基于学习的奖励函数优化其策略 π \pi π。
- 终止条件: 算法会持续运行,直到策略 π \pi π收敛。
第4行: 智能体动作采样
- 智能体根据当前策略 π \pi π选择动作 a k a_k ak: a k ∼ π ( ⋅ ∣ x k ) a_k \sim \pi(\cdot|x_k) ak∼π(⋅∣xk)
- x k x_k xk: 当前的状态。
- π \pi π: 策略,定义在每个状态下选择动作的概率分布。
示例: 在机器人抓取任务中,当前状态 x k x_k xk可能是机械臂和目标物体的相对位置,策略 π \pi π决定移动方向。
第5行: 生成去噪样本
- 使用扩散模型生成 M M M个去噪后的样本: z ~ k 0 : T ∼ p θ ( z k 0 : T ∣ z c ) \tilde{z}_k^{0:T} \sim p_\theta(z_k^{0:T}|z_c) z~k0:T∼pθ(zk0:T∣zc)
- M M M: 样本数,表示模型在相同条件下生成多种可能的轨迹。
- z c z_c zc: 智能体的历史帧。
目的: 通过去噪样本,估计智能体行为与专家轨迹的匹配程度。
第6行: 估计条件熵
- 条件熵( H ( p θ ( ⋅ ∣ z c ) ) H(p_\theta(\cdot|z_c)) H(pθ(⋅∣zc))):
- 衡量在历史帧 z c z_c zc的条件下,扩散模型对后续轨迹的预测不确定性。
- 公式如下:
H ( p θ ( ⋅ ∣ z c ) ) = − ∑ z ~ k 0 : T P θ ( z ~ k 0 : T ∣ z c ) log P θ ( z ~ k 0 : T ∣ z c ) H(p_\theta(\cdot|z_c)) = -\sum_{\tilde{z}_k^{0:T}} P_\theta(\tilde{z}_k^{0:T}|z_c) \log P_\theta(\tilde{z}_k^{0:T}|z_c) H(pθ(⋅∣zc))=−z~k0:T∑Pθ(z~k0:T∣zc)logPθ(z~k0:T∣zc)意义:
- 条件熵低:模型对未来轨迹预测更确定,表示当前行为更接近专家。
- 条件熵高:模型预测不确定性较大,表示当前行为偏离专家轨迹。
示例: 如果机械臂接近目标位置,模型预测未来轨迹可能较确定;如果机械臂偏离目标,预测的轨迹不确定性增加。
第7行: 计算Diffusion Reward
- 根据条件熵计算扩散奖励: r k = − H ( p θ ( ⋅ ∣ z c ) ) r_k = -H(p_\theta(\cdot|z_c)) rk=−H(pθ(⋅∣zc))
- 负条件熵:
- 条件熵低时,奖励高。
- 条件熵高时,奖励低。
- 奖励机制鼓励智能体采取降低条件熵的行为。
作用: 通过奖励信号引导智能体模仿专家行为。
第8行: 环境状态转移
- 智能体执行动作 a k a_k ak,环境状态转移: x k + 1 ∼ T ( x k , a k ) x_{k+1} \sim T(x_k, a_k) xk+1∼T(xk,ak)
- T T T: 环境转移函数,描述当前状态和动作下下一状态的分布。
第9行: 存储过渡数据
- 将当前回合的过渡数据存储: ( x k , a k , r k , x k + 1 ) (x_k, a_k, r_k, x_{k+1}) (xk,ak,rk,xk+1)
- 用于后续策略优化。
第10行: 更新策略 π \pi π
- 根据存储的过渡数据,用强化学习算法更新智能体的策略 π \pi π和随机奖励模型 r rnd r^\text{rnd} rrnd。
核心知识点总结
扩散模型:
- 通过专家视频学习行为分布,生成后续帧用于奖励设计。
- 条件熵用于衡量智能体行为与专家轨迹的接近程度。
扩散奖励:
- 负条件熵作为奖励,指导智能体模仿专家行为。
强化学习与扩散奖励结合:
- 扩散奖励信号提高了智能体在稀疏奖励场景下的学习效率。
实际应用:
- 在机器人操作中,Diffusion Reward可用于引导智能体高效完成复杂任务,例如抓取、装配或路径规划。
这个算法既包括了训练奖励模型的预训练过程,也描述了如何在实验中将该奖励模型应用于强化学习任务。以下是两部分的详细解释:
1. 训练奖励模型(预训练阶段)算法的第1行和第2行描述了奖励模型的训练过程,这是训练阶段的内容,主要目的是利用专家视频构建一个用于奖励生成的模型:
收集专家视频:
- 专家视频用于表示任务的理想完成方式。
- 视频数据是扩散模型的训练数据。
训练扩散模型 p θ p_\theta pθ:
- 利用专家视频训练扩散模型,让其能够根据历史帧生成高质量的后续帧。
- 这一过程属于离线训练,与强化学习过程分离。
作用: 预训练阶段完成后,扩散模型就能有效生成符合专家轨迹的样本,后续强化学习任务中的奖励信号完全依赖这个模型。
2. 强化学习(实验阶段)从第3行开始的循环(
while not converged
)描述了将预训练的扩散模型应用到强化学习任务中的具体过程。这里的内容属于实验阶段:
扩散模型的应用:
- 在每个时间步中,扩散模型基于智能体的历史观测生成后续帧样本。
- 通过这些样本估计条件熵,并将其转化为奖励信号。
- 这个奖励信号被用于指导智能体的行为优化。
强化学习的主流程:
- 智能体在环境中执行动作,根据环境状态和奖励信号优化策略。
- 策略的学习目标是最大化由扩散奖励生成的回报。
作用: 这一阶段是将奖励模型实际用于强化学习任务中,最终目标是验证奖励信号的有效性,以及智能体在稀疏奖励场景下的表现。
总结
训练阶段:
- 包括算法第1行和第2行。
- 扩散模型在离线环境中进行预训练,与实验过程无关。
实验阶段:
- 从第3行开始描述的内容,是强化学习任务的实验过程。
- 智能体在实际任务环境中利用预训练好的扩散模型生成奖励,逐步优化策略。
最终结论: 该算法的核心是通过离线训练奖励模型和在线强化学习优化的结合,将扩散奖励应用到实际任务的强化学习实验中。因此,它不仅仅是训练过程,还包含了将训练好的奖励模型应用到实验中的完整流程。
4.1 基于扩散模型的专家视频建模
扩散模型 [41] 是通过反向扩散过程 [16],逐步对初始分布去噪以建模数据分布的概率模型。这些模型在捕捉高度复杂的分布以及生成具有复杂动态、运动和行为的样本方面,在强化学习文献中展现了其强大的能力 [2,19]。与之前将专家视频建模为规划器的研究不同,我们旨在从在专家视频上训练的扩散模型中学习奖励函数,用于强化学习。这促使我们的视频模型实现快速推理速度并编码足够的时间信息。
潜在扩散过程。 具体来说,我们首先从专家视频中无监督地训练一个编码器,用于压缩高维观察数据。在这里,我们使用VQ-GAN方法 [12] 将图像 x x x表示为一个矢量量化的潜在代码 z = Q ( E ( x ) ) z = Q(E(x)) z=Q(E(x)),其中 E E E是编码器, Q Q Q是逐元素的量化器。整个视频随后由一系列潜在变量 τ = { z 0 , z 1 , . . . , z K − 1 } \tau = \{z_0, z_1, ..., z_{K-1}\} τ={z0,z1,...,zK−1}表示,我们在不引起歧义的情况下重新定义 τ \tau τ。
接下来,前向过程在潜在空间中对每个时间步 t ∈ { 0 , . . . , T } t \in \{0, ..., T\} t∈{0,...,T}的数据分布 z k z_k zk加入噪声 ϵ \epsilon ϵ,生成一个带噪声的样本 z k t z_{k}^t zkt,其中 z k t = α ˉ t z k 0 + 1 − α ˉ t ϵ z_{k}^t = \sqrt{\bar{\alpha}_t} z_{k}^0 + \sqrt{1 - \bar{\alpha}_t} \epsilon zkt=αˉtzk0+1−αˉtϵ, α ˉ t \bar{\alpha}_t αˉt是过去时间步中噪声计划的累积。为了拟合数据分布,我们学习一个参数化的噪声预测器 ϵ θ ( z k t ) \epsilon_\theta(z_{k}^t) ϵθ(zkt),其目标是在前向过程中预测噪声 ϵ \epsilon ϵ。然后,可以通过对初始分布进行迭代去噪来近似和执行参数化的反向过程 p θ ( z k t − 1 ∣ z k t ) p_\theta(z_{k}^{t-1}|z_{k}^t) pθ(zkt−1∣zkt)。
1. 高维数据的压缩:使用VQ-GAN编码器
背景: 在高维观测空间(如RGB图像)中直接操作通常非常复杂且计算昂贵,因此需要将高维数据压缩到一个更紧凑的潜在表示空间中,以便后续处理。
VQ-GAN方法:
VQ-GAN(Vector Quantized Generative Adversarial Network):
- VQ-GAN是一种通过矢量量化(Vector Quantization)来学习数据紧凑表示的方法。
- 将高维数据(如图像 x x x)编码为潜在空间中的离散代码(latent codes)。
核心步骤:
- 编码器 E E E: 提取高维输入数据 x x x的潜在特征。
- 例如,一个 128 × 128 128 \times 128 128×128的图像被映射到低维潜在空间。
- 量化器 Q Q Q: 将编码器生成的连续特征映射为离散的潜在代码 z z z,通过矢量量化选择最接近的离散表示。
- 公式表示: z = Q ( E ( x ) ) z = Q(E(x)) z=Q(E(x))。
- 潜在表示 τ \tau τ:
- 单个视频由多帧组成,使用潜在代码序列 τ = { z 0 , z 1 , . . . , z K − 1 } \tau = \{z_0, z_1, ..., z_{K-1}\} τ={z0,z1,...,zK−1}来表示。
示例: 假设有一个 64 × 64 64 \times 64 64×64的机器人操作图像帧 x x x,VQ-GAN会将其压缩为一个低维的离散向量 z z z,如 z = [ 1 , 45 , 67 ] z = [1, 45, 67] z=[1,45,67]。整个视频帧序列被表示为潜在代码的序列 τ \tau τ。
2. 扩散模型的前向过程
扩散模型简介: 扩散模型是一种生成模型,它通过逐步加入噪声生成带噪的样本(前向过程),然后通过逐步去噪生成无噪声数据(反向过程)。
前向过程: 在潜在空间中,每个时间步 t t t对潜在变量 z k z_k zk进行以下步骤:
加入噪声:
- 对每个初始潜在变量 z k 0 z_k^0 zk0(即无噪声的潜在表示),加入一个噪声 ϵ \epsilon ϵ生成带噪样本 z k t z_k^t zkt:
z k t = α ˉ t z k 0 + 1 − α ˉ t ϵ z_k^t = \sqrt{\bar{\alpha}_t} z_k^0 + \sqrt{1 - \bar{\alpha}_t} \epsilon zkt=αˉtzk0+1−αˉtϵ- 其中:
- α ˉ t \bar{\alpha}_t αˉt表示时间步 t t t的噪声计划的累积。
- ϵ \epsilon ϵ是高斯噪声,用于扰乱数据。
- 随着 t t t的增加, α ˉ t \sqrt{\bar{\alpha}_t} αˉt逐步减小,噪声分量 1 − α ˉ t \sqrt{1 - \bar{\alpha}_t} 1−αˉt逐步增加,最终将数据变成完全随机的噪声。
累积噪声计划:
- α ˉ t \bar{\alpha}_t αˉt的值通常根据预设计划递减,例如线性或指数递减,控制噪声增加的速度。
目的:
- 模拟潜在表示的扰动过程,为后续反向去噪学习提供训练数据。
’
示例: 假设初始潜在代码 z k 0 = [ 1.0 , 0.5 ] z_k^0 = [1.0, 0.5] zk0=[1.0,0.5]:
- 在时间步 t = 1 t=1 t=1,加入小噪声 ϵ = [ − 0.1 , 0.2 ] \epsilon=[-0.1, 0.2] ϵ=[−0.1,0.2],得到 z k 1 = [ 0.9 , 0.7 ] z_k^1 = [0.9, 0.7] zk1=[0.9,0.7]。
- 在 t = 5 t=5 t=5时,加入更强的噪声,可能得到 z k 5 = [ 0.2 , − 0.1 ] z_k^5 = [0.2, -0.1] zk5=[0.2,−0.1]。
3. 噪声预测器的训练
噪声预测器 ϵ θ ( z k t ) \epsilon_\theta(z_k^t) ϵθ(zkt):
- 模型的任务是学习如何预测前向过程中加入的噪声 ϵ \epsilon ϵ。
- 训练目标:
- 通过学习 ϵ θ ( z k t ) \epsilon_\theta(z_k^t) ϵθ(zkt),使模型能够从带噪样本 z k t z_k^t zkt中准确还原噪声 ϵ \epsilon ϵ。
损失函数:
- 常用均方误差(MSE)作为损失函数: L = E z k t , ϵ [ ∥ ϵ − ϵ θ ( z k t ) ∥ 2 ] \mathcal{L} = \mathbb{E}_{z_k^t, \epsilon}\big[\|\epsilon - \epsilon_\theta(z_k^t)\|^2\big] L=Ezkt,ϵ[∥ϵ−ϵθ(zkt)∥2]
- 这里 ϵ \epsilon ϵ是实际加入的噪声, ϵ θ ( z k t ) \epsilon_\theta(z_k^t) ϵθ(zkt)是模型预测的噪声。
示例: 假设 t = 3 t=3 t=3时, z k 3 = [ 0.8 , − 0.3 ] z_k^3 = [0.8, -0.3] zk3=[0.8,−0.3],实际噪声 ϵ = [ 0.1 , − 0.1 ] \epsilon=[0.1, -0.1] ϵ=[0.1,−0.1],模型预测的噪声为 ϵ θ ( z k 3 ) = [ 0.08 , − 0.12 ] \epsilon_\theta(z_k^3) = [0.08, -0.12] ϵθ(zk3)=[0.08,−0.12]。通过MSE计算误差并优化模型。
4. 反向过程
参数化反向过程:
- 学习一个参数化的反向过程 p θ ( z k t − 1 ∣ z k t ) p_\theta(z_k^{t-1}|z_k^t) pθ(zkt−1∣zkt),用于将带噪样本逐步去噪,最终还原为无噪声的潜在代码 z k 0 z_k^0 zk0。
反向过程步骤:
- 从最初的完全随机噪声 z k T z_k^T zkT开始。
- 在每个时间步 t t t,使用预测的噪声 ϵ θ ( z k t ) \epsilon_\theta(z_k^t) ϵθ(zkt)计算去噪样本: z k t − 1 = f ( z k t , ϵ θ ( z k t ) ) z_k^{t-1} = f(z_k^t, \epsilon_\theta(z_k^t)) zkt−1=f(zkt,ϵθ(zkt))
- f f f是去噪函数,根据扩散模型的具体实现可能会有所不同。
- 迭代进行,最终生成无噪声的潜在表示 z k 0 z_k^0 zk0。
目的:
- 从噪声中逐步恢复潜在代码 z k 0 z_k^0 zk0,并重建完整视频帧。
示例: 如果 z k T z_k^T zkT是完全随机噪声,例如 z k T = [ 0.0 , − 0.5 ] z_k^T=[0.0, -0.5] zkT=[0.0,−0.5],反向过程通过去噪逐步恢复为接近原始潜在代码的状态,例如 z k 0 = [ 1.0 , 0.5 ] z_k^0 = [1.0, 0.5] zk0=[1.0,0.5]。
总结知识点
潜在表示:
- 使用VQ-GAN对高维数据进行压缩,将图像 x x x映射为低维的潜在表示 z z z。
- 通过潜在代码 τ \tau τ表示整个视频序列。
扩散模型的前向过程:
- 加入噪声生成带噪样本,为后续的去噪过程提供训练数据。
- 噪声计划( α ˉ t \bar{\alpha}_t αˉt)控制噪声增加的速度。
噪声预测器:
- 训练噪声预测模型 ϵ θ \epsilon_\theta ϵθ,学习还原噪声。
- 使用均方误差(MSE)作为损失函数优化模型。
反向去噪过程:
- 从随机噪声逐步恢复潜在表示 z k 0 z_k^0 zk0,最终生成符合数据分布的样本。
现实应用中的举例
机器人操作任务:
- 高维状态压缩:
- 将机器人抓取任务中的视频帧压缩为潜在代码序列 τ \tau τ。
- 扩散模型的作用:
- 前向过程通过加入噪声生成带噪潜在代码,反向过程通过去噪生成符合机器人行为的潜在表示。
- 奖励信号生成:
- 根据模型生成的潜在表示与专家潜在轨迹的匹配程度,为强化学习提供奖励。
历史帧作为条件。 为了结合视频扩散的能力利用专家视频中的时间信息,我们进一步在反向过程中加入历史帧的条件,即 p θ ( z k t − 1 ∣ z k t , z c ) p_\theta(z_{k}^{t-1}|z_{k}^t, z_c) pθ(zkt−1∣zkt,zc),其中 z c z_c zc是历史帧 ℓ \ell ℓ的拼接,即 [ z k − ℓ , . . . , z k − 1 ] [z_{k-\ell}, ..., z_{k-1}] [zk−ℓ,...,zk−1],以在保持足够历史信息的同时确保较高的计算效率。这也可以被视为匹配专家轨迹和智能体轨迹的分布 [10]。随后,可以从随机采样的噪声开始执行反向过程,生成未来帧的潜在代码,并解码该代码用于视频预测。
在本工作中,我们选择VQ-Diffusion [13] 作为视频扩散模型,因为它在与矢量量化潜在代码的兼容性和性能方面表现出色,但我们的框架原则上可以采用任何现成的视频扩散模型。具体来说,我们首先将每个潜在代码
z
z
z通过其在VQ-GAN的学习代码簿中的索引进行标记,并将拼接的标记作为条件嵌入。然后将嵌入输入到包含多个Transformer
块和带有交叉注意力的Softmax
层的解码器中,为视频扩散提供上下文信息。
这一部分描述了如何在反向过程中引入历史帧作为条件,以利用时间信息更好地生成符合专家行为的视频帧。以下是逐点详细解析:
1. 引入历史帧的条件
条件分布的定义:
- 反向过程公式: p θ ( z k t − 1 ∣ z k t , z c ) p_\theta(z_k^{t-1}|z_k^t, z_c) pθ(zkt−1∣zkt,zc) 其中:
- z k t − 1 z_k^{t-1} zkt−1:前一个时间步 t − 1 t-1 t−1的潜在表示。
- z k t z_k^t zkt:当前时间步 t t t的潜在表示。
- z c z_c zc:历史帧的拼接,表示为 [ z k − ℓ , . . . , z k − 1 ] [z_{k-\ell}, ..., z_{k-1}] [zk−ℓ,...,zk−1],包含 ℓ \ell ℓ个历史帧的潜在表示。
为什么需要历史帧条件?
利用时间信息:
- 历史帧条件提供了时间上的上下文信息,使生成过程能够利用帧间的动态一致性。
- 例如,机械臂接近目标的过程需要参考之前的运动轨迹。
提高计算效率:
- 通过拼接固定数量的历史帧( ℓ \ell ℓ个),限制输入规模,避免处理整个视频历史,保持计算效率。
示例: 假设 z c = [ z k − 3 , z k − 2 , z k − 1 ] z_c = [z_{k-3}, z_{k-2}, z_{k-1}] zc=[zk−3,zk−2,zk−1]表示最近三帧的潜在表示:
- 如果当前帧 z k t z_k^t zkt表示机械臂的状态,历史帧 z c z_c zc包含之前机械臂接近目标的运动轨迹。
- 引入 z c z_c zc作为条件,有助于生成下一帧时考虑到机械臂的动态连续性。
k k k 和 t t t 是两个不同的指标,它们在扩散模型的上下文中表示不同的概念。下面详细解释 k k k 和 t t t 的含义以及它们的区别:
1. k k k 的含义:时间步(序列中的帧索引)
- k k k 表示在视频序列中的帧索引:
- 在强化学习或视频扩散中,视频是由多个时间帧组成的, k k k 是这些帧的时间步索引。
- 例如,对于一个视频序列 τ = { z 0 , z 1 , z 2 , … , z K − 1 } \tau = \{z_0, z_1, z_2, \dots, z_{K-1}\} τ={z0,z1,z2,…,zK−1}, z k z_k zk 是第 k k k 个帧的潜在表示。
- 与历史帧的关系:
- 在历史帧 z c z_c zc的定义中, z c z_c zc 是由时间步 k − ℓ k-\ell k−ℓ 到 k − 1 k-1 k−1 的潜在表示拼接而成:
z c = [ z k − ℓ , z k − ℓ + 1 , … , z k − 1 ] z_c = [z_{k-\ell}, z_{k-\ell+1}, \ldots, z_{k-1}] zc=[zk−ℓ,zk−ℓ+1,…,zk−1]- 这里的 k k k 是当前帧的索引,而 z c z_c zc 表示当前帧之前的 ℓ \ell ℓ 个历史帧。
总结:
- k k k 是视频序列中的帧索引,代表特定的时刻。
2. t t t 的含义:扩散过程中的时间步
- t t t 表示扩散过程的时间步:
- 扩散模型将一个潜在表示(如 z k 0 z_k^0 zk0)逐步加入噪声,生成一系列带噪样本(如 z k t z_k^t zkt)。 t t t 表示扩散过程中的时间步索引。
- 在扩散过程的第 t t t 步, z k t z_k^t zkt 表示加入了 t t t 步噪声后的潜在表示:
z k t = α ˉ t z k 0 + 1 − α ˉ t ϵ z_k^t = \sqrt{\bar{\alpha}_t} z_k^0 + \sqrt{1 - \bar{\alpha}_t} \epsilon zkt=αˉtzk0+1−αˉtϵ- 当 t = 0 t=0 t=0 时, z k 0 z_k^0 zk0 是未加噪的原始潜在表示;
- 当 t = T t=T t=T 时, z k T z_k^T zkT 是完全加入噪声的潜在表示。
- t t t 的范围:
- 通常, t t t 从 0 0 0 到 T T T,表示扩散过程的时间步数。
总结:
- t t t 是扩散模型中的噪声时间步,描述扩散过程的进展程度。
3. k k k 和 t t t 的区别
特性 k k k(帧索引) t t t(扩散时间步) 含义 视频序列中某一帧的索引 扩散过程中的噪声时间步 应用范围 定位某一视频帧 描述扩散模型中加入噪声的步骤 范围 k ∈ { 0 , 1 , … , K − 1 } k \in \{0, 1, \dots, K-1\} k∈{0,1,…,K−1} t ∈ { 0 , 1 , … , T } t \in \{0, 1, \dots, T\} t∈{0,1,…,T} 是否独立 与扩散过程无关 与视频帧索引无关 历史帧的使用 拼接最近的 ℓ \ell ℓ帧,表示为 z c z_c zc 不直接参与历史帧定义 扩散模型中的作用 当前视频帧 z k z_k zk会作为扩散模型的输入条件 决定当前潜在表示 z k t z_k^t zkt的带噪程度
4. 结合上下文的解释
k k k 和 t t t 的区别如下:
k k k(帧索引):
- 表示视频序列中当前帧的索引。
- z c z_c zc 是基于 k k k 的历史帧集合,用来提供上下文信息。
- 例如, z c = [ z k − 3 , z k − 2 , z k − 1 ] z_c = [z_{k-3}, z_{k-2}, z_{k-1}] zc=[zk−3,zk−2,zk−1] 是视频帧索引 k k k 之前的历史帧。
t t t(扩散时间步):
- 表示当前帧 z k z_k zk 在扩散过程中的时间步。
- 例如, z k t z_k^t zkt 是当前帧 z k z_k zk 在扩散时间步 t t t 下的带噪表示。
5. 举例说明
假设有一个机器人操作任务的短视频,由 K = 10 K=10 K=10 帧组成,扩散过程设置为 T = 5 T=5 T=5 个时间步:
帧索引 k k k 的意义:
- 第 6 帧的潜在表示是 z k z_k zk,其中 k = 5 k=5 k=5。
- 历史帧条件 z c = [ z 2 , z 3 , z 4 ] z_c = [z_2, z_3, z_4] zc=[z2,z3,z4],对应 k = 2 , 3 , 4 k=2, 3, 4 k=2,3,4 的潜在表示。
扩散时间步 t t t 的意义:
- 如果第 6 帧的初始潜在表示是 z k 0 z_k^0 zk0:
- 在 t = 1 t=1 t=1 时,加入轻微噪声生成 z k 1 z_k^1 zk1;
- 在 t = 5 t=5 t=5 时,完全加入噪声生成 z k 5 z_k^5 zk5。
结合两者:
- 当前帧 k = 5 k=5 k=5 的带噪潜在表示可以在扩散时间步 t = 3 t=3 t=3 时表示为 z k 3 z_k^3 zk3;
- 历史帧 z c z_c zc 是前几个帧(如 z k − 3 , z k − 2 , z k − 1 z_{k-3}, z_{k-2}, z_{k-1} zk−3,zk−2,zk−1)的初始潜在表示(我理解的就是 z k − 3 0 , z k − 2 0 , z k − 1 0 z_{k-3}^0, z_{k-2}^0, z_{k-1}^0 zk−30,zk−20,zk−10,这里需要注意如果有不同见解欢迎评论区讨论,我也不太确定!!!)。
6. 总结
- k k k 和 t t t 是两个不同的维度:
- k k k 表示视频序列中的帧索引,用于定义历史帧和轨迹。
- t t t 表示扩散模型中的时间步,用于描述带噪潜在表示的生成过程。
- 两者的联系:
- 某一帧 z k z_k zk 在扩散过程的每个时间步 t t t 下都有对应的带噪潜在表示 z k t z_k^t zkt。
- 历史帧 z c z_c zc 中的帧索引 k − ℓ , . . . , k − 1 k-\ell, ..., k-1 k−ℓ,...,k−1 通常指代这些帧的原始潜在表示 z k 0 z_k^0 zk0,不涉及扩散时间步。
2. 匹配专家轨迹与智能体轨迹
概念: 通过条件分布 p θ ( z k t − 1 ∣ z k t , z c ) p_\theta(z_k^{t-1}|z_k^t, z_c) pθ(zkt−1∣zkt,zc),模型学习如何生成智能体轨迹,使其分布尽可能接近专家轨迹分布。
方法:
- 在反向过程中使用历史帧 z c z_c zc指导生成:
- 通过拼接的历史帧提供上下文信息。
- 生成的未来帧不仅与当前帧一致,还能够捕捉历史帧中的动态模式。
- 匹配分布的目标:
- 学习专家轨迹中隐含的动态特性,例如运动连续性或任务执行步骤的模式。
示例: 在机器人抓取任务中:
- 专家轨迹可能表现为机械臂逐渐接近目标物体的连续动作。
- 智能体轨迹通过历史帧条件学习这一分布,使生成的潜在表示与专家轨迹相似。
3. 视频扩散模型:VQ-Diffusion
为什么选择VQ-Diffusion (想知道更多VQ-Diffusion的内容可以阅读《Vector quantized diffusion model for text-to-image synthesis》这篇文章)?
- 兼容性:
- VQ-Diffusion能够很好地处理VQ-GAN生成的矢量量化潜在代码 z z z。
- 性能优势:
- VQ-Diffusion在生成高质量视频方面具有较强的表现,特别是在与潜在表示结合时。
- 灵活性:
- 虽然本工作选择了VQ-Diffusion,但框架可以替换为其他视频扩散模型。
4. 潜在代码的嵌入与解码
潜在代码的标记与嵌入:
标记潜在代码:
- 每个潜在代码 z z z通过其在VQ-GAN学习的代码簿中的索引进行标记。
- 代码簿索引:
- VQ-GAN生成的每个潜在表示对应一个离散的索引,这些索引被用作模型输入。
- 作用:
- 索引提供了潜在表示的离散编码,便于模型理解和处理。
条件嵌入:
- 将拼接的标记(历史帧 z c z_c zc的索引)作为条件嵌入。
- 嵌入提供了输入的上下文信息,增强生成模型对时间信息的感知。
示例: 如果潜在表示 z z z对应的索引为 [ 1 , 45 , 67 ] [1, 45, 67] [1,45,67],拼接后的历史帧条件可能是 [ 1 , 45 , 67 , 12 , 34 ] [1, 45, 67, 12, 34] [1,45,67,12,34],这些索引通过嵌入映射到特征空间。
解码过程:
- 输入到解码器:
- 嵌入后的潜在代码作为输入,被传递到解码器中。
- Transformer块:
- 解码器包含多个Transformer块,用于处理潜在代码及其上下文信息。
- 交叉注意力:
- 通过交叉注意力机制,解码器能够更好地结合历史帧条件,生成符合上下文的未来帧。
- Softmax层:
- 最终使用Softmax层输出生成的潜在代码,用于反向过程中的去噪操作。
示例: 解码器接收潜在代码的嵌入,例如 z c = [ 12 , 34 , 56 ] z_c=[12, 34, 56] zc=[12,34,56],通过交叉注意力生成下一帧的潜在表示 z k t − 1 z_k^{t-1} zkt−1。
5. 从潜在代码到视频帧
反向过程的最终结果:
- 通过从随机噪声中迭代去噪,生成未来帧的潜在代码。
- 解码器将生成的潜在代码解码为完整的视频帧。
示例: 假设生成的潜在代码为 z k 0 = [ 1.2 , − 0.8 ] z_k^0 = [1.2, -0.8] zk0=[1.2,−0.8],通过VQ-GAN解码器解码为机器人抓取物体的图像帧。
总结知识点
历史帧条件:
- 通过拼接 ℓ \ell ℓ个历史帧作为条件,为视频扩散提供时间上下文信息。
- 匹配专家轨迹与智能体轨迹的分布,增强模型对时间一致性的理解。
VQ-Diffusion模型:
- 选择VQ-Diffusion作为视频扩散模型,结合VQ-GAN生成的潜在代码实现高效的潜在空间建模。
Transformer块与交叉注意力:
- 在解码过程中使用Transformer块和交叉注意力,将历史帧条件与当前帧生成过程结合,提高生成质量。
从潜在代码到视频帧的映射:
- 利用学习的潜在表示,通过解码器解码为最终的视频帧。
4.2 条件熵作为奖励
尽管之前的研究探索了将对数似然作为视频预测模型的奖励,例如VIPER [10],但这种方法面临两个主要挑战。首先,它在准确建模具有复杂动态特征的专家视频分布方面存在困难,如表1所示。其次,中等的视频建模能力导致了不理想的奖励。这一问题在图2中表现得尤为明显,即分布内的专家视频和分布外的专家视频之间学习到的奖励出现显著下降,尽管两组视频都展示了最优行为。
Diffusion Reward 的关键见解。 我们通过利用视频扩散模型的强大生成能力来解决这些挑战。我们的观察表明,未见过的历史观测(随机)会增加生成的多样性,而已见过的观测(专家)则会降低多样性,如表2所示。这引出了我们提出的Diffusion Reward的关键见解:以类似专家的轨迹为条件的扩散表现出较低的多样性,此时智能体应该获得更多的奖励;而在非专家的轨迹上则相反。这种多样性区分激励强化学习智能体追逐类似专家的行为。
条件熵的估计。 为了形式化这一思想,我们寻求估计给定历史帧 z c z_c zc的负条件熵,该熵原则上能够捕捉条件生成的多样性:
− H ( p θ ( ⋅ ∣ z c ) ) = E p θ ( ⋅ ∣ z c ) [ log p θ ( z k ∣ z c ) ] . (1) -H(p_\theta(\cdot | z_c)) = \mathbb{E}_{p_\theta(\cdot | z_c)}[\log p_\theta(z_k | z_c)].\tag{1} −H(pθ(⋅∣zc))=Epθ(⋅∣zc)[logpθ(zk∣zc)].(1)
计算公式中的熵(如公式(1)所示)的一个主要挑战是其不可解性,因为我们没有条件分布的显式形式 [43]。因此,我们改为估计熵的变分界限。具体而言,我们首先给出条件对数似然的变分界限如下:
log p θ ( z k 0 ∣ z c ) ≥ E q ( z k 0 : T ∣ z c ) [ log p θ ( z k 0 : T ) q ( z k 1 : T ∣ z k 0 , z c ) ] , (2) \log p_\theta(z_k^0 | z_c) \geq \mathbb{E}_{q(z_k^{0:T} | z_c)} \left[ \log \frac{p_\theta(z_k^{0:T})}{q(z_k^{1:T} | z_k^0, z_c)} \right],\tag{2} logpθ(zk0∣zc)≥Eq(zk0:T∣zc)[logq(zk1:T∣zk0,zc)pθ(zk0:T)],(2)
其中 z k 0 z_k^0 zk0是当前观测 z k z_k zk的去噪预测。该界限可以通过噪声预测器 ϵ θ \epsilon_\theta ϵθ [16,42] 或闭式分布 [20,41](例如离散多变量分布)进行估计。我们选择后者作为估计方法,因为它与VQ-Diffusion的兼容性更好。
接下来,为了估计整体熵,我们从随机采样的噪声中去噪并生成潜在变量 z ~ k 0 : T ∼ p θ ( z k 0 : T ∣ z c ) \tilde{z}_k^{0:T} \sim p_\theta(z_k^{0:T} | z_c) z~k0:T∼pθ(zk0:T∣zc),重复 M M M次。随后,我们使用从随机噪声 z ~ k T \tilde{z}_k^T z~kT(例如随机标记)生成的样本来估计整体条件熵项,其表达式如下:
r c e ( x k − 1 ) = 1 M ∑ j = 1 M log p θ ( z ~ k 0 : T ) q ( z ~ k 1 : T ∣ z ~ k 0 , z c ) , (3) r^{ce}(x_{k-1}) = \frac{1}{M} \sum_{j=1}^{M} \log \frac{p_\theta(\tilde{z}_k^{0:T})}{q(\tilde{z}_k^{1:T} | \tilde{z}_k^0, z_c)}, \tag{3} rce(xk−1)=M1j=1∑Mlogq(z~k1:T∣z~k0,zc)pθ(z~k0:T),(3)
我们在图3中可视化了基于聚合熵的奖励,并在附录中呈现了每个任务的曲线。结果表明,条件熵能够成功捕捉不同视频上生成多样性的变化,这与前文提到的Diffusion Reward的见解相呼应。值得注意的是,我们在此使用标准化的熵奖励 r ˉ c e \bar{r}_{ce} rˉce以减轻超参数调节的负担,因为我们观察到条件熵的尺度在不同任务和领域之间显著变化,这部分归因于对象和环境动态的差异。具体而言,条件熵通过专家视频的经验均值和标准差进行标准化:
r ˉ c e = r c e − mean ( D , r c e ) std ( D , r c e ) . (4) \bar{r}^{ce} = \frac{r^{ce} - \text{mean}(\mathcal{D}, r^{ce})}{\text{std}(\mathcal{D}, r^{ce})}.\tag{4} rˉce=std(D,rce)rce−mean(D,rce).(4)
图3. 奖励分析。(左)我们学习到的奖励在MetaWorld的7个任务中对专家行为的聚合值更高。(右)不同质量的轨迹可以在已见任务和未见任务中区分。(上)轨迹示例。
1. 条件熵的定义与目标
条件熵的定义: 条件熵 H ( p θ ( ⋅ ∣ z c ) ) H(p_\theta(\cdot | z_c)) H(pθ(⋅∣zc))用于衡量给定历史帧 z c z_c zc时,模型生成未来潜在变量 z k z_k zk的不确定性。
公式(1): − H ( p θ ( ⋅ ∣ z c ) ) = E p θ ( ⋅ ∣ z c ) [ log p θ ( z k ∣ z c ) ] . -H(p_\theta(\cdot | z_c)) = \mathbb{E}_{p_\theta(\cdot | z_c)}[\log p_\theta(z_k | z_c)]. −H(pθ(⋅∣zc))=Epθ(⋅∣zc)[logpθ(zk∣zc)].
- − H -H −H: 我们计算的是负条件熵,作为奖励信号。
- p θ ( z k ∣ z c ) p_\theta(z_k | z_c) pθ(zk∣zc): 表示模型在条件 z c z_c zc(历史帧)下生成当前帧 z k z_k zk的概率分布。
- 直观理解:
- 如果生成的分布 p θ ( z k ∣ z c ) p_\theta(z_k | z_c) pθ(zk∣zc)集中(确定性高),条件熵较低,负条件熵较高。
- 如果生成分布较分散(不确定性高),条件熵较高,负条件熵较低。
目标: 条件熵捕捉生成样本的多样性,同时奖励那些在给定条件下生成确定性较高的样本。
2. 条件熵的不可解性及变分界限
为什么条件熵不可直接求解?
- 没有显式形式:
- 条件分布 p θ ( z k ∣ z c ) p_\theta(z_k | z_c) pθ(zk∣zc)通常是隐式定义的,难以直接求解其熵。
- 解决方案:
- 使用变分界限为条件熵提供一个可计算的下界。
变分界限公式: 公式(2): log p θ ( z k 0 ∣ z c ) ≥ E q ( z k 0 : T ∣ z c ) [ log p θ ( z k 0 : T ) q ( z k 1 : T ∣ z k 0 , z c ) ] . \log p_\theta(z_k^0 | z_c) \geq \mathbb{E}_{q(z_k^{0:T} | z_c)} \left[ \log \frac{p_\theta(z_k^{0:T})}{q(z_k^{1:T} | z_k^0, z_c)} \right]. logpθ(zk0∣zc)≥Eq(zk0:T∣zc)[logq(zk1:T∣zk0,zc)pθ(zk0:T)].
符号解释:
- z k 0 z_k^0 zk0:当前帧的去噪潜在表示。
- z k 0 : T z_k^{0:T} zk0:T:从去噪帧到带噪帧的完整潜在变量序列。
- q ( z k 1 : T ∣ z k 0 , z c ) q(z_k^{1:T} | z_k^0, z_c) q(zk1:T∣zk0,zc):后验分布,表示在给定 z k 0 z_k^0 zk0和 z c z_c zc条件下,生成的带噪帧序列。
- p θ ( z k 0 : T ) p_\theta(z_k^{0:T}) pθ(zk0:T):模型生成的潜在变量序列的概率。
直观解释:
- 通过估计 p θ ( z k 0 ∣ z c ) p_\theta(z_k^0 | z_c) pθ(zk0∣zc)的变分下界,我们可以间接估计条件熵。
为什么使用变分界限?
- 它提供了一个可计算的形式,通过优化 q q q来近似 p θ p_\theta pθ。
为什么选择离散多变量分布?
- VQ-Diffusion 的兼容性:
- VQ-Diffusion使用矢量量化(离散潜在变量),离散多变量分布能够直接适配这种结构。
- 实现简单:
- 对于离散潜在空间,计算熵变得更加高效且直观。
3. 条件熵的具体估计方法
(a) 从随机噪声生成潜在变量
为了估计整体熵,采用以下方法:
- 从随机噪声中采样 z ~ k T \tilde{z}_k^T z~kT(带噪潜在表示)。
- 通过反向扩散过程去噪,生成潜在变量序列 z ~ k 0 : T ∼ p θ ( z k 0 : T ∣ z c ) \tilde{z}_k^{0:T} \sim p_\theta(z_k^{0:T} | z_c) z~k0:T∼pθ(zk0:T∣zc)。
- 重复采样 M M M次,生成 M M M组潜在变量序列。
(b) 使用采样结果估计条件熵
公式(3): r c e ( x k − 1 ) = 1 M ∑ j = 1 M log p θ ( z ~ k 0 : T ) q ( z ~ k 1 : T ∣ z ~ k 0 , z c ) . r^{ce}(x_{k-1}) = \frac{1}{M} \sum_{j=1}^{M} \log \frac{p_\theta(\tilde{z}_k^{0:T})}{q(\tilde{z}_k^{1:T} | \tilde{z}_k^0, z_c)}. rce(xk−1)=M1j=1∑Mlogq(z~k1:T∣z~k0,zc)pθ(z~k0:T).
意义:
- r c e ( x k − 1 ) r^{ce}(x_{k-1}) rce(xk−1)是基于条件熵的奖励信号,用于鼓励生成多样性较高的样本。
- 通过采样方法避免显式计算 p θ p_\theta pθ和 q q q。
计算步骤:
- 从模型生成 M M M组潜在变量序列 z ~ k 0 : T \tilde{z}_k^{0:T} z~k0:T。
- 计算每个序列的对数似然比:
log p θ ( z ~ k 0 : T ) q ( z ~ k 1 : T ∣ z ~ k 0 , z c ) . \log \frac{p_\theta(\tilde{z}_k^{0:T})}{q(\tilde{z}_k^{1:T} | \tilde{z}_k^0, z_c)}. logq(z~k1:T∣z~k0,zc)pθ(z~k0:T).- 取所有序列对数似然的均值,作为条件熵的估计值。
4. 熵奖励的标准化
为什么需要标准化?
- 条件熵的绝对值在不同任务和环境中可能差异较大。
- 为了统一奖励尺度,减轻超参数调节的负担,使用标准化方法。
标准化公式: 公式(4): r ˉ c e = r c e − mean ( D , r c e ) std ( D , r c e ) . \bar{r}^{ce} = \frac{r^{ce} - \text{mean}(\mathcal{D}, r^{ce})}{\text{std}(\mathcal{D}, r^{ce})}. rˉce=std(D,rce)rce−mean(D,rce).
符号解释:
- mean ( D , r c e ) \text{mean}(\mathcal{D}, r^{ce}) mean(D,rce):条件熵奖励的经验均值,计算于整个专家视频数据集 D \mathcal{D} D。
- std ( D , r c e ) \text{std}(\mathcal{D}, r^{ce}) std(D,rce):条件熵奖励的经验标准差。
- r ˉ c e \bar{r}^{ce} rˉce:标准化后的熵奖励。
作用:
- 消除任务间的奖励尺度差异,使模型更适应多任务场景。
- 保证奖励信号的稳定性。
5. 举例说明
(a) 条件熵的估计
- 环境:
- 一个机器人抓取任务的视频帧序列,每帧 x k x_k xk通过VQ-GAN编码为潜在表示 z k z_k zk。
- 采样与去噪:
- 从随机噪声 z ~ k T \tilde{z}_k^T z~kT开始,利用扩散模型生成潜在变量序列 z ~ k 0 : T \tilde{z}_k^{0:T} z~k0:T。
- 例如,从 M = 5 M=5 M=5次采样中生成:
- z ~ k 0 : T , 1 , z ~ k 0 : T , 2 , … , z ~ k 0 : T , 5 \tilde{z}_k^{0:T,1}, \tilde{z}_k^{0:T,2}, \dots, \tilde{z}_k^{0:T,5} z~k0:T,1,z~k0:T,2,…,z~k0:T,5。
- 熵估计:
- 对每次采样计算 log p θ ( z ~ k 0 : T ) q ( z ~ k 1 : T ∣ z ~ k 0 , z c ) \log \frac{p_\theta(\tilde{z}_k^{0:T})}{q(\tilde{z}_k^{1:T} | \tilde{z}_k^0, z_c)} logq(z~k1:T∣z~k0,zc)pθ(z~k0:T)。
- 平均后得到 r c e ( x k − 1 ) r^{ce}(x_{k-1}) rce(xk−1)。
(b) 奖励标准化
- 数据集 D \mathcal{D} D包含多个机器人任务的专家视频,条件熵奖励在每个任务中的均值和标准差可能不同。
- 使用公式(4)对奖励信号进行标准化,使得奖励信号在多任务上具有一致性。
6. 结果与意义
条件熵的作用:
- 捕捉生成分布的多样性,鼓励模型在确定性和多样性之间找到平衡。
- 提高Diffusion Reward的有效性。
标准化奖励的效果:
- 减少不同任务间的奖励尺度差异。
- 提高模型在多任务环境中的适应性。
可视化:
- 在图3中,基于条件熵的奖励信号能够成功反映生成分布的变化,说明其在不同视频生成任务中的有效性。
总结
- 条件熵估计:
- 通过变分界限公式和采样方法间接估计条件熵。
- 奖励信号生成:
- 使用条件熵作为奖励,鼓励生成确定性高的样本。
- 标准化:
- 标准化熵奖励解决任务间尺度差异的问题,提高任务适应性和训练稳定性。
4.3 训练细节
我们首先将来自不同任务的专家视频作为整个数据集,用于奖励模型的预训练,这些视频可以通过脚本策略或其他方式生成。然后,我们首先使用VQ-GAN [12] 训练编码器,跨所有领域使用8×8大小的代码簿,并通过附加的感知损失 [55](由判别器计算)来提高感知质量。随后,我们使用VQ-Diffusion [13] 训练条件视频扩散模型,其中历史帧的数量 τ \tau τ在所有任务中设置为2。对于奖励推断,除了Pen任务外,奖励系数在所有任务中均设置为0.95。我们为Pen任务降低了 α \alpha α,因为较高的探索奖励可能会误导机器人手丢掉笔,而较低的奖励则鼓励在手中的操作。此外,在扩散过程中,我们还使用类似DDIM的采样策略 [42] 来加速扩散过程,将去噪步骤设置为10,并在所有任务中将扩散重复次数 M M M设置为1,这在图4中表现出公平的性能并保持较高的推断速度。更多关于超参数、网络架构和下游强化学习的细节可参见附录。
1. 专家视频用于奖励模型预训练
数据集来源:
- 专家视频:不同任务中专家执行的演示视频,作为奖励模型训练的主要数据。
- 生成方式:
- 专家视频可以通过脚本化策略(如预定义的规则)生成。
- 或者通过真实人类演示或经过训练的机器人策略采集。
目的:
- 奖励模型的预训练依赖专家视频捕捉完成任务的关键行为,为后续的强化学习提供指导。
示例: 在机器人操作任务中:
- 一个“抓取物体”任务的专家视频展示了机器人机械臂接近目标、抓取并提升物体的完整过程。
- 奖励模型通过这些视频学习任务的成功标准,例如机械臂接触目标的精确位置和时机。
2. VQ-GAN 的训练
核心组件:
- 编码器:将高维视频帧(如RGB图像)压缩到潜在空间中。
- 代码簿:VQ-GAN 使用一个 8 × 8 8 \times 8 8×8 的离散代码簿,每个代码表示一个潜在变量。
- 感知损失:
- 利用判别器计算感知损失,衡量生成样本与真实数据之间的感知差异。
- 感知损失提升了生成样本的视觉质量,使生成的视频帧更加逼真。
目的:
- 通过VQ-GAN将视频帧映射到潜在空间,为后续扩散模型的训练提供紧凑的潜在表示。
示例: 对于一个 128 × 128 128 \times 128 128×128 的输入帧,VQ-GAN 编码器将其压缩为 16 × 16 16 \times 16 16×16 的潜在表示矩阵,其中每个元素对应代码簿中的一个索引。
3. VQ-Diffusion 的训练
历史帧数量 τ \tau τ:
- 历史帧数量 τ = 2 \tau=2 τ=2,表示扩散模型在生成下一帧时,参考最近两个帧的潜在表示。
- 目的:
- 提供时间上下文,捕捉帧间的动态一致性。
示例: 假设当前帧索引为 k k k,则扩散模型的条件输入为 z k − 2 z_{k-2} zk−2 和 z k − 1 z_{k-1} zk−1,分别表示前两帧的潜在表示。
4. 奖励推断的设置
奖励系数 α \alpha α:
- 奖励系数的作用:
- 控制探索奖励的权重,平衡探索和目标导向行为。
- Pen任务的特殊设置:
- 原因:
- 在Pen任务(机器人操作笔的任务)中,过高的探索奖励可能导致机器人放弃任务目标(如丢掉笔)。
- 因此,将 α \alpha α 调低以鼓励机器人更稳健地操作。
示例: 在Pen任务中:
- 奖励系数降低意味着机器人更关注保持笔的稳定性,而非随机探索其他行为。
5. 扩散采样策略
DDIM 采样策略:
- 定义:
- DDIM(Denoising Diffusion Implicit Models)是一种加速扩散模型采样的策略,通过减少去噪步骤,提高生成效率。
- 去噪步骤设置为10:
- 表示扩散过程中的去噪迭代次数。
- 目的:
- 降低计算开销,同时保持生成质量。
扩散重复次数 M M M:
- 在所有任务中,扩散重复次数 M = 1 M=1 M=1:
- 每次采样生成一个样本,减少重复计算,提升推断速度。
- 性能公平性:
- 图4 表明,即使设置较低的 M M M,扩散模型仍能保持较高的生成质量。
示例: 在机器人操作任务中:
- 假设从噪声开始生成目标潜在变量:
- 使用10步去噪生成目标帧的潜在表示。
- 每次采样生成一个样本,从而高效完成奖励推断。
6. 标准化奖励的必要性
奖励标准化: 奖励信号的分布可能因任务和领域的动态差异而变化(如不同的环境和对象动态)。
- 标准化奖励: r ˉ c e = r c e − mean ( D , r c e ) std ( D , r c e ) . \bar{r}^{ce} = \frac{r^{ce} - \text{mean}(\mathcal{D}, r^{ce})}{\text{std}(\mathcal{D}, r^{ce})}. rˉce=std(D,rce)rce−mean(D,rce).
- 目的:
- 消除奖励尺度差异,使模型更适应不同任务。
示例: 在抓取任务中:
- 奖励信号可能因物体质量或运动路径不同而分布差异显著。
- 标准化后,奖励信号的均值和方差对任务统一,提升训练的稳定性。
7. 训练细节的总结
完整流程:
- 专家视频的收集:
- 通过脚本或人类演示采集用于奖励模型训练的专家视频。
- VQ-GAN 训练:
- 编码视频帧为离散潜在表示,提升感知质量。
- VQ-Diffusion 训练:
- 利用历史帧条件生成未来帧,捕捉帧间动态一致性。
- 奖励推断与标准化:
- 使用条件熵估计生成奖励信号,并通过标准化消除任务间差异。
- 扩散采样策略:
- 使用加速采样策略(DDIM),设置合理去噪步数和扩散次数,提升效率。
设计细节的目的:
- 统一性:
- 跨任务设置统一参数(如 τ = 2 \tau=2 τ=2),保证通用性。
- 效率与质量平衡:
- 使用DDIM加速采样,在性能和推断速度之间取得平衡。
8. 举例应用:Pen任务中的操作
- 任务描述:
- 机器人需要操控笔完成书写或移动。
- 挑战:
- 过高的探索奖励可能导致机器人丢掉笔。
- 解决方案:
- 降低奖励系数 α \alpha α,使机器人更专注于保持笔的稳定性。
- 结果:
- 机器人在保持稳定操作的同时,仍能通过适度的探索优化策略。
9. 总结
- 本节详细描述了奖励模型训练、VQ-GAN 和 VQ-Diffusion 的细节以及扩散采样策略。
- 核心目标是通过设计合理的参数和策略,实现高效的跨任务奖励推断,确保生成质量与速度的平衡。
5 实验
5.1 实验设置
仿真环境。我们旨在展示Diffusion Reward在10个复杂的视觉机器人操作任务中的有效性,包括来自MetaWorld [51]的7个夹爪操作任务和来自Adroit [35]的3个灵巧手操作任务,如图5所示。我们选择这两个仿真环境是因为它们的任务多样性和复杂性。每个任务都与64×64维的RGB图像、±4像素的平移增强 [50] 和0/1稀疏任务奖励相关联。
为了训练奖励模型,我们为每个MetaWorld任务收集了20个专家视频(共5039帧),这些视频通过官方仓库提供的脚本策略生成;为每个Adroit任务收集了50个专家视频(共3346帧),这些视频通过使用高性能强化学习方法 [46] 训练的策略生成。值得注意的是,我们的方法对数据规模的需求较低(参见附录)。对于下游强化学习训练,Adroit任务的交互预算设置为300万次,以确保使用我们的奖励时能够收敛,而MetaWorld任务的预算则根据显著不同的任务复杂性分别设置。
图5. 任务可视化。我们在来自MetaWorld和Adroit的10个具有挑战性的视觉强化学习任务上评估了这些方法,这些任务使用视觉输入和稀疏奖励。选择的任务涵盖了广泛的操作技能。
基线方法。 我们将我们的方法与以下方法进行比较:
- Raw Sparse Reward:使用稀疏任务奖励。这一比较用于测试从视频中添加学习奖励的益处。
- Random Network Distillation (RND, [4]):仅通过寻求新奇性的奖励鼓励探索。这一比较用于测试使用我们预训练奖励为智能体提供奖励的益处。
- Adversarial Motion Priors (AMP, [32]):在线学习一个判别器,根据当前观测区分智能体行为和专家行为。这一比较用于测试学习奖励中编码时间信息和利用寻求新奇性奖励的益处。
- Video Prediction Rewards (VIPER, [10]) 及其标准化变体(VIPER-std):使用VideoGPT [49] 作为视频预测模型,并将智能体观测的预测对数似然作为奖励。这一比较用于测试利用视频扩散模型生成能力以及条件熵作为更具探索性奖励的益处。
为了公平比较,所有方法都与任务奖励相结合,这通常可以提升它们的性能。此外,我们使用DrQv2 [50] 作为强化学习的主干网络,除了奖励预训练(如果存在)外,所有设置保持一致。
5.2 主要结果
与无预训练方法的比较
我们在图1和图7中展示了每种方法在两个仿真领域的成功率学习曲线。结果表明,仅使用稀疏任务奖励能够在相对简单的任务中取得进展,例如Reach和Dial Turn。然而,在更复杂的任务中会遇到显著的挑战,例如灵巧手操作领域中的Door和Hammer任务。
纯粹加入寻求新奇性的奖励(RND)虽然可以预料地增强强化学习智能体的探索能力,尤其是在应对中等复杂的任务(如Coffee Push和Assembly)时表现明显。然而,由于缺乏专家指导的探索,在配置空间极大的灵巧手操作任务中表现仍不理想。
相反,AMP显式地利用专家指导的奖励,激励对类似专家行为的探索。尽管它在简单任务中通常优于RND,但在更复杂的任务(例如Lever Pull和Hammer)中其效率显著下降。
图7. 主要结果。我们的方法和基线方法在MetaWorld的7个夹爪操作任务和Adroit的3个灵巧操作任务(使用图像观测)的成功率。我们的方法在所有任务中均表现出色,且在复杂的Door和Hammer任务中显著优于基线方法。结果为5次运行的平均值,阴影区域表示标准误差。(在线彩色图)
与奖励预训练方法的比较。 上述观察表明,结合专家指导奖励和寻求新奇性的奖励可能表现更优。尽管VIPER奖励结合了这种组合,但其实际表现意外地低于RND和AMP。我们认为,这归因于专家指导奖励(即对数似然)的尺度变化显著,可能削弱了寻求新奇性奖励的效果。VIPER-std缓解了这一问题,在Coffee Push和Pen任务上表现优于未标准化版本,但仍不及我们的方法。
这些结果与第4.2节中列出的两个限制一致,并在图2中进一步验证,表明VIPER在复杂任务中捕捉复杂视频分布方面存在困难。与VIPER形成鲜明对比的是,我们提出的方法不仅利用了扩散模型的建模能力,还使用条件熵作为奖励函数以加速探索。结果显示,在相同训练步骤下,我们的方法在MetaWorld和Adroit上分别比表现最好的基线方法提升了38%和35%。
5.3 零样本奖励泛化
奖励分析 视频扩散模型在生成超出其训练数据范围的样本方面的进展(例如文本到图像的视频生成 [11,15]),激发了我们探索Diffusion Reward在未见任务中泛化潜力的兴趣。为此,我们首先对从MetaWorld中15个不同未见任务中生成的不同质量轨迹的学习回报进行可视化(见图3)。正如预期的那样,与预训练中见过的任务相比,不同质量轨迹之间的差异不那么显著。尽管如此,我们的预训练奖励模型仍然表现出一致的趋势,即类似专家的行为获得了相对较高的学习回报,这得益于视频扩散模型的泛化能力。
强化学习性能 随后,我们将预训练的奖励直接应用于包含不同对象的5个任务(更多细节见附录),无需额外调整。如图8所示,结果表明我们的奖励能够有效引导强化学习的探索,并在所有任务中显著优于其他基线方法。值得注意的是,我们的方法在约束RND探索空间方面表现出色,这是因为它保留了区分类似专家和非专家行为的能力。同时,VIPER在大多数任务上泛化效果不如我们的方法,尤其是在复杂任务(如打开物体)上,这部分归因于采用的视频模型和对数似然奖励的局限性。
这里我们未与AMP进行比较,因为它在线学习奖励,需要目标任务上不可用的专家数据。这些结果不仅验证了我们方法的有效性,还表明采用更大的扩散模型并整合其他模态(例如基于文本的任务规范)具有进一步增强我们方法泛化能力的潜力。
图8. 在5个未见MetaWorld任务上的成功率曲线。Diffusion Reward能够直接泛化到未见任务并生成合理的奖励,显著超过其他基线方法。结果为4次运行的平均值,阴影区域表示标准误差。
5.4 真实机器人评估
奖励分析
我们选择一个真实机器人任务,目标是从桌上拾起一个碗。为了训练和测试我们的奖励模型,我们使用Allegro手、Franka机械臂和RealSense,由人类操作,收集了20个真实机器人视频(10个专家视频和10个随机视频)。如图6所示的可视化结果表明,Diffusion Reward能够适当地为专家视频分配相对较高的奖励,而随机视频则获得较低的奖励,这表明我们的方法在真实世界机器人操作任务中的潜力。更多试验的曲线和与基线相关的奖励分析见附录。
离线强化学习性能
我们在按钮按压任务中使用并行夹爪作为末端执行器进行离线强化学习评估。专家视频用于训练每个奖励模型以供后续离线强化学习训练(更多细节见附录)。我们的结果表明,Diffusion Reward可以生成更具泛化能力的策略,用于按压未见位置的按钮,如图9所示。此外,图10定量展示了Diffusion Reward在现实世界适用性方面显著优于VIPER。
图6. 真实机器人视频的奖励曲线。我们的方法为专家视频分配了比随机视频更高的奖励。
图9. 真实世界任务的泛化能力。使用我们方法训练的策略能够成功按压未见位置的按钮。
图10. 真实世界离线强化学习的成功率。在按钮按压任务中,给定不同的离线数据,我们的方法显著优于基线方法。
5.5 消融研究
如图11所示,我们在前面的实验中对提出框架的关键设计选择进行了消融研究,旨在揭示我们方法定量性能的更多洞见。以下是更详细的分析:
基于扩散模型的条件熵
由于我们的方法已表现出优于VIPER的性能,我们进一步研究奖励类型与视频预测模型的联合影响。具体而言,我们系统地评估了条件熵和对数似然作为奖励信号的所有可能组合,每种组合分别搭配基于扩散或基于Transformer的视频预测模型。注意,所有模型均使用相同的矢量量化编码器,以确保变化仅归因于所选的视频预测模型和奖励类型。
结果与图2中的观察一致,表明了以下两个结论:
- 视频扩散模型在复杂任务中更擅长捕捉专家视频的复杂分布,从而生成更具信息性的奖励;
- 使用条件熵作为奖励在强化学习探索中比使用对数似然更有效,部分原因在于其对奖励预训练阶段未见轨迹的奖励推断具有更强的泛化能力。
图11. 消融研究。Diffusion Reward的不同消融版本的成功率曲线,聚合自Adroit中的Door和Hammer任务。(a) 我们测试生成模型和奖励的组合,展示了使用扩散模型估计条件熵的优势。(b) 我们对去噪步骤数量的选择进行了消融研究。© 我们展示了反向过程中Diffusion Reward固有的随机性有助于强化学习的探索。(d) 我们对奖励系数的选择进行了消融研究。(e) 我们测试了条件帧数量的影响。结果为3个随机种子的平均值,阴影区域表示标准误差。红色表示我们的默认设置。(在线彩色图)
去噪步骤
反向过程中的时间步数量决定了生成帧的质量和多样性 [42]。本研究通过逐步将去噪步骤的数量从2增加到20,探讨其对奖励及后续强化学习性能的影响。结果表明,中间选择(大约10步)表现最佳。这表明,中间选择在生成质量和多样性之间实现了良好的平衡,从而为强化学习探索生成了有效的奖励。此外,我们观察到,随着去噪步骤的增加,奖励推断速度下降,这表明在需要更多去噪步骤的情况下,采用先进技术加速扩散过程具有优势。
扩散过程中的采样噪声
我们假设扩散过程中的随机性可以加速强化学习的探索,这类似于最大熵强化学习的随机特性 [59]。为了验证这一点,我们设计了Diffusion Reward的一个变体,其中在反向过程中故意将采样噪声设为0,从而保证奖励在相同历史观测下是确定性的。结果显示,使用确定性奖励时性能明显下降。值得注意的是,这种性能下降与将扩散模型与对数似然奖励结合时观察到的结果一致,在这种情况下,奖励同样是确定性的。因此,我们的研究结果表明,Diffusion Reward在反向过程中的固有随机性确实有助于强化学习的探索。
奖励系数
α
\alpha
α
奖励系数
α
\alpha
α决定了条件熵奖励与寻求新奇性奖励之间的相对重要性。我们通过逐步将
α
\alpha
α的值从1减少到0,研究了该参数的影响。结果表明,
α
\alpha
α在0.95左右时性能最佳,而过大的值(类似于仅使用RND)和过小的值(类似于不使用RND)都会导致显著的性能下降。这表明,Diffusion Reward的主导地位可能仍会导致陷入局部最优,而我们提出的奖励有效地帮助强化学习智能体缩小寻求新奇性奖励的广泛探索空间。
上下文长度
历史帧的数量决定了视频扩散过程中编码的时间信息的范围,从而影响视频扩散的生成过程和由此产生的奖励推断。为了研究其对下游强化学习的影响,我们测试了不同的上下文长度选择。结果表明,选择1或2个历史帧即可生成高度有效的奖励,这归因于扩散模型固有的强大生成能力。有趣的是,当上下文长度扩展到4或8帧时,性能略有下降。这种现象可能归因于对专家轨迹的潜在过拟合,导致推断的奖励在未见轨迹上的泛化表现次优。
6 结论
在本研究中,我们提出了Diffusion Reward,这是一种新颖的框架,通过预训练的条件视频扩散模型从强化学习任务中提取稠密奖励。我们首先使用专家视频预训练视频扩散模型,观察到预测分布的熵能够很好地区分专家级轨迹和非专家级轨迹。因此,我们将其标准化熵与探索奖励和稀疏环境奖励结合,作为信息丰富的奖励信号。我们在来自MetaWorld和Adroit的10个视觉机器人操作任务中评估了Diffusion Reward,观察到在两个领域内显著的性能提升。此外,我们展示了在真实机器人任务中生成合理奖励的结果以及相关的离线强化学习的优异表现,突出了我们方法在现实世界中的适用性。进一步,我们证明了预训练奖励能够有效地指导强化学习成功解决未见任务,并在很大程度上超越了基线方法。这突显了大规模预训练扩散模型在奖励泛化中的潜力。
局限性与未来工作
未来的工作将利用来自更广泛数据集的更大规模扩散模型,以解决多样化的仿真和真实世界任务。同时,将探索整合其他模态(例如语言)的可能性,以增强Diffusion Reward的泛化能力。此外,还可以改进基于扩散的奖励本身,包括平衡熵奖励与探索奖励的策略、条件熵的估计以及加速奖励推断速度,这些都有望带来更好的结果。
致谢
本工作得到了中国国家重点研发计划(2022ZD0161700)的支持。感谢所有审稿人提出的深刻意见。