Replacing gym's Mujoco envs with brax envs
- Dominant language
- Jupyter Notebook
- Stars
- 3.2k
- Forks
- 349
- PR merge metrics
- No merged PRs in 30d
Description
Had a conversation with @jkterry1 on https://github.com/openai/gym/issues/2366, and it appears brax would also be a great alternative for the mujoco envs replacement.
To help with this transition. I made an attempt to try out brax with pytorch. Here is a basic report: https://wandb.ai/costa-huang/brax/reports/Brax-as-Pybullet-replacement--Vmlldzo5ODI4MDk. The source code is here: https://github.com/vwxyzjn/cleanrl/blob/mybranch/cleanrl/brax/readme.md
One of the biggest issue with the brax adoption is the env normalization:
* gym doesn't have a normalization wrapper
* sb3 has a normalization wrapper but brax does not have a sb3 vector env api
* brax's normalization is implemented in the training side (https://github.com/google/brax/blob/main/brax/training/normalization.py)
I think going forward, probably the best way to fix this is to refactor the brax training side's normalization to the environment side. This in the future will also help throughput with the `JaxToTorchWrapper`. Otherwise, the observation will go from GPU to CPU for gym or sb3's normalization wrapper, then GPU again for torch, which just doesn't make sense.
One small thing is that given the brax environment directly produces the vector env, there is also no way to inject a `ClipActionsWrapper(env)`, which may or may not have a performance impact. That said, this can be implemented in the training side with ease.
Contributor guide
Assessment
This issue has not been assessed yet.