SG Bot

SG-机器人

这是论文《SG-Bot：通过场景图上从粗到细的机器人想象力进行对象重新排列》(ICRA 2024) 的最小实现，arxiv 。

设置

 conda env create - f environment . yml
cd extension
python setup . py install

另请安装 Pytorch。我们使用 Pytorch 1.12.1 和 CUDA 11.6 对其进行测试。

数据集

请参阅此页面下载论文中使用的数据和更多信息。

模块说明

形状自动编码器

我们设置了两个形状自动编码器，称为AtlasNet和AtlastNet2 。 AtlasNet在规范坐标下使用完整形状进行训练，而AtlasNet2在相机框架下进行训练，它提供了目标场景图的先验形状以指导想象。我们还提供在这里下载的训练模型：训练有素的 AtlasNet 和训练有素的 AtlasNet2。

• 用于生成形状

  • 训练AtlasNet 。需要调整--batchSize ， --nepoch以使训练最优。

   cd AtlasNet
  python training/train_AE_AtlasNet.py
  

  • 推理点云[可选]：运行AtlasNet/inference/run_AE_AtlasNet.py 。结果将把生成的点存储在AtlasNet/log/atlasnet_separate_cultery/network下。

  • 获取用于训练 Graph-to-3D 的点特征：运行AtlasNet/inference/create_features_gt.py ，特征存储在objs_features_gt_atlasnet_separate_cultery.json中。 json 文件中的键是对象的名称，例如“cup_1”，值是潜在特征（128 维）。

• 用于产生形状先验

  • 将初始场景中的部分点存储在相机框架下：这旨在训练AtlasNet2 。这些文件可以从这里下载：partial_pcs。您还可以修改文件路径并运行AtlasNet2/auxiliary/generate_partial_pc_for_object.py 。最终输出作为 pickle 文件存储在AtlasNet2/partial_pc_data下。
  • 拆分trainval集： AtlasNet2/auxiliary/generate_partial_pc_for_object.py中的函数generate_train_sample将AtlasNet2/partial_pc_data拆分为train（90%）和test（10%）。文件名存储为AtlasNet2/partial_pc_data_splits.json
  • 训练AtlasNet2 ：过程与AtlasNet相同。

场景生成器

我们基于 Graph-to-3D（GCN-VAE 架构）构建了场景生成器。与原始的 Graph-to-3D 不同，我们利用形状感知场景图来使生成的形状与初始场景中观察到的形状对齐。我们在这里提供经过训练的模型：trained graph_to_3d。

如果你想重新训练网络， --batchSize ， --nepoch ， --exp需要设置适当的数字。

 cd graphto3d
python scripts/train_vaegan.py

更多详细信息可以在原始存储库中找到。

模拟

有两种模式—— robot和oracle 。 robot模式支持机械臂根据想象操纵物体。该模式需要一个抓取姿势预测网络，我们使用Contact-GraspNet。这需要下载tensorflow。

 pip install tensorflow - estimator == 2.7 .0 tensorflow - gpu == 2.7 .0

检查点可以从原始存储库或此处下载。下载检查点后，将它们移至./contact_graspnet 。

oracle模式不需要代理，而是直接将物体摆成相对姿势。要使脚本正常工作，可以修改里面的变量mode ，然后运行：

 python sgbot_pybullet.py

论文中的结果是在oracle模式下的。我们直接使用预定义的场景图作为目标。

真实世界试用 [TODO]

我们提供了录制的rosbag来演示性能。要进行此试验，需要从此处下载 MaskRCNN 检查点。需要安装额外的要求。