Move tendon sparsity fields from mjData to mjModel

PiperOrigin-RevId: 875087590
Change-Id: I1a5489d2d2011795ee38b09d547e68acd72e3cc7

mjx/mujoco/mjx/_src/io.py

@@ -1129,9 +1129,9 @@ def _put_data_jax(
     mujoco.mju_sparse2dense(
         ten_J,
         d.ten_J,
-        d.ten_J_rownnz,
-        d.ten_J_rowadr,
-        d.ten_J_colind,
+        m.ten_J_rownnz,
+        m.ten_J_rowadr,
+        m.ten_J_colind,
     )
   else:
     ten_J = np.zeros((m.ntendon, m.nv))
@@ -1251,6 +1251,9 @@ def _put_data_c(
       for f in types.DataC.fields()
       if hasattr(d, f.name)
   }
+  for f in types.DataC.fields():
+    if not hasattr(d, f.name) and hasattr(m, f.name):
+      impl_fields[f.name] = getattr(m, f.name)
 
   # TODO(stunya): support islanding via C impl.
   impl_fields['solver_niter'] = impl_fields['solver_niter'][0]
@@ -1662,9 +1665,6 @@ def _get_data_into(
             )
           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
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -125,9 +125,9 @@ class SmoothTest(absltest.TestCase):
     mujoco.mju_sparse2dense(
         ten_J,
         d.ten_J,
-        d.ten_J_rownnz,
-        d.ten_J_rowadr,
-        d.ten_J_colind,
+        m.ten_J_rownnz,
+        m.ten_J_rowadr,
+        m.ten_J_colind,
     )
     _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
     _assert_attr_eq(d, dx, 'ten_length')
@@ -407,9 +407,9 @@ class TendonTest(parameterized.TestCase):
     mujoco.mju_sparse2dense(
         ten_J,
         d.ten_J,
-        d.ten_J_rownnz,
-        d.ten_J_rowadr,
-        d.ten_J_colind,
+        m.ten_J_rownnz,
+        m.ten_J_rowadr,
+        m.ten_J_colind,
     )
     _assert_eq(ten_J, dx._impl.ten_J, 'ten_J')
     _assert_eq(d.ten_wrapnum, dx._impl.ten_wrapnum, 'ten_wrapnum')

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

@@ -952,9 +952,12 @@ def put_data(
 
   d.flexedge_J = wp.array(np.tile(mjd.flexedge_J.reshape(-1), (nworld, 1)).reshape((nworld, 1, -1)), dtype=float)
 
-  ten_J = np.zeros((mjm.ntendon, mjm.nv))
-  mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjd.ten_J_rownnz, mjd.ten_J_rowadr, mjd.ten_J_colind.reshape(-1))
-  d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
+  if mjm.ntendon:
+    ten_J = np.zeros((mjm.ntendon, mjm.nv))
+    mujoco.mju_sparse2dense(ten_J, mjd.ten_J.reshape(-1), mjm.ten_J_rownnz, mjm.ten_J_rowadr, mjm.ten_J_colind.reshape(-1))
+    d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), ten_J), dtype=float)
+  else:
+    d.ten_J = wp.array(np.full((nworld, mjm.ntendon, mjm.nv), 0.0), dtype=float)
 
   # TODO(taylorhowell): sparse actuator_moment
   actuator_moment = np.zeros((mjm.nu, mjm.nv))
@@ -1165,9 +1168,9 @@ def get_data_into(
     mujoco.mju_dense2sparse(
         result.ten_J,
         ten_J,
-        result.ten_J_rownnz,
-        result.ten_J_rowadr,
-        result.ten_J_colind,
+        mjm.ten_J_rownnz,
+        mjm.ten_J_rowadr,
+        mjm.ten_J_colind,
     )
   else:
     result.ten_J[:] = d.ten_J.numpy()[world_id]

mjx/mujoco/mjx/warp/forward_test.py

@@ -153,9 +153,9 @@ class ForwardTest(parameterized.TestCase):
       mujoco.mju_sparse2dense(
           ten_J,
           d.ten_J,
-          d.ten_J_rownnz,
-          d.ten_J_rowadr,
-          d.ten_J_colind,
+          m.ten_J_rownnz,
+          m.ten_J_rowadr,
+          m.ten_J_colind,
       )
       tu.assert_eq(dx._impl.ten_J, ten_J, 'ten_J')
       tu.assert_attr_eq(dx._impl, d, 'ten_wrapadr')