warp-lang 1.5.0__py3-none-manylinux2014_x86_64.whl → 1.6.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_sim_grad_bounce_linear.py

@@ -0,0 +1,204 @@
+# Copyright (c) 2025 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import numpy as np
+
+import warp as wp
+import warp.optim
+import warp.sim
+import warp.sim.render
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def update_trajectory_kernel(
+    trajectory: wp.array(dtype=wp.vec3),
+    q: wp.array(dtype=wp.transform),
+    time_step: wp.int32,
+    q_idx: wp.int32,
+):
+    trajectory[time_step] = wp.transform_get_translation(q[q_idx])
+
+
+@wp.kernel
+def trajectory_loss_kernel(
+    trajectory: wp.array(dtype=wp.vec3f),
+    target_trajectory: wp.array(dtype=wp.vec3f),
+    loss: wp.array(dtype=wp.float32),
+):
+    tid = wp.tid()
+    diff = trajectory[tid] - target_trajectory[tid]
+    distance_loss = wp.dot(diff, diff)
+    wp.atomic_add(loss, 0, distance_loss)
+
+
+class BallBounceLinearTest:
+    def __init__(self, gravity=True, rendering=False):
+        # Ball bouncing scenario inspired by https://github.com/NVIDIA/warp/issues/349
+        self.fps = 30
+        self.num_frames = 60
+        self.sim_substeps = 20  # XXX need to use enough substeps to achieve smooth gradients
+        self.frame_dt = 1.0 / self.fps
+        self.sim_dt = self.frame_dt / self.sim_substeps
+        self.sim_duration = self.num_frames * self.frame_dt
+        self.sim_steps = int(self.sim_duration // self.sim_dt)
+
+        self.target_force_linear = 100.0
+
+        if gravity:
+            builder = wp.sim.ModelBuilder(up_vector=wp.vec3(0, 0, 1))
+        else:
+            builder = wp.sim.ModelBuilder(gravity=0.0, up_vector=wp.vec3(0, 0, 1))
+
+        b = builder.add_body(origin=wp.transform((0.5, 0.0, 1.0), wp.quat_identity()), name="ball")
+        builder.add_shape_sphere(
+            body=b, radius=0.1, density=100.0, ke=2000.0, kd=10.0, kf=200.0, mu=0.2, thickness=0.01
+        )
+        builder.set_ground_plane(ke=10, kd=10, kf=0.0, mu=0.2)
+        self.model = builder.finalize(requires_grad=True)
+
+        self.time = np.linspace(0, self.sim_duration, self.sim_steps)
+
+        self.integrator = wp.sim.SemiImplicitIntegrator()
+        if rendering:
+            self.renderer = wp.sim.render.SimRendererOpenGL(self.model, "ball_bounce_linear")
+        else:
+            self.renderer = None
+
+        self.loss = wp.array([0], dtype=wp.float32, requires_grad=True)
+        self.states = [self.model.state() for _ in range(self.sim_steps + 1)]
+        self.target_states = [self.model.state() for _ in range(self.sim_steps + 1)]
+
+        self.target_force = wp.array([0.0, 0.0, 0.0, 0.0, self.target_force_linear, 0.0], dtype=wp.spatial_vectorf)
+
+        self.trajectory = wp.empty(len(self.time), dtype=wp.vec3, requires_grad=True)
+        self.target_trajectory = wp.empty(len(self.time), dtype=wp.vec3)
+
+    def _reset(self):
+        self.loss = wp.array([0], dtype=wp.float32, requires_grad=True)
+
+    def generate_target_trajectory(self):
+        for i in range(self.sim_steps):
+            curr_state = self.target_states[i]
+            next_state = self.target_states[i + 1]
+            curr_state.clear_forces()
+            if i == 0:
+                wp.copy(curr_state.body_f, self.target_force, dest_offset=0, src_offset=0, count=1)
+            wp.sim.collide(self.model, curr_state)
+            self.integrator.simulate(self.model, curr_state, next_state, self.sim_dt)
+            wp.launch(kernel=update_trajectory_kernel, dim=1, inputs=[self.target_trajectory, curr_state.body_q, i, 0])
+
+    def forward(self, force: wp.array):
+        for i in range(self.sim_steps):
+            curr_state = self.states[i]
+            next_state = self.states[i + 1]
+            curr_state.clear_forces()
+            if i == 0:
+                wp.copy(curr_state.body_f, force, dest_offset=0, src_offset=0, count=1)
+            wp.sim.collide(self.model, curr_state)
+            self.integrator.simulate(self.model, curr_state, next_state, self.sim_dt)
+            wp.launch(kernel=update_trajectory_kernel, dim=1, inputs=[self.trajectory, curr_state.body_q, i, 0])
+
+            if self.renderer:
+                self.renderer.begin_frame(self.time[i])
+                self.renderer.render(curr_state)
+                self.renderer.end_frame()
+
+    def step(self, force: wp.array):
+        self.tape = wp.Tape()
+        self._reset()
+        with self.tape:
+            self.forward(force)
+            wp.launch(
+                kernel=trajectory_loss_kernel,
+                dim=len(self.trajectory),
+                inputs=[self.trajectory, self.target_trajectory, self.loss],
+            )
+        self.tape.backward(self.loss)
+        force_grad = force.grad.numpy()[0, 4]
+        self.tape.zero()
+
+        return self.loss.numpy()[0], force_grad
+
+    def evaluate(self, num_samples, plot_results=False):
+        forces = np.linspace(0, self.target_force_linear * 2, num_samples)
+        losses = np.zeros_like(forces)
+        grads = np.zeros_like(forces)
+
+        for i, fx in enumerate(forces):
+            force = wp.array([[0.0, 0.0, 0.0, 0.0, fx, 0.0]], dtype=wp.spatial_vectorf, requires_grad=True)
+            losses[i], grads[i] = self.step(force)
+            if plot_results:
+                print(f"Iteration {i + 1}/{num_samples}")
+                print(f"Force: {fx:.2f}, Loss: {losses[i]:.6f}, Grad: {grads[i]:.6f}")
+
+            assert np.isfinite(losses[i])
+            assert np.isfinite(grads[i])
+            if i > 0:
+                assert grads[i] >= grads[i - 1]
+
+        if plot_results:
+            import matplotlib.pyplot as plt
+
+            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
+
+            # Plot the loss curve
+            ax1.plot(forces, losses, label="Loss")
+            ax1.set_xlabel("Force")
+            ax1.set_ylabel("Loss")
+            ax1.set_title("Loss vs Force")
+            ax1.legend()
+
+            # Make sure the grads are not too large
+            grads = np.clip(grads, -1e4, 1e4)
+
+            # Plot the gradient curve
+            ax2.plot(forces, grads, label="Gradient", color="orange")
+            ax2.set_xlabel("Force")
+            ax2.set_ylabel("Gradient")
+            ax2.set_title("Gradient vs Force")
+            ax2.legend()
+
+            plt.suptitle("Loss and Gradient vs Force")
+            plt.tight_layout(rect=[0, 0, 1, 0.95])
+            plt.show()
+
+        return losses, grads
+
+
+def test_sim_grad_bounce_linear(test, device):
+    with wp.ScopedDevice(device):
+        model = BallBounceLinearTest()
+        model.generate_target_trajectory()
+
+        num_samples = 20
+        losses, grads = model.evaluate(num_samples=num_samples)
+        # gradients must approximate linear behavior with zero crossing in the middle
+        test.assertTrue(np.abs(grads[1:] - grads[:-1]).max() < 1.1)
+        test.assertTrue(np.all(grads[: num_samples // 2] <= 0.0))
+        test.assertTrue(np.all(grads[num_samples // 2 :] >= 0.0))
+        # losses must follow a parabolic behavior
+        test.assertTrue(np.allclose(losses[: num_samples // 2], losses[num_samples // 2 :][::-1], atol=1.0))
+        diffs = losses[1:] - losses[:-1]
+        test.assertTrue(np.all(diffs[: num_samples // 2 - 1] <= 0.0))
+        test.assertTrue(np.all(diffs[num_samples // 2 - 1 :] >= 0.0))
+        # second derivative must be constant positive
+        diffs2 = diffs[1:] - diffs[:-1]
+        test.assertTrue(np.allclose(diffs2, diffs2[0], atol=1e-2))
+        test.assertTrue(np.all(diffs2 >= 0.0))
+
+
+class TestSimGradBounceLinear(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices("basic")
+add_function_test(TestSimGradBounceLinear, "test_sim_grad_bounce_linear", test_sim_grad_bounce_linear, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_smoothstep.py

@@ -31,15 +31,7 @@ class TestData:
 
 TEST_DATA = {
     wp.float32: (
-        TestData(
-            a=1.0,
-            b=2.0,
-            t=1.5,
-            expected=0.5,
-            expected_adj_a=-0.75,
-            expected_adj_b=-0.75,
-            expected_adj_t=1.5,
-        ),
+        TestData(a=1.0, b=2.0, t=1.5, expected=0.5, expected_adj_a=-0.75, expected_adj_b=-0.75, expected_adj_t=1.5),
         TestData(
             a=-1.0,
             b=2.0,
@@ -49,24 +41,8 @@ TEST_DATA = {
             expected_adj_b=-0.09375,
             expected_adj_t=0.375,
         ),
-        TestData(
-            a=0.0,
-            b=1.0,
-            t=9.9,
-            expected=1.0,
-            expected_adj_a=0.0,
-            expected_adj_b=0.0,
-            expected_adj_t=0.0,
-        ),
-        TestData(
-            a=0.0,
-            b=1.0,
-            t=-9.9,
-            expected=0.0,
-            expected_adj_a=0.0,
-            expected_adj_b=0.0,
-            expected_adj_t=0.0,
-        ),
+        TestData(a=0.0, b=1.0, t=9.9, expected=1.0, expected_adj_a=0.0, expected_adj_b=0.0, expected_adj_t=0.0),
+        TestData(a=0.0, b=1.0, t=-9.9, expected=0.0, expected_adj_a=0.0, expected_adj_b=0.0, expected_adj_t=0.0),
     ),
 }
 
@@ -83,72 +59,30 @@ def test_smoothstep(test, device):
 
         return fn
 
-    for data_type in TEST_DATA:
+    for data_type, test_data_set in TEST_DATA.items():
         kernel_fn = make_kernel_fn(data_type)
         kernel = wp.Kernel(
             func=kernel_fn,
             key=f"test_smoothstep{data_type.__name__}_kernel",
         )
 
-        for test_data in TEST_DATA[data_type]:
-            a = wp.array(
-                [test_data.a],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            b = wp.array(
-                [test_data.b],
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-            t = wp.array(
-                [test_data.t],
-                dtype=float,
-                device=device,
-                requires_grad=True,
-            )
-            out = wp.array(
-                [0] * wp.types.type_length(data_type),
-                dtype=data_type,
-                device=device,
-                requires_grad=True,
-            )
-
-            tape = wp.Tape()
-            with tape:
-                wp.launch(
-                    kernel,
-                    dim=1,
-                    inputs=[a, b, t, out],
-                    device=device,
-                )
-
-            assert_np_equal(
-                out.numpy(),
-                np.array([test_data.expected]),
-                tol=1e-6,
-            )
+        for test_data in test_data_set:
+            a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
+            b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
+            t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
+            out = wp.array([0] * wp.types.type_length(data_type), dtype=data_type, device=device, requires_grad=True)
+
+            with wp.Tape() as tape:
+                wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)
+
+            assert_np_equal(out.numpy(), np.array([test_data.expected]), tol=1e-6)
 
             if test_data.check_backwards():
                 tape.backward(out)
 
-                assert_np_equal(
-                    tape.gradients[a].numpy(),
-                    np.array([test_data.expected_adj_a]),
-                    tol=1e-6,
-                )
-                assert_np_equal(
-                    tape.gradients[b].numpy(),
-                    np.array([test_data.expected_adj_b]),
-                    tol=1e-6,
-                )
-                assert_np_equal(
-                    tape.gradients[t].numpy(),
-                    np.array([test_data.expected_adj_t]),
-                    tol=1e-6,
-                )
+                assert_np_equal(tape.gradients[a].numpy(), np.array([test_data.expected_adj_a]), tol=1e-6)
+                assert_np_equal(tape.gradients[b].numpy(), np.array([test_data.expected_adj_b]), tol=1e-6)
+                assert_np_equal(tape.gradients[t].numpy(), np.array([test_data.expected_adj_t]), tol=1e-6)
 
 
 devices = get_test_devices()

warp/tests/test_static.py

@@ -281,20 +281,20 @@ def test_function_lookup(test, device):
 
     outputs = wp.empty(2, dtype=wp.float32)
 
-    for op in op_handlers.keys():
+    for _op, op_func in op_handlers.items():
 
         @wp.kernel
         def operate(input: wp.array(dtype=inputs.dtype, ndim=2), output: wp.array(dtype=wp.float32)):
             tid = wp.tid()
             a, b = input[tid, 0], input[tid, 1]
             # retrieve the right function to use for the captured dtype variable
-            output[tid] = wp.static(op_handlers[op])(a, b)  # noqa: B023
+            output[tid] = wp.static(op_func)(a, b)  # noqa: B023
 
         wp.launch(operate, dim=2, inputs=[inputs], outputs=[outputs])
         outputs_np = outputs.numpy()
         inputs_np = inputs.numpy()
         for i in range(len(outputs_np)):
-            test.assertEqual(outputs_np[i], op_handlers[op](float(inputs_np[i][0]), float(inputs_np[i][1])))
+            test.assertEqual(outputs_np[i], op_func(float(inputs_np[i][0]), float(inputs_np[i][1])))
 
 
 def count_ssa_occurrences(kernel: wp.Kernel, ssas: List[str]) -> Dict[str, int]:
@@ -536,6 +536,21 @@ def test_static_function_hash(test, _):
     test.assertEqual(hash1, hash3)
 
 
+@wp.kernel
+def static_len_query_kernel(v1: wp.vec2):
+    v2 = wp.vec3()
+    m = wp.identity(n=wp.static(len(v1) + len(v2) + 1), dtype=float)
+    wp.expect_eq(wp.ddot(m, m), 6.0)
+
+    t = wp.transform_identity(float)
+    wp.expect_eq(wp.static(len(t)), 7)
+
+
+def test_static_len_query(test, _):
+    v1 = wp.vec2()
+    wp.launch(static_len_query_kernel, 1, inputs=(v1,))
+
+
 devices = get_test_devices()
 
 
@@ -561,6 +576,7 @@ add_function_test(TestStatic, "test_static_if_else_elif", test_static_if_else_el
 
 add_function_test(TestStatic, "test_static_constant_hash", test_static_constant_hash, devices=None)
 add_function_test(TestStatic, "test_static_function_hash", test_static_function_hash, devices=None)
+add_function_test(TestStatic, "test_static_len_query", test_static_len_query, devices=None)
 
 
 if __name__ == "__main__":

warp/tests/test_tape.py

@@ -125,6 +125,30 @@ def test_tape_dot_product(test, device):
     assert_np_equal(tape.gradients[y].numpy(), x.numpy())
 
 
+@wp.kernel
+def assign_chain_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float)):
+    tid = wp.tid()
+    y[tid] = x[tid]
+    z[tid] = y[tid]
+
+
+def test_tape_zero_multiple_outputs(test, device):
+    x = wp.array(np.arange(3), dtype=float, device=device, requires_grad=True)
+    y = wp.zeros_like(x)
+    z = wp.zeros_like(x)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(assign_chain_kernel, dim=3, inputs=[x, y, z], device=device)
+
+    tape.backward(grads={y: wp.ones_like(x)})
+    assert_np_equal(x.grad.numpy(), np.ones(3, dtype=float))
+    tape.zero()
+
+    tape.backward(grads={z: wp.ones_like(x)})
+    assert_np_equal(x.grad.numpy(), np.ones(3, dtype=float))
+
+
 def test_tape_visualize(test, device):
     dim = 8
     tape = wp.Tape()
@@ -163,6 +187,7 @@ class TestTape(unittest.TestCase):
 add_function_test(TestTape, "test_tape_mul_constant", test_tape_mul_constant, devices=devices)
 add_function_test(TestTape, "test_tape_mul_variable", test_tape_mul_variable, devices=devices)
 add_function_test(TestTape, "test_tape_dot_product", test_tape_dot_product, devices=devices)
+add_function_test(TestTape, "test_tape_zero_multiple_outputs", test_tape_zero_multiple_outputs, devices=devices)
 add_function_test(TestTape, "test_tape_visualize", test_tape_visualize, devices=devices)