Jax RL Implementations

Reinforcement Learning in JAX.

  • JaxGCRL  GitHub Stars for JaxGCRL

    Goal-Conditioned Reinforcement Learning with JAX

  • jaxrl  GitHub Stars for jaxrl

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

  • 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

  • skrl  GitHub Stars for skrl

    Modular reinforcement learning library (on PyTorch and JAX) with support for NVIDIA Isaac Gym, Omniverse Isaac Gym and Isaac Lab

  • Stoix  GitHub Stars for Stoix

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