pku-epic / partmanip Goto Github PK

This is the official repository of PartManip: Learning Cross-Category Generalizable Part Manipulation Policy from Point Cloud Observations.

For more information, please visit our project page.

We introduce a large-scale cross-category part manipulation benchmark PartManip with diverse object datasets, realistic settings, and rich annotations. We propose a generalizable vision-based policy learning strategy and boost the performance of part-based object manipulation by a large margin, which can generalize to unseen object categories and novel objects in the real world.

We first train state-based expert policy using our proposed canonicalization to the part coordinate frame and the part-aware reward. We then use the learned expert to collect demonstrations for pre-training the vision-based policy by behavior cloning. After pre-training, we train the vision-based policy to imitate the state-based expert policy using DAgger. We also introduce several point cloud augmentation techniques to boost the generalization ability. For the vision backbone, we introduce 3D Sparse-UNet which has a large expression capability. Furthermore, we introduced an extra domain adversarial learning module for better cross-category generalization.

To obtain our assets, please fill out this form and check the Terms&Conditions in it. Please cite our paper if you use our dataset.

1. install issacgym following the official guide

2. install dependencies

pip install wandb tensorboard ipdb gym tqdm rich opencv_python pytorch3d pyparsing pytorch_lightning addict yapf h5py sorcery pynvml torchdata==0.5.1 einops

3. (Optional for full implementation, not necessary) install GAPartNet environment following the official guide

For training:

run scripts in scripts/train.sh to train different tasks.

An example:

# drawer
python gym/train.py --task=FrankaPoseCabinetBase --algo=pregrasp_ppo \
--task_config=cfg/task_cfg.yaml --algo_config=cfg/algo/pregrasp_ppo.yaml \
--sim_device=cuda:0 --graphics_device_id=0 --seed=0526 --headless \
--group_name state --save_name 0.2_2_0_1_0_2 \
--rot_coef 0.2 --handle_coef 2 --vel_coef 0 --dist_coef 1 --tip_coef 0 --dicouple_coef 2 \
--target open --category drawer --minibatches 2 --nsteps 5 --noptepochs 8 \
--canonical_space handle --pregrasp part --obs state --eval_round 3 --eval_freq 10 \
--action_normalization clip --control ik_abs_axis_angle --asset_num 200_40_40_5

For testing:

run scripts in scripts/test.sh to test different tasks with trained ckpt.

An example:

# drawer
python gym/train.py --task=FrankaPoseCabinetBase --algo=pregrasp_ppo \
--task_config=cfg/task_cfg.yaml --algo_config=cfg/algo/pregrasp_ppo.yaml \
--sim_device=cuda:0 --graphics_device_id=0 --seed=0526 --headless \
--group_name state --save_name 0.2_2_0_1_0_2 \
--rot_coef 0.2 --handle_coef 2 --vel_coef 0 --dist_coef 1 --tip_coef 0 --dicouple_coef 2 \
--target open --category drawer --minibatches 2 --nsteps 5 --noptepochs 8 \
--canonical_space handle --pregrasp part --obs state --eval_round 3 --eval_freq 10 \
--action_normalization clip --control ik_abs_axis_angle --asset_num 200_40_40_5 \
--test --ckpt xxxx.tar

There are some examples of our tested training curves:

A training example of drawer:

A training example of door:

This work and the dataset are licensed under CC BY-NC 4.0.

If you find our work useful in your research, please consider citing:

@article{geng2023partmanip,
    title={PartManip: Learning Cross-Category Generalizable Part Manipulation Policy from Point Cloud Observations},
    author={Geng, Haoran and Li, Ziming and Geng, Yiran and Chen, Jiayi and Dong, Hao and Wang, He},
    journal={arXiv preprint arXiv:2303.16958},
    year={2023}
}

If you have any questions, please open a github issue or contact us:

Haoran Geng: [email protected]

He Wang: [email protected]