google-deepmind / mujoco

Multi-Joint dynamics with Contact. A general purpose physics simulator.
https://mujoco.org
Apache License 2.0
8.25k stars 823 forks source link

`MjSpec` with attachment results in suspicious `sol{ref,imp}` values #2221

Closed hartikainen closed 1 week ago

hartikainen commented 1 week ago

Intro

Hi!

I am a MuJoCo user working on manipulation.

My setup

MuJoCo:

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

What's happening? What did you expect?

When I attach a model to another using MjsSite.attach, the joints' sol{ref,imp} values seem to change. I'm not sure if this is a bug or not, but it feels a little suspicious to me. The following code demonstrates what's happening. (The code is a bit verbose, although still relatively simple.)

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)
    hand_model = hand_spec.compile()
    arm_spec = mujoco.MjSpec.from_file(fr3_mj_description.MJCF_PATH)
    arm_model = arm_spec.compile()

    print("\nHand joints:")
    for joint in range(hand_model.njnt):
        joint_name = hand_model.joint(joint).name
        solref = hand_model.joint(joint).solref.tolist()
        solimp = hand_model.joint(joint).solimp.tolist()
        print(f"  {joint_name=}:\t{solref=}\t{solimp=}")

    print("\nArm joints:")
    for joint in range(arm_model.njnt):
        joint_name = arm_model.joint(joint).name
        solref = arm_model.joint(joint).solref.tolist()
        solimp = arm_model.joint(joint).solimp.tolist()
        print(f"  {joint_name=}:\t{solref=}\t{solimp=}")

    combined_spec = arm_spec.copy()

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

    combined_model = combined_spec.compile()

    joint_keys = {
        # "M0", This probably changes?
        "armature",
        "axis",
        "damping",
        "frictionloss",
        "limited",
        "margin",
        "pos",
        "qpos0",
        "qpos_spring",
        "range",
        "simplenum",
        "solimp",
        "solref",
        "stiffness",
        "type",
    }

    print("\nChecking hand joints:")
    for joint_i in range(hand_model.njnt):
        hand_joint = hand_model.joint(joint_i)
        joint_name = hand_joint.name
        combined_joint = combined_model.joint(f"hand_base:{joint_name}")

        print(f"Checking joint {joint_name}")

        for key in joint_keys:
            hand_joint_value = getattr(hand_joint, key)
            combined_joint_value = getattr(combined_joint, key)

            if not np.allclose(hand_joint_value, combined_joint_value):
                print(
                    f"{key} does not match: {hand_joint_value} != {combined_joint_value}"
                )

    print("\nChecking arm joints:")
    for joint_i in range(arm_model.njnt):
        arm_joint = arm_model.joint(joint_i)
        joint_name = arm_joint.name
        combined_joint = combined_model.joint(f"{joint_name}")

        print(f"Checking joint {joint_name}")

        for key in joint_keys:
            arm_joint_value = getattr(arm_joint, key)
            combined_joint_value = getattr(combined_joint, key)

            if not np.allclose(arm_joint_value, combined_joint_value):
                print(
                    f"{key} does not match: {arm_joint_value} != {combined_joint_value}"
                )

if __name__ == "__main__":
    main()

Output:

Hand joints:
  joint_name='F0/J0':   solref=[[0.02, 1.0]]    solimp=[[0.9, 0.95, 0.001, 0.5, 2.0]]
  joint_name='F0/J1':   solref=[[1.0, 0.02]]    solimp=[[0.95, 0.001, 0.5, 2.0, 0.9]]
  joint_name='F0/J2':   solref=[[0.02, 1.0]]    solimp=[[0.001, 0.5, 2.0, 0.9, 0.95]]
  joint_name='F0/J3':   solref=[[1.0, 0.02]]    solimp=[[0.5, 2.0, 0.9, 0.95, 0.001]]
  joint_name='F1/J0':   solref=[[0.02, 1.0]]    solimp=[[2.0, 0.9, 0.95, 0.001, 0.5]]
  joint_name='F1/J1':   solref=[[1.0, 0.02]]    solimp=[[0.9, 0.95, 0.001, 0.5, 2.0]]
  joint_name='F1/J2':   solref=[[0.02, 1.0]]    solimp=[[0.95, 0.001, 0.5, 2.0, 0.9]]
  joint_name='F1/J3':   solref=[[1.0, 0.02]]    solimp=[[0.001, 0.5, 2.0, 0.9, 0.95]]
  joint_name='F2/J0':   solref=[[0.02, 1.0]]    solimp=[[0.5, 2.0, 0.9, 0.95, 0.001]]
  joint_name='F2/J1':   solref=[[1.0, 0.02]]    solimp=[[2.0, 0.9, 0.95, 0.001, 0.5]]
  joint_name='F2/J2':   solref=[[0.02, 1.0]]    solimp=[[0.9, 0.95, 0.001, 0.5, 2.0]]
  joint_name='F2/J3':   solref=[[1.0, 0.02]]    solimp=[[0.95, 0.001, 0.5, 2.0, 0.9]]

Arm joints:
  joint_name='fr3_joint1':      solref=[[0.02, 1.0]]    solimp=[[0.9, 0.95, 0.001, 0.5, 2.0]]
  joint_name='fr3_joint2':      solref=[[1.0, 0.02]]    solimp=[[0.95, 0.001, 0.5, 2.0, 0.9]]
  joint_name='fr3_joint3':      solref=[[0.02, 1.0]]    solimp=[[0.001, 0.5, 2.0, 0.9, 0.95]]
  joint_name='fr3_joint4':      solref=[[1.0, 0.02]]    solimp=[[0.5, 2.0, 0.9, 0.95, 0.001]]
  joint_name='fr3_joint5':      solref=[[0.02, 1.0]]    solimp=[[2.0, 0.9, 0.95, 0.001, 0.5]]
  joint_name='fr3_joint6':      solref=[[1.0, 0.02]]    solimp=[[0.9, 0.95, 0.001, 0.5, 2.0]]
  joint_name='fr3_joint7':      solref=[[0.02, 1.0]]    solimp=[[0.95, 0.001, 0.5, 2.0, 0.9]]

Checking hand joints:
Checking joint F0/J0
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[9.0e-01 9.5e-01 1.0e-03 5.0e-01 2.0e+00]] != [[1.0e-03 5.0e-01 2.0e+00 9.0e-01 9.5e-01]]
Checking joint F0/J1
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[9.5e-01 1.0e-03 5.0e-01 2.0e+00 9.0e-01]] != [[5.0e-01 2.0e+00 9.0e-01 9.5e-01 1.0e-03]]
Checking joint F0/J2
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[1.0e-03 5.0e-01 2.0e+00 9.0e-01 9.5e-01]] != [[2.0e+00 9.0e-01 9.5e-01 1.0e-03 5.0e-01]]
Checking joint F0/J3
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[5.0e-01 2.0e+00 9.0e-01 9.5e-01 1.0e-03]] != [[9.0e-01 9.5e-01 1.0e-03 5.0e-01 2.0e+00]]
Checking joint F1/J0
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[2.0e+00 9.0e-01 9.5e-01 1.0e-03 5.0e-01]] != [[9.5e-01 1.0e-03 5.0e-01 2.0e+00 9.0e-01]]
Checking joint F1/J1
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[9.0e-01 9.5e-01 1.0e-03 5.0e-01 2.0e+00]] != [[1.0e-03 5.0e-01 2.0e+00 9.0e-01 9.5e-01]]
Checking joint F1/J2
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[9.5e-01 1.0e-03 5.0e-01 2.0e+00 9.0e-01]] != [[5.0e-01 2.0e+00 9.0e-01 9.5e-01 1.0e-03]]
Checking joint F1/J3
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[1.0e-03 5.0e-01 2.0e+00 9.0e-01 9.5e-01]] != [[2.0e+00 9.0e-01 9.5e-01 1.0e-03 5.0e-01]]
Checking joint F2/J0
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[5.0e-01 2.0e+00 9.0e-01 9.5e-01 1.0e-03]] != [[9.0e-01 9.5e-01 1.0e-03 5.0e-01 2.0e+00]]
Checking joint F2/J1
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[2.0e+00 9.0e-01 9.5e-01 1.0e-03 5.0e-01]] != [[9.5e-01 1.0e-03 5.0e-01 2.0e+00 9.0e-01]]
Checking joint F2/J2
solref does not match: [[0.02 1.  ]] != [[1.   0.02]]
solimp does not match: [[9.0e-01 9.5e-01 1.0e-03 5.0e-01 2.0e+00]] != [[1.0e-03 5.0e-01 2.0e+00 9.0e-01 9.5e-01]]
Checking joint F2/J3
solref does not match: [[1.   0.02]] != [[0.02 1.  ]]
solimp does not match: [[9.5e-01 1.0e-03 5.0e-01 2.0e+00 9.0e-01]] != [[5.0e-01 2.0e+00 9.0e-01 9.5e-01 1.0e-03]]

Checking arm joints:
Checking joint fr3_joint1
Checking joint fr3_joint2
Checking joint fr3_joint3
Checking joint fr3_joint4
Checking joint fr3_joint5
Checking joint fr3_joint6
Checking joint fr3_joint7

I was expecting the sol{imp,ref} values to match between each attached joint and its corresponding joint in the non-attached model, but for some reason they seem to get mixed.

It's also a bit surprising to me that even for the non-attached model, some joints have solref == [1, 0.02] and other solref == [0.02, 1] even though none of these parameters are set in the model xml.

Edit: The results of these are so surprising to me that I feel like this has to be expected and I just misunderstand something. I spent a bit of time crawling through the documentation but couldn't find anything that would've explain this.

Steps for reproduction

See above.

Minimal model for reproduction

See above.

Code required for reproduction

See above.

Confirmations

quagla commented 1 week ago

I can't reproduce it in C so this looks like a binding issue (likely not even related to mjSpec). I'm having a deeper look. My MRE is

  def test_solref(self):
    spec = mujoco.MjSpec.from_string("""
      <mujoco>
        <worldbody>
          <body>
            <geom size=".1"/>
            <joint axis="0 0 1"/>
            <body>
              <geom size=".1"/>
              <joint axis="0 0 1"/>
            </body>
          </body>
        </worldbody>
      </mujoco>
    """)

    model = spec.compile()
    for joint in range(model.njnt):
      solref = model.joint(joint).solref
      solimp = model.joint(joint).solimp
      np.testing.assert_array_equal(solref, [[0.02, 1]])
      np.testing.assert_array_equal(solimp, [[0.9, 0.95, 0.001, 0.5, 2.0]])
quagla commented 1 week ago

In the meantime, could you please see if extracting the values with

solref = model.jnt_solref[i]
solimp = model.jnt_solimp[i]

fixes your issue?

hartikainen commented 1 week ago

Indeed, extracting the values directly with integers works as expected. Here's a bit simpler test case that passes:

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)
    hand_model = hand_spec.compile()
    arm_spec = mujoco.MjSpec.from_file(fr3_mj_description.MJCF_PATH)
    arm_model = arm_spec.compile()

    combined_spec = arm_spec.copy()

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

    combined_model = combined_spec.compile()

    np.testing.assert_equal(
        hand_model.jnt_solref,
        combined_model.jnt_solref[None:hand_model.njnt, ...],
    )
    np.testing.assert_equal(
        arm_model.jnt_solref,
        combined_model.jnt_solref[hand_model.njnt:None, ...],
    )
    np.testing.assert_equal(
        hand_model.jnt_solimp,
        combined_model.jnt_solimp[None:hand_model.njnt, ...],
    )
    np.testing.assert_equal(
        arm_model.jnt_solimp,
        combined_model.jnt_solimp[hand_model.njnt:None, ...],
    )

if __name__ == "__main__":
    main()