warp-lang 1.7.0__py3-none-manylinux_2_34_aarch64.whl

warp/tests/run_coverage_serial.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Serial code-coverage runner
+
+This script is used to generate code-coverage reports by running Warp tests.
+It runs in serial so can take over an hour to finish. To generate a coverage
+report in parallel, use the warp/thirdparty./unittest_parallel.py script
+instead with the --coverage option, e.g. python -m warp.tests --coverage
+"""
+
+import coverage
+
+cover = coverage.Coverage(config_file=True, messages=True)
+
+cover.start()
+
+with cover.collect():
+    import unittest_serial  # noqa: E402
+
+    unittest_serial.run_specified()
+
+cover.save()
+
+cover.report()
+
+cover.html_report(title="Warp Testing Code Coverage Report")

warp/tests/sim/disabled_kinematics.py

@@ -0,0 +1,244 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import unittest
+
+import warp as wp
+import warp.examples
+import warp.sim
+from warp.tests.unittest_utils import *
+
+
+def build_ant(num_envs):
+    builder = wp.sim.ModelBuilder()
+    for i in range(num_envs):
+        wp.sim.parse_mjcf(os.path.join(warp.examples.get_asset_directory(), "nv_ant.xml"), builder, up_axis="y")
+
+        coord_count = 15
+        dof_count = 14
+
+        coord_start = i * coord_count
+        dof_start = i * dof_count
+
+        # base
+        p = [i * 2.0, 0.70, 0.0]
+        q = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5)
+        builder.joint_q[coord_start : coord_start + 3] = p
+        builder.joint_q[coord_start + 3 : coord_start + 7] = q
+
+        # joints
+        builder.joint_q[coord_start + 7 : coord_start + coord_count] = [0.0, 1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0]
+        builder.joint_qd[dof_start + 6 : dof_start + dof_count] = [1.0, 1.0, 1.0, -1.0, 1.0, -1.0, 1.0, 1.0]
+
+    return builder
+
+
+def build_complex_joint_mechanism(chain_length):
+    builder = wp.sim.ModelBuilder()
+    com0 = wp.vec3(1.0, 2.0, 3.0)
+    com1 = wp.vec3(4.0, 5.0, 6.0)
+    com2 = wp.vec3(7.0, 8.0, 9.0)
+    ax0 = wp.normalize(wp.vec3(-1.0, 2.0, 3.0))
+    ax1 = wp.normalize(wp.vec3(4.0, -1.0, 2.0))
+    ax2 = wp.normalize(wp.vec3(-3.0, 4.0, -1.0))
+    # declare some transforms with nonzero translation and orientation
+    tf0 = wp.transform(wp.vec3(1.0, 2.0, 3.0), wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), math.pi * 0.25))
+    tf1 = wp.transform(wp.vec3(4.0, 5.0, 6.0), wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), math.pi * 0.5))
+    tf2 = wp.transform(wp.vec3(7.0, 8.0, 9.0), wp.quat_from_axis_angle(wp.vec3(0.0, 0.0, 1.0), math.pi * 0.75))
+
+    parent = -1
+    for _i in range(chain_length):
+        b0 = builder.add_body(com=com0)
+        builder.add_joint_fixed(parent=parent, child=b0, parent_xform=tf1, child_xform=tf0)
+        assert builder.articulation_count == 1
+
+        b1 = builder.add_body(com=com1)
+        builder.add_joint_revolute(parent=b0, child=b1, parent_xform=tf1, child_xform=tf2, axis=ax1)
+        builder.joint_q[-1] = 0.3
+        builder.joint_qd[-1] = 1.0
+
+        b2 = builder.add_body(com=com2)
+        builder.add_joint_universal(parent=b1, child=b2, parent_xform=tf2, child_xform=tf0, axis_0=ax0, axis_1=ax1)
+        builder.joint_q[-2:] = [0.3, 0.5]
+        builder.joint_qd[-2:] = [1.0, -1.0]
+
+        b3 = builder.add_body(com=com0)
+        builder.add_joint_ball(parent=b2, child=b3, parent_xform=tf0, child_xform=tf1)
+        builder.joint_q[-4:] = list(wp.quat_from_axis_angle(ax0, 0.7))
+        builder.joint_qd[-3:] = [1.0, -0.6, 1.5]
+
+        b4 = builder.add_body(com=com1)
+        builder.add_joint_compound(
+            parent=b3,
+            child=b4,
+            parent_xform=tf2,
+            child_xform=tf1,
+            axis_0=(0, 0, 1),
+            axis_1=(1, 0, 0),
+            axis_2=(0, 1, 0),
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.27]
+        builder.joint_qd[-3:] = [1.23, -1.0, 0.5]
+
+        b5 = builder.add_body(com=com2)
+        builder.add_joint_prismatic(
+            parent=b4,
+            child=b5,
+            parent_xform=tf2,
+            child_xform=tf0,
+            axis=ax0,
+        )
+        builder.joint_q[-1] = 0.92
+        builder.joint_qd[-1] = -0.63
+
+        b6 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b5,
+            child=b6,
+            parent_xform=tf0,
+            child_xform=tf2,
+            linear_axes=[ax0, ax1, wp.cross(ax0, ax1)],
+            angular_axes=[ax1, ax2, wp.cross(ax1, ax2)],
+        )
+        builder.joint_q[-6:] = [0.3, 0.5, 0.7, 0.9, 1.1, 1.3]
+        builder.joint_qd[-6:] = [1.0, -1.0, 0.5, 0.8, -0.3, 0.1]
+
+        b7 = builder.add_body(com=com1)
+        builder.add_joint_free(
+            parent=b6,
+            child=b7,
+            parent_xform=tf1,
+            child_xform=tf2,
+        )
+        builder.joint_q[-7:] = [0.5, -0.9, 1.4] + list(wp.quat_rpy(0.3, -0.5, 0.7))
+        builder.joint_qd[-6:] = [1.0, -1.0, 0.5, 0.8, -0.3, 0.1]
+
+        b8 = builder.add_body(com=com2)
+        builder.add_joint_distance(
+            parent=b7,
+            child=b8,
+            parent_xform=tf1,
+            child_xform=tf2,
+        )
+        builder.joint_q[-7:] = [-0.3, -0.7, 0.2] + list(wp.quat_rpy(0.1, 0.1, 0.4))
+        builder.joint_qd[-6:] = [-0.34, 0.5, -0.6, -0.4, 0.2, 0.1]
+
+        # D6 joint that behaves like a fixed joint
+        b9 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b8,
+            child=b9,
+            parent_xform=tf0,
+            child_xform=tf2,
+            linear_axes=[],
+            angular_axes=[],
+        )
+
+        b10 = builder.add_body(com=com0)
+        builder.add_joint_d6(
+            parent=b9,
+            child=b10,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[ax1],
+            angular_axes=[ax2, ax0],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        b11 = builder.add_body(com=com1)
+        builder.add_joint_d6(
+            parent=b10,
+            child=b11,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[ax1, ax0, wp.cross(ax1, ax0)],
+            angular_axes=[],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        b12 = builder.add_body(com=com2)
+        builder.add_joint_d6(
+            parent=b11,
+            child=b12,
+            parent_xform=tf1,
+            child_xform=tf2,
+            linear_axes=[],
+            angular_axes=[ax1, ax2, wp.cross(ax1, ax2)],
+        )
+        builder.joint_q[-3:] = [0.3, 0.5, 0.7]
+        builder.joint_qd[-3:] = [1.0, -1.0, 0.5]
+
+        parent = b12
+
+    return builder
+
+
+def check_fk_ik(builder, device):
+    model = builder.finalize(device)
+    state = model.state()
+
+    q_fk = model.joint_q.numpy()
+    qd_fk = model.joint_qd.numpy()
+
+    wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, state)
+
+    q_ik = wp.zeros_like(model.joint_q)
+    qd_ik = wp.zeros_like(model.joint_qd)
+
+    wp.sim.eval_ik(model, state, q_ik, qd_ik)
+
+    # adjust numpy print settings
+    # np.set_printoptions(precision=4, floatmode="fixed", suppress=True)
+    # print("q:")
+    # print(np.array(q_fk))
+    # print(q_ik.numpy())
+
+    # print("qd:")
+    # print(np.array(qd_fk))
+    # print(qd_ik.numpy())
+
+    assert_np_equal(q_ik.numpy(), q_fk, tol=1e-4)
+    assert_np_equal(qd_ik.numpy(), qd_fk, tol=1e-4)
+
+
+def test_fk_ik_ant(test, device):
+    builder = build_ant(3)
+    check_fk_ik(builder, device)
+
+
+def test_fk_ik_complex_joint_mechanism(test, device):
+    builder = build_complex_joint_mechanism(2)
+    check_fk_ik(builder, device)
+
+
+devices = get_test_devices()
+
+
+class TestKinematics(unittest.TestCase):
+    pass
+
+
+add_function_test(TestKinematics, "test_fk_ik_ant", test_fk_ik_ant, devices=devices)
+add_function_test(
+    TestKinematics, "test_fk_ik_complex_joint_mechanism", test_fk_ik_complex_joint_mechanism, devices=devices
+)
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=False)

warp/tests/sim/flaky_test_sim_grad.py

@@ -0,0 +1,290 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+import warp.sim.render
+from warp.tests.unittest_utils import *
+
+
+@wp.kernel
+def evaluate_loss(
+    joint_q: wp.array(dtype=float),
+    weighting: float,
+    target: float,
+    # output
+    loss: wp.array(dtype=float),
+):
+    tid = wp.tid()
+    d = (target - joint_q[tid * 2 + 1]) ** 2.0
+    wp.atomic_add(loss, 0, weighting * d)
+
+
+@wp.kernel
+def assign_action(action: wp.array(dtype=float), joint_act: wp.array(dtype=float)):
+    tid = wp.tid()
+    joint_act[2 * tid] = action[tid]
+
+
+@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.
+    """
+
+    def f(xs):
+        # call the kernel without taping for finite differences
+        wp_xs = [wp.array(xs[i], ndim=1, dtype=inputs[i].dtype, device=device) for i in range(len(inputs))]
+        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
+    np_xs = []
+    for i in range(len(inputs)):
+        np_xs.append(inputs[i].numpy().flatten().copy())
+
+    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
+            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(ad_grad * fd_grad > 0, True)
+
+    tape.zero()
+
+
+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 = 150
+    sim_substeps = 10
+    dt = 1 / 30
+
+    target = 3.0
+
+    if integrator_type == 0:
+        contact_ke = 1e3
+        contact_kd = 1e1
+    else:
+        contact_ke = 1e3
+        contact_kd = 1e1
+
+    complete_builder = wp.sim.ModelBuilder()
+
+    complete_builder.default_shape_ke = contact_ke
+    complete_builder.default_shape_kd = contact_kd
+
+    for _ in range(num_envs):
+        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_sphere(b0, radius=0.4, density=100.0)
+
+        b1 = builder.add_body(name="passive")
+        builder.add_shape_sphere(b1, radius=0.47, density=100.0)
+
+        if joint_type == 0:
+            builder.add_joint_prismatic(-1, b0)
+            builder.add_joint_prismatic(-1, b1)
+        else:
+            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, 2.0]
+        complete_builder.add_builder(builder)
+
+    assert complete_builder.body_count == 2 * num_envs
+    assert complete_builder.joint_count == 2 * num_envs
+    assert set(complete_builder.shape_collision_group) == set(range(1, num_envs + 1))
+
+    complete_builder.gravity = 0.0
+    model = complete_builder.finalize(device=device, requires_grad=True)
+    model.ground = False
+    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:
+        integrator = wp.sim.SemiImplicitIntegrator()
+        sim_substeps *= 5
+    else:
+        integrator = wp.sim.XPBDIntegrator(iterations=2, rigid_contact_relaxation=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)]
+
+        wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, states[0])
+
+        control_active = model.control()
+        control_nop = model.control()
+
+        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):
+            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(state)
+                renderer.end_frame()
+                render_time += dt
+            for _ in range(sim_substeps):
+                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):
+            # compute generalized coordinates
+            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, weighting, target],
+            outputs=[loss],
+            device=model.device,
+        )
+
+        # if renderer:
+        #     renderer.save()
+
+        return loss
+
+    action_too_far = wp.array(
+        [80.0 for _ in range(num_envs)],
+        device=device,
+        dtype=wp.float32,
+        requires_grad=True,
+    )
+    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, print_grad=print_grad)
+
+    action_too_close = wp.array(
+        [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, print_grad=print_grad)
+
+
+devices = get_test_devices()
+
+
+class TestSimGradients(unittest.TestCase):
+    pass
+
+
+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=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()
+    unittest.main(verbosity=2, failfast=True)