Reproducing MuJoCo Menagerie hand-to-arm attachment example with the new model editing API

[hartikainen]

Intro

Hi!

I am a MuJoCo user working on manipulation.

My setup

MuJoCo:

$ python -c "import mujoco; print(mujoco.__version__)"
3.2.3

API:
Python

OS:

python -c "import platform; print('OS:', platform.system(), platform.release(), platform.machine())"
OS: Darwin 23.3.0 arm64

My question

I've been previously using the dm_control.mjcf API and have been testing out the new native model editing API (in Python) recently.

I'd like to programmatically attach a gripper to my arm model, similarly as done in the MuJoCo Menagerie example.

I've added a couple of my test attempts below. The body/attach documentation mentions explicitly that "the world body cannot be attached", which is why the v1 implementation below doesn't work. One workaround I see would be to attach all the body-elements from the gripper model manually, as I've done in the v2 implementation below. That one is problematic though (see my comments in the function).

I wonder what the recommended way of achieving such model attachment is? It would be great to see a couple of examples on how to do this.

Minimal model and/or code that explain my question

Here's my attempt following the Menagerie example:

import mujoco
import numpy as np
from absl import app
from robot_descriptions import fr3_mj_description, shadow_dexee_mj_description


def find_site_in_body(body: mujoco.MjsBody, name: str) -> mujoco.MjsSite:
    site = body.first_site()
    while site is not None:
        if site.name == name:
            return site
        site = body.next_site(site)

    child_body = body.first_body()
    while child_body is not None:
        site = find_site_in_body(child_body, name)
        if site is not None:
            return site
        child_body = body.next_body(child_body)
    return None


def attach_hand_to_arm(
    arm_spec: mujoco.MjSpec,
    hand_spec: mujoco.MjSpec,
) -> None:
    """Attach a hand to an arm.

    Args:
      arm_spec: The `mujoco.MjSpec` of the arm.
      hand_spec: The `mujoco.MjSpec` of the hand.

    Raises:
      ValueError: If the arm does not have a site named "attachment_site".
    """

    def v0(arm_spec: mujoco.MjSpec, hand_spec: mujoco.MjSpec) -> None:
        # This would be similar to the `dm_control.mjcf`-style attachment but the
        # `attachment_site.attach` does not exist in the new model editing API.
        attachment_site = find_site_in_body(arm_spec.worldbody, "attachment_site")
        if attachment_site is None:
            raise ValueError("No attachment site found in the arm model.")
        attachment_site.attach(hand_spec)

    def v1(arm_spec: mujoco.MjSpec, hand_spec: mujoco.MjSpec) -> None:
        # This one raises an error:
        # ```
        # ValueError: Error: incompatible id in body array, position 10
        # Element name 'world', id 0
        # ```
        # This is expected as the `body/attach` documentation explicitly states that the
        # `worldbody` cannot be attached. https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-attach
        attachment_site_parent_body = arm_spec.find_body(
            "fr3_link7"
        )  # Is there shortcut for this?
        attachment_frame = attachment_site_parent_body.add_frame()
        attachment_frame.attach_body(arm_spec.worldbody, "", "")

    def v2(arm_spec: mujoco.MjSpec, hand_spec: mujoco.MjSpec) -> None:
        # This one succeeds but is lacking because the `nu` and `nq` information doesn't
        # seem to exist in the hand spec. Also, the method of attaching each body manually
        # feels a bit clumsy and indeed the hand is attached to a wrong position with the
        # way done here.

        attachment_site_parent_body = arm_spec.find_body(
            "fr3_link7"
        )  # Is there really no way of finding the parent of the attachment site?
        attachment_frame = attachment_site_parent_body.add_frame()

        body = hand_spec.worldbody.first_body()
        while body is not None:
            attachment_frame.attach_body(body, body.name, "")
            body = hand_spec.worldbody.next_body(body)

        # Expand the ctrl and qpos keyframes to account for the new hand DoFs.
        # TODO(hartikainen): Where do I get the `nu` and `nq` information from?
        nu = nq = 12
        arm_key = next((key for key in arm_spec.key if key.name == "home"), None)
        if arm_key is not None:
            hand_key = next((key for key in hand_spec.key if key.name == "home"), None)
            if hand_key is None:
                arm_key.ctrl = np.concatenate([arm_key.ctrl, np.zeros(nu)])
                arm_key.qpos = np.concatenate([arm_key.qpos, np.zeros(nq)])
            else:
                arm_key.ctrl = np.concatenate([arm_key.ctrl, hand_key.ctrl])
                arm_key.qpos = np.concatenate([arm_key.qpos, hand_key.qpos])

    v2(arm_spec, hand_spec)


def main(unused_argv):
    hand_spec = mujoco.MjSpec()
    hand_spec.from_file(shadow_dexee_mj_description.MJCF_PATH)
    arm_spec = mujoco.MjSpec()
    arm_spec.from_file(fr3_mj_description.MJCF_PATH)

    attach_hand_to_arm(arm_spec, hand_spec)
    mj_model = arm_spec.compile()
    # mujoco.mj_saveModel(mj_model, "/tmp/franka-with-dexee.mjb")
    print(arm_spec.to_xml())


if __name__ == "__main__":
    app.run(main)

v0 fails with:

$ python mj_spec_attach_hand_to_arm.py
Traceback (most recent call last):
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 104, in <module>
    app.run(main)
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 97, in main
    attach_hand_to_arm(arm_spec, hand_spec)
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 88, in attach_hand_to_arm
    v0(arm_spec, hand_spec)
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 43, in v0
    attachment_site.attach(hand_spec)
    ^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'mujoco._specs.MjsSite' object has no attribute 'attach'

v1 fails with:

$ python mj_spec_attach_hand_to_arm.py
Traceback (most recent call last):
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 104, in <module>
    app.run(main)
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
             ^^^^^^^^^^
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 98, in main
    mj_model = arm_spec.compile()
               ^^^^^^^^^^^^^^^^^^
ValueError: Error: incompatible id in body array, position 10
Element name 'world', id 0

v2 technically succeeds but for some reason the code segfaults after the model has been compiled and printed. It's still lacking as detailed in my comments in the v2-function above.

$ python mj_spec_attach_hand_to_arm.py

[xml output]

Fatal Python error: Segmentation fault

Current thread 0x00000001d9a01c40 (most recent call first):
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 254 in _run_main
  File "$HOME/conda/envs/mujoco/lib/python3.11/site-packages/absl/app.py", line 308 in run
  File "$HOME/tmp/scripts/mj_spec_attach_hand_to_arm.py", line 104 in <module>

Extension modules: numpy._core._multiarray_umath, numpy.linalg._umath_linalg (total: 2)
[1]    105 segmentation fault  python /mj_spec_attach_hand_to_arm.py

Confirmations

• I searched the latest documentation thoroughly before posting.
• I searched previous Issues and Discussions, I am certain this has not been raised before.

[yuvaltassa]

I think there might be a bug report and two feature requests here?

1. Feature request: Attach the world body.
2. Feature request: Attach to a site (i.e., using a site as an attachment frame).
3. Bug report: segfault described above.

Is that right?

[hartikainen]

Yeah, assuming that there's no other intended method of achieving this, I think that's right. I can split this into issues for better testability and easier tracking.

[quagla]

1. What happens if you attach it with a prefix? I'm actually debating whether to disable the option of not using prefixes..
2. It's very easy to add a frame at the same position of the site and attach the body to the frame, is this an acceptable solution?
3. Yes please, it would be very nice if you could open an issue with a minimal example, thank you!

[hartikainen]

1. The behavior doesn't change with/without prefix. I guess the prefixes are important if you attach multiple bodies with same identifiers: you could automatically attach unique identifier to repeated names but it's nice for the users to be able to set them manually.
2. Yeah, I was expecting similar behavior as with dm_control.mjcf-API. Their body-to-site attachment behavior is effectively as you describe: a frame with the same position and orientation as the site is created, and the body is then attached to this frame.
3. See MjSpec: model destructor segfaults when attaching a body from spec with plugins #2100

[quagla]

By the way, the keyframe resizing should be done automatically for you (if you don't have nested attachments), see https://github.com/google-deepmind/mujoco/blob/main/test/user/user_api_test.cc#L643

[hartikainen]

I couldn't get this to work. After attach_hand_to_arm, I get arm_spec.key[0].qpos.size == 7 and the compilation fails with:

Traceback (most recent call last):
  File "mj_spec_attach_hand_to_arm_v2.py", line 59, in <module>
    main()
  File "mj_spec_attach_hand_to_arm_v2.py", line 53, in main
    mj_model = arm_spec.compile()
               ^^^^^^^^^^^^^^^^^^
ValueError: Error: incompatible id in body array, position 17
Element name 'base', id 1, line 69

I initially thought that perhaps this is because the arm_spec has a key whereas the hand_spec does not. This is not the case, though, and the result is the same whether the hand_spec has a key or not.

Another problem I noticed when debugging this is a weird behavior when attaching multiple bodies from a spec with a keyframe. See #2101.

[hartikainen]

Just wanted to report back here that this is working as expected from ebe60b9 (not from 3.2.4 that's on PyPI though). Here's a snippet if anyone else is trying to do something similar:

import mujoco
import numpy as np
from robot_descriptions import fr3_mj_description
from robot_descriptions import shadow_dexee_mj_description


def main():
    hand_spec = mujoco.MjSpec.from_file(shadow_dexee_mj_description.MJCF_PATH)
    arm_spec = mujoco.MjSpec.from_file(fr3_mj_description.MJCF_PATH)

    attachment_site = next(s for s in arm_spec.sites if s.name == "attachment_site")
    attachment_site.attach(hand_spec.worldbody, "hand_base", "")

    mj_model = arm_spec.compile()
    mujoco.mj_saveModel(mj_model, "/tmp/franka-with-dexee.mjb")
    print(arm_spec.to_xml())


if __name__ == "__main__":
    main()

And here's what the resulting model looks like:

Thanks guys!