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.