NvidiaWarp-GarmentCode/warp/sim/model.py

# Copyright (c) 2022 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.

"""A module for building simulation models and state.
"""

import copy
import math
from typing import List, Optional, Tuple, Union

import numpy as np

import warp as wp

from .inertia import (
    compute_box_inertia,
    compute_capsule_inertia,
    compute_cone_inertia,
    compute_cylinder_inertia,
    compute_mesh_inertia,
    compute_sphere_inertia,
    transform_inertia,
)

from .utils import compute_vertex_normal

Vec3 = List[float]
Vec4 = List[float]
Quat = List[float]
Mat33 = List[float]
Transform = Tuple[Vec3, Quat]

# Particle flags
PARTICLE_FLAG_ACTIVE = wp.constant(wp.uint32(1 << 0))

# Shape geometry types
GEO_SPHERE = wp.constant(0)
GEO_BOX = wp.constant(1)
GEO_CAPSULE = wp.constant(2)
GEO_CYLINDER = wp.constant(3)
GEO_CONE = wp.constant(4)
GEO_MESH = wp.constant(5)
GEO_SDF = wp.constant(6)
GEO_PLANE = wp.constant(7)
GEO_NONE = wp.constant(8)

# Types of joints linking rigid bodies
JOINT_PRISMATIC = wp.constant(0)
JOINT_REVOLUTE = wp.constant(1)
JOINT_BALL = wp.constant(2)
JOINT_FIXED = wp.constant(3)
JOINT_FREE = wp.constant(4)
JOINT_COMPOUND = wp.constant(5)
JOINT_UNIVERSAL = wp.constant(6)
JOINT_DISTANCE = wp.constant(7)
JOINT_D6 = wp.constant(8)

# Joint axis mode types
JOINT_MODE_LIMIT = wp.constant(0)
JOINT_MODE_TARGET_POSITION = wp.constant(1)
JOINT_MODE_TARGET_VELOCITY = wp.constant(2)


def flag_to_int(flag):
    """Converts a flag to an integer."""
    if type(flag) in wp.types.int_types:
        return flag.value
    return int(flag)


# Material properties pertaining to rigid shape contact dynamics
@wp.struct
class ModelShapeMaterials:
    ke: wp.array(dtype=float)  # The contact elastic stiffness (only used by Euler integrator)
    kd: wp.array(dtype=float)  # The contact damping stiffness (only used by Euler integrator)
    kf: wp.array(dtype=float)  # The contact friction stiffness (only used by Euler integrator)
    mu: wp.array(dtype=float)  # The coefficient of friction
    restitution: wp.array(dtype=float)  # The coefficient of restitution (only used by XPBD integrator)


# Shape properties of geometry
@wp.struct
class ModelShapeGeometry:
    type: wp.array(dtype=wp.int32)  # The type of geometry (GEO_SPHERE, GEO_BOX, etc.)
    is_solid: wp.array(dtype=wp.uint8)  # Indicates whether the shape is solid or hollow
    thickness: wp.array(
        dtype=float
    )  # The thickness of the shape (used for collision detection, and inertia computation of hollow shapes)
    source: wp.array(dtype=wp.uint64)  # Pointer to the source geometry (in case of a mesh, zero otherwise)
    scale: wp.array(dtype=wp.vec3)  # The 3D scale of the shape
    face_filter: wp.array3d(dtype=wp.int32)   # Filters applied to geo mesh shapes (e.g. to exclude parts of mesh from collisions)
    particle_filter_ids: wp.array2d(dtype=wp.int32)   # A selector of facefilters per shape per particle of a model


# 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=-math.inf,
        limit_upper=math.inf,
        limit_ke=100.0,
        limit_kd=10.0,
        target=None,
        target_ke=0.0,
        target_kd=0.0,
        mode=JOINT_MODE_TARGET_POSITION,
    ):
        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.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 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
    """

    def __init__(self, vertices: List[Vec3], indices: List[int], compute_inertia=True, is_solid=True):
        """Construct a Mesh object from a triangle mesh

        The mesh center of mass and inertia tensor will automatically be
        calculated using a density of 1.0. This computation is only valid
        if the mesh is closed (two-manifold).

        Args:
            vertices: List of vertices in the mesh
            indices: List of triangle indices, 3 per-element
            compute_inertia: If True, the mass, inertia tensor and center of mass will be computed assuming density of 1.0
            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 = 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 = wp.mat33(np.eye(3))
            self.mass = 1.0
            self.com = wp.vec3()

    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)
            indices = wp.array(self.indices, dtype=wp.int32)

            self.mesh = wp.Mesh(points=pos, velocities=vel, indices=indices)
            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))


class State:
    """The State object holds all *time-varying* data for a model.

    Time-varying data includes particle positions, velocities, rigid body states, and
    anything that is output from the integrator as derived data, e.g.: forces.

    The exact attributes depend on the contents of the model. State objects should
    generally be created using the :func:`Model.state()` function.

    Attributes:

        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 (array): Tensor of body coordinates
        body_qd (array): Tensor of body velocities
        body_f (array): Tensor of body forces

    """

    def __init__(self):
        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_()

        if self.body_count:
            self.body_f.zero_()

    def flatten(self):
        wp.utils.warn(
            "Model.flatten() will be removed in a future Warp version.",
            DeprecationWarning,
            stacklevel=2,
        )

        arrays = []

        # build a list of all tensor attributes
        for attr, value in self.__dict__.items():
            if isinstance(value, wp.array):
                arrays.append(value)

        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, wp.vec3(), wp.mat33()

    if type == GEO_SPHERE:
        solid = compute_sphere_inertia(density, scale[0])
        if is_solid:
            return solid
        else:
            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 = scale * 2.0
        solid = compute_box_inertia(density, w, h, d)
        if is_solid:
            return solid
        else:
            hollow = compute_box_inertia(density, w - thickness, h - thickness, d - thickness)
            return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
    elif type == GEO_CAPSULE:
        r, h = scale[0], scale[1] * 2.0
        solid = compute_capsule_inertia(density, r, h)
        if is_solid:
            return solid
        else:
            hollow = compute_capsule_inertia(density, r - thickness, h - 2.0 * thickness)
            return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
    elif type == GEO_CYLINDER:
        r, h = scale[0], scale[1] * 2.0
        solid = compute_cylinder_inertia(density, r, h)
        if is_solid:
            return solid
        else:
            hollow = compute_cylinder_inertia(density, r - thickness, h - 2.0 * thickness)
            return solid[0] - hollow[0], solid[1], solid[2] - hollow[2]
    elif type == GEO_CONE:
        r, h = scale[0], scale[1] * 2.0
        solid = compute_cone_inertia(density, r, h)
        if is_solid:
            return solid
        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 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

            sx, sy, sz = scale

            mass_ratio = sx * sy * sz * density
            m_new = m * mass_ratio

            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
            Izz = I[2, 2] * (sx**2 + sy**2) / 2 * mass_ratio
            Ixy = I[0, 1] * sx * sy * mass_ratio
            Ixz = I[0, 2] * sx * sz * mass_ratio
            Iyz = I[1, 2] * sy * sz * mass_ratio

            I_new = wp.mat33([[Ixx, Ixy, Ixz], [Ixy, Iyy, Iyz], [Ixz, Iyz, Izz]])

            return m_new, c_new, I_new
        elif type == GEO_MESH:
            # fall back to computing inertia from mesh geometry
            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))


class Model:
    """Holds the definition of the simulation model

    This class holds the non-time varying description of the system, i.e.:
    all geometry, constraints, and parameters used to describe the simulation.

    Attributes:
        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 (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 (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 (array): Particle enabled state, shape [particle_count], bool
        particle_max_velocity (float): Maximum particle velocity (to prevent instability)

        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
        shape_geo_src (list): List of `wp.Mesh` instances used for rendering of mesh geometry

        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 (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 (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 (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)

        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)
        soft_contact_kf (float): Stiffness of friction force in soft contacts (used by SemiImplicitIntegrator)
        soft_contact_mu (float): Friction coefficient of soft contacts
        soft_contact_restitution (float): Restitution coefficient of soft contacts (used by XPBDIntegrator)

        rigid_contact_margin (float): Contact margin for generation of rigid body contacts
        rigid_contact_torsional_friction (float): Torsional friction coefficient for rigid body contacts (used by XPBDIntegrator)
        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 (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

        particle_count (int): Total number of particles in the system
        body_count (int): Total number of bodies in the system
        shape_count (int): Total number of shapes in the system
        joint_count (int): Total number of joints in the system
        tri_count (int): Total number of triangles in the system
        tet_count (int): Total number of tetrahedra in the system
        edge_count (int): Total number of edges in the system
        spring_count (int): Total number of springs in the system
        contact_count (int): Total number of contacts in the system
        muscle_count (int): Total number of muscles in the system
        articulation_count (int): Total number of articulations in the system
        joint_dof_count (int): Total number of velocity degrees of freedom of all joints in the system
        joint_coord_count (int): Total number of position degrees of freedom of all joints in the system

        device (wp.Device): Device on which the Model was allocated

    Note:
        It is strongly recommended to use the ModelBuilder to construct a simulation rather
        than creating your own Model object directly, however it is possible to do so if
        desired.
    """

    def __init__(self, device=None):
        self.requires_grad = False
        self.num_envs = 0

        self.particle_q = None
        self.particle_qd = None
        self.particle_mass = None
        self.particle_inv_mass = None
        self._particle_radius = None
        self.particle_max_radius = 0.0
        self.particle_ke = 1.0e3
        self.particle_kd = 1.0e2
        self.particle_kf = 1.0e2
        self.particle_mu = 0.5
        self.particle_cohesion = 0.0
        self.particle_adhesion = 0.0
        self.particle_grid = None
        self.particle_flags = None
        self.particle_max_velocity = 1e5 
        self.global_viscous_damping = np.array((0.0, 0.0, 0.0))  # (damping, apply_min, velocity_max)  # TODO Should be removed when cleaning the Euler integrator

        self.attachment_constraint = False
        self.attachment_indices = None
        self.attachment_point = None
        self.attachment_norm = None
        self.attachment_stiffness = None
        self.attachment_damping = None

        self.shape_transform = None
        self.shape_body = None
        self.body_shapes = {}
        self.shape_materials = ModelShapeMaterials()
        self.shape_geo = ModelShapeGeometry()
        self.shape_geo_src = None
        self.shape_geo_src_id = None

        self.shape_collision_group = None
        self.shape_collision_group_map = None
        self.shape_collision_filter_pairs = None
        self.shape_collision_radius = None
        self.shape_ground_collision = None
        self.shape_contact_pairs = None
        self.shape_ground_contact_pairs = None

        self.spring_indices = None
        self.spring_rest_length = None
        self.spring_stiffness = None
        self.spring_damping = None
        self.spring_control = None

        self.tri_indices = None
        self.tri_poses = None
        self.tri_activations = None
        self.tri_materials = None

        self.edge_indices = None
        self.edge_rest_angle = None
        self.edge_bending_properties = None

        self.tet_indices = None
        self.tet_poses = None
        self.tet_activations = None
        self.tet_materials = None

        self.body_q = None
        self.body_qd = None
        self.body_com = None
        self.body_inertia = None
        self.body_inv_inertia = None
        self.body_mass = None
        self.body_inv_mass = None
        self.body_name = None

        self.joint_q = None
        self.joint_qd = None
        self.joint_act = None
        self.joint_type = None
        self.joint_parent = None
        self.joint_child = None
        self.joint_X_p = None
        self.joint_X_c = None
        self.joint_axis = None
        self.joint_armature = None
        self.joint_target = None
        self.joint_target_ke = None
        self.joint_target_kd = None
        self.joint_axis_start = None
        self.joint_axis_dim = None
        self.joint_axis_mode = None
        self.joint_linear_compliance = None
        self.joint_angular_compliance = None
        self.joint_enabled = None
        self.joint_limit_lower = None
        self.joint_limit_upper = None
        self.joint_limit_ke = None
        self.joint_limit_kd = None
        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?
        self.joint_attach_ke = 1.0e3
        self.joint_attach_kd = 1.0e2

        self.soft_contact_margin = 0.2   
        self.soft_contact_ke = 1.0e3
        self.soft_contact_kd = 10.0
        self.soft_contact_kf = 1.0e3
        self.soft_contact_mu = 0.5
        self.soft_contact_restitution = 0.0

        self.rigid_contact_margin = None
        self.rigid_contact_torsional_friction = None
        self.rigid_contact_rolling_friction = None

        # toggles ground contact for all shapes
        self.ground = True
        self.ground_plane = None
        self.up_vector = np.array((0.0, 1.0, 0.0))
        self.up_axis = 1
        self.gravity = np.array((0.0, -9.81, 0.0))

        # for body dragging
        self.cloth_reference_drag = False
        self.cloth_reference_shape_ids = None
        self.cloth_reference_transforms = None
        self.cloth_reference_scale = None
        self.particle_reference_label = None
        self.drag_label_pairs = None
        self.drag_label_pairs_count = 0
        self.cloth_reference_margin = 0.1
        self.cloth_reference_k = 1.0e3
        self.cloth_reference_drag_particles = None
        self.cloth_reference_watertight_whole_shape_index = None

        # for cloth self-collision
        self.particle_shape = None
        self.cloth_self_collision_margin = 0.1
        self.cloth_self_collision_count = None
        self.enable_particle_particle_collisions = False
        self.enable_triangle_particle_collisions = False
        self.enable_edge_edge_collisions = False

        self.particle_normal = None

        self.particle_count = 0
        self.body_count = 0
        self.shape_count = 0
        self.joint_count = 0
        self.joint_axis_count = 0
        self.tri_count = 0
        self.tet_count = 0
        self.edge_count = 0
        self.spring_count = 0
        self.muscle_count = 0
        self.articulation_count = 0
        self.joint_dof_count = 0
        self.joint_coord_count = 0
        self.attachment_count = 0

        self.device = wp.get_device(device)

    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
        the model description.
        """

        s = State()
        if requires_grad is None:
            requires_grad = self.requires_grad

        s.particle_count = self.particle_count
        s.body_count = self.body_count

        # --------------------------------
        # 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)
            s.particle_qd = wp.clone(self.particle_qd)
            s.particle_f = wp.zeros_like(self.particle_qd)

            s.particle_q.requires_grad = requires_grad
            s.particle_qd.requires_grad = requires_grad
            s.particle_f.requires_grad = requires_grad

        # articulations
        if self.body_count:
            s.body_q = wp.clone(self.body_q)
            s.body_qd = wp.clone(self.body_qd)
            s.body_f = wp.zeros_like(self.body_qd)

            s.body_deltas = wp.zeros(
                self.body_count, dtype=wp.spatial_vector, device=s.body_q.device, requires_grad=requires_grad
            )

            s.body_q.requires_grad = requires_grad
            s.body_qd.requires_grad = requires_grad
            s.body_f.requires_grad = requires_grad

        return s

    def allocate_soft_contacts(self, count, requires_grad=False):
        self.soft_contact_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        self.body_cloth_intersection_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        self.soft_contact_particle = wp.zeros(count, dtype=int, device=self.device)
        self.soft_contact_shape = wp.zeros(count, dtype=int, device=self.device)
        self.soft_contact_body_pos = wp.zeros(count, dtype=wp.vec3, device=self.device, requires_grad=requires_grad)
        self.soft_contact_body_vel = wp.zeros(count, dtype=wp.vec3, device=self.device, requires_grad=requires_grad)
        self.soft_contact_normal = wp.zeros(count, dtype=wp.vec3, device=self.device, requires_grad=requires_grad)

    def allocate_particle_tri_self_intersection(self, count, requires_grad=False):
        
        # TODO Proper requires_grad options for the arrays
        self.point_tri_contact_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        self.point_tri_contact_pairs = wp.zeros(count, dtype=wp.vec2i, device=self.device,
                                           requires_grad=requires_grad)
        self.point_tri_contact_filter = wp.zeros(count, dtype=wp.bool, device=self.device)
        self.point_tri_contact_sidedness = wp.zeros(count, dtype=wp.bool, device=self.device,
                                            requires_grad=requires_grad)

    def allocate_edge_self_intersection(self, count, requires_grad=False):

        # TODO Proper requires_grad options for the arrays
        self.edge_contact_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        self.edge_contact_pairs = wp.zeros(count, dtype=wp.vec4i, device=self.device,
                                           requires_grad=requires_grad)
        self.edge_contact_filter = wp.zeros(count, dtype=wp.bool, device=self.device)
        self.edge_contact_normal = wp.zeros(count, dtype=wp.vec3, device=self.device,
                                            requires_grad=requires_grad)

    def allocate_global_self_intersection(self, count, requires_grad=False):

        # TODO Proper requires_grad options for the arrays
        self.particle_self_intersection_count = wp.zeros(1, dtype=wp.int32, device=self.device)  # For simple count
        self.particle_global_self_intersection_count = wp.zeros(1, dtype=wp.int32, device=self.device)  # For global collision resolution
        self.particle_self_intersection_particle = wp.zeros(self.particle_count, dtype=int, device=self.device)
        

    def find_shape_contact_pairs(self):
        # find potential contact pairs based on collision groups and collision mask (pairwise filtering)
        import copy
        import itertools

        filters = copy.copy(self.shape_collision_filter_pairs)
        for a, b in self.shape_collision_filter_pairs:
            filters.add((b, a))
        contact_pairs = []
        # iterate over collision groups (islands)
        for group, shapes in self.shape_collision_group_map.items():
            for shape_a, shape_b in itertools.product(shapes, shapes):
                if shape_a < shape_b and (shape_a, shape_b) not in filters:
                    contact_pairs.append((shape_a, shape_b))
            if group != -1 and -1 in self.shape_collision_group_map:
                # shapes with collision group -1 collide with all other shapes
                for shape_a, shape_b in itertools.product(shapes, self.shape_collision_group_map[-1]):
                    if shape_a < shape_b and (shape_a, shape_b) not in filters:
                        contact_pairs.append((shape_a, shape_b))
        self.shape_contact_pairs = wp.array(np.array(contact_pairs), dtype=wp.int32, device=self.device)
        self.shape_contact_pair_count = len(contact_pairs)
        # find ground contact pairs
        ground_contact_pairs = []
        ground_id = self.shape_count - 1
        for i in range(ground_id):
            if self.shape_ground_collision[i]:
                ground_contact_pairs.append((i, ground_id))
        self.shape_ground_contact_pairs = wp.array(np.array(ground_contact_pairs), dtype=wp.int32, device=self.device)
        self.shape_ground_contact_pair_count = len(ground_contact_pairs)

    def count_contact_points(self):
        """
        Counts the maximum number of contact points that need to be allocated.
        """
        from .collide import count_contact_points

        # calculate the potential number of shape pair contact points
        contact_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        wp.launch(
            kernel=count_contact_points,
            dim=self.shape_contact_pair_count,
            inputs=[
                self.shape_contact_pairs,
                self.shape_geo,
            ],
            outputs=[contact_count],
            device=self.device,
            record_tape=False,
        )
        # count ground contacts
        wp.launch(
            kernel=count_contact_points,
            dim=self.shape_ground_contact_pair_count,
            inputs=[
                self.shape_ground_contact_pairs,
                self.shape_geo,
            ],
            outputs=[contact_count],
            device=self.device,
            record_tape=False,
        )
        count = contact_count.numpy()[0]
        return int(count)

    def allocate_rigid_contacts(self, count=None, requires_grad=False):
        if count is not None:
            self.rigid_contact_max = count
        # serves as counter of the number of active contact points
        self.rigid_contact_count = wp.zeros(1, dtype=wp.int32, device=self.device)
        # contact point ID within the (shape_a, shape_b) contact pair
        self.rigid_contact_point_id = wp.zeros(self.rigid_contact_max, dtype=wp.int32, device=self.device)
        # ID of first rigid body
        self.rigid_contact_body0 = wp.zeros(self.rigid_contact_max, dtype=wp.int32, device=self.device)
        # ID of second rigid body
        self.rigid_contact_body1 = wp.zeros(self.rigid_contact_max, dtype=wp.int32, device=self.device)
        # position of contact point in body 0's frame before the integration step
        self.rigid_contact_point0 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # position of contact point in body 1's frame before the integration step
        self.rigid_contact_point1 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # moment arm before the integration step resulting from thickness displacement added to contact point 0 in body 0's frame (used in XPBD contact friction handling)
        self.rigid_contact_offset0 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # moment arm before the integration step resulting from thickness displacement added to contact point 1 in body 1's frame (used in XPBD contact friction handling)
        self.rigid_contact_offset1 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # contact normal in world frame
        self.rigid_contact_normal = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # combined thickness of both shapes
        self.rigid_contact_thickness = wp.zeros(
            self.rigid_contact_max, dtype=wp.float32, device=self.device, requires_grad=requires_grad
        )
        # ID of the first shape in the contact pair
        self.rigid_contact_shape0 = wp.zeros(self.rigid_contact_max, dtype=wp.int32, device=self.device)
        # ID of the second shape in the contact pair
        self.rigid_contact_shape1 = wp.zeros(self.rigid_contact_max, dtype=wp.int32, device=self.device)

        # temporary variables used during the XPBD solver iterations:
        # world space position of contact point resulting from applying current body 0 transform to its point0
        self.rigid_active_contact_point0 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # world space position of contact point resulting from applying current body 1 transform to its point1
        self.rigid_active_contact_point1 = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # current contact distance (negative penetration depth)
        self.rigid_active_contact_distance = wp.zeros(
            self.rigid_contact_max, dtype=wp.float32, device=self.device, requires_grad=requires_grad
        )
        # contact distance before the solver iterations
        self.rigid_active_contact_distance_prev = wp.zeros(
            self.rigid_contact_max, dtype=wp.float32, device=self.device, requires_grad=requires_grad
        )
        # world space position of point0 before the solver iterations
        self.rigid_active_contact_point0_prev = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # world space position of point1 before the solver iterations
        self.rigid_active_contact_point1_prev = wp.zeros(
            self.rigid_contact_max, dtype=wp.vec3, device=self.device, requires_grad=requires_grad
        )
        # number of contact constraints per rigid body (used for scaling the constraint contributions, a basic version of mass splitting)
        self.rigid_contact_inv_weight = wp.zeros(
            len(self.body_q), dtype=wp.float32, device=self.device, requires_grad=requires_grad
        )
        # number of contact constraints before the solver iterations
        self.rigid_contact_inv_weight_prev = wp.zeros(
            len(self.body_q), dtype=wp.float32, device=self.device, requires_grad=requires_grad
        )

    def flatten(self):
        """Returns a list of Tensors stored by the model

        This function is intended to be used internal-only but can be used to obtain
        a set of all tensors owned by the model.
        """

        tensors = []

        # build a list of all tensor attributes
        for attr, value in self.__dict__.items():
            if wp.is_tensor(value):
                tensors.append(value)

        return tensors

    def collide(self, state: State):
        wp.utils.warn(
            "Model.collide() is not needed anymore and will be removed in a future Warp version.",
            DeprecationWarning,
            stacklevel=2,
        )

    @property
    def soft_contact_max(self):
        """Maximum number of soft contacts that can be registered"""
        return len(self.soft_contact_particle)

    @property
    def edge_contact_max(self):
        """Maximum number of edge-edge contacts that can be registered"""
        return len(self.edge_contact_pairs)

    @property
    def point_tri_contact_max(self):
        """Maximum number of edge-edge contacts that can be registered"""
        return len(self.point_tri_contact_pairs)

    @property
    def soft_contact_distance(self):
        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,
            stacklevel=2,
        )
        return self.particle_max_radius

    @soft_contact_distance.setter
    def soft_contact_distance(self, value):
        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,
            stacklevel=2,
        )

    @property
    def particle_radius(self):
        # Array of per-particle radii
        return self._particle_radius

    @particle_radius.setter
    def particle_radius(self, value):
        if isinstance(value, float):
            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.",
                PendingDeprecationWarning,
                stacklevel=2,
            )
            self._particle_radius.fill_(value)
            self.particle_max_radius = value
        else:
            self._particle_radius = value
            # TODO implement max radius update to be compatible with graph capture
            device = wp.get_device(self.device)
            if not device.is_capturing:
                self.particle_max_radius = self._particle_radius.numpy().max()


class ModelBuilder:
    """A helper class for building simulation models at runtime.

    Use the ModelBuilder to construct a simulation scene. The ModelBuilder
    and builds the scene representation using standard Python data structures (lists),
    this means it is not differentiable. Once :func:`finalize()`
    has been called the ModelBuilder transfers all data to Warp tensors and returns
    an object that may be used for simulation.

    Example
    -------

    .. code-block:: python

        import warp as wp
        import warp.sim

        builder = wp.sim.ModelBuilder()

        # anchor point (zero mass)
        builder.add_particle((0, 1.0, 0.0), (0.0, 0.0, 0.0), 0.0)

        # build chain
        for i in range(1,10):
            builder.add_particle((i, 1.0, 0.0), (0.0, 0.0, 0.0), 1.0)
            builder.add_spring(i-1, i, 1.e+3, 0.0, 0)

        # create model
        model = builder.finalize("cuda")

        state = model.state()
        integrator = wp.sim.SemiImplicitIntegrator()

        for i in range(100):

           state.clear_forces()
           integrator.simulate(model, state, state, dt=1.0/60.0)

    Note:
        It is strongly recommended to use the ModelBuilder to construct a simulation rather
        than creating your own Model object directly, however it is possible to do so if
        desired.
    """

    # Default particle settings
    default_particle_radius = 0.1

    # Default triangle soft mesh settings
    default_tri_ke = 100.0
    default_tri_ka = 100.0
    default_tri_kd = 10.0
    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

    # Default rigid shape contact material properties
    default_shape_ke = 1.0e5
    default_shape_kd = 1000.0
    default_shape_kf = 1000.0
    default_shape_mu = 0.5
    default_shape_restitution = 0.0
    default_shape_density = 1000.0

    # Default joint settings
    default_joint_limit_ke = 100.0
    default_joint_limit_kd = 1.0

    # Default geo settings
    default_geo_thickness = 1e-5

    def __init__(self, up_vector=(0.0, 1.0, 0.0), gravity=-9.80665):
        self.num_envs = 0

        # particles
        self.particle_q = []
        self.particle_qd = []
        self.particle_mass = []
        self.particle_radius = []
        self.particle_flags = []
        self.particle_inds = []
        self.particle_max_velocity = 1e5

        # shapes (each shape has an entry in these arrays)
        # transform from shape to body
        self.shape_transform = []
        # maps from shape index to body index
        self.shape_body = []
        self.shape_geo_type = []
        self.shape_geo_scale = []
        self.shape_geo_src = []
        self.shape_geo_is_solid = []
        self.shape_geo_thickness = []
        self.shape_geo_face_filters = []
        self.shape_geo_particle_filters_id = []
        self.shape_material_ke = []
        self.shape_material_kd = []
        self.shape_material_kf = []
        self.shape_material_mu = []
        self.shape_material_restitution = []
        # collision groups within collisions are handled
        self.shape_collision_group = []
        self.shape_collision_group_map = {}
        self.last_collision_group = 0
        # radius to use for broadphase collision checking
        self.shape_collision_radius = []
        # whether the shape collides with the ground
        self.shape_ground_collision = []

        # filtering to ignore certain collision pairs
        self.shape_collision_filter_pairs = set()

        # geometry
        self.geo_meshes = []
        self.geo_sdfs = []

        # Attachment constraints
        self.attachment_indices = []
        self.attachment_point = []
        self.attachment_norm = []
        self.attachment_stiffness = []
        self.attachment_damping = []

        # springs
        self.spring_indices = []
        self.spring_rest_length = []
        self.spring_stiffness = []
        self.spring_damping = []
        self.spring_control = []

        # triangles
        self.tri_indices = []
        self.tri_poses = []
        self.tri_activations = []
        self.tri_materials = []

        # edges (bending)
        self.edge_indices = []
        self.edge_rest_angle = []
        self.edge_bending_properties = []

        # tetrahedra
        self.tet_indices = []
        self.tet_poses = []
        self.tet_activations = []
        self.tet_materials = []

        # muscles
        self.muscle_start = []
        self.muscle_params = []
        self.muscle_activation = []
        self.muscle_bodies = []
        self.muscle_points = []

        # rigid bodies
        self.body_mass = []
        self.body_inertia = []
        self.body_inv_mass = []
        self.body_inv_inertia = []
        self.body_com = []
        self.body_q = []
        self.body_qd = []
        self.body_name = []
        self.body_shapes = {}  # mapping from body to shapes

        # rigid joints
        self.joint = {}
        self.joint_parent = []  # index of the parent body                      (constant)
        self.joint_parents = {}  # mapping from joint to parent bodies
        self.joint_child = []  # index of the child body                       (constant)
        self.joint_axis = []  # joint axis in child joint frame               (constant)
        self.joint_X_p = []  # frame of joint in parent                      (constant)
        self.joint_X_c = []  # frame of child com (in child coordinates)     (constant)
        self.joint_q = []
        self.joint_qd = []

        self.joint_type = []
        self.joint_name = []
        self.joint_armature = []
        self.joint_target = []
        self.joint_target_ke = []
        self.joint_target_kd = []
        self.joint_axis_mode = []
        self.joint_limit_lower = []
        self.joint_limit_upper = []
        self.joint_limit_ke = []
        self.joint_limit_kd = []
        self.joint_act = []

        self.joint_twist_lower = []
        self.joint_twist_upper = []

        self.joint_linear_compliance = []
        self.joint_angular_compliance = []
        self.joint_enabled = []

        self.joint_q_start = []
        self.joint_qd_start = []
        self.joint_axis_start = []
        self.joint_axis_dim = []
        self.articulation_start = []

        self.joint_dof_count = 0
        self.joint_coord_count = 0
        self.joint_axis_total_count = 0

        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
        # constructor parameters for ground plane shape
        self._ground_params = dict(
            plane=(*up_vector, 0.0),
            width=0.0,
            length=0.0,
            ke=self.default_shape_ke,
            kd=self.default_shape_kd,
            kf=self.default_shape_kf,
            mu=self.default_shape_mu,
            restitution=self.default_shape_restitution
        )

        # Maximum number of soft contacts that can be registered
        self.soft_contact_max = 64 * 1024
        self.edge_contact_max = 3 * 64  # per spring
        self.point_tri_contact_max = 64  # per particle

        # contacts to be generated within the given distance margin to be generated at
        # every simulation substep (can be 0 if only one PBD solver iteration is used)
        self.rigid_contact_margin = 0.1    
        # torsional friction coefficient (only considered by XPBD so far)
        self.rigid_contact_torsional_friction = 0.5
        # rolling friction coefficient (only considered by XPBD so far)
        self.rigid_contact_rolling_friction = 0.001

        # number of rigid contact points to allocate in the model during self.finalize() per environment
        # if setting is None, the number of worst-case number of contacts will be calculated in self.finalize()
        self.num_rigid_contacts_per_env = None

        # for body dragging
        self.cloth_reference_names = []
        self.cloth_reference_shapes = []
        self.cloth_reference_transforms = []
        self.cloth_reference_scale = []
        self.particle_reference_label = []
        self.drag_label_pairs = []

    @property
    def shape_count(self):
        return len(self.shape_geo_type)

    @property
    def body_count(self):
        return len(self.body_q)

    @property
    def joint_count(self):
        return len(self.joint_type)

    @property
    def joint_axis_count(self):
        return len(self.joint_axis)

    @property
    def particle_count(self):
        return len(self.particle_q)

    @property
    def tri_count(self):
        return len(self.tri_poses)

    @property
    def tet_count(self):
        return len(self.tet_poses)

    @property
    def edge_count(self):
        return len(self.edge_rest_angle)

    @property
    def spring_count(self):
        return len(self.spring_rest_length)

    @property
    def muscle_count(self):
        return len(self.muscle_start)

    @property
    def articulation_count(self):
        return len(self.articulation_start)

    # an articulation is a set of contiguous bodies bodies from articulation_start[i] to articulation_start[i+1]
    # these are used for computing forward kinematics e.g.:
    #
    # model.eval_articulation_fk()
    # model.eval_articulation_j()
    # model.eval_articulation_m()
    #
    # articulations are automatically 'closed' when calling finalize

    def add_articulation(self):
        self.articulation_start.append(self.joint_count)

    def add_builder(self, articulation, xform=None, update_num_env_count=True, separate_collision_group=True):
        """Copies a rigid articulation from `articulation`, another `ModelBuilder`.

        Args:
            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
        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]

        if articulation.joint_count:
            joint_X_p = copy.deepcopy(articulation.joint_X_p)
            joint_q = copy.deepcopy(articulation.joint_q)
            if xform is not None:
                for i in range(len(joint_X_p)):
                    if articulation.joint_type[i] == wp.sim.JOINT_FREE:
                        qi = articulation.joint_q_start[i]
                        xform_prev = wp.transform(joint_q[qi : qi + 3], joint_q[qi + 3 : qi + 7])
                        tf = xform * xform_prev
                        joint_q[qi : qi + 3] = tf.p
                        joint_q[qi + 3 : qi + 7] = tf.q
                    elif articulation.joint_parent[i] == -1:
                        joint_X_p[i] = xform * joint_X_p[i]
            self.joint_X_p.extend(joint_X_p)
            self.joint_q.extend(joint_q)

            self.add_articulation()

            # offset the indices
            self.joint_parent.extend([p + self.joint_count if p != -1 else -1 for p in articulation.joint_parent])
            self.joint_child.extend([c + self.joint_count for c in articulation.joint_child])

            self.joint_q_start.extend([c + self.joint_coord_count for c in articulation.joint_q_start])
            self.joint_qd_start.extend([c + self.joint_dof_count for c in articulation.joint_qd_start])

            self.joint_axis_start.extend([a + self.joint_axis_total_count for a in articulation.joint_axis_start])

        joint_children = set(articulation.joint_child)
        for i in range(articulation.body_count):
            if xform is not None and i not in joint_children:
                # rigid body is not attached to a joint, so apply input transform directly
                self.body_q.append(xform * articulation.body_q[i])
            else:
                self.body_q.append(articulation.body_q[i])

        # apply collision group
        if separate_collision_group:
            self.shape_collision_group.extend(
                [self.last_collision_group + 1 for _ in articulation.shape_collision_group]
            )
        else:
            self.shape_collision_group.extend(
                [(g + self.last_collision_group if g > -1 else -1) for g in articulation.shape_collision_group]
            )
        shape_count = self.shape_count
        for i, j in articulation.shape_collision_filter_pairs:
            self.shape_collision_filter_pairs.add((i + shape_count, j + shape_count))
        for group, shapes in articulation.shape_collision_group_map.items():
            if separate_collision_group:
                group = self.last_collision_group + 1
            else:
                group = group + self.last_collision_group if group > -1 else -1
            if group not in self.shape_collision_group_map:
                self.shape_collision_group_map[group] = []
            self.shape_collision_group_map[group].extend([s + shape_count for s in shapes])

        # update last collision group counter
        if separate_collision_group:
            self.last_collision_group += 1
        elif articulation.last_collision_group > -1:
            self.last_collision_group += articulation.last_collision_group

        rigid_articulation_attrs = [
            "body_inertia",
            "body_mass",
            "body_inv_inertia",
            "body_inv_mass",
            "body_com",
            "body_qd",
            "body_name",
            "joint_type",
            "joint_enabled",
            "joint_X_c",
            "joint_armature",
            "joint_axis",
            "joint_axis_dim",
            "joint_axis_mode",
            "joint_name",
            "joint_qd",
            "joint_act",
            "joint_limit_lower",
            "joint_limit_upper",
            "joint_limit_ke",
            "joint_limit_kd",
            "joint_target",
            "joint_target_ke",
            "joint_target_kd",
            "joint_linear_compliance",
            "joint_angular_compliance",
            "shape_geo_type",
            "shape_geo_scale",
            "shape_geo_src",
            "shape_geo_is_solid",
            "shape_geo_thickness",
            "shape_material_ke",
            "shape_material_kd",
            "shape_material_kf",
            "shape_material_mu",
            "shape_material_restitution",
            "shape_collision_radius",
            "shape_ground_collision",
        ]

        for attr in rigid_articulation_attrs:
            getattr(self, attr).extend(getattr(articulation, attr))

        self.joint_dof_count += articulation.joint_dof_count
        self.joint_coord_count += articulation.joint_coord_count
        self.joint_axis_total_count += articulation.joint_axis_total_count

        self.up_vector = articulation.up_vector
        self.gravity = articulation.gravity

        if update_num_env_count:
            self.num_envs += 1

    # register a rigid body and return its index.
    def add_body(
        self,
        origin: Transform = wp.transform(),
        armature: float = 0.0,
        com: Vec3 = wp.vec3(),
        I_m: Mat33 = wp.mat33(),
        m: float = 0.0,
        name: str = None,
    ) -> int:
        """Adds a rigid body to the model.

        Args:
            origin: The location of the body in the world frame
            armature: Artificial inertia added to the body
            com: The center of mass of the body w.r.t its origin
            I_m: The 3x3 inertia tensor of the body (specified relative to the center of mass)
            m: Mass of the body
            name: Name of the body (optional)

        Returns:
            The index of the body in the model

        Note:
            If the mass (m) is zero then the body is treated as kinematic with no dynamics

        """

        body_id = len(self.body_mass)

        # body data
        inertia = I_m + wp.mat33(np.eye(3)) * armature
        self.body_inertia.append(inertia)
        self.body_mass.append(m)
        self.body_com.append(com)

        if m > 0.0:
            self.body_inv_mass.append(1.0 / m)
        else:
            self.body_inv_mass.append(0.0)

        if any(x for x in inertia):
            self.body_inv_inertia.append(wp.inverse(inertia))
        else:
            self.body_inv_inertia.append(inertia)

        self.body_q.append(origin)
        self.body_qd.append(wp.spatial_vector())

        self.body_name.append(name or f"body {body_id}")
        self.body_shapes[body_id] = []
        return body_id

    def add_joint(
        self,
        joint_type: wp.constant,
        parent: int,
        child: int,
        linear_axes: List[JointAxis] = [],
        angular_axes: List[JointAxis] = [],
        name: str = None,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        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:
            self.joint_parents[child] = [parent]
        else:
            self.joint_parents[child].append(parent)
        self.joint_child.append(child)
        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)))
        self.joint_axis_total_count += len(linear_axes) + len(angular_axes)

        self.joint_linear_compliance.append(linear_compliance)
        self.joint_angular_compliance.append(angular_compliance)
        self.joint_enabled.append(enabled)

        def add_axis_dim(dim):
            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)
            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:
            add_axis_dim(dim)

        if joint_type == JOINT_PRISMATIC:
            dof_count = 1
            coord_count = 1
        elif joint_type == JOINT_REVOLUTE:
            dof_count = 1
            coord_count = 1
        elif joint_type == JOINT_BALL:
            dof_count = 3
            coord_count = 4
        elif joint_type == JOINT_FREE or joint_type == JOINT_DISTANCE:
            dof_count = 6
            coord_count = 7
        elif joint_type == JOINT_FIXED:
            dof_count = 0
            coord_count = 0
        elif joint_type == JOINT_UNIVERSAL:
            dof_count = 2
            coord_count = 2
        elif joint_type == JOINT_COMPOUND:
            dof_count = 3
            coord_count = 3
        elif joint_type == JOINT_D6:
            dof_count = len(linear_axes) + len(angular_axes)
            coord_count = dof_count

        for i in range(coord_count):
            self.joint_q.append(0.0)

        for i in range(dof_count):
            self.joint_qd.append(0.0)
            self.joint_act.append(0.0)

        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

        self.joint_q_start.append(self.joint_coord_count)
        self.joint_qd_start.append(self.joint_dof_count)

        self.joint_dof_count += dof_count
        self.joint_coord_count += coord_count

        if collision_filter_parent and parent > -1:
            for child_shape in self.body_shapes[child]:
                for parent_shape in self.body_shapes[parent]:
                    self.shape_collision_filter_pairs.add((parent_shape, child_shape))

        return self.joint_count - 1

    def add_joint_revolute(
        self,
        parent: int,
        child: int,
        parent_xform: wp.transform,
        child_xform: wp.transform,
        axis: Vec3,
        target: float = 0.0,
        target_ke: float = 0.0,
        target_kd: float = 0.0,
        mode: int = JOINT_MODE_TARGET_POSITION,
        limit_lower: float = -2 * math.pi,
        limit_upper: float = 2 * math.pi,
        limit_ke: float = default_joint_limit_ke,
        limit_kd: float = default_joint_limit_kd,
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (: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
            limit_lower: The lower limit of the joint
            limit_upper: The upper limit of the joint
            limit_ke: The stiffness of the joint limit
            limit_kd: The damping of the joint limit
            linear_compliance: The linear compliance of the joint
            angular_compliance: The angular compliance of the joint
            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 added joint

        """
        ax = JointAxis(
            axis=axis,
            limit_lower=limit_lower,
            limit_upper=limit_upper,
            target=target,
            target_ke=target_ke,
            target_kd=target_kd,
            mode=mode,
            limit_ke=limit_ke,
            limit_kd=limit_kd,
        )
        return self.add_joint(
            JOINT_REVOLUTE,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            angular_axes=[ax],
            linear_compliance=linear_compliance,
            angular_compliance=angular_compliance,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_prismatic(
        self,
        parent: int,
        child: int,
        parent_xform: wp.transform,
        child_xform: wp.transform,
        axis: Vec3,
        target: float = 0.0,
        target_ke: float = 0.0,
        target_kd: float = 0.0,
        mode: int = JOINT_MODE_TARGET_POSITION,
        limit_lower: float = -1e4,
        limit_upper: float = 1e4,
        limit_ke: float = default_joint_limit_ke,
        limit_kd: float = default_joint_limit_kd,
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (: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
            limit_lower: The lower limit of the joint
            limit_upper: The upper limit of the joint
            limit_ke: The stiffness of the joint limit
            limit_kd: The damping of the joint limit
            linear_compliance: The linear compliance of the joint
            angular_compliance: The angular compliance of the joint
            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 added joint

        """
        ax = JointAxis(
            axis=axis,
            limit_lower=limit_lower,
            limit_upper=limit_upper,
            target=target,
            target_ke=target_ke,
            target_kd=target_kd,
            mode=mode,
            limit_ke=limit_ke,
            limit_kd=limit_kd,
        )
        return self.add_joint(
            JOINT_PRISMATIC,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            linear_axes=[ax],
            linear_compliance=linear_compliance,
            angular_compliance=angular_compliance,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_ball(
        self,
        parent: int,
        child: int,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (: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
            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

        """
        return self.add_joint(
            JOINT_BALL,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            linear_compliance=linear_compliance,
            angular_compliance=angular_compliance,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_fixed(
        self,
        parent: int,
        child: int,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (: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
            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

        """
        return self.add_joint(
            JOINT_FIXED,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            linear_compliance=linear_compliance,
            angular_compliance=angular_compliance,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_free(
        self,
        child: int,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        parent: int = -1,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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:
            child: The index of the child body
            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 added joint

        """
        return self.add_joint(
            JOINT_FREE,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_distance(
        self,
        parent: int,
        child: int,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        min_distance: float = -1.0,
        max_distance: float = 1.0,
        compliance: float = 0.0,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (: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
            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
        
        .. note:: Distance joints are currently only supported in the :class:`XPBDIntegrator` at the moment.

        """
        ax = JointAxis(
            axis=(1.0, 0.0, 0.0),
            limit_lower=min_distance,
            limit_upper=max_distance,
        )
        return self.add_joint(
            JOINT_DISTANCE,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            linear_axes=[ax],
            linear_compliance=compliance,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_universal(
        self,
        parent: int,
        child: int,
        axis_0: JointAxis,
        axis_1: JointAxis,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (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
            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

        """
        return self.add_joint(
            JOINT_UNIVERSAL,
            parent,
            child,
            angular_axes=[JointAxis(axis_0), JointAxis(axis_1)],
            parent_xform=parent_xform,
            child_xform=child_xform,
            linear_compliance=linear_compliance,
            angular_compliance=angular_compliance,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_compound(
        self,
        parent: int,
        child: int,
        axis_0: JointAxis,
        axis_1: JointAxis,
        axis_2: JointAxis,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        name: str = None,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ) -> int:
        """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 (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 added joint

        """
        return self.add_joint(
            JOINT_COMPOUND,
            parent,
            child,
            angular_axes=[JointAxis(axis_0), JointAxis(axis_1), JointAxis(axis_2)],
            parent_xform=parent_xform,
            child_xform=child_xform,
            name=name,
            collision_filter_parent=collision_filter_parent,
            enabled=enabled,
        )

    def add_joint_d6(
        self,
        parent: int,
        child: int,
        linear_axes: List[JointAxis] = [],
        angular_axes: List[JointAxis] = [],
        name: str = None,
        parent_xform: wp.transform = wp.transform(),
        child_xform: wp.transform = wp.transform(),
        linear_compliance: float = 0.0,
        angular_compliance: float = 0.0,
        collision_filter_parent: bool = True,
        enabled: bool = True,
    ):
        """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
            child: The index of the child body
            linear_axes: A list of linear axes
            angular_axes: A list of angular axes
            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

        """
        return self.add_joint(
            JOINT_D6,
            parent,
            child,
            parent_xform=parent_xform,
            child_xform=child_xform,
            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,
            collision_filter_parent=collision_filter_parent,
            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.

        Args:
            bodies: A list of body indices for each waypoint
            positions: A list of positions of each waypoint in the body's local frame
            f0: Force scaling
            lm: Muscle length
            lt: Tendon length
            lmax: Maximally efficient muscle length

        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)

        self.muscle_start.append(len(self.muscle_bodies))
        self.muscle_params.append((f0, lm, lt, lmax, pen))
        self.muscle_activation.append(0.0)

        for i in range(n):
            self.muscle_bodies.append(bodies[i])
            self.muscle_points.append(positions[i])

        # return the index of the muscle
        return len(self.muscle_start) - 1

    # shapes
    def add_shape_plane(
        self,
        plane: Vec4 = (0.0, 1.0, 0.0, 0.0),
        pos: Vec3 = None,
        rot: Quat = None,
        width: float = 10.0,
        length: float = 10.0,
        body: int = -1,
        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,
        thickness: float = 0.0,
        has_ground_collision: bool = False,
    ):
        """
        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 defined through the `plane` argument is used.

        Args:
            plane: The plane equation in form a*x + b*y + c*z + d = 0
            pos: The position of the plane in world coordinates
            rot: The rotation of the plane in world coordinates
            width: The extent along x of the plane (infinite if 0)
            length: The extent along z of the plane (infinite if 0)
            body: The body index to attach the shape to (-1 by default to keep the plane static)
            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
            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
            normal = np.array(plane[:3])
            normal /= np.linalg.norm(normal)
            pos = plane[3] * normal
            if np.allclose(normal, (0.0, 1.0, 0.0)):
                # no rotation necessary
                rot = (0.0, 0.0, 0.0, 1.0)
            else:
                c = np.cross(normal, (0.0, 1.0, 0.0))
                angle = np.arcsin(np.linalg.norm(c))
                axis = c / np.linalg.norm(c)
                rot = wp.quat_from_axis_angle(axis, angle)
        scale = wp.vec3(width, length, 0.0)

        return self._add_shape(
            body,
            pos,
            rot,
            GEO_PLANE,
            scale,
            None,
            0.0,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_sphere(
        self,
        body,
        pos: Vec3 = (0.0, 0.0, 0.0),
        rot: Quat = (0.0, 0.0, 0.0, 1.0),
        radius: float = 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,
        has_ground_collision: bool = True,
    ):
        """Adds a sphere 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
            radius: The radius of the sphere
            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: Whether the sphere is solid or hollow
            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,
            wp.vec3(pos),
            wp.quat(rot),
            GEO_SPHERE,
            wp.vec3(radius, 0.0, 0.0),
            None,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness + radius,
            is_solid,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_box(
        self,
        body: int,
        pos: Vec3 = (0.0, 0.0, 0.0),
        rot: Quat = (0.0, 0.0, 0.0, 1.0),
        hx: float = 0.5,
        hy: float = 0.5,
        hz: float = 0.5,
        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,
        has_ground_collision: bool = True,
    ):
        """Adds a box 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
            hx: The half-extent along the x-axis
            hy: The half-extent along the y-axis
            hz: The half-extent along the z-axis
            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: Whether the box is solid or hollow
            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,
            wp.vec3(pos),
            wp.quat(rot),
            GEO_BOX,
            wp.vec3(hx, hy, hz),
            None,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness,
            is_solid,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_capsule(
        self,
        body: int,
        pos: Vec3 = (0.0, 0.0, 0.0),
        rot: Quat = (0.0, 0.0, 0.0, 1.0),
        radius: float = 1.0,
        half_height: float = 0.5,
        up_axis: int = 1,
        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,
        has_ground_collision: bool = True,
    ):
        """Adds a capsule 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
            radius: The radius of the capsule
            half_height: The half length of the center cylinder along the up axis
            up_axis: The axis along which the capsule is aligned (0=x, 1=y, 2=z)
            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: Whether the capsule is solid or hollow
            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 = wp.quat(rot)
        sqh = math.sqrt(0.5)
        if up_axis == 0:
            q = wp.mul(q, wp.quat(0.0, 0.0, -sqh, sqh))
        elif up_axis == 2:
            q = wp.mul(q, wp.quat(sqh, 0.0, 0.0, sqh))

        return self._add_shape(
            body,
            wp.vec3(pos),
            wp.quat(q),
            GEO_CAPSULE,
            wp.vec3(radius, half_height, 0.0),
            None,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness + radius,
            is_solid,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_cylinder(
        self,
        body: int,
        pos: Vec3 = (0.0, 0.0, 0.0),
        rot: Quat = (0.0, 0.0, 0.0, 1.0),
        radius: float = 1.0,
        half_height: float = 0.5,
        up_axis: int = 1,
        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,
        has_ground_collision: bool = True,
    ):
        """Adds a cylinder 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
            radius: The radius of the cylinder
            half_height: The half length of the cylinder along the up axis
            up_axis: The axis along which the cylinder is aligned (0=x, 1=y, 2=z)
            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: Whether the cylinder is solid or hollow
            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
        sqh = math.sqrt(0.5)
        if up_axis == 0:
            q = wp.mul(rot, wp.quat(0.0, 0.0, -sqh, sqh))
        elif up_axis == 2:
            q = wp.mul(rot, wp.quat(sqh, 0.0, 0.0, sqh))

        return self._add_shape(
            body,
            wp.vec3(pos),
            wp.quat(q),
            GEO_CYLINDER,
            wp.vec3(radius, half_height, 0.0),
            None,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness,
            is_solid,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_cone(
        self,
        body: int,
        pos: Vec3 = (0.0, 0.0, 0.0),
        rot: Quat = (0.0, 0.0, 0.0, 1.0),
        radius: float = 1.0,
        half_height: float = 0.5,
        up_axis: int = 1,
        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,
        has_ground_collision: bool = True,
    ):
        """Adds a cone 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
            radius: The radius of the cone
            half_height: The half length of the cone along the up axis
            up_axis: The axis along which the cone is aligned (0=x, 1=y, 2=z)
            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: Whether the cone is solid or hollow
            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
        sqh = math.sqrt(0.5)
        if up_axis == 0:
            q = wp.mul(rot, wp.quat(0.0, 0.0, -sqh, sqh))
        elif up_axis == 2:
            q = wp.mul(rot, wp.quat(sqh, 0.0, 0.0, sqh))

        return self._add_shape(
            body,
            wp.vec3(pos),
            wp.quat(q),
            GEO_CONE,
            wp.vec3(radius, half_height, 0.0),
            None,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness,
            is_solid,
            has_ground_collision=has_ground_collision,
        )

    def add_shape_mesh(
        self,
        body: int,
        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 = wp.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,
        has_ground_collision: bool = True,
        face_filters: List = [[]],
        model_particle_filter_ids: List = [],
    ):
        """Adds a triangle mesh 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
            mesh: The mesh 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,
            pos,
            rot,
            GEO_MESH,
            wp.vec3(scale[0], scale[1], scale[2]),
            mesh,
            density,
            ke,
            kd,
            kf,
            mu,
            restitution,
            thickness,
            is_solid,
            has_ground_collision=has_ground_collision,
            face_filters=face_filters,
            model_particle_filter_ids=model_particle_filter_ids
        )

    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.
        """
        if type == GEO_SPHERE:
            return scale[0]
        elif type == GEO_BOX:
            return np.linalg.norm(scale)
        elif type == GEO_CAPSULE or type == GEO_CYLINDER or type == GEO_CONE:
            return scale[0] + scale[1]
        elif type == GEO_MESH:
            vmax = np.max(np.abs(src.vertices), axis=0) * np.max(scale)
            return np.linalg.norm(vmax)
        elif type == GEO_PLANE:
            if scale[0] > 0.0 and scale[1] > 0.0:
                # finite plane
                return np.linalg.norm(scale)
            else:
                return 1.0e6
        else:
            return 10.0

    def _add_shape(
        self,
        body,
        pos,
        rot,
        type,
        scale,
        src,
        density,
        ke,
        kd,
        kf,
        mu,
        restitution,
        thickness=default_geo_thickness,
        is_solid=True,
        collision_group=-1,
        collision_filter_parent=True,
        has_ground_collision=True,
        face_filters=[[]],
        model_particle_filter_ids: List = [],
    ):
        self.shape_body.append(body)
        shape = self.shape_count
        if body in self.body_shapes:
            # no contacts between shapes of the same body
            for same_body_shape in self.body_shapes[body]:
                self.shape_collision_filter_pairs.add((same_body_shape, shape))
            self.body_shapes[body].append(shape)
        else:
            self.body_shapes[body] = [shape]
        self.shape_transform.append(wp.transform(pos, rot))
        self.shape_geo_type.append(type)
        self.shape_geo_scale.append((scale[0], scale[1], scale[2]))
        self.shape_geo_src.append(src)
        self.shape_geo_thickness.append(thickness)
        self.shape_geo_is_solid.append(is_solid)
        self.shape_geo_face_filters.append(face_filters)
        self.shape_geo_particle_filters_id.append(model_particle_filter_ids)

        self.shape_material_ke.append(ke)
        self.shape_material_kd.append(kd)
        self.shape_material_kf.append(kf)
        self.shape_material_mu.append(mu)
        self.shape_material_restitution.append(restitution)
        self.shape_collision_group.append(collision_group)
        if collision_group not in self.shape_collision_group_map:
            self.shape_collision_group_map[collision_group] = []
        self.last_collision_group = max(self.last_collision_group, collision_group)
        self.shape_collision_group_map[collision_group].append(shape)
        self.shape_collision_radius.append(self._shape_radius(type, scale, src))
        if collision_filter_parent and body > -1 and body in self.joint_parents:
            for parent_body in self.joint_parents[body]:
                if parent_body > -1:
                    for parent_shape in self.body_shapes[parent_body]:
                        self.shape_collision_filter_pairs.add((parent_shape, shape))
        if body == -1:
            has_ground_collision = False
        self.shape_ground_collision.append(has_ground_collision)

        (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))
        return shape

    # particles
    def add_particle(
        self, pos: Vec3, vel: Vec3, mass: float, radius: float = None, flags: wp.uint32 = PARTICLE_FLAG_ACTIVE
    ) -> int:
        """Adds a single particle to the model

        Args:
            pos: The initial position of the particle
            vel: The initial velocity of the particle
            mass: The mass of the particle
            flags: The flags that control the dynamical behavior of the particle, see PARTICLE_FLAG_* constants
            radius: The radius of the particle used in collision handling. If None, the radius is set to the default value (default_particle_radius).

        Note:
            Set the mass equal to zero to create a 'kinematic' particle that does is not subject to dynamics.

        Returns:
            The index of the particle in the system
        """
        self.particle_q.append(pos)
        self.particle_qd.append(vel)
        self.particle_mass.append(mass)
        if radius is None:
            radius = self.default_particle_radius
        self.particle_radius.append(radius)
        self.particle_flags.append(flags)

        return len(self.particle_q) - 1

    def add_spring(self, i: int, j, ke: float, kd: float, control: float):
        """Adds a spring between two particles in the system

        Args:
            i: The index of the first particle
            j: The index of the second particle
            ke: The elastic stiffness of the spring
            kd: The damping stiffness of the spring
            control: The actuation level of the spring

        Note:
            The spring is created with a rest-length based on the distance
            between the particles in their initial configuration.

        """
        self.spring_indices.append(i)
        self.spring_indices.append(j)
        self.spring_stiffness.append(ke)
        self.spring_damping.append(kd)
        self.spring_control.append(control)

        # compute rest length
        p = self.particle_q[i]
        q = self.particle_q[j]

        delta = np.subtract(p, q)
        l = np.sqrt(np.dot(delta, delta))

        self.spring_rest_length.append(l)

    def add_sewing_spring(self, i: int, j, orig_len: float, ke: float, kd: float, control: float):
        """Adds a spring between two particles in the system
        Args:
            i: The index of the first particle
            j: The index of the second particle
            ke: The elastic stiffness of the spring
            kd: The damping stiffness of the spring
            control: The actuation level of the spring
        Note:
            The spring is created with a rest-length based on the distance
            between the particles in their initial configuration.
        """
        self.spring_indices.append(i)
        self.spring_indices.append(j)
        self.spring_stiffness.append(ke)
        self.spring_damping.append(kd)
        self.spring_control.append(control)
        self.spring_rest_length.append(orig_len)

    def add_triangle(
        self,
        i: int,
        j: int,
        k: int,
        tri_ke: float = default_tri_ke,
        tri_ka: float = default_tri_ka,
        tri_kd: float = default_tri_kd,
        tri_drag: float = default_tri_drag,
        tri_lift: float = default_tri_lift,
    ) -> float:
        """Adds a triangular FEM element between three particles in the system.

        Triangles are modeled as viscoelastic elements with elastic stiffness and damping
        Parameters specified on the model. See model.tri_ke, model.tri_kd.

        Args:
            i: The index of the first particle
            j: The index of the second particle
            k: The index of the third particle

        Return:
            The area of the triangle

        Note:
            The triangle is created with a rest-length based on the distance
            between the particles in their initial configuration.

        Todo:
            * Expose elastic parameters on a per-element basis

        """
        # compute basis for 2D rest pose
        p = self.particle_q[i]
        q = self.particle_q[j]
        r = self.particle_q[k]

        qp = q - p
        rp = r - p

        # construct basis aligned with the triangle
        n = wp.normalize(wp.cross(qp, rp))
        e1 = wp.normalize(qp)
        e2 = wp.normalize(wp.cross(n, e1))

        R = np.array((e1, e2))
        M = np.array((qp, rp))

        D = R @ M.T

        area = np.linalg.det(D) / 2.0

        if area <= 0.0:
            print("inverted or degenerate triangle element")
            return 0.0
        else:
            inv_D = np.linalg.inv(D)

            self.tri_indices.append((i, j, k))
            self.tri_poses.append(inv_D.tolist())
            self.tri_activations.append(0.0)
            self.tri_materials.append((tri_ke, tri_ka, tri_kd, tri_drag, tri_lift))
            return area

    def add_triangles(
        self,
        i: List[int],
        j: List[int],
        k: List[int],
        tri_ke: Optional[List[float]] = None,
        tri_ka: Optional[List[float]] = None,
        tri_kd: Optional[List[float]] = None,
        tri_drag: Optional[List[float]] = None,
        tri_lift: Optional[List[float]] = None,
    ) -> List[float]:
        """Adds triangular FEM elements between groups of three particles in the system.

        Triangles are modeled as viscoelastic elements with elastic stiffness and damping
        Parameters specified on the model. See model.tri_ke, model.tri_kd.

        Args:
            i: The indices of the first particle
            j: The indices of the second particle
            k: The indices of the third particle

        Return:
            The areas of the triangles

        Note:
            A triangle is created with a rest-length based on the distance
            between the particles in their initial configuration.

        """
        # 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]

        qp = q - p
        rp = r - p

        def normalized(a):
            l = np.linalg.norm(a, axis=-1, keepdims=True)
            l[l == 0] = 1.0
            return a / l

        n = normalized(np.cross(qp, rp))
        e1 = normalized(qp)
        e2 = normalized(np.cross(n, e1))

        R = np.concatenate((e1[..., None], e2[..., None]), axis=-1)
        M = np.concatenate((qp[..., None], rp[..., None]), axis=-1)

        D = np.matmul(R.transpose(0, 2, 1), M)

        areas = np.linalg.det(D) / 2.0
        areas[areas < 0.0] = 0.0
        valid_inds = (areas > 0.0).nonzero()[0]
        if len(valid_inds) < len(areas):
            print("inverted or degenerate triangle elements")

        D[areas == 0.0] = np.eye(2)[None, ...]
        inv_D = np.linalg.inv(D)

        inds = np.concatenate((i[valid_inds, None], j[valid_inds, None], k[valid_inds, None]), axis=-1)

        self.tri_indices.extend(inds.tolist())
        self.tri_poses.extend(inv_D[valid_inds].tolist())
        self.tri_activations.extend([0.0] * len(valid_inds))

        def init_if_none(arr, defaultValue):
            if arr is None:
                return [defaultValue] * len(areas)
            return arr

        tri_ke = init_if_none(tri_ke, self.default_tri_ke)
        tri_ka = init_if_none(tri_ka, self.default_tri_ka)
        tri_kd = init_if_none(tri_kd, self.default_tri_kd)
        tri_drag = init_if_none(tri_drag, self.default_tri_drag)
        tri_lift = init_if_none(tri_lift, self.default_tri_lift)

        self.tri_materials.extend(
            zip(
                np.array(tri_ke)[valid_inds],
                np.array(tri_ka)[valid_inds],
                np.array(tri_kd)[valid_inds],
                np.array(tri_drag)[valid_inds],
                np.array(tri_lift)[valid_inds],
            )
        )
        return areas.tolist()

    def create_triangle_with_edge_lengths(self, a, b, c, tri_indices):
        # Given edge lengths a, b, c of a triangle
        # Check if the edge lengths can form a valid triangle
        if a + b <= c or a + c <= b or b + c <= a:
            print("{}::Error::Invalid edge lengths {}. Not possible to form a triangle "
                  "with vertex indices {}.".format(self.__class__.__name__, [a, b, c], tri_indices))
            raise RuntimeError("{}::Error::Invalid edge lengths {}. "
                                      "Not possible to form a triangle with vertex indices {}.".format(
                self.__class__.__name__, [a, b, c], tri_indices))
        # Calculate angles in radians
        angle_A = np.arccos((b ** 2 + c ** 2 - a ** 2) / (2 * b * c))
        # Coordinates of the vertices
        vertex_A = np.array([0, 0, 0])
        vertex_B = np.array([b, 0, 0])
        vertex_C = np.array([c * np.cos(angle_A), c * np.sin(angle_A), 0])
        return vertex_A, vertex_B, vertex_C

    def add_triangles_stitching(
            self,
            i: List[int],
            j: List[int],
            k: List[int],
            area: float,
            res: float,
            orig_lens: dict = None,
            tri_ke: Optional[List[float]] = None,
            tri_ka: Optional[List[float]] = None,
            tri_kd: Optional[List[float]] = None,
            tri_drag: Optional[List[float]] = None,
            tri_lift: Optional[List[float]] = None,
    ) -> List[float]:
        """Adds triangular FEM elements between groups of three particles in the system.
        Triangles are modeled as viscoelastic elements with elastic stiffness and damping
        Parameters specified on the model. See model.tri_ke, model.tri_kd.
        Args:
            i: The indices of the first particle
            j: The indices of the second particle
            k: The indices of the third particle
        Return:
            The areas of the triangles
        Note:
            A triangle is created with a rest-length based on the distance
            between the particles in their initial configuration.
        """

        if orig_lens is None:
            # compute basis for a equilateral triangle
            p, q, r = self.create_triangle_with_edge_lengths(res, res, res, [])
            qp = q - p
            rp = r - p
            n = wp.normalize(wp.cross(qp, rp))
            e1 = wp.normalize(qp)
            e2 = wp.normalize(wp.cross(n, e1))
            R = np.array((e1, e2))
            M = np.array((qp, rp))
            D = R @ M.T
            area_0 = np.linalg.det(D) / 2.0
            # compute basis for 2D rest pose -> a equilateral triangle with average area
            D = D * np.sqrt(area / area_0)
            # assert abs(np.linalg.det(D) / 2.0 - area) < 1e-6, "D not correct"
            D = np.repeat(D[np.newaxis, :, :], len(i), axis=0)
            areas = np.repeat(area, len(i))
        else:
            areas = []
            Ds = []
            for tri_inds in zip(i, j, k):
                ind1, ind2, ind3 = sorted(tri_inds)
                el1 = orig_lens[(ind1, ind2)]
                el2 = orig_lens[(ind2, ind3)]
                el3 = orig_lens[(ind1, ind3)]
                p, q, r = self.create_triangle_with_edge_lengths(el1, el2, el3, tri_inds)
                qp = q - p
                rp = r - p
                n = wp.normalize(wp.cross(qp, rp))
                e1 = wp.normalize(qp)
                e2 = wp.normalize(wp.cross(n, e1))
                R = np.array((e1, e2))
                M = np.array((qp, rp))
                D = R @ M.T
                area_0 = np.linalg.det(D) / 2.0
                Ds.append(D)
                areas.append(area_0)
            areas = np.array(areas)
            D = np.stack(Ds, axis=0)

        # following is the same as add_triangles
        areas[areas < 0.0] = 0.0
        valid_inds = (areas > 0.0).nonzero()[0]
        if len(valid_inds) < len(areas):
            print("inverted or degenerate triangle elements")
        D[areas == 0.0] = np.eye(2)[None, ...]
        inv_D = np.linalg.inv(D)
        inds = np.concatenate((i[valid_inds, None], j[valid_inds, None], k[valid_inds, None]), axis=-1)
        self.tri_indices.extend(inds.tolist())
        self.tri_poses.extend(inv_D[valid_inds].tolist())
        self.tri_activations.extend([0.0] * len(valid_inds))

        def init_if_none(arr, defaultValue):
            if arr is None:
                return [defaultValue] * len(areas)
            return arr

        tri_ke = init_if_none(tri_ke, self.default_tri_ke)
        tri_ka = init_if_none(tri_ka, self.default_tri_ka)
        tri_kd = init_if_none(tri_kd, self.default_tri_kd)
        tri_drag = init_if_none(tri_drag, self.default_tri_drag)
        tri_lift = init_if_none(tri_lift, self.default_tri_lift)
        self.tri_materials.extend(
            zip(
                np.array(tri_ke)[valid_inds],
                np.array(tri_ka)[valid_inds],
                np.array(tri_kd)[valid_inds],
                np.array(tri_drag)[valid_inds],
                np.array(tri_lift)[valid_inds],
            )
        )
        return areas.tolist()

    def add_tetrahedron(
        self, i: int, j: int, k: int, l: int, k_mu: float = 1.0e3, k_lambda: float = 1.0e3, k_damp: float = 0.0
    ) -> float:
        """Adds a tetrahedral FEM element between four particles in the system.

        Tetrahedra are modeled as viscoelastic elements with a NeoHookean energy
        density based on [Smith et al. 2018].

        Args:
            i: The index of the first particle
            j: The index of the second particle
            k: The index of the third particle
            l: The index of the fourth particle
            k_mu: The first elastic Lame parameter
            k_lambda: The second elastic Lame parameter
            k_damp: The element's damping stiffness

        Return:
            The volume of the tetrahedron

        Note:
            The tetrahedron is created with a rest-pose based on the particle's initial configuration

        """
        # 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])
        s = np.array(self.particle_q[l])

        qp = q - p
        rp = r - p
        sp = s - p

        Dm = np.array((qp, rp, sp)).T
        volume = np.linalg.det(Dm) / 6.0

        if volume <= 0.0:
            print("inverted tetrahedral element")
        else:
            inv_Dm = np.linalg.inv(Dm)

            self.tet_indices.append((i, j, k, l))
            self.tet_poses.append(inv_Dm.tolist())
            self.tet_activations.append(0.0)
            self.tet_materials.append((k_mu, k_lambda, k_damp))

        return volume

    def add_edge(
        self,
        i: int,
        j: int,
        k: int,
        l: int,
        rest: float = None,
        edge_ke: float = default_edge_ke,
        edge_kd: float = default_edge_kd,
    ):
        """Adds a bending edge element between four particles in the system.

        Bending elements are designed to be between two connected triangles. Then
        bending energy is based of [Bridson et al. 2002]. Bending stiffness is controlled
        by the `model.tri_kb` parameter.

        Args:
            i: The index of the first particle
            j: The index of the second particle
            k: The index of the third particle
            l: The index of the fourth particle
            rest: The rest angle across the edge in radians, if not specified it will be computed

        Note:
            The edge lies between the particles indexed by 'k' and 'l' parameters with the opposing
            vertices indexed by 'i' and 'j'. This defines two connected triangles with counter clockwise
            winding: (i, k, l), (j, l, k).

        """
        # compute rest angle
        if rest is None:
            x1 = self.particle_q[i]
            x2 = self.particle_q[j]
            x3 = self.particle_q[k]
            x4 = self.particle_q[l]

            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)

            angle = math.acos(d)
            sign = np.sign(np.dot(np.cross(n2, n1), e))

            rest = angle * sign

        self.edge_indices.append((i, j, k, l))
        self.edge_rest_angle.append(rest)
        self.edge_bending_properties.append((edge_ke, edge_kd))

    def add_edges(
        self,
        i,
        j,
        k,
        l,
        rest: Optional[List[float]] = None,
        edge_ke: Optional[List[float]] = None,
        edge_kd: Optional[List[float]] = None,
    ):
        """Adds bending edge elements between groups of four particles in the system.

        Bending elements are designed to be between two connected triangles. Then
        bending energy is based of [Bridson et al. 2002]. Bending stiffness is controlled
        by the `model.tri_kb` parameter.

        Args:
            i: The indices of the first particle
            j: The indices of the second particle
            k: The indices of the third particle
            l: The indices of the fourth particle
            rest: The rest angles across the edges in radians, if not specified they will be computed

        Note:
            The edge lies between the particles indexed by 'k' and 'l' parameters with the opposing
            vertices indexed by 'i' and 'j'. This defines two connected triangles with counter clockwise
            winding: (i, k, l), (j, l, k).

        """
        if rest is None:
            # compute rest angle
            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]

            def normalized(a):
                l = np.linalg.norm(a, axis=-1, keepdims=True)
                l[l == 0] = 1.0
                return a / l

            n1 = normalized(np.cross(x3 - x1, x4 - x1))
            n2 = normalized(np.cross(x4 - x2, x3 - x2))
            e = normalized(x4 - x3)

            def dot(a, b):
                return (a * b).sum(axis=-1)

            d = np.clip(dot(n2, n1), -1.0, 1.0)

            angle = np.arccos(d)
            sign = np.sign(dot(np.cross(n2, n1), e))

            rest = angle * sign

        inds = np.concatenate((i[:, None], j[:, None], k[:, None], l[:, None]), axis=-1)

        self.edge_indices.extend(inds.tolist())
        self.edge_rest_angle.extend(rest.tolist())

        def init_if_none(arr, defaultValue):
            if arr is None:
                return [defaultValue] * len(i)
            return arr

        edge_ke = init_if_none(edge_ke, self.default_edge_ke)
        edge_kd = init_if_none(edge_kd, self.default_edge_kd)

        self.edge_bending_properties.extend(zip(edge_ke, edge_kd))

    def add_cloth_grid(
        self,
        pos: Vec3,
        rot: Quat,
        vel: Vec3,
        dim_x: int,
        dim_y: int,
        cell_x: float,
        cell_y: float,
        mass: float,
        reverse_winding: bool = False,
        fix_left: bool = False,
        fix_right: bool = False,
        fix_top: bool = False,
        fix_bottom: bool = False,
        tri_ke: float = default_tri_ke,
        tri_ka: float = default_tri_ka,
        tri_kd: float = default_tri_kd,
        tri_drag: float = default_tri_drag,
        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

        Creates a rectangular grid of particles with FEM triangles and bending elements
        automatically.

        Args:
            pos: The position of the cloth in world space
            rot: The orientation of the cloth in world space
            vel: The velocity of the cloth in world space
            dim_x_: The number of rectangular cells along the x-axis
            dim_y: The number of rectangular cells along the y-axis
            cell_x: The width of each cell in the x-direction
            cell_y: The width of each cell in the y-direction
            mass: The mass of each particle
            reverse_winding: Flip the winding of the mesh
            fix_left: Make the left-most edge of particles kinematic (fixed in place)
            fix_right: Make the right-most edge of particles kinematic
            fix_top: Make the top-most edge of particles kinematic
            fix_bottom: Make the bottom-most edge of particles kinematic

        """

        def grid_index(x, y, dim_x):
            return y * dim_x + x

        start_vertex = len(self.particle_q)
        start_tri = len(self.tri_indices)

        for y in range(0, dim_y + 1):
            for x in range(0, dim_x + 1):
                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:
                    m = 0.0
                elif x == dim_x and fix_right:
                    m = 0.0
                elif y == 0 and fix_bottom:
                    m = 0.0
                elif y == dim_y and fix_top:
                    m = 0.0

                self.add_particle(p, vel, m)

                if x > 0 and y > 0:
                    if reverse_winding:
                        tri1 = (
                            start_vertex + grid_index(x - 1, y - 1, dim_x + 1),
                            start_vertex + grid_index(x, y - 1, dim_x + 1),
                            start_vertex + grid_index(x, y, dim_x + 1),
                        )

                        tri2 = (
                            start_vertex + grid_index(x - 1, y - 1, dim_x + 1),
                            start_vertex + grid_index(x, y, dim_x + 1),
                            start_vertex + grid_index(x - 1, y, dim_x + 1),
                        )

                        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)

                    else:
                        tri1 = (
                            start_vertex + grid_index(x - 1, y - 1, dim_x + 1),
                            start_vertex + grid_index(x, y - 1, dim_x + 1),
                            start_vertex + grid_index(x - 1, y, dim_x + 1),
                        )

                        tri2 = (
                            start_vertex + grid_index(x, y - 1, dim_x + 1),
                            start_vertex + grid_index(x, y, dim_x + 1),
                            start_vertex + grid_index(x - 1, y, dim_x + 1),
                        )

                        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_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:
                continue

            self.add_edge(
                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)))   # FIXME This edge does not exist? 
            spring_indices.add((min(e.o0, e.v0), max(e.o0, e.v0)))   # There should ne 5 edges, not 6
            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,
        rot: Quat,
        scale: float,
        vel: Vec3,
        vertices: List[Vec3],
        indices: List[int],
        density: float,
        edge_callback=None,
        face_callback=None,
        tri_ke: float = default_tri_ke,
        tri_ka: float = default_tri_ka,
        tri_kd: float = default_tri_kd,
        tri_drag: float = default_tri_drag,
        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

        Creates one FEM triangle element and one bending element for every face
        and edge in the input triangle mesh

        Args:
            pos: The position of the cloth in world space
            rot: The orientation of the cloth in world space
            vel: The velocity of the cloth in world space
            vertices: A list of vertex positions
            indices: A list of triangle indices, 3 entries per-face
            density: The density per-area of the mesh
            edge_callback: A user callback when an edge is created
            face_callback: A user callback when a face is created

        Note:

            The mesh should be two manifold.
        """
        num_tris = int(len(indices) / 3)

        start_vertex = len(self.particle_q)
        start_tri = len(self.tri_indices)

        # particles
        for v in vertices:
            p = wp.quat_rotate(rot, wp.vec3(v * scale)) + pos

            self.add_particle(p, vel, 0.0)

        # triangles
        inds = start_vertex + np.array(indices)
        self.particle_inds.extend(inds.reshape(-1).tolist())
        inds = inds.reshape(-1, 3)
        areas = self.add_triangles(
            inds[:, 0],
            inds[:, 1],
            inds[:, 2],
            [tri_ke] * num_tris,
            [tri_ka] * num_tris,
            [tri_kd] * num_tris,
            [tri_drag] * num_tris,
            [tri_lift] * num_tris,
        )

        for t in range(num_tris):
            area = areas[t]

            self.particle_mass[inds[t, 0]] += density * area / 3.0
            self.particle_mass[inds[t, 1]] += density * area / 3.0
            self.particle_mass[inds[t, 2]] += density * area / 3.0

        end_tri = len(self.tri_indices)

        adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)

        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],
            edgeinds[:, 2],
            edgeinds[:, 3],
            edge_ke=[edge_ke] * len(edgeinds),
            edge_kd=[edge_kd] * len(edgeinds),
        )

        if add_springs:
            spring_indices = set()
            for i, j, k, l in edgeinds:
                spring_indices.add((min(i, k), max(i, k)))
                spring_indices.add((min(i, l), max(i, l)))
                spring_indices.add((min(k, l), max(k, l)))
                spring_indices.add((min(j, k), max(j, k)))
                spring_indices.add((min(j, l), max(j, l)))

            for i, j in spring_indices:
                self.add_spring(i, j, spring_ke, spring_kd, control=0.0)

    def add_cloth_mesh_sewing(
            self,
            pos: Vec3,
            rot: Quat,
            scale: float,
            vel: Vec3,
            vertices: List[Vec3],
            indices: List[int],
            density: float,
            resolution_scale: float,
            orig_lens: dict = None,
            stitching_vertices: List[int] = None,
            radius: float = 0.1,
            edge_callback=None,
            face_callback=None,
            tri_ke: float = default_tri_ke,
            tri_ka: float = default_tri_ka,
            tri_kd: float = default_tri_kd,
            tri_drag: float = default_tri_drag,
            tri_lift: float = default_tri_lift,
            edge_ke: float = default_edge_ke,
            edge_kd: float = default_edge_kd,
    ):
        """Helper to create a cloth model from a regular triangle mesh
        Creates one FEM triangle element and one bending element for every face
        and edge in the input triangle mesh. # TODO Add specific description
        Args:
            pos: The position of the cloth in world space
            rot: The orientation of the cloth in world space
            vel: The velocity of the cloth in world space
            vertices: A list of vertex positions
            indices: A list of triangle indices, 3 entries per-face
            density: The density per-area of the mesh
            edge_callback: A user callback when an edge is created
            face_callback: A user callback when a face is created
        Note:
            The mesh should be two manifold.
        """
        num_tris = int(len(indices) / 3)
        start_vertex = len(self.particle_q)
        start_tri = len(self.tri_indices)
        # particles
        for v in vertices:
            p = np.array(wp.quat_rotate(rot, wp.vec3(v * scale))) + pos
            self.add_particle(p, vel, radius)
        # triangles
        inds = np.array(indices).reshape(-1, 3)

        # find all stitching triangles
        stitch_vertice_set = set(stitching_vertices)

        stitch_tri_indices = []
        normal_tri_indices = []
        for index, (i, j, k) in enumerate(inds):
            stitch_i = i in stitch_vertice_set  # Could be replaced by "< n_stitches"
            stitch_j = j in stitch_vertice_set
            stitch_k = k in stitch_vertice_set
            if stitch_i or stitch_j or stitch_k:
                stitch_tri_indices.append(index)
            else:
                normal_tri_indices.append(index)

        print("stitch_tri_indices", len(stitch_tri_indices))
        print("normal_tri_indices", len(normal_tri_indices))
        inds = start_vertex + inds
        self.particle_inds.extend(inds.reshape(-1).tolist())
        # normal triangles
        num_normal_tris = len(normal_tri_indices)
        areas_normal = self.add_triangles(
            inds[normal_tri_indices, 0],
            inds[normal_tri_indices, 1],
            inds[normal_tri_indices, 2],
            [tri_ke] * num_normal_tris,
            [tri_ka] * num_normal_tris,
            [tri_kd] * num_normal_tris,
            [tri_drag] * num_normal_tris,
            [tri_lift] * num_normal_tris,
        )

        # stitching triangles
        num_stitch_tris = len(stitch_tri_indices)
        areas_stitching = self.add_triangles_stitching(
            inds[stitch_tri_indices, 0],
            inds[stitch_tri_indices, 1],
            inds[stitch_tri_indices, 2],
            np.mean(areas_normal),
            resolution_scale,
            orig_lens,
            [tri_ke] * num_stitch_tris,
            [tri_ka] * num_stitch_tris,
            [tri_kd] * num_stitch_tris,
            [tri_drag] * num_stitch_tris,
            [tri_lift] * num_stitch_tris,
        )

        areas = np.zeros(num_tris)
        areas[normal_tri_indices] = areas_normal
        areas[stitch_tri_indices] = areas_stitching
        for t in range(num_tris):
            area = areas[t]
            self.particle_mass[inds[t, 0]] += density * area / 3.0
            self.particle_mass[inds[t, 1]] += density * area / 3.0
            self.particle_mass[inds[t, 2]] += density * area / 3.0
        end_tri = len(self.tri_indices)
        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()])
        self.add_edges(
            edgeinds[:, 0],
            edgeinds[:, 1],
            edgeinds[:, 2],
            edgeinds[:, 3],
            edge_ke=[edge_ke] * len(edgeinds),
            edge_kd=[edge_kd] * len(edgeinds),
            rest=np.zeros(len(edgeinds))  # set all the rest angle to 0
        )

    def add_cloth_mesh_sewing_spring(
            self,
            pos: Vec3,
            rot: Quat,
            scale: float,
            vel: Vec3,
            vertices: List[Vec3],
            indices: List[int],
            density: float,
            resolution_scale: float,
            orig_lens: dict = None,
            stitching_vertices: List[int] = None,
            radius: float = 0.1,
            edge_callback=None,
            face_callback=None,
            tri_ke: float = default_tri_ke,
            tri_ka: float = default_tri_ka,
            tri_kd: float = default_tri_kd,
            tri_drag: float = default_tri_drag,
            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
        Creates one FEM triangle element and one bending element for every face
        and edge in the input triangle mesh # TODO Add specific description
        Args:
            pos: The position of the cloth in world space
            rot: The orientation of the cloth in world space
            vel: The velocity of the cloth in world space
            vertices: A list of vertex positions
            indices: A list of triangle indices, 3 entries per-face
            density: The density per-area of the mesh
            edge_callback: A user callback when an edge is created
            face_callback: A user callback when a face is created
        Note:
            The mesh should be two manifold.
        """
        num_tris = int(len(indices) / 3)
        start_vertex = len(self.particle_q)
        start_tri = len(self.tri_indices)
        # particles
        for v in vertices:
            p = wp.quat_rotate(rot, wp.vec3(v * scale)) + pos
            self.add_particle(p, vel, radius)
        # triangles
        inds = np.array(indices).reshape(-1, 3)
        # find all stitching triangles
        stitch_vertices_set = set(stitching_vertices)

        stitch_tri_indices = []
        normal_tri_indices = []
        for index, (i, j, k) in enumerate(inds):
            stitch_i = i in stitch_vertices_set  # Could be replaced by "< n_stitches"
            stitch_j = j in stitch_vertices_set
            stitch_k = k in stitch_vertices_set
            if stitch_i or stitch_j or stitch_k:
                stitch_tri_indices.append(index)
            else:
                normal_tri_indices.append(index)

        inds = start_vertex + inds
        self.particle_inds.extend(inds.reshape(-1).tolist())
        # normal triangles
        num_normal_tris = len(normal_tri_indices)
        areas_normal = self.add_triangles(
            inds[normal_tri_indices, 0],
            inds[normal_tri_indices, 1],
            inds[normal_tri_indices, 2],
            [tri_ke] * num_normal_tris,
            [tri_ka] * num_normal_tris,
            [tri_kd] * num_normal_tris,
            [tri_drag] * num_normal_tris,
            [tri_lift] * num_normal_tris,
        )

        # stitching triangles
        num_stitch_tris = len(stitch_tri_indices)
        areas_stitching = self.add_triangles_stitching(
            inds[stitch_tri_indices, 0],
            inds[stitch_tri_indices, 1],
            inds[stitch_tri_indices, 2],
            np.mean(areas_normal),
            resolution_scale,
            orig_lens,
            [tri_ke] * num_stitch_tris,
            [tri_ka] * num_stitch_tris,
            [tri_kd] * num_stitch_tris,
            [tri_drag] * num_stitch_tris,
            [tri_lift] * num_stitch_tris,
        )

        areas = np.zeros(num_tris)
        areas[normal_tri_indices] = areas_normal
        areas[stitch_tri_indices] = areas_stitching
        for t in range(num_tris):
            area = areas[t]
            self.particle_mass[inds[t, 0]] += density * area / 3.0
            self.particle_mass[inds[t, 1]] += density * area / 3.0
            self.particle_mass[inds[t, 2]] += density * area / 3.0

        end_tri = len(self.tri_indices)
        adj = wp.utils.MeshAdjacency(self.tri_indices[start_tri:end_tri], end_tri - start_tri)
        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],
            edgeinds[:, 2],
            edgeinds[:, 3],
            edge_ke=[edge_ke] * len(edgeinds),
            edge_kd=[edge_kd] * len(edgeinds),
            rest=np.zeros(len(edgeinds))  # set all the rest angle to 0
        )

        if add_springs:
            spring_indices = set()
            #add a spring between each edge
            for i, j, k, l in edgeinds:
                spring_indices.add((min(i, k), max(i, k)))
                spring_indices.add((min(i, l), max(i, l)))
                spring_indices.add((min(k, l), max(k, l)))
                spring_indices.add((min(j, k), max(j, k)))
                spring_indices.add((min(j, l), max(j, l)))
            for (i,j) in spring_indices:
                if orig_lens is None:
                    self.add_spring(i, j, spring_ke, spring_kd, control=0.0)
                elif (i,j) in orig_lens.keys():
                    self.add_sewing_spring(i,j,orig_lens[(i,j)],spring_ke, spring_kd, control=0.0)
                else:
                    self.add_spring(i,j,spring_ke, spring_kd, control=0.0)

    def add_cloth_reference_shape_mesh(
            self,
            mesh: Mesh,
            name: str = None,
            pos: Vec3 = (0.0, 0.0, 0.0),
            rot: Quat = (0.0, 0.0, 0.0, 1.0),
            scale: Vec3 = (1.0, 1.0, 1.0),
    ):
        """Adds a triangle mesh reference shape for garment.
        Args:
            mesh: The mesh object
            pos: The location of the shape with respect to the parent frame
            rot: The rotation of the shape with respect to the parent frame
            scale: Scale to use for the collider
        """
        if name is None:
            name = len(self.cloth_reference_shapes)
        self.cloth_reference_names.append(name)
        self.cloth_reference_transforms.append(wp.transform(pos, rot))
        self.cloth_reference_scale.append(scale)
        self.cloth_reference_shapes.append(mesh)

    def add_cloth_reference_labels(
            self,
            labels: List,
            drag_pairs: List
    ):
        assert len(labels) == len(
            self.particle_q), "ERROR: The number of reference labels must be equal to the number of particles. Labels adding failed."
        assert len(self.cloth_reference_names) != 0, "ERROR: No reference shapes are added yet. Labels adding failed."
        # if len(set(labels)) != len(self.cloth_reference_names):
        #     print("Warning: Number of reference shapes {} != number of types of labels {}.".format(len(self.cloth_reference_names), len(set(labels))))
        name2index = {}
        for idx, name in enumerate(self.cloth_reference_names):
            name2index[name] = idx
        idx_labels = []
        for label in labels:
            if label in name2index:
                idx_labels.append(name2index[label])
            elif label is None or label == "" or label == -1:
                idx_labels.append(-1)
            else:
                print("The label {} is not in the reference shape list. Ignoring...".format(label))
                idx_labels.append(-1)
        self.particle_reference_label = idx_labels
        self.drag_label_pairs = []
        for lab1, lab2 in drag_pairs:
            # Adding for existing labels
            if lab1 in self.cloth_reference_names and lab2 in self.cloth_reference_names:
                self.drag_label_pairs.append(
                    wp.vec2i(name2index[lab1], name2index[lab2])
                )
        return name2index
   
    def add_particle_grid(
        self,
        pos: Vec3,
        rot: Quat,
        vel: Vec3,
        dim_x: int,
        dim_y: int,
        dim_z: int,
        cell_x: float,
        cell_y: float,
        cell_z: float,
        mass: float,
        jitter: float,
        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 = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                    m = mass

                    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
                    else:
                        r = radius_mean
                    self.add_particle(p, vel, m, r)

    def add_soft_grid(
        self,
        pos: Vec3,
        rot: Quat,
        vel: Vec3,
        dim_x: int,
        dim_y: int,
        dim_z: int,
        cell_x: float,
        cell_y: float,
        cell_z: float,
        density: float,
        k_mu: float,
        k_lambda: float,
        k_damp: float,
        fix_left: bool = False,
        fix_right: bool = False,
        fix_top: bool = False,
        fix_bottom: bool = False,
        tri_ke: float = default_tri_ke,
        tri_ka: float = default_tri_ka,
        tri_kd: float = default_tri_kd,
        tri_drag: float = default_tri_drag,
        tri_lift: float = default_tri_lift,
    ):
        """Helper to create a rectangular tetrahedral FEM grid

        Creates a regular grid of FEM tetrahedra and surface triangles. Useful for example
        to create beams and sheets. Each hexahedral cell is decomposed into 5
        tetrahedral elements.

        Args:
            pos: The position of the solid in world space
            rot: The orientation of the solid in world space
            vel: The velocity of the solid in world space
            dim_x_: The number of rectangular cells along the x-axis
            dim_y: The number of rectangular cells along the y-axis
            dim_z: The number of rectangular cells along the z-axis
            cell_x: The width of each cell in the x-direction
            cell_y: The width of each cell in the y-direction
            cell_z: The width of each cell in the z-direction
            density: The density of each particle
            k_mu: The first elastic Lame parameter
            k_lambda: The second elastic Lame parameter
            k_damp: The damping stiffness
            fix_left: Make the left-most edge of particles kinematic (fixed in place)
            fix_right: Make the right-most edge of particles kinematic
            fix_top: Make the top-most edge of particles kinematic
            fix_bottom: Make the bottom-most edge of particles kinematic
        """

        start_vertex = len(self.particle_q)

        mass = cell_x * cell_y * cell_z * density

        for z in range(dim_z + 1):
            for y in range(dim_y + 1):
                for x in range(dim_x + 1):
                    v = wp.vec3(x * cell_x, y * cell_y, z * cell_z)
                    m = mass

                    if fix_left and x == 0:
                        m = 0.0

                    if fix_right and x == dim_x:
                        m = 0.0

                    if fix_top and y == dim_y:
                        m = 0.0

                    if fix_bottom and y == 0:
                        m = 0.0

                    p = wp.quat_rotate(rot, v) + pos

                    self.add_particle(p, vel, m)

        # dict of open faces
        faces = {}

        def add_face(i: int, j: int, k: int):
            key = tuple(sorted((i, j, k)))

            if key not in faces:
                faces[key] = (i, j, k)
            else:
                del faces[key]

        def add_tet(i: int, j: int, k: int, l: int):
            self.add_tetrahedron(i, j, k, l, k_mu, k_lambda, k_damp)

            add_face(i, k, j)
            add_face(j, k, l)
            add_face(i, j, l)
            add_face(i, l, k)

        def grid_index(x, y, z):
            return (dim_x + 1) * (dim_y + 1) * z + (dim_x + 1) * y + x

        for z in range(dim_z):
            for y in range(dim_y):
                for x in range(dim_x):
                    v0 = grid_index(x, y, z) + start_vertex
                    v1 = grid_index(x + 1, y, z) + start_vertex
                    v2 = grid_index(x + 1, y, z + 1) + start_vertex
                    v3 = grid_index(x, y, z + 1) + start_vertex
                    v4 = grid_index(x, y + 1, z) + start_vertex
                    v5 = grid_index(x + 1, y + 1, z) + start_vertex
                    v6 = grid_index(x + 1, y + 1, z + 1) + start_vertex
                    v7 = grid_index(x, y + 1, z + 1) + start_vertex

                    if (x & 1) ^ (y & 1) ^ (z & 1):
                        add_tet(v0, v1, v4, v3)
                        add_tet(v2, v3, v6, v1)
                        add_tet(v5, v4, v1, v6)
                        add_tet(v7, v6, v3, v4)
                        add_tet(v4, v1, v6, v3)

                    else:
                        add_tet(v1, v2, v5, v0)
                        add_tet(v3, v0, v7, v2)
                        add_tet(v4, v7, v0, v5)
                        add_tet(v6, v5, v2, v7)
                        add_tet(v5, v2, v7, v0)

        # add triangles
        for k, v in faces.items():
            self.add_triangle(v[0], v[1], v[2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)

    def add_soft_mesh(
        self,
        pos: Vec3,
        rot: Quat,
        scale: float,
        vel: Vec3,
        vertices: List[Vec3],
        indices: List[int],
        density: float,
        k_mu: float,
        k_lambda: float,
        k_damp: float,
        tri_ke: float = default_tri_ke,
        tri_ka: float = default_tri_ka,
        tri_kd: float = default_tri_kd,
        tri_drag: float = default_tri_drag,
        tri_lift: float = default_tri_lift,
    ):
        """Helper to create a tetrahedral model from an input tetrahedral mesh

        Args:
            pos: The position of the solid in world space
            rot: The orientation of the solid in world space
            vel: The velocity of the solid in world space
            vertices: A list of vertex positions
            indices: A list of tetrahedron indices, 4 entries per-element
            density: The density per-area of the mesh
            k_mu: The first elastic Lame parameter
            k_lambda: The second elastic Lame parameter
            k_damp: The damping stiffness
        """
        num_tets = int(len(indices) / 4)

        start_vertex = len(self.particle_q)

        # dict of open faces
        faces = {}

        def add_face(i, j, k):
            key = tuple(sorted((i, j, k)))

            if key not in faces:
                faces[key] = (i, j, k)
            else:
                del faces[key]

        # add particles
        for v in vertices:
            p = wp.quat_rotate(rot, wp.vec3(v * scale)) + np.array(pos)

            self.add_particle(p, vel, 0.0)

        # add tetrahedra
        for t in range(num_tets):
            v0 = start_vertex + indices[t * 4 + 0]
            v1 = start_vertex + indices[t * 4 + 1]
            v2 = start_vertex + indices[t * 4 + 2]
            v3 = start_vertex + indices[t * 4 + 3]

            volume = self.add_tetrahedron(v0, v1, v2, v3, k_mu, k_lambda, k_damp)

            # distribute volume fraction to particles
            if volume > 0.0:
                self.particle_mass[v0] += density * volume / 4.0
                self.particle_mass[v1] += density * volume / 4.0
                self.particle_mass[v2] += density * volume / 4.0
                self.particle_mass[v3] += density * volume / 4.0

                # build open faces
                add_face(v0, v2, v1)
                add_face(v1, v2, v3)
                add_face(v0, v1, v3)
                add_face(v0, v3, v2)

        # add triangles
        for k, v in faces.items():
            try:
                self.add_triangle(v[0], v[1], v[2], tri_ke, tri_ka, tri_kd, tri_drag, tri_lift)
            except np.linalg.LinAlgError:
                continue

    def add_attachment(
            self, 
            particle_indices: List,
            attachment_point: Union[wp.vec3, List[wp.vec3]],
            attachment_norm: wp.vec3 = None,
            stiffness: float = default_spring_ke,
            damping: float = default_spring_kd
    ):
        """ 
            Attachment constraint allows fixing particle 
            locations to any point in space (equality attachment) or restrict them 
            to be located in a given half-space (inequality attachment)

            Args: 
                particle_indices: list of indices of particles to be constraint
                attachment_point: list of point locations, defining the target attachment locatoin
                    for each particle OR a singular 3D particle location, applied to all
                    particle_indices (useful for inequality constraints)
                attachment_norm (optional): 
                    if not given attachment is treated as an equality constraint 
                        (position matches attachment_point exactly),
                    otherwise norm defines an inequality constraint with 
                    the norm and attachment_point defining the half-space 
                    where each particle should be lying
                stiffness: constraint stiffness
                damping: constraint damping
        """
        self.attachment_indices += particle_indices
        if isinstance(attachment_point, list):
            self.attachment_point += attachment_point
        else:
            self.attachment_point += [attachment_point] * len(particle_indices)
        if attachment_norm is not None: 
            self.attachment_norm += [attachment_norm] * len(particle_indices)
        else:
            self.attachment_norm += [wp.vec3(0, 0, 0)] * len(particle_indices)
        self.attachment_stiffness += [stiffness] * len(particle_indices)
        self.attachment_damping += [damping] * len(particle_indices)

    # incrementally updates rigid body mass with additional mass and inertia expressed at a local to the body
    def _update_body_mass(self, i, m, I, p, q):
        if i == -1:
            return

        # find new COM
        new_mass = self.body_mass[i] + m

        if new_mass == 0.0:  # no mass
            return

        new_com = (self.body_com[i] * self.body_mass[i] + p * m) / new_mass

        # shift inertia to new COM
        com_offset = new_com - self.body_com[i]
        shape_offset = new_com - p

        new_inertia = transform_inertia(
            self.body_mass[i], self.body_inertia[i], com_offset, wp.quat_identity()
        ) + transform_inertia(m, I, shape_offset, q)

        self.body_mass[i] = new_mass
        self.body_inertia[i] = new_inertia
        self.body_com[i] = new_com

        if new_mass > 0.0:
            self.body_inv_mass[i] = 1.0 / new_mass
        else:
            self.body_inv_mass[i] = 0.0

        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

    def set_ground_plane(
        self,
        normal=None,
        offset=0.0,
        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,
    ):
        """
        Creates a ground plane for the world. If the normal is not specified,
        the up_vector of the ModelBuilder is used.
        """
        if normal is None:
            normal = self.up_vector
        self._ground_params = dict(
            plane=(*normal, offset), width=0.0, length=0.0, ke=ke, kd=kd, kf=kf, mu=mu, restitution=restitution
        )

    def _create_ground_plane(self):
        ground_id = self.add_shape_plane(**self._ground_params)
        self._ground_created = True
        # disable ground collisions as they will be treated separately
        for i in range(self.shape_count - 1):
            self.shape_collision_filter_pairs.add((i, ground_id))

    def _fill_face_filters(self, array, max_el_count=None, fill_value=0):
        """Create homogenious array out of current value"""
        # Max 2d dimention
        max_size_2 = 0
        for el in array:
            max_size_2 = max(max_size_2, len(el))
        # Even out 2d dimention
        for el in array:
            if len(el) < max_size_2:
                el += [[]] * (max_size_2 - len(el))

        # Max 3d dimention
        if max_el_count is None:
            max_el_count = 0
            for el in array:
                for subel in el:
                    max_el_count = max(max_el_count, len(subel))
        # Even out innermost dimention
        for el in array: 
            for subel in el:
                if len(subel) < max_el_count:
                    subel += [fill_value] * ((max_el_count) - len(subel))

    def _fill_particle_filters(self, array, max_count, fill_value=-1):
        """Fill up empty sublists of the array of lists s.t. every sublist has max_count of values"""

        for i in range(len(array)):
            if not len(array[i]):
                array[i] = [fill_value] * max_count
            if len(array[i]) != max_count:
                raise ValueError(
                    'Builder::ERROR::Incorrect number of particles '
                    f'{len(array[i])} in particle_collision_filters for shape {i}')

    def finalize(self, device=None, requires_grad=False) -> Model:
        """Convert this builder object to a concrete model for simulation.

        After building simulation elements this method should be called to transfer
        all data to device memory ready for simulation.

        Args:
            device: The simulation device to use, e.g.: 'cpu', 'cuda'
            requires_grad: Whether to enable gradient computation for the model

        Returns:

            A model object.
        """

        # ensure the env count is set correctly
        self.num_envs = max(1, self.num_envs)

        # add ground plane if not already created
        if not self._ground_created:
            self._create_ground_plane()

        # construct particle inv masses
        particle_inv_mass = []

        for m in self.particle_mass:
            if m > 0.0:
                particle_inv_mass.append(1.0 / m)
            else:
                particle_inv_mass.append(0.0)

        with wp.ScopedDevice(device):
            # -------------------------------------
            # construct Model (non-time varying) data

            m = Model(device)
            m.requires_grad = requires_grad

            m.num_envs = self.num_envs

            # ---------------------
            # particles

            # state (initial)
            m.particle_q = wp.array(self.particle_q, dtype=wp.vec3, requires_grad=requires_grad)
            m.particle_qd = wp.array(self.particle_qd, dtype=wp.vec3, requires_grad=requires_grad)
            m.particle_mass = wp.array(self.particle_mass, dtype=wp.float32, requires_grad=requires_grad)
            m.particle_inv_mass = wp.array(particle_inv_mass, dtype=wp.float32, requires_grad=requires_grad)
            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)

            # TODO It's only needed in case of simulating meshes. Under a flag in add cloth_mesh routines
            # build shape of particles  
            m.particle_shape = wp.Mesh(points=wp.array(self.particle_q, dtype=wp.vec3),
                                       indices=wp.array(self.particle_inds, dtype=int))
            # compute per-vertex normal for particle shape 
            m.particle_normal = wp.zeros(len(self.particle_q), dtype=wp.vec3)
            compute_vertex_normal(m.particle_shape, m.particle_normal)

            # ---------------------
            # collision geometry
            m.shape_transform = wp.array(self.shape_transform, dtype=wp.transform, requires_grad=requires_grad)  # TODOLOW Why not in geo?
            m.shape_body = wp.array(self.shape_body, dtype=wp.int32)
            m.body_shapes = self.body_shapes

            # build list of ids for geometry sources (meshes, sdfs)
            geo_sources = []
            finalized_meshes = {}  # do not duplicate meshes
            for geo in self.shape_geo_src:
                geo_hash = hash(geo)  # avoid repeated hash computations
                if geo:
                    if geo_hash not in finalized_meshes:
                        finalized_meshes[geo_hash] = geo.finalize(device=device)
                    geo_sources.append(finalized_meshes[geo_hash])
                else:
                    # add null pointer
                    geo_sources.append(0)

            # Garment per body-part collision references
            # TODO This should probably be under a flag in adding cloth mesh routines
            # build list of ids for cloth reference shapes
            cloth_reference_shapes = []
            for mesh in self.cloth_reference_shapes:
                cloth_reference_shapes.append(mesh.finalize(device=device))
            m.cloth_reference_shapes = self.cloth_reference_shapes   # Keep references to meshes
            m.cloth_reference_shape_ids = wp.array(cloth_reference_shapes, dtype=wp.uint64)
            m.cloth_reference_transforms = wp.array(self.cloth_reference_transforms, dtype=wp.transform,
                                                    requires_grad=requires_grad)
            m.cloth_reference_scale = wp.array(self.cloth_reference_scale, dtype=wp.vec3,
                                               requires_grad=requires_grad)
            m.particle_reference_label = wp.array(self.particle_reference_label, dtype=wp.int32)
            m.cloth_reference_drag_particles = wp.zeros(self.particle_count, dtype=wp.int32)
            m.drag_label_pairs = wp.array(self.drag_label_pairs, dtype=wp.vec2i)
            m.drag_label_pairs_count = len(self.drag_label_pairs) 

            # Overall collision shapes
            m.shape_geo.type = wp.array(self.shape_geo_type, dtype=wp.int32)
            m.shape_geo.source = wp.array(geo_sources, dtype=wp.uint64)
            m.shape_geo.scale = wp.array(self.shape_geo_scale, dtype=wp.vec3, requires_grad=requires_grad)
            m.shape_geo.is_solid = wp.array(self.shape_geo_is_solid, dtype=wp.uint8)
            m.shape_geo.thickness = wp.array(self.shape_geo_thickness, dtype=wp.float32, requires_grad=requires_grad)
            # Collision filters
            self._fill_face_filters(self.shape_geo_face_filters, fill_value=False)  # Even out dimentions
            m.shape_geo.face_filter = wp.array3d(self.shape_geo_face_filters, dtype=wp.int32)

            self._fill_particle_filters(self.shape_geo_particle_filters_id, len(self.particle_q), -1)   # Even out dimentions
            m.shape_geo.particle_filter_ids = wp.array2d(self.shape_geo_particle_filters_id, dtype=wp.int32)

            m.shape_geo_src = self.shape_geo_src  # used for rendering
            m.shape_geo_src_id = geo_sources  # used for smoothing

            m.shape_materials.ke = wp.array(self.shape_material_ke, dtype=wp.float32, requires_grad=requires_grad)
            m.shape_materials.kd = wp.array(self.shape_material_kd, dtype=wp.float32, requires_grad=requires_grad)
            m.shape_materials.kf = wp.array(self.shape_material_kf, dtype=wp.float32, requires_grad=requires_grad)
            m.shape_materials.mu = wp.array(self.shape_material_mu, dtype=wp.float32, requires_grad=requires_grad)
            m.shape_materials.restitution = wp.array(
                self.shape_material_restitution, dtype=wp.float32, requires_grad=requires_grad
            )

            m.shape_collision_filter_pairs = self.shape_collision_filter_pairs
            m.shape_collision_group = self.shape_collision_group
            m.shape_collision_group_map = self.shape_collision_group_map
            m.shape_collision_radius = wp.array(
                self.shape_collision_radius, dtype=wp.float32, requires_grad=requires_grad
            )
            m.shape_ground_collision = self.shape_ground_collision

            # ---------------------
            # Attachment constraints
            # NOTE: Special case implementation with vertical constraint only
            m.attachment_indices = wp.array(self.attachment_indices, dtype=wp.int32)
            m.attachment_point = wp.array(self.attachment_point, dtype=wp.vec3)
            m.attachment_norm = wp.array(self.attachment_norm, dtype=wp.vec3)
            m.vert_attach_constraint_lambdas = wp.array(
                shape=len(self.attachment_indices), dtype=wp.float32, requires_grad=requires_grad
            )
            m.attachment_stiffness = wp.array(
                self.attachment_stiffness, dtype=wp.float32, requires_grad=requires_grad)
            m.attachment_damping = wp.array(
                self.attachment_damping, dtype=wp.float32, requires_grad=requires_grad)
            

            # ---------------------
            # springs

            m.spring_indices = wp.array(self.spring_indices, dtype=wp.int32)
            m.spring_rest_length = wp.array(self.spring_rest_length, dtype=wp.float32, requires_grad=requires_grad)
            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

            m.tri_indices = wp.array(self.tri_indices, dtype=wp.int32)
            m.tri_poses = wp.array(self.tri_poses, dtype=wp.mat22, requires_grad=requires_grad)
            m.tri_activations = wp.array(self.tri_activations, dtype=wp.float32, requires_grad=requires_grad)
            m.tri_materials = wp.array(self.tri_materials, dtype=wp.float32, requires_grad=requires_grad)

            # ---------------------
            # edges

            m.edge_indices = wp.array(self.edge_indices, dtype=wp.int32)
            m.edge_rest_angle = wp.array(self.edge_rest_angle, dtype=wp.float32, requires_grad=requires_grad)
            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

            m.tet_indices = wp.array(self.tet_indices, dtype=wp.int32)
            m.tet_poses = wp.array(self.tet_poses, dtype=wp.mat33, requires_grad=requires_grad)
            m.tet_activations = wp.array(self.tet_activations, dtype=wp.float32, requires_grad=requires_grad)
            m.tet_materials = wp.array(self.tet_materials, dtype=wp.float32, requires_grad=requires_grad)

            # -----------------------
            # muscles

            # close the muscle waypoint indices
            muscle_start = copy.copy(self.muscle_start)
            muscle_start.append(len(self.muscle_bodies))

            m.muscle_start = wp.array(muscle_start, dtype=wp.int32)
            m.muscle_params = wp.array(self.muscle_params, dtype=wp.float32, requires_grad=requires_grad)
            m.muscle_bodies = wp.array(self.muscle_bodies, dtype=wp.int32)
            m.muscle_points = wp.array(self.muscle_points, dtype=wp.vec3, requires_grad=requires_grad)
            m.muscle_activation = wp.array(self.muscle_activation, dtype=wp.float32, requires_grad=requires_grad)

            # --------------------------------------
            # rigid bodies

            m.body_q = wp.array(self.body_q, dtype=wp.transform, requires_grad=requires_grad)
            m.body_qd = wp.array(self.body_qd, dtype=wp.spatial_vector, requires_grad=requires_grad)
            m.body_inertia = wp.array(self.body_inertia, dtype=wp.mat33, requires_grad=requires_grad)
            m.body_inv_inertia = wp.array(self.body_inv_inertia, dtype=wp.mat33, requires_grad=requires_grad)
            m.body_mass = wp.array(self.body_mass, dtype=wp.float32, requires_grad=requires_grad)
            m.body_inv_mass = wp.array(self.body_inv_mass, dtype=wp.float32, requires_grad=requires_grad)
            m.body_com = wp.array(self.body_com, dtype=wp.vec3, requires_grad=requires_grad)
            m.body_name = self.body_name

            # joints
            m.joint_count = self.joint_count
            m.joint_axis_count = self.joint_axis_count
            m.joint_type = wp.array(self.joint_type, dtype=wp.int32)
            m.joint_parent = wp.array(self.joint_parent, dtype=wp.int32)
            m.joint_child = wp.array(self.joint_child, dtype=wp.int32)
            m.joint_X_p = wp.array(self.joint_X_p, dtype=wp.transform, requires_grad=requires_grad)
            m.joint_X_c = wp.array(self.joint_X_c, dtype=wp.transform, requires_grad=requires_grad)
            m.joint_axis_start = wp.array(self.joint_axis_start, dtype=wp.int32)
            m.joint_axis_dim = wp.array(np.array(self.joint_axis_dim), dtype=wp.int32, ndim=2)
            m.joint_axis = wp.array(self.joint_axis, dtype=wp.vec3, requires_grad=requires_grad)
            m.joint_q = wp.array(self.joint_q, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_qd = wp.array(self.joint_qd, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_name = self.joint_name

            # dynamics properties
            # TODO unused joint_armature
            m.joint_armature = wp.array(self.joint_armature, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_target = wp.array(self.joint_target, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_target_ke = wp.array(self.joint_target_ke, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_target_kd = wp.array(self.joint_target_kd, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_axis_mode = wp.array(self.joint_axis_mode, dtype=wp.uint8)
            m.joint_act = wp.array(self.joint_act, dtype=wp.float32, requires_grad=requires_grad)

            m.joint_limit_lower = wp.array(self.joint_limit_lower, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_limit_upper = wp.array(self.joint_limit_upper, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_limit_ke = wp.array(self.joint_limit_ke, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_limit_kd = wp.array(self.joint_limit_kd, dtype=wp.float32, requires_grad=requires_grad)
            m.joint_linear_compliance = wp.array(
                self.joint_linear_compliance, dtype=wp.float32, requires_grad=requires_grad
            )
            m.joint_angular_compliance = wp.array(
                self.joint_angular_compliance, dtype=wp.float32, requires_grad=requires_grad
            )
            m.joint_enabled = wp.array(self.joint_enabled, dtype=wp.int32)

            # 'close' the start index arrays with a sentinel value
            joint_q_start = copy.copy(self.joint_q_start)
            joint_q_start.append(self.joint_coord_count)
            joint_qd_start = copy.copy(self.joint_qd_start)
            joint_qd_start.append(self.joint_dof_count)
            articulation_start = copy.copy(self.articulation_start)
            articulation_start.append(self.joint_count)

            m.joint_q_start = wp.array(joint_q_start, dtype=wp.int32)
            m.joint_qd_start = wp.array(joint_qd_start, dtype=wp.int32)
            m.articulation_start = wp.array(articulation_start, dtype=wp.int32)

            # counts
            m.particle_count = len(self.particle_q)
            m.body_count = len(self.body_q)
            m.shape_count = len(self.shape_geo_type)
            m.tri_count = len(self.tri_poses)
            m.tet_count = len(self.tet_poses)
            m.edge_count = len(self.edge_rest_angle)
            m.spring_count = len(self.spring_rest_length)
            m.muscle_count = len(self.muscle_start)
            m.articulation_count = len(self.articulation_start)
            m.attachment_count = len(self.attachment_indices)

            # contacts
            if m.particle_count:
                m.allocate_soft_contacts(self.soft_contact_max, requires_grad=requires_grad)

                self.edge_contact_max *= self.spring_count
                self.point_tri_contact_max *= m.particle_count
                m.allocate_particle_tri_self_intersection(
                    self.point_tri_contact_max, requires_grad=requires_grad)
                m.allocate_edge_self_intersection(
                    self.edge_contact_max, requires_grad=requires_grad)
                m.allocate_global_self_intersection(
                    m.edge_count, requires_grad=requires_grad
                )


            m.find_shape_contact_pairs()
            if self.num_rigid_contacts_per_env is None:
                contact_count = m.count_contact_points()
            else:
                contact_count = self.num_rigid_contacts_per_env * self.num_envs
            if wp.config.verbose:
                print(f"Allocating {contact_count} rigid contacts.")
            m.allocate_rigid_contacts(contact_count, requires_grad=requires_grad)
            m.rigid_contact_margin = self.rigid_contact_margin
            m.rigid_contact_torsional_friction = self.rigid_contact_torsional_friction
            m.rigid_contact_rolling_friction = self.rigid_contact_rolling_friction

            m.joint_dof_count = self.joint_dof_count
            m.joint_coord_count = self.joint_coord_count

            # store refs to geometry
            m.geo_meshes = self.geo_meshes
            m.geo_sdfs = self.geo_sdfs

            # enable ground plane
            m.ground_plane = wp.array(self._ground_params["plane"], dtype=wp.float32, requires_grad=requires_grad)
            m.gravity = np.array(self.up_vector) * self.gravity

            m.enable_tri_collisions = False

            return m