search

google / brax

Massively parallel rigidbody physics simulation on accelerator hardware.
Apache License 2.0
2.14k stars 234 forks source link

VectorGymWrapper Support with MjxEnv #414

Closed SumeetBatra closed 8 months ago

SumeetBatra commented 8 months ago

Hi folks,

I am trying to use existing brax (0.9.3) wrappers with the new MjxEnv (Mujoco 3). I've set up my code following the examples from the mujoco jupyter notebook for defining an MjxEnv: https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb

And then proceed to wrap the MjxEnv with a VectorGymWrapper. Unfortunately this leads to errors:

AttributeError: 'Model' object has no attribute 'actuator' on line 118 of gym.py:

action = jax.tree_map(np.array, self._env.sys.actuator.ctrl_range)

It looks like self._env.sys contains attributes actuator_xyz ex. actuator_ctrlrange instead of actuator.ctrl_range. I'm guessing this is a bug, unless VecGymWrapper is no longer going to be supported?

btaba commented 8 months ago

Hi @SumeetBatra , thanks for posting. VectorGymWrapper in this repo is written to interface with a brax System, not with mjx.Model which is part of the mujoco/mjx repo. brax.System and mjx.Model are two different APIs. We hope to bring these closer together in the future, but for now they are mostly incompatible. Hope that clarifies

SumeetBatra commented 8 months ago

Got it, thanks. Am I correct in assuming that MjxEnv should still be compatible with other wrappers, such as ActionClipWrapper, TorchWrapper, etc, given that they are more generic, and that MjxEnv subclasses Brax's base Env?

btaba commented 8 months ago

Hi @SumeetBatra , I haven't tried interfacing MjxEnv with those wrappers, but it seems like it might just work, although YMMV