扩散策略

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

Cheng Chi1， Siyuan Feng2， Yilun Du3， Zhenjia Xu1， Eric Cousineau2， Benjamin Burchfiel2， Shuran Song1

1哥伦比亚大学， 2丰田研究院， 3麻省理工学院

🛝 立即尝试！

我们的自包含Google Colab笔记本是体验扩散策略的最简单方式。我们为基于状态的环境和基于视觉的环境提供了单独的笔记本。

🧾 查看我们的实验日志！

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

一个包含重现实验所需的所有参数的config.yaml文件。
每个训练步骤的详细训练/评估logs.json.txt文件。
每次运行的最佳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汇总了所有3次训练运行每50个周期的评估指标。论文中报告的数字对应于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）：

1x UR5-CB3或UR5e（需要RTDE接口）
2x RealSense D415
1x 3Dconnexion SpaceMouse（用于远程操作）
1x Millibar Robotics手动工具更换器（仅需机器人端）
1x 3D打印末端执行器
1x 3D打印T型块
RealSense的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上，同时还会记录一些rollout的视频。

(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'
data/outputs/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 directories, 16 files

🆕 评估预训练检查点

从已发布的训练日志文件夹下载检查点，例如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
data/pusht_eval_output
├── eval_log.json
└── media
    ├── 1fxtno84.mp4
    ├── 224l7jqd.mp4
    ├── 2fo4btlf.mp4
    ├── 2in4cn7a.mp4
    ├── 34b3o2qq.mp4
    └── 3p7jqn32.mp4

1 directory, 7 files

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训练数据集相同。

要训练Diffusion Policy，使用以下配置启动训练：

(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。假设训练已完成，且在 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"停止当前回合。

🗺️ 代码库教程

本代码库的结构基于以下要求：

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

为了实现这一要求，我们：

在任务和方法之间维护了一个简单统一的接口
使任务和方法的实现相互独立

这些设计决策的代价是任务和方法之间存在代码重复。但我们认为，能够在不影响其他部分的情况下添加/修改任务/方法，以及能够通过线性阅读代码来理解任务/方法的好处，远远超过了复制粘贴的成本 😊。

分割

在任务方面，我们有：

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

在策略方面，我们有：

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

接口

低维

LowdimPolicy 接收观察字典：

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

并预测动作字典：

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

LowdimDataset 返回一个字典样本：

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

其 get_normalizer 方法返回一个带有 "obs" 和 "action" 键的 LinearNormalizer。

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

图像

ImagePolicy 接收观察字典：

"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`

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

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

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

`Dataset`

Dataset 对象：

继承自 torch.utils.data.Dataset。
返回符合接口的样本，具体取决于任务是低维还是图像观察。
有一个 get_normalizer 方法，返回符合接口的 LinearNormalizer。

归一化是项目开发过程中很常见的 bug 来源。打印出 LinearNormalizer 中每个键使用的具体 scale 和 bias 向量有时会有所帮助。

我们大多数 Dataset 的实现都使用 ReplayBuffer 和 SequenceSampler 的组合来生成样本。根据 To 和 Ta 正确处理每个演示回合开始和结束时的填充对于获得良好性能很重要。在实现自己的采样方法之前，请阅读我们的 SequenceSampler。

`Policy`

Policy 对象：

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

`EnvRunner`

EnvRunner 对象抽象化了不同任务环境之间的细微差别。

有一个 run 方法，接收用于评估的 Policy 对象，并返回日志和指标的字典。每个值应与 wandb.log 兼容。

为了最大化评估速度，我们通常使用我们对 gym.vector.AsyncVectorEnv 的修改来向量化环境，该修改在单独的进程中运行每个单独的环境（绕过 Python GIL）。

⚠️ 由于子进程在 Linux 上使用 fork 启动，对于在初始化期间创建 OpenGL 上下文的环境（如 robosuite），你需要特别小心，因为一旦被子进程内存空间继承，往往会导致模糊的 bug，如段错误。作为解决方法，你可以提供一个 dummy_env_fn，构造一个不初始化 OpenGL 的环境。

`ReplayBuffer`

ReplayBuffer 是一个关键的数据结构，用于在内存和磁盘上存储演示数据集，具有分块和压缩功能。它大量使用 zarr 格式，但也有一个 numpy 后端，以降低访问开销。在磁盘上，它可以存储为嵌套目录（例如 data/pusht_cchi_v7_replay.zarr）或 zip 文件（例如 data/robomimic/datasets/square/mh/image_abs.hdf5.zarr.zip）。

由于我们的数据集相对较小，通常可以使用 Jpeg2000 压缩将整个基于图像的数据集存储在 RAM 中，这样可以消除训练期间的磁盘 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 是一个无锁的 FILO 数据结构，在我们的真实机器人实现中广泛使用，以利用多个 CPU 核心，同时避免 multiprocessing.Queue 的 pickle 序列化和锁定开销。

例如，我们想从 5 个 RealSense 相机获取最近的 To 帧。我们使用 SingleRealsense 为每个进程启动 1 个 RealSense SDK/pipeline，每个进程持续将捕获的图像写入与主进程共享的 SharedMemoryRingBuffer。由于 SharedMemoryRingBuffer 的 FILO 特性，我们可以在主进程中非常快速地获取最后的 To 帧。

我们还为 FIFO 实现了 SharedMemoryQueue，它在 RTDEInterpolationController 中使用。