Add Gymnasium Wrapper

[vyeevani]

Lots of other frameworks seem to be switching to gymnasium. It'd be nice to also have a gymnasium wrapper in brax.

[btaba]

Hi @vyeevani, is this what you're looking for or something else?

https://github.com/google/brax/blob/16304037a36b1d9c8c0b3084f57d1159627b636b/brax/envs/wrappers/gym.py#L27

[vyeevani]

It’s that except using Gymnasium instead of Gym. The imports would change and also the returns are different for the step API. For example, you’d need to also return truncations (that you can get from the EpisodeWrapper)

It's a simple enough lift, pretty much just headers, error checking, and pulling things out of info.

I think it's worth it because OpenAI's Gym isn't going to be maintained anymore with the aim of switching everyone over to Gymnasium (https://gymnasium.farama.org/)

[btaba]

Ok, yeah the APIs are subtly different. Feel free to open up a PR with brax/envs/wrappers/gymnasium.py

[Surya-77]

Hi there, I noticed this open issue about adding a Gymnasium wrapper for Brax. It's been around 4 months since the last comment inviting @vyeevani to open a PR. I'm curious if there have been any updates or progress on implementing the Gymnasium wrapper? It would be helpful to get a status update on this.