Diffusion Policy实战：5步搞定机器人动作预测（附Python代码）
Diffusion Policy实战：5步搞定机器人动作预测（附Python代码）
最近在机器人圈子里，一个词被反复提及：Diffusion Policy。如果你还在用传统的强化学习（RL）方法，比如PPO或SAC，来调教你的机械臂或移动机器人，可能会发现它们在面对复杂、高维度的连续动作空间时，显得有些力不从心。动作要么过于“僵硬”，缺乏多样性；要么在遇到未曾见过的环境状态时，直接“懵圈”。斯坦福UMI团队提出的Diffusion Policy，就像是为这个痛点量身定制的解药。它把图像生成领域大放异彩的扩散模型（Diffusion Model）搬了过来，让机器人动作的生成过程，从“猜一个答案”变成了“从模糊到清晰”的逐步雕琢。
这听起来很酷，但具体怎么用？网上的论文解读很多，但一涉及到代码落地，很多人就卡在了第一步。这篇文章就是为你准备的实战手册。我们不谈复杂的数学推导，只聚焦于如何用Python，在最短的时间内，搭建一个能跑起来的Diffusion Policy原型，并让它对你的机器人状态做出合理的动作预测。无论你是正在做机械臂抓取研究的工程师，还是探索前沿AI控制算法的开发者，这五个步骤都能帮你快速上手，把前沿论文变成你项目里的实际工具。
1. 环境准备与核心概念速览
在动手写代码之前，我们需要先理清几个关键概念，并准备好“战场”。Diffusion Policy的核心思想，可以类比为一个“去噪”的艺术修复过程。想象一下，你有一张被严重噪声污染的、模糊不清的“理想动作序列”图片。Diffusion Policy的训练，就是学习如何一步步擦除这些噪声，还原出清晰的动作。而在推理时，我们则从一张完全随机的“噪声图片”开始，让训练好的模型按照学到的“修复步骤”，逐步“画”出最终的动作。
这个过程与传统RL策略有本质区别：
- 传统策略（如MLP策略网络）：输入状态，直接输出一个确定的动作。它像一个“点对点”的映射函数。
- Diffusion Policy：输入状态，输出的是一个“去噪过程”。它定义了一个从噪声分布到动作分布的“路径”，因此天生能建模多模态（即同一个状态下可能有多个合理动作）和复杂的动作分布。
为了开始我们的实战，你需要准备以下环境。我强烈建议使用Anaconda或Miniconda来管理环境，避免依赖冲突。
# 创建并激活一个新的conda环境（Python 3.9是一个稳定的选择）
conda create -n diffusion_policy python=3.9 -y
conda activate diffusion_policy
# 安装PyTorch（请根据你的CUDA版本到PyTorch官网选择对应命令）
# 这里以CUDA 11.8为例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 安装核心扩散模型库diffusers，以及一些实用工具
pip install diffusers transformers scikit-learn matplotlib numpy tqdm
提示：如果你的机器没有NVIDIA GPU，可以安装CPU版本的PyTorch，但训练和推理速度会慢很多。对于初步学习和原型验证，CPU版本也是可行的。
2. 构建一个最小化的Diffusion Policy模型
我们不打算一开始就复现论文中复杂的U-Net结构。为了快速理解，我们先构建一个极其简化的“玩具”模型。这个模型的核心是一个“去噪网络”，它负责在扩散过程的每一步，根据当前带噪声的动作和机器人状态，预测出应该被移除的噪声。
让我们先定义一个简单的全连接网络作为我们的去噪模型：
import torch
import torch.nn as nn
import torch.nn.functional as F
class SimpleDenoiser(nn.Module):
"""
一个简单的去噪网络。
输入：带噪声的动作 `noisy_action` 和机器人状态 `state`
输出：预测的噪声 `predicted_noise`
"""
def __init__(self, state_dim=8, action_dim=4, hidden_dim=256):
super().__init__()
# 将状态和带噪声的动作拼接起来作为输入
self.input_layer = nn.Linear(state_dim + action_dim, hidden_dim)
self.hidden_layer1 = nn.Linear(hidden_dim, hidden_dim)
self.hidden_layer2 = nn.Linear(hidden_dim, hidden_dim)
self.output_layer = nn.Linear(hidden_dim, action_dim) # 输出噪声，维度与动作相同
# 使用LayerNorm和Dropout来提升训练稳定性（可选但推荐）
self.norm1 = nn.LayerNorm(hidden_dim)
self.norm2 = nn.LayerNorm(hidden_dim)
self.dropout = nn.Dropout(0.1)
def forward(self, state, noisy_action):
# state: [batch_size, state_dim]
# noisy_action: [batch_size, action_dim]
x = torch.cat([state, noisy_action], dim=-1) # 拼接特征
x = F.relu(self.input_layer(x))
x = self.norm1(x)
x = self.dropout(x)
x = F.relu(self.hidden_layer1(x))
x = self.norm2(x)
x = self.dropout(x)
x = F.relu(self.hidden_layer2(x))
x = self.output_layer(x) # 预测的噪声
return x
这个网络结构非常简单，但它已经具备了Diffusion Policy去噪器的基本形态：根据状态和带噪动作，估算噪声。在实际的复杂任务中（如从图像像素中预测动作），这个网络会被替换为U-Net或Transformer等结构，但输入输出的逻辑是一致的。
接下来，我们需要定义扩散过程的核心调度器。它控制着噪声添加和去除的节奏（即不同时间步的噪声强度）。幸运的是，diffusers库为我们提供了现成的工具。这里我们使用经典的DDPM（Denoising Diffusion Probabilistic Models）调度器。
from diffusers import DDPMScheduler
# 初始化调度器
num_diffusion_timesteps = 1000 # 扩散总步数，步数越多通常生成质量越高，但计算越慢
scheduler = DDPMScheduler(num_train_timesteps=num_diffusion_timesteps,
beta_schedule='linear') # 噪声调度方式
# 我们可以快速查看一下不同时间步的噪声强度（beta）
print(f"Beta schedule sample (first 5): {scheduler.betas[:5]}")
print(f"Alpha bars (sqrt of cumulative product, first 5): {scheduler.alphas_cumprod[:5]}")
alphas_cumprod（常记作ᾱ_t）是一个关键量，它决定了在时间步t，原始信号还保留多少“信息”。当t=0时，ᾱ_t=1，表示完全是原始信号；当t接近T时，ᾱ_t接近0，表示几乎全是噪声。
3. 准备与处理训练数据
任何机器学习模型都离不开数据。对于Diffusion Policy，我们需要的是状态-动作对序列数据。这些数据可以来自：
- 专家演示：人类遥控操作机器人完成任务记录下的
(state, action)。 - 传统控制器：使用PID、MPC等经典方法控制机器人运行产生的数据。
- 其他策略收集：用已有的RL策略在环境中交互收集。
为了演示，我们生成一个简单的模拟数据集。假设我们的机器人状态是8维（可能包含位置、速度等），动作是4维（如关节扭矩或末端执行器速度）。
import numpy as np
from sklearn.preprocessing import StandardScaler
def generate_synthetic_data(num_samples=10000, state_dim=8, action_dim=4):
"""
生成一个简单的线性关系合成数据集。
在实际应用中，你需要用真实数据替换这部分。
"""
np.random.seed(42)
# 生成随机状态
states = np.random.randn(num_samples, state_dim)
# 假设动作是状态的线性函数加上一些噪声（模拟不完美的专家）
true_weights = np.random.randn(state_dim, action_dim)
actions = states @ true_weights + 0.1 * np.random.randn(num_samples, action_dim)
# 数据标准化非常重要，能加速模型收敛
state_scaler = StandardScaler()
action_scaler = StandardScaler()
states_normalized = state_scaler.fit_transform(states)
actions_normalized = action_scaler.fit_transform(actions)
# 转换为PyTorch张量
states_tensor = torch.FloatTensor(states_normalized)
actions_tensor = torch.FloatTensor(actions_normalized)
return states_tensor, actions_tensor, state_scaler, action_scaler
# 生成数据
states, actions, state_scaler, action_scaler = generate_synthetic_data()
print(f"States shape: {states.shape}, Actions shape: {actions.shape}")
现在，我们有了标准化的状态和动作数据。在Diffusion Policy的训练中，我们的目标不是让网络直接预测动作，而是让它学会在任意噪声水平下，如何从“带噪声的动作”中预测出被添加的噪声。
4. 编写训练循环与损失函数
这是整个流程中最关键的一步。我们将按照扩散模型的标准训练流程来训练我们的SimpleDenoiser。
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
# 1. 创建数据集和数据加载器
dataset = TensorDataset(states, actions)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True, pin_memory=True)
# 2. 初始化模型、优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = SimpleDenoiser(state_dim=8, action_dim=4, hidden_dim=256).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
# 3. 训练参数
num_epochs = 200
losses = []
model.train()
for epoch in range(num_epochs):
epoch_loss = 0.0
progress_bar = tqdm(dataloader, desc=f'Epoch {epoch+1}/{num_epochs}')
for batch_states, batch_actions in progress_bar:
batch_states = batch_states.to(device)
batch_actions = batch_actions.to(device)
# --- 扩散模型训练的核心步骤 ---
# a. 为每个样本随机选择一个扩散时间步 t
batch_size = batch_actions.shape[0]
timesteps = torch.randint(0, scheduler.num_train_timesteps, (batch_size,), device=device).long()
# b. 根据时间步 t 和调度器，为干净动作添加噪声
noise = torch.randn_like(batch_actions) # 标准高斯噪声
noisy_actions = scheduler.add_noise(batch_actions, noise, timesteps)
# c. 让去噪网络预测噪声
predicted_noise = model(batch_states, noisy_actions)
# d. 计算损失：预测噪声与真实噪声之间的均方误差 (MSE)
loss = F.mse_loss(predicted_noise, noise)
# e. 反向传播与优化
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) # 梯度裁剪，防止爆炸
optimizer.step()
epoch_loss += loss.item()
progress_bar.set_postfix({'loss': loss.item()})
avg_epoch_loss = epoch_loss / len(dataloader)
losses.append(avg_epoch_loss)
print(f'Epoch {epoch+1} finished. Average Loss: {avg_epoch_loss:.6f}')
# 可选：每50个epoch保存一次模型检查点
if (epoch + 1) % 50 == 0:
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': avg_epoch_loss,
}, f'./checkpoint_epoch_{epoch+1}.pth')
注意：这里的损失函数是预测噪声和真实噪声的MSE。这是扩散模型最基础的训练目标。通过最小化这个损失，模型学会了“如何根据当前状态和带噪动作，判断出混入了多少、什么样的噪声”，从而掌握了“去噪”的能力。
训练完成后，你可以绘制损失曲线来检查收敛情况。一个成功训练的标志是损失值稳步下降并逐渐趋于平缓。
5. 推理：从噪声生成机器人动作
模型训练好后，最激动人心的部分来了：如何用它来为机器人生成动作？推理过程是一个迭代去噪的过程，从纯噪声开始，逐步“减去”模型预测的噪声。
@torch.no_grad()
def generate_action(model, state, scheduler, num_inference_steps=50, action_scaler=None):
"""
使用训练好的Diffusion Policy生成一个动作。
Args:
model: 训练好的去噪模型
state: 当前机器人状态，形状为 [state_dim]
scheduler: 扩散调度器
num_inference_steps: 去噪步数，可以小于训练步数以加速
action_scaler: 用于将标准化动作反标准化回原始空间的scaler
Returns:
action: 生成的机器人动作
"""
model.eval()
device = next(model.parameters()).device
# 1. 将状态处理成合适的形状 [1, state_dim]
if state.dim() == 1:
state = state.unsqueeze(0).to(device)
else:
state = state.to(device)
# 2. 从标准高斯分布中采样一个初始噪声动作
action_shape = (1, model.output_layer.out_features) # [1, action_dim]
noisy_action = torch.randn(action_shape, device=device)
# 3. 设置推理步数，并调整调度器的时间步序列
scheduler.set_timesteps(num_inference_steps, device=device)
# 4. 迭代去噪循环
for t in scheduler.timesteps:
# a. 用模型预测当前时间步的噪声
noise_pred = model(state, noisy_action)
# b. 使用调度器计算“更干净”的样本
# scheduler.step() 会根据预测的噪声、当前 noisy_action 和时间步 t，计算出下一步的 noisy_action
noisy_action = scheduler.step(noise_pred, t, noisy_action).prev_sample
# 5. 循环结束后，noisy_action 就是去噪后的动作（标准化后的）
generated_action_normalized = noisy_action.squeeze(0).cpu() # [action_dim]
# 6. 如果提供了scaler，则反标准化到原始动作空间
if action_scaler is not None:
generated_action = action_scaler.inverse_transform(generated_action_normalized.unsqueeze(0))
generated_action = torch.FloatTensor(generated_action).squeeze(0)
else:
generated_action = generated_action_normalized
return generated_action
# 使用示例：假设我们有一个新的机器人状态
new_state = torch.randn(8) # 一个8维的随机状态，实际中来自传感器
print(f"Current robot state: {new_state}")
# 生成动作
predicted_action = generate_action(model,
new_state,
scheduler,
num_inference_steps=100,
action_scaler=action_scaler)
print(f"Predicted robot action by Diffusion Policy: {predicted_action}")
这个generate_action函数就是Diffusion Policy策略的核心调用接口。你只需要将当前的机器人状态state输入，它就会经过一个完整的迭代去噪过程，输出一个动作。num_inference_steps是一个重要的超参数，步数越多，生成质量通常越高，但耗时也越长。在实际部署中，需要在速度和质量之间做权衡。
6. 进阶技巧与实战调优指南
走通了上面的五步，你已经拥有了一个可工作的Diffusion Policy原型。但要让它在实际机器人任务中表现出色，还需要一些进阶技巧。
技巧一：动作序列预测与历史信息 上面的例子是单步动作预测。但在现实中，机器人控制往往需要一个动作序列。Diffusion Policy的一个强大之处在于可以直接生成未来多步的动作序列。你需要做的只是将模型输出的动作维度从action_dim扩展到action_horizon * action_dim，并在训练时使用动作序列数据。同时，将当前状态与过去若干步的历史状态拼接起来作为输入，能让策略更好地理解动态。
# 伪代码示例：扩展模型以处理动作序列
class SequentialDenoiser(nn.Module):
def __init__(self, state_history_dim, action_horizon, action_dim, hidden_dim):
super().__init__()
# 输入：历史状态 + 带噪声的未来动作序列
self.input_dim = state_history_dim + (action_horizon * action_dim)
self.output_dim = action_horizon * action_dim # 预测整个序列的噪声
# ... 网络层定义 ...
技巧二：条件注入与U-Net结构 对于视觉任务，状态可能是高维的图像。这时，简单的全连接网络就不够用了。你需要使用U-Net结构，它能有效处理空间信息。关键点在于如何将状态条件（如图像特征）注入到U-Net的各个层中。通常采用交叉注意力（Cross-Attention） 或特征拼接（Feature Concatenation） 的方式。
# 伪代码：在U-Net的某个下采样或上采样块中加入条件信息
class ConditionalUNetBlock(nn.Module):
def forward(self, x, timestep_embedding, state_condition):
# x: 特征图， state_condition: 编码后的状态特征
# 方式1：拼接
condition_expanded = state_condition.unsqueeze(-1).unsqueeze(-1).expand_as(x[:, :state_condition.size(1), ...])
x = torch.cat([x, condition_expanded], dim=1)
# 方式2：交叉注意力（更强大）
# query = x, key = state_condition, value = state_condition
# ... 实现注意力机制 ...
return processed_x
技巧三：调度器选择与采样加速 DDPM调度器稳定但采样慢。你可以尝试更快的采样器，如DDIMScheduler或DPMSolverMultistepScheduler，它们能用更少的步数（如20-50步）达到与DDPM上百步相近的效果，极大提升推理速度。
from diffusers import DDIMScheduler
# 切换到DDIM调度器进行快速采样
fast_scheduler = DDIMScheduler.from_config(scheduler.config) # 从DDPM配置继承
fast_scheduler.set_timesteps(num_inference_steps=20) # 只需20步
# 然后使用 generate_action 函数，将 scheduler 参数替换为 fast_scheduler
技巧四：处理真实世界的数据与仿真到实物的迁移 在真实机器人上应用Diffusion Policy，数据是关键。确保你的数据覆盖了任务的各种边缘情况。仿真环境是绝佳的试验场，但仿真到实物（Sim2Real）的差距不容忽视。可以考虑以下策略：
- 数据增强：对状态（如图像）进行随机裁剪、颜色抖动、添加噪声。
- 领域随机化：在仿真中随机化纹理、光照、物理参数，让策略学会关注更本质的特征。
- 在策略中引入不确定性建模：让模型除了预测动作，还预测动作的置信度或方差，在不确定时采取更保守的动作。
我最近在一个机械臂推物体到目标点的任务中应用了Diffusion Policy。最初直接用仿真数据训练的策略在真实机器人上表现不佳，推的力度不是太大就是太小。后来，我在仿真数据中加入了大量的电机噪声模拟和桌面摩擦系数随机化，并且将动作预测从单步改为一个包含速度和力控信息的短序列。重新训练后，策略在真实世界中的鲁棒性得到了显著提升。最关键的一步是，在推理时，我将num_inference_steps从训练时的1000步降到了50步（使用DDIM），在保证实时性的同时，动作的平滑度依然很好。这个调优过程花了一些时间，但看到机器人流畅地完成任务时，感觉一切都值了。
Agent 垂直技术社区，欢迎活跃、内容共建。
更多推荐
13
0- 0
已为社区贡献1条内容
所有评论(0)