JAX (Flax) RL
This repository contains JAX (Flax) implementations of Reinforcement Learning algorithms:
The goal of this repository is to provide simple and clean implementations to build research on top of. Please do not use this repository for baseline results and use the original implementations instead (SAC, AWAC, DrQ).
If you use JAXRL in your work, please cite this repository in publications:
Installation
Prerequisites:
Suggested build environment:
$$
general build dependencies sudo apt-get update; sudo apt-get install make build-essential libssl-dev zlib1g-dev \ libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \ libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev # mujoco dependencies apt-get -y install wget unzip software-properties-common \ libgl1-mesa-dev \ libgl1-mesa-glx \ libglew-dev \ libosmesa6-dev patchelf # mujoco installation curl -OL https://mujoco.org/download/mujoco210-linux-x86_64.tar.gz mkdir ~/.mujoco tar -zxf mujoco210-linux-x86_64.tar.gz -C ~/.mujoco rm mujoco210-linux-x86_64.tar.gz
$$
To install, run
$$
poetry install # For GPU support run pip install "jax[cuda]==0.3.10" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
$$
For further instructions on running this code on GPU, please follow instructions from the official repository.
For MuJoCo inslattion, you may need to add the following lines in the .bashrc
:
$$
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/costa/.mujoco/mujoco210/bin export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia
$$
Development
If you want to modify the code, install following the instructions above.
Troubleshooting
If you experience out-of-memory errors, especially with enabled video saving, please consider reading docs on JAX GPU memory allocation. Also, you can try running with the following environment variable:
$$
XLA_PYTHON_CLIENT_MEM_FRACTION=0.80 python ...
$$
If you run your code on a remote machine and want to save videos for DeepMind Control Suite, please use EGL for rendering:
$$
MUJOCO_GL=egl python train.py --env_name=cheetah-run --save_dir=./tmp/ --save_video
$$
Tensorboard
Launch tensorboard to see training and evaluation logs
$$
tensorboard --logdir=./tmp/
$$
Results
Continous control from states
![gym](https://github.com/ikostrikov/jaxrl/raw/main/learning_curves/images/results.png)
Continous control from pixels
![gym](https://github.com/ikostrikov/jaxrl/raw/main/learning_curves/images/results_drq.png)
Docker
Build
Copy your MuJoCo key to ./vendor
$$
cd remote docker build -t ikostrikov/jaxrl . -f Dockerfile
$$
Test
$$
sudo docker run -v <examples-dir>:/jaxrl/ ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/ # On GPU sudo docker run --rm --gpus all -v <examples-dir>:/jaxrl/ --gpus=all ikostrikov/jaxrl:latest python /jaxrl/train.py --env_name=HalfCheetah-v2 --save_dir=/jaxrl/tmp/
$$
Contributing
When contributing to this repository, please first discuss the change you wish to make via issue. If you are not familiar with pull requests, please read this documentation.
Acknowledgements
Thanks to @evgenii-nikishin for helping with JAX. And @dibyaghosh for helping with vmapped ensembles.