ml-core

OTHER License

Stars
31

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.