haojieshi98 / gpat

[CoRL 2023] Rearrangement Planning for General Part Assembly

Rearrangement Planning for General Part Assembly

Yulong Li¹, Andy Zeng², Shuran Song¹

¹Columbia University, ²Google Deepmind

Project Page | arXiV

Catalog

• Environment
• Dataset
• Training and Testing
• Acknowledgements
• Citation

Environment

conda env create -f environment.yml
conda activate gpat
pip install -e .

Install other dependencies (chamfer and pointops):

cd utils/chamfer; python setup.py install
cd ../pointops; python setup.py install

Dataset

We create our dataset from PartNet. To recreate our data, first download PartNet v0 to dataset/partnet_raw. Also download PartNet meta data to dataset/partnet_dataset. Then run the following command:

python dataset/preprocess.py

For custom datasets, each data point should have the following files in its folder:

target.npy # (5000, 3), required, target point-cloud.
parts.npy # (K, 1000, 3), required, parts point-clouds.
target_100k.npy # (100000, 3), optional, a dense target point-cloud. Optionally include this for better assembly results.
target_100k_labels.npy # (10000), optional, indicates the nearest neighbor of target_100k.npy in target.npy, with values from [0, 5000). Optionally include this for better assembly results.
poses.npy # (K, 7), optional, GT poses for each part. First three coordinates denote (x, y, z) position, last four coordinates denote a quaternion with real-part first. Optionally include this for correct evaluation.
labels.npy # (5000), optional, GT segmentation label of the target, each index takes a value from [0, K). Optionally include this for correct evaluation of segmentation.
eq_class.npy # (K), optional, equivalence classes of the parts. For example, [0,0,1,2,2,2] means that the first two parts are equivalent, and last three parts are equivalent.  Optionally include this for correct evaluation.

Training and Testing

To re-train GPAT,

python learning/gpat/run.py --mode=train --ratio=0.7 --rand --exp=EXPNAME --cuda=CUDAIND

To evaluate with our pretrained GPAT checkpoint, download the checkpoint to logs/pretrained/gpat.pth.

python learning/assembler.py --eval --cat=CATEGORY --exp=EXPNAME --cuda=CUDAIND

Pass in --ratio=1 to test non-exact parts only (or another ratio from [0, 1] to indicate the probability of testing with non-exact parts). Pass in --rand to test targets at random poses. To evaluate with custom dataset, specify the directory containing all the data folders (instructions) with --eval_dir (note that you cannot directly point to a single data folder). For custom model checkpoints, pass in the model path with --model_path. For other parameters, see init_args() in learning_utils.py.

Acknowledgements

• PartNet
• Generative 3D Part Assembly via Dynamic Graph Learning
• PointTransformer
• Zhenjia Xu: Html-Visualization

Citation

If you find this codebase useful, feel free to cite our work!

@inproceedings{li2023rearrangement,
	title={Rearrangement Planning for General Part Assembly},
	author={Li, Yulong and Zeng, Andy and Song, Shuran},
	booktitle={Conference on Robot Learning},
	year={2023},
	organization={PMLR}
}

Feel free to contact Yulong if you have any questions or related ideas!