CoRe: Contrastive Recurrent State-Space Models

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]

Requirements and Installation

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.

  • Download MuJoCo version 2.1 binaries for Linux or macOS.
  • Unzip the downloaded mujoco210 directory into ~/.mujoco/mujoco210.

• Install MuJoCo 2.0 (For robosuite experiments only).

  • Download MuJoCo version 2.0 binaries for Linux or macOS.
  • Unzip the downloaded directory and move it into ~/.mujoco/.
  • Symlink mujoco200_linux (or mujoco200_macos) to mujoco200.

  ln -s ~/.mujoco/mujoco200_linux ~/.mujoco/mujoco200
  

  • Place the license key at ~/.mujoco/mjkey.txt.
  • Add the MuJoCo binaries to LD_LIBRARY_PATH.

  export LD_LIBRARY_PATH=$HOME/.mujoco/mujoco200/bin:$LD_LIBRARY_PATH
  

• Setup EGL GPU rendering (if a GPU is available).

  • To ensure that the GPU is prioritized over the CPU for EGL rendering

  cp 10_nvidia.json /usr/share/glvnd/egl_vendor.d/
  

  • Create a dummy nvidia directory so that mujoco_py builds the extensions needed for GPU rendering.

  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
  

Training

• 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)

  • Using RGB image and proprioceptive inputs.

  conda activate core_robosuite
  MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python train.py --config configs/robosuite/core.yaml
  

  • Using RGB image inputs only.

  conda activate core_robosuite
  MUJOCO_GL=egl CUDA_VISIBLE_DEVICES=0 python train.py --config configs/robosuite/core_imageonly.yaml
  

Citation

@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}
}

License

This code is released under the LICENSE terms.