jason-cky / rl_vae-pytorch

Pipeline to generate images and to train a Variational Autoencoder (VAE) for use in Deep RL environments

Deep RL policies on Pybullet Environments

This repo is a pytorch implementation of training a variational autoencoder (VAE). This is written to train a VAE for use in a RL environment, and contains code to generate images from random exploration of various RL environments from pybullet and RLBench for training.

Beta-VAE can also be trained by setting the Beta value to any value greater than 1.

Dependencies:

• CUDA >= 10.2
• RLBench, only if you want to use RLBench environments to train VAE

How to use

• Clone this repo
• pip install -r requirements.txt

Generating data from openai gym environment

python generate_data.py
usage: generate_data.py [-h] --env ENV --num_samples NUM_SAMPLES
                        [--max_ep_len MAX_EP_LEN] [--seed SEED] [--rlbench]
                        [--view {wrist_rgb,front_rgb,left_shoulder_rgb,right_shoulder_rgb}]

optional arguments:
  -h, --help            show this help message and exit
  --env ENV             environment_id
  --num_samples NUM_SAMPLES
                        specify number of image samples to generate
  --max_ep_len MAX_EP_LEN
                        Maximum length of an episode
  --seed SEED           seed number for reproducibility
  --rlbench             if true, use rlbench environment wrappers
  --view {wrist_rgb,front_rgb,left_shoulder_rgb,right_shoulder_rgb}
                        choose the type of camera view to generate image (only
                        for RLBench envs)

Training VAE

usage: train_vae.py [-h] [--dir DIR] [--seed SEED] [--num_workers NUM_WORKERS]
                    [--batch_size BATCH_SIZE] [--epochs EPOCHS] [--beta BETA]
                    [--lr LR] [--ngpu NGPU] [--save_freq SAVE_FREQ]
                    [--log_freq LOG_FREQ] [--save_dir SAVE_DIR]

optional arguments:
  -h, --help            show this help message and exit
  --dir DIR             path to image folders
  --seed SEED           seed number for reproducibility
  --num_workers NUM_WORKERS
                        number of workers for dataloaders
  --batch_size BATCH_SIZE
                        batch size
  --epochs EPOCHS       Number of epochs
  --beta BETA           Weighing value for KLD in B-VAE
  --lr LR               Learning Rate
  --ngpu NGPU           number of gpus to use
  --save_freq SAVE_FREQ
                        save weights every <x> iterations
  --log_freq LOG_FREQ   log losses every <x> iterations
  --save_dir SAVE_DIR   path to save weights and logs

rl_vae-pytorch's Issues

question

If there are some examples to use vae like this to train rl, it will be better. Your code is only training vae.