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

Add nJmom, the number of non-zeros in sparse actuator_moment matrix, to mjModel.

Browse Source 
PiperOrigin-RevId: 698343986
Change-Id: I3d4a4dea1ec095d4e7b231c2fffc0e6241cee39f
This commit is contained in:
Taylor Howell
2024-11-20 04:29:55 -08:00
committed by Copybara-Service
parent cd84011d7b
commit aae5fd6906
13 changed files with 78 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+3 -2
View File
@@ -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
include/mujoco/mjdata.h
+2 -2
View File
@@ -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)
include/mujoco/mjmodel.h
+1
View File
@@ -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
include/mujoco/mjxmacro.h
+3 -2
View File
@@ -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 ) \
introspect/structs.py
+7 -2
View File
@@ -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',
mjx/mujoco/mjx/_src/io.py
+16 -15
View File
@@ -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
mjx/mujoco/mjx/_src/smooth_test.py
+4 -4
View File
@@ -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')
mjx/mujoco/mjx/_src/types.py
+4 -2
View File
@@ -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)
mjx/mujoco/mjx/integration_test/smooth_test.py
+2 -2
View File
@@ -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,
src/user/user_model.cc
+10 -2
View File
@@ -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;
src/user/user_model.h
+1
View File
@@ -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
test/user/user_model_test.cc
+24
View File
@@ -116,6 +116,30 @@ TEST_F(UserCModelTest, SameFrame) {
mj_deleteModel(model);
}
TEST_F(UserCModelTest, ActuatorSparsity) {
static constexpr char xml[] = R"(
<mujoco>
<worldbody>
<body>
<geom size="1"/>
<joint name="a"/>
<body>
<geom size="1"/>
<joint name="b"/>
</body>
</body>
</worldbody>
<actuator>
<motor joint="a"/>
<motor joint="b"/>
</actuator>
</mujoco>
)";
mjModel* m = LoadModelFromString(xml);
ASSERT_EQ(m->nJmom, 4);
mj_deleteModel(m);
}
// ------------- test automatic inference of nuser_xxx -------------------------
unity/Runtime/Bindings/MjBindings.cs
+1
View File
@@ -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;