warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__py3-none-win_amd64.whl

warp/tests/sim/test_model.py

@@ -1,226 +0,0 @@
-# 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 numpy as np
-
-import warp as wp
-from warp.sim import ModelBuilder
-from warp.tests.unittest_utils import *
-
-
-class TestModel(unittest.TestCase):
-    def test_add_triangles(self):
-        rng = np.random.default_rng(123)
-
-        pts = np.array(
-            [
-                [-0.00585869, 0.34189449, -1.17415233],
-                [-1.894547, 0.1788074, 0.9251329],
-                [-1.26141048, 0.16140787, 0.08823282],
-                [-0.08609255, -0.82722546, 0.65995427],
-                [0.78827592, -1.77375711, -0.55582718],
-            ]
-        )
-        tris = np.array([[0, 3, 4], [0, 2, 3], [2, 1, 3], [1, 4, 3]])
-
-        builder1 = ModelBuilder()
-        builder2 = ModelBuilder()
-        for pt in pts:
-            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
-            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
-
-        # test add_triangle(s) with default arguments:
-        areas = builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2])
-        for i, t in enumerate(tris):
-            area = builder1.add_triangle(t[0], t[1], t[2])
-            self.assertAlmostEqual(area, areas[i], places=6)
-
-        # test add_triangle(s) with non default arguments:
-        tri_ke = rng.standard_normal(size=pts.shape[0])
-        tri_ka = rng.standard_normal(size=pts.shape[0])
-        tri_kd = rng.standard_normal(size=pts.shape[0])
-        tri_drag = rng.standard_normal(size=pts.shape[0])
-        tri_lift = rng.standard_normal(size=pts.shape[0])
-        for i, t in enumerate(tris):
-            builder1.add_triangle(
-                t[0],
-                t[1],
-                t[2],
-                tri_ke[i],
-                tri_ka[i],
-                tri_kd[i],
-                tri_drag[i],
-                tri_lift[i],
-            )
-        builder2.add_triangles(tris[:, 0], tris[:, 1], tris[:, 2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
-
-        assert_np_equal(np.array(builder1.tri_indices), np.array(builder2.tri_indices))
-        assert_np_equal(np.array(builder1.tri_poses), np.array(builder2.tri_poses), tol=1.0e-6)
-        assert_np_equal(np.array(builder1.tri_activations), np.array(builder2.tri_activations))
-        assert_np_equal(np.array(builder1.tri_materials), np.array(builder2.tri_materials))
-
-    def test_add_edges(self):
-        rng = np.random.default_rng(123)
-
-        pts = np.array(
-            [
-                [-0.00585869, 0.34189449, -1.17415233],
-                [-1.894547, 0.1788074, 0.9251329],
-                [-1.26141048, 0.16140787, 0.08823282],
-                [-0.08609255, -0.82722546, 0.65995427],
-                [0.78827592, -1.77375711, -0.55582718],
-            ]
-        )
-        edges = np.array([[0, 4, 3, 1], [3, 2, 4, 1]])
-
-        builder1 = ModelBuilder()
-        builder2 = ModelBuilder()
-        for pt in pts:
-            builder1.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
-            builder2.add_particle(wp.vec3(pt), wp.vec3(), 1.0)
-
-        # test defaults:
-        for i in range(2):
-            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3])
-        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3])
-
-        # test non defaults:
-        rest = rng.standard_normal(size=2)
-        edge_ke = rng.standard_normal(size=2)
-        edge_kd = rng.standard_normal(size=2)
-        for i in range(2):
-            builder1.add_edge(edges[i, 0], edges[i, 1], edges[i, 2], edges[i, 3], rest[i], edge_ke[i], edge_kd[i])
-        builder2.add_edges(edges[:, 0], edges[:, 1], edges[:, 2], edges[:, 3], rest, edge_ke, edge_kd)
-
-        assert_np_equal(np.array(builder1.edge_indices), np.array(builder2.edge_indices))
-        assert_np_equal(np.array(builder1.edge_rest_angle), np.array(builder2.edge_rest_angle), tol=1.0e-4)
-        assert_np_equal(np.array(builder1.edge_bending_properties), np.array(builder2.edge_bending_properties))
-
-    def test_collapse_fixed_joints(self):
-        def add_three_cubes(builder: ModelBuilder, parent_body=-1):
-            unit_cube = {"hx": 0.5, "hy": 0.5, "hz": 0.5, "density": 1.0}
-            b0 = builder.add_body()
-            builder.add_shape_box(body=b0, **unit_cube)
-            builder.add_joint_fixed(parent=parent_body, child=b0, parent_xform=wp.transform(wp.vec3(1.0, 0.0, 0.0)))
-            b1 = builder.add_body()
-            builder.add_shape_box(body=b1, **unit_cube)
-            builder.add_joint_fixed(parent=parent_body, child=b1, parent_xform=wp.transform(wp.vec3(0.0, 1.0, 0.0)))
-            b2 = builder.add_body()
-            builder.add_shape_box(body=b2, **unit_cube)
-            builder.add_joint_fixed(parent=parent_body, child=b2, parent_xform=wp.transform(wp.vec3(0.0, 0.0, 1.0)))
-            return b2
-
-        builder = ModelBuilder()
-        # only fixed joints
-        builder.add_articulation()
-        add_three_cubes(builder)
-        assert builder.joint_count == 3
-        assert builder.body_count == 3
-
-        # fixed joints followed by a non-fixed joint
-        builder.add_articulation()
-        last_body = add_three_cubes(builder)
-        assert builder.joint_count == 6
-        assert builder.body_count == 6
-        assert builder.articulation_count == 2
-        b3 = builder.add_body()
-        builder.add_shape_box(body=b3, hx=0.5, hy=0.5, hz=0.5, density=1.0, pos=wp.vec3(1.0, 2.0, 3.0))
-        builder.add_joint_revolute(parent=last_body, child=b3, axis=wp.vec3(0.0, 1.0, 0.0))
-
-        # a non-fixed joint followed by fixed joints
-        builder.add_articulation()
-        b4 = builder.add_body()
-        builder.add_shape_box(body=b4, hx=0.5, hy=0.5, hz=0.5, density=1.0)
-        builder.add_joint_free(parent=-1, child=b4, parent_xform=wp.transform(wp.vec3(0.0, -1.0, 0.0)))
-        assert builder.joint_count == 8
-        assert builder.body_count == 8
-        assert builder.articulation_count == 3
-        lb = add_three_cubes(builder, parent_body=b4)
-        # add shape with no mass (e.g. visual only shape)
-        builder.add_shape_sphere(lb, density=0.0)
-
-        builder.collapse_fixed_joints()
-
-        assert builder.joint_count == 2
-        assert builder.articulation_count == 2
-        assert builder.articulation_start == [0, 1]
-        assert builder.joint_type == [wp.sim.JOINT_REVOLUTE, wp.sim.JOINT_FREE]
-        assert builder.shape_count == 12
-        assert builder.shape_body == [-1, -1, -1, -1, -1, -1, 0, 1, 1, 1, 1, 1]
-        assert builder.body_count == 2
-        assert builder.body_com[0] == wp.vec3(1.0, 2.0, 3.0)
-        assert builder.body_com[1] == wp.vec3(0.25, 0.25, 0.25)
-        assert builder.body_mass == [1.0, 4.0]
-        assert builder.body_inv_mass == [1.0, 0.25]
-
-        # create another builder, test add_builder function
-        builder2 = ModelBuilder()
-        builder2.add_builder(builder)
-        assert builder2.articulation_count == builder.articulation_count
-        assert builder2.joint_count == builder.joint_count
-        assert builder2.body_count == builder.body_count
-        assert builder2.shape_count == builder.shape_count
-        assert builder2.articulation_start == builder.articulation_start
-        # add the same builder again
-        builder2.add_builder(builder)
-        assert builder2.articulation_count == 2 * builder.articulation_count
-        assert builder2.articulation_start == [0, 1, 2, 3]
-
-    def test_add_builder_with_open_edges(self):
-        builder = wp.sim.ModelBuilder()
-
-        dim_x = 16
-        dim_y = 16
-
-        env_builder = wp.sim.ModelBuilder()
-        env_builder.add_cloth_grid(
-            pos=wp.vec3(0.0, 0.0, 0.0),
-            vel=wp.vec3(0.1, 0.1, 0.0),
-            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.25),
-            dim_x=dim_x,
-            dim_y=dim_y,
-            cell_x=1.0 / dim_x,
-            cell_y=1.0 / dim_y,
-            mass=1.0,
-        )
-
-        num_envs = 2
-        env_offsets = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]])
-
-        builder_open_edge_count = np.sum(np.array(builder.edge_indices) == -1)
-        env_builder_open_edge_count = np.sum(np.array(env_builder.edge_indices) == -1)
-
-        for i in range(num_envs):
-            xform = wp.transform(env_offsets[i], wp.quat_identity())
-            builder.add_builder(
-                env_builder,
-                xform,
-                update_num_env_count=True,
-                separate_collision_group=True,
-            )
-
-        self.assertEqual(
-            np.sum(np.array(builder.edge_indices) == -1),
-            builder_open_edge_count + num_envs * env_builder_open_edge_count,
-            "builder does not have the expected number of open edges",
-        )
-
-
-if __name__ == "__main__":
-    wp.clear_kernel_cache()
-    unittest.main(verbosity=2)

warp/tests/sim/test_sim_grad.py

@@ -1,287 +0,0 @@
-# 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
-
-    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(mode="basic")
-
-
-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)