diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index aeb0f7ed..0e4e21ec 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -367,7 +367,7 @@ class BindModel(object): self.id = ids def __getattr__(self, name: str): - return getattr(self.model, self.prefix + name)[self.id, :] + return getattr(self.model, self.prefix + name)[self.id, ...] def _bind_model(self: Model, obj: Sequence[Any]) -> BindModel: diff --git a/mjx/mujoco/mjx/_src/support_test.py b/mjx/mujoco/mjx/_src/support_test.py index a61374da..da046c70 100644 --- a/mjx/mujoco/mjx/_src/support_test.py +++ b/mjx/mujoco/mjx/_src/support_test.py @@ -219,6 +219,7 @@ class SupportTest(parameterized.TestCase): ) np.testing.assert_array_equal(mx.bind(s.joints).axis, m.jnt_axis) + np.testing.assert_array_equal(mx.bind(s.joints).qposadr, m.jnt_qposadr) for i in range(m.njnt): np.testing.assert_array_equal(m.bind(s.joints[i]).axis, m.jnt_axis[i, :]) np.testing.assert_array_equal(mx.bind(s.joints[i]).axis, m.jnt_axis[i, :])