brax.io.mjcf loader incompatible with MuJoCo
- Dominant language
- Jupyter Notebook
- Stars
- 3.2k
- Forks
- 349
- PR merge metrics
- No merged PRs in 30d
Description
Hello, I was working on a simple cartpole scene in MJCF, and it worked in the MuJoCo simulator and Python import function, but trying to import it using Brax's own `brax.io.mjcf.load()` or `brax.io.mjcf.loads()` doesn't work. I am on MuJoCo `3.3.7` and Brax `0.13.0`.
I've reduced it down to the following code, which is a scene with 3 nested bodies, each with its own `inertial` element:
```python
import mujoco
import brax.io.mjcf
scene = """
"""
print(f"---- testing mujoco (version {mujoco.__version__}) import ----")
print(mujoco.MjModel.from_xml_string(scene))
print(f"---- testing brax (version {brax.__version__}) import ----")
brax.io.mjcf.loads(scene)
```
Which gives the output:
```
Failed to import mujoco_warp: No module named 'warp'
---- testing mujoco (version 3.3.7) import ----
---- testing brax (version 0.13.0) import ----
Traceback (most recent call last):
File "/home/user/dev/cartpole-sim/bug.py", line 28, in
brax.io.mjcf.loads(scene)
File "/home/user/micromamba/envs/brax/lib/python3.12/site-packages/brax/io/mjcf.py", line 509, in loads
mj = mujoco.MjModel.from_xml_string(xml, assets=assets)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: XML Error: Schema violation: unique element 'inertial' found 2 times
Element 'body', line 7
```
However, if I just remove the deepest nested body so that it's only 2 nested elements,
```python
import mujoco
import brax.io.mjcf
scene = """
"""
print(f"---- testing mujoco (version {mujoco.__version__}) import ----")
print(mujoco.MjModel.from_xml_string(scene))
print(f"---- testing brax (version {brax.__version__}) import ----")
brax.io.mjcf.loads(scene)
```
Then the import works, albeit with the warning about Brax's pipelines being not actively maintained:
```
Failed to import warp: No module named 'warp'
Failed to import mujoco_warp: No module named 'warp'
---- testing mujoco (version 3.3.7) import ----
---- testing brax (version 0.13.0) import ----
/home/user/micromamba/envs/brax/lib/python3.12/site-packages/brax/io/mjcf.py:480: UserWarning: Brax System, piplines and environments are not actively being maintained. Please see MJX for a well maintained JAX-based physics engine: https://github.com/google-deepmind/mujoco/tree/main/mjx. For a host of environments that use MJX, see: https://github.com/google-deepmind/mujoco_playground.
warnings.warn(
```
Having `inertial` elements to specify the masses of nested bodies seems like a rather basic thing that one would want to do with MJCF. What's the situation with the Brax MJCF importer pipeline? It also seems weird that it would work with 2 level nesting but not 3 level nesting. Is the Brax MJCF importer supposed to be now deprecated, and using scene importing logic from MuJoCo itself is preferred? The Brax loading pipeline is still used in some of the tutorial Colab notebooks, it seems.
I am interested to know how the whole thing with deprecating parts of Brax and moving to MJX backend is going, and whether stripping out the aforementioned code and/or updating the documentation (i.e. Colab notebooks I suppose) is a priority - as I am interested in learning more about Brax and potentially contributing.
Contributor guide
Assessment
This issue has not been assessed yet.