Todo: Look into adding support for Brax for RL envs

[lebrice]

Todo: Look into adding support for Google's cool-looking Brax framework to get massively parallelized mujoco-like environments with low overhead. (and no closed-source code! 😄 )

• Check if Brax has a gym-compatible API Yes, it does!
• Test it out standalone, check for conflicting packages
• Look into conversions between pytorch and Jax tensors.. Maybe we could get the environment's observations directly on the GPU! (Edit: Done through Add wrapper for easy conversion to/from PyTorch Tensors google/brax#23)
• Add support for the Brax envs
• Add the brax learners as Methods in Sequoia, making sure they still work end-to-end.
• ... More to come

[erikfrey]

You may find this issue interesting in which we discuss converting between pytorch and jax: google/brax#2