Connect constraint anchor computation does not account for joint reference positions. Affects MuJoCo CPU and MuJoCo Warp. #1270
Description
Summary:
The connect constraint couples two bodies and has two anchors, one for each body. The first anchor is set by the user. The second anchor is computed by mujoco with the rule that the constraint is satisfied when the bodies are posed at the reference joint positions (https://mujoco.readthedocs.io/en/stable/XMLreference.html#equality)
The computation of the second anchor ignores the reference joint positions and instead computes the second anchor with all reference joint positions having value 0. This affects Mujoco C and Mujoco Warp (see test below).
Diagnosis:
The root of the problem is in set_const_0. We first set qpos=qpos0
wp.launch(_copy_qpos0_to_qpos, dim=(d.nworld, m.nq), inputs=[m.qpos0], outputs=[d.qpos])
We then compute the body positions as a deviation from qpos0 using qpos-qpos0. For example, for prismatic joints we do this:
xpos += xaxis * (qpos[qadr] - qpos0[worldid % qpos0.shape[0], qadr])
If qpos equals qpos0 then we always add 0 to xpos. We have merely computed the body positions at qpos0={0}. qpos0 is therefore ignored.
The issue described here also affects MuJoCo CPU. This is demonstrated in the repro below.
Repro:
I modified test_set_const_eq_data_connect to demonstrate the problem.
def test_set_const_eq_data_connect(self):
"""Test set_const_0 recomputes eq_data for CONNECT constraints."""
mjm, mjd, m, d = test_data.fixture(
xml="""
<mujoco>
<option gravity="0 0 0">
<flag contact="disable"/>
</option>
<worldbody>
<body name="b1" pos="1 0 0">
<joint type="slide" axis="1 0 0" ref="1"/>
<geom type="sphere" size="0.1" mass="1"/>
</body>
<body name="b2" pos="2 0 0">
<joint type="slide" axis="1 0 0" ref="4"/>
<geom type="sphere" size="0.1" mass="1"/>
</body>
</worldbody>
<equality>
<connect body1="b1" body2="b2" anchor="0.5 0 0"/>
</equality>
</mujoco>
"""
)
# Verify qpos0 reflects the reference poses from the XML.
np.testing.assert_allclose(m.qpos0.numpy()[0][0], 1.0)
np.testing.assert_allclose(m.qpos0.numpy()[0][1], 4.0)
np.testing.assert_allclose(m.qpos0.numpy()[0], mjm.qpos0)
# Overwrite eq_data[3:6] with wrong values.
# This isn't strictly necessary but it helps confirm that
# setConst and set_const_0 actually compute a value for the
# 2nd anchor in the connect constraint.
wrong_values = [99.0, 99.0, 99.0]
mjm.eq_data[0, 3:6] = wrong_values
eq_data_np = m.eq_data.numpy().copy()
eq_data_np[0, 0, 3:6] = wrong_values
m.eq_data.assign(eq_data_np)
mujoco.mj_setConst(mjm, mjd)
mjwarp.set_const_0(m, d)
# The connect constraint should be satisfied at the joint reference positions.
# Given the prismatic joints along X and the joint reference positions (1,4) we
# can compute the anchor that must be associated with body2:
# qpos0[0] + anchor1.x = qpos[1] + anchor2.x
# 1.0 + 0.5 = 4.0 + anchor2.x
# so
# anchor2.x = -2.5
print("The 2nd anchor of the connect constraint should have value (-2.5, 0, 0)")
print("mjm.eq_data[3:6} has value]")
print(mjm.eq_data[0, 3:6])
print("m.eq_data[3:6} has value]")
print(m.eq_data.numpy()[0, 0, 3:6])