Make JacobianType.auto switch to sparse in put_model (#234)

adenzler-nvidia · web-flow · commit dc2d4d9e1ae8 · 2025-05-12T09:22:37.000+02:00
* make jac_auto switch to sparse

* remove old test

* review feedback
diff --git a/mujoco_warp/_src/io.py b/mujoco_warp/_src/io.py
@@ -64,7 +64,7 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
 
   # TODO(team): remove after solver._update_gradient for Newton solver utilizes tile operations for islands
   nv_max = 60
-  if mjm.nv > nv_max and (not mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_SPARSE):
+  if mjm.nv > nv_max and mjm.opt.jacobian == mujoco.mjtJacobian.mjJAC_DENSE:
     raise ValueError(f"Dense is unsupported for nv > {nv_max} (nv = {mjm.nv}).")
 
   # calculate some fields that cannot be easily computed inline
diff --git a/mujoco_warp/_src/io_test.py b/mujoco_warp/_src/io_test.py
@@ -89,11 +89,6 @@ def test_geom_type(self):
     with self.assertRaises(NotImplementedError):
       mjwarp.put_model(mjm)
 
-  def test_dense(self):
-    with self.assertRaises(ValueError):
-      # dense not supported yet for large nv
-      test_util.fixture("humanoid/n_humanoids.xml")
-
   def test_actuator_trntype(self):
     mjm = mujoco.MjModel.from_xml_string("""
       <mujoco>
@@ -243,6 +238,22 @@ def test_option_physical_constants(self):
     with self.assertRaises(NotImplementedError):
       mjwarp.put_model(mjm)
 
+  def test_jacobian_auto(self):
+    mjm = mujoco.MjModel.from_xml_string("""
+      <mujoco>
+        <option jacobian="auto"/>
+        <worldbody>
+          <replicate count="11">
+          <body>          
+            <geom type="sphere" size=".1"/>
+            <freejoint/>
+            </body>
+          </replicate>
+        </worldbody> 
+      </mujoco>
+    """)
+    mjwarp.put_model(mjm)
+
 
 if __name__ == "__main__":
   wp.init()
