Sparse ten_J and ten_J_colind memory

PiperOrigin-RevId: 868834376
Change-Id: Iec92d9e09e3134666895e54a78cb02439d74a4de

mjx/mujoco/mjx/_src/io.py

@@ -768,8 +768,8 @@ def _make_data_c(
       'flexedge_length': (nflexedge, float_),
       'ten_J_rownnz': (m.ntendon, np.int32),
       'ten_J_rowadr': (m.ntendon, np.int32),
-      'ten_J_colind': (m.nJten, np.int32),
-      'ten_J': (m.nJten, float_),
+      'ten_J_colind': (m.ntendon, m.nv, np.int32),
+      'ten_J': (m.ntendon, m.nv, float_),
       'ten_wrapadr': (m.ntendon, np.int32),
       'ten_wrapnum': (m.ntendon, np.int32),
       'wrap_obj': (m.nwrap, 2, np.int32),
@@ -1099,7 +1099,7 @@ def _put_data_jax(
   # MJX does not support islanding, so only transfer the first solver_niter
   impl_fields['solver_niter'] = impl_fields['solver_niter'][0]
 
-  # convert sparse actuator_moment to dense matrix
+  # convert sparse representation of actuator_moment to dense matrix
   moment = np.zeros((m.nu, m.nv))
   mujoco.mju_sparse2dense(
       moment,
@@ -1110,22 +1110,6 @@ def _put_data_jax(
   )
   impl_fields['actuator_moment'] = moment
 
-  # convert ten_J to dense matrix
-  if mujoco.mj_isSparse(m):
-    ten_J = np.zeros((m.ntendon, m.nv))
-    mujoco.mju_sparse2dense(
-        ten_J,
-        d.ten_J,
-        d.ten_J_rownnz,
-        d.ten_J_rowadr,
-        d.ten_J_colind,
-    )
-  elif m.ntendon:
-    ten_J = d.ten_J.reshape((m.ntendon, m.nv))
-  else:
-    ten_J = np.zeros((m.ntendon, m.nv))
-  impl_fields['ten_J'] = ten_J
-
   contact, contact_map = _put_contact(d.contact, dim, efc_address)
 
   # pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
@@ -1584,29 +1568,6 @@ def _get_data_into(
         result_i.actuator_moment[:] = actuator_moment
         continue
 
-      # MuJoCo ten_J is sparse, MJX uses a dense representation.
-      if field.name == 'ten_J':
-        ten_j_rownnz = np.zeros(m.ntendon, dtype=np.int32)
-        ten_j_rowadr = np.zeros(m.ntendon, dtype=np.int32)
-        ten_j_colind = np.zeros(m.nJten, dtype=np.int32)
-        ten_j = np.zeros(m.nJten)
-        if m.ntendon:
-          if d_i.impl == types.Impl.JAX:
-            mujoco.mju_dense2sparse(
-                ten_j,
-                d_i._impl.ten_J,
-                ten_j_rownnz,
-                ten_j_rowadr,
-                ten_j_colind,
-            )
-          else:
-            ten_j = d_i._impl.ten_J
-        result_i.ten_J_rownnz[:] = ten_j_rownnz
-        result_i.ten_J_rowadr[:] = ten_j_rowadr
-        result_i.ten_J_colind[:] = ten_j_colind
-        result_i.ten_J[:] = ten_j
-        continue
-
       if hasattr(d_i._impl, field.name):
         value = getattr(d_i._impl, field.name)
       else:

mjx/mujoco/mjx/_src/smooth_test.py

@@ -121,9 +121,7 @@ class SmoothTest(absltest.TestCase):
     mujoco.mj_forward(m, d)
     # tendon
     dx = jax.jit(mjx.tendon)(mx, mjx.put_data(m, d))
-    # with dense jacobian mode, d.ten_J is already dense (ntendon*nv,), just reshape
-    ten_J = d.ten_J.reshape((m.ntendon, m.nv))
-    _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
+    _assert_attr_eq(d, dx._impl, 'ten_J')
     _assert_attr_eq(d, dx, 'ten_length')
     # transmission
     dx = jax.jit(mjx.transmission)(mx, dx)
@@ -397,20 +395,7 @@ class TendonTest(parameterized.TestCase):
     dx = jax.jit(mjx.forward)(mx, dx)
 
     _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
-    # convert ten_J for comparison based on jacobian mode
-    if mujoco.mj_isSparse(m):
-      ten_J = np.zeros((m.ntendon, m.nv))
-      mujoco.mju_sparse2dense(
-          ten_J,
-          d.ten_J,
-          d.ten_J_rownnz,
-          d.ten_J_rowadr,
-          d.ten_J_colind,
-      )
-    else:
-      # dense mode: just reshape
-      ten_J = d.ten_J.reshape((m.ntendon, m.nv))
-    _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
+    _assert_eq(d.ten_J, dx._impl.ten_J, 'ten_J')
     _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')
     _assert_eq(d.ten_wrapadr, dx._impl.ten_wrapadr, 'ten_wrapadr')
     _assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')

mjx/mujoco/mjx/_src/types.py

@@ -722,7 +722,6 @@ class Model(PyTreeNode):
     nC: number of non-zeros in C matrix
     nD: number of non-zeros in D matrix
     nJmom: number of non-zeros in Jacobian momentum matrix
-    nJten: number of non-zeros in sparse tendon Jacobian
     ngravcomp: number of bodies with nonzero gravcomp
     nuserdata: number of elements in userdata
     nsensordata: number of elements in sensor data vector
@@ -773,7 +772,6 @@ class Model(PyTreeNode):
   nC: int  # pylint:disable=invalid-name
   nD: int  # pylint:disable=invalid-name
   nJmom: int  # pylint:disable=invalid-name
-  nJten: int  # pylint:disable=invalid-name
   ngravcomp: int
   nuserdata: int
   nsensordata: int

mjx/mujoco/mjx/third_party/mujoco_warp/_src/io.py

@@ -1169,18 +1169,7 @@ def get_data_into(
 
   # tendon
   result.ten_length[:] = d.ten_length.numpy()[world_id]
-  # TODO(team): remove after mjwarp depends on mujoco > 3.4.0 in pyproject.toml
-  if BLEEDING_EDGE_MUJOCO:
-    ten_J = d.ten_J.numpy()[world_id]
-    mujoco.mju_dense2sparse(
-        result.ten_J,
-        ten_J,
-        result.ten_J_rownnz,
-        result.ten_J_rowadr,
-        result.ten_J_colind,
-    )
-  else:
-    result.ten_J[:] = d.ten_J.numpy()[world_id]
+  result.ten_J[:] = d.ten_J.numpy()[world_id]
   result.ten_wrapadr[:] = d.ten_wrapadr.numpy()[world_id]
   result.ten_wrapnum[:] = d.ten_wrapnum.numpy()[world_id]
   result.wrap_obj[:] = d.wrap_obj.numpy()[world_id]

mjx/mujoco/mjx/warp/forward_test.py

@@ -149,8 +149,7 @@ class ForwardTest(parameterized.TestCase):
         tu.assert_attr_eq(dx, d, 'cam_xpos')
         tu.assert_eq(dx.cam_xmat, d.cam_xmat.reshape((-1, 3, 3)), 'cam_xmat')
       tu.assert_attr_eq(dx, d, 'ten_length')
-      ten_J = d.ten_J.reshape((m.ntendon, m.nv))
-      tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
+      tu.assert_attr_eq(dx._impl, d, 'ten_J')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapnum')
       tu.assert_attr_eq(dx._impl, d, 'wrap_xpos')