跳到主要内容

Diffusion Policy 模型微调

概述​

Diffusion Policy 是一种基于扩散模型的视觉运动策略学习方法，将扩散模型的生成能力应用于机器人控制领域。该方法通过学习动作分布的扩散过程，能够生成多样化且高质量的机器人动作序列，在复杂的机器人操作任务中表现出色。

核心特点​

  • 扩散生成：使用扩散模型生成连续的动作序列
  • 多模态动作：能够处理具有多种解决方案的任务
  • 高质量输出：生成平滑、自然的机器人动作
  • 鲁棒性强：对噪声和扰动具有良好的鲁棒性
  • 表达能力强：能够学习复杂的动作分布

先决条件​

系统要求​

  • 操作系统：Linux（推荐 Ubuntu 20.04+）或 macOS
  • Python 版本：3.8+
  • GPU：NVIDIA GPU（推荐 RTX 3080 或更高），至少 10GB 显存
  • 内存：至少 32GB RAM
  • 存储空间：至少 50GB 可用空间

环境准备​

1. 安装 LeRobot​

# 克隆 LeRobot 仓库
git clone https://github.com/huggingface/lerobot.git
cd lerobot

# 创建虚拟环境
conda create -n lerobot python=3.10
conda activate lerobot

# 安装依赖
pip install -e .

2. 安装 Diffusion Policy 特定依赖​

# 安装扩散模型相关依赖
pip install diffusers
pip install accelerate
pip install transformers
pip install einops
pip install wandb

# 安装数值计算库
pip install scipy
pip install scikit-learn

# 登录 Weights & Biases（可选）
wandb login

Diffusion Policy 架构​

核心组件​

  1. 视觉编码器：提取图像特征
  2. 状态编码器：处理机器人状态信息
  3. 条件编码器：融合视觉和状态信息
  4. 扩散网络：学习动作分布的扩散过程
  5. 噪声调度器：控制扩散过程的噪声水平

扩散过程​

  1. 前向过程：逐步向动作序列添加噪声
  2. 反向过程：从噪声中逐步恢复动作序列
  3. 条件生成：基于观察条件生成动作
  4. 采样策略：使用 DDPM 或 DDIM 采样

数据准备​

LeRobot 格式数据​

Diffusion Policy 需要使用 LeRobot 格式的数据集：

your_dataset/
├── data/
│   ├── chunk-001/
│   │   ├── observation.images.cam_high.png
│   │   ├── observation.images.cam_low.png
│   │   ├── observation.state.npy
│   │   ├── action.npy
│   │   └── ...
│   └── chunk-002/
│       └── ...
├── meta.json
├── stats.safetensors
└── videos/
    ├── episode_000000.mp4
    └── ...

数据质量要求​

  • 最少 100 个 episode 用于基本训练
  • 推荐 500+ episode 以获得最佳效果
  • 动作序列应该平滑连续
  • 包含多样化的任务场景
  • 高质量的视觉观察数据

微调训练​

基本训练命令​

# 设置环境变量
export HF_USER="your-huggingface-username"
export CUDA_VISIBLE_DEVICES=0

# 启动 Diffusion Policy 训练
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 64 \
  --steps 100000 \
  --output_dir outputs/train/diffusion_policy_finetuned \
  --job_name diffusion_policy_finetuning \
  --policy.device cuda \
  --policy.horizon 16 \
  --policy.n_action_steps 8 \
  --policy.n_obs_steps 2 \
  --policy.num_inference_steps 100 \
  --policy.optimizer_lr 1e-4 \
  --policy.optimizer_weight_decay 1e-6 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --save_freq 10000 \
  --wandb.enable true

高级训练配置​

多步预测配置​

# 针对长序列预测的配置
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 32 \
  --steps 150000 \
  --output_dir outputs/train/diffusion_policy_long_horizon \
  --job_name diffusion_policy_long_horizon \
  --policy.device cuda \
  --policy.horizon 32 \
  --policy.n_action_steps 16 \
  --policy.n_obs_steps 4 \
  --policy.num_inference_steps 100 \
  --policy.beta_schedule squaredcos_cap_v2 \
  --policy.clip_sample true \
  --policy.prediction_type epsilon \
  --policy.optimizer_lr 1e-4 \
  --policy.scheduler_name cosine \
  --policy.scheduler_warmup_steps 5000 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

内存优化配置​

# 针对显存较小的 GPU
lerobot-train \
  --policy.type diffusion \
  --policy.pretrained_path lerobot/diffusion_policy \
  --dataset.repo_id ${HF_USER}/your_dataset \
  --batch_size 16 \
  --steps 200000 \
  --output_dir outputs/train/diffusion_policy_memory_opt \
  --job_name diffusion_policy_memory_optimized \
  --policy.device cuda \
  --policy.horizon 16 \
  --policy.n_action_steps 8 \
  --policy.num_inference_steps 50 \
  --policy.optimizer_lr 5e-5 \
  --policy.use_amp true \
  --num_workers 2 \
  --policy.push_to_hub false \
  --save_checkpoint true \
  --wandb.enable true

参数详细说明​

核心参数​

参数含义推荐值说明
--policy.type策略类型diffusionDiffusion Policy 模型类型
--policy.pretrained_path预训练模型路径lerobot/diffusion_policyLeRobot 官方模型（可选）
--dataset.repo_id数据集仓库 ID${HF_USER}/dataset你的 HuggingFace 数据集
--batch_size批处理大小64根据显存调整，RTX 3080 推荐 32-64
--steps训练步数100000扩散模型通常需要更多训练步数
--output_dir输出目录outputs/train/diffusion_policy_finetuned模型保存路径
--job_name任务名称diffusion_policy_finetuning用于日志和实验跟踪（可选）

Diffusion Policy 特定参数​

参数含义推荐值说明
--policy.horizon预测时间范围16预测的动作序列长度
--policy.n_action_steps执行动作步数8每次执行的动作数量
--policy.n_obs_steps观察步数2历史观察的数量
--policy.num_inference_steps推理步数100扩散采样的步数（训练时不生效）
--policy.beta_schedule噪声调度squaredcos_cap_v2噪声添加的调度策略
--policy.clip_sample样本裁剪true是否裁剪生成的样本
--policy.clip_sample_range裁剪范围1.0样本裁剪的范围
--policy.prediction_type预测类型epsilon预测噪声或样本
--policy.num_train_timesteps训练时间步100前向扩散的步数

网络架构参数​

参数含义推荐值说明
--policy.vision_backbone视觉骨干网络resnet18图像特征提取网络
--policy.crop_shape图像裁剪尺寸84 84输入图像的裁剪尺寸
--policy.crop_is_random随机裁剪true训练时是否随机裁剪
--policy.use_group_norm使用组归一化true替换批归一化
--policy.spatial_softmax_num_keypoints空间软最大关键点数32空间软最大层的关键点数
--policy.down_dims下采样维度512 1024 2048U-Net 下采样路径的维度
--policy.kernel_size卷积核大小51D 卷积的核大小
--policy.n_groups组归一化组数8GroupNorm 的组数
--policy.diffusion_step_embed_dim步骤嵌入维度128扩散步骤的嵌入维度

训练参数​

参数含义推荐值说明
--policy.optimizer_lr学习率1e-4扩散模型推荐的学习率
--policy.optimizer_weight_decay权重衰减1e-6正则化参数
--policy.optimizer_betasAdam优化器beta0.95 0.999Adam 优化器的 beta 参数
--policy.optimizer_epsAdam epsilon1e-8数值稳定性参数
--policy.scheduler_name学习率调度器cosine余弦退火调度
--policy.scheduler_warmup_steps预热步数500学习率预热
--policy.use_amp混合精度true节省显存
--num_workers数据加载线程4根据CPU核心数调整
--policy.push_to_hub推送到Hubfalse是否上传模型到HuggingFace（需要repo_id）
--save_checkpoint保存检查点true是否保存训练检查点
--save_freq保存频率10000检查点保存间隔

训练监控和调试​

Weights & Biases 集成​

# 详细的 W&B 配置
lerobot-train \
  --policy.type diffusion \
  --dataset.repo_id your-name/your-dataset \
  --batch_size 64 \
  --steps 100000 \
  --policy.push_to_hub false \
  --wandb.enable true \
  --wandb.project diffusion_policy_experiments \
  --wandb.notes "Diffusion Policy training with long horizon" \
  # ... 其他参数

关键指标监控​

训练过程中需要关注的指标：

  • Diffusion Loss：扩散模型的总体损失
  • MSE Loss：均方误差损失
  • Learning Rate：学习率变化曲线
  • Gradient Norm：梯度范数
  • Inference Time：推理时间
  • Sample Quality：生成样本的质量

训练可视化​

# visualization.py
import torch
import matplotlib.pyplot as plt
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy

def visualize_diffusion_process(model_path, observation):
    # 加载模型
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    # 生成动作序列的扩散过程
    with torch.no_grad():
        # 初始噪声
        noise = torch.randn(1, policy.horizon, policy.action_dim, device="cuda")
        
        # 扩散采样过程
        actions_sequence = []
        for t in range(policy.num_inference_steps):
            # 预测噪声
            noise_pred = policy.unet(noise, t, observation)
            
            # 更新样本
            noise = policy.scheduler.step(noise_pred, t, noise).prev_sample
            
            # 保存中间结果
            if t % 10 == 0:
                actions_sequence.append(noise.cpu().numpy())
        
        final_actions = noise.cpu().numpy()
    
    # 可视化扩散过程
    fig, axes = plt.subplots(2, 3, figsize=(15, 10))
    
    for i, actions in enumerate(actions_sequence[:6]):
        ax = axes[i//3, i%3]
        ax.plot(actions[0, :, 0], label='Action Dim 0')
        ax.plot(actions[0, :, 1], label='Action Dim 1')
        ax.set_title(f'Diffusion Step {i*10}')
        ax.legend()
    
    plt.tight_layout()
    plt.savefig('diffusion_process.png')
    plt.show()
    
    return final_actions

if __name__ == "__main__":
    model_path = "outputs/train/diffusion_policy_finetuned/checkpoints/last"
    
    # 模拟观察
    observation = {
        "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 7, device="cuda")
    }
    
    actions = visualize_diffusion_process(model_path, observation)
    print(f"Generated actions shape: {actions.shape}")

模型评估​

离线评估​

# offline_evaluation.py
import torch
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from lerobot.datasets.lerobot_dataset import LeRobotDataset

def evaluate_diffusion_policy(model_path, dataset_path):
    # 加载模型
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    # 加载测试数据集
    dataset = LeRobotDataset(dataset_path, split="test")
    
    total_mse_loss = 0
    total_mae_loss = 0
    num_samples = 0
    
    with torch.no_grad():
        for batch in dataset:
            # 模型预测
            prediction = policy(batch)
            
            # 计算损失
            target_actions = batch['action']
            predicted_actions = prediction['action']
            
            mse_loss = torch.mean((predicted_actions - target_actions) ** 2)
            mae_loss = torch.mean(torch.abs(predicted_actions - target_actions))
            
            total_mse_loss += mse_loss.item()
            total_mae_loss += mae_loss.item()
            num_samples += 1
    
    avg_mse_loss = total_mse_loss / num_samples
    avg_mae_loss = total_mae_loss / num_samples
    
    print(f"Average MSE Loss: {avg_mse_loss:.4f}")
    print(f"Average MAE Loss: {avg_mae_loss:.4f}")
    
    return avg_mse_loss, avg_mae_loss

def evaluate_action_diversity(model_path, observation, num_samples=10):
    # 评估动作多样性
    policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
    policy.eval()
    
    actions_list = []
    
    with torch.no_grad():
        for _ in range(num_samples):
            prediction = policy(observation)
            actions_list.append(prediction['action'].cpu().numpy())
    
    actions_array = np.array(actions_list)  # [num_samples, horizon, action_dim]
    
    # 计算动作多样性指标
    action_std = np.std(actions_array, axis=0)  # [horizon, action_dim]
    avg_std = np.mean(action_std)
    
    print(f"Average action standard deviation: {avg_std:.4f}")
    
    return avg_std, actions_array

if __name__ == "__main__":
    model_path = "outputs/train/diffusion_policy_finetuned/checkpoints/last"
    dataset_path = "path/to/your/test/dataset"
    
    # 离线评估
    evaluate_diffusion_policy(model_path, dataset_path)
    
    # 多样性评估
    observation = {
        "observation.images.cam_high": torch.randn(1, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 7, device="cuda")
    }
    
    evaluate_action_diversity(model_path, observation)

在线评估（机器人环境）​

# robot_evaluation.py
import torch
import numpy as np
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy

class DiffusionPolicyController:
    def __init__(self, model_path, num_inference_steps=50):
        self.policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        self.num_inference_steps = num_inference_steps
        self.action_queue = []
        self.current_obs_history = []
        
    def get_action(self, observations):
        # 更新观察历史
        self.current_obs_history.append(observations)
        if len(self.current_obs_history) > self.policy.n_obs_steps:
            self.current_obs_history.pop(0)
        
        # 如果动作队列为空或需要重新规划，生成新的动作序列
        if len(self.action_queue) == 0 or self.should_replan():
            with torch.no_grad():
                # 构建输入
                batch = self.prepare_observation_batch()
                
                # 设置推理步数
                self.policy.scheduler.set_timesteps(self.num_inference_steps)
                
                # 生成动作序列
                prediction = self.policy(batch)
                actions = prediction['action'].cpu().numpy()[0]  # [horizon, action_dim]
                
                # 更新动作队列
                self.action_queue = list(actions[:self.policy.n_action_steps])
        
        # 返回下一个动作
        return self.action_queue.pop(0)
    
    def should_replan(self):
        # 简单的重新规划策略：当动作队列剩余不足一半时重新规划
        return len(self.action_queue) < self.policy.n_action_steps // 2
    
    def prepare_observation_batch(self):
        batch = {}
        
        # 处理图像观察
        if "observation.images.cam_high" in self.current_obs_history[-1]:
            images = []
            for obs in self.current_obs_history:
                image = obs["observation.images.cam_high"]
                image_tensor = self.preprocess_image(image)
                images.append(image_tensor)
            
            # 如果历史不足，重复最后一个观察
            while len(images) < self.policy.n_obs_steps:
                images.insert(0, images[0])
            
            batch["observation.images.cam_high"] = torch.stack(images).unsqueeze(0)
        
        # 处理状态观察
        if "observation.state" in self.current_obs_history[-1]:
            states = []
            for obs in self.current_obs_history:
                state = torch.tensor(obs["observation.state"], dtype=torch.float32)
                states.append(state)
            
            # 如果历史不足，重复最后一个状态
            while len(states) < self.policy.n_obs_steps:
                states.insert(0, states[0])
            
            batch["observation.state"] = torch.stack(states).unsqueeze(0)
        
        return batch
    
    def preprocess_image(self, image):
        # 图像预处理逻辑
        image_tensor = torch.tensor(image).permute(2, 0, 1).float() / 255.0
        return image_tensor

# 使用示例
if __name__ == "__main__":
    controller = DiffusionPolicyController(
        model_path="outputs/train/diffusion_policy_finetuned/checkpoints/last",
        num_inference_steps=50
    )
    
    # 模拟机器人控制循环
    for step in range(100):
        # 获取当前观察
        observations = {
            "observation.images.cam_high": np.random.randint(0, 255, (224, 224, 3)),
            "observation.state": np.random.randn(7)
        }
        
        # 获取动作
        action = controller.get_action(observations)
        
        # 执行动作
        print(f"Step {step}: Action = {action}")
        
        # 这里应该将动作发送给实际的机器人
        # robot.execute_action(action)

部署和优化​

推理加速​

# fast_inference.py
import torch
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from diffusers import DDIMScheduler

class FastDiffusionInference:
    def __init__(self, model_path, num_inference_steps=10):
        self.policy = DiffusionPolicy.from_pretrained(model_path, device="cuda")
        self.policy.eval()
        
        # 使用 DDIM 调度器进行快速采样
        self.policy.scheduler = DDIMScheduler.from_config(self.policy.scheduler.config)
        self.num_inference_steps = num_inference_steps
        
        # 预热模型
        self.warmup()
    
    def warmup(self):
        # 使用虚拟数据预热模型
        dummy_batch = {
            "observation.images.cam_high": torch.randn(1, 2, 3, 224, 224, device="cuda"),
            "observation.state": torch.randn(1, 2, 7, device="cuda")
        }
        
        with torch.no_grad():
            for _ in range(5):
                _ = self.predict(dummy_batch)
    
    @torch.no_grad()
    def predict(self, observations):
        # 设置推理步数
        self.policy.scheduler.set_timesteps(self.num_inference_steps)
        
        # 快速推理
        prediction = self.policy(observations)
        return prediction['action'].cpu().numpy()

if __name__ == "__main__":
    fast_inference = FastDiffusionInference(
        "outputs/train/diffusion_policy_finetuned/checkpoints/last",
        num_inference_steps=10
    )
    
    # 测试推理速度
    import time
    
    observations = {
        "observation.images.cam_high": torch.randn(1, 2, 3, 224, 224, device="cuda"),
        "observation.state": torch.randn(1, 2, 7, device="cuda")
    }
    
    start_time = time.time()
    for _ in range(100):
        action = fast_inference.predict(observations)
    end_time = time.time()
    
    avg_inference_time = (end_time - start_time) / 100
    print(f"Average inference time: {avg_inference_time:.4f} seconds")
    print(f"Inference frequency: {1/avg_inference_time:.2f} Hz")

最佳实践​

数据收集建议​

  1. 平滑动作：确保演示数据中的动作序列平滑连续
  2. 多样化场景：收集不同起始状态和目标的数据
  3. 高质量标注：确保动作标注的准确性
  4. 充足数据量：扩散模型通常需要更多数据

训练优化建议​

  1. 噪声调度：选择合适的噪声调度策略
  2. 推理步数：平衡质量和速度，选择合适的推理步数
  3. 学习率调度：使用余弦退火或阶梯式衰减
  4. 正则化：适当使用权重衰减

部署优化建议​

  1. 快速采样：使用 DDIM 或其他快速采样方法
  2. 模型压缩：使用知识蒸馏或量化技术
  3. 并行推理：利用 GPU 并行能力
  4. 缓存优化：缓存中间计算结果

常见问题 (FAQ)​

Q: Diffusion Policy 与其他策略学习方法相比有什么优势？​

A: Diffusion Policy 的主要优势包括：

  • 多模态生成：能够处理具有多种解决方案的任务
  • 高质量输出：生成平滑、自然的动作序列
  • 鲁棒性强：对噪声和扰动具有良好的鲁棒性
  • 表达能力强：能够学习复杂的动作分布

Q: 如何选择合适的推理步数？​

A: 推理步数的选择需要平衡质量和速度：

  • 高质量：100-1000 步，适合离线评估
  • 实时应用：10-50 步，适合在线控制
  • 快速原型：5-10 步，适合快速测试

Q: 训练需要多长时间？​

A: 训练时间取决于多个因素：

  • 数据集大小：500 episodes 约需 12-24 小时（RTX 3080）
  • 模型复杂度：更大的模型需要更长时间
  • 推理步数：更多步数增加训练时间
  • 收敛要求：通常需要 100000-200000 步

Q: 如何提升生成动作的质量？​

A: 提升动作质量的方法：

  • 增加推理步数：更多步数通常产生更好的结果
  • 优化噪声调度：选择合适的噪声添加策略
  • 数据质量：确保训练数据的高质量
  • 模型架构：使用更大或更深的网络
  • 正则化技术：适当的正则化防止过拟合

Q: 如何处理实时性要求？​

A: 满足实时性要求的方法：

  • 快速采样：使用 DDIM 或 DPM-Solver
  • 减少推理步数：在质量和速度间找到平衡
  • 模型蒸馏：训练更小的学生模型
  • 并行推理：利用多 GPU 或批处理
  • 预计算：提前计算部分结果

相关资源​

更新日志​

  • 2024-01: 初始版本发布
  • 2024-02: 添加快速采样支持
  • 2024-03: 优化内存使用和训练效率
  • 2024-04: 添加多样性评估和部署优化