SG Bot

SG-Bot

Esta é uma implementação mínima do artigo SG-Bot: Object Rearrangement via Coarse-to-Fine Robotic Imagination on Scene Graphs (ICRA 2024), arxiv .

Configurar

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

Instale também o Pytorch. Testamos com Pytorch 1.12.1 com CUDA 11.6.

Conjunto de dados

Consulte esta página para baixar os dados utilizados no artigo e obter mais informações.

Instrução do Módulo

Autoencoders de forma

Configuramos dois autoencoders de formato chamados AtlasNet e AtlastNet2 . AtlasNet é treinado com formas completas sob coordenadas canônicas, enquanto AtlasNet2 é treinado sob o quadro da câmera, que fornece forma antes do gráfico da cena do objetivo para guiar a imaginação. Também fornecemos modelos treinados baixados aqui: AtlasNet treinado e AtlasNet2 treinado.

• Para gerar formas

  • Treine AtlasNet . Precisa ajustar --batchSize ， --nepoch para tornar o treinamento ideal.

   cd AtlasNet
  python training/train_AE_AtlasNet.py
  

  • Nuvens de pontos de inferência [opcional]: execute AtlasNet/inference/run_AE_AtlasNet.py . Os resultados armazenariam os pontos gerados em AtlasNet/log/atlasnet_separate_cultery/network .

  • Obtenha recurso de ponto para treinamento de gráfico para 3D: execute AtlasNet/inference/create_features_gt.py e os recursos são armazenados em objs_features_gt_atlasnet_separate_cultery.json . As chaves no arquivo json são os nomes dos objetos, por exemplo, "cup_1", e os valores são os recursos latentes (128 dimensões).

• Para produzir anteriores de forma

  • Armazene pontos parciais nas cenas iniciais sob o quadro da câmera: Isto visa treinar AtlasNet2 . Os arquivos podem ser baixados aqui: parcial_pcs. Você também pode modificar o caminho do arquivo e executar AtlasNet2/auxiliary/generate_partial_pc_for_object.py . A saída final é armazenada como arquivos pickle em AtlasNet2/partial_pc_data .
  • Divida o conjunto trainval: A função generate_train_sample em AtlasNet2/auxiliary/generate_partial_pc_for_object.py divide AtlasNet2/partial_pc_data em train (90%) e teste (10%). Os nomes dos arquivos são armazenados como AtlasNet2/partial_pc_data_splits.json
  • Treinar AtlasNet2 : O procedimento é o mesmo do AtlasNet .

Gerador de cena

Construímos o gerador de cena baseado em Graph-to-3D, uma arquitetura GCN-VAE. Diferente do Graph-to-3D original, aproveitamos um gráfico de cena com reconhecimento de forma para alinhar as formas geradas com as formas observadas na cena inicial. Fornecemos o modelo treinado disponível aqui: graph_to_3d treinado.

Se você deseja treinar novamente a rede, --batchSize , --nepoch , --exp precisa ser definido com números adequados.

 cd graphto3d
python scripts/train_vaegan.py

Mais detalhes podem ser encontrados no repositório original.

Simulação

Existem dois modos - robot e oracle . O modo robot suporta um braço robótico manipulando os objetos de acordo com a imaginação. Este modo precisa de uma rede de previsão de pose de preensão, que usamos Contact-GraspNet. Isso precisa do download do tensorflow.

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

Os pontos de verificação podem ser baixados do repositório original ou aqui. Após baixar os pontos de verificação, mova-os para ./contact_graspnet .

O modo oracle não precisa de um agente, apenas coloca diretamente os objetos em poses relativas. Para fazer o script funcionar, pode-se modificar a variável mode interna e então executar:

 python sgbot_pybullet.py

Os resultados no artigo estão no modo oracle . Usamos diretamente o gráfico de cena predefinido como objetivo.

Teste do mundo real [TODO]

Fornecemos um rosbag gravado para demonstrar o desempenho. Para conduzir este teste, o ponto de verificação MaskRCNN precisa ser baixado aqui. Requisitos adicionais precisam ser instalados.