warp-lang 1.5.1__py3-none-win_amd64.whl → 1.6.0__py3-none-win_amd64.whl

warp/tape.py

@@ -43,7 +43,6 @@ class Tape:
 
     def __init__(self):
         self.gradients = {}
-        self.const_gradients = set()
         self.launches = []
         self.scopes = []
 
@@ -106,7 +105,6 @@ class Tape:
                 else:
                     # ensure we can capture this backward pass in a CUDA graph
                     a.grad.assign(g)
-                self.const_gradients.add(a)
 
         # run launches backwards
         for launch in reversed(self.launches):
@@ -115,13 +113,13 @@ class Tape:
 
             else:
                 # kernel option takes precedence over module option
-                kernel_enable_backward = launch[0].options.get("enable_backward")
-                if kernel_enable_backward is False:
+                enable_backward = launch[0].options.get("enable_backward")
+                if enable_backward is False:
                     msg = f"Running the tape backwards may produce incorrect gradients because recorded kernel {launch[0].key} is configured with the option 'enable_backward=False'."
                     wp.utils.warn(msg)
-                elif kernel_enable_backward is None:
-                    module_enable_backward = launch[0].module.options.get("enable_backward")
-                    if module_enable_backward is False:
+                elif enable_backward is None:
+                    enable_backward = launch[0].module.options.get("enable_backward")
+                    if enable_backward is False:
                         msg = f"Running the tape backwards may produce incorrect gradients because recorded kernel {launch[0].key} is defined in a module with the option 'enable_backward=False' set."
                         wp.utils.warn(msg)
 
@@ -144,18 +142,19 @@ class Tape:
                 for a in outputs:
                     adj_outputs.append(self.get_adjoint(a))
 
-                wp.launch(
-                    kernel=kernel,
-                    dim=dim,
-                    inputs=inputs,
-                    outputs=outputs,
-                    adj_inputs=adj_inputs,
-                    adj_outputs=adj_outputs,
-                    device=device,
-                    adjoint=True,
-                    max_blocks=max_blocks,
-                    block_dim=block_dim,
-                )
+                if enable_backward:
+                    wp.launch(
+                        kernel=kernel,
+                        dim=dim,
+                        inputs=inputs,
+                        outputs=outputs,
+                        adj_inputs=adj_inputs,
+                        adj_outputs=adj_outputs,
+                        device=device,
+                        adjoint=True,
+                        max_blocks=max_blocks,
+                        block_dim=block_dim,
+                    )
 
     # record a kernel launch on the tape
     def record_launch(self, kernel, dim, max_blocks, inputs, outputs, device, block_dim=0, metadata=None):
@@ -222,9 +221,9 @@ class Tape:
     # returns the adjoint of a kernel parameter
     def get_adjoint(self, a):
         if not wp.types.is_array(a) and not isinstance(a, wp.codegen.StructInstance):
-            # if input is a simple type (e.g.: float, vec3, etc) then
-            # no gradient needed (we only return gradients through arrays and structs)
-            return a
+            # if input is a simple type (e.g.: float, vec3, etc) or a non-Warp array,
+            # then no gradient needed (we only return gradients through Warp arrays and structs)
+            return None
 
         elif wp.types.is_array(a) and a.grad:
             # keep track of all gradients used by the tape (for zeroing)
@@ -267,13 +266,12 @@ class Tape:
         Zero out all gradients recorded on the tape.
         """
         for a, g in self.gradients.items():
-            if a not in self.const_gradients:
-                if isinstance(a, wp.codegen.StructInstance):
-                    for name in g._cls.vars:
-                        if isinstance(g._cls.vars[name].type, wp.array) and g._cls.vars[name].requires_grad:
-                            getattr(g, name).zero_()
-                else:
-                    g.zero_()
+            if isinstance(a, wp.codegen.StructInstance):
+                for name in g._cls.vars:
+                    if isinstance(g._cls.vars[name].type, wp.array) and g._cls.vars[name].requires_grad:
+                        getattr(g, name).zero_()
+            else:
+                g.zero_()
 
     def _reset_array_read_flags(self):
         """
@@ -517,7 +515,7 @@ class GraphvizTapeVisitor(TapeVisitor):
         node_attrs = f"label=<{label}>"
         if "caller" in launch_data:
             caller = launch_data["caller"]
-            node_attrs += f",tooltip=\"{self.sanitize(caller['file'])}:{caller['lineno']} ({caller['func']})\""
+            node_attrs += f',tooltip="{self.sanitize(caller["file"])}:{caller["lineno"]} ({caller["func"]})"'
 
         self.graphviz_lines.append(f"{chart_indent}{kernel_launch_id} [{node_attrs}];")
 

warp/tests/aux_test_module_unload.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 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.
+
+"""Dummy module used in test_reload.py"""
+
+import warp as wp
+
+
+@wp.kernel
+def k():
+    pass

warp/tests/{test_sim_grad.py → flaky_test_sim_grad.py}

@@ -11,6 +11,7 @@ import numpy as np
 
 import warp as wp
 import warp.sim
+import warp.sim.render
 from warp.tests.unittest_utils import *
 
 
@@ -23,8 +24,7 @@ def evaluate_loss(
     loss: wp.array(dtype=float),
 ):
     tid = wp.tid()
-    # wp.atomic_add(loss, 0, weighting * (target - joint_q[tid * 2 + 1]) ** 2.0)
-    d = wp.abs(target - joint_q[tid * 2 + 1])
+    d = (target - joint_q[tid * 2 + 1]) ** 2.0
     wp.atomic_add(loss, 0, weighting * d)
 
 
@@ -34,7 +34,13 @@ def assign_action(action: wp.array(dtype=float), joint_act: wp.array(dtype=float
     joint_act[2 * tid] = action[tid]
 
 
-def gradcheck(func, inputs, device, eps=1e-1, tol=1e-2):
+@wp.kernel
+def assign_force(action: wp.array(dtype=float), body_f: wp.array(dtype=wp.spatial_vector)):
+    tid = wp.tid()
+    body_f[2 * tid] = wp.spatial_vector(0.0, 0.0, 0.0, action[tid], 0.0, 0.0)
+
+
+def gradcheck(func, inputs, device, eps=1e-1, tol=1e-2, print_grad=False):
     """
     Checks that the gradient of the Warp kernel is correct by comparing it to the
     numerical gradient computed using finite differences.
@@ -46,56 +52,64 @@ def gradcheck(func, inputs, device, eps=1e-1, tol=1e-2):
         output = func(*wp_xs)
         return output.numpy()[0]
 
+    # compute analytical gradient
+    tape = wp.Tape()
+    with tape:
+        output = func(*inputs)
+
+    tape.backward(loss=output)
+
     # compute numerical gradient
-    numerical_grad = []
     np_xs = []
     for i in range(len(inputs)):
         np_xs.append(inputs[i].numpy().flatten().copy())
-        numerical_grad.append(np.zeros_like(np_xs[-1]))
-        inputs[i].requires_grad = True
 
-    for i in range(len(np_xs)):
+    for i in range(len(inputs)):
+        fd_grad = np.zeros_like(np_xs[i])
         for j in range(len(np_xs[i])):
             np_xs[i][j] += eps
             y1 = f(np_xs)
             np_xs[i][j] -= 2 * eps
             y2 = f(np_xs)
             np_xs[i][j] += eps
-            numerical_grad[i][j] = (y1 - y2) / (2 * eps)
-
-    # compute analytical gradient
-    tape = wp.Tape()
-    with tape:
-        output = func(*inputs)
-
-    tape.backward(loss=output)
-
-    # compare gradients
-    for i in range(len(inputs)):
-        grad = tape.gradients[inputs[i]]
-        assert_np_equal(grad.numpy(), numerical_grad[i], tol=tol)
+            fd_grad[j] = (y1 - y2) / (2 * eps)
+
+        # compare gradients
+        ad_grad = tape.gradients[inputs[i]].numpy()
+        if print_grad:
+            print("grad ad:", ad_grad)
+            print("grad fd:", fd_grad)
+        assert_np_equal(ad_grad, fd_grad, tol=tol)
         # ensure the signs match
-        assert np.allclose(grad.numpy() * numerical_grad[i] > 0, True)
+        assert np.allclose(ad_grad * fd_grad > 0, True)
 
     tape.zero()
 
 
-def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
-    # Two boxes on a rail (prismatic or D6 joint), one is pushed, the other is passive.
+def test_sphere_pushing_on_rails(
+    test,
+    device,
+    joint_type,
+    integrator_type,
+    apply_force=False,
+    static_contacts=True,
+    print_grad=False,
+):
+    # Two spheres on a rail (prismatic or D6 joint), one is pushed, the other is passive.
     # The absolute distance to a target is measured and gradients are compared for
     # a push that is too far and too close.
     num_envs = 2
-    num_steps = 200
-    sim_substeps = 2
+    num_steps = 150
+    sim_substeps = 10
     dt = 1 / 30
 
-    target = 5.0
+    target = 3.0
 
     if integrator_type == 0:
-        contact_ke = 1e5
-        contact_kd = 1e3
+        contact_ke = 1e3
+        contact_kd = 1e1
     else:
-        contact_ke = 1e5
+        contact_ke = 1e3
         contact_kd = 1e1
 
     complete_builder = wp.sim.ModelBuilder()
@@ -104,16 +118,16 @@ def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
     complete_builder.default_shape_kd = contact_kd
 
     for _ in range(num_envs):
-        builder = wp.sim.ModelBuilder()
+        builder = wp.sim.ModelBuilder(gravity=0.0)
 
         builder.default_shape_ke = complete_builder.default_shape_ke
         builder.default_shape_kd = complete_builder.default_shape_kd
 
         b0 = builder.add_body(name="pusher")
-        builder.add_shape_box(b0, density=1000.0)
+        builder.add_shape_sphere(b0, radius=0.4, density=100.0)
 
         b1 = builder.add_body(name="passive")
-        builder.add_shape_box(b1, hx=0.4, hy=0.4, hz=0.4, density=1000.0)
+        builder.add_shape_sphere(b1, radius=0.47, density=100.0)
 
         if joint_type == 0:
             builder.add_joint_prismatic(-1, b0)
@@ -122,7 +136,7 @@ def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
             builder.add_joint_d6(-1, b0, linear_axes=[wp.sim.JointAxis((1.0, 0.0, 0.0))])
             builder.add_joint_d6(-1, b1, linear_axes=[wp.sim.JointAxis((1.0, 0.0, 0.0))])
 
-        builder.joint_q[-2:] = [0.0, 1.0]
+        builder.joint_q[-2:] = [0.0, 2.0]
         complete_builder.add_builder(builder)
 
     assert complete_builder.body_count == 2 * num_envs
@@ -135,6 +149,15 @@ def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
     model.joint_attach_ke = 32000.0 * 16
     model.joint_attach_kd = 500.0 * 4
 
+    model.shape_geo.scale.requires_grad = False
+    model.shape_geo.thickness.requires_grad = False
+
+    if static_contacts:
+        wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, model)
+        model.rigid_contact_margin = 10.0
+        state = model.state()
+        wp.sim.collide(model, state)
+
     if integrator_type == 0:
         integrator = wp.sim.FeatherstoneIntegrator(model, update_mass_matrix_every=num_steps * sim_substeps)
     elif integrator_type == 1:
@@ -143,40 +166,57 @@ def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
     else:
         integrator = wp.sim.XPBDIntegrator(iterations=2, rigid_contact_relaxation=1.0)
 
-    # renderer = wp.sim.render.SimRenderer(model, "test_sim_grad.usd", scaling=1.0)
+    # renderer = wp.sim.render.SimRendererOpenGL(model, "test_sim_grad.usd", scaling=1.0)
     renderer = None
     render_time = 0.0
 
+    if renderer:
+        renderer.render_sphere("target", pos=wp.vec3(target, 0, 0), rot=wp.quat_identity(), radius=0.1, color=(1, 0, 0))
+
     def rollout(action: wp.array) -> wp.array:
         nonlocal render_time
         states = [model.state() for _ in range(num_steps * sim_substeps + 1)]
 
-        if not isinstance(integrator, wp.sim.FeatherstoneIntegrator):
-            # apply initial generalized coordinates
-            wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, states[0])
+        wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, states[0])
 
         control_active = model.control()
         control_nop = model.control()
 
-        wp.launch(
-            assign_action,
-            dim=num_envs,
-            inputs=[action],
-            outputs=[control_active.joint_act],
-            device=model.device,
-        )
+        if not apply_force:
+            wp.launch(
+                assign_action,
+                dim=num_envs,
+                inputs=[action],
+                outputs=[control_active.joint_act],
+                device=model.device,
+            )
 
         i = 0
         for step in range(num_steps):
-            wp.sim.collide(model, states[i])
-            control = control_active if step < 10 else control_nop
+            state = states[i]
+            if not static_contacts:
+                wp.sim.collide(model, state)
+            if apply_force:
+                control = control_nop
+            else:
+                control = control_active if step < 10 else control_nop
             if renderer:
                 renderer.begin_frame(render_time)
-                renderer.render(states[i])
+                renderer.render(state)
                 renderer.end_frame()
                 render_time += dt
             for _ in range(sim_substeps):
-                integrator.simulate(model, states[i], states[i + 1], dt / sim_substeps, control)
+                state = states[i]
+                next_state = states[i + 1]
+                if apply_force and step < 10:
+                    wp.launch(
+                        assign_force,
+                        dim=num_envs,
+                        inputs=[action],
+                        outputs=[state.body_f],
+                        device=model.device,
+                    )
+                integrator.simulate(model, state, next_state, dt / sim_substeps, control)
                 i += 1
 
         if not isinstance(integrator, wp.sim.FeatherstoneIntegrator):
@@ -184,39 +224,40 @@ def test_box_pushing_on_rails(test, device, joint_type, integrator_type):
             wp.sim.eval_ik(model, states[-1], states[-1].joint_q, states[-1].joint_qd)
 
         loss = wp.zeros(1, requires_grad=True, device=device)
+        weighting = 1.0
         wp.launch(
             evaluate_loss,
             dim=num_envs,
-            inputs=[states[-1].joint_q, 1.0, target],
+            inputs=[states[-1].joint_q, weighting, target],
             outputs=[loss],
             device=model.device,
         )
 
-        if renderer:
-            renderer.save()
+        # if renderer:
+        #     renderer.save()
 
         return loss
 
     action_too_far = wp.array(
-        [5000.0 for _ in range(num_envs)],
+        [80.0 for _ in range(num_envs)],
         device=device,
         dtype=wp.float32,
         requires_grad=True,
     )
-    tol = 1e-2
+    tol = 2e-1
     if isinstance(integrator, wp.sim.XPBDIntegrator):
         # Euler, XPBD do not yield as accurate gradients, but at least the
         # signs should match
         tol = 0.1
-    gradcheck(rollout, [action_too_far], device=device, eps=0.2, tol=tol)
+    gradcheck(rollout, [action_too_far], device=device, eps=0.2, tol=tol, print_grad=print_grad)
 
     action_too_close = wp.array(
-        [1500.0 for _ in range(num_envs)],
+        [40.0 for _ in range(num_envs)],
         device=device,
         dtype=wp.float32,
         requires_grad=True,
     )
-    gradcheck(rollout, [action_too_close], device=device, eps=0.2, tol=tol)
+    gradcheck(rollout, [action_too_close], device=device, eps=0.2, tol=tol, print_grad=print_grad)
 
 
 devices = get_test_devices()
@@ -226,15 +267,15 @@ class TestSimGradients(unittest.TestCase):
     pass
 
 
-for int_type, int_name in enumerate(["featherstone", "semiimplicit"]):
-    for jt_type, jt_name in enumerate(["prismatic", "d6"]):
-        test_name = f"test_box_pushing_on_rails_{int_name}_{jt_name}"
+for jt_type, jt_name in enumerate(["prismatic", "d6"]):
+    test_name = f"test_sphere_pushing_on_rails_{jt_name}"
 
-        def test_fn(self, device, jt_type=jt_type, int_type=int_type):
-            return test_box_pushing_on_rails(self, device, jt_type, int_type)
-
-        add_function_test(TestSimGradients, test_name, test_fn, devices=devices)
+    def test_fn(self, device, jt_type=jt_type, int_type=1):
+        return test_sphere_pushing_on_rails(
+            self, device, jt_type, int_type, apply_force=True, static_contacts=True, print_grad=False
+        )
 
+    add_function_test(TestSimGradients, test_name, test_fn, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_array.py

@@ -2766,6 +2766,76 @@ def test_indexing_types(test, device):
     )
 
 
+def test_alloc_strides(test, device):
+    def test_transposed(shape, dtype):
+        # allocate without specifying strides
+        a1 = wp.zeros(shape, dtype=dtype)
+
+        # allocate with contiguous strides
+        strides = wp.types.strides_from_shape(shape, dtype)
+        a2 = wp.zeros(shape, dtype=dtype, strides=strides)
+
+        # allocate with transposed (reversed) shape/strides
+        rshape = shape[::-1]
+        rstrides = strides[::-1]
+        a3 = wp.zeros(rshape, dtype=dtype, strides=rstrides)
+
+        # ensure that correct capacity was allocated
+        assert a2.capacity == a1.capacity
+        assert a3.capacity == a1.capacity
+
+    with wp.ScopedDevice(device):
+        shapes = [(5, 5), (5, 3), (3, 5), (2, 3, 4), (4, 2, 3), (3, 2, 4)]
+        for shape in shapes:
+            with test.subTest(msg=f"shape={shape}"):
+                test_transposed(shape, wp.int8)
+                test_transposed(shape, wp.float32)
+                test_transposed(shape, wp.vec3)
+
+
+def test_casting(test, device):
+    idxs = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)
+    idxs = wp.array(idxs, device=device).reshape((-1, 3))
+    idxs = wp.array(idxs, shape=idxs.shape[0], dtype=wp.vec3i, device=device)
+    assert idxs.dtype is wp.vec3i
+    assert idxs.shape == (4,)
+    assert idxs.strides == (12,)
+
+
+@wp.kernel
+def array_len_kernel(
+    a1: wp.array(dtype=int),
+    a2: wp.array(dtype=float, ndim=3),
+    out: wp.array(dtype=int),
+):
+    length = len(a1)
+    wp.expect_eq(len(a1), 123)
+    out[0] = len(a1)
+
+    length = len(a2)
+    wp.expect_eq(len(a2), 2)
+    out[1] = len(a2)
+
+
+def test_array_len(test, device):
+    a1 = wp.zeros(123, dtype=int, device=device)
+    a2 = wp.zeros((2, 3, 4), dtype=float, device=device)
+    out = wp.empty(2, dtype=int, device=device)
+    wp.launch(
+        array_len_kernel,
+        dim=(1,),
+        inputs=(
+            a1,
+            a2,
+        ),
+        outputs=(out,),
+        device=device,
+    )
+
+    test.assertEqual(out.numpy()[0], 123)
+    test.assertEqual(out.numpy()[1], 2)
+
+
 devices = get_test_devices()
 
 
@@ -2835,6 +2905,10 @@ add_function_test(TestArray, "test_array_from_int32_domain", test_array_from_int
 add_function_test(TestArray, "test_array_from_int64_domain", test_array_from_int64_domain, devices=devices)
 add_function_test(TestArray, "test_indexing_types", test_indexing_types, devices=devices)
 
+add_function_test(TestArray, "test_alloc_strides", test_alloc_strides, devices=devices)
+add_function_test(TestArray, "test_casting", test_casting, devices=devices)
+add_function_test(TestArray, "test_array_len", test_array_len, devices=devices)
+
 try:
     import torch