SG Bot

SG-Bot

Esta es una implementación mínima del artículo SG-Bot: Reordenamiento de objetos mediante imaginación robótica gruesa a fina en gráficos de escena (ICRA 2024), arxiv .

Configuración

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

Instale también Pytorch. Lo probamos con Pytorch 1.12.1 con CUDA 11.6.

Conjunto de datos

Consulte esta página para descargar los datos utilizados en el artículo y obtener más información.

Instrucción del módulo

Codificadores automáticos de formas

Configuramos dos codificadores automáticos de formas llamados AtlasNet y AtlastNet2 . AtlasNet se entrena con formas completas bajo coordenadas canónicas, mientras que AtlasNet2 se entrena bajo el marco de la cámara, que proporciona formas anteriores al gráfico de la escena objetivo para guiar la imaginación. También proporcionamos modelos entrenados descargados aquí: AtlasNet entrenado y AtlasNet2 entrenado.

• Para generar formas

  • Entrena AtlasNet . Es necesario ajustar --batchSize ， --nepoch para que el entrenamiento sea óptimo.

   cd AtlasNet
  python training/train_AE_AtlasNet.py
  

  • Nubes de puntos de inferencia [opcional]: ejecute AtlasNet/inference/run_AE_AtlasNet.py . Los resultados almacenarían los puntos generados en AtlasNet/log/atlasnet_separate_cultery/network .

  • Obtenga una característica de puntos para entrenar Graph-to-3D: ejecute AtlasNet/inference/create_features_gt.py y las características se almacenarán en objs_features_gt_atlasnet_separate_cultery.json . Las claves en el archivo json son el nombre de los objetos, por ejemplo, "cup_1", y los valores son las características latentes (128 dimensiones).

• Para producir formas previas

  • Almacene puntos parciales en las escenas iniciales debajo del marco de la cámara: esto tiene como objetivo entrenar AtlasNet2 . Los archivos se pueden descargar desde aquí: parcial_pcs. También puede modificar la ruta del archivo y ejecutar AtlasNet2/auxiliary/generate_partial_pc_for_object.py . El resultado final se almacena como archivos pickle en AtlasNet2/partial_pc_data .
  • Divida el conjunto trainval: la función generate_train_sample en AtlasNet2/auxiliary/generate_partial_pc_for_object.py divide AtlasNet2/partial_pc_data en tren (90%) y prueba (10%). Los nombres de los archivos se almacenan como AtlasNet2/partial_pc_data_splits.json
  • Entrenar AtlasNet2 : El procedimiento es el mismo que AtlasNet .

Generador de escenas

Construimos el generador de escenas basado en Graph-to-3D, una arquitectura GCN-VAE. A diferencia del Graph-to-3D original, aprovechamos un gráfico de escena con reconocimiento de formas para alinear las formas generadas con las formas observadas en la escena inicial. Proporcionamos el modelo entrenado disponible aquí: Graph_to_3d entrenado.

Si desea volver a entrenar la red, --batchSize , --nepoch , --exp debe configurarse con los números adecuados.

 cd graphto3d
python scripts/train_vaegan.py

Se pueden encontrar más detalles en el repositorio original.

Simulación

Hay dos modos: robot y oracle . El modo robot admite un brazo robótico que manipula los objetos según la imaginación. Este modo necesita una red de predicción de postura de agarre, que utilizamos Contact-GraspNet. Esto necesita descargar tensorflow.

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

Los puntos de control se pueden descargar desde el repositorio original o aquí. Después de descargar los puntos de control, muévalos a ./contact_graspnet .

El modo oracle no necesita un agente, simplemente coloca los objetos directamente en poses relativas. Para que el script funcione, se puede modificar el mode variable interno y luego ejecutar:

 python sgbot_pybullet.py

Los resultados en el documento están en modo oracle . Utilizamos directamente el gráfico de escena predefinido como objetivo.

Prueba en el mundo real [TODO]

Proporcionamos un rosbag grabado para demostrar la actuación. Para realizar esta prueba, es necesario descargar el punto de control MaskRCNN desde aquí. Es necesario instalar requisitos adicionales.