Bug in `brax` MJCF Parsing

[alberthli]

The brax model loading function doesn't correctly parse some MJCF files. I'm testing the Franka Panda model from the mujoco_menagerie repo, and the following fails (after removing the line <option integrator="implicitfast" impratio="10"/> from the MJCF file):

from brax.io.mjcf import load
model = load("panda_nohand.xml")  # ValueError: XML Error: World body cannot have inertia

However, directly loading the model with mujoco and then using the load_model function works:

import mujoco
from brax.io.mjcf import load_model
mj_model = mujoco.MjModel.from_xml_path("panda_nohand.xml")
model = load_model(mj_model)

[btaba]

Hi @alberthli , thanks for the bug report! The ValueError: XML Error: World body cannot have inertia bug is coming from mujoco.MjModel.from_xml_string vs the mujoco.MjModel.from_xml_path that you are referencing above. This is implemented in MJ. Since we edit the xml (to fuse bodies without joints), we need to use from_xml_string in brax

I would remove the world body inertia for usage in brax https://github.com/deepmind/mujoco_menagerie/blob/main/franka_emika_panda/panda.xml#L121-L122

[willthibault]

Hi @btaba, I am experiencing a similar error when loading a model for my MjxEnv training:

  File "/usr/local/lib/python3.9/dist-packages/brax/envs/__init__.py", line 63, in get_environment
    return _envs[env_name](**kwargs)
  File "/home/workspace/train.py", line 141, in __init__
    self.brax_sys = mjcf.load(path).replace(dt=self._dt)
  File "/usr/local/lib/python3.9/dist-packages/brax/io/mjcf.py", line 540, in load
    mj = load_mjmodel(path)
  File "/usr/local/lib/python3.9/dist-packages/brax/io/mjcf.py", line 534, in load_mjmodel
    mj = mujoco.MjModel.from_xml_string(xml, assets=assets)
ValueError: XML Error: Schema violation: unique element 'inertial' found 19 times

Essentially, I am trying to freeze certain joints in my model. Normally, I can remove the from the MuJoCo XML and that is sufficient, but this results in multiple inertial elements rather than one inertial element in the frozen part.

Is there a way I can keep the inertial impact of the frozen joints? Maybe there is a better way to do this. I don't want the frozen joints to impact the size of the model I'm training (ex. having the joints actuated and setting the same position control value every step).

Thanks!

[btaba]

Hi @willthibault ,it sounds like you want a way to combine inertial elements so that you can remove joints, but keep the same inertial effect of all the bodies. Please correct me if I'm wrong.

The easiest way I can think of is to modify the joint limit range, or add joint equality constraints.

https://github.com/google-deepmind/mujoco_menagerie/blob/8d283c99f7de77ae5a703b21504d07593a461a37/aloha/aloha.xml#L281-L284

However, as you mention, this will impact the runtime of your model, since joint equality constraints need to be computed.

The harder way is to sum the mass, get the new center of mass, and then use the parallel axis theorem to recompute the total inertia about the new center of mass. Hope this helps.

[willthibault]

Hi @btaba, thanks for the reply.

Yes, that is exactly what I am trying to accomplish.

I will give the joint limit range and joint equality constraints a try to see what sort of impact that has.

I also considered recomputing the inertia, but was hoping for something flexible so I could essentially freeze the joints at different configurations.

As I understand, the action space in Brax is decided based on the number of actuated joints in the model. Is there a way to shrink the action space to only a set number of actuated joints in the model or change the action space to be something other than joint commands? Maybe this is out of the scope of Brax, but just wanted to check.

Thanks again!

[btaba]

RE action space: are you trying to remove an actuator from the model, without removing an actuator from the XML? This could work with torque actuators, where you just don't apply torque to some joints.

For actuators other than joint actuators, see transmission in MJ https://mujoco.readthedocs.io/en/stable/computation/index.html#transmission (namely site actuators which are implemented in MJX, which effectively apply a force/torque at a location on a body)

[willthibault]

Yes, that is what I mean. I'm hoping to reduce the action size here and set the "removed" joints to a value during a physics step.

The site actuators look useful. Presumably the use of refsite for end effector control would be useful for an environment like the Fetch pick and place, correct? Must the actions always be some form of actuator value or could they be simply a set of values that could be used to compute actuator values for the physics step?

[btaba]

RE reducing action size, if you are using position control, it makes sense to set the ctrl to a fixed target value if you want the joint to be stationary. I would suggest implementing this kind of logic in the environment code

RE: site, yes they'd be useful for pick-and-place. The environment is the place to implement custom logic around the ctrl. It's up to the environment to decide what the action_size should be and what to do with it. Ultimately, within the environment, you should send the simulator a ctrl that corresponds to the actuators defined in the XML. Hope this helps

[willthibault]

Thank you very much clarification, this is very helpful!