Make actuator_length, cdof and cdof_dot public.
PiperOrigin-RevId: 842182467 Change-Id: Id14a1c8266e13a6f95e1c87ea07da993b614b62e
This commit is contained in:
committed by
Copybara-Service
parent
f65c1cdfd1
commit
1483aefe5f
@@ -2,6 +2,19 @@
|
||||
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)
|
||||
--------------------------------
|
||||
|
||||
|
||||
@@ -182,7 +182,7 @@ def fwd_actuation(m: Model, d: Data) -> Data:
|
||||
m.actuator_gainprm,
|
||||
m.actuator_biastype,
|
||||
m.actuator_biasprm,
|
||||
d._impl.actuator_length,
|
||||
d.actuator_length,
|
||||
d._impl.actuator_velocity,
|
||||
ctrl_act,
|
||||
jp.array(m.actuator_lengthrange),
|
||||
|
||||
@@ -552,6 +552,7 @@ def _make_data_public_fields(m: types.Model) -> Dict[str, Any]:
|
||||
'cam_xmat': (m.ncam, 3, 3, float_),
|
||||
'subtree_com': (m.nbody, 3, float_),
|
||||
'actuator_force': (m.nu, float_),
|
||||
'actuator_length': (m.nu, float_),
|
||||
'qfrc_bias': (m.nv, float_),
|
||||
'qfrc_gravcomp': (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_inverse': (m.nv, float_),
|
||||
'cvel': (m.nbody, 6, float_),
|
||||
'cdof': (m.nv, 6, float_),
|
||||
'cdof_dot': (m.nv, 6, float_),
|
||||
'ten_length': (m.ntendon, float_),
|
||||
}
|
||||
zero_fields = {
|
||||
@@ -619,14 +622,12 @@ def _make_data_jax(
|
||||
|
||||
zero_impl_fields = {
|
||||
'solver_niter': (int_,),
|
||||
'cdof': (m.nv, 6, float_),
|
||||
'cinert': (m.nbody, 10, float_),
|
||||
'ten_wrapadr': (m.ntendon, np.int32),
|
||||
'ten_wrapnum': (m.ntendon, np.int32),
|
||||
'ten_J': (m.ntendon, m.nv, float_),
|
||||
'wrap_obj': (m.nwrap, 2, np.int32),
|
||||
'wrap_xpos': (m.nwrap, 6, float_),
|
||||
'actuator_length': (m.nu, float_),
|
||||
'actuator_moment': (m.nu, m.nv, float_),
|
||||
'crb': (m.nbody, 10, 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_),
|
||||
'ten_velocity': (m.ntendon, float_),
|
||||
'actuator_velocity': (m.nu, float_),
|
||||
'cdof_dot': (m.nv, 6, float_),
|
||||
'cacc': (m.nbody, 6, float_),
|
||||
'cfrc_int': (m.nbody, 6, float_),
|
||||
'cfrc_ext': (m.nbody, 6, float_),
|
||||
@@ -713,7 +713,6 @@ def _make_data_c(
|
||||
nbvhdynamic = get(m, 'nbvhdynamic')
|
||||
zero_impl_fields = {
|
||||
'solver_niter': (int_,),
|
||||
'cdof': (m.nv, 6, float_),
|
||||
'cinert': (m.nbody, 10, float_),
|
||||
'light_xpos': (m.nlight, 3, float_),
|
||||
'light_xdir': (m.nlight, 3, float_),
|
||||
@@ -732,7 +731,6 @@ def _make_data_c(
|
||||
'ten_wrapnum': (m.ntendon, np.int32),
|
||||
'wrap_obj': (m.nwrap, 2, np.int32),
|
||||
'wrap_xpos': (m.nwrap, 6, float_),
|
||||
'actuator_length': (m.nu, float_),
|
||||
'moment_rownnz': (m.nu, np.int32),
|
||||
'moment_rowadr': (m.nu, np.int32),
|
||||
'moment_colind': (m.nJmom, np.int32),
|
||||
@@ -762,7 +760,6 @@ def _make_data_c(
|
||||
'qLU': (m.nD, float_),
|
||||
'qfrc_spring': (m.nv, float_),
|
||||
'qfrc_damper': (m.nv, float_),
|
||||
'cdof_dot': (m.nv, 6, float_),
|
||||
'cacc': (m.nbody, 6, float_),
|
||||
'cfrc_int': (m.nbody, 6, float_),
|
||||
'cfrc_ext': (m.nbody, 6, float_),
|
||||
|
||||
@@ -402,10 +402,10 @@ class DataIOTest(parameterized.TestCase):
|
||||
self.assertEqual(d.geom_xpos.shape, (6, 3))
|
||||
self.assertEqual(d.geom_xmat.shape, (6, 3, 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.crb.shape, (nbody, 10))
|
||||
self.assertEqual(d._impl.actuator_length.shape, (1,))
|
||||
self.assertEqual(d.actuator_length.shape, (1,))
|
||||
if impl == 'jax':
|
||||
self.assertEqual(d._impl.actuator_moment.shape, (1, nv))
|
||||
elif impl == 'c':
|
||||
@@ -422,7 +422,7 @@ class DataIOTest(parameterized.TestCase):
|
||||
self.assertEqual(d._impl.efc_D.shape, (nefc,))
|
||||
self.assertEqual(d._impl.actuator_velocity.shape, (1,))
|
||||
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_passive.shape, (nv,))
|
||||
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.xpos, d.xpos)
|
||||
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
|
||||
self.assertFalse(
|
||||
|
||||
@@ -175,7 +175,7 @@ def sensor_pos(m: Model, d: Data) -> Data:
|
||||
elif sensor_type == SensorType.TENDONPOS:
|
||||
sensor = d.ten_length[objid]
|
||||
elif sensor_type == SensorType.ACTUATORPOS:
|
||||
sensor = d._impl.actuator_length[objid]
|
||||
sensor = d.actuator_length[objid]
|
||||
elif sensor_type == SensorType.BALLQUAT:
|
||||
jnt_qposadr = m.jnt_qposadr[objid, None] + np.arange(4)[None]
|
||||
quat = d.qpos[jnt_qposadr]
|
||||
|
||||
@@ -212,7 +212,7 @@ def com_pos(m: Model, d: Data) -> Data:
|
||||
d.xanchor,
|
||||
d.xaxis,
|
||||
)
|
||||
d = d.tree_replace({'_impl.cdof': cdof})
|
||||
d = d.tree_replace({'cdof': cdof})
|
||||
|
||||
return d
|
||||
|
||||
@@ -305,8 +305,8 @@ def crb(m: Model, d: Data) -> Data:
|
||||
d = d.tree_replace({'_impl.crb': crb_body})
|
||||
|
||||
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)
|
||||
qm = support.make_m(m, crb_cdof, d._impl.cdof, m.dof_armature)
|
||||
crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
|
||||
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
|
||||
d = d.tree_replace({'_impl.qM': qm})
|
||||
return d
|
||||
|
||||
@@ -445,11 +445,11 @@ def com_vel(m: Model, d: Data) -> Data:
|
||||
'jvv',
|
||||
'bv',
|
||||
m.jnt_type,
|
||||
d._impl.cdof,
|
||||
d.cdof,
|
||||
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
|
||||
|
||||
@@ -576,7 +576,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
|
||||
return cacc
|
||||
|
||||
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):
|
||||
@@ -594,7 +594,7 @@ def rne(m: Model, d: Data, flg_acc: bool = False) -> Data:
|
||||
return cfrc
|
||||
|
||||
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)
|
||||
|
||||
@@ -806,8 +806,8 @@ def rne_postconstraint(m: Model, d: Data) -> Data:
|
||||
)
|
||||
|
||||
# cacc = cacc_parent + cdofdot * qvel + cdof * qacc
|
||||
cacc_vel = d._impl.cdof_dot.T @ (mask * d.qvel)
|
||||
cacc_acc = d._impl.cdof.T @ (mask * d.qacc)
|
||||
cacc_vel = d.cdof_dot.T @ (mask * d.qvel)
|
||||
cacc_acc = d.cdof.T @ (mask * d.qacc)
|
||||
cacc = cacc_parent + cacc_vel + cacc_acc
|
||||
|
||||
# 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))
|
||||
|
||||
d = d.tree_replace(
|
||||
{'_impl.actuator_length': length, '_impl.actuator_moment': moment}
|
||||
{'actuator_length': length, '_impl.actuator_moment': moment}
|
||||
)
|
||||
return d
|
||||
|
||||
|
||||
@@ -81,7 +81,7 @@ class SmoothTest(absltest.TestCase):
|
||||
dx = jax.jit(mjx.com_pos)(mx, mjx.put_data(m, d))
|
||||
_assert_attr_eq(d, dx, 'subtree_com')
|
||||
_assert_attr_eq(d, dx._impl, 'cinert')
|
||||
_assert_attr_eq(d, dx._impl, 'cdof')
|
||||
_assert_attr_eq(d, dx, 'cdof')
|
||||
# camlight
|
||||
dx = jax.jit(mjx.camlight)(mx, mjx.put_data(m, d))
|
||||
_assert_attr_eq(d, dx, 'cam_xpos')
|
||||
@@ -100,7 +100,7 @@ class SmoothTest(absltest.TestCase):
|
||||
# com_vel
|
||||
dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
|
||||
_assert_attr_eq(d, dx, 'cvel')
|
||||
_assert_attr_eq(d, dx._impl, 'cdof_dot')
|
||||
_assert_attr_eq(d, dx, 'cdof_dot')
|
||||
# rne
|
||||
dx = jax.jit(mjx.rne)(mx, mjx.put_data(m, d))
|
||||
_assert_attr_eq(d, dx, 'qfrc_bias')
|
||||
@@ -125,7 +125,7 @@ class SmoothTest(absltest.TestCase):
|
||||
_assert_attr_eq(d, dx, 'ten_length')
|
||||
# transmission
|
||||
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
|
||||
moment = np.zeros((m.nu, m.nv))
|
||||
@@ -196,7 +196,7 @@ class SmoothTest(absltest.TestCase):
|
||||
|
||||
mujoco.mj_transmission(m, d)
|
||||
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
|
||||
moment = np.zeros((m.nu, m.nv))
|
||||
|
||||
@@ -149,9 +149,9 @@ def jac(
|
||||
mask = mask[jp.array(m.dof_bodyid)] > 0
|
||||
|
||||
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)
|
||||
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
|
||||
|
||||
@@ -169,8 +169,8 @@ def jac_dot(
|
||||
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])
|
||||
|
||||
cdof = d._impl.cdof
|
||||
cdof_dot = d._impl.cdof_dot
|
||||
cdof = d.cdof
|
||||
cdof_dot = d.cdof_dot
|
||||
|
||||
# check for quaternion
|
||||
jnt_type = m.jnt_type[m.dof_jntid]
|
||||
|
||||
@@ -1021,7 +1021,6 @@ class DataC(PyTreeNode):
|
||||
plugin_data: jax.Array
|
||||
light_xpos: jax.Array
|
||||
light_xdir: jax.Array
|
||||
cdof: jax.Array
|
||||
cinert: jax.Array
|
||||
flexvert_xpos: jax.Array
|
||||
flexelem_aabb: jax.Array
|
||||
@@ -1039,7 +1038,6 @@ class DataC(PyTreeNode):
|
||||
ten_J: jax.Array # pylint:disable=invalid-name
|
||||
wrap_obj: jax.Array
|
||||
wrap_xpos: jax.Array
|
||||
actuator_length: jax.Array
|
||||
moment_rownnz: jax.Array # pylint:disable=invalid-name
|
||||
moment_rowadr: 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
|
||||
ten_velocity: jax.Array
|
||||
actuator_velocity: jax.Array
|
||||
cdof_dot: jax.Array
|
||||
|
||||
qfrc_spring: jax.Array
|
||||
qfrc_damper: jax.Array
|
||||
subtree_linvel: jax.Array
|
||||
@@ -1095,14 +1093,12 @@ class DataJAX(PyTreeNode):
|
||||
nefc: int
|
||||
ncon: int
|
||||
solver_niter: jax.Array
|
||||
cdof: jax.Array
|
||||
cinert: jax.Array
|
||||
ten_wrapadr: jax.Array
|
||||
ten_wrapnum: jax.Array
|
||||
ten_J: jax.Array # pylint:disable=invalid-name
|
||||
wrap_obj: jax.Array
|
||||
wrap_xpos: jax.Array
|
||||
actuator_length: jax.Array
|
||||
actuator_moment: jax.Array
|
||||
crb: jax.Array
|
||||
qM: jax.Array # pylint:disable=invalid-name
|
||||
@@ -1111,7 +1107,7 @@ class DataJAX(PyTreeNode):
|
||||
qLDiagInv: jax.Array # pylint:disable=invalid-name
|
||||
ten_velocity: jax.Array
|
||||
actuator_velocity: jax.Array
|
||||
cdof_dot: jax.Array
|
||||
|
||||
cacc: jax.Array
|
||||
cfrc_int: jax.Array
|
||||
cfrc_ext: jax.Array
|
||||
@@ -1171,12 +1167,15 @@ class Data(PyTreeNode):
|
||||
cam_xmat: jax.Array
|
||||
subtree_com: jax.Array
|
||||
cvel: jax.Array
|
||||
cdof: jax.Array
|
||||
cdof_dot: jax.Array
|
||||
qfrc_bias: jax.Array
|
||||
qfrc_gravcomp: jax.Array
|
||||
qfrc_fluid: jax.Array
|
||||
qfrc_passive: jax.Array
|
||||
qfrc_actuator: jax.Array
|
||||
actuator_force: jax.Array
|
||||
actuator_length: jax.Array
|
||||
qfrc_smooth: jax.Array
|
||||
qacc_smooth: jax.Array
|
||||
qfrc_constraint: jax.Array
|
||||
|
||||
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
|
||||
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
||||
)
|
||||
|
||||
|
||||
@ffi.format_args_for_warp
|
||||
def _collision_shim(
|
||||
# Model
|
||||
|
||||
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
|
||||
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
||||
)
|
||||
|
||||
|
||||
@ffi.format_args_for_warp
|
||||
def _forward_shim(
|
||||
# Model
|
||||
@@ -952,14 +951,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
|
||||
'act': d.act.shape,
|
||||
'act_dot': d.act_dot.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_velocity': d._impl.actuator_velocity.shape,
|
||||
'cacc': d._impl.cacc.shape,
|
||||
'cam_xmat': d.cam_xmat.shape,
|
||||
'cam_xpos': d.cam_xpos.shape,
|
||||
'cdof': d._impl.cdof.shape,
|
||||
'cdof_dot': d._impl.cdof_dot.shape,
|
||||
'cdof': d.cdof.shape,
|
||||
'cdof_dot': d.cdof_dot.shape,
|
||||
'cfrc_ext': d._impl.cfrc_ext.shape,
|
||||
'cfrc_int': d._impl.cfrc_int.shape,
|
||||
'cinert': d._impl.cinert.shape,
|
||||
@@ -1539,14 +1538,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
|
||||
d.act,
|
||||
d.act_dot,
|
||||
d.actuator_force,
|
||||
d._impl.actuator_length,
|
||||
d.actuator_length,
|
||||
d._impl.actuator_moment,
|
||||
d._impl.actuator_velocity,
|
||||
d._impl.cacc,
|
||||
d.cam_xmat,
|
||||
d.cam_xpos,
|
||||
d._impl.cdof,
|
||||
d._impl.cdof_dot,
|
||||
d.cdof,
|
||||
d.cdof_dot,
|
||||
d._impl.cfrc_ext,
|
||||
d._impl.cfrc_int,
|
||||
d._impl.cinert,
|
||||
@@ -1672,14 +1671,14 @@ def _forward_jax_impl(m: types.Model, d: types.Data):
|
||||
'act': out[0],
|
||||
'act_dot': out[1],
|
||||
'actuator_force': out[2],
|
||||
'_impl.actuator_length': out[3],
|
||||
'actuator_length': out[3],
|
||||
'_impl.actuator_moment': out[4],
|
||||
'_impl.actuator_velocity': out[5],
|
||||
'_impl.cacc': out[6],
|
||||
'cam_xmat': out[7],
|
||||
'cam_xpos': out[8],
|
||||
'_impl.cdof': out[9],
|
||||
'_impl.cdof_dot': out[10],
|
||||
'cdof': out[9],
|
||||
'cdof_dot': out[10],
|
||||
'_impl.cfrc_ext': out[11],
|
||||
'_impl.cfrc_int': out[12],
|
||||
'_impl.cinert': out[13],
|
||||
@@ -1834,7 +1833,6 @@ _e = mjwarp.Constraint(
|
||||
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
||||
)
|
||||
|
||||
|
||||
@ffi.format_args_for_warp
|
||||
def _step_shim(
|
||||
# Model
|
||||
@@ -2748,14 +2746,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
|
||||
'act': d.act.shape,
|
||||
'act_dot': d.act_dot.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_velocity': d._impl.actuator_velocity.shape,
|
||||
'cacc': d._impl.cacc.shape,
|
||||
'cam_xmat': d.cam_xmat.shape,
|
||||
'cam_xpos': d.cam_xpos.shape,
|
||||
'cdof': d._impl.cdof.shape,
|
||||
'cdof_dot': d._impl.cdof_dot.shape,
|
||||
'cdof': d.cdof.shape,
|
||||
'cdof_dot': d.cdof_dot.shape,
|
||||
'cfrc_ext': d._impl.cfrc_ext.shape,
|
||||
'cfrc_int': d._impl.cfrc_int.shape,
|
||||
'cinert': d._impl.cinert.shape,
|
||||
@@ -3337,14 +3335,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
|
||||
d.act,
|
||||
d.act_dot,
|
||||
d.actuator_force,
|
||||
d._impl.actuator_length,
|
||||
d.actuator_length,
|
||||
d._impl.actuator_moment,
|
||||
d._impl.actuator_velocity,
|
||||
d._impl.cacc,
|
||||
d.cam_xmat,
|
||||
d.cam_xpos,
|
||||
d._impl.cdof,
|
||||
d._impl.cdof_dot,
|
||||
d.cdof,
|
||||
d.cdof_dot,
|
||||
d._impl.cfrc_ext,
|
||||
d._impl.cfrc_int,
|
||||
d._impl.cinert,
|
||||
@@ -3470,14 +3468,14 @@ def _step_jax_impl(m: types.Model, d: types.Data):
|
||||
'act': out[0],
|
||||
'act_dot': out[1],
|
||||
'actuator_force': out[2],
|
||||
'_impl.actuator_length': out[3],
|
||||
'actuator_length': out[3],
|
||||
'_impl.actuator_moment': out[4],
|
||||
'_impl.actuator_velocity': out[5],
|
||||
'_impl.cacc': out[6],
|
||||
'cam_xmat': out[7],
|
||||
'cam_xpos': out[8],
|
||||
'_impl.cdof': out[9],
|
||||
'_impl.cdof_dot': out[10],
|
||||
'cdof': out[9],
|
||||
'cdof_dot': out[10],
|
||||
'_impl.cfrc_ext': out[11],
|
||||
'_impl.cfrc_int': out[12],
|
||||
'_impl.cinert': out[13],
|
||||
|
||||
@@ -133,7 +133,7 @@ class ForwardTest(parameterized.TestCase):
|
||||
if m.nsite:
|
||||
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_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, d, 'subtree_com')
|
||||
if m.nlight:
|
||||
@@ -161,7 +161,7 @@ class ForwardTest(parameterized.TestCase):
|
||||
|
||||
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))
|
||||
mujoco.mju_sparse2dense(
|
||||
actuator_moment,
|
||||
@@ -175,7 +175,7 @@ class ForwardTest(parameterized.TestCase):
|
||||
# fwd_velocity
|
||||
tu.assert_attr_eq(dx._impl, d, 'actuator_velocity')
|
||||
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_damper')
|
||||
tu.assert_attr_eq(dx, d, 'qfrc_gravcomp')
|
||||
|
||||
@@ -42,7 +42,6 @@ _e = mjwarp.Constraint(
|
||||
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
||||
)
|
||||
|
||||
|
||||
@ffi.format_args_for_warp
|
||||
def _kinematics_shim(
|
||||
# Model
|
||||
@@ -310,7 +309,6 @@ _e = mjwarp.Constraint(
|
||||
**{f.name: None for f in dataclasses.fields(mjwarp.Constraint) if f.init}
|
||||
)
|
||||
|
||||
|
||||
@ffi.format_args_for_warp
|
||||
def _tendon_shim(
|
||||
# Model
|
||||
@@ -397,7 +395,7 @@ def _tendon_shim(
|
||||
|
||||
def _tendon_jax_impl(m: types.Model, d: types.Data):
|
||||
output_dims = {
|
||||
'cdof': d._impl.cdof.shape,
|
||||
'cdof': d.cdof.shape,
|
||||
'geom_xmat': d.geom_xmat.shape,
|
||||
'geom_xpos': d.geom_xpos.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_site_pair_adr,
|
||||
m.wrap_type,
|
||||
d._impl.cdof,
|
||||
d.cdof,
|
||||
d.geom_xmat,
|
||||
d.geom_xpos,
|
||||
d.qpos,
|
||||
@@ -469,7 +467,7 @@ def _tendon_jax_impl(m: types.Model, d: types.Data):
|
||||
d._impl.wrap_xpos,
|
||||
)
|
||||
d = d.tree_replace({
|
||||
'_impl.cdof': out[0],
|
||||
'cdof': out[0],
|
||||
'geom_xmat': out[1],
|
||||
'geom_xpos': out[2],
|
||||
'qpos': out[3],
|
||||
|
||||
@@ -217,12 +217,9 @@ class ModelWarp(PyTreeNode):
|
||||
|
||||
class DataWarp(PyTreeNode):
|
||||
"""Derived fields from Data."""
|
||||
actuator_length: jax.Array
|
||||
actuator_moment: jax.Array
|
||||
actuator_velocity: jax.Array
|
||||
cacc: jax.Array
|
||||
cdof: jax.Array
|
||||
cdof_dot: jax.Array
|
||||
cfrc_ext: jax.Array
|
||||
cfrc_int: jax.Array
|
||||
cinert: jax.Array
|
||||
|
||||
Reference in New Issue
Block a user