Txing

欢迎来到 | 伽蓝之堂

0%

Imitation Learning | Diffusion Policy: Visuomotor Policy Learning via Action Diffusion

Posted on In
Symbols count in article: 11k Reading time ≈ 10 mins.
Diffusion Policy机器人末端将轨迹生成视为一个条件去噪扩散过程，比当前的sota算法实现了46.9%的性能提升，达到了机器人操作的最好性能。

  • Diffusion Policy机器人末端将轨迹生成视为一个条件去噪扩散过程，比当前的sota算法实现了46.9%的性能提升，达到了机器人操作的最好性能

  • Diffusion Policy学习的是行为分布评分函数的梯度，然后通过一系列的stochastic Langevin dynamics steps，实现优雅地处理多模态动作分布，适用于高维动作空间，并表现出令人印象深刻的训练稳定性。

  • 项目网站：https://diffusion-policy.cs.columbia.edu/

  • 轨迹规划示例

1 INTRODUCTION

  • 当前策略学习的问题：
    • 行为的多模态分布
    • 序列的正确性
    • 高精确度
  • 此前的工作尝试使用不同的行为表征方式，下图展示了不同的行为形式：
    • 混合高斯表征：量化行为表征，形成显式分布；
    • 策略表征：从显式到隐式，以更好捕捉多模态行为分布；
  • Policy Representations
  • 在本文中，设计了一个 行为空间上的conditional denoising diffusion process，来学习视觉运动策略，取名为 Diffusion Policy
    • Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239, 2020
  • 该策略不是直接输出行为，而是推断action-score gradient，经过K次去噪迭代；该方式允许机器人策略继承 diffusion 模型的几个关键属性，从而显著提高性能：
    • Expressing multimodal action distributions
      • 学习行为分数函数的梯度以及在梯度场上执行 Stochastic Langevin Dynamics sampling，diffusion policy可以表达任意的正则化分布，包括多模态的行为分布；
        • Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019
        • Mcmc using hamiltonian dynamics. Handbook of markov chain monte carlo, 2011
    • High-dimensional output space
      • 该属性允许策略联合推断未来动作的序列而不是单步动作，这对于时序动作一致性和避免短视的规划至关重要；
    • Stable training
      • 训练基于能量的策略通常需要负采样（negative samples，梯度只更新网络的一小部分参数）来估计一个棘手的归一化常数，这会导致训练不稳定
        • Implicit behavioral cloning. In 5th Annual Conference on Robot Learning, 2021.
  • 主要贡献：
    1. Closed-loop action sequences
      1. diffusion policy预测高维行为序列，结合滚动时域控制实现稳健执行，该设计允许策略以闭环形式连续re-plan机器人的行为，并且保存行为一致性；实现长期规划和实时响应的平衡；（开环规划，执行一部分后重规划，近似闭环效果）
    2. Visual conditioning
      1. 设计了一个 vision-conditioned diffusion policy ，视觉观测作为条件而非联合数据分布的一部分；因此策略只进行1次视觉特征提取，并且图像不参与去噪的迭代过程，提升了计算效率；
    3. Time-series diffusion transformer
      1. 提出了一个 transformer-based diffusion network 来最小化CNN-based模型中常常出现的 over-smoothing effects，并在需要速度控制和高频变化行为的任务中取得sota的效果
      2. over-smoothing effects
        • 定义：Over-smoothing 是指随着图神经网络层数的增加，节点特征变得越来越相似，最终在高层次上收敛到一个相似或相同的状态。这导致不同节点之间的特征区分度降低，使得GNN难以捕捉到复杂的图结构。
        • 原因：这种现象通常是由于多层GNN在特征传播过程中重复应用相同的平滑操作导致的。每一层都将一个节点的特征与其邻居的特征混合，多层之后，所有节点的特征趋于一致。
        • 后果：这会降低GNN在节点分类和其他任务中的性能，因为不同节点的区别性减少了。
  • 后续实验在4个基准的12个任务上进行了评估；任务包括：
    • 仿真和实际系统；simulated and real-world environments
    • 不同自由度机器人；2DoF to 6DoF actions
    • 单任务和多任务基准；single- and multi-task benchmarks
    • 全驱动和欠驱动；fully- and underactuated systems
    • 刚性与流动性物体；rigid and fluid objects
    • 示教数据由单人和多人收集；demonstration data collected by single and multiple users.
  • 4个 benchmarks:
    1. Implicit behavioral cloning. In 5th Annual Conference on Robot Learning, 2021.
    2. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint arXiv:1910.11956, 2019.
    3. What matters in learning from offline human demonstrations for robot manipulation. In 5th Annual Conference on Robot Learning, 2021.
    4. Behavior transformers: Cloning modes with one stone. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  • 结果发现所有基准任务重都有一致的性能提升，平均提升了46.9%

### 2 DIFFUSION POLICY FORMULATION

  • 将机器人的视觉运动策略形式化为去噪扩散概率模型（Denoising Diffusion Probabilistic Models (DDPMs)）

A. Denoising Diffusion Probabilistic Models

  • 去噪过程被称为 Stochastic Langevin Dynamics (随机郎之万动力学)

    • Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 681–688, 2011.

  • 郎之万动力学仅使用评分函数 就可以从概率密度 中生成样本： 其中， 是一个固定的步长，初始值 属于一个先验分布；；当 , 时，在一些regularity conditions下，就等同于来自的样本

    因此，只要使得评分函数逼近，就可以直接用郎之万动力学获取 分布中的样本

  • Realworld Benchmarks

  • 首先从Gaussian noise中采样出 ，DDPM迭代进行 次去噪过程从而产生一系列噪声干扰逐渐减小的中间行为，分别表示为 直到期望的无噪声输出 生成，迭代过程如下： 其中，表示噪声预测网络，是网络参数；表示在每次迭代中添加的高斯噪声；

  • 上述公式也可以被截石位单个噪声梯度下降步骤： 其中，噪声预测网络用于预测梯度场 ，表示学习率；

B. DDPM Training

  • 训练时，首先从数据集中随机抽取未修改的样本 ，对每个样本，随机选择一个去噪迭代次数 ；在迭代的第轮，再采样一个具有适当方差（with appropriate variance）的随机噪声

  • 之后噪声预测网络预测样本加噪后的噪声，损失函数定义如下，

    • Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239, 2020. 1

  • 同时还需要最小化数据分布 和从DDPM获得的分布 的KL散度的变分下界（variational lower bound）

C. Diffusion for Visuomotor Policy Learning

  • Diffusion Policy Overview
  • DDPM主要用于图像的生成，我们需要将其变为对机器人行为的生成，需要执行的修改包括：
    1. 将输出x变为机器人的行为；
    2. 需要将去噪过程依据输入观测量进行条件化
Closed-loop action-sequence prediction

  • 一个有效的行为应该具有时域的一致性；在长期规划上具有平滑性；能对对意料之外的障碍物做出反应；

  • 本文使用一个扩散模型+滚动时域控制完成了上述目标，实现了稳健的动作执行；

  • 在第步的时候，将最近步的观测数据作为输入，用于预测步的行为，其中，只有步的行为会被执行；这种方式可以有效保持行为迭代过程中的一致性

Visual observation conditioning

  • 使用一个DDPM来近似条件分布 ，而不是用于近似联合分布 ，这样可以允许模型在观测数据的条件下预测未来行为，并且没有推断未来状态的负担

    • 这篇文章是近似的联合分布：Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, 2022

  • 将2.A中的公式修改为：

  • 相应的训练损失也变为：

  • 将观测值从去噪过程的输出中排出可以有效提高计算速度，便于后续的实时计算和控制；

3. KEY DESIGN DECISION

A. Network Architecture Option

  • noise prediction network 的实现考虑了两种网络结构：基于CNN的和基于Transformer的；

    CNN-based Diffusion Policy

    • 在下文基础上，采用了1D的时序CNN，并做了一些改进；
      • Planning with diffusion for flexible behavior synthesis. Proceedings of the 39th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 17–23 Jul 2022.
    • 改进1：仅通过使用特征线性调制（Feature-wise Linear Modulatio, FiLM）以及去噪迭代，对观测特征 上的动作生成过程进行条件化处理，实现对条件分布 进行建模
    • 改进2：只预测行为轨迹，而不是串联的观察-行为轨迹
    • 改进3：由于与我们利用滚动预测窗口的框架不兼容，我们删除了基于修复的目标状态条件（？）
    • 基于CNN的骨架在行为序列快速变化的case中表现不好，这可能是因为时间卷积的归纳偏好，使得模型偏好低频的信号；

    Time-series diffusion transformer

    • 为了降低CNN模型中over-smoothing效应的影响，这里采用了一个 novel transformer-based DDPM，参考了minGPT的结构来进行行为预测
      • Behavior transformers: Cloning k modes with one stone. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022
    • 带噪声的行为 作为输入token传入Transformer decoder blocks；第k次扩散迭代的正弦embedding作为第一个token输入；
    • 观测量通过共享的MLP变换为观测值嵌入序列，然后作为输入特征传递到 Transformer decoder 堆栈中
    • 梯度 通过decoder堆栈的每个相应输出token预测
    • 多数的高性能模型使用的都是Transformer backbone，特别是在任务复杂，动作变化频率高的时候；但是本研究对比发现Transformer对超参数更敏感，一般的解决方案是改进Transformer的训练技术或者直接增加数据量；

    Recommendations

    • In general, we recommend starting with the CNN-based diffusion policy implementation as the first attempt at a new task. If performance is low due to task complexity or high-rate action changes, then the Time-series Diffusion Transformer formulation can be used to potentially improve performance at the cost of additional tuning.

B. Visual Encoder

  • 视觉编码器将原始图像序列映射到一个潜在embedding ，并且可以通过端到端的方式训练扩散策略

  • 不同的摄像机视图使用不同的编码器，每个时间步中的图像被独立编码，然后连接起来形成

  • 使用标准的 ResNet18 (without pretraining) 作为编码器，并进行以下修改：

    1. 用空间softmax池化替换全局平均池化，以保持空间信息
      • What matters in learning from offline human demonstrations for robot manipulation. In 5th Annual Conference on Robot Learning, 2021.
    2. 用GroupNorm替换BatchNorm以进行稳定训练， Exponential Moving Average 和归一化层一起使用时非常关键
      • Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3–19, 2018.
      • Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729–9738, 2020

C. Noise Schedule

  • 噪声的排期通过 以及高斯噪声 组成关于 的函数来定义。
  • noise schedule 控制扩散策略捕捉动作信号的高频和低频特征的程度。在控制任务中，根据经验发现 iDDPM 中提出的Square Cosine Schedule 最适合本文的任务。
    • Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pages 8162–8171. PMLR, 2021

D. Accelerating Inference for Real-time Control

  • 使用 Denoising Diffusion Implicit Models (DDIM) approach 解耦了在训练和推断中的去噪迭代的次数；因此能够在推断中使用更少的迭代次数，加速推断过程；
  • 实验中，using DDIM with 100 training iterations and 10 inference iterations enables 0.1s inference latency on an Nvidia 3080 GPU.

4. INTRIGUING PROPERTIES OF DIFFUSION POLICY

A. Model Multi-Modal Action Distributions

  • 行为的多模态在很多behavior cloning的文章中被讨论过，diffusion policy在这个问题上具有优势
    • Implicit behavioral cloning. In 5th Annual Conference on Robot Learning, 2021
  • diffusion policy处理多模态行为的能力来自两个方面：
    • stochastic initialization：Stochastic Langevin Dynamics过程中，初始样本 来自标准的Gaussian分布，有助于为最终动作预测指定不同的可能的收敛盆地；
    • stochastic sampling procedure：之后的随机优化过程中，又在多轮迭代中添加了高斯扰动，最终促成了独立的行为采样；
  • Multimodal behavior
    可见diffusion policy多余多模态行为没有特定的偏好，而其他的算法要么偏好特定的行为模态，要么出现了模态混淆（行为没有一致性）

B. Synergy with Position Control

  • （diffusion policy + 位置控制）效果比（diffusion policy + 速度控制）效果好，可能的原因有两个：
    1. 行为模态在位置控制的情况下会更加显著，但由于diffusion policy能够比现有方法更好地表达行为的多模态，所以速度控制对算法的影响在diffusion policy的情况下会轻一些；
    2. 位置控制受到的符合误差效应（compounding error effects）的影响更小，因此更加适合行为序列预测；
  • Velocity v.s. Position Control

C. Benefits of Action-Sequence Prediction

  • 序列预测在多数的策略学习工作中都是尽量避免的，因为存在在高维输出空间中高效采样的困难；
  • 相比之下，DDPM可以很好地扩展输出尺寸，而不会牺牲模型的表现力，这在许多图像生成应用程序中得到了验证。利用这种能力，扩散策略以高维动作序列的形式表示动作，这自然解决了以下问题:
    • Temporal action consistency：原本连续的动作在每一步独立生成，导致行为在多个模态之间来回跳转；
    • Robustness to idle actions：示教暂停时会出现无意义的动作，并导致一系列相同的位置动作或接近零速度动作；这在遥操作中很常见；如果每一步独立生成，模型会很容易过拟合这些无意义的动作或停顿；

D. Training Stability

  • 使用Energy-Based Model (EBM)表征行为分布： 其中， 是比较棘手的正则化常数；

  • 为了训练EBM隐式策略，使用了一个 InfoNCE-style loss： 其中，负样本集合 用于估计归一化常数 ；实际计算中，不精确的负采样会造成EBM的训练不稳定

    • Improved contrastive divergence training of energy based models. arXiv preprint arXiv:2012.01316, 2020.
    • Conditional energy-based models for implicit policies: The gap between theory and practice. arXiv preprint arXiv:2207.05824, 2022.

  • 扩散策略和DDPMs通过对同一行动分布的得分函数进行建模，回避了估计的问题 其中，当，噪声预测网络 j近似于负评分函数 ，并且独立于归一化常数

    • Compositional visual generation with composable diffusion models. arXiv preprint arXiv:2206.01714, 2022

  • 因此，无论是推断还是训练过程中，diffusion policy都不用估计 ，使得diffusion policy的训练更加稳定；

5. EVALUATION

A. Simulation Environments and datasets

  1. Robomimic
    • What matters in learning from offline human demonstrations for robot manipulation. In 5th Annual Conference on Robot Learning, 2021
  2. Push-T
  3. Multimodal Block Pushing
  4. Franka Kitchen

B. Evaluation Methodology

C. Key Findings

  • Diffusion Policy can express short-horizon multimodality

  • Diffusion Policy can express long-horizon multimodality

  • Diffusion Policy can better leverage position control

  • The tradeoff in action horizon

  • Robustness against latency

  • Diffusion Policy is stable to train

Multi-Stage Tasks (State Observation)
Realworld Push-T Experiment
Realworld Push-T Comparisons
Robustness Test for Diffusion Policy
6DoF Mug Flipping Task

总结

diffusion policy的确是个很具有创新性的点子，在性能上的提升也非常多。但我比较好奇的点在于观测量的视角设置有什么要求？生成的定位误差大时怎么执行？行为怎么定义？是关节角度还是末端位置坐标，如果是坐标，又是在什么坐标系下？具体内容还需要结合代码和附录来分析（也可以参考这篇文章Planning with Diffusion for Flexible Behavior Synthesis）。总之，这篇文章创新点很好，机器人操作领域的发展很长一段时间都是滞后于cv和nlp的，本文将最新的技术思想迁移了过来，有可能推动机器人操作迈上新台阶。