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

Make actuator_length, cdof and cdof_dot public.

Browse Source 
PiperOrigin-RevId: 842182467
Change-Id: Id14a1c8266e13a6f95e1c87ea07da993b614b62e
This commit is contained in:
Google DeepMind
2025-12-09 04:56:57 -08:00
committed by Copybara-Service
parent f65c1cdfd1
commit 1483aefe5f
14 changed files with 69 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+13
View File
@@ -2,6 +2,19 @@
Changelog Changelog
========= =========
Upcoming version (not yet released)
-----------------------------------
General
^^^^^^^
MJX
^^^
1. Added ``actuator_length``, ``cdof`` and ``cdof_dof`` fields to ``mjx.Data``.
Bug fixes
^^^^^^^^^
Version 3.4.0 (December 5, 2025) Version 3.4.0 (December 5, 2025)
-------------------------------- --------------------------------
mjx/mujoco/mjx/_src/forward.py
+1 -1
View File
@@ -182,7 +182,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
m.actuator_gainprm, m.actuator_gainprm,
m.actuator_biastype, m.actuator_biastype,
m.actuator_biasprm, m.actuator_biasprm,
d._impl.actuator_length, d.actuator_length,
d._impl.actuator_velocity, d._impl.actuator_velocity,
ctrl_act, ctrl_act,
jp.array(m.actuator_lengthrange), jp.array(m.actuator_lengthrange),
mjx/mujoco/mjx/_src/io.py
+3 -6
View File
@@ -552,6 +552,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
'cam_xmat': (m.ncam, 3, 3, float_), 'cam_xmat': (m.ncam, 3, 3, float_),
'subtree_com': (m.nbody, 3, float_), 'subtree_com': (m.nbody, 3, float_),
'actuator_force': (m.nu, float_), 'actuator_force': (m.nu, float_),
'actuator_length': (m.nu, float_),
'qfrc_bias': (m.nv, float_), 'qfrc_bias': (m.nv, float_),
'qfrc_gravcomp': (m.nv, float_), 'qfrc_gravcomp': (m.nv, float_),
'qfrc_fluid': (m.nv, float_), 'qfrc_fluid': (m.nv, float_),
@@ -562,6 +563,8 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
'qfrc_constraint': (m.nv, float_), 'qfrc_constraint': (m.nv, float_),
'qfrc_inverse': (m.nv, float_), 'qfrc_inverse': (m.nv, float_),
'cvel': (m.nbody, 6, float_), 'cvel': (m.nbody, 6, float_),
'cdof': (m.nv, 6, float_),
'cdof_dot': (m.nv, 6, float_),
'ten_length': (m.ntendon, float_), 'ten_length': (m.ntendon, float_),
} }
zero_fields = { zero_fields = {
@@ -619,14 +622,12 @@ def _make_data_jax(
zero_impl_fields = { zero_impl_fields = {
'solver_niter': (int_,), 'solver_niter': (int_,),
'cdof': (m.nv, 6, float_),
'cinert': (m.nbody, 10, float_), 'cinert': (m.nbody, 10, float_),
'ten_wrapadr': (m.ntendon, np.int32), 'ten_wrapadr': (m.ntendon, np.int32),
'ten_wrapnum': (m.ntendon, np.int32), 'ten_wrapnum': (m.ntendon, np.int32),
'ten_J': (m.ntendon, m.nv, float_), 'ten_J': (m.ntendon, m.nv, float_),
'wrap_obj': (m.nwrap, 2, np.int32), 'wrap_obj': (m.nwrap, 2, np.int32),
'wrap_xpos': (m.nwrap, 6, float_), 'wrap_xpos': (m.nwrap, 6, float_),
'actuator_length': (m.nu, float_),
'actuator_moment': (m.nu, m.nv, float_), 'actuator_moment': (m.nu, m.nv, float_),
'crb': (m.nbody, 10, float_), 'crb': (m.nbody, 10, float_),
'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_), 'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
@@ -635,7 +636,6 @@ def _make_data_jax(
'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_), 'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_),
'ten_velocity': (m.ntendon, float_), 'ten_velocity': (m.ntendon, float_),
'actuator_velocity': (m.nu, float_), 'actuator_velocity': (m.nu, float_),
'cdof_dot': (m.nv, 6, float_),
'cacc': (m.nbody, 6, float_), 'cacc': (m.nbody, 6, float_),
'cfrc_int': (m.nbody, 6, float_), 'cfrc_int': (m.nbody, 6, float_),
'cfrc_ext': (m.nbody, 6, float_), 'cfrc_ext': (m.nbody, 6, float_),
@@ -713,7 +713,6 @@ def _make_data_c(
nbvhdynamic = get(m, 'nbvhdynamic') nbvhdynamic = get(m, 'nbvhdynamic')
zero_impl_fields = { zero_impl_fields = {
'solver_niter': (int_,), 'solver_niter': (int_,),
'cdof': (m.nv, 6, float_),
'cinert': (m.nbody, 10, float_), 'cinert': (m.nbody, 10, float_),
'light_xpos': (m.nlight, 3, float_), 'light_xpos': (m.nlight, 3, float_),
'light_xdir': (m.nlight, 3, float_), 'light_xdir': (m.nlight, 3, float_),
@@ -732,7 +731,6 @@ def _make_data_c(
'ten_wrapnum': (m.ntendon, np.int32), 'ten_wrapnum': (m.ntendon, np.int32),
'wrap_obj': (m.nwrap, 2, np.int32), 'wrap_obj': (m.nwrap, 2, np.int32),
'wrap_xpos': (m.nwrap, 6, float_), 'wrap_xpos': (m.nwrap, 6, float_),
'actuator_length': (m.nu, float_),
'moment_rownnz': (m.nu, np.int32), 'moment_rownnz': (m.nu, np.int32),
'moment_rowadr': (m.nu, np.int32), 'moment_rowadr': (m.nu, np.int32),
'moment_colind': (m.nJmom, np.int32), 'moment_colind': (m.nJmom, np.int32),
@@ -762,7 +760,6 @@ def _make_data_c(
'qLU': (m.nD, float_), 'qLU': (m.nD, float_),
'qfrc_spring': (m.nv, float_), 'qfrc_spring': (m.nv, float_),
'qfrc_damper': (m.nv, float_), 'qfrc_damper': (m.nv, float_),
'cdof_dot': (m.nv, 6, float_),
'cacc': (m.nbody, 6, float_), 'cacc': (m.nbody, 6, float_),
'cfrc_int': (m.nbody, 6, float_), 'cfrc_int': (m.nbody, 6, float_),
'cfrc_ext': (m.nbody, 6, float_), 'cfrc_ext': (m.nbody, 6, float_),
mjx/mujoco/mjx/_src/io_test.py
+4 -4
View File
@@ -402,10 +402,10 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d.geom_xpos.shape, (6, 3)) self.assertEqual(d.geom_xpos.shape, (6, 3))
self.assertEqual(d.geom_xmat.shape, (6, 3, 3)) self.assertEqual(d.geom_xmat.shape, (6, 3, 3))
self.assertEqual(d.subtree_com.shape, (nbody, 3)) self.assertEqual(d.subtree_com.shape, (nbody, 3))
self.assertEqual(d._impl.cdof.shape, (nv, 6)) self.assertEqual(d.cdof.shape, (nv, 6))
self.assertEqual(d._impl.cinert.shape, (nbody, 10)) self.assertEqual(d._impl.cinert.shape, (nbody, 10))
self.assertEqual(d._impl.crb.shape, (nbody, 10)) self.assertEqual(d._impl.crb.shape, (nbody, 10))
self.assertEqual(d._impl.actuator_length.shape, (1,)) self.assertEqual(d.actuator_length.shape, (1,))
if impl == 'jax': if impl == 'jax':
self.assertEqual(d._impl.actuator_moment.shape, (1, nv)) self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
elif impl == 'c': elif impl == 'c':
@@ -422,7 +422,7 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d._impl.efc_D.shape, (nefc,)) self.assertEqual(d._impl.efc_D.shape, (nefc,))
self.assertEqual(d._impl.actuator_velocity.shape, (1,)) self.assertEqual(d._impl.actuator_velocity.shape, (1,))
self.assertEqual(d.cvel.shape, (nbody, 6)) self.assertEqual(d.cvel.shape, (nbody, 6))
self.assertEqual(d._impl.cdof_dot.shape, (nv, 6)) self.assertEqual(d.cdof_dot.shape, (nv, 6))
self.assertEqual(d.qfrc_bias.shape, (nv,)) self.assertEqual(d.qfrc_bias.shape, (nv,))
self.assertEqual(d.qfrc_passive.shape, (nv,)) self.assertEqual(d.qfrc_passive.shape, (nv,))
self.assertEqual(d._impl.efc_aref.shape, (nefc,)) self.assertEqual(d._impl.efc_aref.shape, (nefc,))
@@ -477,7 +477,7 @@ class DataIOTest(parameterized.TestCase):
np.testing.assert_allclose(dx.qpos, d.qpos) np.testing.assert_allclose(dx.qpos, d.qpos)
np.testing.assert_allclose(dx.xpos, d.xpos) np.testing.assert_allclose(dx.xpos, d.xpos)
np.testing.assert_allclose(dx.cvel, d.cvel) np.testing.assert_allclose(dx.cvel, d.cvel)
np.testing.assert_allclose(dx._impl.cdof_dot, d.cdof_dot) np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
# check that there are no weak types # check that there are no weak types
self.assertFalse( self.assertFalse(
mjx/mujoco/mjx/_src/sensor.py
+1 -1
View File
@@ -175,7 +175,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
elif sensor_type == SensorType.TENDONPOS: elif sensor_type == SensorType.TENDONPOS:
sensor = d.ten_length[objid] sensor = d.ten_length[objid]
elif sensor_type == SensorType.ACTUATORPOS: elif sensor_type == SensorType.ACTUATORPOS:
sensor = d._impl.actuator_length[objid] sensor = d.actuator_length[objid]
elif sensor_type == SensorType.BALLQUAT: elif sensor_type == SensorType.BALLQUAT:
jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None] jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
quat = d.qpos[jnt_qposadr] quat = d.qpos[jnt_qposadr]
mjx/mujoco/mjx/_src/smooth.py
+10 -10
View File
@@ -212,7 +212,7 @@ def com_pos(m: Model, d: Data) -> Data:
d.xanchor, d.xanchor,
d.xaxis, d.xaxis,
) )
d = d.tree_replace({'_impl.cdof': cdof}) d = d.tree_replace({'cdof': cdof})
return d return d
@@ -305,8 +305,8 @@ def crb(m: Model, d: Data) -> Data:
d = d.tree_replace({'_impl.crb': crb_body}) d = d.tree_replace({'_impl.crb': crb_body})
crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0) crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d._impl.cdof) crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature) qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.tree_replace({'_impl.qM': qm}) d = d.tree_replace({'_impl.qM': qm})
return d return d
@@ -445,11 +445,11 @@ def com_vel(m: Model, d: Data) -> Data:
'jvv', 'jvv',
'bv', 'bv',
m.jnt_type, m.jnt_type,
d._impl.cdof, d.cdof,
d.qvel, d.qvel,
) )
d = d.tree_replace({'cvel': cvel, '_impl.cdof_dot': cdof_dot}) d = d.tree_replace({'cvel': cvel, 'cdof_dot': cdof_dot})
return d return d
@@ -576,7 +576,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
return cacc return cacc
cacc = scan.body_tree( cacc = scan.body_tree(
m, cacc_fn, 'vvvv', 'b', d._impl.cdof_dot, d.qvel, d._impl.cdof, d.qacc m, cacc_fn, 'vvvv', 'b', d.cdof_dot, d.qvel, d.cdof, d.qacc
) )
def frc(cinert, cacc, cvel): def frc(cinert, cacc, cvel):
@@ -594,7 +594,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
return cfrc return cfrc
cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True) cfrc = scan.body_tree(m, cfrc_fn, 'b', 'b', loc_cfrc, reverse=True)
qfrc_bias = jax.vmap(jp.dot)(d._impl.cdof, cfrc[jp.array(m.dof_bodyid)]) qfrc_bias = jax.vmap(jp.dot)(d.cdof, cfrc[jp.array(m.dof_bodyid)])
d = d.replace(qfrc_bias=qfrc_bias) d = d.replace(qfrc_bias=qfrc_bias)
@@ -806,8 +806,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
) )
# cacc = cacc_parent + cdofdot * qvel + cdof * qacc # cacc = cacc_parent + cdofdot * qvel + cdof * qacc
cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel) cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
cacc_acc = d._impl.cdof.T @ (mask * d.qacc) cacc_acc = d.cdof.T @ (mask * d.qacc)
cacc = cacc_parent + cacc_vel + cacc_acc cacc = cacc_parent + cacc_vel + cacc_acc
# cfrc_body = cinert * cacc + cvel x (cinert * cvel) # cfrc_body = cinert * cacc + cvel x (cinert * cvel)
@@ -1310,7 +1310,7 @@ def transmission(m: Model, d: Data) -> Data:
moment = moment.reshape((m.nu, m.nv)) moment = moment.reshape((m.nu, m.nv))
d = d.tree_replace( d = d.tree_replace(
{'_impl.actuator_length': length, '_impl.actuator_moment': moment} {'actuator_length': length, '_impl.actuator_moment': moment}
) )
return d return d
mjx/mujoco/mjx/_src/smooth_test.py
+4 -4
View File
@@ -81,7 +81,7 @@ class SmoothTest(absltest.TestCase):
dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d)) dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'subtree_com') _assert_attr_eq(d, dx, 'subtree_com')
_assert_attr_eq(d, dx._impl, 'cinert') _assert_attr_eq(d, dx._impl, 'cinert')
_assert_attr_eq(d, dx._impl, 'cdof') _assert_attr_eq(d, dx, 'cdof')
# camlight # camlight
dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d)) dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'cam_xpos') _assert_attr_eq(d, dx, 'cam_xpos')
@@ -100,7 +100,7 @@ class SmoothTest(absltest.TestCase):
# com_vel # com_vel
dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d)) dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'cvel') _assert_attr_eq(d, dx, 'cvel')
_assert_attr_eq(d, dx._impl, 'cdof_dot') _assert_attr_eq(d, dx, 'cdof_dot')
# rne # rne
dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d)) dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx, 'qfrc_bias') _assert_attr_eq(d, dx, 'qfrc_bias')
@@ -125,7 +125,7 @@ class SmoothTest(absltest.TestCase):
_assert_attr_eq(d, dx, 'ten_length') _assert_attr_eq(d, dx, 'ten_length')
# transmission # transmission
dx = jax.jit(mjx.transmission)(mx, dx) dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx._impl, 'actuator_length') _assert_attr_eq(d, dx, 'actuator_length')
# convert sparse actuator_moment to dense representation # convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv)) moment = np.zeros((m.nu, m.nv))
@@ -196,7 +196,7 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_transmission(m, d) mujoco.mj_transmission(m, d)
dx = jax.jit(mjx.transmission)(mx, dx) dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx._impl, 'actuator_length') _assert_attr_eq(d, dx, 'actuator_length')
# convert sparse actuator_moment to dense representation # convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv)) moment = np.zeros((m.nu, m.nv))
mjx/mujoco/mjx/_src/support.py
+4 -4
View File
@@ -149,9 +149,9 @@ def jac(
mask = mask[jp.array(m.dof_bodyid)] > 0 mask = mask[jp.array(m.dof_bodyid)] > 0
offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]] offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d._impl.cdof) # pytype: disable=attribute-error jacp = jax.vmap(lambda a, b=offset: a[3:] + jp.cross(a[:3], b))(d.cdof) # pytype: disable=attribute-error
jacp = jax.vmap(jp.multiply)(jacp, mask) jacp = jax.vmap(jp.multiply)(jacp, mask)
jacr = jax.vmap(jp.multiply)(d._impl.cdof[:, :3], mask) # pytype: disable=attribute-error jacr = jax.vmap(jp.multiply)(d.cdof[:, :3], mask) # pytype: disable=attribute-error
return jacp, jacr return jacp, jacr
@@ -169,8 +169,8 @@ def jac_dot(
offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]] offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
pvel_lin = d.cvel[body_id][3:] - jp.cross(offset, d.cvel[body_id][:3]) pvel_lin = d.cvel[body_id][3:] - jp.cross(offset, d.cvel[body_id][:3])
cdof = d._impl.cdof cdof = d.cdof
cdof_dot = d._impl.cdof_dot cdof_dot = d.cdof_dot
# check for quaternion # check for quaternion
jnt_type = m.jnt_type[m.dof_jntid] jnt_type = m.jnt_type[m.dof_jntid]
mjx/mujoco/mjx/_src/types.py
+5 -6
View File
@@ -1021,7 +1021,6 @@ class DataC(PyTreeNode):
plugin_data: jax.Array plugin_data: jax.Array
light_xpos: jax.Array light_xpos: jax.Array
light_xdir: jax.Array light_xdir: jax.Array
cdof: jax.Array
cinert: jax.Array cinert: jax.Array
flexvert_xpos: jax.Array flexvert_xpos: jax.Array
flexelem_aabb: jax.Array flexelem_aabb: jax.Array
@@ -1039,7 +1038,6 @@ class DataC(PyTreeNode):
ten_J: jax.Array # pylint:disable=invalid-name ten_J: jax.Array # pylint:disable=invalid-name
wrap_obj: jax.Array wrap_obj: jax.Array
wrap_xpos: jax.Array wrap_xpos: jax.Array
actuator_length: jax.Array
moment_rownnz: jax.Array # pylint:disable=invalid-name moment_rownnz: jax.Array # pylint:disable=invalid-name
moment_rowadr: jax.Array # pylint:disable=invalid-name moment_rowadr: jax.Array # pylint:disable=invalid-name
moment_colind: jax.Array # pylint:disable=invalid-name moment_colind: jax.Array # pylint:disable=invalid-name
@@ -1059,7 +1057,7 @@ class DataC(PyTreeNode):
flexedge_velocity: jax.Array flexedge_velocity: jax.Array
ten_velocity: jax.Array ten_velocity: jax.Array
actuator_velocity: jax.Array actuator_velocity: jax.Array
cdof_dot: jax.Array
qfrc_spring: jax.Array qfrc_spring: jax.Array
qfrc_damper: jax.Array qfrc_damper: jax.Array
subtree_linvel: jax.Array subtree_linvel: jax.Array
@@ -1095,14 +1093,12 @@ class DataJAX(PyTreeNode):
nefc: int nefc: int
ncon: int ncon: int
solver_niter: jax.Array solver_niter: jax.Array
cdof: jax.Array
cinert: jax.Array cinert: jax.Array
ten_wrapadr: jax.Array ten_wrapadr: jax.Array
ten_wrapnum: jax.Array ten_wrapnum: jax.Array
ten_J: jax.Array # pylint:disable=invalid-name ten_J: jax.Array # pylint:disable=invalid-name
wrap_obj: jax.Array wrap_obj: jax.Array
wrap_xpos: jax.Array wrap_xpos: jax.Array
actuator_length: jax.Array
actuator_moment: jax.Array actuator_moment: jax.Array
crb: jax.Array crb: jax.Array
qM: jax.Array # pylint:disable=invalid-name qM: jax.Array # pylint:disable=invalid-name
@@ -1111,7 +1107,7 @@ class DataJAX(PyTreeNode):
qLDiagInv: jax.Array # pylint:disable=invalid-name qLDiagInv: jax.Array # pylint:disable=invalid-name
ten_velocity: jax.Array ten_velocity: jax.Array
actuator_velocity: jax.Array actuator_velocity: jax.Array
cdof_dot: jax.Array
cacc: jax.Array cacc: jax.Array
cfrc_int: jax.Array cfrc_int: jax.Array
cfrc_ext: jax.Array cfrc_ext: jax.Array
@@ -1171,12 +1167,15 @@ class Data(PyTreeNode):
cam_xmat: jax.Array cam_xmat: jax.Array
subtree_com: jax.Array subtree_com: jax.Array
cvel: jax.Array cvel: jax.Array
cdof: jax.Array
cdof_dot: jax.Array
qfrc_bias: jax.Array qfrc_bias: jax.Array
qfrc_gravcomp: jax.Array qfrc_gravcomp: jax.Array
qfrc_fluid: jax.Array qfrc_fluid: jax.Array
qfrc_passive: jax.Array qfrc_passive: jax.Array
qfrc_actuator: jax.Array qfrc_actuator: jax.Array
actuator_force: jax.Array actuator_force: jax.Array
actuator_length: jax.Array
qfrc_smooth: jax.Array qfrc_smooth: jax.Array
qacc_smooth: jax.Array qacc_smooth: jax.Array
qfrc_constraint: jax.Array qfrc_constraint: jax.Array
mjx/mujoco/mjx/warp/collision_driver.py
-1
View File
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _collision_shim( def _collision_shim(
# Model # Model
mjx/mujoco/mjx/warp/forward.py
+18 -20
View File
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _forward_shim( def _forward_shim(
# Model # Model
@@ -952,14 +951,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'act': d.act.shape, 'act': d.act.shape,
'act_dot': d.act_dot.shape, 'act_dot': d.act_dot.shape,
'actuator_force': d.actuator_force.shape, 'actuator_force': d.actuator_force.shape,
'actuator_length': d._impl.actuator_length.shape, 'actuator_length': d.actuator_length.shape,
'actuator_moment': d._impl.actuator_moment.shape, 'actuator_moment': d._impl.actuator_moment.shape,
'actuator_velocity': d._impl.actuator_velocity.shape, 'actuator_velocity': d._impl.actuator_velocity.shape,
'cacc': d._impl.cacc.shape, 'cacc': d._impl.cacc.shape,
'cam_xmat': d.cam_xmat.shape, 'cam_xmat': d.cam_xmat.shape,
'cam_xpos': d.cam_xpos.shape, 'cam_xpos': d.cam_xpos.shape,
'cdof': d._impl.cdof.shape, 'cdof': d.cdof.shape,
'cdof_dot': d._impl.cdof_dot.shape, 'cdof_dot': d.cdof_dot.shape,
'cfrc_ext': d._impl.cfrc_ext.shape, 'cfrc_ext': d._impl.cfrc_ext.shape,
'cfrc_int': d._impl.cfrc_int.shape, 'cfrc_int': d._impl.cfrc_int.shape,
'cinert': d._impl.cinert.shape, 'cinert': d._impl.cinert.shape,
@@ -1539,14 +1538,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
d.act, d.act,
d.act_dot, d.act_dot,
d.actuator_force, d.actuator_force,
d._impl.actuator_length, d.actuator_length,
d._impl.actuator_moment, d._impl.actuator_moment,
d._impl.actuator_velocity, d._impl.actuator_velocity,
d._impl.cacc, d._impl.cacc,
d.cam_xmat, d.cam_xmat,
d.cam_xpos, d.cam_xpos,
d._impl.cdof, d.cdof,
d._impl.cdof_dot, d.cdof_dot,
d._impl.cfrc_ext, d._impl.cfrc_ext,
d._impl.cfrc_int, d._impl.cfrc_int,
d._impl.cinert, d._impl.cinert,
@@ -1672,14 +1671,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
'act': out[0], 'act': out[0],
'act_dot': out[1], 'act_dot': out[1],
'actuator_force': out[2], 'actuator_force': out[2],
'_impl.actuator_length': out[3], 'actuator_length': out[3],
'_impl.actuator_moment': out[4], '_impl.actuator_moment': out[4],
'_impl.actuator_velocity': out[5], '_impl.actuator_velocity': out[5],
'_impl.cacc': out[6], '_impl.cacc': out[6],
'cam_xmat': out[7], 'cam_xmat': out[7],
'cam_xpos': out[8], 'cam_xpos': out[8],
'_impl.cdof': out[9], 'cdof': out[9],
'_impl.cdof_dot': out[10], 'cdof_dot': out[10],
'_impl.cfrc_ext': out[11], '_impl.cfrc_ext': out[11],
'_impl.cfrc_int': out[12], '_impl.cfrc_int': out[12],
'_impl.cinert': out[13], '_impl.cinert': out[13],
@@ -1834,7 +1833,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _step_shim( def _step_shim(
# Model # Model
@@ -2748,14 +2746,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'act': d.act.shape, 'act': d.act.shape,
'act_dot': d.act_dot.shape, 'act_dot': d.act_dot.shape,
'actuator_force': d.actuator_force.shape, 'actuator_force': d.actuator_force.shape,
'actuator_length': d._impl.actuator_length.shape, 'actuator_length': d.actuator_length.shape,
'actuator_moment': d._impl.actuator_moment.shape, 'actuator_moment': d._impl.actuator_moment.shape,
'actuator_velocity': d._impl.actuator_velocity.shape, 'actuator_velocity': d._impl.actuator_velocity.shape,
'cacc': d._impl.cacc.shape, 'cacc': d._impl.cacc.shape,
'cam_xmat': d.cam_xmat.shape, 'cam_xmat': d.cam_xmat.shape,
'cam_xpos': d.cam_xpos.shape, 'cam_xpos': d.cam_xpos.shape,
'cdof': d._impl.cdof.shape, 'cdof': d.cdof.shape,
'cdof_dot': d._impl.cdof_dot.shape, 'cdof_dot': d.cdof_dot.shape,
'cfrc_ext': d._impl.cfrc_ext.shape, 'cfrc_ext': d._impl.cfrc_ext.shape,
'cfrc_int': d._impl.cfrc_int.shape, 'cfrc_int': d._impl.cfrc_int.shape,
'cinert': d._impl.cinert.shape, 'cinert': d._impl.cinert.shape,
@@ -3337,14 +3335,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
d.act, d.act,
d.act_dot, d.act_dot,
d.actuator_force, d.actuator_force,
d._impl.actuator_length, d.actuator_length,
d._impl.actuator_moment, d._impl.actuator_moment,
d._impl.actuator_velocity, d._impl.actuator_velocity,
d._impl.cacc, d._impl.cacc,
d.cam_xmat, d.cam_xmat,
d.cam_xpos, d.cam_xpos,
d._impl.cdof, d.cdof,
d._impl.cdof_dot, d.cdof_dot,
d._impl.cfrc_ext, d._impl.cfrc_ext,
d._impl.cfrc_int, d._impl.cfrc_int,
d._impl.cinert, d._impl.cinert,
@@ -3470,14 +3468,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
'act': out[0], 'act': out[0],
'act_dot': out[1], 'act_dot': out[1],
'actuator_force': out[2], 'actuator_force': out[2],
'_impl.actuator_length': out[3], 'actuator_length': out[3],
'_impl.actuator_moment': out[4], '_impl.actuator_moment': out[4],
'_impl.actuator_velocity': out[5], '_impl.actuator_velocity': out[5],
'_impl.cacc': out[6], '_impl.cacc': out[6],
'cam_xmat': out[7], 'cam_xmat': out[7],
'cam_xpos': out[8], 'cam_xpos': out[8],
'_impl.cdof': out[9], 'cdof': out[9],
'_impl.cdof_dot': out[10], 'cdof_dot': out[10],
'_impl.cfrc_ext': out[11], '_impl.cfrc_ext': out[11],
'_impl.cfrc_int': out[12], '_impl.cfrc_int': out[12],
'_impl.cinert': out[13], '_impl.cinert': out[13],
mjx/mujoco/mjx/warp/forward_test.py
+3 -3
View File
@@ -133,7 +133,7 @@ class ForwardTest(parameterized.TestCase):
if m.nsite: if m.nsite:
tu.assert_attr_eq(dx, d, 'site_xpos') tu.assert_attr_eq(dx, d, 'site_xpos')
tu.assert_eq(dx.site_xmat, d.site_xmat.reshape((-1, 3, 3)), 'site_xmat') tu.assert_eq(dx.site_xmat, d.site_xmat.reshape((-1, 3, 3)), 'site_xmat')
tu.assert_attr_eq(dx._impl, d, 'cdof') tu.assert_attr_eq(dx, d, 'cdof')
tu.assert_attr_eq(dx._impl, d, 'cinert') tu.assert_attr_eq(dx._impl, d, 'cinert')
tu.assert_attr_eq(dx, d, 'subtree_com') tu.assert_attr_eq(dx, d, 'subtree_com')
if m.nlight: if m.nlight:
@@ -161,7 +161,7 @@ class ForwardTest(parameterized.TestCase):
tu.assert_contact_eq(d, dx, worldid=i) tu.assert_contact_eq(d, dx, worldid=i)
tu.assert_attr_eq(dx._impl, d, 'actuator_length') tu.assert_attr_eq(dx, d, 'actuator_length')
actuator_moment = np.zeros((m.nu, m.nv)) actuator_moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense( mujoco.mju_sparse2dense(
actuator_moment, actuator_moment,
@@ -175,7 +175,7 @@ class ForwardTest(parameterized.TestCase):
# fwd_velocity # fwd_velocity
tu.assert_attr_eq(dx._impl, d, 'actuator_velocity') tu.assert_attr_eq(dx._impl, d, 'actuator_velocity')
tu.assert_attr_eq(dx, d, 'cvel') tu.assert_attr_eq(dx, d, 'cvel')
tu.assert_attr_eq(dx._impl, d, 'cdof_dot') tu.assert_attr_eq(dx, d, 'cdof_dot')
tu.assert_attr_eq(dx._impl, d, 'qfrc_spring') tu.assert_attr_eq(dx._impl, d, 'qfrc_spring')
tu.assert_attr_eq(dx._impl, d, 'qfrc_damper') tu.assert_attr_eq(dx._impl, d, 'qfrc_damper')
tu.assert_attr_eq(dx, d, 'qfrc_gravcomp') tu.assert_attr_eq(dx, d, 'qfrc_gravcomp')
mjx/mujoco/mjx/warp/smooth.py
+3 -5
View File
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _kinematics_shim( def _kinematics_shim(
# Model # Model
@@ -310,7 +309,6 @@ _e = mjwarp.Constraint(
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init} **{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
) )
@ffi.format_args_for_warp @ffi.format_args_for_warp
def _tendon_shim( def _tendon_shim(
# Model # Model
@@ -397,7 +395,7 @@ def _tendon_shim(
def _tendon_jax_impl(m: types.Model, d: types.Data): def _tendon_jax_impl(m: types.Model, d: types.Data):
output_dims = { output_dims = {
'cdof': d._impl.cdof.shape, 'cdof': d.cdof.shape,
'geom_xmat': d.geom_xmat.shape, 'geom_xmat': d.geom_xmat.shape,
'geom_xpos': d.geom_xpos.shape, 'geom_xpos': d.geom_xpos.shape,
'qpos': d.qpos.shape, 'qpos': d.qpos.shape,
@@ -455,7 +453,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
m._impl.wrap_pulley_scale, m._impl.wrap_pulley_scale,
m._impl.wrap_site_pair_adr, m._impl.wrap_site_pair_adr,
m.wrap_type, m.wrap_type,
d._impl.cdof, d.cdof,
d.geom_xmat, d.geom_xmat,
d.geom_xpos, d.geom_xpos,
d.qpos, d.qpos,
@@ -469,7 +467,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
d._impl.wrap_xpos, d._impl.wrap_xpos,
) )
d = d.tree_replace({ d = d.tree_replace({
'_impl.cdof': out[0], 'cdof': out[0],
'geom_xmat': out[1], 'geom_xmat': out[1],
'geom_xpos': out[2], 'geom_xpos': out[2],
'qpos': out[3], 'qpos': out[3],
mjx/mujoco/mjx/warp/types.py
-3
View File
@@ -217,12 +217,9 @@ class ModelWarp(PyTreeNode):
class DataWarp(PyTreeNode): class DataWarp(PyTreeNode):
"""Derived fields from Data.""" """Derived fields from Data."""
actuator_length: jax.Array
actuator_moment: jax.Array actuator_moment: jax.Array
actuator_velocity: jax.Array actuator_velocity: jax.Array
cacc: jax.Array cacc: jax.Array
cdof: jax.Array
cdof_dot: jax.Array
cfrc_ext: jax.Array cfrc_ext: jax.Array
cfrc_int: jax.Array cfrc_int: jax.Array
cinert: jax.Array cinert: jax.Array