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! :smile: )

• [X] Check if Brax has a gym-compatible API Yes, it does!
• [ ] Test it out standalone, check for conflicting packages
• [x] Look into conversions between pytorch and Jax tensors.. Maybe we could get the environment's observations directly on the GPU! (Edit: Done through https://github.com/google/brax/pull/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: https://github.com/google/brax/issues/2