[RFC] MJX environment prototype
#834
Conversation
`tests/env/mujoco/test_mojoco_v3.py`
gymnasium/envs/mujoco/f.py
Outdated
| - exclude_current_positions_from_observation | ||
| ) | ||
|
|
||
| self.observation_space = gymnasium.spaces.Box( # TODO use jnp when and if `Box` supports jax natively |
There was a problem hiding this comment.
Yeah, I guess we'll need to make spaces compatible with jax to be able to sample in jit compiled training function
gymnasium/envs/mujoco/f.py
Outdated
| self.camera_id = camera_id | ||
| self.camera_name = camera_name | ||
|
|
||
| return MujocoRenderer( |
There was a problem hiding this comment.
No it is not, and the render functions will not be JITable for a while (if ever), because of the structure of the deepming/mujoco project
gymnasium/envs/mujoco/f.py
Outdated
| ) | ||
|
|
||
|
|
||
| # TODO add vector environment |
There was a problem hiding this comment.
These should be super simple. It's essentially vmap step and reset
|
Thanks for this work. I have a few question regarding the WIP. In particular is there any training script that support the MJX Envs. I am asking this, since the return structure of some of the functions, and no mention of truncated in the environments in the PR suggest that, currently something other way of using these environments. |
|
I have paused development, waiting for
mjx==3.2& a few updates onFuncEnvrelated proposal (read first): #833
Feel free to comment on anything you want, but in particular check the
NOTEcommentsChecklist
MJXclassInvertedDoublePendulum,InvertedPendulum)HalfCheetah,Hopper,Walker2D)Reacher,Pusher)Swimmer)Ant)Humanoid,HumanoidStandup)