From a51f346059bbefad31b6c00572a4f137e053a6cd Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Fri, 1 Nov 2024 08:05:14 -0700
Subject: [PATCH] Use sparse (uncompressed) actuator_moment in mj_transmission.

PiperOrigin-RevId: 692179704
Change-Id: Ic30ac5a98dc13de2028e378df65dc88ba3912bf5
---
 mjx/mujoco/mjx/_src/io.py                     |  47 ++++++-
 mjx/mujoco/mjx/_src/smooth_test.py            |  24 +++-
 mjx/mujoco/mjx/_src/types.py                  |   6 +
 .../mjx/integration_test/smooth_test.py       |  25 +++-
 python/LQR.ipynb                              |  10 +-
 src/engine/engine_core_smooth.c               | 128 +++++++++++-------
 src/engine/engine_derivative.c                |  11 +-
 src/engine/engine_forward.c                   |  12 +-
 src/engine/engine_io.c                        |   3 -
 src/engine/engine_print.c                     |   9 +-
 src/engine/engine_setconst.c                  |  28 +++-
 src/engine/engine_util_sparse.h               |   4 +-
 test/engine/engine_derivative_test.cc         |   4 +-
 13 files changed, 230 insertions(+), 81 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py
index 00ba75e5..28b84bbe 100644
--- a/mjx/mujoco/mjx/_src/io.py
+++ b/mjx/mujoco/mjx/_src/io.py
@@ -274,6 +274,9 @@ def make_data(
         'wrap_obj': (m.nwrap, 2, jp.int32),
         'wrap_xpos': (m.nwrap, 6, float),
         '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),
         '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),
@@ -427,6 +430,25 @@ def get_data_into(
         result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address]
         continue
 
+      # MuJoCo actuator_moment is sparse, MJX uses a dense representation.
+      if field.name == 'actuator_moment' and m.nu:
+        moment_rownnz = np.zeros(m.nu, dtype=int)
+        moment_rowadr = np.zeros(m.nu, dtype=int)
+        moment_colind = np.zeros(m.nu * m.nv, dtype=int)
+        actuator_moment = np.zeros(m.nu * m.nv)
+        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))
+        continue
+
       value = getattr(d_i, field.name)
 
       if field.name in ('nefc', 'ncon'):
@@ -532,6 +554,17 @@ def put_data(
   # MJX does not support islanding, so only transfer the first solver_niter
   fields['solver_niter'] = fields['solver_niter'][0]
 
+  # convert sparse representation of actuator_moment to dense matrix
+  moment = np.zeros((m.nu, m.nv))
+  mujoco.mju_sparse2dense(
+      moment,
+      d.actuator_moment.reshape(-1),
+      d.moment_rownnz,
+      d.moment_rowadr,
+      d.moment_colind.reshape(-1),
+  )
+  fields['actuator_moment'] = moment
+
   contact, contact_map = _make_contact(d.contact, dim, efc_address)
 
   # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
@@ -539,12 +572,14 @@ def put_data(
   # neither: it contains zeros for inactive constraints, and efc_J is always
   # (nefc, nv).  this may change in the future.
   if mujoco.mj_isSparse(m):
-    nr = d.efc_J_rownnz.shape[0]
-    efc_j = np.zeros((nr, m.nv))
-    for i in range(nr):
-      rowadr = d.efc_J_rowadr[i]
-      for j in range(d.efc_J_rownnz[i]):
-        efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
+    efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv))
+    mujoco.mju_sparse2dense(
+        efc_j,
+        fields['efc_J'],
+        d.efc_J_rownnz,
+        d.efc_J_rowadr,
+        d.efc_J_colind,
+    )
     fields['efc_J'] = efc_j
   else:
     fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
index fc14ac65..39339297 100644
--- a/mjx/mujoco/mjx/_src/smooth_test.py
+++ b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -117,7 +117,17 @@ class SmoothTest(absltest.TestCase):
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
     _assert_attr_eq(d, dx, 'actuator_length')
-    _assert_attr_eq(d, dx, 'actuator_moment')
+
+    # convert sparse actuator_moment to dense representation
+    moment = np.zeros((m.nu, m.nv))
+    mujoco.mju_sparse2dense(
+        moment,
+        d.actuator_moment.reshape(-1),
+        d.moment_rownnz,
+        d.moment_rowadr,
+        d.moment_colind.reshape(-1),
+    )
+    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
 
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""
@@ -178,7 +188,17 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_transmission(m, d)
     dx = jax.jit(mjx.transmission)(mx, dx)
     _assert_attr_eq(d, dx, 'actuator_length')
-    _assert_attr_eq(d, dx, 'actuator_moment')
+
+    # convert sparse actuator_moment to dense representation
+    moment = np.zeros((m.nu, m.nv))
+    mujoco.mju_sparse2dense(
+        moment,
+        d.actuator_moment.reshape(-1),
+        d.moment_rownnz,
+        d.moment_rowadr,
+        d.moment_colind.reshape(-1),
+    )
+    _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
 
   def test_subtree_vel(self):
     """Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel."""
diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py
index c36fa66c..e1df9681 100644
--- a/mjx/mujoco/mjx/_src/types.py
+++ b/mjx/mujoco/mjx/_src/types.py
@@ -1228,6 +1228,9 @@ class Data(PyTreeNode):
     wrap_obj: geom id; -1: site; -2: pulley                     (nwrap*2,)
     wrap_xpos: Cartesian 3D points in all path                  (nwrap*2, 3)
     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)
     crb: com-based composite inertia and mass                   (nbody, 10)
     qM: total inertia                                if sparse: (nM,)
@@ -1350,6 +1353,9 @@ class Data(PyTreeNode):
   wrap_obj: jax.Array
   wrap_xpos: jax.Array
   actuator_length: jax.Array
+  moment_rownnz: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  moment_rowadr: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
+  moment_colind: jax.Array = _restricted_to('mujoco')  # pylint:disable=invalid-name
   actuator_moment: jax.Array
   crb: jax.Array
   qM: jax.Array  # pylint:disable=invalid-name
diff --git a/mjx/mujoco/mjx/integration_test/smooth_test.py b/mjx/mujoco/mjx/integration_test/smooth_test.py
index 5359245e..8d6b2137 100644
--- a/mjx/mujoco/mjx/integration_test/smooth_test.py
+++ b/mjx/mujoco/mjx/integration_test/smooth_test.py
@@ -68,10 +68,27 @@ class TransmissionIntegrationTest(parameterized.TestCase):
     mujoco.mj_transmission(m, d)
     dx = transmission_jit_fn(mx, dx)
 
-    for field in ['actuator_length', 'actuator_moment']:
-      _assert_attr_eq(
-          d, dx, field, seed, f'transmission{seed}', atol=1e-4
-      )
+    _assert_attr_eq(
+        d, dx, 'actuator_length', seed, f'transmission{seed}', atol=1e-4
+    )
+
+    # convert sparse actuator_moment to dense representation
+    moment = np.zeros((m.nu, m.nv))
+    mujoco.mju_sparse2dense(
+        moment,
+        d.actuator_moment.reshape(-1),
+        d.moment_rownnz,
+        d.moment_rowadr,
+        d.moment_colind.reshape(-1),
+    )
+    _assert_eq(
+        moment,
+        dx.actuator_moment,
+        'actuator_moment',
+        seed,
+        f'transmission{seed}',
+        atol=1e-4,
+    )
 
 
 if __name__ == '__main__':
diff --git a/python/LQR.ipynb b/python/LQR.ipynb
index b683444a..adee43ed 100644
--- a/python/LQR.ipynb
+++ b/python/LQR.ipynb
@@ -491,7 +491,15 @@
    },
    "outputs": [],
    "source": [
-    "ctrl0 = np.atleast_2d(qfrc0) @ np.linalg.pinv(data.actuator_moment)\n",
+    "actuator_moment = np.zeros((model.nu, model.nv))\n",
+    "mujoco.mju_sparse2dense(\n",
+    "    actuator_moment,\n",
+    "    data.actuator_moment,\n",
+    "    data.moment_rownnz,\n",
+    "    data.moment_rowadr,\n",
+    "    data.moment_colind,\n",
+    ")\n",
+    "ctrl0 = np.atleast_2d(qfrc0) @ np.linalg.pinv(actuator_moment)\n",
     "ctrl0 = ctrl0.flatten()  # Save the ctrl setpoint.\n",
     "print('control setpoint:', ctrl0)"
    ]
diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index 1bdcfda1..e9bf33cc 100644
--- a/src/engine/engine_core_smooth.c
+++ b/src/engine/engine_core_smooth.c
@@ -857,6 +857,9 @@ void mj_transmission(const mjModel* m, mjData* d) {
   // outputs
   mjtNum* length = d->actuator_length;
   mjtNum* moment = d->actuator_moment;
+  int *rownnz = d->moment_rownnz;
+  int *rowadr = d->moment_rowadr;
+  int *colind = d->moment_colind;
 
   // allocate Jacbians
   mj_markStack(d);
@@ -875,6 +878,10 @@ void mj_transmission(const mjModel* m, mjData* d) {
 
   // compute lengths and moments
   for (int i=0; i < nu; i++) {
+    rownnz[i] = 0;
+    rowadr[i] = i == 0 ? 0 : rowadr[i-1] + rownnz[i-1];
+    int adr = rowadr[i];
+
     // extract info
     int id = m->actuator_trnid[2*i];
     mjtNum* gear = m->actuator_gear+6*i;
@@ -885,18 +892,19 @@ void mj_transmission(const mjModel* m, mjData* d) {
     case mjTRN_JOINTINPARENT:           // joint, force in parent frame
       // slide and hinge joint: scalar gear
       if (m->jnt_type[id] == mjJNT_SLIDE || m->jnt_type[id] == mjJNT_HINGE) {
+        // sparsity
+        rownnz[i]++;
+        colind[adr] = m->jnt_dofadr[id];
+
         length[i] = d->qpos[m->jnt_qposadr[id]]*gear[0];
-        moment[i*nv + m->jnt_dofadr[id]] = gear[0];
+        moment[adr] = gear[0];
       }
 
       // ball joint: 3D wrench gear
       else if (m->jnt_type[id] == mjJNT_BALL) {
-        // j: qpos start address
-        int j = m->jnt_qposadr[id];
-
         // axis: expmap representation of quaternion
         mjtNum axis[3], quat[4];
-        mju_copy4(quat, d->qpos+j);
+        mju_copy4(quat, d->qpos+m->jnt_qposadr[id]);
         mju_normalize4(quat);
         mju_quat2Vel(axis, quat, 1);
 
@@ -912,11 +920,17 @@ void mj_transmission(const mjModel* m, mjData* d) {
         // length: axis*gearAxis
         length[i] = mju_dot3(axis, gearAxis);
 
-        // j: dof start address
-        j = m->jnt_dofadr[id];
+        // dof start address
+        int jnt_dofadr = m->jnt_dofadr[id];
+
+        // sparsity
+        for (int j = 0; j < 3; j++) {
+          colind[adr+j] = jnt_dofadr + j;
+        }
+        rownnz[i] += 3;
 
         // moment: gearAxis
-        mju_copy3(moment+i*nv+j, gearAxis);
+        mju_copy3(moment+adr, gearAxis);
       }
 
       // free joint: 6D wrench gear
@@ -924,35 +938,30 @@ void mj_transmission(const mjModel* m, mjData* d) {
         // cannot compute meaningful length, set to 0
         length[i] = 0;
 
-        // j: qpos start address
-        int j = m->jnt_qposadr[id];
-
-        // vec: translational components
-        mjtNum vec[3];
-        mju_copy3(vec, d->qpos+j);
-
-        // axis: expmap representation of quaternion
-        mjtNum axis[3], quat[4];
-        mju_quat2Vel(axis, d->qpos+j+3, 1);
-        mju_copy4(quat, d->qpos+j+3);
-        mju_normalize4(quat);
-        mju_quat2Vel(axis, quat, 1);
-
         // gearAxis: rotate to world frame if necessary
         mjtNum gearAxis[3];
         if (m->actuator_trntype[i] == mjTRN_JOINT) {
           mju_copy3(gearAxis, gear+3);
         } else {
+          mjtNum quat[4];
+          mju_copy4(quat, d->qpos+m->jnt_qposadr[id]+3);
+          mju_normalize4(quat);
           mju_negQuat(quat, quat);
           mju_rotVecQuat(gearAxis, gear+3, quat);
         }
 
-        // j: dof start address
-        j = m->jnt_dofadr[id];
+        // dof start address
+        int jnt_dofadr = m->jnt_dofadr[id];
+
+        // sparsity
+        for (int j = 0; j < 6; j++) {
+          colind[adr+j] = jnt_dofadr + j;
+        }
+        rownnz[i] += 6;
 
         // moment: gear(tran), gearAxis
-        mju_copy3(moment+i*nv+j, gear);
-        mju_copy3(moment+i*nv+j+3, gearAxis);
+        mju_copy3(moment+adr, gear);
+        mju_copy3(moment+adr+3, gearAxis);
       }
       break;
 
@@ -1000,20 +1009,26 @@ void mj_transmission(const mjModel* m, mjData* d) {
         mj_jacSite(m, d, jac, 0, id);
         mju_subFrom(jac, jacS, 3*nv);
 
+        // sparsity
+        for (int j = 0; j < nv; j++) {
+          colind[adr+j] = j;
+        }
+        rownnz[i] += nv;
+
         // clear moment
-        mju_zero(moment+i*nv, nv);
+        mju_zero(moment + adr, nv);
 
         // apply chain rule
         for (int j=0; j < nv; j++) {
           for (int k=0; k < 3; k++) {
-            moment[i*nv+j] += dlda[k]*jacA[k*nv+j] + dldv[k]*jac[k*nv+j];
+            moment[adr+j] += dlda[k]*jacA[k*nv+j] + dldv[k]*jac[k*nv+j];
           }
         }
 
         // scale by gear ratio
         length[i] *= gear[0];
         for (int j = 0; j < nv; j++) {
-          moment[i*nv + j] *= gear[0];
+          moment[adr+j] *= gear[0];
         }
       }
       break;
@@ -1022,20 +1037,32 @@ void mj_transmission(const mjModel* m, mjData* d) {
       length[i] = d->ten_length[id]*gear[0];
 
       // moment: sparse or dense
-      if (mj_isSparse(m)) {
-        // clear moment
-        mju_zero(moment+i*nv, nv);
+      if (issparse) {
+        // sparsity
+        int ten_J_rownnz = d->ten_J_rownnz[id];
+        int ten_J_rowadr = d->ten_J_rowadr[id];
+        rownnz[i] += ten_J_rownnz;
+        mju_copyInt(colind + adr, d->ten_J_colind + ten_J_rowadr, ten_J_rownnz);
 
-        int end = d->ten_J_rowadr[id] + d->ten_J_rownnz[id];
-        for (int j=d->ten_J_rowadr[id]; j < end; j++) {
-          moment[i*nv + d->ten_J_colind[j]] = d->ten_J[j] * gear[0];
-        }
+        mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz);
       } else {
-        mju_scl(moment + i*nv, d->ten_J + id*nv, gear[0], nv);
+        // sparsity
+        for (int j = 0; j < nv; j++) {
+          colind[adr+j] = j;
+        }
+        rownnz[i] += nv;
+
+        mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv);
       }
       break;
 
     case mjTRN_SITE:                    // site
+      // sparsity
+      for (int j = 0; j < nv; j++) {
+        colind[adr+j] = j;
+      }
+      rownnz[i] += nv;
+
       // get site translation (jac) and rotation (jacS) Jacobians in global frame
       mj_jacSite(m, d, jac, jacS, id);
 
@@ -1050,9 +1077,9 @@ void mj_transmission(const mjModel* m, mjData* d) {
         mju_mulMatVec3(wrench+3, d->site_xmat+9*id, gear+3);  // rotation
 
         // moment: global Jacobian projected on wrench
-        mju_mulMatTVec(moment+i*nv, jac, wrench, 3, nv);     // translation
-        mju_mulMatTVec(jac, jacS, wrench+3, 3, nv);          // rotation
-        mju_addTo(moment+i*nv, jac, nv);                     // add the two
+        mju_mulMatTVec(moment+adr, jac, wrench, 3, nv);       // translation
+        mju_mulMatTVec(jac, jacS, wrench+3, 3, nv);           // rotation
+        mju_addTo(moment+adr, jac, nv);                       // add the two
       }
 
       // reference site defined
@@ -1089,7 +1116,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
         }
 
         // clear moment
-        mju_zero(moment+i*nv, nv);
+        mju_zero(moment+adr, nv);
 
         // translational transmission
         if (!mju_isZero(gear, 3)) {
@@ -1121,7 +1148,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
           mju_mulMatVec3(wrench, d->site_xmat+9*refid, gear);
 
           // moment: global Jacobian projected on wrench
-          mju_mulMatTVec(moment+i*nv, jac, wrench, 3, nv);
+          mju_mulMatTVec(moment+adr, jac, wrench, 3, nv);
         }
 
         // rotational transmission
@@ -1162,18 +1189,24 @@ void mj_transmission(const mjModel* m, mjData* d) {
           // moment_tmp: global Jacobian projected on wrench, add to moment
           if (!moment_tmp) moment_tmp = mj_stackAllocNum(d, nv);
           mju_mulMatTVec(moment_tmp, jacS, wrench, 3, nv);
-          mju_addTo(moment+i*nv, moment_tmp, nv);
+          mju_addTo(moment+adr, moment_tmp, nv);
         }
       }
 
       break;
 
     case mjTRN_BODY:                  // body (adhesive contacts)
+      // sparsity
+      for (int j = 0; j < nv; j++) {
+        colind[adr+j] = j;
+      }
+      rownnz[i] += nv;
+
       // cannot compute meaningful length, set to 0
       length[i] = 0;
 
       // clear moment
-      mju_zero(moment+i*nv, nv);
+      mju_zero(moment+adr, nv);
 
       // moment is average of all contact normal Jacobians
       {
@@ -1257,15 +1290,16 @@ void mj_transmission(const mjModel* m, mjData* d) {
         // moment is average over contact normal Jacobians, make negative for adhesion
         if (counter) {
           // accumulate active contact Jacobians into moment
-          mj_mulJacTVec(m, d, moment+i*nv, efc_force);
+          mj_mulJacTVec(m, d, moment+adr, efc_force);
 
           // add Jacobians from excluded contacts
-          mju_addTo(moment+i*nv, moment_exclude, nv);
+          mju_addTo(moment+adr, moment_exclude, nv);
 
           // normalize by total contacts, flip sign
-          mju_scl(moment+i*nv, moment+i*nv, -1.0/counter, nv);
+          mju_scl(moment+adr, moment+adr, -1.0/counter, nv);
         }
       }
+
       break;
 
     default:
diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c
index 3f7eb63f..15d48149 100644
--- a/src/engine/engine_derivative.c
+++ b/src/engine/engine_derivative.c
@@ -827,6 +827,10 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
     return;
   }
 
+  // allocate dense actuator_moment row
+  mj_markStack(d);
+  mjtNum* moment = mj_stackAllocNum(d, nv);
+
   // process actuators
   for (int i=0; i < nu; i++) {
     // skip if disabled
@@ -870,9 +874,14 @@ void mjd_actuator_vel(const mjModel* m, mjData* d) {
 
     // add
     if (bias_vel != 0) {
-      addJTBJ(m, d, d->actuator_moment+i*nv, &bias_vel, 1);
+      mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + i,
+                       d->moment_rowadr + i, d->moment_colind);
+      addJTBJ(m, d, moment, &bias_vel, 1);
     }
   }
+
+  // free space
+  mj_freeStack(d);
 }
 
 
diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c
index ae4b7c2d..e77639e5 100644
--- a/src/engine/engine_forward.c
+++ b/src/engine/engine_forward.c
@@ -208,8 +208,11 @@ void mj_fwdVelocity(const mjModel* m, mjData* d) {
     mju_mulMatVec(d->ten_velocity, d->ten_J, d->qvel, m->ntendon, m->nv);
   }
 
-  // actuator velocity: always dense
-  mju_mulMatVec(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu, m->nv);
+  // actuator velocity: always sparse
+  if (!mjDISABLED(mjDSBL_ACTUATION)) {
+    mju_mulMatVecSparse(d->actuator_velocity, d->actuator_moment, d->qvel, m->nu,
+                        d->moment_rownnz, d->moment_rowadr, d->moment_colind, NULL);
+  }
 
   // com-based velocities, passive forces, constraint references
   mj_comVel(m, d);
@@ -270,7 +273,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
   TM_START;
   int nv = m->nv, nu = m->nu;
   mjtNum gain, bias, tau;
-  mjtNum *prm, *moment = d->actuator_moment, *force = d->actuator_force;
+  mjtNum *prm, *force = d->actuator_force;
 
   // clear actuator_force
   mju_zero(force, nu);
@@ -475,7 +478,8 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
   clampVec(force, m->actuator_forcerange, m->actuator_forcelimited, nu, NULL);
 
   // qfrc_actuator = moment' * force
-  mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv);
+  mju_mulMatTVecSparse(d->qfrc_actuator, d->actuator_moment, force, nu, nv,
+                       d->moment_rownnz, d->moment_rowadr, d->moment_colind);
 
   // actuator-level gravity compensation
   if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c
index 6891da68..57684692 100644
--- a/src/engine/engine_io.c
+++ b/src/engine/engine_io.c
@@ -1857,9 +1857,6 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
   mju_zero(d->mocap_pos, 3*m->nmocap);
   mju_zero(d->mocap_quat, 4*m->nmocap);
 
-  // zero out actuator_moment, mj_transmission touches it selectively
-  mju_zero(d->actuator_moment, m->nv*m->nu);
-
   // copy qpos0 from model
   if (m->qpos0) {
     memcpy(d->qpos, m->qpos0, m->nq*sizeof(mjtNum));
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index 626ced94..96b41546 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -93,7 +93,7 @@ static void printSparse(const char* str, const mjtNum* mat, int nr,
                         const int* rownnz, const int* rowadr,
                         const int* colind, FILE* fp, const char* float_format) {
   // if no data, or too many rows to be visually useful, return
-  if (!mat || nr > 300) {
+  if (!mat || !nr || nr > 300) {
     return;
   }
   fprintf(fp, "%s\n", str);
@@ -147,7 +147,7 @@ static void printSparsity(const char* str, int nr, int nc,
 // print vector
 static void printVector(const char* str, const mjtNum* data, int n, FILE* fp,
                         const char* float_format) {
-  if (!data) {
+  if (!data || !n) {
     return;
   }
   // print str
@@ -1005,7 +1005,10 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   }
 
   printArray("ACTUATOR_LENGTH", m->nu, 1, d->actuator_length, fp, float_format);
-  printArray("ACTUATOR_MOMENT", m->nu, m->nv, d->actuator_moment, fp, float_format);
+  printSparsity("actuator_moments", m->nu, m->nv,
+                d->moment_rowadr, d->moment_rownnz, d->moment_colind, fp);
+  printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz,
+              d->moment_rowadr, d->moment_colind, fp, float_format);
   printArray("CRB", m->nbody, 10, d->crb, fp, float_format);
 
   if (M) {
diff --git a/src/engine/engine_setconst.c b/src/engine/engine_setconst.c
index f451353b..c1926673 100644
--- a/src/engine/engine_setconst.c
+++ b/src/engine/engine_setconst.c
@@ -29,6 +29,7 @@
 #include "engine/engine_util_blas.h"
 #include "engine/engine_util_errmem.h"
 #include "engine/engine_util_misc.h"
+#include "engine/engine_util_sparse.h"
 #include "engine/engine_util_spatial.h"
 
 
@@ -66,6 +67,7 @@ static void set0(mjModel* m, mjData* d) {
   mj_markStack(d);
   mjtNum* jac = mj_stackAllocNum(d, 6*nv);
   mjtNum* tmp = mj_stackAllocNum(d, 6*nv);
+  mjtNum* moment = mj_stackAllocNum(d, nv);
   int* cammode = 0;
   int* lightmode = 0;
 
@@ -278,7 +280,9 @@ static void set0(mjModel* m, mjData* d) {
 
     // compute actuator_acc0
     for (int i=0; i < m->nu; i++) {
-      mj_solveM(m, d, tmp, d->actuator_moment+i*nv, 1);
+      mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + i,
+                       d->moment_rowadr + i, d->moment_colind);
+      mj_solveM(m, d, tmp, moment, 1);
       m->actuator_acc0[i] = mju_norm(tmp, nv);
     }
   } else {
@@ -395,13 +399,16 @@ static void set0(mjModel* m, mjData* d) {
     // === interpret biasprm[2] > 0 as dampratio for position-like actuators
 
     // "reflected" inertia (inversely scaled by transmission squared)
-    mjtNum* transmission = d->actuator_moment + i*nv;
+    int rownnz = d->moment_rownnz[i];
+    int rowadr = d->moment_rowadr[i];
+    mjtNum* transmission = d->actuator_moment + rowadr;
     mjtNum mass = 0;
-    for (int j=0; j < nv; j++) {
+    for (int j=0; j < rownnz; j++) {
       mjtNum trn = mju_abs(transmission[j]);
       mjtNum trn2 = trn*trn;  // transmission squared
       if (trn2 > mjMINVAL) {
-        mass += m->dof_M0[j] / trn2;
+        int dof = d->moment_colind[rowadr + j];
+        mass += m->dof_M0[dof] / trn2;
       }
     }
 
@@ -598,11 +605,16 @@ static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
   // step1: compute inertia and actuator moments
   mj_step1(m, d);
 
+  // dense actuator_moment row
+  mj_markStack(d);
+  mjtNum* moment = mj_stackAllocNum(d, nv);
+  mju_sparse2dense(moment, d->actuator_moment, 1, nv, d->moment_rownnz + index,
+                   d->moment_rowadr + index, d->moment_colind);
+
   // set force to generate desired acceleration
-  mj_solveM(m, d, d->qfrc_applied, d->actuator_moment+index*nv, 1);
+  mj_solveM(m, d, d->qfrc_applied, moment, 1);
   mjtNum nrm = mju_norm(d->qfrc_applied, nv);
-  mju_scl(d->qfrc_applied, d->actuator_moment+index*nv,
-          (2*side-1)*opt->accel/mjMAX(mjMINVAL, nrm), nv);
+  mju_scl(d->qfrc_applied, moment, (2*side-1)*opt->accel/mjMAX(mjMINVAL, nrm), nv);
 
   // impose maxforce
   nrm = mju_norm(d->qfrc_applied, nv);
@@ -613,6 +625,8 @@ static mjtNum evalAct(const mjModel* m, mjData* d, int index, int side,
   // step2: apply force
   mj_step2(m, d);
 
+  mj_freeStack(d);
+
   // return actuator length
   return d->actuator_length[index];
 }
diff --git a/src/engine/engine_util_sparse.h b/src/engine/engine_util_sparse.h
index fa2647c3..d1dc2d5d 100644
--- a/src/engine/engine_util_sparse.h
+++ b/src/engine/engine_util_sparse.h
@@ -39,8 +39,8 @@ MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
                            int* rownnz, int* rowadr, int* colind, int nnz);
 
 // convert matrix from sparse to dense
-MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
-                            const int* rownnz, const int* rowadr, const int* colind);
+MJAPI void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc, const int* rownnz,
+                            const int* rowadr, const int* colind);
 
 // multiply sparse matrix and dense vector:  res = mat * vec
 MJAPI void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
diff --git a/test/engine/engine_derivative_test.cc b/test/engine/engine_derivative_test.cc
index 4984a311..51dc26b7 100644
--- a/test/engine/engine_derivative_test.cc
+++ b/test/engine/engine_derivative_test.cc
@@ -31,6 +31,7 @@
 #include "src/engine/engine_io.h"
 #include "src/engine/engine_util_blas.h"
 #include "src/engine/engine_util_errmem.h"
+#include "src/engine/engine_util_sparse.h"
 #include "test/fixture.h"
 
 namespace mujoco {
@@ -475,7 +476,8 @@ static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
   if (B) {
     mjtNum *Bc = mj_stackAllocNum(d, nu*nv);
     mjtNum *BcT = mj_stackAllocNum(d, nv*nu);
-    mju_copy(Bc, d->actuator_moment, nv*nu);
+    mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
+                     d->moment_rowadr, d->moment_colind);
     mj_solveLD(m, Bc, nu, d->qH, d->qHDiagInv);
     mju_transpose(BcT, Bc, nu, nv);
     mju_scl(B, BcT, dt*dt, nu*nv);