RL Environments in JAX

Reinforcement Learning environments written in JAX.

  • Brax  GitHub Stars for Brax

    Massively parallel rigidbody physics simulation on accelerator hardware.

  • Evorl  GitHub Stars for Evorl

    EvoRL is a fully GPU-accelerated framework for Evolutionary Reinforcement Learning, implemented with JAX

  • 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.

  • Navix  GitHub Stars for Navix

    Accelerated minigrid environments with JAX