A simple and modular implementation of the Conservative Q Learning and Soft Actor Critic algorithm in Jax and Flax.

This repository is a reimplementation of my other codebase of the same algorithms in Pytorch.

1. Install and use the included Ananconda environment

$ conda env create -f environment.yml
$ source activate JaxCQL

You'll need to get your own MuJoCo key if you want to use MuJoCo.

2. Add this repo directory to your PYTHONPATH environment variable.

export PYTHONPATH="$PYTHONPATH:$(pwd)"

You can run SAC experiments using the following command:

python -m JaxCQL.sac_main \
    --env 'HalfCheetah-v2' \
    --logging.output_dir './experiment_output'

All available command options can be seen in JaxCQL/conservative_sac_main.py and JaxCQL/conservative_sac.py.

You can run CQL experiments using the following command:

python -m JaxCQL.conservative_sac_main \
    --env 'halfcheetah-medium-v0' \
    --logging.output_dir './experiment_output'

All available command options can be seen in JaxCQL/sac_main.py and JaxCQL/sac.py.

You can visualize the experiment metrics with viskit:

python -m viskit './experiment_output'

and simply navigate to http://localhost:5000/

This codebase can also log to W&B online visualization platform. To log to W&B, you first need to set your W&B API key environment variable:

export WANDB_API_KEY='YOUR W&B API KEY HERE'

Then you can run experiments with W&B logging turned on:

python -m JaxCQL.conservative_sac_main \
    --env 'halfcheetah-medium-v0' \
    --logging.output_dir './experiment_output' \
    --logging.online

In order to save your time and compute resources, I've done a sweep of JaxCQL on certain D4RL environments with various min Q weight values. The results can be seen here. You can choose the environment to visualize by filtering on env. The results for each cql.cql_min_q_weight on each env is repeated and average across 3 random seeds.

If you find this open source release useful, please reference in your paper:

@article{geng2022jaxcql,
  title={JaxCQL: a simple implementation of SAC and CQL in JAX},
  author={Xinyang Geng},
  year={2022},
  url={https://github.com/young-geng/JaxCQL}
}

The project organization is inspired by TD3. The SAC implementation is based on rlkit. THe CQL implementation is based on CQL. The viskit visualization is taken from viskit, which is taken from rllab.