哥伦比亚大学突破性的方法- Diffusion Policy：利用Action Diffusion进行视觉运动策略学习

Diffusion Policy：利用Action Diffusion进行视觉运动策略学习

1. 论文概述

这篇论文提出了一种名为Diffusion Policy的新方法，通过将机器人的视觉运动策略表示为条件扩散去噪过程来生成机器人行为。作者在15个不同的机器人操作基准测试任务上对Diffusion Policy进行了评估，发现它始终优于现有的机器人学习方法，平均提升了46.9%。

Diffusion Policy通过学习action分布的score function梯度，在推理过程中通过一系列随机Langevin动力学步骤迭代优化这个梯度场。研究发现，当用于机器人策略时，扩散公式产生了强大的优势，包括优雅地处理多模态action分布、适用于高维action空间以及表现出令人印象深刻的训练稳定性。

为了在物理机器人上充分发挥扩散模型在视觉运动策略学习中的潜力，本文提出了一些关键的技术贡献，包括结合时间衰减控制(receding horizon control)、视觉条件和时间序列扩散transformer。

2. 方法原理

2.1 将DDPM用于视觉运动策略学习

作者通过以下两个主要修改，将去噪扩散概率模型(DDPM)应用于机器人视觉运动策略学习：

1. 将输出x改为表示机器人动作。
2. 使去噪过程以输入观测O_t为条件。

在具体实现上，在时间步t，策略将最新的T_o步观测数据O_t作为输入，并输出T_a步动作A_t。作者将T_o定义为观测视界，T_p定义为动作预测视界，T_a定义为动作执行视界。这鼓励了时间上的动作一致性，同时保持响应性。

2.2 视觉编码器

视觉编码器将原始图像序列映射到潜在嵌入O_t中，并与扩散策略端到端训练。不同的相机视图使用独立的编码器，每个时间步的图像独立编码，然后串联形成O_t。

2.3 加速实时控制的推理

为了实现实时闭环控制，需要快速的推理速度。作者采用去噪扩散隐式模型(DDIM)方法，将训练和推理中的去噪迭代次数解耦，从而允许算法在推理时使用更少的迭代来加速。

3. Diffusion Policy的特点

• 表达多模态action分布：通过学习action score函数的梯度并在该梯度场上执行随机Langevin动力学采样，Diffusion Policy可以表达任意可归一化的分布，包括多模态action分布。

• 高维输出空间：扩散模型已经在高维输出空间中表现出色。这使得策略能够联合推断未来的一系列动作，而不是单步动作，这对于鼓励时间上的动作一致性和避免近视规划至关重要。

• 训练稳定：EBM训练通常需要负采样来估计难以处理的归一化常数，而这会导致训练不稳定。Diffusion Policy通过学习能量函数的梯度来绕过这一要求，从而在保持分布表达能力的同时实现稳定训练。

• 与位置控制的协同作用：Diffusion Policy与位置控制action空间的结合始终优于速度控制。位置控制比速度控制更不容易受到复合误差的影响，因此更适合动作序列预测。

4. 性能评估

作者在15个任务上系统地评估了Diffusion Policy，这些任务来自4个不同的基准测试。评估包括模拟和真实环境、2DoF到6DoF的动作、单任务和多任务基准、完全执行器和欠驱动系统，以及刚性和流体物体。

实证结果表明，在所有测试的基准上，Diffusion Policy的性能始终优于先前的最佳方法，平均成功率提高了46.9%。

此外，作者还进行了详细的分析，仔细研究了所提出算法的特点以及关键设计决策的影响。

论文解读

1. 引言

在机器人领域,策略学习从演示(Policy Learning from Demonstration)可以看作是从观察到动作的监督回归任务。但是,预测机器人动作的独特性质,如多模态分布的存在、序列相关性和高精度要求,使得这个任务与其他监督学习问题相比更具挑战性。

本文提出了一种名为Diffusion Policy的新型机器人视觉运动策略,通过在机器人动作空间上的"条件去噪扩散过程"来生成行为。在这个公式中,策略推断动作得分梯度,以视觉观察为条件,进行K次去噪迭代。这使得机器人策略继承了扩散模型的几个关键性质,从而显著提高了性能。

2. 扩散策略公式

2.1 去噪扩散概率模型(DDPM)

DDPM是一类生成模型,其中输出生成被建模为去噪过程,通常称为随机Langevin动力学。从高斯噪声x_K采样开始,DDPM执行K次去噪迭代,产生一系列具有递减噪声水平的中间动作,直到形成所需的无噪声输出x_0。这个过程可以用以下方程来表示:

def diffusion_step(x, k, theta, alpha, sigma):
    eps_theta = eps_theta_fn(x, k)  # 噪声预测网络
    x_prev = alpha * (x - (1 - alpha) / torch.sqrt(1 - alpha_bar[k]) * eps_theta) + torch.sqrt(1 - alpha) * torch.randn_like(x)
    return x_prev

其中,eps_theta是参数为θ的噪声预测网络,alpha, sigma是噪声调度的函数。

2.2 用于视觉运动策略学习的DDPM

为了将DDPM用于机器人视觉运动策略学习,作者做了两个主要修改:

1. 将输出x改为表示机器人动作。
2. 使去噪过程以输入观测O_t为条件。

具体来说,在时间步t,策略将最新的T_o步观测数据O_t作为输入,并预测T_p步动作,其中T_a步动作在机器人上执行,而无需重新规划。这鼓励了时间上的动作一致性,同时保持响应性。修改后的方程如下:

def diffusion_step_conditioned(x, k, theta, alpha, sigma, o):
    eps_theta = eps_theta_fn(o, x, k)  # 以观测为条件的噪声预测网络
    x_prev = alpha * (x - (1 - alpha) / torch.sqrt(1 - alpha_bar[k]) * eps_theta) + torch.sqrt(1 - alpha) * torch.randn_like(x)
    return x_prev

相应地,训练损失也修改为:

def loss_fn(x0, o, k, theta):
    eps_k = torch.randn_like(x0)
    xk = torch.sqrt(alpha_bar[k]) * x0 + torch.sqrt(1 - alpha_bar[k]) * eps_k
    eps_theta = eps_theta_fn(o, xk, k)
    return torch.mean((eps_theta - eps_k)**2)

3. 关键设计决策

3.1 网络架构选择

噪声预测网络eps_theta的选择是第一个设计决策。本文检验了两种常见的网络架构类型:卷积神经网络(CNN)和Transformer,并比较了它们的性能和训练特点。

• 基于CNN的Diffusion Policy
  采用了Janner等人的1D时间CNN,并做了一些修改,如只建模条件分布p(A_t|O_t),移除了基于inpainting的目标状态条件等。

• 基于Transformer的Diffusion Policy
  为了减少CNN模型中的过度平滑效应,作者引入了一种新颖的基于Transformer的DDPM。带噪声的动作A_k_t作为Transformer解码器块的输入令牌传入,扩散迭代k的正弦嵌入作为第一个令牌。

3.2 视觉编码器

视觉编码器将原始图像序列映射到潜在嵌入O_t中,并与扩散策略端到端训练。不同的相机视图使用单独的编码器,并独立编码每个时间步的图像,然后串联形成O_t。作者使用标准的ResNet-18作为编码器,并进行了一些修改,如用空间softmax池化替换全局平均池化,用GroupNorm替换BatchNorm等。

3.3 噪声调度

噪声调度由函数σ、α、γ和加性高斯噪声ϵ_k定义,控制着扩散策略捕捉动作信号的高频和低频特性。在实验中,作者发现iDDPM中提出的Square Cosine Schedule效果最好。

3.4 加速实时控制的推理

为了实现实时闭环控制,推理速度非常关键。去噪扩散隐式模型(DDIM)方法将训练和推理中的去噪迭代次数解耦,允许算法在推理时使用更少的迭代来加速。在实际实验中,使用100次训练迭代和10次推理迭代的DDIM可以在Nvidia 3080 GPU上实现0.1秒的推理延迟。

4. Diffusion Policy的有趣特性

• 表达多模态动作分布
  通过学习动作得分函数的梯度并在该梯度场上执行随机Langevin动力学采样,Diffusion Policy可以自然而精确地表达多模态分布。

• 与位置控制的协同作用
  Diffusion Policy与位置控制动作空间的结合始终优于速度控制。消融研究表明,选择位置控制作为扩散策略的动作空间显著优于速度控制。

• 动作视界的权衡
  动作视界大于1有助于策略预测一致的动作并补偿演示中的闲置部分,但视界过长会因反应时间慢而降低性能。实验确认了这种权衡,并发现8步的动作视界对于测试的大多数任务来说是最优的。

• 对延迟的稳健性
  通过采用衰减视界位置控制来预测未来的一系列动作,Diffusion Policy有助于解决图像处理、策略推理和网络延迟引起的延迟问题。

• 训练稳定性
  Diffusion Policy的最优超参数在不同任务之间大多一致。相比之下,隐式行为克隆(IBC)容易出现训练不稳定。这一特性得益于Diffusion Policy通过学习能量函数的梯度来绕过估计归一化常数,从而实现稳定训练。

5. 性能评估

作者在4个基准测试的15个任务上系统地评估了Diffusion Policy。这套评估包括模拟和真实环境、2DoF到6DoF动作空间、单任务和多任务基准、完全执行器和欠驱动系统,以及刚性和流体对象。实证结果表明,Diffusion Policy在所有测试的基准上始终优于现有方法,平均成功率提高了46.9%。

此外,作者还进行了详细的分析,仔细研究了所提出算法的特点以及关键设计决策的影响。如多模态行为的案例研究、动作空间选择的消融研究、动作视界权衡的实验确认、网络延迟下的稳健性测试等。

6. 代码示例

下面是一些PyTorch的示例代码片段,帮助读者更好地理解Diffusion Policy的关键组件:

条件去噪扩散步骤

def diffusion_step_conditioned(x, k, theta, alpha, sigma, o):
    eps_theta = eps_theta_fn(o, x, k)  # 以观测为条件的噪声预测网络
    x_prev = alpha * (x - (1 - alpha) / torch.sqrt(1 - alpha_bar[k]) * eps_theta) + torch.sqrt(1 - alpha) * torch.randn_like(x)
    return x_prev

训练损失

def loss_fn(x0, o, k, theta):
    eps_k = torch.randn_like(x0)
    xk = torch.sqrt(alpha_bar[k]) * x0 + torch.sqrt(1 - alpha_bar[k]) * eps_k
    eps_theta = eps_theta_fn(o, xk, k)
    return torch.mean((eps_theta - eps_k)**2)

视觉编码器

class VisualEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.resnet = models.resnet18(pretrained=False)
        self.resnet.fc = nn.Identity()
        self.resnet.avgpool = nn.Sequential(
            nn.Conv2d(512, 512, 3, stride=1, padding=1),
            nn.GroupNorm(32, 512),
            nn.SiLU(),
            SpatialSoftmax(),
        )

    def forward(self, x):
        return self.resnet(x)

噪声调度

def cosine_beta_schedule(timesteps, s=0.008):
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0.0001, 0.9999)

7. 总结

本论文工作提出了一种名为Diffusion Policy的新型视觉运动策略学习方法，它利用了扩散模型强大的生成建模能力。通过引入时间衰减控制、视觉条件和时间序列扩散transformer等关键技术贡献，Diffusion Policy在各种复杂的机器人操作任务上实现了最佳性能，展示了其在机器人领域的广阔应用前景。