Diffusion Policy实战：如何用扩散模型提升机器人操作精度（附代码示例）

最近在机器人实验室里，我们团队遇到了一个挺有意思的瓶颈：让机械臂完成一些看似简单的抓取和放置任务时，传统的模仿学习方法总是差那么一点意思。动作要么过于“平均”，缺乏灵性；要么在遇到稍微复杂的场景时，决策就变得犹豫不决，动作分布呈现出明显的多峰性——简单说，就是机器人“不知道选哪个动作好”。这让我想起了图像生成领域里扩散模型（Diffusion Model）的惊艳表现，它不正是擅长从噪声中“去伪存真”，生成丰富多样的高质量样本吗？那么，能否把这种生成能力“嫁接”到机器人动作序列的生成上呢？

答案是肯定的，这正是Diffusion Policy的核心思想。它不再像传统方法那样直接预测一个单一的动作，而是将动作序列的生成建模为一个“去噪”过程。想象一下，你给机器人一个模糊的、充满噪声的动作草图，然后让它一步步地“擦除”不合理的部分，最终得到一个清晰、精确、且能应对多种可能性的动作计划。这种方法尤其适合机器人操作任务，因为现实世界本身就充满了不确定性，一个优秀的策略需要能处理这种“多峰”的动作分布。对于机器人开发者、AI研究员以及任何希望将前沿生成式AI落地到实体智能体上的朋友来说，理解并实践Diffusion Policy，或许能为你打开一扇新的大门。

1. Diffusion Policy：为何是机器人模仿学习的“破局者”？

在深入代码之前，我们有必要先厘清Diffusion Policy究竟解决了什么根本问题。传统的机器人模仿学习，无论是基于行为克隆（Behavioral Cloning）的LSTM-GMM，还是隐式行为克隆（IBC），亦或是行为Transformer（BET），都面临着一个共同的挑战：动作分布的多峰性（Multi-modality）。

举个简单的例子，教机器人把桌上的一个方块推到指定区域。从起点到终点，可能存在多条等效的路径：可以直线推过去，也可以先向左再向右绕个弧线。传统的确定性策略（如MLP）会输出一个“平均”动作，可能导致动作不自然甚至无效；而一些概率性方法在拟合这种复杂分布时，要么训练不稳定，要么表达能力有限。

Diffusion Policy的巧妙之处在于，它借用了扩散模型在生成式AI中的范式。其核心流程可以概括为：

1. 前向扩散过程：将一个干净的动作序列逐步添加噪声，直至变成纯高斯噪声。
2. 反向去噪过程：训练一个神经网络，学习如何从带噪声的动作序列和当前的环境观测（如图像、状态）中，预测出所添加的噪声。
3. 推理生成：在部署时，从一个随机噪声开始，利用训练好的网络进行多步迭代去噪，最终“生成”一个合理的动作序列。

这个过程带来了几个关键优势：

• 强大的多峰分布建模能力：扩散模型本身就是为了生成多样化的高质量样本而设计的，天生适合捕捉动作空间中的多种可能性。
• 高维连续输出：可以直接输出未来一段时间窗口内的高维动作序列（如多个时间步的关节角度或末端执行器位姿），实现时序上的规划。
• 训练稳定性：不同于一些需要困难负样本采样的能量基模型，扩散模型的训练目标（预测噪声）相对更直接和稳定。

下表对比了Diffusion Policy与几种主流模仿学习方法的核心差异：

特性	Diffusion Policy	LSTM-GMM	隐式行为克隆 (IBC)	行为Transformer (BET)
多峰处理能力	优秀，通过生成式去噪过程自然实现	中等，通过混合高斯模型显式建模	依赖基于能量的模型，采样可能不稳定	中等，依赖离散化或VQ-VAE
输出形式	连续的高维动作序列	连续的参数化分布（高斯混合）	连续的隐式分布	离散的token序列或连续序列
训练稳定性	高，目标函数平滑	中等	较低，对负采样敏感	高
时序建模	内置，通过处理序列噪声	通过LSTM循环网络	通常为单步	通过Transformer注意力机制
实时性考量	需要多步迭代，对计算有要求	单步前向传播，较快	需要迭代优化，可能较慢	单步前向传播，较快

提示：虽然Diffusion Policy在精度和泛化性上表现突出，但其多步迭代的去噪过程意味着更高的计算开销。在部署到真实机器人时，需要在性能提升和实时性要求之间做出权衡，或通过模型压缩、蒸馏等技术进行优化。

2. 从零搭建Diffusion Policy训练管道

理论很美好，但我们需要脚踏实地。这一部分，我们将聚焦于如何准备数据、构建模型并启动训练。我将以PyTorch框架为例，展示一个简化但完整的训练流程。假设我们的任务是基于视觉的机械臂抓取。

2.1 数据准备与预处理

机器人模仿学习的数据通常是一系列(观测， 动作)对。观测可以是图像、机器人本体状态（关节角度、力传感器读数等），动作则是执行器指令。

import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
import h5py  # 假设数据存储在HDF5文件中

class RobotDataset(Dataset):
    """
    一个简单的机器人操作数据集类。
    假设数据格式：每个episode包含N个时间步的观测和动作。
    观测: {'image': (T, C, H, W), 'state': (T, D_state)}
    动作: (T, D_action)
    """
    def __init__(self, hdf5_path, episode_length=100, action_horizon=10, observation_horizon=2):
        super().__init__()
        self.data = h5py.File(hdf5_path, 'r')
        self.episode_keys = list(self.data.keys())
        self.episode_length = episode_length
        self.action_horizon = action_horizon  # 模型预测的动作序列长度
        self.observation_horizon = observation_horizon  # 模型观察的历史长度

    def __len__(self):
        return len(self.episode_keys) * (self.episode_length - self.action_horizon - self.observation_horizon + 1)

    def __getitem__(self, idx):
        ep_idx = idx // (self.episode_length - self.action_horizon - self.observation_horizon + 1)
        t_start = idx % (self.episode_length - self.action_horizon - self.observation_horizon + 1)

        ep_key = self.episode_keys[ep_idx]
        episode = self.data[ep_key]

        # 获取观测序列
        # 图像观测：取最近 observation_horizon 帧
        images = episode['images'][t_start:t_start+self.observation_horizon]  # (O, C, H, W)
        # 状态观测
        states = episode['states'][t_start:t_start+self.observation_horizon]   # (O, D_state)
        # 拼接观测
        # 注意：这里简化处理，实际中图像需要经过编码器
        obs = {
            'image': torch.from_numpy(images).float(),
            'state': torch.from_numpy(states).float()
        }

        # 获取目标动作序列
        # 模型需要预测从当前时刻开始的未来 action_horizon 步动作
        target_actions = episode['actions'][t_start:t_start+self.action_horizon]  # (A, D_action)
        target_actions = torch.from_numpy(target_actions).float()

        return obs, target_actions

数据预处理的关键在于确定action_horizon和observation_horizon。前者决定了模型规划的未来步长，后者决定了模型决策所依赖的历史上下文长度。这两个超参数对性能影响显著。

2.2 核心模型构建：噪声预测网络

Diffusion Policy的核心是一个噪声预测网络（通常是一个U-Net或Transformer）。它的输入是带噪声的动作序列、扩散过程的时间步索引以及条件信息（即观测），输出是预测的噪声。

import torch.nn as nn
import torch.nn.functional as F

class ConditionalUNet(nn.Module):
    """
    一个简化的条件U-Net，用于预测噪声。
    输入: noisy_action_sequence, diffusion_timestep, conditioning_observation
    输出: predicted_noise
    """
    def __init__(self, action_dim, obs_feat_dim, hidden_dim=256):
        super().__init__()
        self.action_dim = action_dim
        self.time_embedding = nn.Sequential(
            nn.Linear(1, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.obs_embedding = nn.Linear(obs_feat_dim, hidden_dim)

        # 简化的U-Net骨干（这里用MLP示意，实际应用会更复杂）
        self.input_proj = nn.Linear(action_dim + hidden_dim*2, hidden_dim)
        self.mid_block = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
        )
        self.output_layer = nn.Linear(hidden_dim, action_dim)

    def forward(self, noisy_actions, timesteps, obs_features):
        """
        noisy_actions: (B, A, action_dim) 带噪声的动作序列
        timesteps: (B,) 扩散时间步
        obs_features: (B, obs_feat_dim) 编码后的观测特征
        """
        B, A, D = noisy_actions.shape
        # 1. 时间步嵌入
        t_emb = self.time_embedding(timesteps.unsqueeze(-1).float() / 1000.0) # (B, hidden_dim)
        t_emb = t_emb.unsqueeze(1).repeat(1, A, 1) # (B, A, hidden_dim)

        # 2. 观测条件嵌入
        c_emb = self.obs_embedding(obs_features) # (B, hidden_dim)
        c_emb = c_emb.unsqueeze(1).repeat(1, A, 1) # (B, A, hidden_dim)

        # 3. 拼接所有输入
        x = torch.cat([noisy_actions, t_emb, c_emb], dim=-1) # (B, A, action_dim + hidden_dim*2)
        x = x.reshape(B*A, -1)
        x = self.input_proj(x)
        x = F.silu(x)
        x = self.mid_block(x)
        noise_pred = self.output_layer(x)
        noise_pred = noise_pred.reshape(B, A, D)
        return noise_pred

在实际的Diffusion Policy实现中（如原论文的CNN或Transformer版本），网络结构会复杂得多，特别是用于处理图像观测的视觉编码器（如ResNet）和处理时序的扩散Transformer（Diffusion Transformer）。

2.3 训练循环与损失函数

扩散模型的训练目标非常直观：最小化预测噪声与真实噪声之间的差距。

def train_one_epoch(model, dataloader, optimizer, device, noise_scheduler):
    model.train()
    total_loss = 0
    for batch_idx, (obs, target_actions) in enumerate(dataloader):
        # 将数据移至设备
        obs = {k: v.to(device) for k, v in obs.items()}
        target_actions = target_actions.to(device) # (B, A, D)

        B, A, D = target_actions.shape

        # --- 扩散过程 ---
        # 1. 采样随机时间步
        timesteps = torch.randint(0, noise_scheduler.num_train_timesteps, (B,), device=device).long()
        # 2. 采样随机噪声
        noise = torch.randn_like(target_actions)
        # 3. 根据时间步向干净动作添加噪声
        noisy_actions = noise_scheduler.add_noise(target_actions, noise, timesteps)

        # --- 编码观测条件 ---
        # 此处简化：假设obs['image']已通过一个预训练或并行的CNN编码器得到特征
        # obs_features = vision_encoder(obs['image'])
        # 为示例，我们使用状态特征代替
        obs_features = obs['state'].mean(dim=1) # (B, D_state) 简化处理

        # --- 模型前向传播，预测噪声 ---
        predicted_noise = model(noisy_actions, timesteps, obs_features)

        # --- 计算损失 ---
        loss = F.mse_loss(predicted_noise, noise)

        # --- 反向传播与优化 ---
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)

这里的noise_scheduler负责管理扩散过程中的噪声调度（如线性、余弦等），它定义了不同时间步下噪声的强度。常用的库如diffusers提供了丰富的调度器实现。

3. 推理部署：从噪声到机器人动作

模型训练好后，如何在机器人上运行？推理过程是一个迭代的去噪采样过程。

@torch.no_grad()
def sample_action_sequence(model, obs, noise_scheduler, num_inference_steps=50, action_horizon=10, action_dim=7):
    """
    给定当前观测，采样生成动作序列。
    obs: 当前观测字典，包含图像和状态
    """
    model.eval()
    device = next(model.parameters()).device

    # 1. 编码观测条件
    # obs_features = vision_encoder(obs['image'])
    obs_features = obs['state'].mean(dim=0, keepdim=True) # (1, D_state) 简化

    # 2. 初始化随机噪声动作序列
    sample = torch.randn((1, action_horizon, action_dim), device=device)

    # 3. 设置采样调度器
    noise_scheduler.set_timesteps(num_inference_steps, device=device)

    # 4. 迭代去噪
    for t in noise_scheduler.timesteps:
        # 扩展时间步和条件以匹配样本维度
        timestep = torch.tensor([t], device=device).long()
        # 预测噪声
        noise_pred = model(sample, timestep, obs_features)
        # 使用调度器计算去噪后的样本
        sample = noise_scheduler.step(noise_pred, t, sample).prev_sample

    # 5. 生成的样本即为预测的动作序列
    predicted_actions = sample.squeeze(0).cpu().numpy() # (A, D)
    # 通常只执行序列中的第一个动作，然后重新规划（receding horizon control）
    action_to_execute = predicted_actions[0]
    return action_to_execute, predicted_actions

在实际的机器人控制中，我们通常采用滚动时域控制（Receding Horizon Control）。即每次只执行生成动作序列的第一个动作，然后获取新的观测，再重新运行整个采样过程。这形成了一个闭环，使策略能够应对动态变化的环境。

注意：多步迭代采样是计算瓶颈。为了满足实时控制频率（如10-100Hz），可能需要使用更快的采样器（如DDIM），减少采样步数，或使用蒸馏技术将扩散模型压缩为一步生成的模型。

4. 关键调优技巧与实战避坑指南

纸上得来终觉浅，绝知此事要躬行。在将Diffusion Policy应用到真实任务时，以下几个方面的调优至关重要，它们往往决定了项目的成败。

4.1 观测编码与特征融合

观测信息的质量直接决定了策略的条件生成效果。

• 视觉编码器：如果使用图像，一个在大型数据集（如ImageNet）上预训练的ResNet backbone是很好的起点。是否需要微调取决于你的数据量和任务特异性。特征提取层不宜过深，避免引入过多与机器人操作无关的语义信息。
• 多模态融合：如何融合图像特征和向量状态（如关节角度、力觉）？简单的拼接（concatenation）有时就足够有效。更复杂的方法可以使用交叉注意力（Cross-Attention）机制，让扩散模型在去噪过程中动态关注观测的不同部分。

  # 一个简单的特征融合示例（在模型内部）
  # image_feat: (B, F_img)
  # state_feat: (B, F_state)
  fused_feat = torch.cat([image_feat, state_feat], dim=-1)
  fused_feat = self.fusion_proj(fused_feat) # 投影到统一维度
  

4.2 扩散过程参数化

扩散过程的设计对生成质量和速度有巨大影响。

• 噪声调度（Noise Schedule）：决定了噪声如何随时间步添加。余弦调度（cosine schedule）通常比线性调度在图像生成上效果更好，在机器人任务中也值得尝试。它在前向过程开始和结束时变化较慢，中间变化较快。
• 采样步数：训练时通常使用较多的步数（如1000步），但推理时可以通过加速采样器（DDIM、DPM-Solver）大幅减少到50步甚至20步，在精度和速度间取得平衡。
• 分类器自由引导（Classifier-Free Guidance）：这是一种在生成过程中增强条件控制的技术。在训练时，以一定概率（如10%）将条件信息（观测）置零（无条件训练）。在推理时，通过一个引导尺度（guidance scale）来调整有条件预测和无条件预测的权重，可以显著提升生成动作与条件的对齐程度。虽然增加了训练开销，但效果提升往往非常明显。

4.3 处理高频与动态任务

原论文指出，基于CNN的扩散策略在输出上可能过于平滑，对于需要高频、敏捷动作的任务（如灵巧操作）可能不足。为此，作者提出了Time-series Diffusion Transformer。

• 核心思想：将整个动作序列和历史观测序列视为一个时间序列，用Transformer架构进行全局建模。其自注意力机制能更好地捕捉长距离依赖和瞬间变化。
• 实现关键：需要为时间步、观测序列和噪声动作序列设计合适的嵌入和位置编码。计算开销比CNN版本大，但对复杂任务更有效。

4.4 仿真到实物的迁移

如果你主要在仿真中训练，希望部署到真实机器人，需要注意：

• 域随机化（Domain Randomization）：在仿真训练时，随机化纹理、光照、物体质量、摩擦系数等，以增加策略的鲁棒性。
• 观测对齐：确保仿真和真实环境的观测空间尽可能一致。例如，相机视角、图像分辨率、颜色通道等。
• 动作空间与动力学：仿真与实物的电机响应、延迟存在差异。可以考虑在动作输出层加入低通滤波器，或在训练数据中引入少量的动作延迟噪声。

我在一个桌面杂物清理的任务上尝试过Diffusion Policy。最初直接用仿真数据训练的策略在实物上表现不佳，机械臂的动作显得“僵硬”且容易在接触物体时振荡。后来，我们在仿真中加入了大量的电机噪声和随机延迟，并对输出的关节速度指令进行了平滑滤波。同时，我们收集了少量实物演示数据（约50条轨迹），与仿真数据混合进行微调。这个过程虽然增加了工作量，但最终策略在实物上的抓取成功率从不到40%提升到了85%以上。最关键的是，策略展现出了令人惊喜的“试探性”行为：当第一次抓取滑脱时，它会稍微调整姿态再次尝试，这种多峰策略的能力是传统方法很难自然表现出来的。