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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (131) hide show
  1. warp/__init__.py +5 -0
  2. warp/autograd.py +414 -191
  3. warp/bin/libwarp-clang.dylib +0 -0
  4. warp/bin/libwarp.dylib +0 -0
  5. warp/build.py +40 -12
  6. warp/build_dll.py +13 -6
  7. warp/builtins.py +1077 -481
  8. warp/codegen.py +250 -122
  9. warp/config.py +65 -21
  10. warp/context.py +500 -149
  11. warp/examples/assets/square_cloth.usd +0 -0
  12. warp/examples/benchmarks/benchmark_gemm.py +27 -18
  13. warp/examples/benchmarks/benchmark_interop_paddle.py +3 -3
  14. warp/examples/benchmarks/benchmark_interop_torch.py +3 -3
  15. warp/examples/core/example_marching_cubes.py +1 -1
  16. warp/examples/core/example_mesh.py +1 -1
  17. warp/examples/core/example_torch.py +18 -34
  18. warp/examples/core/example_wave.py +1 -1
  19. warp/examples/fem/example_apic_fluid.py +1 -0
  20. warp/examples/fem/example_mixed_elasticity.py +1 -1
  21. warp/examples/optim/example_bounce.py +1 -1
  22. warp/examples/optim/example_cloth_throw.py +1 -1
  23. warp/examples/optim/example_diffray.py +4 -15
  24. warp/examples/optim/example_drone.py +1 -1
  25. warp/examples/optim/example_softbody_properties.py +392 -0
  26. warp/examples/optim/example_trajectory.py +1 -3
  27. warp/examples/optim/example_walker.py +5 -0
  28. warp/examples/sim/example_cartpole.py +0 -2
  29. warp/examples/sim/example_cloth_self_contact.py +314 -0
  30. warp/examples/sim/example_granular_collision_sdf.py +4 -5
  31. warp/examples/sim/example_jacobian_ik.py +0 -2
  32. warp/examples/sim/example_quadruped.py +5 -2
  33. warp/examples/tile/example_tile_cholesky.py +79 -0
  34. warp/examples/tile/example_tile_convolution.py +2 -2
  35. warp/examples/tile/example_tile_fft.py +2 -2
  36. warp/examples/tile/example_tile_filtering.py +3 -3
  37. warp/examples/tile/example_tile_matmul.py +4 -4
  38. warp/examples/tile/example_tile_mlp.py +12 -12
  39. warp/examples/tile/example_tile_nbody.py +191 -0
  40. warp/examples/tile/example_tile_walker.py +319 -0
  41. warp/math.py +147 -0
  42. warp/native/array.h +12 -0
  43. warp/native/builtin.h +0 -1
  44. warp/native/bvh.cpp +149 -70
  45. warp/native/bvh.cu +287 -68
  46. warp/native/bvh.h +195 -85
  47. warp/native/clang/clang.cpp +6 -2
  48. warp/native/crt.h +1 -0
  49. warp/native/cuda_util.cpp +35 -0
  50. warp/native/cuda_util.h +5 -0
  51. warp/native/exports.h +40 -40
  52. warp/native/intersect.h +17 -0
  53. warp/native/mat.h +57 -3
  54. warp/native/mathdx.cpp +19 -0
  55. warp/native/mesh.cpp +25 -8
  56. warp/native/mesh.cu +153 -101
  57. warp/native/mesh.h +482 -403
  58. warp/native/quat.h +40 -0
  59. warp/native/solid_angle.h +7 -0
  60. warp/native/sort.cpp +85 -0
  61. warp/native/sort.cu +34 -0
  62. warp/native/sort.h +3 -1
  63. warp/native/spatial.h +11 -0
  64. warp/native/tile.h +1189 -664
  65. warp/native/tile_reduce.h +8 -6
  66. warp/native/vec.h +41 -0
  67. warp/native/warp.cpp +8 -1
  68. warp/native/warp.cu +263 -40
  69. warp/native/warp.h +19 -5
  70. warp/optim/linear.py +22 -4
  71. warp/render/render_opengl.py +132 -59
  72. warp/render/render_usd.py +10 -2
  73. warp/sim/__init__.py +6 -1
  74. warp/sim/collide.py +289 -32
  75. warp/sim/import_urdf.py +20 -5
  76. warp/sim/integrator_euler.py +25 -7
  77. warp/sim/integrator_featherstone.py +147 -35
  78. warp/sim/integrator_vbd.py +842 -40
  79. warp/sim/model.py +173 -112
  80. warp/sim/render.py +2 -2
  81. warp/stubs.py +249 -116
  82. warp/tape.py +28 -30
  83. warp/tests/aux_test_module_unload.py +15 -0
  84. warp/tests/{test_sim_grad.py → flaky_test_sim_grad.py} +104 -63
  85. warp/tests/test_array.py +100 -0
  86. warp/tests/test_assert.py +242 -0
  87. warp/tests/test_codegen.py +14 -61
  88. warp/tests/test_collision.py +8 -8
  89. warp/tests/test_examples.py +16 -1
  90. warp/tests/test_grad_debug.py +87 -2
  91. warp/tests/test_hash_grid.py +1 -1
  92. warp/tests/test_ipc.py +116 -0
  93. warp/tests/test_launch.py +77 -26
  94. warp/tests/test_mat.py +213 -168
  95. warp/tests/test_math.py +47 -1
  96. warp/tests/test_matmul.py +11 -7
  97. warp/tests/test_matmul_lite.py +4 -4
  98. warp/tests/test_mesh.py +84 -60
  99. warp/tests/test_mesh_query_aabb.py +165 -0
  100. warp/tests/test_mesh_query_point.py +328 -286
  101. warp/tests/test_mesh_query_ray.py +134 -121
  102. warp/tests/test_mlp.py +2 -2
  103. warp/tests/test_operators.py +43 -0
  104. warp/tests/test_overwrite.py +6 -5
  105. warp/tests/test_quat.py +77 -0
  106. warp/tests/test_reload.py +29 -0
  107. warp/tests/test_sim_grad_bounce_linear.py +204 -0
  108. warp/tests/test_static.py +16 -0
  109. warp/tests/test_tape.py +25 -0
  110. warp/tests/test_tile.py +134 -191
  111. warp/tests/test_tile_load.py +399 -0
  112. warp/tests/test_tile_mathdx.py +61 -8
  113. warp/tests/test_tile_mlp.py +17 -17
  114. warp/tests/test_tile_reduce.py +24 -18
  115. warp/tests/test_tile_shared_memory.py +66 -17
  116. warp/tests/test_tile_view.py +165 -0
  117. warp/tests/test_torch.py +35 -0
  118. warp/tests/test_utils.py +36 -24
  119. warp/tests/test_vec.py +110 -0
  120. warp/tests/unittest_suites.py +29 -4
  121. warp/tests/unittest_utils.py +30 -11
  122. warp/thirdparty/unittest_parallel.py +5 -2
  123. warp/types.py +419 -111
  124. warp/utils.py +9 -5
  125. {warp_lang-1.5.1.dist-info → warp_lang-1.6.1.dist-info}/METADATA +86 -45
  126. {warp_lang-1.5.1.dist-info → warp_lang-1.6.1.dist-info}/RECORD +129 -118
  127. {warp_lang-1.5.1.dist-info → warp_lang-1.6.1.dist-info}/WHEEL +1 -1
  128. warp/examples/benchmarks/benchmark_tile.py +0 -179
  129. warp/native/tile_gemm.h +0 -341
  130. {warp_lang-1.5.1.dist-info → warp_lang-1.6.1.dist-info}/LICENSE.md +0 -0
  131. {warp_lang-1.5.1.dist-info → warp_lang-1.6.1.dist-info}/top_level.txt +0 -0
warp/sim/collide.py CHANGED
@@ -16,6 +16,7 @@ from warp.sim.model import Model
16
16
 
17
17
  from .model import PARTICLE_FLAG_ACTIVE, ModelShapeGeometry
18
18
 
19
+ # types of triangle's closest point to a point
19
20
  TRI_CONTACT_FEATURE_VERTEX_A = wp.constant(0)
20
21
  TRI_CONTACT_FEATURE_VERTEX_B = wp.constant(1)
21
22
  TRI_CONTACT_FEATURE_VERTEX_C = wp.constant(2)
@@ -24,6 +25,45 @@ TRI_CONTACT_FEATURE_EDGE_AC = wp.constant(4)
24
25
  TRI_CONTACT_FEATURE_EDGE_BC = wp.constant(5)
25
26
  TRI_CONTACT_FEATURE_FACE_INTERIOR = wp.constant(6)
26
27
 
28
+ # constants used to access TriMeshCollisionDetector.resize_flags
29
+ VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(0)
30
+ TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(1)
31
+ EDGE_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(2)
32
+ TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX = wp.constant(3)
33
+
34
+
35
+ @wp.func
36
+ def build_orthonormal_basis(n: wp.vec3):
37
+ """
38
+ Builds an orthonormal basis given a normal vector `n`. Return the two axes that are perpendicular to `n`.
39
+
40
+ :param n: A 3D vector (list or array-like) representing the normal vector
41
+ """
42
+ b1 = wp.vec3()
43
+ b2 = wp.vec3()
44
+ if n[2] < 0.0:
45
+ a = 1.0 / (1.0 - n[2])
46
+ b = n[0] * n[1] * a
47
+ b1[0] = 1.0 - n[0] * n[0] * a
48
+ b1[1] = -b
49
+ b1[2] = n[0]
50
+
51
+ b2[0] = b
52
+ b2[1] = n[1] * n[1] * a - 1.0
53
+ b2[2] = -n[1]
54
+ else:
55
+ a = 1.0 / (1.0 + n[2])
56
+ b = -n[0] * n[1] * a
57
+ b1[0] = 1.0 - n[0] * n[0] * a
58
+ b1[1] = b
59
+ b1[2] = -n[0]
60
+
61
+ b2[0] = b
62
+ b2[1] = 1.0 - n[1] * n[1] * a
63
+ b2[2] = -n[1]
64
+
65
+ return b1, b2
66
+
27
67
 
28
68
  @wp.func
29
69
  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(
797
837
  num_actual_contacts = 8 + 4
798
838
  else:
799
839
  num_contacts = 8
840
+ num_actual_contacts = 8
800
841
  elif actual_type_a == wp.sim.GEO_MESH:
801
842
  mesh_a = wp.mesh_get(geo.source[actual_shape_a])
802
843
  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 =
1527
1568
  # generate soft contacts for particles and shapes except ground plane (last shape)
1528
1569
  if model.particle_count and model.shape_count > 1:
1529
1570
  if requires_grad:
1530
- model.soft_contact_body_pos = wp.clone(model.soft_contact_body_pos)
1531
- model.soft_contact_body_vel = wp.clone(model.soft_contact_body_vel)
1532
- model.soft_contact_normal = wp.clone(model.soft_contact_normal)
1571
+ model.soft_contact_body_pos = wp.empty_like(model.soft_contact_body_pos)
1572
+ model.soft_contact_body_vel = wp.empty_like(model.soft_contact_body_vel)
1573
+ model.soft_contact_normal = wp.empty_like(model.soft_contact_normal)
1533
1574
  # clear old count
1534
1575
  model.soft_contact_count.zero_()
1535
1576
  wp.launch(
@@ -1626,12 +1667,12 @@ def collide(model, state, edge_sdf_iter: int = 10, iterate_mesh_vertices: bool =
1626
1667
 
1627
1668
  if model.shape_contact_pair_count or model.ground and model.shape_ground_contact_pair_count:
1628
1669
  if requires_grad:
1629
- model.rigid_contact_point0 = wp.clone(model.rigid_contact_point0)
1630
- model.rigid_contact_point1 = wp.clone(model.rigid_contact_point1)
1631
- model.rigid_contact_offset0 = wp.clone(model.rigid_contact_offset0)
1632
- model.rigid_contact_offset1 = wp.clone(model.rigid_contact_offset1)
1633
- model.rigid_contact_normal = wp.clone(model.rigid_contact_normal)
1634
- model.rigid_contact_thickness = wp.clone(model.rigid_contact_thickness)
1670
+ model.rigid_contact_point0 = wp.empty_like(model.rigid_contact_point0)
1671
+ model.rigid_contact_point1 = wp.empty_like(model.rigid_contact_point1)
1672
+ model.rigid_contact_offset0 = wp.empty_like(model.rigid_contact_offset0)
1673
+ model.rigid_contact_offset1 = wp.empty_like(model.rigid_contact_offset1)
1674
+ model.rigid_contact_normal = wp.empty_like(model.rigid_contact_normal)
1675
+ model.rigid_contact_thickness = wp.empty_like(model.rigid_contact_thickness)
1635
1676
  model.rigid_contact_count = wp.zeros_like(model.rigid_contact_count)
1636
1677
  model.rigid_contact_tids = wp.zeros_like(model.rigid_contact_tids)
1637
1678
  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 =
1679
1720
  )
1680
1721
 
1681
1722
 
1723
+ @wp.func
1724
+ def compute_tri_aabb(
1725
+ v1: wp.vec3,
1726
+ v2: wp.vec3,
1727
+ v3: wp.vec3,
1728
+ ):
1729
+ lower = wp.min(wp.min(v1, v2), v3)
1730
+ upper = wp.max(wp.max(v1, v2), v3)
1731
+
1732
+ return lower, upper
1733
+
1734
+
1682
1735
  @wp.kernel
1683
1736
  def compute_tri_aabbs(
1684
1737
  pos: wp.array(dtype=wp.vec3),
@@ -1692,8 +1745,10 @@ def compute_tri_aabbs(
1692
1745
  v2 = pos[tri_indices[t_id, 1]]
1693
1746
  v3 = pos[tri_indices[t_id, 2]]
1694
1747
 
1695
- lower_bounds[t_id] = wp.min(wp.min(v1, v2), v3)
1696
- upper_bounds[t_id] = wp.max(wp.max(v1, v2), v3)
1748
+ lower, upper = compute_tri_aabb(v1, v2, v3)
1749
+
1750
+ lower_bounds[t_id] = lower
1751
+ upper_bounds[t_id] = upper
1697
1752
 
1698
1753
 
1699
1754
  @wp.kernel
@@ -1755,11 +1810,11 @@ def init_triangle_collision_data_kernel(
1755
1810
  @wp.kernel
1756
1811
  def vertex_triangle_collision_detection_kernel(
1757
1812
  query_radius: float,
1758
- mesh_id: wp.uint64,
1813
+ bvh_id: wp.uint64,
1759
1814
  pos: wp.array(dtype=wp.vec3),
1760
1815
  tri_indices: wp.array(dtype=wp.int32, ndim=2),
1761
1816
  vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32),
1762
- vertex_colliding_triangles_buffer_size: wp.array(dtype=wp.int32),
1817
+ vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32),
1763
1818
  triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32),
1764
1819
  triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32),
1765
1820
  # outputs
@@ -1786,13 +1841,13 @@ def vertex_triangle_collision_detection_kernel(
1786
1841
  and vertex_colliding_triangles_count.
1787
1842
 
1788
1843
  Attributes:
1789
- mesh_id (int): the mesh id you want to collide with
1844
+ bvh_id (int): the bvh id you want to collide with
1790
1845
  query_radius (float): the contact radius. vertex-triangle pairs whose distance are less than this will get detected
1791
1846
  pos (array): positions of all the vertices that make up triangles
1792
1847
  vertex_colliding_triangles (array): flattened buffer of vertices' collision triangles
1793
1848
  vertex_colliding_triangles_count (array): number of triangles each vertex collides
1794
1849
  vertex_colliding_triangles_offsets (array): where each vertex' collision buffer starts
1795
- vertex_colliding_triangles_buffer_size (array): size of each vertex' collision buffer, will be modified if resizing is needed
1850
+ vertex_colliding_triangles_buffer_sizes (array): size of each vertex' collision buffer, will be modified if resizing is needed
1796
1851
  vertex_colliding_triangles_min_dist (array): each vertex' min distance to all (non-neighbor) triangles
1797
1852
  triangle_colliding_vertices (array): positions of all the triangles' collision vertices
1798
1853
  triangle_colliding_vertices_count (array): number of triangles each vertex collides
@@ -1808,12 +1863,12 @@ def vertex_triangle_collision_detection_kernel(
1808
1863
  lower = wp.vec3(v[0] - query_radius, v[1] - query_radius, v[2] - query_radius)
1809
1864
  upper = wp.vec3(v[0] + query_radius, v[1] + query_radius, v[2] + query_radius)
1810
1865
 
1811
- query = wp.mesh_query_aabb(mesh_id, lower, upper)
1866
+ query = wp.bvh_query_aabb(bvh_id, lower, upper)
1812
1867
 
1813
1868
  tri_index = wp.int32(0)
1814
1869
  vertex_num_collisions = wp.int32(0)
1815
1870
  min_dis_to_tris = query_radius
1816
- while wp.mesh_query_aabb_next(query, tri_index):
1871
+ while wp.bvh_query_next(query, tri_index):
1817
1872
  t1 = tri_indices[tri_index, 0]
1818
1873
  t2 = tri_indices[tri_index, 1]
1819
1874
  t3 = tri_indices[tri_index, 2]
@@ -1837,7 +1892,7 @@ def vertex_triangle_collision_detection_kernel(
1837
1892
  if vertex_num_collisions < vertex_buffer_size:
1838
1893
  vertex_colliding_triangles[vertex_buffer_offset + vertex_num_collisions] = tri_index
1839
1894
  else:
1840
- resize_flags[0] = 1
1895
+ resize_flags[VERTEX_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1841
1896
 
1842
1897
  vertex_num_collisions = vertex_num_collisions + 1
1843
1898
 
@@ -1850,14 +1905,14 @@ def vertex_triangle_collision_detection_kernel(
1850
1905
  # record v-f collision to triangle
1851
1906
  triangle_colliding_vertices[tri_buffer_offset + tri_num_collisions] = v_index
1852
1907
  else:
1853
- resize_flags[1] = 1
1908
+ resize_flags[TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1854
1909
 
1855
1910
  vertex_colliding_triangles_count[v_index] = vertex_num_collisions
1856
1911
  vertex_colliding_triangles_min_dist[v_index] = min_dis_to_tris
1857
1912
 
1858
1913
 
1859
1914
  @wp.kernel
1860
- def edge_edge_collision_detection_kernel(
1915
+ def edge_colliding_edges_detection_kernel(
1861
1916
  query_radius: float,
1862
1917
  bvh_id: wp.uint64,
1863
1918
  pos: wp.array(dtype=wp.vec3),
@@ -1876,7 +1931,7 @@ def edge_edge_collision_detection_kernel(
1876
1931
  query_radius (float):
1877
1932
  pos (array): positions of all the vertices that make up edges
1878
1933
  edge_colliding_triangles (array): flattened buffer of edges' collision edges
1879
- edge_colliding_edges_counts (array): number of edges each edge collides
1934
+ edge_colliding_edges_count (array): number of edges each edge collides
1880
1935
  edge_colliding_triangles_offsets (array): where each edge's collision buffer starts
1881
1936
  edge_colliding_triangles_buffer_size (array): size of each edge's collision buffer, will be modified if resizing is needed
1882
1937
  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(
1927
1982
  if edge_num_collisions < edge_buffer_size:
1928
1983
  edge_colliding_edges[edge_buffer_offset + edge_num_collisions] = colliding_edge_index
1929
1984
  else:
1930
- resize_flags[1] = 1
1985
+ resize_flags[EDGE_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
1931
1986
 
1932
1987
  edge_num_collisions = edge_num_collisions + 1
1933
1988
 
@@ -1935,19 +1990,134 @@ def edge_edge_collision_detection_kernel(
1935
1990
  edge_colliding_edges_min_dist[e_index] = min_dis_to_edges
1936
1991
 
1937
1992
 
1993
+ @wp.kernel
1994
+ def triangle_triangle_collision_detection_kernel(
1995
+ bvh_id: wp.uint64,
1996
+ pos: wp.array(dtype=wp.vec3),
1997
+ tri_indices: wp.array(dtype=wp.int32, ndim=2),
1998
+ triangle_intersecting_triangles_offsets: wp.array(dtype=wp.int32),
1999
+ # outputs
2000
+ triangle_intersecting_triangles: wp.array(dtype=wp.int32),
2001
+ triangle_intersecting_triangles_count: wp.array(dtype=wp.int32),
2002
+ resize_flags: wp.array(dtype=wp.int32),
2003
+ ):
2004
+ tri_index = wp.tid()
2005
+ t1_v1 = tri_indices[tri_index, 0]
2006
+ t1_v2 = tri_indices[tri_index, 1]
2007
+ t1_v3 = tri_indices[tri_index, 2]
2008
+
2009
+ v1 = pos[t1_v1]
2010
+ v2 = pos[t1_v2]
2011
+ v3 = pos[t1_v3]
2012
+
2013
+ lower, upper = compute_tri_aabb(v1, v2, v3)
2014
+
2015
+ buffer_offset = triangle_intersecting_triangles_offsets[tri_index]
2016
+ buffer_size = triangle_intersecting_triangles_offsets[tri_index + 1] - buffer_offset
2017
+
2018
+ query = wp.bvh_query_aabb(bvh_id, lower, upper)
2019
+ tri_index_2 = wp.int32(0)
2020
+ intersection_count = wp.int32(0)
2021
+ while wp.bvh_query_next(query, tri_index_2):
2022
+ t2_v1 = tri_indices[tri_index_2, 0]
2023
+ t2_v2 = tri_indices[tri_index_2, 1]
2024
+ t2_v3 = tri_indices[tri_index_2, 2]
2025
+
2026
+ # filter out intersection test with neighbor triangles
2027
+ if (
2028
+ vertex_adjacent_to_triangle(t1_v1, t2_v1, t2_v2, t2_v3)
2029
+ or vertex_adjacent_to_triangle(t1_v2, t2_v1, t2_v2, t2_v3)
2030
+ or vertex_adjacent_to_triangle(t1_v3, t2_v1, t2_v2, t2_v3)
2031
+ ):
2032
+ continue
2033
+
2034
+ u1 = pos[t2_v1]
2035
+ u2 = pos[t2_v2]
2036
+ u3 = pos[t2_v3]
2037
+
2038
+ if wp.intersect_tri_tri(v1, v2, v3, u1, u2, u3):
2039
+ if intersection_count < buffer_size:
2040
+ triangle_intersecting_triangles[buffer_offset + intersection_count] = tri_index_2
2041
+ else:
2042
+ resize_flags[TRI_TRI_COLLISION_BUFFER_OVERFLOW_INDEX] = 1
2043
+ intersection_count = intersection_count + 1
2044
+
2045
+ triangle_intersecting_triangles_count[tri_index] = intersection_count
2046
+
2047
+
2048
+ @wp.struct
2049
+ class TriMeshCollisionInfo:
2050
+ vertex_colliding_triangles: wp.array(dtype=wp.int32)
2051
+ vertex_colliding_triangles_offsets: wp.array(dtype=wp.int32)
2052
+ vertex_colliding_triangles_buffer_sizes: wp.array(dtype=wp.int32)
2053
+ vertex_colliding_triangles_count: wp.array(dtype=wp.int32)
2054
+ vertex_colliding_triangles_min_dist: wp.array(dtype=float)
2055
+
2056
+ triangle_colliding_vertices: wp.array(dtype=wp.int32)
2057
+ triangle_colliding_vertices_offsets: wp.array(dtype=wp.int32)
2058
+ triangle_colliding_vertices_buffer_sizes: wp.array(dtype=wp.int32)
2059
+ triangle_colliding_vertices_count: wp.array(dtype=wp.int32)
2060
+ triangle_colliding_vertices_min_dist: wp.array(dtype=float)
2061
+
2062
+ edge_colliding_edges: wp.array(dtype=wp.int32)
2063
+ edge_colliding_edges_offsets: wp.array(dtype=wp.int32)
2064
+ edge_colliding_edges_buffer_sizes: wp.array(dtype=wp.int32)
2065
+ edge_colliding_edges_count: wp.array(dtype=wp.int32)
2066
+ edge_colliding_edges_min_dist: wp.array(dtype=float)
2067
+
2068
+
2069
+ @wp.func
2070
+ def get_vertex_colliding_triangles_count(col_info: TriMeshCollisionInfo, v: int):
2071
+ return wp.min(col_info.vertex_colliding_triangles_count[v], col_info.vertex_colliding_triangles_buffer_sizes[v])
2072
+
2073
+
2074
+ @wp.func
2075
+ def get_vertex_colliding_triangles(col_info: TriMeshCollisionInfo, v: int, i_collision: int):
2076
+ offset = col_info.vertex_colliding_triangles_offsets[v]
2077
+ return col_info.vertex_colliding_triangles[offset + i_collision]
2078
+
2079
+
2080
+ @wp.func
2081
+ def get_triangle_colliding_vertices_count(col_info: TriMeshCollisionInfo, tri: int):
2082
+ return wp.min(
2083
+ col_info.triangle_colliding_vertices_count[tri], col_info.triangle_colliding_vertices_buffer_sizes[tri]
2084
+ )
2085
+
2086
+
2087
+ @wp.func
2088
+ def get_triangle_colliding_vertices(col_info: TriMeshCollisionInfo, tri: int, i_collision: int):
2089
+ offset = col_info.triangle_colliding_vertices_offsets[tri]
2090
+ return col_info.triangle_colliding_vertices[offset + i_collision]
2091
+
2092
+
2093
+ @wp.func
2094
+ def get_edge_colliding_edges_count(col_info: TriMeshCollisionInfo, e: int):
2095
+ return wp.min(col_info.edge_colliding_edges_count[e], col_info.edge_colliding_edges_buffer_sizes[e])
2096
+
2097
+
2098
+ @wp.func
2099
+ def get_edge_colliding_edges(col_info: TriMeshCollisionInfo, e: int, i_collision: int):
2100
+ offset = col_info.edge_colliding_edges_offsets[e]
2101
+ return col_info.edge_colliding_edges[offset + i_collision]
2102
+
2103
+
1938
2104
  class TriMeshCollisionDetector:
1939
2105
  def __init__(
1940
2106
  self,
1941
2107
  model: Model,
2108
+ vertex_positions=None,
1942
2109
  vertex_collision_buffer_pre_alloc=8,
1943
2110
  vertex_collision_buffer_max_alloc=256,
1944
2111
  triangle_collision_buffer_pre_alloc=16,
1945
2112
  triangle_collision_buffer_max_alloc=256,
1946
2113
  edge_collision_buffer_pre_alloc=8,
1947
2114
  edge_collision_buffer_max_alloc=256,
2115
+ triangle_triangle_collision_buffer_pre_alloc=8,
2116
+ triangle_triangle_collision_buffer_max_alloc=256,
1948
2117
  edge_edge_parallel_epsilon=1e-5,
1949
2118
  ):
1950
2119
  self.model = model
2120
+ self.vertex_positions = model.particle_q if vertex_positions is None else vertex_positions
1951
2121
  self.device = model.device
1952
2122
  self.vertex_collision_buffer_pre_alloc = vertex_collision_buffer_pre_alloc
1953
2123
  self.vertex_collision_buffer_max_alloc = vertex_collision_buffer_max_alloc
@@ -1955,10 +2125,21 @@ class TriMeshCollisionDetector:
1955
2125
  self.triangle_collision_buffer_max_alloc = triangle_collision_buffer_max_alloc
1956
2126
  self.edge_collision_buffer_pre_alloc = edge_collision_buffer_pre_alloc
1957
2127
  self.edge_collision_buffer_max_alloc = edge_collision_buffer_max_alloc
2128
+ self.triangle_triangle_collision_buffer_pre_alloc = triangle_triangle_collision_buffer_pre_alloc
2129
+ self.triangle_triangle_collision_buffer_max_alloc = triangle_triangle_collision_buffer_max_alloc
1958
2130
 
1959
2131
  self.edge_edge_parallel_epsilon = edge_edge_parallel_epsilon
1960
2132
 
1961
- self.mesh_tris = wp.Mesh(model.particle_q, model.tri_indices.flatten())
2133
+ self.lower_bounds_tris = wp.array(shape=(model.tri_count,), dtype=wp.vec3, device=model.device)
2134
+ self.upper_bounds_tris = wp.array(shape=(model.tri_count,), dtype=wp.vec3, device=model.device)
2135
+ wp.launch(
2136
+ kernel=compute_tri_aabbs,
2137
+ inputs=[self.vertex_positions, model.tri_indices, self.lower_bounds_tris, self.upper_bounds_tris],
2138
+ dim=model.tri_count,
2139
+ device=model.device,
2140
+ )
2141
+
2142
+ self.bvh_tris = wp.Bvh(self.lower_bounds_tris, self.upper_bounds_tris)
1962
2143
 
1963
2144
  # collision detections results
1964
2145
 
@@ -2021,14 +2202,44 @@ class TriMeshCollisionDetector:
2021
2202
  self.upper_bounds_edges = wp.array(shape=(model.edge_count,), dtype=wp.vec3, device=model.device)
2022
2203
  wp.launch(
2023
2204
  kernel=compute_edge_aabbs,
2024
- inputs=[model.particle_q, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2205
+ inputs=[self.vertex_positions, model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2025
2206
  dim=model.edge_count,
2026
2207
  device=model.device,
2027
2208
  )
2028
2209
 
2029
2210
  self.bvh_edges = wp.Bvh(self.lower_bounds_edges, self.upper_bounds_edges)
2030
2211
 
2031
- self.resize_flags = wp.zeros(shape=(3,), dtype=wp.int32, device=self.device)
2212
+ self.resize_flags = wp.zeros(shape=(4,), dtype=wp.int32, device=self.device)
2213
+
2214
+ self.collision_info = self.get_collision_data()
2215
+
2216
+ # data for triangle-triangle intersection; they will only be initialized on demand, as triangle-triangle intersection is not needed for simulation
2217
+ self.triangle_intersecting_triangles = None
2218
+ self.triangle_intersecting_triangles_count = None
2219
+ self.triangle_intersecting_triangles_offsets = None
2220
+
2221
+ def get_collision_data(self):
2222
+ collision_info = TriMeshCollisionInfo()
2223
+
2224
+ collision_info.vertex_colliding_triangles = self.vertex_colliding_triangles
2225
+ collision_info.vertex_colliding_triangles_offsets = self.vertex_colliding_triangles_offsets
2226
+ collision_info.vertex_colliding_triangles_buffer_sizes = self.vertex_colliding_triangles_buffer_sizes
2227
+ collision_info.vertex_colliding_triangles_count = self.vertex_colliding_triangles_count
2228
+ collision_info.vertex_colliding_triangles_min_dist = self.vertex_colliding_triangles_min_dist
2229
+
2230
+ collision_info.triangle_colliding_vertices = self.triangle_colliding_vertices
2231
+ collision_info.triangle_colliding_vertices_offsets = self.triangle_colliding_vertices_offsets
2232
+ collision_info.triangle_colliding_vertices_buffer_sizes = self.triangle_colliding_vertices_buffer_sizes
2233
+ collision_info.triangle_colliding_vertices_count = self.triangle_colliding_vertices_count
2234
+ collision_info.triangle_colliding_vertices_min_dist = self.triangle_colliding_vertices_min_dist
2235
+
2236
+ collision_info.edge_colliding_edges = self.edge_colliding_edges
2237
+ collision_info.edge_colliding_edges_offsets = self.edge_colliding_edges_offsets
2238
+ collision_info.edge_colliding_edges_buffer_sizes = self.edge_colliding_edges_buffer_sizes
2239
+ collision_info.edge_colliding_edges_count = self.edge_colliding_edges_count
2240
+ collision_info.edge_colliding_edges_min_dist = self.edge_colliding_edges_min_dist
2241
+
2242
+ return collision_info
2032
2243
 
2033
2244
  def compute_collision_buffer_offsets(
2034
2245
  self, buffer_sizes: wp.array(dtype=wp.int32), offsets: wp.array(dtype=wp.int32)
@@ -2040,17 +2251,26 @@ class TriMeshCollisionDetector:
2040
2251
 
2041
2252
  offsets.assign(offsets_np)
2042
2253
 
2043
- def refit(self):
2254
+ def refit(self, new_pos=None):
2255
+ if new_pos is not None:
2256
+ self.vertex_positions = new_pos
2257
+
2044
2258
  self.refit_triangles()
2045
2259
  self.refit_edges()
2046
2260
 
2047
2261
  def refit_triangles(self):
2048
- self.mesh_tris.refit()
2262
+ wp.launch(
2263
+ kernel=compute_tri_aabbs,
2264
+ inputs=[self.vertex_positions, self.model.tri_indices, self.lower_bounds_tris, self.upper_bounds_tris],
2265
+ dim=self.model.tri_count,
2266
+ device=self.model.device,
2267
+ )
2268
+ self.bvh_tris.refit()
2049
2269
 
2050
2270
  def refit_edges(self):
2051
2271
  wp.launch(
2052
2272
  kernel=compute_edge_aabbs,
2053
- inputs=[self.model.particle_q, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2273
+ inputs=[self.vertex_positions, self.model.edge_indices, self.lower_bounds_edges, self.upper_bounds_edges],
2054
2274
  dim=self.model.edge_count,
2055
2275
  device=self.model.device,
2056
2276
  )
@@ -2075,8 +2295,8 @@ class TriMeshCollisionDetector:
2075
2295
  kernel=vertex_triangle_collision_detection_kernel,
2076
2296
  inputs=[
2077
2297
  query_radius,
2078
- self.mesh_tris.id,
2079
- self.model.particle_q,
2298
+ self.bvh_tris.id,
2299
+ self.vertex_positions,
2080
2300
  self.model.tri_indices,
2081
2301
  self.vertex_colliding_triangles_offsets,
2082
2302
  self.vertex_colliding_triangles_buffer_sizes,
@@ -2098,11 +2318,11 @@ class TriMeshCollisionDetector:
2098
2318
 
2099
2319
  def edge_edge_collision_detection(self, query_radius):
2100
2320
  wp.launch(
2101
- kernel=edge_edge_collision_detection_kernel,
2321
+ kernel=edge_colliding_edges_detection_kernel,
2102
2322
  inputs=[
2103
2323
  query_radius,
2104
2324
  self.bvh_edges.id,
2105
- self.model.particle_q,
2325
+ self.vertex_positions,
2106
2326
  self.model.edge_indices,
2107
2327
  self.edge_colliding_edges_offsets,
2108
2328
  self.edge_colliding_edges_buffer_sizes,
@@ -2117,3 +2337,40 @@ class TriMeshCollisionDetector:
2117
2337
  dim=self.model.edge_count,
2118
2338
  device=self.model.device,
2119
2339
  )
2340
+
2341
+ def triangle_triangle_intersection_detection(self):
2342
+ if self.triangle_intersecting_triangles is None:
2343
+ self.triangle_intersecting_triangles = wp.zeros(
2344
+ shape=(self.model.tri_count * self.triangle_triangle_collision_buffer_pre_alloc,),
2345
+ dtype=wp.int32,
2346
+ device=self.device,
2347
+ )
2348
+
2349
+ if self.triangle_intersecting_triangles_count is None:
2350
+ self.triangle_intersecting_triangles_count = wp.array(
2351
+ shape=(self.model.tri_count,), dtype=wp.int32, device=self.device
2352
+ )
2353
+
2354
+ if self.triangle_intersecting_triangles_offsets is None:
2355
+ buffer_sizes = np.full((self.model.tri_count,), self.triangle_triangle_collision_buffer_pre_alloc)
2356
+ offsets = np.zeros((self.model.tri_count + 1,), dtype=np.int32)
2357
+ offsets[1:] = np.cumsum(buffer_sizes)
2358
+
2359
+ self.triangle_intersecting_triangles_offsets = wp.array(offsets, dtype=wp.int32, device=self.device)
2360
+
2361
+ wp.launch(
2362
+ kernel=triangle_triangle_collision_detection_kernel,
2363
+ inputs=[
2364
+ self.bvh_tris.id,
2365
+ self.vertex_positions,
2366
+ self.model.tri_indices,
2367
+ self.triangle_intersecting_triangles_offsets,
2368
+ ],
2369
+ outputs=[
2370
+ self.triangle_intersecting_triangles,
2371
+ self.triangle_intersecting_triangles_count,
2372
+ self.resize_flags,
2373
+ ],
2374
+ dim=self.model.tri_count,
2375
+ device=self.model.device,
2376
+ )
warp/sim/import_urdf.py CHANGED
@@ -12,6 +12,7 @@ from typing import Union
12
12
  import numpy as np
13
13
 
14
14
  import warp as wp
15
+ import warp.sim
15
16
  from warp.sim.model import Mesh
16
17
 
17
18
 
@@ -20,7 +21,7 @@ def parse_urdf(
20
21
  builder,
21
22
  xform=None,
22
23
  floating=False,
23
- base_joint: Union[dict, str] = None,
24
+ base_joint: Union[dict, str, None] = None,
24
25
  density=1000.0,
25
26
  stiffness=100.0,
26
27
  damping=10.0,
@@ -511,11 +512,25 @@ def parse_urdf(
511
512
 
512
513
  builder.add_joint_d6(
513
514
  linear_axes=[
514
- wp.sim.JointAxis(
515
- u, limit_lower=lower * scale, limit_upper=upper * scale, limit_ke=limit_ke, limit_kd=limit_kd
515
+ warp.sim.JointAxis(
516
+ u,
517
+ limit_lower=lower * scale,
518
+ limit_upper=upper * scale,
519
+ limit_ke=limit_ke,
520
+ limit_kd=limit_kd,
521
+ target_ke=stiffness,
522
+ target_kd=joint_damping,
523
+ mode=joint_mode,
516
524
  ),
517
- wp.sim.JointAxis(
518
- v, limit_lower=lower * scale, limit_upper=upper * scale, limit_ke=limit_ke, limit_kd=limit_kd
525
+ warp.sim.JointAxis(
526
+ v,
527
+ limit_lower=lower * scale,
528
+ limit_upper=upper * scale,
529
+ limit_ke=limit_ke,
530
+ limit_kd=limit_kd,
531
+ target_ke=stiffness,
532
+ target_kd=joint_damping,
533
+ mode=joint_mode,
519
534
  ),
520
535
  ],
521
536
  **joint_params,
@@ -871,6 +871,7 @@ def eval_rigid_contacts(
871
871
  contact_shape0: wp.array(dtype=int),
872
872
  contact_shape1: wp.array(dtype=int),
873
873
  force_in_world_frame: bool,
874
+ friction_smoothing: float,
874
875
  # outputs
875
876
  body_f: wp.array(dtype=wp.spatial_vector),
876
877
  ):
@@ -924,6 +925,8 @@ def eval_rigid_contacts(
924
925
  n = contact_normal[tid]
925
926
  bx_a = contact_point0[tid]
926
927
  bx_b = contact_point1[tid]
928
+ r_a = wp.vec3(0.0)
929
+ r_b = wp.vec3(0.0)
927
930
  if body_a >= 0:
928
931
  X_wb_a = body_q[body_a]
929
932
  X_com_a = body_com[body_a]
@@ -990,12 +993,16 @@ def eval_rigid_contacts(
990
993
  # ft = wp.vec3(vx, 0.0, vz)
991
994
 
992
995
  # Coulomb friction (smooth, but gradients are numerically unstable around |vt| = 0)
993
- # ft = wp.normalize(vt)*wp.min(kf*wp.length(vt), abs(mu*d*ke))
994
996
  ft = wp.vec3(0.0)
995
997
  if d < 0.0:
996
- ft = wp.normalize(vt) * wp.min(kf * wp.length(vt), -mu * (fn + fd))
998
+ # use a smooth vector norm to avoid gradient instability at/around zero velocity
999
+ vs = wp.norm_huber(vt, delta=friction_smoothing)
1000
+ if vs > 0.0:
1001
+ fr = vt / vs
1002
+ ft = fr * wp.min(kf * vs, -mu * (fn + fd))
997
1003
 
998
1004
  f_total = n * (fn + fd) + ft
1005
+ # f_total = n * (fn + fd)
999
1006
  # f_total = n * fn
1000
1007
 
1001
1008
  if body_a >= 0:
@@ -1761,7 +1768,7 @@ def eval_tetrahedral_forces(model: Model, state: State, control: Control, partic
1761
1768
  )
1762
1769
 
1763
1770
 
1764
- def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
1771
+ def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array, friction_smoothing: float = 1.0):
1765
1772
  if model.rigid_contact_max and (
1766
1773
  model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count
1767
1774
  ):
@@ -1782,6 +1789,7 @@ def eval_body_contact_forces(model: Model, state: State, particle_f: wp.array):
1782
1789
  model.rigid_contact_shape0,
1783
1790
  model.rigid_contact_shape1,
1784
1791
  False,
1792
+ friction_smoothing,
1785
1793
  ],
1786
1794
  outputs=[state.body_f],
1787
1795
  device=model.device,
@@ -1880,7 +1888,15 @@ def eval_muscle_forces(model: Model, state: State, control: Control, body_f: wp.
1880
1888
  )
1881
1889
 
1882
1890
 
1883
- def compute_forces(model: Model, state: State, control: Control, particle_f: wp.array, body_f: wp.array, dt: float):
1891
+ def compute_forces(
1892
+ model: Model,
1893
+ state: State,
1894
+ control: Control,
1895
+ particle_f: wp.array,
1896
+ body_f: wp.array,
1897
+ dt: float,
1898
+ friction_smoothing: float = 1.0,
1899
+ ):
1884
1900
  # damped springs
1885
1901
  eval_spring_forces(model, state, particle_f)
1886
1902
 
@@ -1906,7 +1922,7 @@ def compute_forces(model: Model, state: State, control: Control, particle_f: wp.
1906
1922
  eval_particle_ground_contact_forces(model, state, particle_f)
1907
1923
 
1908
1924
  # body contacts
1909
- eval_body_contact_forces(model, state, particle_f)
1925
+ eval_body_contact_forces(model, state, particle_f, friction_smoothing=friction_smoothing)
1910
1926
 
1911
1927
  # particle shape contact
1912
1928
  eval_particle_body_contact_forces(model, state, particle_f, body_f, body_f_in_world_frame=False)
@@ -1941,12 +1957,14 @@ class SemiImplicitIntegrator(Integrator):
1941
1957
 
1942
1958
  """
1943
1959
 
1944
- def __init__(self, angular_damping: float = 0.05):
1960
+ def __init__(self, angular_damping: float = 0.05, friction_smoothing: float = 1.0):
1945
1961
  """
1946
1962
  Args:
1947
1963
  angular_damping (float, optional): Angular damping factor. Defaults to 0.05.
1964
+ 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.
1948
1965
  """
1949
1966
  self.angular_damping = angular_damping
1967
+ self.friction_smoothing = friction_smoothing
1950
1968
 
1951
1969
  def simulate(self, model: Model, state_in: State, state_out: State, dt: float, control: Control = None):
1952
1970
  with wp.ScopedTimer("simulate", False):
@@ -1962,7 +1980,7 @@ class SemiImplicitIntegrator(Integrator):
1962
1980
  if control is None:
1963
1981
  control = model.control(clone_variables=False)
1964
1982
 
1965
- compute_forces(model, state_in, control, particle_f, body_f, dt)
1983
+ compute_forces(model, state_in, control, particle_f, body_f, dt, friction_smoothing=self.friction_smoothing)
1966
1984
 
1967
1985
  self.integrate_bodies(model, state_in, state_out, dt, self.angular_damping)
1968
1986