Regression Planning Networks
A reference implementation of the Regression Planning Networks.
This repo also includes the simulated environments used in the paper: GridWorld and Kitchen3D. GridWorld is built on gym-minigrid. Kitchen3D is simulated with PyBullet. Most of the PyBullet utilities are adapted from Caelan Garrett's ss-pybullet. The Appendix of the original paper is also included here for reference.
To cite the code, use bibtex:
@inproceedings{xu2019rpn,
title={Regression Planning Networks},
author={Xu, Danfei and Martín-Martín, Roberto and Huang, De-An and Zhu, Yuke and Savarese, Silvio and Fei-Fei, Li},
booktitle={Thirty-third Conference on Neural Information Processing Systems (NeurIPS)},
year={2019}
}
Installation and Setup
Requirement: Python >= 3.6
Note: The code DOES NOT work with Python 3.5
$ git clone https://github.com/danfeiX/rpn.git
$ git submodule update --init --recursive
$ pip install -r requirements.txt
$ cd rpn
$ mkdir checkpoints
$ mkdir data
Kitchen3D
Dataset
To reproduce the main experiments, first create a demonstration dataset with three ingredients and two dishes
(I=3, D=2).
$ python create_dataset_pb.py --problem pb_cook_meal_3i_2d_3m_iter \
--dataset data/pb_cook_3i_2d_3m_10_s0.group \
--num_episodes 10 --num_tasks -1 --seed 0 \
--teleport --num_chunks 100 --num_workers 10
The entire process takes around 3 hours with 10 works (--num_workers 10) on a machine with i7-7700K CPU @ 4.20GHz.
The dataset will be written to data/pb_cook_3i_2d_3m_10_s0.group.
To create a smaller validation set, run:
$ python create_dataset_pb.py --problem pb_cook_meal_3i_2d_3m_iter \
--dataset data/pb_cook_3i_2d_3m_1_s1.group \
--num_episodes 1 --num_tasks -1 --seed 1 \
--teleport --num_chunks 10 --num_workers 10
Train
To train our full RPN model, run:
$ python train.py --dataset data/pb_cook_3i_2d_3m_10_s0.group \
--testset data/pb_cook_3i_2d_3m_1_s1.group \
--exp rpn \
--run experiment \
--batch_size 128 \
--eval_batch_size 128 \
--eval_freq 1 --save_freq 1 \
--num_epoch 10 --config configs/pbv_rpn.json \
--num_workers 1
Checkpoints will be written to checkpoints/rpn/experiment/ each epoch. Both text and tensorboard event will be written
to checkpoints/rpn/experiment/log/.
Evaluation
To evaluate a trained model on the Kitchen3D task with I=6 and D=3, run:
$ python eval_pb.py --restore_path checkpoints/rpn/experiment/ckpt_ep6.pt \
--problem pb_cook_meal_6i_3d_6m_shuffle \
--num_eval 1000 --seed 1 \
--display
Aside: Our naming convention for the cooking tasks is pb_cook_meal_<I>i_<D>d_<M>m_<shuffle/iter>, where <I> is the
number of ingredients involved, <D> is the number of dishes to cook, <M> is the maximum number of ingredients that
can be placed on a plate (usually set to be equal to <I>), and <shuffle/iter> is whether to shuffle the list of
all possible tasks under the current task specification.