Problem implementing an MJX environment from an .xml file (NotImplementedError)

[jamesheald]

I'm following the google colab MJX tutorial (https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/mjx/tutorial.ipynb) but instead of creating a Humanoid environment, I'm trying to create an environment from MyoSuite (https://github.com/MyoHub/myosuite) - Myosuite environments are built from Mujoco and wrapped in an OpenAI gym API.

Specifically, in the 'Humanoid Env' tab of the colab tutorial, I replace the path to the humanoid.xml file with a path to the .xml file of a MyoSuite environment by replacing the following lines of code:

path = epath.Path(epath.resource_path('mujoco')) / (
        'mjx/benchmark/model/humanoid'
    )
    mj_model = mujoco.MjModel.from_xml_path(
        (path / 'humanoid.xml').as_posix())

with:

    path = epath.Path(epath.resource_path('myosuite')) / (
        'simhive/myo_sim/finger'
    )
    mj_model = mujoco.MjModel.from_xml_path(
        (path / 'myofinger_v0.xml').as_posix()) 

I then register the environment under the name 'myofinger_mjx', call 'env = envs.get_environment('myofinger_mjx')' and get the following error message and Tracback:

/usr/local/lib/python3.10/dist-packages/mujoco/mjx/_src/device.py:108: UserWarning: Solver mjSOL_NEWTON is not supported, reverting to CG.
  warnings.warn(f'Solver {name} is not supported, reverting to CG.')
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
[<ipython-input-8-b333cd37f0b8>](https://localhost:8080/#) in <cell line: 3>()
      1 # instantiate the environment
      2 env_name = 'myofinger_mjx'
----> 3 env = envs.get_environment(env_name)
      4 
      5 # define the jit reset/step functions

4 frames
[/usr/local/lib/python3.10/dist-packages/brax/envs/__init__.py](https://localhost:8080/#) in get_environment(env_name, **kwargs)
     61     env: an environment
     62   """
---> 63   return _envs[env_name](**kwargs)
     64 
     65 

[<ipython-input-7-95c14c60b39d>](https://localhost:8080/#) in __init__(self, forward_reward_weight, ctrl_cost_weight, healthy_reward, terminate_when_unhealthy, healthy_z_range, reset_noise_scale, exclude_current_positions_from_observation, **kwargs)
     33         'physics_steps_per_control_step', physics_steps_per_control_step)
     34 
---> 35     super().__init__(mj_model=mj_model, **kwargs)
     36 
     37     self._forward_reward_weight = forward_reward_weight

[<ipython-input-6-ca9bba043d3a>](https://localhost:8080/#) in __init__(self, mj_model, physics_steps_per_control_step)
     18     self.model = mj_model
     19     self.data = mujoco.MjData(mj_model)
---> 20     self.sys = mjx.device_put(mj_model)
     21     self._physics_steps_per_control_step = physics_steps_per_control_step
     22 

[/usr/local/lib/python3.10/dist-packages/mujoco/mjx/_src/device.py](https://localhost:8080/#) in device_put(value)
    184 
    185   if isinstance(value, mujoco.MjModel):
--> 186     _validate(value)  # type: ignore
    187 
    188   init_kwargs = {}

[/usr/local/lib/python3.10/dist-packages/mujoco/mjx/_src/device.py](https://localhost:8080/#) in _validate(m)
    118     unsupported_typs = typs - set(_TYPE_MAP[mj_type])
    119     if unsupported_typs:
--> 120       raise NotImplementedError(f'{unsupported_typs} not implemented.')
    121 
    122   # check condim

NotImplementedError: {2} not implemented.

Can anybody help me to understand what is going on here?

Here is a link to my google colab code to reproduce the error: https://colab.research.google.com/drive/17QJDvNaFYzluhSDe2X3L5j81oCIm_3kd?usp=sharing).

Thanks in advance.

[btaba]

Hi @jamesheald , we have a recent PR that makes the error more readable https://github.com/google-deepmind/mujoco/commit/e4a9f535a4fe05c2120ae9cd7afb17665f917e6b

If you run with that change, the error will be able to give a better idea of what to change in the XML. Also see https://mujoco.readthedocs.io/en/stable/mjx.html#feature-parity

[jamesheald]

Thanks @btaba.

I now get the following message: 'NotImplementedError: [<mjtBias.mjBIAS_MUSCLE: 2>] not implemented.'

I assume this is because, under feature parity, the feature MUSCLE in category transmission is not supported (MyoSuite is a musculoskeletal model)? It looks like this feature is not in development either, which is a shame.

For what it's worth, I'd like to request support for features used by MyoSuite musculoskeletal models.

[btaba]

Thanks @jamesheald , noted! Feel free to send us a PR as well. I think this will be easier to implement once we have another transmission type (besides joint) implemented in MJX

[Vittorio-Caggiano]

+1 ... Would love to have this feature and integrate the possibility of GPU/TPU acceleration in MyoSuite!

Also, there is a growing community using MuJoCo for musculoskeletal modeling more holistically that would love to benefit from gpu/tpu acceleration.

[jamesheald]

Thanks @btaba, I will try to send a PR at some point.

[P-Schumacher]

+1 on the want for muscle implementations :)

[schlagercollin]

+1 for a request for muscle implementations!

It would also be helpful to know what blockers (if any) exist for the implementation. It sounds like there might be a few steps before muscles could be supported by MJX (e.g. other transmission types).

Thanks for all your hard work on supporting this package! 🙂

CC: @rick-warren @SSaltEng @jsmerelFB

[vikashplus]

+1 support for Muscles support

MSK systems are quite high dimensional. IMO, this is exactly where acceleration provided by MJX will shine. There is limited hope otherwise.

[rick-warren]

+1 for muscle support :)

[erikfrey]

Hello! We've heard you loud and clear and will prioritize support for muscles. Tendons will be done first, and muscles will come some time after that. I will update the documentation so that muscles are marked as in development.

[schlagercollin]

Hello! We've heard you loud and clear and will prioritize support for muscles. Tendons will be done first, and muscles will come some time after that. I will update the documentation so that muscles are marked as in development.

Thanks for the follow-up, @erikfrey! Awesome work so far. Really excited to get back into the JAX ecosystem and leverage the benefits of MJX for our projects. Cheers!

[kevinzakka]

Let's consolidate MJX feature request + feature tracking in #1179.