ManiBox: Enhancing Spatial Grasping Generalization via Scalable Simulation Data Generation

Paper
Project Page
Other Contributors: Xuezhou Xu, Chengyang Ying, Xinyi Mao

The code will be updated within a week. TODO:

more policies
data generation in simulator

Installation Instructions

conda deactivate
conda create -n manibox python=3.9
conda activate manibox

pip install -e .

# for student inference code, you should install these in `isaac lab conda env`:
pip install einops

Dataset

The full space version dataset is in Download link. It has a total of 38,150 trajectories, which you should rename to integration.pkl in any directory.

Student Training Instructions

The train.py will will read integration.pkl in the --dataset directory as a dataset. integration.pkl is a dict containing three keys, image_data, qpos_data, action_data, with shape (num_episodes, episode_len , dim).

# BBOX RNN
python ManiBox/train.py --policy_class RNN --batch_size 128 --dataset ../ --num_episodes 38000  --loss_function l1  --rnn_layers 3 --rnn_hidden_dim 512 --actor_hidden_dim 512 --num_epochs 50 --lr 2e-3 --gradient_accumulation_steps 1 > train.log 2>&1

Deployment Instructions

python ManiBox/inference_real_world.py  --ckpt_dir /PATH/TO/ManiBox/ckpt/2024-xx-xx_xx-xx-xxRNN --policy_class RNN --ckpt_name policy_best.ckpt

Other Instructions

TODO...

Teacher Policy Training Instructions

# PPO Training
python source/standalone/workflows/rsl_rl/train.py --task Isaac-Lift-Cube-MobileAloha-v0  --num_envs 4096  --headless

Data Collection Instructions

# Collect data in simulator
python source/standalone/workflows/rsl_rl/play_collect_data.py --task Isaac-Lift-Cube-MobileAloha-Play-v0 \
--num_envs 40  --load_run  2024-09-12_17-42-24 --headless --enable_cameras

Simulator Inference Instruction

# Student policy inference in simulator
python source/standalone/workflows/rsl_rl/student_inference_orbit_multi_envs.py --task Isaac-Lift-Cube-MobileAloha-Play-v0 \
--ckpt_dir "\PATH\TO\Your\CKPT" --policy_class "RNN" --ckpt_name policy_best.ckpt --nheads 48 --num_train_step 38000 --seed 0

Acknowledgement

BibTeX

If you find our work useful for your project, please consider citing the following paper.

@article{tan2024manibox,
  title={ManiBox: Enhancing Spatial Grasping Generalization via Scalable Simulation Data Generation},
  author={Tan, Hengkai and Xu, Xuezhou and Ying, Chengyang and Mao, Xinyi and Liu, Songming and Zhang, Xingxing and Su, Hang and Zhu, Jun},
  journal={arXiv preprint arXiv:2411.01850},
  year={2024}
}

thkkk/manibox