Official Docs

Links to official documentations for JAX, RL and related tools.

  • Equinox

    Equinox is your one-stop JAX library, for everything you need that isn't already in core JAX.

  • Flax

    Neural Networks for JAX

  • Jax

    Official JAX docs

  • Optax

    Optax is a gradient processing and optimization library for JAX.