Visual Imitation Made Easy

[Project Page] [Paper] [Video]

Sarah Young¹, Dhiraj Gandhi², Shubham Tulsiani², Abhinav Gupta^{2 3}, Pieter Abbeel¹, Lerrel Pinto^{1 4}

¹University of California, Berkeley, ²Facebook AI Research, ³Carnegie Mellon University, ⁴New York University

Data

Pushing data is located here.

Stacking data is located here.

Usage

Below are example scripts for training and testing on provided sample data.

Setup

1. Clone repo.

git clone https://github.com/sarahisyoung/Visual-Imitation-Made-Easy.git

2. Create and activate conda environment.

conda env create -f environment.yml
conda activate trashbot

3. Set env path.

export PYTHONPATH=$PYTHONPATH:path_to_proj/

Training

To train with custom data, see this for details on data processing.

1. Download the data (~45 GB) from dropbox by running this script. Feel free to comment out training/validation URL in the script to download just a small sample of the data.

   python download_data.py --task push

2. To train:

   python train.py --task push --train_dir data/train --val_dir data/val --test_dir data/test --save_dir results

Testing

1. This command predicts on a folder of images. Output visualization is saved to specified folder.

python push_test.py --model results/exp1/policy_earlystop.pt --image_folder test_data/ --output predicted_data/