-
Notifications
You must be signed in to change notification settings - Fork 232
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Gym API? #14
Comments
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 https://github.com/google/brax/blob/main/brax/envs/wrappers.py |
Thanks @erikfrey for the quick response. Interesting. I'd suggest creating a simple example notebook/script where this gets used.
|
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, 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 |
OpenAI/StableBaseline In Thank you! |
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 |
Thank you for the quick reply and for pointing to that example. |
Hey, noob question here:
Is there an example showing how to interface with brax the same way you would with a gym Env?
If not, then why? Is there any particular reason why you make your own Env class rather than create a new
gym.Env
subclass?The text was updated successfully, but these errors were encountered: