search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.14k stars 234 forks source link

Add Gymnasium Wrapper #421

Open vyeevani opened 7 months ago

vyeevani commented 7 months ago

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

btaba commented 7 months ago

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 commented 7 months ago

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 commented 7 months ago

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

Surya-77 commented 3 months ago

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.