StarAI——Lerobot机械臂算法复现之DiffusionPolicy篇

StarAI——Lerobot机械臂算法复现之DiffusionPolicy篇

理论部分

实践部分

0.系统要求

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

1.录制数据集

数据质量要求

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

2.改进算法

如果按照lerobot封装的源代码去训练的话，效果会很差很差很差很差！！！！网络上不少开发者对lerobot封装的dp策略也是嗤之以鼻，太难用了。通过我和dp斗智斗勇了一个星期之后，我终于知道为什么效果这么差了！

在diffusion的参数配置文件中有一个这样的参数：

crop_shape: tuple[int, int] | None = (84, 84)

crop？？？？裁剪？？？这是什么操作，我们继续追溯，在模型架构文件中，DiffusionRgbEncoder类中找到了这个参数的用法

这什么玩意？？？如如果启动了crop，就会对图像进行随机/中心裁减。难怪dp训练出来的效果这么差，我录制的数据集设置的图像大小为640*480，按照它这个裁减的方法（84，84），99%的视觉信息都丢失了！！

不要crop，不要裁减。视觉信息在运动控制中是非常重要的！！！

于是我修改了代码，从裁减（crop）修改到了缩放（reshape）,具体操作如下：

src/lerobot/policies/smolvla/configuration_smolvla.py：

修改变量，注释crop_is_random

    # Architecture / modeling.
    # Vision backbone.
    vision_backbone: str = "resnet18"
    resize_shape: tuple[int, int] | None = (224, 224)
    #crop_is_random: bool = False

修改验证函数：

    def validate_features(self) -> None:
        if len(self.image_features) == 0 and self.env_state_feature is None:
            raise ValueError("You must provide at least one image or the environment state among the inputs.")

        if self.resize_shape is not None:
            # 对于resize_shape，不需要检查是否在原图范围内，因为缩放可以放大或缩小
            pass  # 移除裁剪范围的检查，缩放可以处理任意尺寸

        # Check that all input images have the same shape.
        if len(self.image_features) > 0:
            first_image_key, first_image_ft = next(iter(self.image_features.items()))
            for key, image_ft in self.image_features.items():
                if image_ft.shape != first_image_ft.shape:
                    raise ValueError(
                        f"`{key}` does not match `{first_image_key}`, but we expect all image shapes to match."
                    )

src/lerobot/policies/diffusion/modeling_diffusion.py：

找到DiffusionRgbEncoder类，这是视觉编码器的架构：

修改初始化函数：

    def __init__(self, config: DiffusionConfig):
        super().__init__()
        # Set up optional preprocessing.
        if config.resize_shape is not None:
            self.do_resize = True
            # Use resize transform
            self.resize = torchvision.transforms.Resize(config.resize_shape)
        else:
            self.do_resize = False

修改与crop有关参数：

        images_shape = next(iter(config.image_features.values())).shape
        dummy_shape_h_w = config.resize_shape if config.resize_shape is not None else images_shape[1:]
        dummy_shape = (1, images_shape[0], *dummy_shape_h_w)

怎么会犯如此低级的错误？？带着疑惑，我打开了Lerobot主枝的代码一探究竟。一看，果然新版的lerobot代码已经改了，并且思路和主播的一样

但是这也不怪seeed官方，毕竟落后了97个提交，跟最新版总会有些差错（偷笑）

3.训练

accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=/mnt/sdb/dataset/zzh/DataSet/pick_and_place \
  --policy.type=diffusion \
  --policy.device=cuda \
  --policy.repo_id=starai/my_policy \
  --batch_size=32 \
  --optimizer.lr=1e-04 \
  --num_workers=8 \
  --output_dir=outputs/train/diffusion_viola_low \
  --job_name=diffusion_viola_low \
  --wandb.enable=false \
  --steps=20000 \
  --save_checkpoint=true \
  --save_freq=10000 2>&1 | tee training_dplow_log.txt

3.1参数讲解

3.1.1核心参数

3.1.2Diffusion Policy 特定参数

3.1.3网络架构参数

3.1.4训练参数

4.评估

将crop修改为reshape之后，整个模型直接变聪明了，任务成功率嘎嘎提升，换成谁都要来一句：DiffusionPolicy真香！！！

lerobot-record  \
  --robot.type=lerobot_robot_viola \
  --robot.port=/dev/ttyUSB1 \
  --robot.cameras="{ up: {type: opencv, index_or_path: /dev/video6, width: 640, height: 480, fps: 30, fourcc: "MJPG"},front: {type: opencv, index_or_path: /dev/video8, width: 640, height: 480, fps: 30, fourcc: "MJPG"}}" \
  --robot.id=my_awesome_staraiviola_arm \
  --display_data=true \
  --dataset.episode_time_s=100 \
  --dataset.num_episodes=10 \
  --dataset.repo_id=starai/eval_dp \
  --dataset.single_task="Put the yellow square into the drawer." \
  --dataset.push_to_hub=False \
  --policy.path=/home/zzh/lerobot/outputs/train/dp/pretrained_model_resize

常见问题 (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 或批处理
预计算：提前计算部分结果