brax.io.mjcf loader incompatible with MuJoCo #649
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:
import mujoco
import brax.io.mjcf
scene = """
<mujoco>
<worldbody>
<body name="cart">
<inertial pos="0 0 0" mass="1" diaginertia="0.001 0.001 0.001"/>
<joint name="slider" type="slide" axis="1 0 0" range="-1 1" frictionloss="0"/>
<geom name="red_box" type="box" size=".05 .05 .05" rgba="1 0 0 1"/>
<body name="pole">
<joint type="hinge" axis="0 1 0" pos="0 0 0" frictionloss="0.01"/>
<geom type="cylinder" size="0.01" rgba="0 1 0 1" fromto="0 0 0 0 0 -0.27"/>
<inertial pos="0 0 -0.13" mass="0.040" diaginertia="0.001 0.001 0.001"/>
<body name="mass" pos="0 0 -0.27">
<geom type="sphere" size="0.02" pos="0 0 0"/>
<inertial pos="0 0 0" mass="0.075" diaginertia="0.001 0.001 0.001"/>
</body>
</body>
</body>
</worldbody>
</mujoco>
"""
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 ----
<mujoco._structs.MjModel object at 0x7f9d7ba23630>
---- testing brax (version 0.13.0) import ----
Traceback (most recent call last):
File "/home/user/dev/cartpole-sim/bug.py", line 28, in <module>
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,
import mujoco
import brax.io.mjcf
scene = """
<mujoco>
<worldbody>
<body name="cart">
<inertial pos="0 0 0" mass="1" diaginertia="0.001 0.001 0.001"/>
<joint name="slider" type="slide" axis="1 0 0" range="-1 1" frictionloss="0"/>
<geom name="red_box" type="box" size=".05 .05 .05" rgba="1 0 0 1"/>
<body name="pole">
<joint type="hinge" axis="0 1 0" pos="0 0 0" frictionloss="0.01"/>
<geom type="cylinder" size="0.01" rgba="0 1 0 1" fromto="0 0 0 0 0 -0.27"/>
<inertial pos="0 0 -0.13" mass="0.040" diaginertia="0.001 0.001 0.001"/>
</body>
</body>
</worldbody>
</mujoco>
"""
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 ----
<mujoco._structs.MjModel object at 0x7f72e665f130>
---- 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.