warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/sim/particles.py

@@ -11,6 +11,7 @@ from .model import PARTICLE_FLAG_ACTIVE
 
 @wp.func
 def particle_force(n: wp.vec3, v: wp.vec3, c: float, k_n: float, k_d: float, k_f: float, k_mu: float):
+    # compute normal and tangential friction force for a single contact
     vn = wp.dot(n, v)
     jn = c * k_n
     jd = min(vn, 0.0) * k_d
@@ -89,7 +90,7 @@ def eval_particle_forces_kernel(
 
 
 def eval_particle_forces(model, state, forces):
-    if model.particle_max_radius > 0.0:
+    if model.particle_count > 1 and model.particle_max_radius > 0.0:
         wp.launch(
             kernel=eval_particle_forces_kernel,
             dim=model.particle_count,

warp/sim/render.py

@@ -23,6 +23,7 @@ NAN = wp.constant(-1.0e8)
 def compute_contact_points(
     body_q: wp.array(dtype=wp.transform),
     shape_body: wp.array(dtype=int),
+    contact_count: wp.array(dtype=int),
     contact_shape0: wp.array(dtype=int),
     contact_shape1: wp.array(dtype=int),
     contact_point0: wp.array(dtype=wp.vec3),
@@ -32,6 +33,11 @@ def compute_contact_points(
     contact_pos1: wp.array(dtype=wp.vec3),
 ):
     tid = wp.tid()
+    count = contact_count[0]
+    if tid >= count:
+        contact_pos0[tid] = wp.vec3(NAN, NAN, NAN)
+        contact_pos1[tid] = wp.vec3(NAN, NAN, NAN)
+        return
     shape_a = contact_shape0[tid]
     shape_b = contact_shape1[tid]
     if shape_a == shape_b:
@@ -124,11 +130,12 @@ def CreateSimRenderer(renderer):
                 shape_geo_thickness = model.shape_geo.thickness.numpy()
                 shape_geo_is_solid = model.shape_geo.is_solid.numpy()
                 shape_transform = model.shape_transform.numpy()
+                shape_visible = model.shape_visible.numpy()
 
                 p = np.zeros(3, dtype=np.float32)
                 q = np.array([0.0, 0.0, 0.0, 1.0], dtype=np.float32)
                 scale = np.ones(3)
-                color = np.ones(3)
+                color = (1.0, 1.0, 1.0)
                 # loop over shapes excluding the ground plane
                 for s in range(model.shape_count - 1):
                     geo_type = shape_geo_type[s]
@@ -136,8 +143,6 @@ def CreateSimRenderer(renderer):
                     geo_thickness = float(shape_geo_thickness[s])
                     geo_is_solid = bool(shape_geo_is_solid[s])
                     geo_src = shape_geo_src[s]
-                    if self.use_unique_colors:
-                        color = self._get_new_color()
                     name = f"shape_{s}"
 
                     # shape transform in body frame
@@ -147,6 +152,9 @@ def CreateSimRenderer(renderer):
                     else:
                         body = None
 
+                    if self.use_unique_colors and body is not None:
+                        color = self._get_new_color()
+
                     # shape transform in body frame
                     X_bs = wp.transform_expand(shape_transform[s])
                     # check whether we can instance an already created shape with the same geometry
@@ -167,25 +175,25 @@ def CreateSimRenderer(renderer):
                             )
 
                         elif geo_type == warp.sim.GEO_SPHERE:
-                            shape = self.render_sphere(name, p, q, geo_scale[0], parent_body=body, is_template=True)
+                            shape = self.render_sphere(name, p, q, geo_scale[0], parent_body=body, is_template=True, color=color)
 
                         elif geo_type == warp.sim.GEO_CAPSULE:
                             shape = self.render_capsule(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True
+                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
                             )
 
                         elif geo_type == warp.sim.GEO_CYLINDER:
                             shape = self.render_cylinder(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True
+                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
                             )
 
                         elif geo_type == warp.sim.GEO_CONE:
                             shape = self.render_cone(
-                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True
+                                name, p, q, geo_scale[0], geo_scale[1], parent_body=body, is_template=True, color=color
                             )
 
                         elif geo_type == warp.sim.GEO_BOX:
-                            shape = self.render_box(name, p, q, geo_scale, parent_body=body, is_template=True)
+                            shape = self.render_box(name, p, q, geo_scale, parent_body=body, is_template=True, color=color)
 
                         elif geo_type == warp.sim.GEO_MESH:
                             if not geo_is_solid:
@@ -210,7 +218,9 @@ def CreateSimRenderer(renderer):
 
                         self.geo_shape[geo_hash] = shape
 
-                    self.add_shape_instance(name, shape, body, X_bs.p, X_bs.q, scale)
+                    if shape_visible[s]:
+                        # TODO support dynamic visibility
+                        self.add_shape_instance(name, shape, body, X_bs.p, X_bs.q, scale, custom_index=s, visible=shape_visible[s])
                     self.instance_count += 1
 
                 if self.show_joints and model.joint_count:
@@ -274,7 +284,7 @@ def CreateSimRenderer(renderer):
                             self.instance_count += 1
 
             if model.ground:
-                self.render_ground()
+                self.render_ground(plane=model.ground_plane_params)
 
             if hasattr(self, "complete_setup"):
                 self.complete_setup()
@@ -283,6 +293,12 @@ def CreateSimRenderer(renderer):
             return tab10_color_map(self.instance_count)
 
         def render(self, state: warp.sim.State):
+            """
+            Updates the renderer with the given simulation state.
+
+            Args:
+                state (warp.sim.State): The simulation state to render.
+            """
             if self.skip_rendering:
                 return
 
@@ -290,15 +306,20 @@ def CreateSimRenderer(renderer):
                 particle_q = state.particle_q.numpy()
 
                 # render particles
-                self.render_points("particles", particle_q, radius=self.model.particle_radius.numpy())
+                self.render_points("particles", particle_q, radius=self.model.particle_radius.numpy(), colors=((0.8, 0.3, 0.2),) * len(particle_q))
 
                 # render tris
                 if self.model.tri_count:
-                    self.render_mesh("surface", particle_q, self.model.tri_indices.numpy().flatten())
+                    self.render_mesh(
+                        "surface",
+                        particle_q,
+                        self.model.tri_indices.numpy().flatten(),
+                        colors=(((0.75, 0.25, 0.0),) * len(particle_q)),
+                    )
 
                 # render springs
                 if self.model.spring_count:
-                    self.render_line_list("springs", particle_q, self.model.spring_indices.numpy().flatten(), [], 0.05)
+                    self.render_line_list("springs", particle_q, self.model.spring_indices.numpy().flatten(), (0.25, 0.5, 0.25), 0.02)
 
             # render muscles
             if self.model.muscle_count:
@@ -348,6 +369,7 @@ def CreateSimRenderer(renderer):
                         inputs=[
                             state.body_q,
                             self.model.shape_body,
+                            self.model.rigid_contact_count,
                             self.model.rigid_contact_shape0,
                             self.model.rigid_contact_shape1,
                             self.model.rigid_contact_point0,

warp/sim/utils.py

@@ -1,13 +1,20 @@
 import warp as wp
+import numpy as np
 
-PI = wp.constant(3.14159265359)
-PI_2 = wp.constant(1.57079632679)
+from typing import Tuple, List
 
 
 @wp.func
 def velocity_at_point(qd: wp.spatial_vector, r: wp.vec3):
     """
     Returns the velocity of a point relative to the frame with the given spatial velocity.
+
+    Args:
+        qd (spatial_vector): The spatial velocity of the frame.
+        r (vec3): The position of the point relative to the frame. 
+
+    Returns:
+        vec3: The velocity of the point.
     """
     return wp.cross(wp.spatial_top(qd), r) + wp.spatial_bottom(qd)
 
@@ -52,7 +59,7 @@ def quat_decompose(q: wp.quat):
     phi = wp.atan2(R[1, 2], R[2, 2])
     sinp = -R[0, 2]
     if wp.abs(sinp) >= 1.0:
-        theta = 1.57079632679 * wp.sign(sinp)
+        theta = wp.HALF_PI * wp.sign(sinp)
     else:
         theta = wp.asin(-R[0, 2])
     psi = wp.atan2(R[0, 1], R[0, 0])
@@ -63,7 +70,7 @@ def quat_decompose(q: wp.quat):
 @wp.func
 def quat_to_rpy(q: wp.quat):
     """
-    Convert a quaternion into euler angles (roll, pitch, yaw)
+    Convert a quaternion into Euler angles (roll, pitch, yaw)
     roll is rotation around x in radians (counterclockwise)
     pitch is rotation around y in radians (counterclockwise)
     yaw is rotation around z in radians (counterclockwise)
@@ -90,11 +97,27 @@ def quat_to_rpy(q: wp.quat):
 @wp.func
 def quat_to_euler(q: wp.quat, i: int, j: int, k: int) -> wp.vec3:
     """
-    Convert a quaternion into euler angles
-    i, j, k are the indices in [1,2,3] of the axes to use
-    (i != j, j != k)
+    Convert a quaternion into Euler angles.
+
+    :math:`i, j, k` are the indices in :math:`[0, 1, 2]` of the axes to use
+    (:math:`i \\neq j, j \\neq k`).
+
+    Reference: https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0276302
+
+    Args:
+        q (quat): The quaternion to convert
+        i (int): The index of the first axis
+        j (int): The index of the second axis
+        k (int): The index of the third axis
+
+    Returns:
+        vec3: The Euler angles (in radians)
     """
-    # https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0276302
+    # i, j, k are actually assumed to follow 1-based indexing but
+    # we want to be compatible with quat_from_euler
+    i += 1
+    j += 1
+    k += 1
     not_proper = True
     if i == k:
         not_proper = False
@@ -116,17 +139,54 @@ def quat_to_euler(q: wp.quat, i: int, j: int, k: int) -> wp.vec3:
     t3 = 0.0
     if wp.abs(t2) < 1e-6:
         t3 = 2.0 * tp - t1
-    elif wp.abs(t2 - PI_2) < 1e-6:
+    elif wp.abs(t2 - wp.HALF_PI) < 1e-6:
         t3 = 2.0 * tm + t1
     else:
         t1 = tp - tm
         t3 = tp + tm
     if not_proper:
-        t2 -= PI_2
+        t2 -= wp.HALF_PI
         t3 *= e
     return wp.vec3(t1, t2, t3)
 
 
+@wp.func
+def quat_from_euler(e: wp.vec3, i: int, j: int, k: int) -> wp.quat:
+    """
+    Convert Euler angles to a quaternion.
+
+    :math:`i, j, k` are the indices in :math:`[0, 1, 2]` of the axes in which the Euler angles are provided
+    (:math:`i \\neq j, j \\neq k`), e.g. (0, 1, 2) for Euler sequence XYZ.
+
+    Args:
+        e (vec3): The Euler angles (in radians)
+        i (int): The index of the first axis
+        j (int): The index of the second axis
+        k (int): The index of the third axis
+
+    Returns:
+        quat: The quaternion
+    """
+    # Half angles
+    half_e = e / 2.0
+
+    # Precompute sines and cosines of half angles
+    cr = wp.cos(half_e[i])
+    sr = wp.sin(half_e[i])
+    cp = wp.cos(half_e[j])
+    sp = wp.sin(half_e[j])
+    cy = wp.cos(half_e[k])
+    sy = wp.sin(half_e[k])
+
+    # Components of the quaternion based on the rotation sequence
+    return wp.quat(
+        (cy * sr * cp - sy * cr * sp),
+        (cy * cr * sp + sy * sr * cp),
+        (sy * cr * cp - cy * sr * sp),
+        (cy * cr * cp + sy * sr * sp),
+    )
+
+
 @wp.func
 def transform_twist(t: wp.transform, x: wp.spatial_vector):
     # Frank & Park definition 3.20, pg 100
@@ -181,6 +241,35 @@ def transform_inertia(t: wp.transform, I: wp.spatial_matrix):
     return wp.mul(wp.mul(wp.transpose(T), I), T)
 
 
+@wp.func
+def boltzmann(a: float, b: float, alpha: float):
+    e1 = wp.exp(alpha * a)
+    e2 = wp.exp(alpha * b)
+    return (a * e1 + b * e2) / (e1 + e2)
+
+
+@wp.func
+def smooth_max(a: float, b: float, eps: float):
+    d = a - b
+    return 0.5 * (a + b + wp.sqrt(d * d + eps))
+
+
+@wp.func
+def smooth_min(a: float, b: float, eps: float):
+    d = a - b
+    return 0.5 * (a + b - wp.sqrt(d * d + eps))
+
+
+@wp.func
+def leaky_max(a: float, b: float):
+    return smooth_max(a, b, 1e-5)
+
+
+@wp.func
+def leaky_min(a: float, b: float):
+    return smooth_min(a, b, 1e-5)
+
+
 @wp.func
 def vec_min(a: wp.vec3, b: wp.vec3):
     return wp.vec3(wp.min(a[0], b[0]), wp.min(a[1], b[1]), wp.min(a[2], b[2]))
@@ -191,9 +280,131 @@ def vec_max(a: wp.vec3, b: wp.vec3):
     return wp.vec3(wp.max(a[0], b[0]), wp.max(a[1], b[1]), wp.max(a[2], b[2]))
 
 
+@wp.func
+def vec_leaky_min(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(leaky_min(a[0], b[0]), leaky_min(a[1], b[1]), leaky_min(a[2], b[2]))
+
+
+@wp.func
+def vec_leaky_max(a: wp.vec3, b: wp.vec3):
+    return wp.vec3(leaky_max(a[0], b[0]), leaky_max(a[1], b[1]), leaky_max(a[2], b[2]))
+
+
 @wp.func
 def vec_abs(a: wp.vec3):
     return wp.vec3(wp.abs(a[0]), wp.abs(a[1]), wp.abs(a[2]))
 
 
-        
+def load_mesh(filename: str, method: str = None):
+    """
+    Loads a 3D triangular surface mesh from a file.
+
+    Args:
+        filename (str): The path to the 3D model file (obj, and other formats supported by the different methods) to load.
+        method (str): The method to use for loading the mesh (default None). Can be either `"trimesh"`, `"meshio"`, `"pcu"`, or `"openmesh"`. If None, every method is tried and the first successful mesh import where the number of vertices is greater than 0 is returned.
+
+    Returns:
+        Tuple of (mesh_points, mesh_indices), where mesh_points is a Nx3 numpy array of vertex positions (float32),
+        and mesh_indices is a Mx3 numpy array of vertex indices (int32) for the triangular faces.
+    """
+    import os
+
+    if not os.path.exists(filename):
+        raise ValueError(f"File not found: {filename}")
+
+    def load_mesh_with_method(method):
+        if method == "meshio":
+            import meshio
+
+            m = meshio.read(filename)
+            mesh_points = np.array(m.points)
+            mesh_indices = np.array(m.cells[0].data, dtype=np.int32)
+        elif method == "openmesh":
+            import openmesh
+
+            m = openmesh.read_trimesh(filename)
+            mesh_points = np.array(m.points())
+            mesh_indices = np.array(m.face_vertex_indices(), dtype=np.int32)
+        elif method == "pcu":
+            import point_cloud_utils as pcu
+
+            mesh_points, mesh_indices = pcu.load_mesh_vf(filename)
+            mesh_indices = mesh_indices.flatten()
+        else:
+            import trimesh
+
+            m = trimesh.load(filename)
+            if hasattr(m, "geometry"):
+                # multiple meshes are contained in a scene; combine to one mesh
+                mesh_points = []
+                mesh_indices = []
+                index_offset = 0
+                for geom in m.geometry.values():
+                    vertices = np.array(geom.vertices, dtype=np.float32)
+                    faces = np.array(geom.faces.flatten(), dtype=np.int32)
+                    mesh_points.append(vertices)
+                    mesh_indices.append(faces + index_offset)
+                    index_offset += len(vertices)
+                mesh_points = np.concatenate(mesh_points, axis=0)
+                mesh_indices = np.concatenate(mesh_indices)
+            else:
+                # a single mesh
+                mesh_points = np.array(m.vertices, dtype=np.float32)
+                mesh_indices = np.array(m.faces.flatten(), dtype=np.int32)
+        return mesh_points, mesh_indices
+
+    if method is None:
+        methods = ["trimesh", "meshio", "pcu", "openmesh"]
+        for method in methods:
+            try:
+                mesh = load_mesh_with_method(method)
+                if mesh is not None and len(mesh[0]) > 0:
+                    return mesh
+            except Exception:
+                pass
+        raise ValueError(f"Failed to load mesh using any of the methods: {methods}")
+    else:
+        mesh = load_mesh_with_method(method)
+        if mesh is None or len(mesh[0]) == 0:
+            raise ValueError(f"Failed to load mesh using method {method}")
+        return mesh
+
+
+def visualize_meshes(
+    meshes: List[Tuple[list, list]], num_cols=0, num_rows=0, titles=[], scale_axes=True, show_plot=True
+):
+    # render meshes in a grid with matplotlib
+    import matplotlib.pyplot as plt
+    from mpl_toolkits.mplot3d import Axes3D
+
+    num_cols = min(num_cols, len(meshes))
+    num_rows = min(num_rows, len(meshes))
+    if num_cols and not num_rows:
+        num_rows = int(np.ceil(len(meshes) / num_cols))
+    elif num_rows and not num_cols:
+        num_cols = int(np.ceil(len(meshes) / num_rows))
+    else:
+        num_cols = len(meshes)
+        num_rows = 1
+
+    vertices = [np.array(v).reshape((-1, 3)) for v, _ in meshes]
+    faces = [np.array(f, dtype=np.int32).reshape((-1, 3)) for _, f in meshes]
+    if scale_axes:
+        ranges = np.array([v.max(axis=0) - v.min(axis=0) for v in vertices])
+        max_range = ranges.max()
+        mid_points = np.array([v.max(axis=0) + v.min(axis=0) for v in vertices]) * 0.5
+
+    fig = plt.figure(figsize=(12, 6))
+    for i, (vertices, faces) in enumerate(meshes):
+        ax = fig.add_subplot(num_rows, num_cols, i + 1, projection="3d")
+        if i < len(titles):
+            ax.set_title(titles[i])
+        ax.plot_trisurf(vertices[:, 0], vertices[:, 1], vertices[:, 2], triangles=faces, edgecolor="k")
+        if scale_axes:
+            mid = mid_points[i]
+            ax.set_xlim(mid[0] - max_range, mid[0] + max_range)
+            ax.set_ylim(mid[1] - max_range, mid[1] + max_range)
+            ax.set_zlim(mid[2] - max_range, mid[2] + max_range)
+    if show_plot:
+        plt.show()
+    return fig

warp/stubs.py

@@ -56,6 +56,8 @@ from warp.context import get_device, set_device, synchronize_device
 from warp.context import (
     zeros,
     zeros_like,
+    ones,
+    ones_like,
     full,
     full_like,
     clone,
@@ -74,9 +76,15 @@ from warp.context import Kernel, Function, Launch
 from warp.context import Stream, get_stream, set_stream, synchronize_stream
 from warp.context import Event, record_event, wait_event, wait_stream
 from warp.context import RegisteredGLBuffer
+from warp.context import is_mempool_supported, is_mempool_enabled, set_mempool_enabled
+from warp.context import set_mempool_release_threshold, get_mempool_release_threshold
+from warp.context import is_mempool_access_supported, is_mempool_access_enabled, set_mempool_access_enabled
+from warp.context import is_peer_access_supported, is_peer_access_enabled, set_peer_access_enabled
 
 from warp.tape import Tape
 from warp.utils import ScopedTimer, ScopedDevice, ScopedStream
+from warp.utils import ScopedMempool, ScopedMempoolAccess, ScopedPeerAccess
+from warp.utils import ScopedCapture
 from warp.utils import transform_expand, quat_between_vectors
 
 from warp.torch import from_torch, to_torch

warp/tape.py

@@ -95,7 +95,11 @@ class Tape:
         # existing code before we added wp.array.grad attribute
         if grads:
             for a, g in grads.items():
-                a.grad = g
+                if a.grad is None:
+                    a.grad = g
+                else:
+                    # ensure we can capture this backward pass in a CUDA graph
+                    a.grad.assign(g)
                 self.const_gradients.add(a)
 
         # run launches backwards
@@ -104,6 +108,17 @@ class Tape:
                 launch()
 
             else:
+                # kernel option takes precedence over module option
+                kernel_enable_backward = launch[0].options.get("enable_backward")
+                if kernel_enable_backward is False:
+                    msg = f"Running the tape backwards may produce incorrect gradients because recorded kernel {launch[0].key} is configured with the option 'enable_backward=False'."
+                    wp.utils.warn(msg)
+                elif kernel_enable_backward is None:
+                    module_enable_backward = launch[0].module.options.get("enable_backward")
+                    if module_enable_backward is False:
+                        msg = f"Running the tape backwards may produce incorrect gradients because recorded kernel {launch[0].key} is defined in a module with the option 'enable_backward=False' set."
+                        wp.utils.warn(msg)
+
                 kernel = launch[0]
                 dim = launch[1]
                 max_blocks = launch[2]

warp/tests/aux_test_grad_customs.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""This file is used to test importing a user-defined function with a custom gradient"""
+
+import warp as wp
+
+wp.init()
+
+
+@wp.func
+def aux_custom_fn(x: float, y: float):
+    return x * 3.0 + y / 3.0, y**2.5
+
+
+@wp.func_grad(aux_custom_fn)
+def aux_custom_fn_grad(x: float, y: float, adj_ret0: float, adj_ret1: float):
+    wp.adjoint[x] += x * adj_ret0 * 42.0 + y * adj_ret1 * 10.0
+    wp.adjoint[y] += y * adj_ret1 * 3.0