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

warp/sim/collide.py

@@ -837,6 +837,7 @@ def count_contact_points(
                 num_actual_contacts = 8 + 4
         else:
             num_contacts = 8
+            num_actual_contacts = 8
     elif actual_type_a == wp.sim.GEO_MESH:
         mesh_a = wp.mesh_get(geo.source[actual_shape_a])
         num_contacts_a = mesh_a.points.shape[0]
@@ -1744,7 +1745,7 @@ def compute_tri_aabbs(
     v2 = pos[tri_indices[t_id, 1]]
     v3 = pos[tri_indices[t_id, 2]]
 
-    upper, lower = compute_tri_aabb(v1, v2, v3)
+    lower, upper = compute_tri_aabb(v1, v2, v3)
 
     lower_bounds[t_id] = lower
     upper_bounds[t_id] = upper
@@ -1809,7 +1810,7 @@ def init_triangle_collision_data_kernel(
 @wp.kernel
 def vertex_triangle_collision_detection_kernel(
     query_radius: float,
-    mesh_id: wp.uint64,
+    bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
     tri_indices: wp.array(dtype=wp.int32, ndim=2),
     vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
@@ -1840,7 +1841,7 @@ def vertex_triangle_collision_detection_kernel(
         and vertex_colliding_triangles_count.
 
     Attributes:
-        mesh_id (int): the mesh id you want to collide with
+        bvh_id (int): the bvh id you want to collide with
         query_radius (float): the contact radius. vertex-triangle pairs whose distance are less than this will get detected
         pos (array): positions of all the vertices that make up triangles
         vertex_colliding_triangles (array): flattened buffer of vertices' collision triangles
@@ -1862,12 +1863,12 @@ def vertex_triangle_collision_detection_kernel(
     lower = wp.vec3(v[0] - query_radius, v[1] - query_radius, v[2] - query_radius)
     upper = wp.vec3(v[0] + query_radius, v[1] + query_radius, v[2] + query_radius)
 
-    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
 
     tri_index = wp.int32(0)
     vertex_num_collisions = wp.int32(0)
     min_dis_to_tris = query_radius
-    while wp.mesh_query_aabb_next(query, tri_index):
+    while wp.bvh_query_next(query, tri_index):
         t1 = tri_indices[tri_index, 0]
         t2 = tri_indices[tri_index, 1]
         t3 = tri_indices[tri_index, 2]
@@ -1991,7 +1992,7 @@ def edge_colliding_edges_detection_kernel(
 
 @wp.kernel
 def triangle_triangle_collision_detection_kernel(
-    mesh_id: wp.uint64,
+    bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
     tri_indices: wp.array(dtype=wp.int32, ndim=2),
     triangle_intersecting_triangles_offsets: wp.array(dtype=wp.int32),
@@ -2014,10 +2015,10 @@ def triangle_triangle_collision_detection_kernel(
     buffer_offset = triangle_intersecting_triangles_offsets[tri_index]
     buffer_size = triangle_intersecting_triangles_offsets[tri_index + 1] - buffer_offset
 
-    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
     tri_index_2 = wp.int32(0)
     intersection_count = wp.int32(0)
-    while wp.mesh_query_aabb_next(query, tri_index_2):
+    while wp.bvh_query_next(query, tri_index_2):
         t2_v1 = tri_indices[tri_index_2, 0]
         t2_v2 = tri_indices[tri_index_2, 1]
         t2_v3 = tri_indices[tri_index_2, 2]
@@ -2129,7 +2130,16 @@ class TriMeshCollisionDetector:
 
         self.edge_edge_parallel_epsilon = edge_edge_parallel_epsilon
 
-        self.mesh_tris = wp.Mesh(model.particle_q, model.tri_indices.flatten())
+        self.lower_bounds_tris = wp.array(shape=(model.tri_count,), dtype=wp.vec3, device=model.device)
+        self.upper_bounds_tris = wp.array(shape=(model.tri_count,), dtype=wp.vec3, device=model.device)
+        wp.launch(
+            kernel=compute_tri_aabbs,
+            inputs=[self.vertex_positions, model.tri_indices, self.lower_bounds_tris, self.upper_bounds_tris],
+            dim=model.tri_count,
+            device=model.device,
+        )
+
+        self.bvh_tris = wp.Bvh(self.lower_bounds_tris, self.upper_bounds_tris)
 
         # collision detections results
 
@@ -2244,15 +2254,18 @@ class TriMeshCollisionDetector:
     def refit(self, new_pos=None):
         if new_pos is not None:
             self.vertex_positions = new_pos
-            # this will automatically apply refit
-            self.mesh_tris.points = new_pos
-        else:
-            self.refit_triangles()
 
+        self.refit_triangles()
         self.refit_edges()
 
     def refit_triangles(self):
-        self.mesh_tris.refit()
+        wp.launch(
+            kernel=compute_tri_aabbs,
+            inputs=[self.vertex_positions, self.model.tri_indices, self.lower_bounds_tris, self.upper_bounds_tris],
+            dim=self.model.tri_count,
+            device=self.model.device,
+        )
+        self.bvh_tris.refit()
 
     def refit_edges(self):
         wp.launch(
@@ -2282,7 +2295,7 @@ class TriMeshCollisionDetector:
             kernel=vertex_triangle_collision_detection_kernel,
             inputs=[
                 query_radius,
-                self.mesh_tris.id,
+                self.bvh_tris.id,
                 self.vertex_positions,
                 self.model.tri_indices,
                 self.vertex_colliding_triangles_offsets,
@@ -2348,7 +2361,7 @@ class TriMeshCollisionDetector:
         wp.launch(
             kernel=triangle_triangle_collision_detection_kernel,
             inputs=[
-                self.mesh_tris.id,
+                self.bvh_tris.id,
                 self.vertex_positions,
                 self.model.tri_indices,
                 self.triangle_intersecting_triangles_offsets,

warp/sim/import_urdf.py

@@ -12,6 +12,7 @@ from typing import Union
 import numpy as np
 
 import warp as wp
+import warp.sim
 from warp.sim.model import Mesh
 
 
@@ -20,7 +21,7 @@ def parse_urdf(
     builder,
     xform=None,
     floating=False,
-    base_joint: Union[dict, str] = None,
+    base_joint: Union[dict, str, None] = None,
     density=1000.0,
     stiffness=100.0,
     damping=10.0,
@@ -511,11 +512,25 @@ def parse_urdf(
 
             builder.add_joint_d6(
                 linear_axes=[
-                    wp.sim.JointAxis(
-                        u, limit_lower=lower * scale, limit_upper=upper * scale, limit_ke=limit_ke, limit_kd=limit_kd
+                    warp.sim.JointAxis(
+                        u,
+                        limit_lower=lower * scale,
+                        limit_upper=upper * scale,
+                        limit_ke=limit_ke,
+                        limit_kd=limit_kd,
+                        target_ke=stiffness,
+                        target_kd=joint_damping,
+                        mode=joint_mode,
                     ),
-                    wp.sim.JointAxis(
-                        v, limit_lower=lower * scale, limit_upper=upper * scale, limit_ke=limit_ke, limit_kd=limit_kd
+                    warp.sim.JointAxis(
+                        v,
+                        limit_lower=lower * scale,
+                        limit_upper=upper * scale,
+                        limit_ke=limit_ke,
+                        limit_kd=limit_kd,
+                        target_ke=stiffness,
+                        target_kd=joint_damping,
+                        mode=joint_mode,
                     ),
                 ],
                 **joint_params,

warp/sim/integrator_featherstone.py

@@ -1177,14 +1177,7 @@ def create_inertia_matrix_kernel(num_joints, num_dofs):
             P_body = wp.tile_matmul(M_body, J_body)
 
             # assign to the P slice
-            wp.tile_assign(
-                P,
-                P_body,
-                offset=(
-                    i * 6,
-                    0,
-                ),
-            )
+            wp.tile_assign(P, P_body, offset=(i * 6, 0))
 
         # compute H = J^T*P
         H = wp.tile_matmul(wp.tile_transpose(J), P)
@@ -1205,9 +1198,9 @@ def create_batched_cholesky_kernel(num_dofs):
 
         a = wp.tile_load(A[articulation], shape=(num_dofs, num_dofs), storage="shared")
         r = wp.tile_load(R[articulation], shape=num_dofs, storage="shared")
-        wp.tile_diag_add(a, r)
-        wp.tile_cholesky(a)
-        wp.tile_store(L[articulation], wp.tile_transpose(a))
+        a_r = wp.tile_diag_add(a, r)
+        l = wp.tile_cholesky(a_r)
+        wp.tile_store(L[articulation], wp.tile_transpose(l))
 
     return eval_tiled_dense_cholesky_batched
 

warp/sim/model.py

@@ -7,6 +7,8 @@
 
 """A module for building simulation models and state."""
 
+from __future__ import annotations
+
 import copy
 import math
 from typing import List, Optional, Tuple
@@ -262,47 +264,51 @@ class Mesh:
 
 
 class State:
-    """The State object holds all *time-varying* data for a model.
+    """Time-varying state data for a :class:`Model`.
 
-    Time-varying data includes particle positions, velocities, rigid body states, and
+    Time-varying state data includes particle positions, velocities, rigid body states, and
     anything that is output from the integrator as derived data, e.g.: forces.
 
     The exact attributes depend on the contents of the model. State objects should
     generally be created using the :func:`Model.state()` function.
+    """
 
-    Attributes:
+    def __init__(self):
+        self.particle_q: Optional[wp.array] = None
+        """Array of 3D particle positions with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-        particle_q (array): Array of 3D particle positions, shape [particle_count], :class:`vec3`
-        particle_qd (array): Array of 3D particle velocities, shape [particle_count], :class:`vec3`
-        particle_f (array): Array of 3D particle forces, shape [particle_count], :class:`vec3`
+        self.particle_qd: Optional[wp.array] = None
+        """Array of 3D particle velocities with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-        body_q (array): Array of body coordinates (7-dof transforms) in maximal coordinates, shape [body_count], :class:`transform`
-        body_qd (array): Array of body velocities in maximal coordinates (first 3 entries represent angular velocity, last 3 entries represent linear velocity), shape [body_count], :class:`spatial_vector`
-        body_f (array): Array of body forces in maximal coordinates (first 3 entries represent torque, last 3 entries represent linear force), shape [body_count], :class:`spatial_vector`
+        self.particle_f: Optional[wp.array] = None
+        """Array of 3D particle forces with shape ``(particle_count,)`` and type :class:`vec3`."""
 
-            Note:
+        self.body_q: Optional[wp.array] = None
+        """Array of body coordinates (7-dof transforms) in maximal coordinates with shape ``(body_count,)`` and type :class:`transform`."""
 
-                :attr:`body_f` represents external wrenches in world frame and denotes wrenches measured w.r.t. to the body's center of mass for all integrators except :class:`FeatherstoneIntegrator` which assumes the wrenches are measured w.r.t. world origin.
-
-        joint_q (array): Array of generalized joint coordinates, shape [joint_coord_count], float
-        joint_qd (array): Array of generalized joint velocities, shape [joint_dof_count], float
+        self.body_qd: Optional[wp.array] = None
+        """Array of body velocities in maximal coordinates (first three entries represent angular velocity,
+        last three entries represent linear velocity) with shape ``(body_count,)`` and type :class:`spatial_vector`.
+        """
 
-    """
+        self.body_f: Optional[wp.array] = None
+        """Array of body forces in maximal coordinates (first three entries represent torque, last three
+        entries represent linear force) with shape ``(body_count,)`` and type :class:`spatial_vector`.
 
-    def __init__(self):
-        self.particle_q = None
-        self.particle_qd = None
-        self.particle_f = None
+        .. note::
+            :attr:`body_f` represents external wrenches in world frame and denotes wrenches measured w.r.t.
+            to the body's center of mass for all integrators except :class:`FeatherstoneIntegrator`, which
+            assumes the wrenches are measured w.r.t. world origin.
+        """
 
-        self.body_q = None
-        self.body_qd = None
-        self.body_f = None
+        self.joint_q: Optional[wp.array] = None
+        """Array of generalized joint coordinates with shape ``(joint_coord_count,)`` and type ``float``."""
 
-        self.joint_q = None
-        self.joint_qd = None
+        self.joint_qd: Optional[wp.array] = None
+        """Array of generalized joint velocities with shape ``(joint_dof_count,)`` and type ``float``."""
 
-    def clear_forces(self):
-        """Clears all forces (for particles and bodies) in the state object."""
+    def clear_forces(self) -> None:
+        """Clear all forces (for particles and bodies) in the state object."""
         with wp.ScopedTimer("clear_forces", False):
             if self.particle_count:
                 self.particle_f.zero_()
@@ -311,7 +317,7 @@ class State:
                 self.body_f.zero_()
 
     @property
-    def requires_grad(self):
+    def requires_grad(self) -> bool:
         """Indicates whether the state arrays have gradient computation enabled."""
         if self.particle_q:
             return self.particle_q.requires_grad
@@ -320,28 +326,28 @@ class State:
         return False
 
     @property
-    def body_count(self):
+    def body_count(self) -> int:
         """The number of bodies represented in the state."""
         if self.body_q is None:
             return 0
         return len(self.body_q)
 
     @property
-    def particle_count(self):
+    def particle_count(self) -> int:
         """The number of particles represented in the state."""
         if self.particle_q is None:
             return 0
         return len(self.particle_q)
 
     @property
-    def joint_coord_count(self):
+    def joint_coord_count(self) -> int:
         """The number of generalized joint position coordinates represented in the state."""
         if self.joint_q is None:
             return 0
         return len(self.joint_q)
 
     @property
-    def joint_dof_count(self):
+    def joint_dof_count(self) -> int:
         """The number of generalized joint velocity coordinates represented in the state."""
         if self.joint_qd is None:
             return 0
@@ -349,37 +355,34 @@ class State:
 
 
 class Control:
+    """Time-varying control data for a :class:`Model`.
+
+    Time-varying control data includes joint control inputs, muscle activations,
+    and activation forces for triangle and tetrahedral elements.
+
+    The exact attributes depend on the contents of the model. Control objects
+    should generally be created using the :func:`Model.control()` function.
     """
-    The Control object holds all *time-varying* control data for a model.
 
-    Time-varying control data includes joint control inputs, muscle activations, and activation forces for triangle and tetrahedral elements.
+    def __init__(self, model: Model):
+        self.model: Model = model
+        """Model to use as a reference for the control inputs."""
 
-    The exact attributes depend on the contents of the model. Control objects should generally be created using the :func:`Model.control()` function.
+        self.joint_act: Optional[wp.array] = None
+        """Array of joint control inputs with shape ``(joint_axis_count,)`` and type ``float``."""
 
-    Attributes:
+        self.tri_activations: Optional[wp.array] = None
+        """Array of triangle element activations with shape ``(tri_count,)`` and type ``float``."""
 
-        joint_act (array): Array of joint control inputs, shape [joint_axis_count], float
-        tri_activations (array): Array of triangle element activations, shape [tri_count], float
-        tet_activations (array): Array of tetrahedral element activations, shape [tet_count], float
-        muscle_activations (array): Array of muscle activations, shape [muscle_count], float
+        self.tet_activations: Optional[wp.array] = None
+        """Array of tetrahedral element activations with shape with shape ``(tet_count,) and type ``float``."""
 
-    """
+        self.muscle_activations: Optional[wp.array] = None
+        """Array of muscle activations with shape ``(muscle_count,)`` and type ``float``."""
 
-    def __init__(self, model):
-        """
-        Args:
-            model (Model): The model to use as a reference for the control inputs
-        """
-        self.model = model
-        self.joint_act = None
-        self.tri_activations = None
-        self.tet_activations = None
-        self.muscle_activations = None
+    def reset(self) -> None:
+        """Reset the control inputs to their initial state defined in :attr:`model`."""
 
-    def reset(self):
-        """
-        Resets the control inputs to their initial state defined in :class:`Model`.
-        """
         if self.joint_act is not None:
             self.joint_act.assign(self.model.joint_act)
         if self.tri_activations is not None:
@@ -1396,6 +1399,8 @@ class ModelBuilder:
                 # apply offset transform to root bodies
                 if xform is not None:
                     self.shape_transform.append(xform * builder.shape_transform[s])
+                else:
+                    self.shape_transform.append(builder.shape_transform[s])
 
         for b, shapes in builder.body_shapes.items():
             self.body_shapes[b + start_body_idx] = [s + start_shape_idx for s in shapes]
@@ -1418,18 +1423,16 @@ class ModelBuilder:
 
             # offset the indices
             self.articulation_start.extend([a + self.joint_count for a in builder.articulation_start])
-            self.joint_parent.extend([p + self.joint_count if p != -1 else -1 for p in builder.joint_parent])
-            self.joint_child.extend([c + self.joint_count for c in builder.joint_child])
+            self.joint_parent.extend([p + self.body_count if p != -1 else -1 for p in builder.joint_parent])
+            self.joint_child.extend([c + self.body_count for c in builder.joint_child])
 
             self.joint_q_start.extend([c + self.joint_coord_count for c in builder.joint_q_start])
             self.joint_qd_start.extend([c + self.joint_dof_count for c in builder.joint_qd_start])
 
             self.joint_axis_start.extend([a + self.joint_axis_total_count for a in builder.joint_axis_start])
 
-        joint_children = set(builder.joint_child)
         for i in range(builder.body_count):
-            if xform is not None and i not in joint_children:
-                # rigid body is not attached to a joint, so apply input transform directly
+            if xform is not None:
                 self.body_q.append(xform * builder.body_q[i])
             else:
                 self.body_q.append(builder.body_q[i])

warp/sim/render.py

@@ -213,7 +213,7 @@ def CreateSimRenderer(renderer):
                                 pos=p,
                                 rot=q,
                                 scale=geo_scale,
-                                colors=[color],
+                                colors=color,
                                 parent_body=body,
                                 is_template=True,
                             )
@@ -328,7 +328,7 @@ def CreateSimRenderer(renderer):
                         "surface",
                         particle_q,
                         self.model.tri_indices.numpy().flatten(),
-                        colors=(((0.75, 0.25, 0.0),) * len(particle_q)),
+                        colors=(0.75, 0.25, 0.0),
                     )
 
                 # render springs

warp/stubs.py

@@ -1785,7 +1785,7 @@ def rand_init(seed: int32, offset: int32) -> uint32:
 
 @over
 def randi(state: uint32) -> int:
-    """Return a random integer in the range [0, 2^32)."""
+    """Return a random integer in the range [-2^31, 2^31)."""
     ...
 
 

warp/tests/test_array.py

@@ -2836,6 +2836,31 @@ def test_array_len(test, device):
     test.assertEqual(out.numpy()[1], 2)
 
 
+def test_cuda_interface_conversion(test, device):
+    class MyArrayInterface:
+        def __init__(self, data):
+            self.data = np.array(data)
+            self.__array_interface__ = self.data.__array_interface__
+            self.__cuda_array_interface__ = self.data.__array_interface__
+            self.__len__ = self.data.__len__
+
+    array = MyArrayInterface((1, 2, 3))
+    wp_array = wp.array(array, dtype=wp.int8, device=device)
+    assert wp_array.ptr != 0
+
+    array = MyArrayInterface((1, 2, 3))
+    wp_array = wp.array(array, dtype=wp.float32, device=device)
+    assert wp_array.ptr != 0
+
+    array = MyArrayInterface((1, 2, 3))
+    wp_array = wp.array(array, dtype=wp.vec3, device=device)
+    assert wp_array.ptr != 0
+
+    array = MyArrayInterface((1, 2, 3, 4))
+    wp_array = wp.array(array, dtype=wp.mat22, device=device)
+    assert wp_array.ptr != 0
+
+
 devices = get_test_devices()
 
 
@@ -2908,6 +2933,7 @@ add_function_test(TestArray, "test_indexing_types", test_indexing_types, 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)
+add_function_test(TestArray, "test_cuda_interface_conversion", test_cuda_interface_conversion, devices=devices)
 
 try:
     import torch

warp/tests/test_collision.py

@@ -18,7 +18,7 @@ from warp.tests.unittest_utils import *
 @wp.kernel
 def vertex_triangle_collision_detection_brute_force(
     query_radius: float,
-    mesh_id: wp.uint64,
+    bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
     tri_indices: wp.array(dtype=wp.int32, ndim=2),
     vertex_colliding_triangles: wp.array(dtype=wp.int32),
@@ -67,7 +67,7 @@ def vertex_triangle_collision_detection_brute_force(
 @wp.kernel
 def validate_vertex_collisions(
     query_radius: float,
-    mesh_id: wp.uint64,
+    bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
     tri_indices: wp.array(dtype=wp.int32, ndim=2),
     vertex_colliding_triangles: wp.array(dtype=wp.int32),
@@ -107,7 +107,7 @@ def validate_vertex_collisions(
 @wp.kernel
 def validate_triangle_collisions(
     query_radius: float,
-    mesh_id: wp.uint64,
+    bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
     tri_indices: wp.array(dtype=wp.int32, ndim=2),
     triangle_colliding_vertices: wp.array(dtype=wp.int32),
@@ -303,7 +303,7 @@ class Example:
                 kernel=vertex_triangle_collision_detection_brute_force,
                 inputs=[
                     query_radius,
-                    self.collision_detector.mesh_tris.id,
+                    self.collision_detector.bvh_tris.id,
                     self.collision_detector.model.particle_q,
                     self.collision_detector.model.tri_indices,
                     self.collision_detector.vertex_colliding_triangles,
@@ -337,7 +337,7 @@ class Example:
                 kernel=validate_vertex_collisions,
                 inputs=[
                     query_radius,
-                    self.collision_detector.mesh_tris.id,
+                    self.collision_detector.bvh_tris.id,
                     self.collision_detector.model.particle_q,
                     self.collision_detector.model.tri_indices,
                     self.collision_detector.vertex_colliding_triangles,
@@ -355,7 +355,7 @@ class Example:
                 kernel=validate_triangle_collisions,
                 inputs=[
                     query_radius,
-                    self.collision_detector.mesh_tris.id,
+                    self.collision_detector.bvh_tris.id,
                     self.collision_detector.model.particle_q,
                     self.collision_detector.model.tri_indices,
                     self.collision_detector.triangle_colliding_vertices,

warp/tests/test_examples.py

@@ -234,7 +234,13 @@ add_example_test(
     devices=test_devices,
     test_options={"height": 512, "width": 1024, "headless": True},
 )
-add_example_test(TestCoreExamples, name="core.example_sph", devices=test_devices, test_options_cpu={"num_frames": 1})
+add_example_test(
+    TestCoreExamples,
+    name="core.example_sph",
+    devices=test_devices,
+    test_options_cpu={"num_frames": 1},
+    test_options_cuda={"test_timeout": 600},
+)
 add_example_test(
     TestCoreExamples,
     name="core.example_torch",