OTHER License
This code implements the CoRe model and reproduces experimental results found in Robust Robotic Control from Pixels using Contrastive Recurrent State-Space models NeurIPS Deep Reinforcement Learning Workshop 2021 Nitish Srivastava, Walter Talbott, Martin Bertran Lopez, Shuangfei Zhai & Joshua M. Susskind [paper]
Clone this repository and then execute the following steps. See setup.sh
for an example of how to run these steps on a Ubuntu 18.04 machine.
Install dependencies.
apt install -y libgl1-mesa-dev libgl1-mesa-glx libglew-dev \
libosmesa6-dev software-properties-common net-tools unzip \
virtualenv wget xpra xserver-xorg-dev libglfw3-dev patchelf xvfb ffmpeg
Download the DAVIS 2017
dataset. Make sure to select the 2017 TrainVal - Images and Annotations (480p). The training images will be used as distracting backgrounds. The DAVIS
directory should be in the same directory as the code. Check that ls ./DAVIS/JPEGImages/480p/...
shows 90 video directories.
Install MuJoCo 2.1.
mujoco210
directory into ~/.mujoco/mujoco210
.Install MuJoCo 2.0 (For robosuite experiments only).
~/.mujoco/
.mujoco200_linux
(or mujoco200_macos
) to mujoco200
.ln -s ~/.mujoco/mujoco200_linux ~/.mujoco/mujoco200
~/.mujoco/mjkey.txt
.LD_LIBRARY_PATH
.export LD_LIBRARY_PATH=$HOME/.mujoco/mujoco200/bin:$LD_LIBRARY_PATH
Setup EGL GPU rendering (if a GPU is available).
cp 10_nvidia.json /usr/share/glvnd/egl_vendor.d/
mkdir -p /usr/lib/nvidia-000
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia-000
Create a conda environment.
For Distracting Control Suite
conda env create -f conda_env.yml
For Robosuite
conda env create -f conda_env_robosuite.yml
The CoRe model can be trained on the Distracting Control Suite as follows:
conda activate core
MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python train.py --config configs/dcs/core.yaml
The training artifacts, including tensorboard logs and videos of validation rollouts will be written in ./artifacts/
.
To change the distraction setting, modify the difficulty
parameter in configs/dcs/core.yaml
. Possible values are ['easy', 'medium', 'hard', 'none', 'hard_bg']
.
To change the domain, modify the domain
parameter in configs/dcs/core.yaml
. Possible values are ['ball_in_cup', 'cartpole', 'cheetah', 'finger', 'reacher', 'walker']
.
To train on Robosuite (Door Task, Franka Panda Arm)
conda activate core_robosuite
MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python train.py --config configs/robosuite/core.yaml
conda activate core_robosuite
MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python train.py --config configs/robosuite/core_imageonly.yaml
@article{srivastava2021core,
title={Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Models},
author={Nitish Srivastava and Walter Talbott and Martin Bertran Lopez and Shuangfei Zhai and Josh Susskind},
journal={NeurIPS Deep Reinforcement Learning Workshop},
year={2021}
}
This code is released under the LICENSE terms.