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

warp/sim/collide.py

@@ -16,6 +16,7 @@ from warp.sim.model import Model
 
 from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry
 
+# types of triangle's closest point to a point
 TRI_CONTACT_FEATURE_VERTEX_A = wp.constant(0)
 TRI_CONTACT_FEATURE_VERTEX_B = wp.constant(1)
 TRI_CONTACT_FEATURE_VERTEX_C = wp.constant(2)
@@ -24,6 +25,45 @@ TRI_CONTACT_FEATURE_EDGE_AC = wp.constant(4)
 TRI_CONTACT_FEATURE_EDGE_BC = wp.constant(5)
 TRI_CONTACT_FEATURE_FACE_INTERIOR = wp.constant(6)
 
+# constants used to access TriMeshCollisionDetector.resize_flags
+VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(0)
+TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(1)
+EDGE_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(2)
+TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(3)
+
+
+@wp.func
+def build_orthonormal_basis(n: wp.vec3):
+    """
+    Builds an orthonormal basis given a normal vector `n`. Return the two axes that are perpendicular to `n`.
+
+    :param n: A 3D vector (list or array-like) representing the normal vector
+    """
+    b1 = wp.vec3()
+    b2 = wp.vec3()
+    if n[2] < 0.0:
+        a = 1.0 / (1.0 - n[2])
+        b = n[0] * n[1] * a
+        b1[0] = 1.0 - n[0] * n[0] * a
+        b1[1] = -b
+        b1[2] = n[0]
+
+        b2[0] = b
+        b2[1] = n[1] * n[1] * a - 1.0
+        b2[2] = -n[1]
+    else:
+        a = 1.0 / (1.0 + n[2])
+        b = -n[0] * n[1] * a
+        b1[0] = 1.0 - n[0] * n[0] * a
+        b1[1] = b
+        b1[2] = -n[0]
+
+        b2[0] = b
+        b2[1] = 1.0 - n[1] * n[1] * a
+        b2[2] = -n[1]
+
+    return b1, b2
+
 
 @wp.func
 def triangle_closest_point_barycentric(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
@@ -797,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]
@@ -1527,9 +1568,9 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
         # generate soft contacts for particles and shapes except ground plane (last shape)
         if model.particle_count and model.shape_count > 1:
             if requires_grad:
-                model.soft_contact_body_pos = wp.clone(model.soft_contact_body_pos)
-                model.soft_contact_body_vel = wp.clone(model.soft_contact_body_vel)
-                model.soft_contact_normal = wp.clone(model.soft_contact_normal)
+                model.soft_contact_body_pos = wp.empty_like(model.soft_contact_body_pos)
+                model.soft_contact_body_vel = wp.empty_like(model.soft_contact_body_vel)
+                model.soft_contact_normal = wp.empty_like(model.soft_contact_normal)
             # clear old count
             model.soft_contact_count.zero_()
             wp.launch(
@@ -1626,12 +1667,12 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
 
         if model.shape_contact_pair_count or model.ground and model.shape_ground_contact_pair_count:
             if requires_grad:
-                model.rigid_contact_point0 = wp.clone(model.rigid_contact_point0)
-                model.rigid_contact_point1 = wp.clone(model.rigid_contact_point1)
-                model.rigid_contact_offset0 = wp.clone(model.rigid_contact_offset0)
-                model.rigid_contact_offset1 = wp.clone(model.rigid_contact_offset1)
-                model.rigid_contact_normal = wp.clone(model.rigid_contact_normal)
-                model.rigid_contact_thickness = wp.clone(model.rigid_contact_thickness)
+                model.rigid_contact_point0 = wp.empty_like(model.rigid_contact_point0)
+                model.rigid_contact_point1 = wp.empty_like(model.rigid_contact_point1)
+                model.rigid_contact_offset0 = wp.empty_like(model.rigid_contact_offset0)
+                model.rigid_contact_offset1 = wp.empty_like(model.rigid_contact_offset1)
+                model.rigid_contact_normal = wp.empty_like(model.rigid_contact_normal)
+                model.rigid_contact_thickness = wp.empty_like(model.rigid_contact_thickness)
                 model.rigid_contact_count = wp.zeros_like(model.rigid_contact_count)
                 model.rigid_contact_tids = wp.zeros_like(model.rigid_contact_tids)
                 model.rigid_contact_shape0 = wp.empty_like(model.rigid_contact_shape0)
@@ -1679,6 +1720,18 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
             )
 
 
+@wp.func
+def compute_tri_aabb(
+    v1: wp.vec3,
+    v2: wp.vec3,
+    v3: wp.vec3,
+):
+    lower = wp.min(wp.min(v1, v2), v3)
+    upper = wp.max(wp.max(v1, v2), v3)
+
+    return lower, upper
+
+
 @wp.kernel
 def compute_tri_aabbs(
     pos: wp.array(dtype=wp.vec3),
@@ -1692,8 +1745,10 @@ def compute_tri_aabbs(
     v2 = pos[tri_indices[t_id, 1]]
     v3 = pos[tri_indices[t_id, 2]]
 
-    lower_bounds[t_id] = wp.min(wp.min(v1, v2), v3)
-    upper_bounds[t_id] = wp.max(wp.max(v1, v2), v3)
+    lower, upper = compute_tri_aabb(v1, v2, v3)
+
+    lower_bounds[t_id] = lower
+    upper_bounds[t_id] = upper
 
 
 @wp.kernel
@@ -1755,11 +1810,11 @@ 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),
-    vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
+    vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32),
     triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32),
     triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32),
     # outputs
@@ -1786,13 +1841,13 @@ 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
         vertex_colliding_triangles_count (array): number of triangles each vertex collides
         vertex_colliding_triangles_offsets (array): where each vertex' collision buffer starts
-        vertex_colliding_triangles_buffer_size (array): size of each vertex' collision buffer, will be modified if resizing is needed
+        vertex_colliding_triangles_buffer_sizes (array): size of each vertex' collision buffer, will be modified if resizing is needed
         vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
         triangle_colliding_vertices (array): positions of all the triangles' collision vertices
         triangle_colliding_vertices_count (array): number of triangles each vertex collides
@@ -1808,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]
@@ -1837,7 +1892,7 @@ def vertex_triangle_collision_detection_kernel(
             if vertex_num_collisions < vertex_buffer_size:
                 vertex_colliding_triangles[vertex_buffer_offset + vertex_num_collisions] = tri_index
             else:
-                resize_flags[0] = 1
+                resize_flags[VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
 
             vertex_num_collisions = vertex_num_collisions + 1
 
@@ -1850,14 +1905,14 @@ def vertex_triangle_collision_detection_kernel(
                 # record v-f collision to triangle
                 triangle_colliding_vertices[tri_buffer_offset + tri_num_collisions] = v_index
             else:
-                resize_flags[1] = 1
+                resize_flags[TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
 
     vertex_colliding_triangles_count[v_index] = vertex_num_collisions
     vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
 
 
 @wp.kernel
-def edge_edge_collision_detection_kernel(
+def edge_colliding_edges_detection_kernel(
     query_radius: float,
     bvh_id: wp.uint64,
     pos: wp.array(dtype=wp.vec3),
@@ -1876,7 +1931,7 @@ def edge_edge_collision_detection_kernel(
     query_radius (float):
     pos (array): positions of all the vertices that make up edges
     edge_colliding_triangles (array): flattened buffer of edges' collision edges
-    edge_colliding_edges_counts (array): number of edges each edge collides
+    edge_colliding_edges_count (array): number of edges each edge collides
     edge_colliding_triangles_offsets (array): where each edge's collision buffer starts
     edge_colliding_triangles_buffer_size (array): size of each edge's collision buffer, will be modified if resizing is needed
     edge_min_dis_to_triangles (array): each vertex' min distance to all (non-neighbor) triangles
@@ -1927,7 +1982,7 @@ def edge_edge_collision_detection_kernel(
             if edge_num_collisions < edge_buffer_size:
                 edge_colliding_edges[edge_buffer_offset + edge_num_collisions] = colliding_edge_index
             else:
-                resize_flags[1] = 1
+                resize_flags[EDGE_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
 
             edge_num_collisions = edge_num_collisions + 1
 
@@ -1935,19 +1990,134 @@ def edge_edge_collision_detection_kernel(
     edge_colliding_edges_min_dist[e_index] = min_dis_to_edges
 
 
+@wp.kernel
+def triangle_triangle_collision_detection_kernel(
+    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),
+    # outputs
+    triangle_intersecting_triangles: wp.array(dtype=wp.int32),
+    triangle_intersecting_triangles_count: wp.array(dtype=wp.int32),
+    resize_flags: wp.array(dtype=wp.int32),
+):
+    tri_index = wp.tid()
+    t1_v1 = tri_indices[tri_index, 0]
+    t1_v2 = tri_indices[tri_index, 1]
+    t1_v3 = tri_indices[tri_index, 2]
+
+    v1 = pos[t1_v1]
+    v2 = pos[t1_v2]
+    v3 = pos[t1_v3]
+
+    lower, upper = compute_tri_aabb(v1, v2, v3)
+
+    buffer_offset = triangle_intersecting_triangles_offsets[tri_index]
+    buffer_size = triangle_intersecting_triangles_offsets[tri_index + 1] - buffer_offset
+
+    query = wp.bvh_query_aabb(bvh_id, lower, upper)
+    tri_index_2 = wp.int32(0)
+    intersection_count = wp.int32(0)
+    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]
+
+        # filter out intersection test with neighbor triangles
+        if (
+            vertex_adjacent_to_triangle(t1_v1, t2_v1, t2_v2, t2_v3)
+            or vertex_adjacent_to_triangle(t1_v2, t2_v1, t2_v2, t2_v3)
+            or vertex_adjacent_to_triangle(t1_v3, t2_v1, t2_v2, t2_v3)
+        ):
+            continue
+
+        u1 = pos[t2_v1]
+        u2 = pos[t2_v2]
+        u3 = pos[t2_v3]
+
+        if wp.intersect_tri_tri(v1, v2, v3, u1, u2, u3):
+            if intersection_count < buffer_size:
+                triangle_intersecting_triangles[buffer_offset + intersection_count] = tri_index_2
+            else:
+                resize_flags[TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
+            intersection_count = intersection_count + 1
+
+    triangle_intersecting_triangles_count[tri_index] = intersection_count
+
+
+@wp.struct
+class TriMeshCollisionInfo:
+    vertex_colliding_triangles: wp.array(dtype=wp.int32)
+    vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32)
+    vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32)
+    vertex_colliding_triangles_count: wp.array(dtype=wp.int32)
+    vertex_colliding_triangles_min_dist: wp.array(dtype=float)
+
+    triangle_colliding_vertices: wp.array(dtype=wp.int32)
+    triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32)
+    triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32)
+    triangle_colliding_vertices_count: wp.array(dtype=wp.int32)
+    triangle_colliding_vertices_min_dist: wp.array(dtype=float)
+
+    edge_colliding_edges: wp.array(dtype=wp.int32)
+    edge_colliding_edges_offsets: wp.array(dtype=wp.int32)
+    edge_colliding_edges_buffer_sizes: wp.array(dtype=wp.int32)
+    edge_colliding_edges_count: wp.array(dtype=wp.int32)
+    edge_colliding_edges_min_dist: wp.array(dtype=float)
+
+
+@wp.func
+def get_vertex_colliding_triangles_count(col_info: TriMeshCollisionInfo, v: int):
+    return wp.min(col_info.vertex_colliding_triangles_count[v], col_info.vertex_colliding_triangles_buffer_sizes[v])
+
+
+@wp.func
+def get_vertex_colliding_triangles(col_info: TriMeshCollisionInfo, v: int, i_collision: int):
+    offset = col_info.vertex_colliding_triangles_offsets[v]
+    return col_info.vertex_colliding_triangles[offset + i_collision]
+
+
+@wp.func
+def get_triangle_colliding_vertices_count(col_info: TriMeshCollisionInfo, tri: int):
+    return wp.min(
+        col_info.triangle_colliding_vertices_count[tri], col_info.triangle_colliding_vertices_buffer_sizes[tri]
+    )
+
+
+@wp.func
+def get_triangle_colliding_vertices(col_info: TriMeshCollisionInfo, tri: int, i_collision: int):
+    offset = col_info.triangle_colliding_vertices_offsets[tri]
+    return col_info.triangle_colliding_vertices[offset + i_collision]
+
+
+@wp.func
+def get_edge_colliding_edges_count(col_info: TriMeshCollisionInfo, e: int):
+    return wp.min(col_info.edge_colliding_edges_count[e], col_info.edge_colliding_edges_buffer_sizes[e])
+
+
+@wp.func
+def get_edge_colliding_edges(col_info: TriMeshCollisionInfo, e: int, i_collision: int):
+    offset = col_info.edge_colliding_edges_offsets[e]
+    return col_info.edge_colliding_edges[offset + i_collision]
+
+
 class TriMeshCollisionDetector:
     def __init__(
         self,
         model: Model,
+        vertex_positions=None,
         vertex_collision_buffer_pre_alloc=8,
         vertex_collision_buffer_max_alloc=256,
         triangle_collision_buffer_pre_alloc=16,
         triangle_collision_buffer_max_alloc=256,
         edge_collision_buffer_pre_alloc=8,
         edge_collision_buffer_max_alloc=256,
+        triangle_triangle_collision_buffer_pre_alloc=8,
+        triangle_triangle_collision_buffer_max_alloc=256,
         edge_edge_parallel_epsilon=1e-5,
     ):
         self.model = model
+        self.vertex_positions = model.particle_q if vertex_positions is None else vertex_positions
         self.device = model.device
         self.vertex_collision_buffer_pre_alloc = vertex_collision_buffer_pre_alloc
         self.vertex_collision_buffer_max_alloc = vertex_collision_buffer_max_alloc
@@ -1955,10 +2125,21 @@ class TriMeshCollisionDetector:
         self.triangle_collision_buffer_max_alloc = triangle_collision_buffer_max_alloc
         self.edge_collision_buffer_pre_alloc = edge_collision_buffer_pre_alloc
         self.edge_collision_buffer_max_alloc = edge_collision_buffer_max_alloc
+        self.triangle_triangle_collision_buffer_pre_alloc = triangle_triangle_collision_buffer_pre_alloc
+        self.triangle_triangle_collision_buffer_max_alloc = triangle_triangle_collision_buffer_max_alloc
 
         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
 
@@ -2021,14 +2202,44 @@ class TriMeshCollisionDetector:
         self.upper_bounds_edges = wp.array(shape=(model.edge_count,), dtype=wp.vec3, device=model.device)
         wp.launch(
             kernel=compute_edge_aabbs,
-            inputs=[model.particle_q, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
+            inputs=[self.vertex_positions, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
             dim=model.edge_count,
             device=model.device,
         )
 
         self.bvh_edges = wp.Bvh(self.lower_bounds_edges, self.upper_bounds_edges)
 
-        self.resize_flags = wp.zeros(shape=(3,), dtype=wp.int32, device=self.device)
+        self.resize_flags = wp.zeros(shape=(4,), dtype=wp.int32, device=self.device)
+
+        self.collision_info = self.get_collision_data()
+
+        # data for triangle-triangle intersection; they will only be initialized on demand, as triangle-triangle intersection is not needed for simulation
+        self.triangle_intersecting_triangles = None
+        self.triangle_intersecting_triangles_count = None
+        self.triangle_intersecting_triangles_offsets = None
+
+    def get_collision_data(self):
+        collision_info = TriMeshCollisionInfo()
+
+        collision_info.vertex_colliding_triangles = self.vertex_colliding_triangles
+        collision_info.vertex_colliding_triangles_offsets = self.vertex_colliding_triangles_offsets
+        collision_info.vertex_colliding_triangles_buffer_sizes = self.vertex_colliding_triangles_buffer_sizes
+        collision_info.vertex_colliding_triangles_count = self.vertex_colliding_triangles_count
+        collision_info.vertex_colliding_triangles_min_dist = self.vertex_colliding_triangles_min_dist
+
+        collision_info.triangle_colliding_vertices = self.triangle_colliding_vertices
+        collision_info.triangle_colliding_vertices_offsets = self.triangle_colliding_vertices_offsets
+        collision_info.triangle_colliding_vertices_buffer_sizes = self.triangle_colliding_vertices_buffer_sizes
+        collision_info.triangle_colliding_vertices_count = self.triangle_colliding_vertices_count
+        collision_info.triangle_colliding_vertices_min_dist = self.triangle_colliding_vertices_min_dist
+
+        collision_info.edge_colliding_edges = self.edge_colliding_edges
+        collision_info.edge_colliding_edges_offsets = self.edge_colliding_edges_offsets
+        collision_info.edge_colliding_edges_buffer_sizes = self.edge_colliding_edges_buffer_sizes
+        collision_info.edge_colliding_edges_count = self.edge_colliding_edges_count
+        collision_info.edge_colliding_edges_min_dist = self.edge_colliding_edges_min_dist
+
+        return collision_info
 
     def compute_collision_buffer_offsets(
         self, buffer_sizes: wp.array(dtype=wp.int32), offsets: wp.array(dtype=wp.int32)
@@ -2040,17 +2251,26 @@ class TriMeshCollisionDetector:
 
         offsets.assign(offsets_np)
 
-    def refit(self):
+    def refit(self, new_pos=None):
+        if new_pos is not None:
+            self.vertex_positions = new_pos
+
         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(
             kernel=compute_edge_aabbs,
-            inputs=[self.model.particle_q, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
+            inputs=[self.vertex_positions, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
             dim=self.model.edge_count,
             device=self.model.device,
         )
@@ -2075,8 +2295,8 @@ class TriMeshCollisionDetector:
             kernel=vertex_triangle_collision_detection_kernel,
             inputs=[
                 query_radius,
-                self.mesh_tris.id,
-                self.model.particle_q,
+                self.bvh_tris.id,
+                self.vertex_positions,
                 self.model.tri_indices,
                 self.vertex_colliding_triangles_offsets,
                 self.vertex_colliding_triangles_buffer_sizes,
@@ -2098,11 +2318,11 @@ class TriMeshCollisionDetector:
 
     def edge_edge_collision_detection(self, query_radius):
         wp.launch(
-            kernel=edge_edge_collision_detection_kernel,
+            kernel=edge_colliding_edges_detection_kernel,
             inputs=[
                 query_radius,
                 self.bvh_edges.id,
-                self.model.particle_q,
+                self.vertex_positions,
                 self.model.edge_indices,
                 self.edge_colliding_edges_offsets,
                 self.edge_colliding_edges_buffer_sizes,
@@ -2117,3 +2337,40 @@ class TriMeshCollisionDetector:
             dim=self.model.edge_count,
             device=self.model.device,
         )
+
+    def triangle_triangle_intersection_detection(self):
+        if self.triangle_intersecting_triangles is None:
+            self.triangle_intersecting_triangles = wp.zeros(
+                shape=(self.model.tri_count * self.triangle_triangle_collision_buffer_pre_alloc,),
+                dtype=wp.int32,
+                device=self.device,
+            )
+
+        if self.triangle_intersecting_triangles_count is None:
+            self.triangle_intersecting_triangles_count = wp.array(
+                shape=(self.model.tri_count,), dtype=wp.int32, device=self.device
+            )
+
+        if self.triangle_intersecting_triangles_offsets is None:
+            buffer_sizes = np.full((self.model.tri_count,), self.triangle_triangle_collision_buffer_pre_alloc)
+            offsets = np.zeros((self.model.tri_count + 1,), dtype=np.int32)
+            offsets[1:] = np.cumsum(buffer_sizes)
+
+            self.triangle_intersecting_triangles_offsets = wp.array(offsets, dtype=wp.int32, device=self.device)
+
+        wp.launch(
+            kernel=triangle_triangle_collision_detection_kernel,
+            inputs=[
+                self.bvh_tris.id,
+                self.vertex_positions,
+                self.model.tri_indices,
+                self.triangle_intersecting_triangles_offsets,
+            ],
+            outputs=[
+                self.triangle_intersecting_triangles,
+                self.triangle_intersecting_triangles_count,
+                self.resize_flags,
+            ],
+            dim=self.model.tri_count,
+            device=self.model.device,
+        )

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_euler.py

@@ -871,6 +871,7 @@ def eval_rigid_contacts(
     contact_shape0: wp.array(dtype=int),
     contact_shape1: wp.array(dtype=int),
     force_in_world_frame: bool,
+    friction_smoothing: float,
     # outputs
     body_f: wp.array(dtype=wp.spatial_vector),
 ):
@@ -924,6 +925,8 @@ def eval_rigid_contacts(
     n = contact_normal[tid]
     bx_a = contact_point0[tid]
     bx_b = contact_point1[tid]
+    r_a = wp.vec3(0.0)
+    r_b = wp.vec3(0.0)
     if body_a >= 0:
         X_wb_a = body_q[body_a]
         X_com_a = body_com[body_a]
@@ -990,12 +993,16 @@ def eval_rigid_contacts(
     # ft = wp.vec3(vx, 0.0, vz)
 
     # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
-    # ft = wp.normalize(vt)*wp.min(kf*wp.length(vt), abs(mu*d*ke))
     ft = wp.vec3(0.0)
     if d < 0.0:
-        ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), -mu * (fn + fd))
+        # use a smooth vector norm to avoid gradient instability at/around zero velocity
+        vs = wp.norm_huber(vt, delta=friction_smoothing)
+        if vs > 0.0:
+            fr = vt / vs
+            ft = fr * wp.min(kf * vs, -mu * (fn + fd))
 
     f_total = n * (fn + fd) + ft
+    # f_total = n * (fn + fd)
     # f_total = n * fn
 
     if body_a >= 0:
@@ -1761,7 +1768,7 @@ def eval_tetrahedral_forces(model: Model, state: State, control: Control, partic
         )
 
 
-def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
+def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array, friction_smoothing: float = 1.0):
     if model.rigid_contact_max and (
         model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
     ):
@@ -1782,6 +1789,7 @@ def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
                 model.rigid_contact_shape0,
                 model.rigid_contact_shape1,
                 False,
+                friction_smoothing,
             ],
             outputs=[state.body_f],
             device=model.device,
@@ -1880,7 +1888,15 @@ def eval_muscle_forces(model: Model, state: State, control: Control, body_f: wp.
         )
 
 
-def compute_forces(model: Model, state: State, control: Control, particle_f: wp.array, body_f: wp.array, dt: float):
+def compute_forces(
+    model: Model,
+    state: State,
+    control: Control,
+    particle_f: wp.array,
+    body_f: wp.array,
+    dt: float,
+    friction_smoothing: float = 1.0,
+):
     # damped springs
     eval_spring_forces(model, state, particle_f)
 
@@ -1906,7 +1922,7 @@ def compute_forces(model: Model, state: State, control: Control, particle_f: wp.
     eval_particle_ground_contact_forces(model, state, particle_f)
 
     # body contacts
-    eval_body_contact_forces(model, state, particle_f)
+    eval_body_contact_forces(model, state, particle_f, friction_smoothing=friction_smoothing)
 
     # particle shape contact
     eval_particle_body_contact_forces(model, state, particle_f, body_f, body_f_in_world_frame=False)
@@ -1941,12 +1957,14 @@ class SemiImplicitIntegrator(Integrator):
 
     """
 
-    def __init__(self, angular_damping: float = 0.05):
+    def __init__(self, angular_damping: float = 0.05, friction_smoothing: float = 1.0):
         """
         Args:
             angular_damping (float, optional): Angular damping factor. Defaults to 0.05.
+            friction_smoothing (float, optional): The delta value for the Huber norm (see :func:`warp.math.norm_huber`) used for the friction velocity normalization. Defaults to 1.0.
         """
         self.angular_damping = angular_damping
+        self.friction_smoothing = friction_smoothing
 
     def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
         with wp.ScopedTimer("simulate", False):
@@ -1962,7 +1980,7 @@ class SemiImplicitIntegrator(Integrator):
             if control is None:
                 control = model.control(clone_variables=False)
 
-            compute_forces(model, state_in, control, particle_f, body_f, dt)
+            compute_forces(model, state_in, control, particle_f, body_f, dt, friction_smoothing=self.friction_smoothing)
 
             self.integrate_bodies(model, state_in, state_out, dt, self.angular_damping)