MJX Training Implementation #1
Merged
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
budzianowski
approved these changes
May 24, 2024
| @@ -0,0 +1,10 @@ | |||
| from brax import envs | |||
There was a problem hiding this comment.
If you can omit putting stuff here, that would be preferable.
| @@ -0,0 +1,152 @@ | |||
| import jax | |||
There was a problem hiding this comment.
Adding one line explanation would be useful.
| @@ -0,0 +1,152 @@ | |||
| import jax | |||
| import jax.numpy as jp | |||
| Observations of the environment. | ||
| """ | ||
| position = data.qpos | ||
| if self._exclude_current_positions_from_observation: |
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.
The following is copied from kscalelabs/sim#14
This PR introduces a new way to massively scale up locomotion training. Building upon Brax, it ultimately uses the MJX physics engine for simulation.
Structure
Specifically, this PR includes the following directories:
Envsincludes two types of Brax environments: DefaultHumanoidEnv and StompyEnv. Each environment includes a main class which implements the Brax environment interface and utilizes MJX for all physics calculations. One important thing to note is that reward functions are modular, allowing for quick experimentation.Experimentsincludes two .yaml files that include sample configurations for model training.Utilsinclude default values, rendering rollouts, etc.Weightscurrently include default humanoid weights (for locomotion) that should work out of the boxtrain.py and play.py both integrate with wandb. train.py utilizes the Brax implementation of PPO for now, but can be easily customized if needed.
Performance Samples
Training Curves
Example humanoid robot walking in MJX
https://github.com/kscalelabs/sim/assets/43460304/8e12b0e6-48ea-4af0-8283-1dc4880767b4
Humanoid trained in MJX, eval in CPU-based MuJoCo
https://github.com/kscalelabs/sim/assets/43460304/7f158aeb-6bc9-4056-bd1d-12882adbd13c