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.
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: )