Magi RL library in JAX

Installation | Agents | Examples | Contributing | Documentation

Magi is a RL library in JAX that is fully compatible with Acme.

In addition to the features provided by Acme, Magi offers implementation of RL agents that are not found in the Acme repository as well as providing useful tools for integrating experiment logging services such as WandB.

Note: Magi is in alpha development so expect breaking changes!

Magi currently depends on HEAD version of dm-acme instead of the latest release version on PyPI which is fairly old.

Installation

Create a new Python virtual environment

python3 -m venv venv
source venv/bin/activate

Install dependencies with the following commands.

pip install -U pip setuptools wheel
# Magi depends on latest version of dm-acme.
# The dependencies in setup.py are abstract which allows you to pin
# a specific version of dm-acme.
# The following command installs the latest version of dm-acme
pip install 'git+https://github.com/deepmind/acme.git#egg=dm-acme[jax,tf,examples]'
# Install magi in editable mode, with additional dependencies.
# In case you need to run examples on GPU, you should install the
# GPU version of JAX with a command like the following
pip install 'jax[cuda]<0.4' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
pip install -e '.[jax]'

The base installation for magi does not list TensorFlow/JAX as a dependency. However, note that JAX requires platform-specific installation (CPU/GPU and CUDA versions). Furthermore, Acme depends on Reverb and LaunchPad which requires them to be pinned against specific versions of TensorFlow. This should be handled if you use install dm-acme with [jax,tf] extras. However, you can also use install with different versions of TensorFlow/Reverb/Launchpad. In that case, you should omit the extras and find compatible versions and pin those versions accordingly.

If for some reason installation fails, first check out GitHub Actions badge to see if this fails on the latest CI run. If the CI is successful, then it's likely that there are some issues to setting up your own environment. Refer to .github/workflows/ci.yaml as the official source for how to set up the environment.

Agents

magi includes popular RL algorithm implementation such as SAC, DrQ, SAC-AE and PETS. Refer to magi/agents for a full list of agents.

Examples

Check out examples where we include examples of using our RL agents on popular benchmark tasks.

Testing

On Linux, you can run tests with

nox test

Contributing

Refer to CONTRIBUTING.md.

Acknowledgements

Magi is inspired by many of the open-source RL projects out there. Here is a (non-exhaustive) list of related libraries and packages that Magi references:

License

Apache License 2.0

Citation

If you use Magi in your work, please cite us according to the CITATION file. You may learn more about the CITATION file from here.

ethanluoyc/magi

ethanluoyc

Reviews

Repository Details