warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/tests/sim/test_coloring.py

@@ -0,0 +1,258 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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 numpy as np
+
+import warp as wp
+import warp.examples
+import warp.sim
+from warp.sim.graph_coloring import (
+    ColoringAlgorithm,
+    construct_trimesh_graph_edges,
+    convert_to_color_groups,
+    validate_graph_coloring,
+)
+from warp.tests.unittest_utils import *
+
+
+def color_lattice_grid(num_x, num_y):
+    colors = []
+    for _ in range(4):
+        colors.append([])
+
+    for xi in range(num_x + 1):
+        for yi in range(num_y + 1):
+            node_dx = yi * (num_x + 1) + xi
+
+            a = 1 if xi % 2 else 0
+            b = 1 if yi % 2 else 0
+
+            c = b * 2 + a
+
+            colors[c].append(node_dx)
+
+    color_groups = [np.array(group) for group in colors]
+
+    return color_groups
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_coloring_trimesh(test, device):
+    from pxr import Usd, UsdGeom
+
+    with wp.ScopedDevice(device):
+        usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+        vertices = np.array(usd_geom.GetPointsAttr().Get())
+        faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+        builder = wp.sim.ModelBuilder()
+
+        builder.add_cloth_mesh(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            scale=1.0,
+            vertices=[wp.vec3(p) for p in vertices],
+            indices=faces.flatten(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+        )
+
+        model = builder.finalize()
+
+        particle_colors = wp.empty(shape=(model.particle_count), dtype=int, device="cpu")
+
+        edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
+
+        # coloring without bending
+        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+            model.particle_count,
+            edge_indices_cpu.__ctype__(),
+            ColoringAlgorithm.GREEDY.value,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, particle_colors],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+            model.particle_count,
+            edge_indices_cpu.__ctype__(),
+            ColoringAlgorithm.MCS.value,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, particle_colors],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+        # coloring with bending
+        edge_indices_cpu_with_bending = construct_trimesh_graph_edges(model.edge_indices, True)
+        num_colors_greedy = wp.context.runtime.core.graph_coloring(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            ColoringAlgorithm.GREEDY.value,
+            particle_colors.__ctype__(),
+        )
+        wp.context.runtime.core.balance_coloring(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            num_colors_greedy,
+            1.1,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu_with_bending, particle_colors],
+            dim=edge_indices_cpu_with_bending.shape[0],
+            device="cpu",
+        )
+
+        num_colors_mcs = wp.context.runtime.core.graph_coloring(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            ColoringAlgorithm.MCS.value,
+            particle_colors.__ctype__(),
+        )
+        max_min_ratio = wp.context.runtime.core.balance_coloring(
+            model.particle_count,
+            edge_indices_cpu_with_bending.__ctype__(),
+            num_colors_mcs,
+            1.1,
+            particle_colors.__ctype__(),
+        )
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu_with_bending, particle_colors],
+            dim=edge_indices_cpu_with_bending.shape[0],
+            device="cpu",
+        )
+
+        color_categories_balanced = convert_to_color_groups(num_colors_mcs, particle_colors)
+
+        color_sizes = np.array([c.shape[0] for c in color_categories_balanced], dtype=np.float32)
+        test.assertTrue(np.max(color_sizes) / np.min(color_sizes) <= max_min_ratio)
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_combine_coloring(test, device):
+    from pxr import Usd, UsdGeom
+
+    with wp.ScopedDevice(device):
+        builder1 = wp.sim.ModelBuilder()
+        usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+        usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+        vertices = np.array(usd_geom.GetPointsAttr().Get())
+        faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+        builder1.add_cloth_mesh(
+            pos=wp.vec3(0.0, 0.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            scale=1.0,
+            vertices=[wp.vec3(p) for p in vertices],
+            indices=faces.flatten(),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            density=0.02,
+        )
+
+        builder1.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            fix_left=True,
+        )
+        builder1.color()
+
+        builder2 = wp.sim.ModelBuilder()
+        builder2.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            # to include bending in coloring
+            edge_ke=100000,
+            fix_left=True,
+        )
+        builder2.color()
+
+        builder3 = wp.sim.ModelBuilder()
+        builder3.add_cloth_grid(
+            pos=wp.vec3(0.0, 4.0, 0.0),
+            rot=wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.0),
+            vel=wp.vec3(0.0, 0.0, 0.0),
+            dim_x=50,
+            dim_y=100,
+            cell_x=0.1,
+            cell_y=0.1,
+            mass=0.1,
+            fix_left=True,
+        )
+
+        builder3.set_coloring(
+            color_lattice_grid(50, 100),
+        )
+
+        builder1.add_builder(builder2)
+        builder1.add_builder(builder3)
+
+        model = builder2.finalize()
+
+        particle_number_colored = np.full((model.particle_count), -1, dtype=int)
+        particle_colors = np.full((model.particle_count), -1, dtype=int)
+        for color, color_group in enumerate(model.particle_coloring):
+            particle_number_colored[color_group.numpy()] += 1
+            particle_colors[color_group.numpy()] = color
+
+        # all particles has been colored exactly once
+        assert_np_equal(particle_number_colored, 0)
+
+        edge_indices_cpu = wp.array(model.edge_indices.numpy()[:, 2:], dtype=int, device="cpu")
+        wp.launch(
+            kernel=validate_graph_coloring,
+            inputs=[edge_indices_cpu, wp.array(particle_colors, dtype=int, device="cpu")],
+            dim=edge_indices_cpu.shape[0],
+            device="cpu",
+        )
+
+
+devices = get_test_devices()
+
+
+class TestColoring(unittest.TestCase):
+    pass
+
+
+add_function_test(TestColoring, "test_coloring_trimesh", test_coloring_trimesh, devices=devices)
+add_function_test(TestColoring, "test_combine_coloring", test_combine_coloring, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/sim/test_model.py

@@ -0,0 +1,224 @@
+# 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
+        add_three_cubes(builder, parent_body=b4)
+
+        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 == 11
+        assert builder.shape_body == [-1, -1, -1, -1, -1, -1, 0, 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_bounce_linear.py

@@ -0,0 +1,212 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 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)