RL Environments in JAX

Reinforcement Learning environments written in JAX.

  • Brax  GitHub Stars for Brax

    Massively parallel rigidbody physics simulation on accelerator hardware.

  • Evojax  GitHub Stars for Evojax

    EvoJAX is a scalable, general purpose, hardware-accelerated neuroevolution toolkit.

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

  • JaxMARL  GitHub Stars for JaxMARL

    Multi-Agent Reinforcement Learning with 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

  • Pgx  GitHub Stars for Pgx

    Vectorized RL game environments in JAX

  • popjym  GitHub Stars for popjym

    POPGym Library in JAX

  • XLand-MiniGrid  GitHub Stars for XLand-MiniGrid

    JAX-accelerated Meta-Reinforcement Learning Environments Inspired by XLand and MiniGrid