diff --git a/doc/includes/references.h b/doc/includes/references.h index ac6c4b35..199599f4 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -263,8 +263,8 @@ struct mjData_ { mjtNum* actuator_length; // actuator lengths (nu x 1) int* moment_rownnz; // number of non-zeros in actuator_moment row (nu x 1) int* moment_rowadr; // row start address in colind array (nu x 1) - int* moment_colind; // column indices in sparse Jacobian (nu x nv) - mjtNum* actuator_moment; // actuator moments (nu x nv) + int* moment_colind; // column indices in sparse Jacobian (nJmom x 1) + mjtNum* actuator_moment; // actuator moments (nJmom x 1) // computed by mj_fwdPosition/mj_crb mjtNum* crb; // com-based composite inertia and mass (nbody x 10) @@ -956,6 +956,7 @@ struct mjModel_ { int nB; // number of non-zeros in sparse body-dof matrix int nC; // number of non-zeros in sparse reduced dof-dof matrix int nD; // number of non-zeros in sparse dof-dof matrix + int nJmom; // number of non-zeros in sparse actuator_moment matrix int ntree; // number of kinematic trees under world body int ngravcomp; // number of bodies with nonzero gravcomp int nemax; // number of potential equality-constraint rows diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index 96b37acc..7ef96b12 100644 --- a/include/mujoco/mjdata.h +++ b/include/mujoco/mjdata.h @@ -291,8 +291,8 @@ struct mjData_ { mjtNum* actuator_length; // actuator lengths (nu x 1) int* moment_rownnz; // number of non-zeros in actuator_moment row (nu x 1) int* moment_rowadr; // row start address in colind array (nu x 1) - int* moment_colind; // column indices in sparse Jacobian (nu x nv) - mjtNum* actuator_moment; // actuator moments (nu x nv) + int* moment_colind; // column indices in sparse Jacobian (nJmom x 1) + mjtNum* actuator_moment; // actuator moments (nJmom x 1) // computed by mj_fwdPosition/mj_crb mjtNum* crb; // com-based composite inertia and mass (nbody x 10) diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h index 3cc59311..ebd752aa 100644 --- a/include/mujoco/mjmodel.h +++ b/include/mujoco/mjmodel.h @@ -663,6 +663,7 @@ struct mjModel_ { int nB; // number of non-zeros in sparse body-dof matrix int nC; // number of non-zeros in sparse reduced dof-dof matrix int nD; // number of non-zeros in sparse dof-dof matrix + int nJmom; // number of non-zeros in sparse actuator_moment matrix int ntree; // number of kinematic trees under world body int ngravcomp; // number of bodies with nonzero gravcomp int nemax; // number of potential equality-constraint rows diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 33b83f7b..05d51606 100644 --- a/include/mujoco/mjxmacro.h +++ b/include/mujoco/mjxmacro.h @@ -137,6 +137,7 @@ X ( nB ) \ X ( nC ) \ X ( nD ) \ + X ( nJmom ) \ XMJV( ntree ) \ X ( ngravcomp ) \ X ( nemax ) \ @@ -625,8 +626,8 @@ X ( mjtNum, actuator_length, nu, 1 ) \ X ( int, moment_rownnz, nu, 1 ) \ X ( int, moment_rowadr, nu, 1 ) \ - X ( int, moment_colind, nu, MJ_M(nv) ) \ - X ( mjtNum, actuator_moment, nu, MJ_M(nv) ) \ + X ( int, moment_colind, nJmom, 1 ) \ + X ( mjtNum, actuator_moment, nJmom, 1 ) \ X ( mjtNum, crb, nbody, 10 ) \ X ( mjtNum, qM, nM, 1 ) \ X ( mjtNum, qLD, nM, 1 ) \ diff --git a/introspect/structs.py b/introspect/structs.py index c50bea6e..9b8c4c47 100644 --- a/introspect/structs.py +++ b/introspect/structs.py @@ -1178,6 +1178,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([ type=ValueType(name='int'), doc='number of non-zeros in sparse dof-dof matrix', ), + StructFieldDecl( + name='nJmom', + type=ValueType(name='int'), + doc='number of non-zeros in sparse actuator_moment matrix', + ), StructFieldDecl( name='ntree', type=ValueType(name='int'), @@ -5166,7 +5171,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='int'), ), doc='column indices in sparse Jacobian', - array_extent=('nu', 'nv'), + array_extent=('nJmom',), ), StructFieldDecl( name='actuator_moment', @@ -5174,7 +5179,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([ inner_type=ValueType(name='mjtNum'), ), doc='actuator moments', - array_extent=('nu', 'nv'), + array_extent=('nJmom',), ), StructFieldDecl( name='crb', diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 200d742e..fedb014e 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -276,7 +276,7 @@ def make_data( 'actuator_length': (m.nu, float), 'moment_rownnz': (m.nu, jp.int32), 'moment_rowadr': (m.nu, jp.int32), - 'moment_colind': (m.nu, m.nv, jp.int32), + 'moment_colind': (m.nJmom, jp.int32), '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), @@ -431,22 +431,23 @@ def get_data_into( continue # MuJoCo actuator_moment is sparse, MJX uses a dense representation. - if field.name == 'actuator_moment' and m.nu: + if field.name == 'actuator_moment': moment_rownnz = np.zeros(m.nu, dtype=np.int32) moment_rowadr = np.zeros(m.nu, dtype=np.int32) - moment_colind = np.zeros(m.nu * m.nv, dtype=np.int32) - actuator_moment = np.zeros(m.nu * m.nv) - mujoco.mju_dense2sparse( - actuator_moment, - d.actuator_moment, - moment_rownnz, - moment_rowadr, - moment_colind, - ) + moment_colind = np.zeros(m.nJmom, dtype=np.int32) + actuator_moment = np.zeros(m.nJmom) + if m.nu: + mujoco.mju_dense2sparse( + actuator_moment, + d.actuator_moment, + moment_rownnz, + moment_rowadr, + moment_colind, + ) result_i.moment_rownnz[:] = moment_rownnz result_i.moment_rowadr[:] = moment_rowadr - result_i.moment_colind[:] = moment_colind.reshape((m.nu, m.nv)) - result_i.actuator_moment[:] = actuator_moment.reshape((m.nu, m.nv)) + result_i.moment_colind[:] = moment_colind + result_i.actuator_moment[:] = actuator_moment continue value = getattr(d_i, field.name) @@ -558,10 +559,10 @@ def put_data( moment = np.zeros((m.nu, m.nv)) mujoco.mju_sparse2dense( moment, - d.actuator_moment.reshape(-1), + d.actuator_moment, d.moment_rownnz, d.moment_rowadr, - d.moment_colind.reshape(-1), + d.moment_colind, ) fields['actuator_moment'] = moment diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index 39339297..11cf85fd 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -122,10 +122,10 @@ class SmoothTest(absltest.TestCase): moment = np.zeros((m.nu, m.nv)) mujoco.mju_sparse2dense( moment, - d.actuator_moment.reshape(-1), + d.actuator_moment, d.moment_rownnz, d.moment_rowadr, - d.moment_colind.reshape(-1), + d.moment_colind, ) _assert_eq(moment, dx.actuator_moment, 'actuator_moment') @@ -193,10 +193,10 @@ class SmoothTest(absltest.TestCase): moment = np.zeros((m.nu, m.nv)) mujoco.mju_sparse2dense( moment, - d.actuator_moment.reshape(-1), + d.actuator_moment, d.moment_rownnz, d.moment_rowadr, - d.moment_colind.reshape(-1), + d.moment_colind, ) _assert_eq(moment, dx.actuator_moment, 'actuator_moment') diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index dcfbdb6f..89dab88e 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -536,6 +536,7 @@ class Model(PyTreeNode): nM: number of non-zeros in sparse inertia matrix nD: number of non-zeros in sparse dof-dof matrix nB: number of non-zeros in sparse body-dof matrix + nJmom: number of non-zeros in sparse actuator_moment matrix ntree: number of kinematic trees under world body ngravcomp: number of bodies with nonzero gravcomp nuserdata: size of userdata array @@ -855,6 +856,7 @@ class Model(PyTreeNode): nM: int # pylint:disable=invalid-name nD: int # pylint:disable=invalid-name nB: int # pylint:disable=invalid-name + nJmom: int ntree: int = _restricted_to('mujoco') ngravcomp: int nuserdata: int @@ -1236,8 +1238,8 @@ class Data(PyTreeNode): actuator_length: actuator lengths (nu,) moment_rownnz: number of non-zeros in actuator_moment row (nu,) moment_rowadr: row start address in colind array (nu,) - moment_colind: column indices in sparse Jacobian (nu, nv) - actuator_moment: actuator moments (nu, nv) + moment_colind: column indices in sparse Jacobian (nJmom,) + actuator_moment: actuator moments (nJmom,) crb: com-based composite inertia and mass (nbody, 10) qM: total inertia if sparse: (nM,) if dense: (nv, nv) diff --git a/mjx/mujoco/mjx/integration_test/smooth_test.py b/mjx/mujoco/mjx/integration_test/smooth_test.py index 8d6b2137..5ec8426a 100644 --- a/mjx/mujoco/mjx/integration_test/smooth_test.py +++ b/mjx/mujoco/mjx/integration_test/smooth_test.py @@ -76,10 +76,10 @@ class TransmissionIntegrationTest(parameterized.TestCase): moment = np.zeros((m.nu, m.nv)) mujoco.mju_sparse2dense( moment, - d.actuator_moment.reshape(-1), + d.actuator_moment, d.moment_rownnz, d.moment_rowadr, - d.moment_colind.reshape(-1), + d.moment_colind, ) _assert_eq( moment, diff --git a/src/user/user_model.cc b/src/user/user_model.cc index a3e43f45..927a8b0b 100644 --- a/src/user/user_model.cc +++ b/src/user/user_model.cc @@ -777,6 +777,7 @@ void mjCModel::Clear() { nM = 0; nD = 0; nB = 0; + nJmom = 0; njmax = -1; nconmax = -1; nmocap = 0; @@ -2490,7 +2491,7 @@ void mjCModel::CopyTree(mjModel* m) { } } } - m->nC = nC = 2 * nOD + nv; + m->nC = nC = 2 * nOD + nv; } // copy plugin data @@ -4156,6 +4157,13 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) { // copy objects outsite kinematic tree (including keyframes) CopyObjects(m); + // compute nJmom + for (int i = 0; i < nu; i++) { + // dense rows + nJmom += nv; + } + m->nJmom = nJmom; + // scale mass if (compiler.settotalmass>0) { mj_setTotalmass(m, compiler.settotalmass); @@ -4292,7 +4300,7 @@ bool mjCModel::CopyBack(const mjModel* m) { neq!=m->neq || ntendon!=m->ntendon || nwrap!=m->nwrap || nsensor!=m->nsensor || nnumeric!=m->nnumeric || nnumericdata!=m->nnumericdata || ntext!=m->ntext || ntextdata!=m->ntextdata || nnames!=m->nnames || nM!=m->nM || nD!=m->nD || nC!=m->nC || - nB!=m->nB || nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax || + nB!=m->nB || nJmom!=m->nJmom ||nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax || npaths!=m->npaths) { errInfo = mjCError(0, "incompatible models in CopyBack"); return false; diff --git a/src/user/user_model.h b/src/user/user_model.h index 41865ac6..d4d7dffa 100644 --- a/src/user/user_model.h +++ b/src/user/user_model.h @@ -118,6 +118,7 @@ class mjCModel_ : public mjsElement { int nB; // number of non-zeros in sparse body-dof matrix int nC; // number of non-zeros in reduced sparse dof-dof matrix int nD; // number of non-zeros in sparse dof-dof matrix + int nJmom; // number of non-zeros in sparse actuator_moment matrix // statistics, as computed by mj_setConst double meaninertia_auto; // mean diagonal inertia, as computed by mj_setConst diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc index 0874a14c..53c6277b 100644 --- a/test/user/user_model_test.cc +++ b/test/user/user_model_test.cc @@ -116,6 +116,30 @@ TEST_F(UserCModelTest, SameFrame) { mj_deleteModel(model); } +TEST_F(UserCModelTest, ActuatorSparsity) { + static constexpr char xml[] = R"( + + + + + + + + + + + + + + + + + )"; + mjModel* m = LoadModelFromString(xml); + ASSERT_EQ(m->nJmom, 4); + mj_deleteModel(m); +} + // ------------- test automatic inference of nuser_xxx ------------------------- diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index a8d063be..8858f99c 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/unity/Runtime/Bindings/MjBindings.cs @@ -5244,6 +5244,7 @@ public unsafe struct mjModel_ { public int nB; public int nC; public int nD; + public int nJmom; public int ntree; public int ngravcomp; public int nemax;