This is a partial implementation of the paper Deep Imitation Learning for Complex Manipulation Tasks from Virtual Reality Teleoperation written for Danfei Xu's Spring 2023 class CS 8803: Deep Learning for Robotics at Georgia Tech. More info about this class can be found here.
A google slides presentation about the project can be found here.
This repo has 3 functions.
- Use virtual reality to collect training data of a simulated pybullet environment. Video of Data Collection
- Use this training data to train the model described in the above paper.
- Run this trained model in the simulated enviornment to validate performance. Video of Model Running
datacollector.ipynb - a notebook for generating training data and running a pretrained model.
trainer.ipynb - a notebook for training the model with collected data.
VRNet.py - a pytorch implementation of the imitation learning model described in the above paper. Note that this contains a SpatialSoftmax implementation from here.
model.pt - a pretrained model which can place the cube on the tray after starting on the cube.
RGB Image:
Depth Image:
The full training set is available upon request. If there's interest I can upload it to google drive or kaggle.