chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add qpos to MJX bindings.

Browse Source 
PiperOrigin-RevId: 720311208
Change-Id: Iaf5f905e206256f38cd71949bbebd5151b2a2041
This commit is contained in:
Alessio Quaglino
2025-01-27 14:35:16 -08:00
committed by Copybara-Service
parent 8c22181156
commit 3dcfad3293
2 changed files with 36 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+23 -4
View File
@@ -25,6 +25,7 @@ from mujoco.mjx._src import scan
from mujoco.mjx._src.types import ConeType
from mujoco.mjx._src.types import Data
from mujoco.mjx._src.types import JacobianType
from mujoco.mjx._src.types import JointType
from mujoco.mjx._src.types import Model
# pylint: enable=g-importing-member
import numpy as np
@@ -433,20 +434,38 @@ class BindData(object):
return name
else:
raise AttributeError('ctrl is not available for this type')
if name == 'qpos':
if self.prefix == 'jnt_':
return name
else:
raise AttributeError('qpos is not available for this type')
else:
return self.prefix + name
def __getattr__(self, name: str):
if name == 'sensordata':
adr = self.model.sensor_adr[self.id]
num = self.model.sensor_dim[self.id]
if name == 'sensordata' or name == 'qpos':
adr = num = 0
if name == 'sensordata':
adr = self.model.sensor_adr[self.id]
num = self.model.sensor_dim[self.id]
elif name == 'qpos':
adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id]
num = (
(typ == JointType.FREE) * JointType.FREE.qpos_width()
+ (typ == JointType.BALL) * JointType.BALL.qpos_width()
+ (typ == JointType.HINGE) * JointType.HINGE.qpos_width()
+ (typ == JointType.SLIDE) * JointType.SLIDE.qpos_width()
)
if isinstance(self.id, list):
idx = []
for a, n in zip(adr, num):
idx.extend(a + j for j in range(n))
return getattr(self.data, name)[idx, ...]
else:
elif num > 1:
return getattr(self.data, name)[adr : adr + num, ...]
else:
return getattr(self.data, name)[adr, ...]
return getattr(self.data, self.__getname(name))[self.id, ...]
def set(self, name: str, value: jax.Array) -> Data:
mjx/mujoco/mjx/_src/support_test.py
+13 -3
View File
@@ -161,15 +161,15 @@ class SupportTest(parameterized.TestCase):
xml = """
<mujoco model="test_bind_model">
<worldbody>
<body pos="1 2 3" name="body1">
<body pos="10 20 30" name="body1">
<joint axis="1 0 0" type="slide" name="joint1"/>
<geom size="1 2 3" type="box" name="geom1"/>
</body>
<body pos="4 5 6" name="body2">
<body pos="40 50 60" name="body2">
<joint axis="0 1 0" type="slide" name="joint2"/>
<geom size="4 5 6" type="box" name="geom2"/>
</body>
<body pos="7 8 9" name="body3">
<body pos="70 80 90" name="body3">
<joint axis="0 0 1" type="slide" name="joint3"/>
<geom size="7 8 9" type="box" name="geom3"/>
</body>
@@ -223,6 +223,10 @@ class SupportTest(parameterized.TestCase):
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, :])
np.testing.assert_array_almost_equal(
dx.bind(mx, s.joints[i]).qpos,
d.qpos[m.jnt_qposadr[i]], decimal=6
)
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, d.ctrl)
for i in range(m.nu):
@@ -256,6 +260,12 @@ class SupportTest(parameterized.TestCase):
np.testing.assert_array_equal(dx5.bind(mx, s.actuators).ctrl, [0, 7, 0])
np.testing.assert_array_equal(dx.bind(mx, s.actuators).ctrl, [0, 0, 0])
np.testing.assert_array_almost_equal(d.qpos, [0, 0, -3.924e-05])
np.testing.assert_array_almost_equal(dx.bind(mx, s.joints).qpos, d.qpos)
dx6 = dx.bind(mx, s.joints[1:]).set('qpos', [8, 0])
np.testing.assert_array_equal(dx6.bind(mx, s.joints).qpos, [0, 8, 0])
np.testing.assert_array_almost_equal(dx.bind(mx, s.joints).qpos, d.qpos)
# test invalid name
with self.assertRaises(AttributeError):
print(dx.bind(mx, s.geoms).ctrl)