Gym API? about brax HOT 6 CLOSED

Hello Fabrice,

JAX's functional semantics make it a bit tricky to implement the gym API where the Env itself is supposed to maintain state, and still work for end-to-end training.

But the good news is that you can easily create gym Environments by calling create_gym_env in place of create_env, which will wrap the environment with a gym.Env. See:

https://github.com/google/brax/blob/main/brax/envs/wrappers.py

from brax.

Thanks @erikfrey for the quick response.

Interesting. I'd suggest creating a simple example notebook/script where this gets used.
Other questions:

• Does this GymWrapper support rendering? If not, are there plans for adding support for it (which seems to already be present for the other envs?)
• What is the right way of "scaling up"/vectorizing this GymWrapper? (I assume you wouldn't want users to pass it to one of the gym.vector.VectorEnv subclasses, for instance?)
• How well would you expect this GymWrapper to "scale" compared to the jax-only version?

from brax.

Right now we only support visualization via the three.js webview. Actual rendering - as in a python function that returns an image buffer - is on our roadmap! But don't have an ETA yet.

re: vectorizing, we have only used to gym wrapper for CPU training on multi-process envs, so no batching dimension in the gym env. That said, VectorEnv class looks good. We could extend the create_gym_env to return VectorEnv if you specify a batch dimension. Feel free to submit a PR if that would be helpful to you.

performance: as long as you're careful to avoid copying data between CPU and accelerator (for example, by ensuring your learner is jitted too, and that the actions you pass to step are on device), then there should be no runtime performance penalty to using a gym wrapper, just a very small jit time cost for a bit of extra function tracing.

from brax.

OpenAI/StableBaseline VectorEnv automatically calls reset on a worker when the episode finished.

In brax however I guess that because how vmap works one needs to check if a given episode in the batch is done and treat this on the agent side (e.g.: ignore data from finished episodes in the batch) since there is no way to reset a single episode in the batch, right?

Thank you!

from brax.

Hi @floringogianu - yes, we leave it up to the user to decide how to reset. For the case of training, we actually do reset individual episodes within a batch. We support this in the training env wrapper, see here:

https://github.com/google/brax/blob/main/brax/training/env.py#L57

from brax.

Thank you for the quick reply and for pointing to that example.

from brax.