# 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 `): 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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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 `): The transform of the joint in the parent body's local frame child_xform (:ref:`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