GitHub Repositories

A collection of useful GitHub repositories for JAX-based Reinforcement Learning.

Jax RL Implementations

Reinforcement Learning in JAX.

  • Mava  GitHub Stars for Mava

    A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX.

  • PureJaxRL  GitHub Stars for PureJaxRL

    Really Fast End-to-End Jax RL Implementations.

  • RLax  GitHub Stars for RLax

    Building blocks for implementing RL agents.

  • SBX  GitHub Stars for SBX

    Stable Baselines Jax (SB3 + Jax) RL algorithms

  • Stoix  GitHub Stars for Stoix

    A research-friendly codebase for fast experimentation of single-agent reinforcement learning in JAX.


Network Libraries

Neural networks and scientific computing in JAX.

  • Equinox  GitHub Stars for Equinox

    Elegant easy-to-use neural networks + scientific computing in JAX.

  • Flax  GitHub Stars for Flax

    High-performance neural network library for JAX.


Training

Repositories tailored to optimize and enhance the training process in JAX.

  • EasyDeL  GitHub Stars for EasyDeL

    Accelerate, Optimize performance with streamlined training and serving options with JAX.

  • Flashbax  GitHub Stars for Flashbax

    Accelerated Replay Buffers in JAX.


Optimization Libraries

Gradient processing and optimization in JAX.

  • Optax  GitHub Stars for Optax

    Optax is a gradient processing and optimization library for JAX.


RL Environments in JAX

Reinforcement Learning environments written in JAX.

  • Brax  GitHub Stars for Brax

    Massively parallel rigidbody physics simulation on accelerator hardware.

  • Gymnax  GitHub Stars for Gymnax

    RL Environments in JAX.

  • Jumanji  GitHub Stars for Jumanji

    A diverse suite of scalable reinforcement learning environments in JAX.

  • Kinetix  GitHub Stars for Kinetix

    Reinforcement learning on general 2D physics environments in JAX.


Tools

Additional repositories that may be useful.

  • Haliax  GitHub Stars for Haliax

    Named Tensors for Legible Deep Learning in JAX.

  • Treescope  GitHub Stars for Treescope

    An interactive HTML pretty-printer for machine learning research in IPython notebooks.