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

Fix qvel bind in MJX when using multiple joints.

Browse Source 
Reported by #2402.

PiperOrigin-RevId: 725237347
Change-Id: I279c73e5cf2b959086292b739972afe5f70ceffb
This commit is contained in:
Alessio Quaglino
2025-02-10 09:08:24 -08:00
committed by Copybara-Service
parent 4156350708
commit 94115ecf08
1 changed files with 13 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/support.py
+13 -15
View File
@@ -438,16 +438,16 @@ class BindData(object):
return name return name
else: else:
raise AttributeError('ctrl is not available for this type') raise AttributeError('ctrl is not available for this type')
if name == 'qpos': if name == 'qpos' or name == 'qvel':
if self.prefix == 'jnt_': if self.prefix == 'jnt_':
return name return name
else: else:
raise AttributeError('qpos is not available for this type') raise AttributeError('qpos and qvel are not available for this type')
else: else:
return self.prefix + name return self.prefix + name
def __getattr__(self, name: str): def __getattr__(self, name: str):
if name == 'sensordata' or name == 'qpos': if name in ('sensordata', 'qpos', 'qvel'):
adr = num = 0 adr = num = 0
if name == 'sensordata': if name == 'sensordata':
adr = self.model.sensor_adr[self.id] adr = self.model.sensor_adr[self.id]
@@ -455,12 +455,11 @@ class BindData(object):
elif name == 'qpos': elif name == 'qpos':
adr = self.model.jnt_qposadr[self.id] adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id] typ = self.model.jnt_type[self.id]
num = ( num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
(typ == JointType.FREE) * JointType.FREE.qpos_width() elif name == 'qvel':
+ (typ == JointType.BALL) * JointType.BALL.qpos_width() adr = self.model.jnt_dofadr[self.id]
+ (typ == JointType.HINGE) * JointType.HINGE.qpos_width() typ = self.model.jnt_type[self.id]
+ (typ == JointType.SLIDE) * JointType.SLIDE.qpos_width() num = sum((typ == jt) * jt.dof_width() for jt in JointType)
)
if isinstance(self.id, list): if isinstance(self.id, list):
idx = [] idx = []
for a, n in zip(adr, num): for a, n in zip(adr, num):
@@ -484,12 +483,11 @@ class BindData(object):
if name == 'qpos': if name == 'qpos':
adr = self.model.jnt_qposadr[self.id] adr = self.model.jnt_qposadr[self.id]
typ = self.model.jnt_type[self.id] typ = self.model.jnt_type[self.id]
num = ( num = sum((typ == jt) * jt.qpos_width() for jt in JointType)
(typ == JointType.FREE) * JointType.FREE.qpos_width() elif name == 'qvel':
+ (typ == JointType.BALL) * JointType.BALL.qpos_width() adr = self.model.jnt_dofadr[self.id]
+ (typ == JointType.HINGE) * JointType.HINGE.qpos_width() typ = self.model.jnt_type[self.id]
+ (typ == JointType.SLIDE) * JointType.SLIDE.qpos_width() num = sum((typ == jt) * jt.dof_width() for jt in JointType)
)
elif isinstance(self.id, list): elif isinstance(self.id, list):
adr = self.id adr = self.id
num = [1 for _ in range(len(self.id))] num = [1 for _ in range(len(self.id))]