Add support for binding to arrays of mjs element.

Fixes #2402.

PiperOrigin-RevId: 726449592
Change-Id: I0cde44889837a71a222b20cf91fa56568b4af689
This commit is contained in:
Alessio Quaglino
2025-02-13 06:00:31 -08:00
committed by Copybara-Service
parent 1b4258d274
commit 5b924fec98
3 changed files with 101 additions and 3 deletions
+35
View File
@@ -1100,5 +1100,40 @@ class SpecsTest(absltest.TestCase):
with self.assertRaisesRegex(ValueError, 'Frame not found.'):
parent.attach(child4, frame='invalid_frame', prefix='child3-')
def test_bind(self):
spec = mujoco.MjSpec.from_string("""
<mujoco>
<worldbody>
<body name="main">
<geom name="main" size="0.15 0.15 0.15" mass="1" type="box"/>
<freejoint/>
<body name="box">
<joint name="box" type="hinge" range="-1 +1"/>
<geom name="box" size="0.15 0.15 0.15" mass="1" type="box"/>
</body>
<body name="sphere">
<joint name="sphere" type="hinge" range="-1 +1"/>
<geom name="sphere" size="0.15 0.15 0.15" mass="1" type="box"/>
</body>
</body>
</worldbody>
</mujoco>
""")
joint_box = spec.joint('box')
joint_sphere = spec.joint('sphere')
joints = [joint_box, joint_sphere]
mj_model = spec.compile()
mj_data = mujoco.MjData(mj_model)
np.testing.assert_array_equal(mj_data.bind(joint_box).qpos, 0)
np.testing.assert_array_equal(mj_model.bind(joint_box).qposadr, 7)
np.testing.assert_array_equal(mj_data.bind(joints).qpos, [0, 0])
np.testing.assert_array_equal(mj_model.bind(joints).qposadr, [7, 8])
np.testing.assert_array_equal(mj_data.bind([]).qpos, [])
np.testing.assert_array_equal(mj_model.bind([]).qposadr, [])
with self.assertRaisesRegex(
AttributeError, "object has no attribute 'invalid'"
):
print(mj_model.bind(joints).invalid)
if __name__ == '__main__':
absltest.main()