Initial support for Jax env #2955
Closed
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
… to fix bug if environment doesn't use np_random in reset
…n the opposite case than was intended to (openai#2871)" This reverts commit 519dfd9.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
With the future addition of Brax then this PR adds a new type of Env,
JaxEnvalong withJaxStatethat follows the functional API suggested in #2954This PR introductions three new classes
JaxState- This is the dataclass object that contains all of the information required for the environment to take a step. The state is created during theresetfunction.JaxEnv- This is the subclass ofEnvthat requires functions for theresetandstepfunctions. Using these functions, they can be JIT hardware accelerated by default and allows easier vectorisationVectoriseJaxEnv- AllowsJaxEnvto be passed which will contain the statelessresetandstepfunctions. Using these functions, we can usejax.vmaporjax.pmapfor parallelisation of theresetandstep. This implementation does not follow use the currentVectorEnvdue to the number of unnecessary components however this can be changed