Code used for our 2022 RSS paper, Play it by Ear: Learning Skills amidst Occlusion through Audio-Visual Imitation Learning
- Run
conda env create -f conda_env.yml
. - Install system-compatible versions of cudatoolkit, pytorch, torchvision, and torchaudio
The demo collection program is writeDemos_episodes.py
with a .yaml
configuration under the same name.
The environmentName
parameter sets the environment. Pick between BlockedPickPlace
(occluded pick-place task) and IndicatorBoxBlock
(occluded lift task).
Set the demo_root
parameter to demo directory. Other parameters to change include episodeLength
and episodes
. After running, the code will save the demos to a .pkl
file that will be loaded by training code.
The imitation learning code is imitationtrain_memory.py
with a .yaml
configuration under the same name.
Like in the demo collection, the environmentName
parameter sets the environment. Pick between BlockedPickPlace
(occluded pick-place task) and IndicatorBoxBlock
(occluded lift task).
Set both demo_root
and actor_root
to be the demo directory and the results directory, respectively. They should be outer directories, i.e. if your demo is saved in demos/IndicatorBoxBlock_50/demo.pkl
, the demo_root
should be demos/
.
The same code can also accomodate balanced batch training. Use the balanced_batches
flag and modify the corrections_file
to reference the replay buffer of corrections. If priority
is enabled, we will only sample active corrections from the corrections buffer. Otherwise, we will sample uniformly from the whole episode.
The intervention code is sim_intervention_episodes.py
with a .yaml
configuration under the same name.
Like in simulation, set both demo_root
and actor_root
to be the demo directory and the results directory, respectively. Set the desired environment.
Set the orig_seed
to the seed of the imitation-trained model, and set the checkpoint
to the saved model version. The locations of the files should be automatically found.
The core/
directory contains all the code for the model we used, as well as some critical utility functions and replay buffer implementations.
The custom_environments/
directory contains the code and .xml
models for the Robosuite environments we used.
All the code that can be run directly is in the main folder.
This codebase is adapted from Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels (https://github.com/denisyarats/drq)