Optimize test suite runtime (~18% faster) (#1689)

Author: adenzler-nvidia

newton/tests/test_collision_pipeline.py

@@ -276,7 +276,7 @@ def test_collision_pipeline(
         shape_type_b=shape_type_b,
         broad_phase=broad_phase,
     )
-    for _ in range(200):
+    for _ in range(100):
         setup.step()
         setup.render()
     setup.test(test_level_a, 0, tolerance=tolerance)
@@ -386,7 +386,7 @@ def test_mesh_mesh_sdf_modes(
         sdf_max_resolution_a=sdf_max_resolution_a,
         sdf_max_resolution_b=sdf_max_resolution_b,
     )
-    for _ in range(200):
+    for _ in range(100):
         setup.step()
         setup.render()
     setup.test(TestLevel.VELOCITY_YZ, 0, tolerance=tolerance)
@@ -552,10 +552,12 @@ def _contact_pairs(broad_phase):
             contacts = pipeline.contacts()
             pipeline.collide(state, contacts)
             n = contacts.rigid_contact_count.numpy()[0]
+            shape0_np = contacts.rigid_contact_shape0.numpy()
+            shape1_np = contacts.rigid_contact_shape1.numpy()
             pairs = set()
             for i in range(n):
-                s0 = int(contacts.rigid_contact_shape0.numpy()[i])
-                s1 = int(contacts.rigid_contact_shape1.numpy()[i])
+                s0 = int(shape0_np[i])
+                s1 = int(shape1_np[i])
                 pairs.add((min(s0, s1), max(s0, s1)))
             return pairs
 

newton/tests/test_equality_constraints.py

@@ -71,7 +71,7 @@ def test_multiple_constraints(self):
         self.state_0, self.state_1 = self.model.state(), self.model.state()
         newton.eval_fk(self.model, self.model.joint_q, self.model.joint_qd, self.state_0)
 
-        for _ in range(1000):
+        for _ in range(200):
             for _ in range(10):
                 self.state_0.clear_forces()
                 self.solver.step(self.state_0, self.state_1, self.control, None, self.sim_dt)

newton/tests/test_fixed_tendon.py

@@ -16,6 +16,8 @@
 import unittest
 from enum import IntEnum
 
+import warp as wp
+
 import newton
 from newton.solvers import SolverMuJoCo
 
@@ -99,7 +101,25 @@ def test_single_mujoco_fixed_tendon_length_behaviour(self):
         joint_start_positions = [0.5, 0.0, 0.5, 0.0]
         state_in.joint_q.assign(joint_start_positions)
 
-        for _i in range(0, 200):
+        device = model.device
+        use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device)
+        if use_cuda_graph:
+            # warmup (2 steps for full ping-pong cycle)
+            solver.step(state_in=state_in, state_out=state_out, contacts=contacts, control=control, dt=dt)
+            solver.step(state_in=state_out, state_out=state_in, contacts=contacts, control=control, dt=dt)
+            with wp.ScopedCapture(device) as capture:
+                solver.step(state_in=state_in, state_out=state_out, contacts=contacts, control=control, dt=dt)
+                solver.step(state_in=state_out, state_out=state_in, contacts=contacts, control=control, dt=dt)
+            graph = capture.graph
+
+        remaining = 200 - (4 if use_cuda_graph else 0)
+        for _i in range(remaining // 2 if use_cuda_graph else remaining):
+            if use_cuda_graph:
+                wp.capture_launch(graph)
+            else:
+                solver.step(state_in=state_in, state_out=state_out, contacts=contacts, control=control, dt=dt)
+                state_in, state_out = state_out, state_in
+        if use_cuda_graph and remaining % 2 == 1:
             solver.step(state_in=state_in, state_out=state_out, contacts=contacts, control=control, dt=dt)
             state_in, state_out = state_out, state_in
 

newton/tests/test_heightfield.py

@@ -270,7 +270,25 @@ def test_heightfield_collision(self):
         control = model.control()
         sim_dt = 1.0 / 240.0
 
-        for _ in range(500):
+        device = model.device
+        use_cuda_graph = device.is_cuda and wp.is_mempool_enabled(device)
+        if use_cuda_graph:
+            # warmup (2 steps for full ping-pong cycle)
+            solver.step(state_in, state_out, control, None, sim_dt)
+            solver.step(state_out, state_in, control, None, sim_dt)
+            with wp.ScopedCapture(device) as capture:
+                solver.step(state_in, state_out, control, None, sim_dt)
+                solver.step(state_out, state_in, control, None, sim_dt)
+            graph = capture.graph
+
+        remaining = 500 - (4 if use_cuda_graph else 0)
+        for _ in range(remaining // 2 if use_cuda_graph else remaining):
+            if use_cuda_graph:
+                wp.capture_launch(graph)
+            else:
+                solver.step(state_in, state_out, control, None, sim_dt)
+                state_in, state_out = state_out, state_in
+        if use_cuda_graph and remaining % 2 == 1:
             solver.step(state_in, state_out, control, None, sim_dt)
             state_in, state_out = state_out, state_in
 

newton/tests/test_hydroelastic.py

@@ -367,19 +367,19 @@ def test_mujoco_hydroelastic_penetration_depth(test, device):
     substeps = 10
     sim_time = 3.0
     num_frames = int(sim_time / sim_dt)
+    total_steps = num_frames * substeps
 
-    for _ in range(num_frames):
-        for _ in range(substeps):
-            state_0.clear_forces()
-            # Apply external force to upper boxes
-            forces = np.zeros(model.body_count * 6, dtype=np.float32)
-            for body_idx in upper_body_indices:
-                forces[body_idx * 6 + 2] = -external_force
-            state_0.body_f.assign(forces)
-
-            collision_pipeline.collide(state_0, contacts)
-            solver.step(state_0, state_1, control, contacts, sim_dt / substeps)
-            state_0, state_1 = state_1, state_0
+    # Pre-compute forces as a Warp array
+    forces_np = np.zeros(model.body_count * 6, dtype=np.float32)
+    for body_idx in upper_body_indices:
+        forces_np[body_idx * 6 + 2] = -external_force
+    precomputed_forces = wp.array(forces_np.reshape(model.body_count, 6), dtype=wp.spatial_vector, device=device)
+
+    for _ in range(total_steps):
+        wp.copy(state_0.body_f, precomputed_forces)
+        collision_pipeline.collide(state_0, contacts)
+        solver.step(state_0, state_1, control, contacts, sim_dt / substeps)
+        state_0, state_1 = state_1, state_0
 
     # Check that upper cubes are near their original positions
     body_q = state_0.body_q.numpy()

newton/tests/test_import_mjcf.py

@@ -29,7 +29,7 @@
 from newton.solvers import SolverMuJoCo
 
 
-class TestImportMjcf(unittest.TestCase):
+class TestImportMjcfBasic(unittest.TestCase):
     def test_humanoid_mjcf(self):
         builder = newton.ModelBuilder()
         builder.default_shape_cfg.ke = 123.0
@@ -610,6 +610,8 @@ def test_replace_3d_hinge_with_ball_joint(self):
         # note we need to swap quaternion order wxyz -> xyzw
         np.testing.assert_allclose(joint_x_p.q, [0, 0, 0.7071068, 0.7071068], atol=1e-6)
 
+
+class TestImportMjcfGeometry(unittest.TestCase):
     def test_cylinder_shapes_preserved(self):
         """Test that cylinder geometries are properly imported as cylinders, not capsules."""
         # Create MJCF content with cylinder geometry
@@ -1829,6 +1831,8 @@ def test_single_mujoco_fixed_tendon_auto_springlength(self):
             msg=f"Expected tendon_length0: {expected_tendon_length0}, Measured: {measured_tendon_length0}",
         )
 
+
+class TestImportMjcfSolverParams(unittest.TestCase):
     def test_solimplimit_parsing(self):
         """Test that solimplimit attribute is parsed correctly from MJCF."""
         mjcf = """<?xml version="1.0" ?>
@@ -2902,6 +2906,8 @@ def test_default_inheritance(self):
         else:
             self.fail("Model should have mujoco.condim attribute")
 
+
+class TestImportMjcfActuatorsFrames(unittest.TestCase):
     def test_actuatorfrcrange_parsing(self):
         """Test that actuatorfrcrange is parsed from MJCF joint attributes and applied to joint effort limits."""
         mjcf_content = """<?xml version="1.0" encoding="utf-8"?>
@@ -3741,6 +3747,8 @@ def test_joint_anchor_with_rotated_child_body(self):
         # In xyzw format: [0, 0, sin(45°), cos(45°)] = [0, 0, 0.7071, 0.7071]
         np.testing.assert_allclose(joint_X_p[3:7], [0, 0, 0.7071068, 0.7071068], atol=1e-5)
 
+
+class TestImportMjcfComposition(unittest.TestCase):
     def test_floating_true_creates_free_joint(self):
         """Test that floating=True creates a free joint for the root body."""
         mjcf_content = """<?xml version="1.0" encoding="utf-8"?>