google / google/brax

Replacing gym's Mujoco envs with brax envs

Open
#49 70 comments 8 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.