SNeRL

This is the official PyTorch implementation of the paper "SNeRL: Semantic-aware Neural Radiance Fields for Reinforcement Learning" (ICML 2023). [paper]

Dataset

We provide dataset for two environments (Drawer, Soccer).

Drawer

Soccer

Setup Instructions

0. Create a conda environment:

   conda create -n snerl python=3.9
   conda activate snerl

1. Install MuJoCo and task environments:

   cd metaworld
   pip install -e .
   cd ..

2. install pytorch (use tested on pytorch 1.12.1 with CUDA 11.3)

3. install additional dependencies:

   pip install scikit-image
   pip install tensorboard
   pip install termcolor
   pip install imageio
   pip install imageio-ffmpeg
   pip install opencv-python
   pip install matplotlib
   pip isntall tqdm
   pip install timm
   pip install configargparse

Usage

Our code does not include the dataset generator for nerf pretraining. Please prepare your dataset for nerf pretraining.

Pretrain Encoder

cd nerf_pretrain
python run_nerf.py --config configs/{env_name}.txt

Train Donstream RL

0. Locate pretained model in './encoder_pretrained/{env_name}/snerl.tar'

1. Use the following commands to train RL agents:

window-open-v2

CUDA_VISIBLE_DEVICES=0 python snerl/train.py --env_name window-open-v2 --encoder_type nerf --save_tb --frame_stack 2 --eval_freq 10000 --batch_size 128 --save_video --save_model --image_size 128 --camera_name cam_1_1 cam_7_4 cam_14_2 --multiview 3 --encoder_name 'snerl' --seed 1

drawer-open-v2

CUDA_VISIBLE_DEVICES=0 python snerl/train.py --env_name drawer-open-v2 --encoder_type nerf --save_tb --frame_stack 2 --eval_freq 10000 --batch_size 128 --save_video --save_model --image_size 128 --camera_name cam_1_1 cam_7_4 cam_14_2 --multiview 3 --encoder_name 'snerl' --seed 1

hammer-v2

CUDA_VISIBLE_DEVICES=0 python snerl/train.py --env_name hammer-v2 --encoder_type nerf --save_tb --frame_stack 2 --eval_freq 10000 --batch_size 128 --save_video --save_model --image_size 128 --camera_name cam_1_1 cam_7_4 cam_14_2 --multiview 3 --encoder_name 'snerl' --seed 1

soccer-v2

CUDA_VISIBLE_DEVICES=0 python snerl/train.py --env_name soccer-v2 --encoder_type nerf --save_tb --frame_stack 2 --eval_freq 10000 --batch_size 128 --save_video --save_model --image_size 128 --camera_name cam_1_1 cam_7_4 cam_14_2 --multiview 3 --encoder_name 'snerl' --seed 1

Reference

Our code sourced and modified from official implementation of CURL and NeRF.

Citation

If you use this repo in your research, please consider citing the paper as follows.

@inproceedings{shim2023snerl,
  title={SNeRL: Semantic-aware Neural Radiance Fields for Reinforcement Learning},
  author={Shim, Dongseok and Lee, Seungjae and Kim, H Jin},
  booktitle=International Conference on Machine Learning},
  pages={},
  year={2023},
  organization={PMLR}
}