PyTorch实现基于Lie群的扩散模型——6D抓取姿态生成

1. 项目介绍

本项目是一个基于PyTorch的开源项目，实现了在Lie群上的扩散模型，用于6D抓取姿态的生成。该模型能够学习抓取分布，并将其作为抓取和运动优化问题的代价函数。项目的核心是SE(3)-DiffusionFields模型，该模型通过扩散过程学习得到平滑的代价函数，进而用于抓取和运动的联合优化。

2. 项目快速启动

环境搭建

首先，创建一个conda环境，并激活它：

conda create -f environment.yml
conda activate se3dif_env

然后，安装必要的依赖库：

pip install -e .

数据准备

在se3dif/utils/directory_utils.py中设置数据根目录，确保数据集和训练模型的位置正确。

数据集的准备指示可以在scripts/create_data中找到。如果需要，可以从上级目录下载已处理的数据。

模型训练

训练一个基于点云的模型：

python scripts/train/train_pointcloud_6d_grasp_diffusion.py

或者，训练一个基于部分点云的模型：

python scripts/train/train_partial_pointcloud_6d_grasp_diffusion.py

抓取姿态采样

给定整个物体的点云，生成抓取姿态：

python scripts/sample/generate_pointcloud_6d_grasp_poses.py --n_grasps 10 --obj_id 0 --obj_class 'ScrewDriver'

对于特定物体（如杯子）的模型，可以生成专门的抓取姿态：

python scripts/sample/generate_pointcloud_6d_grasp_poses.py --n_grasps 10 --obj_id 10 --obj_class 'Mug' --model 'grasp_dif_mugs'

给定部分点云，生成抓取姿态：

python scripts/sample/generate_partial_pointcloud_6d_grasp_poses.py --n_grasps 10 --obj_id 12 --obj_class 'Mug'

3. 应用案例和最佳实践

• 数据增强：为了提高模型的泛化能力，可以在训练前对数据集进行增强，如旋转、缩放和裁剪点云。
• 模型评估：在模拟环境中（如Isaac Gym）评估训练好的模型，并计算抓取姿态的地球移动距离（EMD）来衡量模型性能。

4. 典型生态项目

• 抓取数据集：如Acronym和ShapeNet，这些数据集提供了丰富的抓取和物体形状数据，用于训练和评估模型。
• 机器人仿真工具：如Isaac Gym，它是一个用于机器人研究的快速、灵活的仿真环境，可以用来模拟和测试抓取策略。

以上是本项目的基本介绍和快速启动指南，希望对您有所帮助！