Video-Language Critic: Transferable Reward Functions for Language-Conditioned Robotics

June 4, 2024 ยท View on GitHub

Minttu Alakuijala1, Reginald McLean2, Isaac Woungang2, Nariman Farsad2, Samuel Kaski1,3, Pekka Marttinen1, Kai Yuan 4

1Department of Computer Science, Aalto University, 2Department of Computer Science, Toronto Metropolitan University, 3Department of Computer Science, University of Manchester, 4Intel Corporation

This is the official training reinforcement learning training code for VLC, a reward model trained using contrastive learning and a temporal ranking objective for downstream robotics tasks.

Table of Contents

Installation

VLC Experiments

License

Citation

Acknowledgement

Installation

If you do not have Video Language Critic (VLC) installed, please install this repo first.

We do recommend creating an environment via miniconda as it can handle the requirements for headless rendering.

conda create -n vlc_rl python=3.10 pip

After cloning this repo, in the same environment that VLC is installed in:

cd vlc_rl
git submodule update --init --recursive
cd Metaworld
pip install -e .
cd ../Gymnasium
pip install -e .
cd ../
pip install -r requirements.txt
pip install 'jax[cuda12]==0.4.23'

Next we will need to setup the paths.py file

cp paths_template.py paths.py

then edit paths.py with your favourite text editor, replacing the paths with your local paths. MW_DATA_DIR should point to the data from VLC. The data can be downloaded from this Google Drive, specifically the data.zip file. The REWARD_CKPT_DIR should point to a folder that contains the VLC checkpoint files, also from this Google Drive in the VLC Checkpoints folder. EXP_DIR will be the output of the experiment.

All needed data can be downloaded with the following commands:
MW_DATA: gdown 1u0aNE2pnsWTC669tVwZ0zv-5Zk42-S-R
REWARD_CKPT: gdown 1hWpQO7nWIM5IhY4fRI0LxxABTmQ7LIvo

VLC Experiments

Replication

To run a VLC experiment, use the following command replacing ENV_NAME with the environment name and FILE_NAME with the VLC checkpoint you would like to use.

python sac_jax.py --env-id ENV_NAME --seed 12345 --reward-normalization-gymnasium --reward-normalization-offset --predict-for-partial-videos --stretch-partial-videos --sparse-reward-weight=50 --vlm-reward-weight=1 --evaluation-frequency 20000 --vlc-ckpt FILE_NAME

Rendering

This repo currently uses EGL which should be supported on current NVIDIA GPU's. Untested on AMD/Intel GPUs. To enable Mujoco to use EGL, run the following commands in a terminal on a headless server or add them to your .bashrc file:

export MUJOCO_GL=egl
export PYOPENGL_PLATFORM=egl

License

The source code in this repository is licensed under the X License.

Citation

If you find this repository or paper useful for your research, please cite

@article{alakuijala2024videolanguage,
      title={Video-Language Critic: Transferable Reward Functions for Language-Conditioned Robotics}, 
      author={Minttu Alakuijala and Reginald McLean and Isaac Woungang and Nariman Farsad and Samuel Kaski and Pekka Marttinen and Kai Yuan},
      year={2024},
      eprint={2405.19988},
      archivePrefix={arXiv},
      primaryClass={cs.RO}
}

Acknowledgement

This implementation is based off CleanRL. Thanks to Evangelos and Frank for the base implementation!