/event-vae-rl

Visuomotor policies from event-based cameras through representation learning and reinforcement learning. Accompanies our paper: https://arxiv.org/abs/2103.00806

Primary LanguagePythonMIT LicenseMIT

Representation Learning and Reinforcement Learning from Event Cameras

This repository provides a PyTorch codebase to train and evaluate models as seen in the paper

Representation Learning for Event-based Visuomotor Policies
Sai Vemprala, Sami Mian, Ashish Kapoor
NeurIPS 2021 (Spotlight) Paper

Please visit our blogpost for a summary on the idea and approach!

License and Citation

This project is licensed under the terms of the MIT license. By using the software, you are agreeing to the terms of the license agreement. If you use this code in your research, please cite our work as follows:

@misc{vemprala2021representation,
      title={Representation Learning for Event-based Visuomotor Policies}, 
      author={Sai Vemprala and Sami Mian and Ashish Kapoor},
      year={2021},
      eprint={2103.00806},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Structure

The repository contains code to train an event variational autoencoder (eVAE) in event_vae. event_rl contains code to take pretrained eVAE representations and train RL policies to perform obstacle avoidance in AirSim. This uses a basic gym wrapper over AirSim, which can be extended to other kinds of downstream tasks. The event data is generated by an 'event simulator' which simulates event firings from successive images captured by AirSim. Code in event_rl uses stable-baselines3 for the RL algorithm implementation, and uses the PPO algorithm.

Sample data for training the VAE, and some pretrained weights/RL policies can be found under the Releases section.

Getting started

The following steps will get you set up with the required packages:

  1. Clone our repo: git clone https://github.com/microsoft/event-vae-rl
  2. Install dependencies:
cd event-vae-rl
conda create -n evrl python=3.8
conda activate evrl
pip install -r requirements.txt

Event VAE

To train an event VAE, the code expects a file of event data. We provide some sample files to train representations over (See https://github.com/microsoft/event-vae-rl/releases/tag/v1.0-d/), or you can generate event data files yourself from AirSim through AirSim's event simulator feature.

Given a file of events, you can run the train_vae.py script as follows:

  • With polarity and temporal coding: python train_vae.py --input_file <path_to_event_data> --data_len 3 --tcode
  • Without polarity and temporal coding: python train_vae.py --input_file <path_to_event_data> --data_len 2

See test.py for a starting point on how to test trained eVAE models.

Event RL

The obstacle course environment used for training and testing policies is under the Releases section, accessible as an AirSim binary. Both Windows and Linux binaries are available.

To run the AirSim environment:

  1. Download the binaries from the releases section and unzip it.
  2. Copy the settings file event_rl/settings.json to C:\Users\$USERNAME\Documents\AirSim\settings.json.
  3. Navigate to the downloaded binary folder and run ./run_env.bat.

Our sample task involves a drone navigating in an obstacle course in two dimensions, with three discrete actions: forward, left or right.

The environment contains several 'maps', accessible through the UE console command Open <mapname> (UE console can be opened with the tilde key ~). We also provide evaluation maps with different obstacle textures and shapes, focused on evaluating the out-of-distribution generalization capability of the policies trained with the eVAE. The available maps are as follows:

  • Train:Training environment with pole-shaped obstacles and four lanes of increasing difficulty.
  • Test: Larger poles environment for evaluation
  • ChangeTexture: Test environment with obstacles of different texture.
  • ChangeShape: Test environment with elliptical obstacles.
  • ChangeTextureShape: Test environment with triangular obstacles with dynamic texture.

Given pretrained eVAE weights, a policy can be trained as follows, once the right AirSim environment is up and running:

python train.py --obs_type event_stream --data_len 3 --tc --rep_weights <path-to-evae-weights> --ls <latent_vector_size>

Through test.py, a trained policy can be used to test in different environments. To test the robustness of trained policies, the environment can be initialized with a couple of extra arguments in test.py, which are event noise (random events firing without any activity), and sparsity (events not firing although there is activity).

Contributions

This project welcomes contributions and suggestions. Most contributions require you to agree to a Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.

When you submit a pull request, a CLA bot will automatically determine whether you need to provide a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions provided by the bot. You will only need to do this once across all repos using our CLA.

This project has adopted the Microsoft Open Source Code of Conduct. For more information see the Code of Conduct FAQ or contact opencode@microsoft.com with any additional questions or comments.