warp-lang 0.10.1__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/sim/model.py

@@ -8,21 +8,23 @@
 """A module for building simulation models and state.
 """
 
-from .inertia import transform_inertia
-from .inertia import compute_cone_inertia
-from .inertia import compute_cylinder_inertia
-from .inertia import compute_capsule_inertia
-from .inertia import compute_box_inertia
-from .inertia import compute_sphere_inertia
-from .inertia import compute_mesh_inertia
+import copy
+import math
+from typing import List, Optional, Tuple
 
-import warp as wp
 import numpy as np
 
-import math
-import copy
+import warp as wp
 
-from typing import List, Optional, Tuple
+from .inertia import (
+    compute_box_inertia,
+    compute_capsule_inertia,
+    compute_cone_inertia,
+    compute_cylinder_inertia,
+    compute_mesh_inertia,
+    compute_sphere_inertia,
+    transform_inertia,
+)
 
 Vec3 = List[float]
 Vec4 = List[float]
@@ -92,11 +94,27 @@ class ModelShapeGeometry:
 
 # Axis (linear or angular) of a joint that can have bounds and be driven towards a target
 class JointAxis:
+    """
+    Describes a joint axis that can have limits and be driven towards a target.
+
+    Attributes:
+
+        axis (3D vector or JointAxis): The 3D axis that this JointAxis object describes, or alternatively another JointAxis object to copy from
+        limit_lower (float): The lower limit of the joint axis
+        limit_upper (float): The upper limit of the joint axis
+        limit_ke (float): The elastic stiffness of the joint axis limits, only respected by :class:`SemiImplicitIntegrator`
+        limit_kd (float): The damping stiffness of the joint axis limits, only respected by :class:`SemiImplicitIntegrator`
+        target (float): The target position or velocity (depending on the mode, see `Joint modes`_) of the joint axis
+        target_ke (float): The proportional gain of the joint axis target drive PD controller
+        target_kd (float): The derivative gain of the joint axis target drive PD controller
+        mode (int): The mode of the joint axis, see `Joint modes`_
+    """
+
     def __init__(
         self,
         axis,
-        limit_lower=-np.inf,
-        limit_upper=np.inf,
+        limit_lower=-math.inf,
+        limit_upper=math.inf,
         limit_ke=100.0,
         limit_kd=10.0,
         target=None,
@@ -104,30 +122,92 @@ class JointAxis:
         target_kd=0.0,
         mode=JOINT_MODE_TARGET_POSITION,
     ):
-        self.axis = np.array(wp.normalize(np.array(axis, dtype=np.float32)))
-        self.limit_lower = limit_lower
-        self.limit_upper = limit_upper
-        self.limit_ke = limit_ke
-        self.limit_kd = limit_kd
-        if target is not None:
-            self.target = target
-        elif limit_lower > 0.0 or limit_upper < 0.0:
-            self.target = 0.5 * (limit_lower + limit_upper)
+        if isinstance(axis, JointAxis):
+            self.axis = axis.axis
+            self.limit_lower = axis.limit_lower
+            self.limit_upper = axis.limit_upper
+            self.limit_ke = axis.limit_ke
+            self.limit_kd = axis.limit_kd
+            self.target = axis.target
+            self.target_ke = axis.target_ke
+            self.target_kd = axis.target_kd
+            self.mode = axis.mode
         else:
-            self.target = 0.0
-        self.target_ke = target_ke
-        self.target_kd = target_kd
-        self.mode = mode
+            self.axis = wp.normalize(wp.vec3(axis))
+            self.limit_lower = limit_lower
+            self.limit_upper = limit_upper
+            self.limit_ke = limit_ke
+            self.limit_kd = limit_kd
+            if target is not None:
+                self.target = target
+            elif limit_lower > 0.0 or limit_upper < 0.0:
+                self.target = 0.5 * (limit_lower + limit_upper)
+            else:
+                self.target = 0.0
+            self.target_ke = target_ke
+            self.target_kd = target_kd
+            self.mode = mode
+
+
+class SDF:
+    """Describes a signed distance field for simulation
+    
+    Attributes:
+
+        volume (Volume): The volume defining the SDF
+        I (Mat33): 3x3 inertia matrix of the SDF
+        mass (float): The total mass of the SDF
+        com (Vec3): The center of mass of the SDF
+    """
+    def __init__(self, volume=None, I=wp.mat33(np.eye(3)), mass=1.0, com=wp.vec3()):
+        self.volume = volume
+        self.I = I
+        self.mass = mass
+        self.com = com
+
+        # Need to specify these for now
+        self.has_inertia = True
+        self.is_solid = True
+
+    def finalize(self, device=None):
+        return self.volume.id
+
+    def __hash__(self):
+        return hash((self.volume.id))
 
 
 class Mesh:
     """Describes a triangle collision mesh for simulation
+    
+    Example mesh creation from a triangle OBJ mesh file:
+    ====================================================
+
+    .. code-block:: python
+
+        import numpy as np
+        import warp as wp
+        import warp.sim
+
+        def load_mesh(self, filename, use_meshio=True):
+            if use_meshio:
+                import meshio
+                m = meshio.read(filename)
+                mesh_points = np.array(m.points)
+                mesh_indices = np.array(m.cells[0].data, dtype=np.int32).flatten()
+            else:
+                import openmesh
+                m = openmesh.read_trimesh(filename)
+                mesh_points = np.array(m.points())
+                mesh_indices = np.array(m.face_vertex_indices(), dtype=np.int32).flatten()
+            return wp.sim.Mesh(mesh_points, mesh_indices)
+        
+        mesh = load_mesh("mesh.obj")
 
     Attributes:
 
-        vertices (List[Vec3]): Mesh vertices
-        indices (List[int]): Mesh indices
-        I (Mat33): Inertia tensor of the mesh assuming density of 1.0 (around the center of mass)
+        vertices (List[Vec3]): Mesh 3D vertices points
+        indices (List[int]): Mesh indices as a flattened list of vertex indices describing triangles
+        I (Mat33): 3x3 inertia matrix of the mesh assuming density of 1.0 (around the center of mass)
         mass (float): The total mass of the body assuming density of 1.0
         com (Vec3): The center of mass of the body
     """
@@ -146,20 +226,28 @@ class Mesh:
             is_solid: If True, the mesh is assumed to be a solid during inertia computation, otherwise it is assumed to be a hollow surface
         """
 
-        self.vertices = vertices
-        self.indices = indices
+        self.vertices = np.array(vertices).reshape(-1, 3)
+        self.indices = np.array(indices, dtype=np.int32).flatten()
         self.is_solid = is_solid
         self.has_inertia = compute_inertia
 
         if compute_inertia:
             self.mass, self.com, self.I, _ = compute_mesh_inertia(1.0, vertices, indices, is_solid=is_solid)
         else:
-            self.I = np.eye(3, dtype=np.float32)
+            self.I = wp.mat33(np.eye(3))
             self.mass = 1.0
-            self.com = np.array((0.0, 0.0, 0.0))
+            self.com = wp.vec3()
 
-    # construct simulation ready buffers from points
     def finalize(self, device=None):
+        """
+        Constructs a simulation-ready :class:`Mesh` object from the mesh data and returns its ID.
+        
+        Args:
+            device: The device on which to allocate the mesh buffers
+
+        Returns:
+            The ID of the simulation-ready :class:`Mesh`
+        """
         with wp.ScopedDevice(device):
             pos = wp.array(self.vertices, dtype=wp.vec3)
             vel = wp.zeros_like(pos)
@@ -169,6 +257,9 @@ class Mesh:
             return self.mesh.id
 
     def __hash__(self):
+        """
+        Computes a hash of the mesh data for use in caching. The hash considers the mesh vertices, indices, and whether the mesh is solid or not.
+        """
         return hash((tuple(np.array(self.vertices).flatten()), tuple(np.array(self.indices).flatten()), self.is_solid))
 
 
@@ -183,11 +274,16 @@ class State:
 
     Attributes:
 
-        particle_q (wp.array): Tensor of particle positions
-        particle_qd (wp.array): Tensor of particle velocities
+        particle_count (int): Number of particles
+        body_count (int): Number of rigid bodies
+
+        particle_q (array): Tensor of particle positions
+        particle_qd (array): Tensor of particle velocities
+        particle_f (array): Tensor of particle forces
 
-        body_q (wp.array): Tensor of body coordinates
-        body_qd (wp.array): Tensor of body velocities
+        body_q (array): Tensor of body coordinates
+        body_qd (array): Tensor of body velocities
+        body_f (array): Tensor of body forces
 
     """
 
@@ -195,6 +291,14 @@ class State:
         self.particle_count = 0
         self.body_count = 0
 
+        self.particle_q = None
+        self.particle_qd = None
+        self.particle_f = None
+
+        self.body_q = None
+        self.body_qd = None
+        self.body_f = None
+
     def clear_forces(self):
         if self.particle_count:
             self.particle_f.zero_()
@@ -203,25 +307,25 @@ class State:
             self.body_f.zero_()
 
     def flatten(self):
-        """Returns a list of Tensors stored by the state
-
-        This function is intended to be used internal-only but can be used to obtain
-        a set of all tensors owned by the state.
-        """
+        wp.utils.warn(
+            "Model.flatten() will be removed in a future Warp version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
 
-        tensors = []
+        arrays = []
 
         # build a list of all tensor attributes
         for attr, value in self.__dict__.items():
-            if wp.is_tensor(value):
-                tensors.append(value)
+            if isinstance(value, wp.array):
+                arrays.append(value)
 
-        return tensors
+        return arrays
 
 
 def compute_shape_mass(type, scale, src, density, is_solid, thickness):
     if density == 0.0 or type == GEO_PLANE:  # zero density means fixed
-        return 0.0, np.zeros(3), np.zeros((3, 3))
+        return 0.0, wp.vec3(), wp.mat33()
 
     if type == GEO_SPHERE:
         solid = compute_sphere_inertia(density, scale[0])
@@ -231,7 +335,7 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
             hollow = compute_sphere_inertia(density, scale[0] - thickness)
             return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
     elif type == GEO_BOX:
-        w, h, d = np.array(scale[:3]) * 2.0
+        w, h, d = scale * 2.0
         solid = compute_box_inertia(density, w, h, d)
         if is_solid:
             return solid
@@ -262,17 +366,16 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
         else:
             hollow = compute_cone_inertia(density, r - thickness, h - 2.0 * thickness)
             return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
-    elif type == GEO_MESH:
+    elif type == GEO_MESH or type == GEO_SDF:
         if src.has_inertia and src.mass > 0.0 and src.is_solid == is_solid:
             m, c, I = src.mass, src.com, src.I
 
-            s = np.array(scale[:3])
-            sx, sy, sz = s
+            sx, sy, sz = scale
 
             mass_ratio = sx * sy * sz * density
             m_new = m * mass_ratio
 
-            c_new = c * s
+            c_new = wp.cw_mul(c, scale)
 
             Ixx = I[0, 0] * (sy**2 + sz**2) / 2 * mass_ratio
             Iyy = I[1, 1] * (sx**2 + sz**2) / 2 * mass_ratio
@@ -281,12 +384,12 @@ def compute_shape_mass(type, scale, src, density, is_solid, thickness):
             Ixz = I[0, 2] * sx * sz * mass_ratio
             Iyz = I[1, 2] * sy * sz * mass_ratio
 
-            I_new = np.array([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])
+            I_new = wp.mat33([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])
 
             return m_new, c_new, I_new
-        else:
+        elif type == GEO_MESH:
             # fall back to computing inertia from mesh geometry
-            vertices = np.array(src.vertices) * np.array(scale[:3])
+            vertices = np.array(src.vertices) * np.array(scale)
             m, c, I, vol = compute_mesh_inertia(density, vertices, src.indices, is_solid, thickness)
             return m, c, I
     raise ValueError("Unsupported shape type: {}".format(type))
@@ -302,23 +405,24 @@ class Model:
         requires_grad (float): Indicates whether the model was finalized with gradient computation enabled
         num_envs (int): Number of articulation environments that were added to the ModelBuilder via `add_builder`
 
-        particle_q (wp.array): Particle positions, shape [particle_count, 3], float
-        particle_qd (wp.array): Particle velocities, shape [particle_count, 3], float
-        particle_mass (wp.array): Particle mass, shape [particle_count], float
-        particle_inv_mass (wp.array): Particle inverse mass, shape [particle_count], float
-        particle_radius (wp.array): Particle radius, shape [particle_count], float
+        particle_q (array): Particle positions, shape [particle_count, 3], float
+        particle_qd (array): Particle velocities, shape [particle_count, 3], float
+        particle_mass (array): Particle mass, shape [particle_count], float
+        particle_inv_mass (array): Particle inverse mass, shape [particle_count], float
+        particle_radius (array): Particle radius, shape [particle_count], float
         particle_max_radius (float): Maximum particle radius (useful for HashGrid construction)
-        particle_ke (wp.array): Particle normal contact stiffness (used by SemiImplicitIntegrator), shape [particle_count], float
-        particle_kd (wp.array): Particle normal contact damping (used by SemiImplicitIntegrator), shape [particle_count], float
-        particle_kf (wp.array): Particle friction force stiffness (used by SemiImplicitIntegrator), shape [particle_count], float
-        particle_mu (wp.array): Particle friction coefficient, shape [particle_count], float
-        particle_cohesion (wp.array): Particle cohesion strength, shape [particle_count], float
-        particle_adhesion (wp.array): Particle adhesion strength, shape [particle_count], float
+        particle_ke (array): Particle normal contact stiffness (used by :class:`SemiImplicitIntegrator`), shape [particle_count], float
+        particle_kd (array): Particle normal contact damping (used by :class:`SemiImplicitIntegrator`), shape [particle_count], float
+        particle_kf (array): Particle friction force stiffness (used by :class:`SemiImplicitIntegrator`), shape [particle_count], float
+        particle_mu (array): Particle friction coefficient, shape [particle_count], float
+        particle_cohesion (array): Particle cohesion strength, shape [particle_count], float
+        particle_adhesion (array): Particle adhesion strength, shape [particle_count], float
         particle_grid (HashGrid): HashGrid instance used for accelerated simulation of particle interactions
-        particle_flags (wp.array): Particle enabled state, shape [particle_count], bool
+        particle_flags (array): Particle enabled state, shape [particle_count], bool
+        particle_max_velocity (float): Maximum particle velocity (to prevent instability)
 
-        shape_transform (wp.array): Rigid shape transforms, shape [shape_count, 7], float
-        shape_body (wp.array): Rigid shape body index, shape [shape_count], int
+        shape_transform (array): Rigid shape transforms, shape [shape_count, 7], float
+        shape_body (array): Rigid shape body index, shape [shape_count], int
         body_shapes (dict): Mapping from body index to list of attached shape indices
         shape_materials (ModelShapeMaterials): Rigid shape contact materials, shape [shape_count], float
         shape_shape_geo (ModelShapeGeometry): Shape geometry properties (geo type, scale, thickness, etc.), shape [shape_count, 3], float
@@ -327,73 +431,72 @@ class Model:
         shape_collision_group (list): Collision group of each shape, shape [shape_count], int
         shape_collision_group_map (dict): Mapping from collision group to list of shape indices
         shape_collision_filter_pairs (set): Pairs of shape indices that should not collide
-        shape_collision_radius (wp.array): Collision radius of each shape used for bounding sphere broadphase collision checking, shape [shape_count], float
+        shape_collision_radius (array): Collision radius of each shape used for bounding sphere broadphase collision checking, shape [shape_count], float
         shape_ground_collision (list): Indicates whether each shape should collide with the ground, shape [shape_count], bool
-        shape_contact_pairs (wp.array): Pairs of shape indices that may collide, shape [contact_pair_count, 2], int
-        shape_ground_contact_pairs (wp.array): Pairs of shape, ground indices that may collide, shape [ground_contact_pair_count, 2], int
-
-        spring_indices (wp.array): Particle spring indices, shape [spring_count*2], int
-        spring_rest_length (wp.array): Particle spring rest length, shape [spring_count], float
-        spring_stiffness (wp.array): Particle spring stiffness, shape [spring_count], float
-        spring_damping (wp.array): Particle spring damping, shape [spring_count], float
-        spring_control (wp.array): Particle spring activation, shape [spring_count], float
-
-        tri_indices (wp.array): Triangle element indices, shape [tri_count*3], int
-        tri_poses (wp.array): Triangle element rest pose, shape [tri_count, 2, 2], float
-        tri_activations (wp.array): Triangle element activations, shape [tri_count], float
-        tri_materials (wp.array): Triangle element materials, shape [tri_count, 5], float
-
-        edge_indices (wp.array): Bending edge indices, shape [edge_count*2], int
-        edge_rest_angle (wp.array): Bending edge rest angle, shape [edge_count], float
-        edge_bending_properties (wp.array): Bending edge stiffness and damping parameters, shape [edge_count, 2], float
-
-        tet_indices (wp.array): Tetrahedral element indices, shape [tet_count*4], int
-        tet_poses (wp.array): Tetrahedral rest poses, shape [tet_count, 3, 3], float
-        tet_activations (wp.array): Tetrahedral volumetric activations, shape [tet_count], float
-        tet_materials (wp.array): Tetrahedral elastic parameters in form :math:`k_{mu}, k_{lambda}, k_{damp}`, shape [tet_count, 3]
-
-        body_q (wp.array): Poses of rigid bodies used for state initialization, shape [body_count, 7], float
-        body_qd (wp.array): Velocities of rigid bodies used for state initialization, shape [body_count, 6], float
-        body_com (wp.array): Rigid body center of mass (in local frame), shape [body_count, 7], float
-        body_inertia (wp.array): Rigid body inertia tensor (relative to COM), shape [body_count, 3, 3], float
-        body_inv_inertia (wp.array): Rigid body inverse inertia tensor (relative to COM), shape [body_count, 3, 3], float
-        body_mass (wp.array): Rigid body mass, shape [body_count], float
-        body_inv_mass (wp.array): Rigid body inverse mass, shape [body_count], float
+        shape_contact_pairs (array): Pairs of shape indices that may collide, shape [contact_pair_count, 2], int
+        shape_ground_contact_pairs (array): Pairs of shape, ground indices that may collide, shape [ground_contact_pair_count, 2], int
+
+        spring_indices (array): Particle spring indices, shape [spring_count*2], int
+        spring_rest_length (array): Particle spring rest length, shape [spring_count], float
+        spring_stiffness (array): Particle spring stiffness, shape [spring_count], float
+        spring_damping (array): Particle spring damping, shape [spring_count], float
+        spring_control (array): Particle spring activation, shape [spring_count], float
+
+        tri_indices (array): Triangle element indices, shape [tri_count*3], int
+        tri_poses (array): Triangle element rest pose, shape [tri_count, 2, 2], float
+        tri_activations (array): Triangle element activations, shape [tri_count], float
+        tri_materials (array): Triangle element materials, shape [tri_count, 5], float
+
+        edge_indices (array): Bending edge indices, shape [edge_count*4], int
+        edge_rest_angle (array): Bending edge rest angle, shape [edge_count], float
+        edge_bending_properties (array): Bending edge stiffness and damping parameters, shape [edge_count, 2], float
+
+        tet_indices (array): Tetrahedral element indices, shape [tet_count*4], int
+        tet_poses (array): Tetrahedral rest poses, shape [tet_count, 3, 3], float
+        tet_activations (array): Tetrahedral volumetric activations, shape [tet_count], float
+        tet_materials (array): Tetrahedral elastic parameters in form :math:`k_{mu}, k_{lambda}, k_{damp}`, shape [tet_count, 3]
+
+        body_q (array): Poses of rigid bodies used for state initialization, shape [body_count, 7], float
+        body_qd (array): Velocities of rigid bodies used for state initialization, shape [body_count, 6], float
+        body_com (array): Rigid body center of mass (in local frame), shape [body_count, 7], float
+        body_inertia (array): Rigid body inertia tensor (relative to COM), shape [body_count, 3, 3], float
+        body_inv_inertia (array): Rigid body inverse inertia tensor (relative to COM), shape [body_count, 3, 3], float
+        body_mass (array): Rigid body mass, shape [body_count], float
+        body_inv_mass (array): Rigid body inverse mass, shape [body_count], float
         body_name (list): Rigid body names, shape [body_count], str
 
-        joint_q (wp.array): Generalized joint positions used for state initialization, shape [joint_coord_count], float
-        joint_qd (wp.array): Generalized joint velocities used for state initialization, shape [joint_dof_count], float
-        joint_act (wp.array): Generalized joint actuation force, shape [joint_dof_count], float
-        joint_type (wp.array): Joint type, shape [joint_count], int
-        joint_parent (wp.array): Joint parent body indices, shape [joint_count], int
-        joint_child (wp.array): Joint child body indices, shape [joint_count], int
-        joint_X_p (wp.array): Joint transform in parent frame, shape [joint_count, 7], float
-        joint_X_c (wp.array): Joint mass frame in child frame, shape [joint_count, 7], float
-        joint_axis (wp.array): Joint axis in child frame, shape [joint_axis_count, 3], float
-        joint_armature (wp.array): Armature for each joint, shape [joint_count], float
-        joint_target (wp.array): Joint target position/velocity (depending on joint axis mode), shape [joint_axis_count], float
-        joint_target_ke (wp.array): Joint stiffness, shape [joint_axis_count], float
-        joint_target_kd (wp.array): Joint damping, shape [joint_axis_count], float
-        joint_axis_start (wp.array): Start index of the first axis per joint, shape [joint_count], int
-        joint_axis_dim (wp.array): Number of linear and angular axes per joint, shape [joint_count, 2], int
-        joint_axis_mode (wp.array): Joint axis mode, shape [joint_axis_count], int
-        joint_linear_compliance (wp.array): Joint linear compliance, shape [joint_count], float
-        joint_angular_compliance (wp.array): Joint linear compliance, shape [joint_count], float
-        joint_enabled (wp.array): Joint enabled, shape [joint_count], int
-        joint_limit_lower (wp.array): Joint lower position limits, shape [joint_count], float
-        joint_limit_upper (wp.array): Joint upper position limits, shape [joint_count], float
-        joint_limit_ke (wp.array): Joint position limit stiffness (used by SemiImplicitIntegrator), shape [joint_count], float
-        joint_limit_kd (wp.array): Joint position limit damping (used by SemiImplicitIntegrator), shape [joint_count], float
-        joint_twist_lower (wp.array): Joint lower twist limit, shape [joint_count], float
-        joint_twist_upper (wp.array): Joint upper twist limit, shape [joint_count], float
-        joint_q_start (wp.array): Start index of the first position coordinate per joint, shape [joint_count], int
-        joint_qd_start (wp.array): Start index of the first velocity coordinate per joint, shape [joint_count], int
+        joint_q (array): Generalized joint positions used for state initialization, shape [joint_coord_count], float
+        joint_qd (array): Generalized joint velocities used for state initialization, shape [joint_dof_count], float
+        joint_act (array): Generalized joint actuation force, shape [joint_dof_count], float
+        joint_type (array): Joint type, shape [joint_count], int
+        joint_parent (array): Joint parent body indices, shape [joint_count], int
+        joint_child (array): Joint child body indices, shape [joint_count], int
+        joint_X_p (array): Joint transform in parent frame, shape [joint_count, 7], float
+        joint_X_c (array): Joint mass frame in child frame, shape [joint_count, 7], float
+        joint_axis (array): Joint axis in child frame, shape [joint_axis_count, 3], float
+        joint_armature (array): Armature for each joint, shape [joint_count], float
+        joint_target (array): Joint target position/velocity (depending on joint axis mode), shape [joint_axis_count], float
+        joint_target_ke (array): Joint stiffness, shape [joint_axis_count], float
+        joint_target_kd (array): Joint damping, shape [joint_axis_count], float
+        joint_axis_start (array): Start index of the first axis per joint, shape [joint_count], int
+        joint_axis_dim (array): Number of linear and angular axes per joint, shape [joint_count, 2], int
+        joint_axis_mode (array): Joint axis mode, shape [joint_axis_count], int
+        joint_linear_compliance (array): Joint linear compliance, shape [joint_count], float
+        joint_angular_compliance (array): Joint linear compliance, shape [joint_count], float
+        joint_enabled (array): Joint enabled, shape [joint_count], int
+        joint_limit_lower (array): Joint lower position limits, shape [joint_count], float
+        joint_limit_upper (array): Joint upper position limits, shape [joint_count], float
+        joint_limit_ke (array): Joint position limit stiffness (used by SemiImplicitIntegrator), shape [joint_count], float
+        joint_limit_kd (array): Joint position limit damping (used by SemiImplicitIntegrator), shape [joint_count], float
+        joint_twist_lower (array): Joint lower twist limit, shape [joint_count], float
+        joint_twist_upper (array): Joint upper twist limit, shape [joint_count], float
+        joint_q_start (array): Start index of the first position coordinate per joint, shape [joint_count], int
+        joint_qd_start (array): Start index of the first velocity coordinate per joint, shape [joint_count], int
+        articulation_start (array): Articulation start index, shape [articulation_count], int
         joint_name (list): Joint names, shape [joint_count], str
         joint_attach_ke (float): Joint attachment force stiffness (used by SemiImplicitIntegrator)
         joint_attach_kd (float): Joint attachment force damping (used by SemiImplicitIntegrator)
 
-        articulation_start (wp.array): Articulation start index, shape [articulation_count], int
-
         soft_contact_margin (float): Contact margin for generation of soft contacts
         soft_contact_ke (float): Stiffness of soft contacts (used by SemiImplicitIntegrator)
         soft_contact_kd (float): Damping of soft contacts (used by SemiImplicitIntegrator)
@@ -406,7 +509,7 @@ class Model:
         rigid_contact_rolling_friction (float): Rolling friction coefficient for rigid body contacts (used by XPBDIntegrator)
 
         ground (bool): Whether the ground plane and ground contacts are enabled
-        ground_plane (wp.array): Ground plane 3D normal and offset, shape [4], float
+        ground_plane (array): Ground plane 3D normal and offset, shape [4], float
         up_vector (np.ndarray): Up vector of the world, shape [3], float
         up_axis (int): Up axis, 0 for x, 1 for y, 2 for z
         gravity (np.ndarray): Gravity vector, shape [3], float
@@ -451,6 +554,7 @@ class Model:
         self.particle_adhesion = 0.0
         self.particle_grid = None
         self.particle_flags = None
+        self.particle_max_velocity = 1e5
 
         self.shape_transform = None
         self.shape_body = None
@@ -522,6 +626,8 @@ class Model:
         self.joint_twist_lower = None
         self.joint_twist_upper = None
         self.joint_q_start = None
+        self.joint_qd_start = None
+        self.articulation_start = None
         self.joint_name = None
 
         # todo: per-joint values?
@@ -562,7 +668,7 @@ class Model:
 
         self.device = wp.get_device(device)
 
-    def state(self, requires_grad=False) -> State:
+    def state(self, requires_grad=None) -> State:
         """Returns a state object for the model
 
         The returned state will be initialized with the initial configuration given in
@@ -570,6 +676,8 @@ class Model:
         """
 
         s = State()
+        if requires_grad is None:
+            requires_grad = self.requires_grad
 
         s.particle_count = self.particle_count
         s.body_count = self.body_count
@@ -577,15 +685,6 @@ class Model:
         # --------------------------------
         # dynamic state (input, output)
 
-        s.particle_q = None
-        s.particle_qd = None
-        s.particle_f = None
-
-        s.body_q = None
-        s.body_qd = None
-        s.body_f = None
-        s.body_deltas = None
-
         # particles
         if self.particle_count:
             s.particle_q = wp.clone(self.particle_q)
@@ -622,8 +721,8 @@ class Model:
 
     def find_shape_contact_pairs(self):
         # find potential contact pairs based on collision groups and collision mask (pairwise filtering)
-        import itertools
         import copy
+        import itertools
 
         filters = copy.copy(self.shape_collision_filter_pairs)
         for a, b in self.shape_collision_filter_pairs:
@@ -775,9 +874,7 @@ class Model:
         return tensors
 
     def collide(self, state: State):
-        import warnings
-
-        warnings.warn(
+        wp.utils.warn(
             "Model.collide() is not needed anymore and will be removed in a future Warp version.",
             DeprecationWarning,
             stacklevel=2,
@@ -790,9 +887,7 @@ class Model:
 
     @property
     def soft_contact_distance(self):
-        import warnings
-
-        warnings.warn(
+        wp.utils.warn(
             "Model.soft_contact_distance is deprecated and will be removed in a future Warp version. "
             "Particles now have individual radii, returning `Model.particle_max_radius`.",
             DeprecationWarning,
@@ -802,9 +897,7 @@ class Model:
 
     @soft_contact_distance.setter
     def soft_contact_distance(self, value):
-        import warnings
-
-        warnings.warn(
+        wp.utils.warn(
             "Model.soft_contact_distance is deprecated and will be removed in a future Warp version. "
             "Particles now have individual radii, see `Model.particle_radius`.",
             DeprecationWarning,
@@ -813,18 +906,16 @@ class Model:
 
     @property
     def particle_radius(self):
-        """Array of per-particle radii"""
+        # Array of per-particle radii
         return self._particle_radius
 
     @particle_radius.setter
     def particle_radius(self, value):
         if isinstance(value, float):
-            import warnings
-
-            warnings.warn(
+            wp.utils.warn(
                 "Model.particle_radius is an array of per-particle radii, assigning with a scalar value "
                 "is deprecated and will be removed in a future Warp version.",
-                DeprecationWarning,
+                PendingDeprecationWarning,
                 stacklevel=2,
             )
             self._particle_radius.fill_(value)
@@ -891,6 +982,10 @@ class ModelBuilder:
     default_tri_drag = 0.0
     default_tri_lift = 0.0
 
+    # Default distance constraint properties
+    default_spring_ke = 100.0
+    default_spring_kd = 0.0
+
     # Default edge bending properties
     default_edge_ke = 100.0
     default_edge_kd = 0.0
@@ -919,6 +1014,7 @@ class ModelBuilder:
         self.particle_mass = []
         self.particle_radius = []
         self.particle_flags = []
+        self.particle_max_velocity = 1e5
 
         # shapes (each shape has an entry in these arrays)
         # transform from shape to body
@@ -1034,8 +1130,8 @@ class ModelBuilder:
         self.joint_coord_count = 0
         self.joint_axis_total_count = 0
 
-        self.up_vector = up_vector
-        self.up_axis = np.argmax(np.abs(up_vector))
+        self.up_vector = wp.vec3(up_vector)
+        self.up_axis = wp.vec3(np.argmax(np.abs(up_vector)))
         self.gravity = gravity
         # indicates whether a ground plane has been created
         self._ground_created = False
@@ -1126,15 +1222,25 @@ class ModelBuilder:
         """Copies a rigid articulation from `articulation`, another `ModelBuilder`.
 
         Args:
-            articulation: a model builder to add rigid articulation from.
-            xform: offset transform applied to root bodies.
-            update_num_env_count: if True, the number of environments is incremented by 1.
-            separate_collision_group: if True, the shapes from the articulation will all be put into a single new collision group, otherwise, only the shapes in collision group > -1 will be moved to a new group.
+            articulation (ModelBuilder): a model builder to add rigid articulation from.
+            xform (:ref:`transform <transform>`): offset transform applied to root bodies.
+            update_num_env_count (bool): if True, the number of environments is incremented by 1.
+            separate_collision_group (bool): if True, the shapes from the articulation will all be put into a single new collision group, otherwise, only the shapes in collision group > -1 will be moved to a new group.
         """
 
         start_body_idx = self.body_count
         start_shape_idx = self.shape_count
-        self.shape_body.extend([b + start_body_idx for b in articulation.shape_body])
+        for s, b in enumerate(articulation.shape_body):
+            if b > -1:
+                new_b = b + start_body_idx
+                self.shape_body.append(new_b)
+                self.shape_transform.append(articulation.shape_transform[s])
+            else:
+                self.shape_body.append(-1)
+                # apply offset transform to root bodies
+                if xform is not None:
+                    self.shape_transform.append(xform * articulation.shape_transform[s])
+
         for b, shapes in articulation.body_shapes.items():
             self.body_shapes[b + start_body_idx] = [s + start_shape_idx for s in shapes]
 
@@ -1227,7 +1333,6 @@ class ModelBuilder:
             "joint_target_kd",
             "joint_linear_compliance",
             "joint_angular_compliance",
-            "shape_transform",
             "shape_geo_type",
             "shape_geo_scale",
             "shape_geo_src",
@@ -1260,8 +1365,8 @@ class ModelBuilder:
         self,
         origin: Transform = wp.transform(),
         armature: float = 0.0,
-        com: Vec3 = np.zeros(3),
-        I_m: Mat33 = np.zeros((3, 3)),
+        com: Vec3 = wp.vec3(),
+        I_m: Mat33 = wp.mat33(),
         m: float = 0.0,
         name: str = None,
     ) -> int:
@@ -1286,7 +1391,7 @@ class ModelBuilder:
         body_id = len(self.body_mass)
 
         # body data
-        inertia = I_m + np.eye(3) * armature
+        inertia = I_m + wp.mat33(np.eye(3)) * armature
         self.body_inertia.append(inertia)
         self.body_mass.append(m)
         self.body_com.append(com)
@@ -1296,8 +1401,8 @@ class ModelBuilder:
         else:
             self.body_inv_mass.append(0.0)
 
-        if inertia.any():
-            self.body_inv_inertia.append(np.linalg.inv(inertia))
+        if any(x for x in inertia):
+            self.body_inv_inertia.append(wp.inverse(inertia))
         else:
             self.body_inv_inertia.append(inertia)
 
@@ -1323,6 +1428,29 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
+        """
+        Generic method to add any type of joint to this ModelBuilder.
+
+        Args:
+            joint_type: The type of joint to add (see `Joint types`_)
+            parent: The index of the parent body (-1 is the world)
+            child: The index of the child body
+            linear_axes: The linear axes (see :class:`JointAxis`) of the joint
+            angular_axes: The angular axes (see :class:`JointAxis`) of the joint
+            name: The name of the joint
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
+            linear_compliance: The linear compliance of the joint
+            angular_compliance: The angular compliance of the joint
+            collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies
+            enabled: Whether the joint is enabled
+
+        Returns:
+            The index of the added joint
+        """
+        if len(self.articulation_start) == 0:
+            # automatically add an articulation if none exists
+            self.add_articulation()
         self.joint_type.append(joint_type)
         self.joint_parent.append(parent)
         if child not in self.joint_parents:
@@ -1330,8 +1458,8 @@ class ModelBuilder:
         else:
             self.joint_parents[child].append(parent)
         self.joint_child.append(child)
-        self.joint_X_p.append([*parent_xform])
-        self.joint_X_c.append([*child_xform])
+        self.joint_X_p.append(wp.transform(parent_xform))
+        self.joint_X_c.append(wp.transform(child_xform))
         self.joint_name.append(name or f"joint {self.joint_count}")
         self.joint_axis_start.append(len(self.joint_axis))
         self.joint_axis_dim.append((len(linear_axes), len(angular_axes)))
@@ -1341,7 +1469,7 @@ class ModelBuilder:
         self.joint_angular_compliance.append(angular_compliance)
         self.joint_enabled.append(enabled)
 
-        for dim in linear_axes:
+        def add_axis_dim(dim):
             self.joint_axis.append(dim.axis)
             self.joint_axis_mode.append(dim.mode)
             self.joint_target.append(dim.target)
@@ -1349,18 +1477,21 @@ class ModelBuilder:
             self.joint_target_kd.append(dim.target_kd)
             self.joint_limit_ke.append(dim.limit_ke)
             self.joint_limit_kd.append(dim.limit_kd)
-            self.joint_limit_lower.append(dim.limit_lower)
-            self.joint_limit_upper.append(dim.limit_upper)
+            if np.isfinite(dim.limit_lower):
+                self.joint_limit_lower.append(dim.limit_lower)
+            else:
+                self.joint_limit_lower.append(-1e6)
+            if np.isfinite(dim.limit_upper):
+                self.joint_limit_upper.append(dim.limit_upper)
+            else:
+                self.joint_limit_upper.append(1e6)
+            # self.joint_limit_lower.append(dim.limit_lower)
+            # self.joint_limit_upper.append(dim.limit_upper)
+
+        for dim in linear_axes:
+            add_axis_dim(dim)
         for dim in angular_axes:
-            self.joint_axis.append(dim.axis)
-            self.joint_axis_mode.append(dim.mode)
-            self.joint_target.append(dim.target)
-            self.joint_target_ke.append(dim.target_ke)
-            self.joint_target_kd.append(dim.target_kd)
-            self.joint_limit_ke.append(dim.limit_ke)
-            self.joint_limit_kd.append(dim.limit_kd)
-            self.joint_limit_lower.append(dim.limit_lower)
-            self.joint_limit_upper.append(dim.limit_upper)
+            add_axis_dim(dim)
 
         if joint_type == JOINT_PRISMATIC:
             dof_count = 1
@@ -1371,7 +1502,7 @@ class ModelBuilder:
         elif joint_type == JOINT_BALL:
             dof_count = 3
             coord_count = 4
-        elif joint_type == JOINT_FREE:
+        elif joint_type == JOINT_FREE or joint_type == JOINT_DISTANCE:
             dof_count = 6
             coord_count = 7
         elif joint_type == JOINT_FIXED:
@@ -1383,10 +1514,6 @@ class ModelBuilder:
         elif joint_type == JOINT_COMPOUND:
             dof_count = 3
             coord_count = 3
-        elif joint_type == JOINT_DISTANCE:
-            # todo use free joint dofs?
-            dof_count = 0
-            coord_count = 0
         elif joint_type == JOINT_D6:
             dof_count = len(linear_axes) + len(angular_axes)
             coord_count = dof_count
@@ -1398,7 +1525,7 @@ class ModelBuilder:
             self.joint_qd.append(0.0)
             self.joint_act.append(0.0)
 
-        if joint_type == JOINT_FREE or joint_type == JOINT_BALL:
+        if joint_type == JOINT_FREE or joint_type == JOINT_DISTANCE or joint_type == JOINT_BALL:
             # ensure that a valid quaternion is used for the angular dofs
             self.joint_q[-1] = 1.0
 
@@ -1436,14 +1563,14 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a revolute joint to the model
+        """Adds a revolute (hinge) joint to the model. It has one degree of freedom.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            child_xform: The transform of the joint in the child body's local frame
-            axis: The axis of rotation in the parent body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
+            axis (3D vector or JointAxis): The axis of rotation in the parent body's local frame, can be a JointAxis object whose settings will be used instead of the other arguments
             target: The target angle (in radians) of the joint
             target_ke: The stiffness of the joint target
             target_kd: The damping of the joint target
@@ -1458,7 +1585,7 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         ax = JointAxis(
@@ -1507,14 +1634,14 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a prismatic joint to the model
+        """Adds a prismatic (sliding) joint to the model. It has one degree of freedom.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            child_xform: The transform of the joint in the child body's local frame
-            axis: The axis of rotation in the parent body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
+            axis (3D vector or JointAxis): The axis of rotation in the parent body's local frame, can be a JointAxis object whose settings will be used instead of the other arguments
             target: The target position of the joint
             target_ke: The stiffness of the joint target
             target_kd: The damping of the joint target
@@ -1529,7 +1656,7 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         ax = JointAxis(
@@ -1569,13 +1696,13 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a ball joint to the model
+        """Adds a ball (spherical) joint to the model. Its position is defined by a 4D quaternion (xyzw) and its velocity is a 3D vector.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             linear_compliance: The linear compliance of the joint
             angular_compliance: The angular compliance of the joint
             name: The name of the joint
@@ -1583,7 +1710,7 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
@@ -1611,13 +1738,14 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a fixed joint to the model
+        """Adds a fixed (static) joint to the model. It has no degrees of freedom.
+        See :meth:`collapse_fixed_joints` for a helper function that removes these fixed joints and merges the connecting bodies to simplify the model and improve stability.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             linear_compliance: The linear compliance of the joint
             angular_compliance: The angular compliance of the joint
             name: The name of the joint
@@ -1625,7 +1753,7 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
@@ -1651,19 +1779,20 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a free joint to the model
+        """Adds a free joint to the model.
+        It has 7 positional degrees of freedom (first 3 linear and then 4 angular dimensions for the orientation quaternion in `xyzw` notation) and 6 velocity degrees of freedom (first 3 angular and then 3 linear velocity dimensions).
 
         Args:
-            parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
+            parent: The index of the parent body (-1 by default to use the world frame, e.g. to make the child body and its children a floating-base mechanism)
             name: The name of the joint
             collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
@@ -1689,13 +1818,14 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a distance joint to the model
+        """Adds a distance joint to the model. The distance joint constraints the distance between the joint anchor points on the two bodies (see :ref:`FK-IK`) it connects to the interval [`min_distance`, `max_distance`].
+        It has 7 positional degrees of freedom (first 3 linear and then 4 angular dimensions for the orientation quaternion in `xyzw` notation) and 6 velocity degrees of freedom (first 3 angular and then 3 linear velocity dimensions).
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             min_distance: The minimum distance between the bodies (no limit if negative)
             max_distance: The maximum distance between the bodies (no limit if negative)
             compliance: The compliance of the joint
@@ -1703,7 +1833,9 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
+        
+        .. note:: Distance joints are currently only supported in the :class:`XPBDIntegrator` at the moment.
 
         """
         ax = JointAxis(
@@ -1737,15 +1869,15 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a universal joint to the model
+        """Adds a universal joint to the model. U-joints have two degrees of freedom, one for each axis.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            axis_0: The first axis of the joint
-            axis_1: The second axis of the joint
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            axis_0 (3D vector or JointAxis): The first axis of the joint, can be a JointAxis object whose settings will be used instead of the other arguments
+            axis_1 (3D vector or JointAxis): The second axis of the joint, can be a JointAxis object whose settings will be used instead of the other arguments
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             linear_compliance: The linear compliance of the joint
             angular_compliance: The angular compliance of the joint
             name: The name of the joint
@@ -1753,14 +1885,14 @@ class ModelBuilder:
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
             JOINT_UNIVERSAL,
             parent,
             child,
-            angular_axes=[axis_0, axis_1],
+            angular_axes=[JointAxis(axis_0), JointAxis(axis_1)],
             parent_xform=parent_xform,
             child_xform=child_xform,
             linear_compliance=linear_compliance,
@@ -1783,29 +1915,32 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ) -> int:
-        """Adds a compound joint to the model
+        """Adds a compound joint to the model, which has 3 degrees of freedom, one for each axis.
+        Similar to the ball joint (see :meth:`add_ball_joint`), the compound joint allows bodies to move in a 3D rotation relative to each other,
+        except that the rotation is defined by 3 axes instead of a quaternion.
+        Depending on the choice of axes, the orientation can be specified through Euler angles, e.g. `z-x-z` or `x-y-x`, or through a Tait-Bryan angle sequence, e.g. `z-y-x` or `x-y-z`.
 
         Args:
             parent: The index of the parent body
             child: The index of the child body
-            axis_0: The first axis of the joint
-            axis_1: The second axis of the joint
-            axis_2: The third axis of the joint
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            axis_0 (3D vector or JointAxis): The first axis of the joint, can be a JointAxis object whose settings will be used instead of the other arguments
+            axis_1 (3D vector or JointAxis): The second axis of the joint, can be a JointAxis object whose settings will be used instead of the other arguments
+            axis_2 (3D vector or JointAxis): The third axis of the joint, can be a JointAxis object whose settings will be used instead of the other arguments
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             name: The name of the joint
             collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
             JOINT_COMPOUND,
             parent,
             child,
-            angular_axes=[axis_0, axis_1, axis_2],
+            angular_axes=[JointAxis(axis_0), JointAxis(axis_1), JointAxis(axis_2)],
             parent_xform=parent_xform,
             child_xform=child_xform,
             name=name,
@@ -1827,7 +1962,7 @@ class ModelBuilder:
         collision_filter_parent: bool = True,
         enabled: bool = True,
     ):
-        """Adds a generic joint with custom linear and angular axes.
+        """Adds a generic joint with custom linear and angular axes. The number of axes determines the number of degrees of freedom of the joint.
 
         Args:
             parent: The index of the parent body
@@ -1835,15 +1970,15 @@ class ModelBuilder:
             linear_axes: A list of linear axes
             angular_axes: A list of angular axes
             name: The name of the joint
-            parent_xform: The transform of the joint in the parent body's local frame
-            xform: The transform of the joint in the child body's local frame
+            parent_xform (:ref:`transform <transform>`): The transform of the joint in the parent body's local frame
+            child_xform (:ref:`transform <transform>`): The transform of the joint in the child body's local frame
             linear_compliance: The linear compliance of the joint
             angular_compliance: The angular compliance of the joint
             collision_filter_parent: Whether to filter collisions between shapes of the parent and child bodies
             enabled: Whether the joint is enabled
 
         Returns:
-            The index of the child body
+            The index of the added joint
 
         """
         return self.add_joint(
@@ -1852,8 +1987,8 @@ class ModelBuilder:
             child,
             parent_xform=parent_xform,
             child_xform=child_xform,
-            linear_axes=linear_axes,
-            angular_axes=angular_axes,
+            linear_axes=[JointAxis(a) for a in linear_axes],
+            angular_axes=[JointAxis(a) for a in angular_axes],
             linear_compliance=linear_compliance,
             angular_compliance=angular_compliance,
             name=name,
@@ -1861,11 +1996,257 @@ class ModelBuilder:
             enabled=enabled,
         )
 
+    def collapse_fixed_joints(self, verbose=wp.config.verbose):
+        """Removes fixed joints from the model and merges the bodies they connect. This is useful for simplifying the model for faster and more stable simulation."""
+
+        body_data = {}
+        body_children = {-1: []}
+        visited = {}
+        for i in range(self.body_count):
+            name = self.body_name[i]
+            body_data[i] = {
+                "shapes": self.body_shapes[i],
+                "q": self.body_q[i],
+                "qd": self.body_qd[i],
+                "mass": self.body_mass[i],
+                "inertia": self.body_inertia[i],
+                "inv_mass": self.body_inv_mass[i],
+                "inv_inertia": self.body_inv_inertia[i],
+                "com": self.body_com[i],
+                "name": name,
+                "original_id": i,
+            }
+            visited[i] = False
+            body_children[i] = []
+
+        joint_data = {}
+        for i in range(self.joint_count):
+            name = self.joint_name[i]
+            parent = self.joint_parent[i]
+            child = self.joint_child[i]
+            body_children[parent].append(child)
+
+            q_start = self.joint_q_start[i]
+            qd_start = self.joint_qd_start[i]
+            if i < self.joint_count - 1:
+                q_dim = self.joint_q_start[i + 1] - q_start
+                qd_dim = self.joint_qd_start[i + 1] - qd_start
+            else:
+                q_dim = len(self.joint_q) - q_start
+                qd_dim = len(self.joint_qd) - qd_start
+
+            data = {
+                "type": self.joint_type[i],
+                # 'armature': self.joint_armature[i],
+                "q": self.joint_q[q_start : q_start + q_dim],
+                "qd": self.joint_qd[qd_start : qd_start + qd_dim],
+                "act": self.joint_act[qd_start : qd_start + qd_dim],
+                "q_start": q_start,
+                "qd_start": qd_start,
+                "linear_compliance": self.joint_linear_compliance[i],
+                "angular_compliance": self.joint_angular_compliance[i],
+                "name": name,
+                "parent_xform": wp.transform_expand(self.joint_X_p[i]),
+                "child_xform": wp.transform_expand(self.joint_X_c[i]),
+                "enabled": self.joint_enabled[i],
+                "axes": [],
+                "axis_dim": self.joint_axis_dim[i],
+                "parent": parent,
+                "child": child,
+                "original_id": i,
+            }
+            num_lin_axes, num_ang_axes = self.joint_axis_dim[i]
+            start_ax = self.joint_axis_start[i]
+            for j in range(start_ax, start_ax + num_lin_axes + num_ang_axes):
+                data["axes"].append(
+                    {
+                        "axis": self.joint_axis[j],
+                        "axis_mode": self.joint_axis_mode[j],
+                        "target": self.joint_target[j],
+                        "target_ke": self.joint_target_ke[j],
+                        "target_kd": self.joint_target_kd[j],
+                        "limit_ke": self.joint_limit_ke[j],
+                        "limit_kd": self.joint_limit_kd[j],
+                        "limit_lower": self.joint_limit_lower[j],
+                        "limit_upper": self.joint_limit_upper[j],
+                    }
+                )
+
+            joint_data[(parent, child)] = data
+
+        # sort body children so we traverse the tree in the same order as the bodies are listed
+        for children in body_children.values():
+            children.sort(key=lambda x: body_data[x]["original_id"])
+
+        retained_joints = []
+        retained_bodies = []
+        body_remap = {-1: -1}
+
+        # depth first search over the joint graph
+        def dfs(parent_body: int, child_body: int, incoming_xform: wp.transform, last_dynamic_body: int):
+            nonlocal visited
+            nonlocal retained_joints
+            nonlocal retained_bodies
+            nonlocal body_data
+            nonlocal body_remap
+
+            joint = joint_data[(parent_body, child_body)]
+            if joint["type"] == JOINT_FIXED:
+                joint_xform = joint["parent_xform"] * wp.transform_inverse(joint["child_xform"])
+                incoming_xform = incoming_xform * joint_xform
+                parent_name = self.body_name[parent_body] if parent_body > -1 else "world"
+                child_name = self.body_name[child_body]
+                last_dynamic_body_name = self.body_name[last_dynamic_body] if last_dynamic_body > -1 else "world"
+                if verbose:
+                    print(
+                        f'Remove fixed joint {joint["name"]} between {parent_name} and {child_name}, '
+                        f"merging {child_name} into {last_dynamic_body_name}"
+                    )
+                child_id = body_data[child_body]["original_id"]
+                for shape in self.body_shapes[child_id]:
+                    self.shape_transform[shape] = incoming_xform * self.shape_transform[shape]
+                    if verbose:
+                        print(
+                            f"  Shape {shape} moved to body {last_dynamic_body_name} with transform {self.shape_transform[shape]}"
+                        )
+                    if last_dynamic_body > -1:
+                        self.shape_body[shape] = body_data[last_dynamic_body]["id"]
+                        # add inertia to last_dynamic_body
+                        m = body_data[child_body]["mass"]
+                        com = body_data[child_body]["com"]
+                        inertia = body_data[child_body]["inertia"]
+                        body_data[last_dynamic_body]["inertia"] += wp.sim.transform_inertia(
+                            m, inertia, incoming_xform.p, incoming_xform.q
+                        )
+                        body_data[last_dynamic_body]["mass"] += m
+                        source_m = body_data[last_dynamic_body]["mass"]
+                        source_com = body_data[last_dynamic_body]["com"]
+                        body_data[last_dynamic_body]["com"] = (m * com + source_m * source_com) / (m + source_m)
+                        body_data[last_dynamic_body]["shapes"].append(shape)
+                        # indicate to recompute inverse mass, inertia for this body
+                        body_data[last_dynamic_body]["inv_mass"] = None
+                    else:
+                        self.shape_body[shape] = -1
+            else:
+                joint["parent_xform"] = incoming_xform * joint["parent_xform"]
+                joint["parent"] = last_dynamic_body
+                last_dynamic_body = child_body
+                incoming_xform = wp.transform()
+                retained_joints.append(joint)
+                new_id = len(retained_bodies)
+                body_data[child_body]["id"] = new_id
+                retained_bodies.append(child_body)
+                for shape in body_data[child_body]["shapes"]:
+                    self.shape_body[shape] = new_id
+
+            visited[parent_body] = True
+            if visited[child_body] or child_body not in body_children:
+                return
+            for child in body_children[child_body]:
+                if not visited[child]:
+                    dfs(child_body, child, incoming_xform, last_dynamic_body)
+
+        for body in body_children[-1]:
+            if not visited[body]:
+                dfs(-1, body, wp.transform(), -1)
+
+        # repopulate the model
+        self.body_name.clear()
+        self.body_q.clear()
+        self.body_qd.clear()
+        self.body_mass.clear()
+        self.body_inertia.clear()
+        self.body_com.clear()
+        self.body_inv_mass.clear()
+        self.body_inv_inertia.clear()
+        self.body_shapes.clear()
+        for i in retained_bodies:
+            body = body_data[i]
+            new_id = len(self.body_name)
+            body_remap[body["original_id"]] = new_id
+            self.body_name.append(body["name"])
+            self.body_q.append(list(body["q"]))
+            self.body_qd.append(list(body["qd"]))
+            m = body["mass"]
+            inertia = body["inertia"]
+            self.body_mass.append(m)
+            self.body_inertia.append(inertia)
+            self.body_com.append(body["com"])
+            if body["inv_mass"] is None:
+                # recompute inverse mass and inertia
+                if m > 0.0:
+                    self.body_inv_mass.append(1.0 / m)
+                    self.body_inv_inertia.append(np.linalg.inv(inertia))
+                else:
+                    self.body_inv_mass.append(0.0)
+                    self.body_inv_inertia.append(np.zeros((3, 3)))
+            else:
+                self.body_inv_mass.append(body["inv_mass"])
+                self.body_inv_inertia.append(body["inv_inertia"])
+            self.body_shapes[new_id] = body["shapes"]
+            body_remap[body["original_id"]] = new_id
+
+        # sort joints so they appear in the same order as before
+        retained_joints.sort(key=lambda x: x["original_id"])
+
+        self.joint_name.clear()
+        self.joint_type.clear()
+        self.joint_parent.clear()
+        self.joint_child.clear()
+        self.joint_q.clear()
+        self.joint_qd.clear()
+        self.joint_q_start.clear()
+        self.joint_qd_start.clear()
+        self.joint_enabled.clear()
+        self.joint_linear_compliance.clear()
+        self.joint_angular_compliance.clear()
+        self.joint_X_p.clear()
+        self.joint_X_c.clear()
+        self.joint_axis.clear()
+        self.joint_axis_mode.clear()
+        self.joint_target.clear()
+        self.joint_target_ke.clear()
+        self.joint_target_kd.clear()
+        self.joint_limit_lower.clear()
+        self.joint_limit_upper.clear()
+        self.joint_limit_ke.clear()
+        self.joint_limit_kd.clear()
+        self.joint_axis_dim.clear()
+        self.joint_axis_start.clear()
+        self.joint_act.clear()
+        for joint in retained_joints:
+            self.joint_name.append(joint["name"])
+            self.joint_type.append(joint["type"])
+            self.joint_parent.append(body_remap[joint["parent"]])
+            self.joint_child.append(body_remap[joint["child"]])
+            self.joint_q_start.append(len(self.joint_q))
+            self.joint_qd_start.append(len(self.joint_qd))
+            self.joint_q.extend(joint["q"])
+            self.joint_qd.extend(joint["qd"])
+            self.joint_act.extend(joint["act"])
+            self.joint_enabled.append(joint["enabled"])
+            self.joint_linear_compliance.append(joint["linear_compliance"])
+            self.joint_angular_compliance.append(joint["angular_compliance"])
+            self.joint_X_p.append(list(joint["parent_xform"]))
+            self.joint_X_c.append(list(joint["child_xform"]))
+            self.joint_axis_dim.append(joint["axis_dim"])
+            self.joint_axis_start.append(len(self.joint_axis))
+            for axis in joint["axes"]:
+                self.joint_axis.append(axis["axis"])
+                self.joint_axis_mode.append(axis["axis_mode"])
+                self.joint_target.append(axis["target"])
+                self.joint_target_ke.append(axis["target_ke"])
+                self.joint_target_kd.append(axis["target_kd"])
+                self.joint_limit_lower.append(axis["limit_lower"])
+                self.joint_limit_upper.append(axis["limit_upper"])
+                self.joint_limit_ke.append(axis["limit_ke"])
+                self.joint_limit_kd.append(axis["limit_kd"])
+
     # muscles
     def add_muscle(
         self, bodies: List[int], positions: List[Vec3], f0: float, lm: float, lt: float, lmax: float, pen: float
     ) -> float:
-        """Adds a muscle-tendon activation unit
+        """Adds a muscle-tendon activation unit.
 
         Args:
             bodies: A list of body indices for each waypoint
@@ -1878,6 +2259,8 @@ class ModelBuilder:
         Returns:
             The index of the muscle in the model
 
+        .. note:: The simulation support for muscles is in progress and not yet fully functional.
+
         """
 
         n = len(bodies)
@@ -1913,7 +2296,7 @@ class ModelBuilder:
         """
         Adds a plane collision shape.
         If pos and rot are defined, the plane is assumed to have its normal as (0, 1, 0).
-        Otherwise, the plane equation is used.
+        Otherwise, the plane equation defined through the `plane` argument is used.
 
         Args:
             plane: The plane equation in form a*x + b*y + c*z + d = 0
@@ -1930,6 +2313,9 @@ class ModelBuilder:
             thickness: The thickness of the plane (0 by default) for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
         if pos is None or rot is None:
             # compute position and rotation from plane equation
@@ -1944,7 +2330,7 @@ class ModelBuilder:
                 angle = np.arcsin(np.linalg.norm(c))
                 axis = c / np.linalg.norm(c)
                 rot = wp.quat_from_axis_angle(axis, angle)
-        scale = (width, length, 0.0)
+        scale = wp.vec3(width, length, 0.0)
 
         return self._add_shape(
             body,
@@ -1982,7 +2368,7 @@ class ModelBuilder:
         """Adds a sphere collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             radius: The radius of the sphere
@@ -1996,14 +2382,17 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow sphere, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
         return self._add_shape(
             body,
-            pos,
-            rot,
+            wp.vec3(pos),
+            wp.quat(rot),
             GEO_SPHERE,
-            (radius, 0.0, 0.0, 0.0),
+            wp.vec3(radius, 0.0, 0.0),
             None,
             density,
             ke,
@@ -2037,7 +2426,7 @@ class ModelBuilder:
         """Adds a box collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             hx: The half-extent along the x-axis
@@ -2053,14 +2442,17 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow box, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
         return self._add_shape(
             body,
-            pos,
-            rot,
+            wp.vec3(pos),
+            wp.quat(rot),
             GEO_BOX,
-            (hx, hy, hz, 0.0),
+            wp.vec3(hx, hy, hz),
             None,
             density,
             ke,
@@ -2094,7 +2486,7 @@ class ModelBuilder:
         """Adds a capsule collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             radius: The radius of the capsule
@@ -2110,21 +2502,24 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow capsule, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
-        q = rot
+        q = wp.quat(rot)
         sqh = math.sqrt(0.5)
         if up_axis == 0:
-            q = wp.mul(rot, wp.quat(0.0, 0.0, -sqh, sqh))
+            q = wp.mul(q, wp.quat(0.0, 0.0, -sqh, sqh))
         elif up_axis == 2:
-            q = wp.mul(rot, wp.quat(sqh, 0.0, 0.0, sqh))
+            q = wp.mul(q, wp.quat(sqh, 0.0, 0.0, sqh))
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CAPSULE,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2158,7 +2553,7 @@ class ModelBuilder:
         """Adds a cylinder collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             radius: The radius of the cylinder
@@ -2174,6 +2569,9 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow cylinder, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
         q = rot
@@ -2185,10 +2583,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CYLINDER,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2222,7 +2620,7 @@ class ModelBuilder:
         """Adds a cone collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             radius: The radius of the cone
@@ -2238,6 +2636,9 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow cone, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
         q = rot
@@ -2249,10 +2650,10 @@ class ModelBuilder:
 
         return self._add_shape(
             body,
-            pos,
-            q,
+            wp.vec3(pos),
+            wp.quat(q),
             GEO_CONE,
-            (radius, half_height, 0.0, 0.0),
+            wp.vec3(radius, half_height, 0.0),
             None,
             density,
             ke,
@@ -2268,10 +2669,10 @@ class ModelBuilder:
     def add_shape_mesh(
         self,
         body: int,
-        pos: Vec3 = (0.0, 0.0, 0.0),
-        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        pos: Vec3 = wp.vec3(0.0, 0.0, 0.0),
+        rot: Quat = wp.quat(0.0, 0.0, 0.0, 1.0),
         mesh: Mesh = None,
-        scale: Vec3 = (1.0, 1.0, 1.0),
+        scale: Vec3 = wp.vec3(1.0, 1.0, 1.0),
         density: float = default_shape_density,
         ke: float = default_shape_ke,
         kd: float = default_shape_kd,
@@ -2285,7 +2686,7 @@ class ModelBuilder:
         """Adds a triangle mesh collision shape to a body.
 
         Args:
-            body: The index of the parent body this shape belongs to
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
             pos: The location of the shape with respect to the parent frame
             rot: The rotation of the shape with respect to the parent frame
             mesh: The mesh object
@@ -2300,6 +2701,9 @@ class ModelBuilder:
             thickness: Thickness to use for computing inertia of a hollow mesh, and for collision handling
             has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
 
+        Returns:
+            The index of the added shape
+
         """
 
         return self._add_shape(
@@ -2307,7 +2711,7 @@ class ModelBuilder:
             pos,
             rot,
             GEO_MESH,
-            (scale[0], scale[1], scale[2], 0.0),
+            wp.vec3(scale[0], scale[1], scale[2]),
             mesh,
             density,
             ke,
@@ -2320,6 +2724,61 @@ class ModelBuilder:
             has_ground_collision=has_ground_collision,
         )
 
+    def add_shape_sdf(
+        self,
+        body: int,
+        pos: Vec3 = (0.0, 0.0, 0.0),
+        rot: Quat = (0.0, 0.0, 0.0, 1.0),
+        sdf: SDF = None,
+        scale: Vec3 = (1.0, 1.0, 1.0),
+        density: float = default_shape_density,
+        ke: float = default_shape_ke,
+        kd: float = default_shape_kd,
+        kf: float = default_shape_kf,
+        mu: float = default_shape_mu,
+        restitution: float = default_shape_restitution,
+        is_solid: bool = True,
+        thickness: float = default_geo_thickness,
+    ):
+        """Adds SDF collision shape to a body.
+
+        Args:
+            body: The index of the parent body this shape belongs to (use -1 for static shapes)
+            pos: The location of the shape with respect to the parent frame
+            rot: The rotation of the shape with respect to the parent frame
+            sdf: The sdf object
+            scale: Scale to use for the collider
+            density: The density of the shape
+            ke: The contact elastic stiffness
+            kd: The contact damping stiffness
+            kf: The contact friction stiffness
+            mu: The coefficient of friction
+            restitution: The coefficient of restitution
+            is_solid: If True, the mesh is solid, otherwise it is a hollow surface with the given wall thickness
+            thickness: Thickness to use for computing inertia of a hollow mesh, and for collision handling
+            has_ground_collision: If True, the mesh will collide with the ground plane if `Model.ground` is True
+
+        Returns:
+            The index of the added shape
+
+        """
+        return self._add_shape(
+            body,
+            wp.vec3(pos),
+            wp.quat(rot),
+            GEO_SDF,
+            wp.vec3(scale[0], scale[1], scale[2]),
+            sdf,
+            density,
+            ke,
+            kd,
+            kf,
+            mu,
+            restitution,
+            thickness,
+            is_solid,
+        )
+
     def _shape_radius(self, type, scale, src):
         """
         Calculates the radius of a sphere that encloses the shape, used for broadphase collision detection.
@@ -2399,7 +2858,7 @@ class ModelBuilder:
 
         (m, c, I) = compute_shape_mass(type, scale, src, density, is_solid, thickness)
 
-        self._update_body_mass(body, m, I, np.array(pos) + c, np.array(rot))
+        self._update_body_mass(body, m, I, pos + c, rot)
         return shape
 
     # particles
@@ -2494,9 +2953,9 @@ class ModelBuilder:
 
         """
         # compute basis for 2D rest pose
-        p = np.array(self.particle_q[i])
-        q = np.array(self.particle_q[j])
-        r = np.array(self.particle_q[k])
+        p = self.particle_q[i]
+        q = self.particle_q[j]
+        r = self.particle_q[k]
 
         qp = q - p
         rp = r - p
@@ -2634,7 +3093,7 @@ class ModelBuilder:
             The volume of the tetrahedron
 
         Note:
-            The tetrahedron is created with a rest-pose based on the particle's initial configruation
+            The tetrahedron is created with a rest-pose based on the particle's initial configuration
 
         """
         # compute basis for 2D rest pose
@@ -2693,13 +3152,13 @@ class ModelBuilder:
         """
         # compute rest angle
         if rest is None:
-            x1 = np.array(self.particle_q[i])
-            x2 = np.array(self.particle_q[j])
-            x3 = np.array(self.particle_q[k])
-            x4 = np.array(self.particle_q[l])
+            x1 = self.particle_q[i]
+            x2 = self.particle_q[j]
+            x3 = self.particle_q[k]
+            x4 = self.particle_q[l]
 
-            n1 = wp.normalize(np.cross(x3 - x1, x4 - x1))
-            n2 = wp.normalize(np.cross(x4 - x2, x3 - x2))
+            n1 = wp.normalize(wp.cross(x3 - x1, x4 - x1))
+            n2 = wp.normalize(wp.cross(x4 - x2, x3 - x2))
             e = wp.normalize(x4 - x3)
 
             d = np.clip(np.dot(n2, n1), -1.0, 1.0)
@@ -2805,6 +3264,9 @@ class ModelBuilder:
         tri_lift: float = default_tri_lift,
         edge_ke: float = default_edge_ke,
         edge_kd: float = default_edge_kd,
+        add_springs: bool = False,
+        spring_ke: float = default_spring_ke,
+        spring_kd: float = default_spring_kd,
     ):
         """Helper to create a regular planar cloth grid
 
@@ -2836,8 +3298,8 @@ class ModelBuilder:
 
         for y in range(0, dim_y + 1):
             for x in range(0, dim_x + 1):
-                g = np.array((x * cell_x, y * cell_y, 0.0))
-                p = np.array(wp.quat_rotate(rot, g)) + pos
+                g = wp.vec3(x * cell_x, y * cell_y, 0.0)
+                p = wp.quat_rotate(rot, g) + pos
                 m = mass
 
                 if x == 0 and fix_left:
@@ -2884,13 +3346,14 @@ class ModelBuilder:
                         self.add_triangle(*tri1, tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
                         self.add_triangle(*tri2, tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
 
-        end_vertex = len(self.particle_q)
         end_tri = len(self.tri_indices)
 
         # bending constraints, could create these explicitly for a grid but this
         # is a good test of the adjacency structure
         adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)
 
+        spring_indices = set()
+
         for k, e in adj.edges.items():
             # skip open edges
             if e.f0 == -1 or e.f1 == -1:
@@ -2900,6 +3363,19 @@ class ModelBuilder:
                 e.o0, e.o1, e.v0, e.v1, edge_ke=edge_ke, edge_kd=edge_kd
             )  # opposite 0, opposite 1, vertex 0, vertex 1
 
+            spring_indices.add((min(e.o0, e.o1), max(e.o0, e.o1)))
+            spring_indices.add((min(e.o0, e.v0), max(e.o0, e.v0)))
+            spring_indices.add((min(e.o0, e.v1), max(e.o0, e.v1)))
+
+            spring_indices.add((min(e.o1, e.v0), max(e.o1, e.v0)))
+            spring_indices.add((min(e.o1, e.v1), max(e.o1, e.v1)))
+
+            spring_indices.add((min(e.v0, e.v1), max(e.v0, e.v1)))
+
+        if add_springs:
+            for i, j in spring_indices:
+                self.add_spring(i, j, spring_ke, spring_kd, control=0.0)
+
     def add_cloth_mesh(
         self,
         pos: Vec3,
@@ -2918,6 +3394,9 @@ class ModelBuilder:
         tri_lift: float = default_tri_lift,
         edge_ke: float = default_edge_ke,
         edge_kd: float = default_edge_kd,
+        add_springs: bool = False,
+        spring_ke: float = default_spring_ke,
+        spring_kd: float = default_spring_kd,
     ):
         """Helper to create a cloth model from a regular triangle mesh
 
@@ -2945,7 +3424,7 @@ class ModelBuilder:
 
         # particles
         for v in vertices:
-            p = np.array(wp.quat_rotate(rot, v * scale)) + pos
+            p = wp.quat_rotate(rot, v * scale) + pos
 
             self.add_particle(p, vel, 0.0)
 
@@ -2974,7 +3453,10 @@ class ModelBuilder:
 
         adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)
 
-        edgeinds = np.array([[e.o0, e.o1, e.v0, e.v1] for k, e in adj.edges.items()])
+        edgeinds = np.fromiter(
+            (x for e in adj.edges.values() if e.f0 != -1 and e.f1 != -1 for x in (e.o0, e.o1, e.v0, e.v1)),
+            int,
+        ).reshape(-1, 4)
         self.add_edges(
             edgeinds[:, 0],
             edgeinds[:, 1],
@@ -2984,6 +3466,21 @@ class ModelBuilder:
             edge_kd=[edge_kd] * len(edgeinds),
         )
 
+        if add_springs:
+            spring_indices = set()
+            for i, j, k, l in edgeinds:
+                spring_indices.add((min(i, j), max(i, j)))
+                spring_indices.add((min(i, k), max(i, k)))
+                spring_indices.add((min(i, l), max(i, l)))
+
+                spring_indices.add((min(j, k), max(j, k)))
+                spring_indices.add((min(j, l), max(j, l)))
+
+                spring_indices.add((min(k, l), max(k, l)))
+
+            for i, j in spring_indices:
+                self.add_spring(i, j, spring_ke, spring_kd, control=0.0)
+
     def add_particle_grid(
         self,
         pos: Vec3,
@@ -3000,13 +3497,14 @@ class ModelBuilder:
         radius_mean: float = default_particle_radius,
         radius_std: float = 0.0,
     ):
+        rng = np.random.default_rng()
         for z in range(dim_z):
             for y in range(dim_y):
                 for x in range(dim_x):
-                    v = np.array((x * cell_x, y * cell_y, z * cell_z))
+                    v = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                     m = mass
 
-                    p = np.array(wp.quat_rotate(rot, v)) + pos + np.random.rand(3) * jitter
+                    p = wp.quat_rotate(rot, v) + pos + wp.vec3(rng.random(3) * jitter)
 
                     if radius_std > 0.0:
                         r = radius_mean + np.random.randn() * radius_std
@@ -3072,7 +3570,7 @@ class ModelBuilder:
         for z in range(dim_z + 1):
             for y in range(dim_y + 1):
                 for x in range(dim_x + 1):
-                    v = np.array((x * cell_x, y * cell_y, z * cell_z))
+                    v = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                     m = mass
 
                     if fix_left and x == 0:
@@ -3087,7 +3585,7 @@ class ModelBuilder:
                     if fix_bottom and y == 0:
                         m = 0.0
 
-                    p = np.array(wp.quat_rotate(rot, v)) + pos
+                    p = wp.quat_rotate(rot, v) + pos
 
                     self.add_particle(p, vel, m)
 
@@ -3177,7 +3675,6 @@ class ModelBuilder:
         num_tets = int(len(indices) / 4)
 
         start_vertex = len(self.particle_q)
-        start_tri = len(self.tri_indices)
 
         # dict of open faces
         faces = {}
@@ -3255,8 +3752,8 @@ class ModelBuilder:
         else:
             self.body_inv_mass[i] = 0.0
 
-        if new_inertia.any():
-            self.body_inv_inertia[i] = np.linalg.inv(new_inertia)
+        if any(x for x in new_inertia):
+            self.body_inv_inertia[i] = wp.inverse(new_inertia)
         else:
             self.body_inv_inertia[i] = new_inertia
 
@@ -3338,6 +3835,7 @@ class ModelBuilder:
             m._particle_radius = wp.array(self.particle_radius, dtype=wp.float32, requires_grad=requires_grad)
             m.particle_flags = wp.array([flag_to_int(f) for f in self.particle_flags], dtype=wp.uint32)
             m.particle_max_radius = np.max(self.particle_radius) if len(self.particle_radius) > 0 else 0.0
+            m.particle_max_velocity = self.particle_max_velocity
 
             # hash-grid for particle interactions
             m.particle_grid = wp.HashGrid(128, 128, 128)
@@ -3393,6 +3891,9 @@ class ModelBuilder:
             m.spring_stiffness = wp.array(self.spring_stiffness, dtype=wp.float32, requires_grad=requires_grad)
             m.spring_damping = wp.array(self.spring_damping, dtype=wp.float32, requires_grad=requires_grad)
             m.spring_control = wp.array(self.spring_control, dtype=wp.float32, requires_grad=requires_grad)
+            m.spring_constraint_lambdas = wp.array(
+                shape=len(self.spring_rest_length), dtype=wp.float32, requires_grad=requires_grad
+            )
 
             # ---------------------
             # triangles
@@ -3410,6 +3911,9 @@ class ModelBuilder:
             m.edge_bending_properties = wp.array(
                 self.edge_bending_properties, dtype=wp.float32, requires_grad=requires_grad
             )
+            m.edge_constraint_lambdas = wp.array(
+                shape=len(self.edge_rest_angle), dtype=wp.float32, requires_grad=requires_grad
+            )
 
             # ---------------------
             # tetrahedra