ikostrikov / jaxrl

JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

JAX (Flax) RL

This repository contains JAX (Flax) implementations of Reinforcement Learning algorithms:

• Soft Actor Critic with learnable temperature
• Advantage Weighted Actor Critic
• Image Augmentation Is All You Need(only [K=1, M=1])
• Deep Deterministic Policy Gradient with Clipped Double Q-Learning
• Randomized Ensembled Double Q-Learning: Learning Fast Without a Model
• Behavioral Cloning

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:

@misc{jaxrl,
  author = {Kostrikov, Ilya},
  doi = {10.5281/zenodo.5535154},
  month = {10},
  title = {{JAXRL: Implementations of Reinforcement Learning algorithms in JAX}},
  url = {https://github.com/ikostrikov/jaxrl},
  year = {2021}
}

You can find an updated version of this repository here.

Changelog

January 10th, 2022

• Added an implementation of Randomized Ensembled Double Q-Learning: Learning Fast Without a Model

July 20th, 2021

• Added an implementation of Deep Deterministic Policy Gradient with Clipped Double Q-Learning

May 19th, 2021

• Added an implementation of Soft Actor Critic v1

April 29th, 2021

• Added an implementation of data augmentation from Image Augmentation Is All You Need

Installation

Prerequisites:

• Python 3.8-3.9 (not yet 3.10)
• Poetry
• patchelf

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.

Examples

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

Continous control from pixels

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.

jaxrl's Issues

Assert always evaluates to True

In these lines, because the assert statement is broken into multiple lines and hence covered with parentheses, the assert evaluates the tuple itself. So it will always evaluate to True.

jaxrl/jaxrl/networks/common.py

Lines 80 to 81 in 3674862

	assert (loss_fn is not None or grads is not None,
	'Either a loss function or grads must be specified.')

jaxrl/jaxrl/networks/common.py

Lines 89 to 90 in 3674862

	assert (has_aux,
	'When grads are provided, expects no aux outputs.')

This can be fixed by changing it to

assert loss_fn is not None or grads is not None, \
                'Either a loss function or grads must be specified.'

assert has_aux, \
                    'When grads are provided, expects no aux outputs.'

A question about temperature module

Hi,

Thanks a lot for the great implementation!

I have a question regarding the temperature, I saw that in
https://github.com/ikostrikov/jax-rl/blob/main/jax_rl/agents/sac/temperature.py#L16

You are computing the temperature (entropy coefficient) by this instead of directly taking an exp on the log-temperature. Is there a particular reason that you implement it this way as compared to other implementation (for example, the pytorch DrQ, soft learning etc.)

Thanks!

Comment on installation of MuJoCo and GPU version of JAX

Remove TF dependency

TF is a heavy library that might need a lot of dependencies and time to install. Removing the dependency on tensorflow_probability is not essential but might be good too.

[question] Can you add GAIL to this?

Great work btw

Discount factor in main.py might be misleading

Right now it indicates whether the env is done due to a timelimit, while in the update we use 0.99.

Add reward plotting script

Garden variety installation error

Hey @ikostrikov really cool project!

I was trying to install it but hit some installation difficulties. It should be an easy fix if I look around more, but maybe it's worth introducing a proper dependency manager like poetry. I filed PR #20 to close this issue. Let me know if you'd be interested in it and I can add more docs :)

➜  jaxrl git:(main) ✗ . /home/costa/.cache/pypoetry/virtualenvs/jaxrl-m03KgMHV-py3.9/bin/activate
(jaxrl-m03KgMHV-py3.9) ➜  jaxrl git:(main) ✗ pip install -e .
Looking in indexes: https://pypi.org/simple, https://pypi.ngc.nvidia.com
Obtaining file:///home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
    Preparing wheel metadata ... done
Installing collected packages: jaxrl
  Running setup.py develop for jaxrl
    ERROR: Command errored out with exit status 1:
     command: /home/costa/.cache/pypoetry/virtualenvs/jaxrl-m03KgMHV-py3.9/bin/python -c 'import io, os, sys, setuptools, tokenize; sys.argv[0] = '"'"'/home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl/setup.py'"'"'; __file__='"'"'/home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl/setup.py'"'"';f = getattr(tokenize, '"'"'open'"'"', open)(__file__) if os.path.exists(__file__) else io.StringIO('"'"'from setuptools import setup; setup()'"'"');code = f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' develop --no-deps
         cwd: /home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl/
    Complete output (3 lines):
    Traceback (most recent call last):
      File "<string>", line 1, in <module>
    ModuleNotFoundError: No module named 'setuptools'
    ----------------------------------------
ERROR: Command errored out with exit status 1: /home/costa/.cache/pypoetry/virtualenvs/jaxrl-m03KgMHV-py3.9/bin/python -c 'import io, os, sys, setuptools, tokenize; sys.argv[0] = '"'"'/home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl/setup.py'"'"'; __file__='"'"'/home/costa/Documents/go/src/github.com/vwxyzjn/jaxrl/setup.py'"'"';f = getattr(tokenize, '"'"'open'"'"', open)(__file__) if os.path.exists(__file__) else io.StringIO('"'"'from setuptools import setup; setup()'"'"');code = f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' develop --no-deps Check the logs for full command output.