PyTorch Code for "Generalization in Dexterous Manipulation via Geometry-Aware Multi-Task Learning"

Generalization in Dexterous Manipulation via
Geometry-Aware Multi-Task Learning

[Project Page] [Paper]

Wenlong Huang¹, Igor Mordatch², Pieter Abbeel¹, Deepak Pathak³

¹University of California, Berkeley, ²Google Brain, ³Carnegie Mellon University

This is a PyTorch implementation of our Geometry-Aware Multi-Task Policy. The codebase also includes a suite of dexterous manipulation environments with 114 diverse real-world objects built upon Gym and MuJoCo.

We show that a single generalist policy can perform in-hand manipulation of over 100 geometrically-diverse real-world objects and generalize to new objects with unseen shape or size. Interestingly, we find that multi-task learning with object point cloud representations not only generalizes better but even outperforms the single-object specialist policies on both training as well as held-out test objects.

If you find this work useful in your research, please cite using the following BibTeX:

@article{huang2021geometry,
  title={Generalization in Dexterous Manipulation via Geometry-Aware Multi-Task Learning},
  author={Huang, Wenlong and Mordatch, Igor and Abbeel, Pieter and Pathak, Deepak},
  journal={arXiv preprint arXiv:2111.03062},
  year={2021}
}

Setup

Requirements

• Python=3.6.9
• CUDA=10.2
• CUDNN=7.6.5
• MuJoCo=1.50 (Installation Instructions)

Setup Instructions

git clone https://github.com/huangwl18/geometry-dex.git
cd geometry-dex/
conda create --name geometry-dex-env python=3.6.9
conda activate geometry-dex-env
pip install --upgrade pip
pip install -r requirements.txt
bash install-baselines.sh

Running Code

Below are some flags and parameters for run_ddpg.py that you may find useful for reference:

The code also uses WandB. You may wish to run wandb login in terminal to record to your account or choose to run anonymously.

WARNING: Due to the large number of total environments, generating videos during training can be slow and memory intensive. You may wish to train the policy without generating videos by passing video_count=0. After training completes, simply run run_ddpg.py with flags --eval and --video_count=1 to visualize the policy. See example below.

Training

To train Vanilla Multi-Task DDPG policy:

python run_ddpg.py --expID 1 --video_count 0 --n_cycles 40000 --chunk 10

To train Geometry-Aware Multi-Task DDPG policy, first pretrain PointNet encoder:

python train_pointnet.py --expID 2

Then train the policy:

python run_ddpg.py --expID 3 --video_count 0 --n_cycles 40000 --chunk 10 --point_cloud --pointnet_load_path 2 --no_save_buffer

Note we don't save replay buffer here because it is slow as it contains sampled point clouds. If you wish to resume training in the future, do not pass --no_save_buffer above.

Evaluation / Visualization

To evaluate a trained policy and generate video visualizations, run the same command used to train the policy but with additional flags --eval --video_count=<VIDEO_COUNT> --load_path=<LOAD_EXPID>. Replace <VIDEO_COUNT> with 1 if you wish to enable visualization and 0 otherwise. Replace <LOAD_EXPID> with the Experiment ID of the trained policy. For a Geometry-Aware Multi-Task DDPG policy trained using above command, run the following for evaluation and visualization:

python run_ddpg.py --expID 4 --video_count 1 --n_cycles 40000 --chunk 10 --point_cloud --pointnet_load_path 2 --no_save_buffer --eval --load_path 3

Trained Models

We will be releasing trained model files for our Geometry-Aware Policy and single-task oracle policies for each individual object. Stay tuned! Early access can be requested via email.

Provided Environments

Training Envs
e_toy_airplane	knife	flat_screwdriver	elephant	apple
scissors	i_cups	cup	foam_brick	pudding_box
wristwatch	padlock	power_drill	binoculars	b_lego_duplo
ps_controller	mouse	hammer	f_lego_duplo	piggy_bank
can	extra_large_clamp	peach	a_lego_duplo	racquetball
tuna_fish_can	a_cups	pan	strawberry	d_toy_airplane
wood_block	small_marker	sugar_box	ball	torus
i_toy_airplane	chain	j_cups	c_toy_airplane	airplane
nine_hole_peg_test	water_bottle	c_cups	medium_clamp	large_marker
h_cups	b_colored_wood_blocks	j_lego_duplo	f_toy_airplane	toothbrush
tennis_ball	mug	sponge	k_lego_duplo	phillips_screwdriver
f_cups	c_lego_duplo	d_marbles	d_cups	camera
d_lego_duplo	golf_ball	k_toy_airplane	b_cups	softball
wine_glass	chips_can	cube	master_chef_can	alarm_clock
gelatin_box	h_lego_duplo	baseball	light_bulb	banana
rubber_duck	headphones	i_lego_duplo	b_toy_airplane	pitcher_base
j_toy_airplane	g_lego_duplo	cracker_box	orange	e_cups

Test Envs
rubiks_cube	dice	bleach_cleanser	pear	e_lego_duplo
pyramid	stapler	flashlight	large_clamp	a_toy_airplane
tomato_soup_can	fork	cell_phone	m_lego_duplo	toothpaste
flute	stanford_bunny	a_marbles	potted_meat_can	timer
lemon	utah_teapot	train	g_cups	l_lego_duplo
bowl	door_knob	mustard_bottle	plum

Acknowledgement

The code is adapted from this open-sourced implementation of DDPG + HER. The object meshes are from the YCB Dataset and the ContactDB Dataset. We use SubprocChunkVecEnv from this pull request of OpenAI Baselines to speedup vectorized environments.