扩散策略

[项目页面] [论文] [数据] [Colab（状态）] [Colab（视觉）]

程驰¹, 冯思源², 杜一伦³, 徐振佳¹, 埃里克·考斯尼奥², 本杰明·布楚菲尔德², 宋书然¹

¹哥伦比亚大学， ²丰田研究所， ³麻省理工学院

尝试一下！

使用自包含的Google Colab笔记本是体验Diffusion Policy最简单的方法。我们为 基于状态的环境 和 基于视觉的环境 提供了单独的笔记本。

查看实验日志！

对于生成paper中表I、II和IV的每个实验，我们提供：

1. 包含重现实验所需所有参数的config.yaml文件。
2. 每个训练步骤详细训练/评估logs.json.txt。
3. 每次运行的最佳epoch=*-test_mean_score=*.ckpt和最后一个latest.ckpt时期的检查点。

实验日志以嵌套目录的形式托管在我们的网站上： https://diffusion-policy.cs.columbia.edu/data/experiments/<image|low_dim>/<task>/<method>/

在每个实验目录中，您可能会找到：

.
├── config.yaml
├── metrics
│   └── logs.json.txt
├── train_0
│   ├── checkpoints
│   │   ├── epoch=0300-test_mean_score=1.000.ckpt
│   │   └── latest.ckpt
│   └── logs.json.txt
├── train_1
│   ├── checkpoints
│   │   ├── epoch=0250-test_mean_score=1.000.ckpt
│   │   └── latest.ckpt
│   └── logs.json.txt
└── train_2
    ├── checkpoints
    │   ├── epoch=0250-test_mean_score=1.000.ckpt
    │   └── latest.ckpt
    └── logs.json.txt

metrics/logs.json.txt 文件聚合了使用multirun_metrics.py每50个周期从所有3个训练运行中收集的评估指标。论文中报告的数字对应于max和k_min_train_loss聚合键。

要下载子目录中的所有文件，请使用：

$ wget --recursive --no-parent --no-host-directories --relative --reject="index.html*" https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/square_ph/diffusion_policy_cnn/

安装

模拟仿真

为了重现我们的模拟基准测试结果，在具有Nvidia GPU的Linux机器上安装我们的Conda环境。在Ubuntu 20.04上，为mujoco安装以下apt包：

$ sudo apt install -y libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf

推荐使用Mambaforge代替标准的Anaconda分发以加快安装速度：

$ mamba env create -f conda_environment.yaml

但也可以使用Conda：

$ conda env create -f conda_environment.yaml

conda_environment_macos.yaml文件仅用于MacOS上的开发，并不完全支持基准测试。

实物机器人

硬件（针对Push-T）：

• 1个UR5-CB3 或 UR5e （需要RTDE接口）
• 2个RealSense D415
• 1个3Dconnexion SpaceMouse（用于远程操作）
• 1个[Millibar Robotics手动工具交换器](https://www.millibar.com/manual-tool-changer/）（只需要机器人侧）
• 1个3D打印的末端执行器
• 1个3D打印T块
• 实感相机所需的USB-C线缆和螺丝

软件：

• Ubuntu 20.04.3（已测试）
• Mujoco依赖项：sudo apt install libosmesa6-dev libgl1-mesa-glx libglfw3 patchelf
• RealSense SDK
• Spacemouse依赖项：sudo apt install libspnav-dev spacenavd; sudo systemctl start spacenavd
• Conda环境 mamba env create -f conda_environment_real.yaml

重现模拟基准测试结果

下载训练数据

在仓库根目录下创建数据子目录：

[diffusion_policy]$ mkdir data && cd data

从https://diffusion-policy.cs.columbia.edu/data/training/ 下载相应的zip文件：

[data]$ wget https://diffusion-policy.cs.columbia.edu/data/training/pusht.zip

解压训练数据：

[data]$ unzip pusht.zip && rm -f pusht.zip && cd ..

获取对应实验的配置文件：

[diffusion_policy]$ wget -O image_pusht_diffusion_policy_cnn.yaml https://diffusion-policy.cs.columbia.edu/data/experiments/image/pusht/diffusion_policy_cnn/config.yaml

运行单个种子

激活Conda环境，如果还没有的话，登录到wandb。

[diffusion_policy]$ conda activate robodiff
(robodiff)[diffusion_policy]$ wandb login

在GPU 0上用种子42启动训练。

(robodiff)[diffusion_policy]$ python train.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml training.seed=42 training.device=cuda:0 hydra.run.dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}'

这将在data/outputs/yyyy.mm.dd/hh.mm.ss_<方法名>_<任务名>格式的目录中创建一个文件夹，其中会写入配置、日志和检查点。政策每50个周期都会被评估一次，成功率作为test/mean_score记录在wandb上，同时还会为一些回放视频。

(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image -I wandb
data/outputs/2023.03.01/20.02.03_train_diffusion_unet_hybrid_pusht_image
├── checkpoints
│   ├── epoch=0000-test_mean_score=0.134.ckpt
│   └── latest.ckpt
├── .hydra
│   ├── config.yaml
│   ├── hydra.yaml
│   └── overrides.yaml
├── logs.json.txt
├── media
│   ├── 2k5u6wli.mp4
│   ├── 2kvovxms.mp4
│   ├── 2pxd9f6b.mp4
│   ├── 2q5gjt5f.mp4
│   ├── 2sawbf6m.mp4
│   └── 538ubl79.mp4
└── train.log

3 directories, 13 files

以上就是如何设置和运行实验的基本步骤。

多种子运行

启动本地 Ray 集群。对于大规模实验，可能需要设置带有自动扩展功能的 AWS 集群。

(robodiff)[diffusion_policy]$ export CUDA_VISIBLE_DEVICES=0,1,2  # 选择要由 Ray 集群管理的 GPU
(robodiff)[diffusion_policy]$ ray start --head --num-gpus=3

启动一个 Ray 客户端，它将启动 3 个训练工作者（对应 3 个种子）和 1 个指标监控工作者。

(robodiff)[diffusion_policy]$ python ray_train_multirun.py --config-dir=. --config-name=image_pusht_diffusion_policy_cnn.yaml --seeds=42,43,44 --monitor_key=test/mean_score --multi_run.run_dir='data/outputs/${now:%Y.%m.%d}/${now:%H.%M.%S}_${name}_${task_name}' multi_run.wandb_name_base='${now:%Y.%m.%d-%H.%M.%S}_${name}_${task_name}'

除了每个训练工作者单独写入的 WandB 日志外，指标监控工作者还会将所有 3 次训练运行的指标汇总到 WandB 项目 diffusion_policy_metrics 中。本地配置、日志和检查点会按照与我们训练日志相同的目录结构写入 data/outputs/yyyy.mm.dd/hh.mm.ss_<方法名>_<任务名>：

(robodiff)[diffusion_policy]$ tree data/outputs/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image -I 'wandb|media'
数据/输出/2023.03.01/22.13.58_train_diffusion_unet_hybrid_pusht_image
├── config.yaml
├── metrics
│   ├── logs.json.txt
│   ├── metrics.json
│   └── metrics.log
├── train_0
│   ├── checkpoints
│   │   ├── epoch=0000-test_mean_score=0.174.ckpt
│   │   └── latest.ckpt
│   ├── logs.json.txt
│   └── train.log
├── train_1
│   ├── checkpoints
│   │   ├── epoch=0000-test_mean_score=0.131.ckpt
│   │   └── latest.ckpt
│   ├── logs.json.txt
│   └── train.log
└── train_2
    ├── checkpoints
    │   ├── epoch=0000-test_mean_score=0.105.ckpt
    │   └── latest.ckpt
    ├── logs.json.txt
    └── train.log

7 个目录，16 个文件

🆕 预训练检查点评估

从发布的训练日志文件夹下载一个检查点，例如：https://diffusion-policy.cs.columbia.edu/data/experiments/low_dim/pusht/diffusion_policy_cnn/train_0/checkpoints/epoch=0550-test_mean_score=0.969.ckpt。

运行评估脚本：

(robodiff)[diffusion_policy]$ python eval.py --checkpoint data/0550-test_mean_score=0.969.ckpt --output_dir data/pusht_eval_output --device cuda:0

这会产生如下目录结构：

(robodiff)[diffusion_policy]$ tree data/pusht_eval_output
数据/pusht_eval_output
├── eval_log.json
└── 媒体
    ├── 1fxtno84.mp4
    ├── 224l7jqd.mp4
    ├── 2fo4btlf.mp4
    ├── 2in4cn7a.mp4
    ├── 34b3o2qq.mp4
    └── 3p7jqn32.mp4

1 个目录，7 个文件

eval_log.json 包含在训练期间记录到 WandB 的指标：

(robodiff)[diffusion_policy]$ cat data/pusht_eval_output/eval_log.json
{
  "test/mean_score": 0.9150393806777066,
  "test/sim_max_reward_4300000": 1.0,
  "test/sim_max_reward_4300001": 0.9872969750774386,
  ...
  "train/sim_video_1": "data/pusht_eval_output//media/2fo4btlf.mp4"
}

🦾 真实机器人上的演示、训练和评估

确保你的 UR5 机器人正在运行并接受网络接口的命令（随时可以触及紧急停止按钮），RealSense 相机已连接到工作站（使用 realsense-viewer 进行测试），SpaceMouse 已连接并且 spacenavd 服务正在运行（通过 systemctl status spacenavd 验证）。

开始演示收集脚本。按下 "C" 开始录制。使用 SpaceMouse 控制机器人移动。按下 "S" 停止录制。

(robodiff)[diffusion_policy]$ python demo_real_robot.py -o data/demo_pusht_real --robot_ip 192.168.0.204

这会在 data/demo_pusht_real 产生一个与我们的示例真实 Push-T 训练数据集相同结构的演示数据集。

为了训练扩散策略，用配置启动训练：

(robodiff)[diffusion_policy]$ python train.py --config-name=train_diffusion_unet_real_image_workspace --task.dataset_path=data/demo_pusht_real

如果相机设置不同，请编辑 [diffusion_policy/config/task/real_pusht_image.yaml](./diffusion_policy/config/task/real_pusht_image.yaml)。

假设训练已完成，并且你有一个位于 data/outputs/blah/checkpoints/latest.ckpt 的检查点，用以下方式启动评估脚本：

python eval_real_robot.py -i data/outputs/blah/checkpoints/latest.ckpt -o data/eval_pusht_real --robot_ip 192.168.0.204

按 "C" 开始评估（将控制权交给策略）。按 "S" 结束当前回合。

🗺️ 代码库教程

此代码库根据以下需求组织：

1. 实现 N 个任务和 M 种方法仅需 O(N+M) 量级的代码，而不是 O(N*M)。
2. 在保持最大灵活性的同时。

为了满足这些要求，我们：

1. 维护了简单统一的任务和方法接口。
2. 使任务和方法的实现彼此独立。

这些设计决策以任务和方法之间的代码重复为代价。然而，我们认为能够添加/修改任务/方法而不影响其余部分，并能通过线性阅读代码理解任务/方法的优点超过了复制和粘贴的成本 😊。

分割

在任务侧，我们有：

• Dataset：将（第三方）数据集适配到接口中。
• EnvRunner：执行接受接口并生成日志和指标的策略。
• config/task/<task_name>.yaml：包含构建 Dataset 和 EnvRunner 所需的所有信息。
• （可选）Env：一个兼容 gym==0.21.0 的类，封装了任务环境。

在策略侧，我们有：

• Policy：实现了根据接口进行推断以及部分训练过程。
• Workspace：管理方法的训练和评估（交织）生命周期。
• config/<workspace_name>.yaml：包含构建 Policy 和 Workspace 所需的所有信息。

界面

低维度

LowdimPolicy（在./diffusion_policy/policy/base_lowdim_policy.py）接收观察字典：

• "obs": 形状为(B, To, Do)的张量

并预测动作字典：

• "action": 形状为(B, Ta, Da)的张量

LowdimDataset（在./diffusion_policy/dataset/base_dataset.py）返回一个样本字典：

• "obs": 形状为(To, Do)的张量
• "action": 形状为(Ta, Da)的张量

其get_normalizer方法返回一个具有键"obs", "action"的LinearNormalizer（在./diffusion_policy/model/common/normalizer.py）。

策略在GPU上使用LinearNormalizer的副本处理归一化。LinearNormalizer的参数作为策略权重检查点的一部分保存。

图像

ImagePolicy（在./diffusion_policy/policy/base_image_policy.py）接收观察字典：

• "key0": 形状为(B, To, *)的张量
• "key1": 形状如(B, To, H, W, 3)（[0,1]范围的float32类型）

并预测动作字典：

• "action": 形状为(B, Ta, Da)的张量

ImageDataset返回一个样本字典：

• "obs": 字典包含
  • "key0": 形状为(To, *)的张量
  • "key1": 形状为(To, H, W, 3)的张量
• "action": 形状为(Ta, Da)的张量

其get_normalizer方法返回键为"key0", "key1", "action"的LinearNormalizer。

示例

To = 3
Ta = 4
T = 6
|o|o|o|
| | |a|a|a|a|
|o|o|
| |a|a|a|a|a|
| | | | |a|a|

术语对照（论文中的表示）：

• 观察时域：To|n_obs_steps
• 动作时域：Ta|n_action_steps
• 预测时域：T|horizon

经典的（例如MDP）单步观察/动作形式是当To=1和Ta=1的特殊情况。

🛠️ 关键组件

工作区

Workspace对象封装了运行实验所需的所有状态和代码。

• 继承自BaseWorkspace（在./diffusion_policy/workspace/base_workspace.py）。
• 单个由hydra生成的OmegaConf配置对象应包含构建Workspace对象和运行实验所需的所有信息。该配置对应于config/<workspace_name>.yaml加上 Hydra 覆盖。
• run方法包含了实验的完整流程。
• 检查点在Workspace级别发生。所有实现为对象属性的训练状态都通过save_checkpoint方法自动保存。
• 实验的所有其他状态应在run方法中实现为局部变量。

训练入口点是train.py，它使用@hydra.main装饰器。阅读Hydra的官方文档以了解命令行参数和配置覆盖。例如，task=<task_name>参数会替换配置中的task子树，并用config/task/<task_name>.yaml的内容选择要为此实验运行的任务。

数据集

Dataset对象：

• 继承自torch.utils.data.Dataset。
• 根据任务是否有低维或图像观察，返回符合界面的样本。
• 具有一个返回符合界面的LinearNormalizer的方法get_normalizer。

在项目开发过程中，归一化是常见的错误源。有时打印出用于LinearNormalizer每个键的具体scale和bias向量是有帮助的。

我们的大多数Dataset实现都结合了ReplayBuffer和SequenceSampler（在./diffusion_policy/common/sampler.py）来生成样本。正确处理根据To和Ta在每个演示序列开头和结尾的填充对于良好性能至关重要。在实现自己的采样方法之前，请先阅读我们提供的SequenceSampler。

策略

Policy对象：

• 继承自BaseLowdimPolicy或BaseImagePolicy。
• 具有一个方法predict_action，给定观察字典，预测符合界面的动作。
• 具有一个方法set_normalizer，接受一个LinearNormalizer并在策略内部处理观察/动作的归一化。
• （可选）可能具有一个方法compute_loss，接受一批数据并返回要优化的损失。
• （可选）通常每个Policy类对应一个Workspace类，由于不同方法之间的训练和评估过程的差异。

环境执行器

EnvRunner对象抽象了不同任务环境之间的微小差异。

• 具有一个方法run，接受一个策略对象进行评估，并返回包含日志和指标的字典。每个值应与wandb.log兼容。

为了最大化评估速度，我们通常使用对gym.vector.AsyncVectorEnv（在./diffusion_policy/gym_util/async_vector_env.py）的修改版进行环境向量化，使每个个体环境在单独进程中运行（绕过python的全局解释器锁）。

警告：由于在Linux上使用fork启动子进程，对于在初始化期间创建OpenGL上下文的环境（如robosuite），需要特别小心。一旦继承到子进程的内存空间，往往会导致诸如段错误等隐晦的bug。作为解决办法，可以提供一个dummy_env_fn，它构建环境而不初始化OpenGL。

回放缓冲区

ReplayBuffer（在./diffusion_policy/common/replay_buffer.py）是关键数据结构，用于以内存和磁盘方式存储带有分块和压缩的演示数据集。它大量使用了zarr格式，但也具有降低访问开销的numpy后端。

在磁盘上，它可以作为一个嵌套目录存储（如data/pusht_cchi_v7_replay.zarr）或zip文件（如data/robomimic/datasets/square/mh/image_abs.hdf5.zarr.zip）。

由于我们的数据集相对较小，经常可以用Jpeg2000压缩存储整个基于图像的数据集（在./diffusion_policy/codecs/imagecodecs_numcodecs.py），从而消除训练过程中的磁盘IO，但会增加CPU负载。

示例：

data/pusht_cchi_v7_replay.zarr
 ├── data
 │   ├── action (25650, 2) float32
 │   ├── img (25650, 96, 96, 3) float32
 │   ├── keypoint (25650, 9, 2) float32
 │   ├── n_contacts (25650, 1) float32
 │   └── state (25650, 5) float32
 └── meta
     └── episode_ends (206,) int64

data下的每个数组沿着第一个维度（时间）存储来自所有序列的一个数据字段。meta/episode_ends数组存储沿着第一个维度的每个序列结束索引。

SharedMemoryRingBuffer

SharedMemoryRingBuffer(./diffusion_policy/shared_memory/shared_memory_ring_buffer.py)是一种无锁FILO（先进后出）数据结构，在我们的真实机器人实现中广泛使用，用于在多个CPU核心之间并行操作，避免了pickle序列化和多进程队列锁定带来的开销。

例如，假设我们想要获取5个RealSense相机最新的To帧。每个进程启动一个使用SingleRealsense(./diffusion_policy/real_world/single_realsense.py)的RealSense SDK/管道，它们不断将捕获的图像写入与主进程共享的SharedMemoryRingBuffer。由于SharedMemoryRingBuffer的FILO特性，我们可以在主进程中非常快速地获取最后To帧。

此外，我们也实现了SharedMemoryQueue(./diffusion_policy/shared_memory/shared_memory_queue.py)以满足FIFO（先进先出）的需求，它被用在RTDEInterpolationController(./diffusion_policy/real_world/rtde_interpolation_controller.py)中。

RealEnv

与OpenAI Gym不同，我们的策略与环境异步交互。在RealEnv(./diffusion_policy/real_world/real_env.py)中，gym的step方法被拆分为两个方法：get_obs和exec_actions。

get_obs方法从SharedMemoryRingBuffer返回最新的观测值及其对应的时间戳，这个方法可以在评估期间的任意时刻调用。

exec_actions方法接收一序列的动作和期望执行时间的时标。一旦调用，这些动作会被简单地添加到RTDEInterpolationController的队列中，然后方法立即返回，不阻塞执行过程。

添加任务

参考并模仿：

• diffusion_policy/dataset/pusht_image_dataset.py
• diffusion_policy/env_runner/pusht_image_runner.py
• diffusion_policy/config/task/pusht_image.yaml

确保shape_meta符合你的任务输入和输出形状。同时，确认env_runner._target_和dataset._target_指向你新添加的类。训练时，在train.py的参数中添加task=<你的任务名称>。

添加方法

参考并模仿：

• diffusion_policy/workspace/train_diffusion_unet_image_workspace.py
• diffusion_policy/policy/diffusion_unet_image_policy.py
• diffusion_policy/config/train_diffusion_unet_image_workspace.yaml

确保工作区yaml文件的_target_指向你创建的新工作区类。

许可证

此仓库遵循MIT许可。详情见LICENSE。

致谢

• 我们的ConditionalUnet1D(./diffusion_policy/model/diffusion/conditional_unet1d.py)实现改编自Planning with Diffusion。
• 我们的TransformerForDiffusion(./diffusion_policy/model/diffusion/transformer_for_diffusion.py)实现改编自MinGPT。
• BET基线改编自其原始仓库。
• IBC基线改编自Kevin Zakka的重实现。
• 使用了Robomimic的任务以及ObservationEncoder。
• Push-T任务改编自IBC。
• Block Pushing任务改编自BET 和 IBC。
• Kitchen任务改编自BET 和 Relay Policy Learning。
• 我们shared_memory数据结构深受shared-ndarray2启发。