# 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. import warp as wp from .model import ( PARTICLE_FLAG_ACTIVE, ModelShapeMaterials, ModelShapeGeometry, JOINT_MODE_TARGET_POSITION, JOINT_MODE_TARGET_VELOCITY, JOINT_MODE_LIMIT, ) from .utils import velocity_at_point, vec_min, vec_max, vec_abs from .collide import triangle_closest_point_barycentric from .integrator_euler import integrate_bodies, integrate_particles from .utils import mesh_query_inside_ @wp.kernel def solve_particle_ground_contacts( particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_flags: wp.array(dtype=wp.uint32), ke: float, kd: float, kf: float, mu: float, ground: wp.array(dtype=float), dt: float, relaxation: float, delta: wp.array(dtype=wp.vec3), ): tid = wp.tid() if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0: return wi = invmass[tid] if wi == 0.0: return x = particle_x[tid] v = particle_v[tid] n = wp.vec3(ground[0], ground[1], ground[2]) c = wp.min(wp.dot(n, x) + ground[3] - particle_radius[tid], 0.0) if c > 0.0: return # normal lambda_n = c delta_n = n * lambda_n # friction vn = wp.dot(n, v) vt = v - n * vn lambda_f = wp.max(mu * lambda_n, 0.0 - wp.length(vt) * dt) delta_f = wp.normalize(vt) * lambda_f wp.atomic_add(delta, tid, (delta_f - delta_n) / wi * relaxation) @wp.kernel def apply_soft_restitution_ground( particle_x_new: wp.array(dtype=wp.vec3), particle_v_new: wp.array(dtype=wp.vec3), particle_x_old: wp.array(dtype=wp.vec3), particle_v_old: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_flags: wp.array(dtype=wp.uint32), restitution: float, ground: wp.array(dtype=float), dt: float, relaxation: float, particle_v_out: wp.array(dtype=wp.vec3), ): tid = wp.tid() if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0: return wi = invmass[tid] if wi == 0.0: return # x_new = particle_x_new[tid] v_new = particle_v_new[tid] x_old = particle_x_old[tid] v_old = particle_v_old[tid] n = wp.vec3(ground[0], ground[1], ground[2]) c = wp.dot(n, x_old) + ground[3] - particle_radius[tid] if c > 0.0: return rel_vel_old = wp.dot(n, v_old) rel_vel_new = wp.dot(n, v_new) dv = n * wp.max(-rel_vel_new + wp.max(-restitution * rel_vel_old, 0.0), 0.0) wp.atomic_add(particle_v_out, tid, dv / wi * relaxation) @wp.kernel def solve_particle_shape_contacts( particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), particle_invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_flags: wp.array(dtype=wp.uint32), body_q: wp.array(dtype=wp.transform), body_qd: wp.array(dtype=wp.spatial_vector), body_com: wp.array(dtype=wp.vec3), body_m_inv: wp.array(dtype=float), body_I_inv: wp.array(dtype=wp.mat33), shape_body: wp.array(dtype=int), shape_materials: ModelShapeMaterials, particle_mu: float, particle_ka: float, contact_count: wp.array(dtype=int), contact_particle: wp.array(dtype=int), contact_shape: wp.array(dtype=int), contact_body_pos: wp.array(dtype=wp.vec3), contact_body_vel: wp.array(dtype=wp.vec3), contact_normal: wp.array(dtype=wp.vec3), contact_max: int, dt: float, relaxation: float, # outputs delta: wp.array(dtype=wp.vec3), body_delta: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() count = min(contact_max, contact_count[0]) if tid >= count: return shape_index = contact_shape[tid] body_index = shape_body[shape_index] particle_index = contact_particle[tid] if (particle_flags[particle_index] & PARTICLE_FLAG_ACTIVE) == 0: return px = particle_x[particle_index] pv = particle_v[particle_index] X_wb = wp.transform_identity() X_com = wp.vec3() if body_index >= 0: X_wb = body_q[body_index] X_com = body_com[body_index] # body position in world space bx = wp.transform_point(X_wb, contact_body_pos[tid]) r = bx - wp.transform_point(X_wb, X_com) n = contact_normal[tid] c = wp.dot(n, px - bx) - particle_radius[particle_index] if c > particle_ka: return # take average material properties of shape and particle parameters mu = 0.5 * (particle_mu + shape_materials.mu[shape_index]) # body velocity body_v_s = wp.spatial_vector() if body_index >= 0: body_v_s = body_qd[body_index] body_w = wp.spatial_top(body_v_s) body_v = wp.spatial_bottom(body_v_s) # compute the body velocity at the particle position bv = body_v + wp.cross(body_w, r) + wp.transform_vector(X_wb, contact_body_vel[tid]) # relative velocity v = pv - bv # normal lambda_n = c delta_n = n * lambda_n # friction vn = wp.dot(n, v) vt = v - n * vn # compute inverse masses w1 = particle_invmass[particle_index] w2 = 0.0 if body_index >= 0: angular = wp.cross(r, n) q = wp.transform_get_rotation(X_wb) rot_angular = wp.quat_rotate_inv(q, angular) I_inv = body_I_inv[body_index] w2 = body_m_inv[body_index] + wp.dot(rot_angular, I_inv * rot_angular) denom = w1 + w2 if denom == 0.0: return lambda_f = wp.max(mu * lambda_n, -wp.length(vt) * dt) delta_f = wp.normalize(vt) * lambda_f delta_total = (delta_f - delta_n) / denom * relaxation wp.atomic_add(delta, particle_index, delta_total) if body_index >= 0: delta_t = wp.cross(r, delta_total) wp.atomic_sub(body_delta, body_index, wp.spatial_vector(delta_t, delta_total)) @wp.kernel def solve_particle_particle_contacts( grid: wp.uint64, particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), particle_invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_flags: wp.array(dtype=wp.uint32), k_mu: float, k_cohesion: float, max_radius: float, dt: float, relaxation: float, # outputs deltas: wp.array(dtype=wp.vec3), ): tid = wp.tid() # order threads by cell i = wp.hash_grid_point_id(grid, tid) if i == -1: # hash grid has not been built yet return if (particle_flags[i] & PARTICLE_FLAG_ACTIVE) == 0: return x = particle_x[i] v = particle_v[i] radius = particle_radius[i] w1 = particle_invmass[i] # particle contact query = wp.hash_grid_query(grid, x, radius + max_radius + k_cohesion) index = int(0) delta = wp.vec3(0.0) while wp.hash_grid_query_next(query, index): if (particle_flags[index] & PARTICLE_FLAG_ACTIVE) != 0 and index != i: # compute distance to point n = x - particle_x[index] d = wp.length(n) err = d - radius - particle_radius[index] # compute inverse masses w2 = particle_invmass[index] denom = w1 + w2 if err <= k_cohesion and denom > 0.0: n = n / d vrel = v - particle_v[index] # normal lambda_n = err delta_n = n * lambda_n # friction vn = wp.dot(n, vrel) vt = v - n * vn lambda_f = wp.max(k_mu * lambda_n, -wp.length(vt) * dt) delta_f = wp.normalize(vt) * lambda_f delta += (delta_f - delta_n) / denom wp.atomic_add(deltas, i, delta * relaxation) @wp.kernel def solve_particle_triangle_self_contacts( particle_x: wp.array(dtype=wp.vec3), particle_invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_shape: wp.uint64, point_tri_contact_count: wp.array(dtype=int), point_tri_contact_pairs: wp.array(dtype=wp.vec2i), point_tri_contact_filter: wp.array(dtype=bool), point_tri_contact_sidedness: wp.array(dtype=bool), # outputs delta: wp.array(dtype=wp.vec3), ): """For a given pair of particle-triangle, checks if a collision happened and resolve""" tid = wp.tid() # current point-triangle pair if tid >= point_tri_contact_count[0] or point_tri_contact_filter[tid]: # to be skipped return # Info pair = point_tri_contact_pairs[tid] particle_index, face_index = pair[0], pair[1] p = particle_x[particle_index] radius = particle_radius[particle_index] # NOTE == Thickness wpoint = particle_invmass[particle_index] mesh = wp.mesh_get(particle_shape) # Parameters # TODO from outside EPSILON = 0.00001 relaxation = 1. p_delta = wp.vec3(0.0) p0_idx = mesh.indices[face_index * 3 + 0] p1_idx = mesh.indices[face_index * 3 + 1] p2_idx = mesh.indices[face_index * 3 + 2] p0 = particle_x[p0_idx] p1 = particle_x[p1_idx] p2 = particle_x[p2_idx] w0 = particle_invmass[p0_idx] w1 = particle_invmass[p1_idx] w2 = particle_invmass[p2_idx] # Closest point bary_closest = triangle_closest_point_barycentric(p0, p1, p2, p) baryAlpha, baryBeta, baryGamma = bary_closest[0], bary_closest[1], bary_closest[2] pC = baryAlpha * p0 + baryBeta * p1 + baryGamma * p2 ######## #Set up constraints: # https://github.com/vasumahesh1/azura/blob/master/Source/Samples/3_ClothSim/Shaders/SolvingPass_Cloth_GenerateSelfCollisions.cs.slang normal = wp.cross(p1 - p0, p2 - p0) n_norm = wp.length(normal) n_hat = wp.normalize(normal) # Apply contact sidedness if required if point_tri_contact_sidedness[tid]: n_hat = -1.0 * n_hat # NOTE: Using pC instead of p0 as in the paper (suspect typo) # Because it makes the most sence to compare with closest point c_def = wp.dot(n_hat, p - pC) - 2.0 * radius #Apply constaints: # https://github.com/vasumahesh1/azura/blob/master/Source/Samples/3_ClothSim/Shaders/SolvingPass_Cloth_ApplyConstraints.cs.slang # https://media.contentapi.ea.com/content/dam/eacom/frostbite/files/gdc2018-chrislewin-clothselfcollisionwithpredictivecontacts.pdf if c_def < 0.0: a = n_hat[0] b = n_hat[1] c = n_hat[2] a2 = a * a b2 = b * b c2 = c * c N_n_hat = wp.vec3( a - (a2 * a) - (a * b2) - (a * c2), (-a2 * b) + b - (b2 * b) - (b * c2), (-a2 * c) - (b2 * c) + c - (c2 * c) ) N_n_hat = N_n_hat / n_norm grad_c_p = n_hat grad_c_p0 = wp.cross((p1 - p2), N_n_hat) - n_hat grad_c_p1 = wp.cross((p2 - p0), N_n_hat) grad_c_p2 = wp.cross((p1 - p0), N_n_hat) denom = w0 * wp.dot(grad_c_p0, grad_c_p0) + w1 * wp.dot(grad_c_p1, grad_c_p1) + w2 * wp.dot(grad_c_p2, grad_c_p2) + wpoint if abs(denom) > EPSILON: # NOTE: No need for alpha/damping for this constraint # as it should be as hard as possible dlambda = (-1.0 * c_def) / denom # Constraint value p_delta = wpoint * dlambda * grad_c_p # TODO Account for friction wp.atomic_add(delta, particle_index, relaxation * p_delta) @wp.kernel def solve_edge_edge_self_contact( particle_x: wp.array(dtype=wp.vec3), particle_invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), edge_contact_count: wp.array(dtype=int), edge_contact_pairs: wp.array(dtype=wp.vec4i), edge_contact_filter: wp.array(dtype=bool), edge_contact_normal: wp.array(dtype=wp.vec3), # outputs delta: wp.array(dtype=wp.vec3), ): """For a given pair of edges, checks if a collision happened and resolve""" tid = wp.tid() # current edge pair if tid >= edge_contact_count[0] or edge_contact_filter[tid]: # to be skipped return # Info pair = edge_contact_pairs[tid] p_a_idx, p_b_idx, e0_idx, e1_idx = pair[0], pair[1], pair[2], pair[3] p_a = particle_x[p_a_idx] w_a = particle_invmass[p_a_idx] p_b = particle_x[p_b_idx] w_b = particle_invmass[p_b_idx] e_0 = particle_x[e0_idx] w_0 = particle_invmass[e0_idx] e_1 = particle_x[e1_idx] w_1 = particle_invmass[e1_idx] # -- Parameters -- # TODO Outside radius = max(particle_radius[p_a_idx], particle_radius[p_b_idx]) relaxation = 1. # Closest points # We are going with de-facto closest point and using only general direction from the saved info out = wp.closest_point_edge_edge(p_a, p_b, e_0, e_1, 0.001) t_alpha, t_beta, _ = out[0], out[1], out[2] p_alpha = p_a * (1. - t_alpha) + p_b * t_alpha p_beta = e_0 * (1. - t_beta) + e_1 * t_beta lDir = wp.normalize(p_beta - p_alpha) # Compare projection on connecting vector (as it was at the beginning of the frame) orig_lDir = edge_contact_normal[tid] if wp.dot(lDir, orig_lDir) < 0: lDir = -1. * lDir # Calculate constraint constraint = wp.dot(p_beta - p_alpha, lDir) - 2. * radius if constraint < 0.0: # Resolve grad_pa = -(1. - t_alpha) * lDir grad_pb = - t_alpha * lDir grad_e0 = -(1. - t_beta) * lDir grad_e1 = - t_beta * lDir denom = (w_a * wp.dot(grad_pa, grad_pa) + w_b * wp.dot(grad_pb, grad_pb) + w_0 * wp.dot(grad_e0, grad_e0) + w_1 * wp.dot(grad_e1, grad_e1)) if abs(denom) > 1e-5: # NOTE: No need for alpha/damping for this constraint # as it should be as hard as possible dlambda = (-1.0 * constraint) / denom # Constraint value delta_a = w_a * dlambda * grad_pa delta_b = w_b * dlambda * grad_pb # TODO Account for friction wp.atomic_add(delta, p_a_idx, relaxation * delta_a) wp.atomic_add(delta, p_b_idx, relaxation * delta_b) @wp.kernel def replace_mesh_points( shape: wp.uint64, vertices: wp.array(dtype=wp.vec3) ): tid = wp.tid() mesh = wp.mesh_get(shape) mesh.points[tid] = vertices[tid] @wp.kernel #distance constraints / stretching def solve_springs( x: wp.array(dtype=wp.vec3), v: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), spring_indices: wp.array(dtype=int), spring_rest_lengths: wp.array(dtype=float), spring_stiffness: wp.array(dtype=float), spring_damping: wp.array(dtype=float), dt: float, lambdas: wp.array(dtype=float), delta: wp.array(dtype=wp.vec3), ): tid = wp.tid() i = spring_indices[tid * 2 + 0] j = spring_indices[tid * 2 + 1] ke = spring_stiffness[tid] kd = spring_damping[tid] rest = spring_rest_lengths[tid] xi = x[i] xj = x[j] vi = v[i] vj = v[j] xij = xi - xj vij = vi - vj l = wp.length(xij) if l == 0.0: return n = xij / l c = l - rest grad_c_xi = n grad_c_xj = -1.0 * n wi = invmass[i] wj = invmass[j] denom = wi + wj # Note strict inequality for damping -- 0 damping is ok if denom <= 0.0 or ke <= 0.0 or kd < 0.0: return alpha= 1.0 / (ke * dt * dt) gamma = kd / (ke * dt) grad_c_dot_v = dt * wp.dot(grad_c_xi, vij) # Note: dt because from the paper we want x_i - x^n, not v... dlambda = -1.0 * (c + alpha* lambdas[tid] + gamma * grad_c_dot_v) / ((1.0 + gamma) * denom + alpha) dxi = wi * dlambda * grad_c_xi dxj = wj * dlambda * grad_c_xj lambdas[tid] = lambdas[tid] + dlambda wp.atomic_add(delta, i, dxi) wp.atomic_add(delta, j, dxj) @wp.kernel def bending_constraint( x: wp.array(dtype=wp.vec3), v: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), indices: wp.array2d(dtype=int), rest: wp.array(dtype=float), bending_properties: wp.array2d(dtype=float), dt: float, lambdas: wp.array(dtype=float), delta: wp.array(dtype=wp.vec3), ): tid = wp.tid() eps = 1.0e-6 ke = bending_properties[tid, 0] kd = bending_properties[tid, 1] i = indices[tid, 0] j = indices[tid, 1] k = indices[tid, 2] l = indices[tid, 3] if i == -1 or j == -1 or k == -1 or l == -1: return rest_angle = rest[tid] x1 = x[i] x2 = x[j] x3 = x[k] x4 = x[l] v1 = v[i] v2 = v[j] v3 = v[k] v4 = v[l] w1 = invmass[i] w2 = invmass[j] w3 = invmass[k] w4 = invmass[l] n1 = wp.cross(x3 - x1, x4 - x1) # normal to face 1 n2 = wp.cross(x4 - x2, x3 - x2) # normal to face 2 n1_length = wp.length(n1) n2_length = wp.length(n2) if n1_length < eps or n2_length < eps: return n1 /= n1_length n2 /= n2_length cos_theta = wp.dot(n1, n2) e = x4 - x3 e_hat = wp.normalize(e) e_length = wp.length(e) derivative_flip = wp.sign(wp.dot(wp.cross(n1, n2), e)) derivative_flip *= -1.0 angle = wp.acos(cos_theta) grad_x1 = n1 * e_length * derivative_flip grad_x2 = n2 * e_length * derivative_flip grad_x3 = (n1 * wp.dot(x1 - x4, e_hat) + n2 * wp.dot(x2 - x4, e_hat)) * derivative_flip grad_x4 = (n1 * wp.dot(x3 - x1, e_hat) + n2 * wp.dot(x3 - x2, e_hat)) * derivative_flip c = angle - rest_angle denominator = (w1 * wp.length_sq(grad_x1) + w2 * wp.length_sq(grad_x2) + w3 * wp.length_sq(grad_x3) + w4 * wp.length_sq(grad_x4)) # Note strict inequality for damping -- 0 damping is ok if denominator <= 0.0 or ke <= 0.0 or kd < 0.0: return alpha = 1.0 / (ke * dt * dt) gamma = kd / (ke * dt) # NOTE: eq. (1) in smallsteps grad_dot_v = dt * (wp.dot(grad_x1, v1) + wp.dot(grad_x2, v2) + wp.dot(grad_x3, v3) + wp.dot(grad_x4, v4)) dlambda = -1.0 * (c + alpha * lambdas[tid] + gamma * grad_dot_v) / ((1.0 + gamma) * denominator + alpha) delta0 = w1 * dlambda * grad_x1 delta1 = w2 * dlambda * grad_x2 delta2 = w3 * dlambda * grad_x3 delta3 = w4 * dlambda * grad_x4 lambdas[tid] = lambdas[tid] + dlambda wp.atomic_add(delta, i, delta0) wp.atomic_add(delta, j, delta1) wp.atomic_add(delta, k, delta2) wp.atomic_add(delta, l, delta3) @wp.kernel def solve_tetrahedra( x: wp.array(dtype=wp.vec3), v: wp.array(dtype=wp.vec3), inv_mass: wp.array(dtype=float), indices: wp.array(dtype=int, ndim=2), pose: wp.array(dtype=wp.mat33), activation: wp.array(dtype=float), materials: wp.array(dtype=float, ndim=2), dt: float, relaxation: float, delta: wp.array(dtype=wp.vec3), ): tid = wp.tid() i = indices[tid, 0] j = indices[tid, 1] k = indices[tid, 2] l = indices[tid, 3] act = activation[tid] k_mu = materials[tid, 0] k_lambda = materials[tid, 1] k_damp = materials[tid, 2] x0 = x[i] x1 = x[j] x2 = x[k] x3 = x[l] v0 = v[i] v1 = v[j] v2 = v[k] v3 = v[l] w0 = inv_mass[i] w1 = inv_mass[j] w2 = inv_mass[k] w3 = inv_mass[l] x10 = x1 - x0 x20 = x2 - x0 x30 = x3 - x0 v10 = v1 - v0 v20 = v2 - v0 v30 = v3 - v0 Ds = wp.mat33(x10, x20, x30) Dm = pose[tid] inv_rest_volume = wp.determinant(Dm) * 6.0 rest_volume = 1.0 / inv_rest_volume # F = Xs*Xm^-1 F = Ds * Dm f1 = wp.vec3(F[0, 0], F[1, 0], F[2, 0]) f2 = wp.vec3(F[0, 1], F[1, 1], F[2, 1]) f3 = wp.vec3(F[0, 2], F[1, 2], F[2, 2]) # C_sqrt # tr = wp.dot(f1, f1) + wp.dot(f2, f2) + wp.dot(f3, f3) # r_s = wp.sqrt(abs(tr - 3.0)) # C = r_s # if (r_s == 0.0): # return # if (tr < 3.0): # r_s = 0.0 - r_s # dCdx = F*wp.transpose(Dm)*(1.0/r_s) # alpha = 1.0 + k_mu / k_lambda # C_Neo r_s = wp.sqrt(wp.dot(f1, f1) + wp.dot(f2, f2) + wp.dot(f3, f3)) if r_s == 0.0: return # tr = wp.dot(f1, f1) + wp.dot(f2, f2) + wp.dot(f3, f3) # if (tr < 3.0): # r_s = -r_s r_s_inv = 1.0 / r_s C = r_s dCdx = F * wp.transpose(Dm) * r_s_inv alpha = 1.0 + k_mu / k_lambda # C_Spherical # r_s = wp.sqrt(wp.dot(f1, f1) + wp.dot(f2, f2) + wp.dot(f3, f3)) # r_s_inv = 1.0/r_s # C = r_s - wp.sqrt(3.0) # dCdx = F*wp.transpose(Dm)*r_s_inv # alpha = 1.0 # C_D # r_s = wp.sqrt(wp.dot(f1, f1) + wp.dot(f2, f2) + wp.dot(f3, f3)) # C = r_s*r_s - 3.0 # dCdx = F*wp.transpose(Dm)*2.0 # alpha = 1.0 grad1 = wp.vec3(dCdx[0, 0], dCdx[1, 0], dCdx[2, 0]) grad2 = wp.vec3(dCdx[0, 1], dCdx[1, 1], dCdx[2, 1]) grad3 = wp.vec3(dCdx[0, 2], dCdx[1, 2], dCdx[2, 2]) grad0 = (grad1 + grad2 + grad3) * (0.0 - 1.0) denom = ( wp.dot(grad0, grad0) * w0 + wp.dot(grad1, grad1) * w1 + wp.dot(grad2, grad2) * w2 + wp.dot(grad3, grad3) * w3 ) multiplier = C / (denom + 1.0 / (k_mu * dt * dt * rest_volume)) delta0 = grad0 * multiplier delta1 = grad1 * multiplier delta2 = grad2 * multiplier delta3 = grad3 * multiplier # hydrostatic part J = wp.determinant(F) C_vol = J - alpha # dCdx = wp.mat33(wp.cross(f2, f3), wp.cross(f3, f1), wp.cross(f1, f2))*wp.transpose(Dm) # grad1 = wp.vec3(dCdx[0,0], dCdx[1,0], dCdx[2,0]) # grad2 = wp.vec3(dCdx[0,1], dCdx[1,1], dCdx[2,1]) # grad3 = wp.vec3(dCdx[0,2], dCdx[1,2], dCdx[2,2]) # grad0 = (grad1 + grad2 + grad3)*(0.0 - 1.0) s = inv_rest_volume / 6.0 grad1 = wp.cross(x20, x30) * s grad2 = wp.cross(x30, x10) * s grad3 = wp.cross(x10, x20) * s grad0 = -(grad1 + grad2 + grad3) denom = ( wp.dot(grad0, grad0) * w0 + wp.dot(grad1, grad1) * w1 + wp.dot(grad2, grad2) * w2 + wp.dot(grad3, grad3) * w3 ) multiplier = C_vol / (denom + 1.0 / (k_lambda * dt * dt * rest_volume)) delta0 += grad0 * multiplier delta1 += grad1 * multiplier delta2 += grad2 * multiplier delta3 += grad3 * multiplier # apply forces wp.atomic_sub(delta, i, delta0 * w0 * relaxation) wp.atomic_sub(delta, j, delta1 * w1 * relaxation) wp.atomic_sub(delta, k, delta2 * w2 * relaxation) wp.atomic_sub(delta, l, delta3 * w3 * relaxation) @wp.kernel def apply_particle_deltas( x_orig: wp.array(dtype=wp.vec3), x_pred: wp.array(dtype=wp.vec3), particle_flags: wp.array(dtype=wp.uint32), delta: wp.array(dtype=wp.vec3), dt: float, v_max: float, x_out: wp.array(dtype=wp.vec3), v_out: wp.array(dtype=wp.vec3), ): tid = wp.tid() if (particle_flags[tid] & PARTICLE_FLAG_ACTIVE) == 0: return x0 = x_orig[tid] xp = x_pred[tid] # constraint deltas d = delta[tid] x_new = xp + d v_new = (x_new - x0) / dt # enforce velocity limit to prevent instability v_new_mag = wp.length(v_new) if v_new_mag > v_max: v_new *= v_max / v_new_mag x_new = x0 + v_new * dt # NOTE: Update x in accordance with damped velocity for more stability x_out[tid] = x_new v_out[tid] = v_new @wp.kernel def apply_body_deltas( q_in: wp.array(dtype=wp.transform), qd_in: wp.array(dtype=wp.spatial_vector), body_com: wp.array(dtype=wp.vec3), body_I: wp.array(dtype=wp.mat33), body_inv_m: wp.array(dtype=float), body_inv_I: wp.array(dtype=wp.mat33), deltas: wp.array(dtype=wp.spatial_vector), constraint_inv_weights: wp.array(dtype=float), dt: float, # outputs q_out: wp.array(dtype=wp.transform), qd_out: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() inv_m = body_inv_m[tid] if inv_m == 0.0: q_out[tid] = q_in[tid] qd_out[tid] = qd_in[tid] return inv_I = body_inv_I[tid] tf = q_in[tid] delta = deltas[tid] p0 = wp.transform_get_translation(tf) q0 = wp.transform_get_rotation(tf) weight = 1.0 if constraint_inv_weights: if constraint_inv_weights[tid] > 0.0: weight = 1.0 / constraint_inv_weights[tid] dp = wp.spatial_bottom(delta) * (inv_m * weight) dq = wp.spatial_top(delta) * weight dq = wp.quat_rotate(q0, inv_I * wp.quat_rotate_inv(q0, dq)) # update orientation q1 = q0 + 0.5 * wp.quat(dq * dt * dt, 0.0) * q0 q1 = wp.normalize(q1) # update position com = body_com[tid] x_com = p0 + wp.quat_rotate(q0, com) p1 = x_com + dp * dt * dt p1 -= wp.quat_rotate(q1, com) q_out[tid] = wp.transform(p1, q1) v0 = wp.spatial_bottom(qd_in[tid]) w0 = wp.spatial_top(qd_in[tid]) # update linear and angular velocity v1 = v0 + dp * dt # angular part (compute in body frame) wb = wp.quat_rotate_inv(q0, w0 + dq * dt) tb = -wp.cross(wb, body_I[tid] * wb) # coriolis forces w1 = wp.quat_rotate(q0, wb + inv_I * tb * dt) qd_out[tid] = wp.spatial_vector(w1, v1) @wp.kernel def apply_body_delta_velocities( qd_in: wp.array(dtype=wp.spatial_vector), deltas: wp.array(dtype=wp.spatial_vector), qd_out: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() qd_out[tid] = qd_in[tid] + deltas[tid] @wp.kernel def apply_joint_torques( body_q: wp.array(dtype=wp.transform), body_com: wp.array(dtype=wp.vec3), joint_q_start: wp.array(dtype=int), joint_qd_start: wp.array(dtype=int), joint_type: wp.array(dtype=int), joint_parent: wp.array(dtype=int), joint_child: wp.array(dtype=int), joint_X_p: wp.array(dtype=wp.transform), joint_X_c: wp.array(dtype=wp.transform), joint_axis_start: wp.array(dtype=int), joint_axis_dim: wp.array(dtype=int, ndim=2), joint_axis: wp.array(dtype=wp.vec3), joint_act: wp.array(dtype=float), body_f: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() type = joint_type[tid] if type == wp.sim.JOINT_FIXED: return if type == wp.sim.JOINT_FREE: return if type == wp.sim.JOINT_DISTANCE: return if type == wp.sim.JOINT_BALL: return # rigid body indices of the child and parent id_c = joint_child[tid] id_p = joint_parent[tid] X_pj = joint_X_p[tid] X_cj = joint_X_c[tid] X_wp = X_pj pose_p = X_pj com_p = wp.vec3(0.0) # parent transform and moment arm if id_p >= 0: pose_p = body_q[id_p] X_wp = pose_p * X_wp com_p = body_com[id_p] r_p = wp.transform_get_translation(X_wp) - wp.transform_point(pose_p, com_p) # child transform and moment arm pose_c = body_q[id_c] X_wc = pose_c com_c = body_com[id_c] r_c = wp.transform_get_translation(X_wc) - wp.transform_point(pose_c, com_c) # local joint rotations q_p = wp.transform_get_rotation(X_wp) q_c = wp.transform_get_rotation(X_wc) # joint properties (for 1D joints) q_start = joint_q_start[tid] qd_start = joint_qd_start[tid] axis_start = joint_axis_start[tid] lin_axis_count = joint_axis_dim[tid, 0] ang_axis_count = joint_axis_dim[tid, 1] # total force/torque on the parent t_total = wp.vec3() f_total = wp.vec3() # handle angular constraints if type == wp.sim.JOINT_REVOLUTE: axis = joint_axis[axis_start] act = joint_act[qd_start] a_p = wp.transform_vector(X_wp, axis) t_total += act * a_p elif type == wp.sim.JOINT_PRISMATIC: axis = joint_axis[axis_start] act = joint_act[qd_start] a_p = wp.transform_vector(X_wp, axis) f_total += act * a_p elif type == wp.sim.JOINT_COMPOUND: # q_off = wp.transform_get_rotation(X_cj) # q_pc = wp.quat_inverse(q_off)*wp.quat_inverse(q_p)*q_c*q_off # # decompose to a compound rotation each axis # angles = quat_decompose(q_pc) # # reconstruct rotation axes # axis_0 = wp.vec3(1.0, 0.0, 0.0) # q_0 = wp.quat_from_axis_angle(axis_0, angles[0]) # axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0)) # q_1 = wp.quat_from_axis_angle(axis_1, angles[1]) # axis_2 = wp.quat_rotate(q_1*q_0, wp.vec3(0.0, 0.0, 1.0)) # q_w = q_p*q_off # t_total += joint_act[qd_start+0] * wp.quat_rotate(q_w, axis_0) # t_total += joint_act[qd_start+1] * wp.quat_rotate(q_w, axis_1) # t_total += joint_act[qd_start+2] * wp.quat_rotate(q_w, axis_2) axis_0 = joint_axis[axis_start + 0] axis_1 = joint_axis[axis_start + 1] axis_2 = joint_axis[axis_start + 2] t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0) t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1) t_total += joint_act[qd_start + 2] * wp.transform_vector(X_wp, axis_2) elif type == wp.sim.JOINT_UNIVERSAL: # q_off = wp.transform_get_rotation(X_cj) # q_pc = wp.quat_inverse(q_off)*wp.quat_inverse(q_p)*q_c*q_off # # decompose to a compound rotation each axis # angles = quat_decompose(q_pc) # # reconstruct rotation axes # axis_0 = wp.vec3(1.0, 0.0, 0.0) # q_0 = wp.quat_from_axis_angle(axis_0, angles[0]) # axis_1 = wp.quat_rotate(q_0, wp.vec3(0.0, 1.0, 0.0)) # q_1 = wp.quat_from_axis_angle(axis_1, angles[1]) # axis_2 = wp.quat_rotate(q_1*q_0, wp.vec3(0.0, 0.0, 1.0)) # q_w = q_p*q_off # free axes # t_total += joint_act[qd_start+0] * wp.quat_rotate(q_w, axis_0) # t_total += joint_act[qd_start+1] * wp.quat_rotate(q_w, axis_1) axis_0 = joint_axis[axis_start + 0] axis_1 = joint_axis[axis_start + 1] t_total += joint_act[qd_start + 0] * wp.transform_vector(X_wp, axis_0) t_total += joint_act[qd_start + 1] * wp.transform_vector(X_wp, axis_1) elif type == wp.sim.JOINT_D6: # unroll for loop to ensure joint actions remain differentiable # (since differentiating through a dynamic for loop that updates a local variable is not supported) if lin_axis_count > 0: axis = joint_axis[axis_start + 0] act = joint_act[qd_start + 0] a_p = wp.transform_vector(X_wp, axis) f_total += act * a_p if lin_axis_count > 1: axis = joint_axis[axis_start + 1] act = joint_act[qd_start + 1] a_p = wp.transform_vector(X_wp, axis) f_total += act * a_p if lin_axis_count > 2: axis = joint_axis[axis_start + 2] act = joint_act[qd_start + 2] a_p = wp.transform_vector(X_wp, axis) f_total += act * a_p if ang_axis_count > 0: axis = joint_axis[axis_start + lin_axis_count + 0] act = joint_act[qd_start + lin_axis_count + 0] a_p = wp.transform_vector(X_wp, axis) t_total += act * a_p if ang_axis_count > 1: axis = joint_axis[axis_start + lin_axis_count + 1] act = joint_act[qd_start + lin_axis_count + 1] a_p = wp.transform_vector(X_wp, axis) t_total += act * a_p if ang_axis_count > 2: axis = joint_axis[axis_start + lin_axis_count + 2] act = joint_act[qd_start + lin_axis_count + 2] a_p = wp.transform_vector(X_wp, axis) t_total += act * a_p else: print("joint type not handled in apply_joint_torques") # write forces if id_p >= 0: wp.atomic_add(body_f, id_p, wp.spatial_vector(t_total + wp.cross(r_p, f_total), f_total)) wp.atomic_sub(body_f, id_c, wp.spatial_vector(t_total + wp.cross(r_c, f_total), f_total)) @wp.func def update_joint_axis_mode(mode: wp.uint8, axis: wp.vec3, input_axis_mode: wp.vec3ub): # update the 3D axis mode flags given the axis vector and mode of this axis mode_x = wp.max(wp.uint8(wp.nonzero(axis[0])) * mode, input_axis_mode[0]) mode_y = wp.max(wp.uint8(wp.nonzero(axis[1])) * mode, input_axis_mode[1]) mode_z = wp.max(wp.uint8(wp.nonzero(axis[2])) * mode, input_axis_mode[2]) return wp.vec3ub(mode_x, mode_y, mode_z) @wp.func def update_joint_axis_limits(axis: wp.vec3, limit_lower: float, limit_upper: float, input_limits: wp.spatial_vector): # update the 3D linear/angular limits (spatial_vector [lower, upper]) given the axis vector and limits lo_temp = axis * limit_lower up_temp = axis * limit_upper lo = vec_min(lo_temp, up_temp) up = vec_max(lo_temp, up_temp) input_lower = wp.spatial_top(input_limits) input_upper = wp.spatial_bottom(input_limits) lower = vec_min(input_lower, lo) upper = vec_max(input_upper, up) return wp.spatial_vector(lower, upper) @wp.func def update_joint_axis_target_ke_kd( axis: wp.vec3, target: float, target_ke: float, target_kd: float, input_target_ke_kd: wp.mat33 ): # update the 3D linear/angular target, target_ke, and target_kd (mat33 [target, ke, kd]) given the axis vector and target, target_ke, target_kd axis_target = input_target_ke_kd[0] axis_ke = input_target_ke_kd[1] axis_kd = input_target_ke_kd[2] stiffness = axis * target_ke axis_target += stiffness * target # weighted target (to be normalized later by sum of target_ke) axis_ke += vec_abs(stiffness) axis_kd += vec_abs(axis * target_kd) return wp.mat33( axis_target[0], axis_target[1], axis_target[2], axis_ke[0], axis_ke[1], axis_ke[2], axis_kd[0], axis_kd[1], axis_kd[2], ) @wp.kernel def solve_body_joints( body_q: wp.array(dtype=wp.transform), body_qd: wp.array(dtype=wp.spatial_vector), body_com: wp.array(dtype=wp.vec3), body_inv_m: wp.array(dtype=float), body_inv_I: wp.array(dtype=wp.mat33), joint_type: wp.array(dtype=int), joint_enabled: wp.array(dtype=int), joint_parent: wp.array(dtype=int), joint_child: wp.array(dtype=int), joint_X_p: wp.array(dtype=wp.transform), joint_X_c: wp.array(dtype=wp.transform), joint_limit_lower: wp.array(dtype=float), joint_limit_upper: wp.array(dtype=float), joint_axis_start: wp.array(dtype=int), joint_axis_dim: wp.array(dtype=int, ndim=2), joint_axis_mode: wp.array(dtype=wp.uint8), joint_axis: wp.array(dtype=wp.vec3), joint_target: wp.array(dtype=float), joint_target_ke: wp.array(dtype=float), joint_target_kd: wp.array(dtype=float), joint_linear_compliance: wp.array(dtype=float), joint_angular_compliance: wp.array(dtype=float), angular_relaxation: float, linear_relaxation: float, dt: float, deltas: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() type = joint_type[tid] if joint_enabled[tid] == 0 or type == wp.sim.JOINT_FREE: return # rigid body indices of the child and parent id_c = joint_child[tid] id_p = joint_parent[tid] X_pj = joint_X_p[tid] X_cj = joint_X_c[tid] X_wp = X_pj m_inv_p = 0.0 I_inv_p = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) pose_p = X_pj com_p = wp.vec3(0.0) vel_p = wp.vec3(0.0) omega_p = wp.vec3(0.0) # parent transform and moment arm if id_p >= 0: pose_p = body_q[id_p] X_wp = pose_p * X_wp com_p = body_com[id_p] m_inv_p = body_inv_m[id_p] I_inv_p = body_inv_I[id_p] vel_p = wp.spatial_bottom(body_qd[id_p]) omega_p = wp.spatial_top(body_qd[id_p]) # child transform and moment arm pose_c = body_q[id_c] X_wc = pose_c * X_cj com_c = body_com[id_c] m_inv_c = body_inv_m[id_c] I_inv_c = body_inv_I[id_c] vel_c = wp.spatial_bottom(body_qd[id_c]) omega_c = wp.spatial_top(body_qd[id_c]) if m_inv_p == 0.0 and m_inv_c == 0.0: # connection between two immovable bodies return # accumulate constraint deltas lin_delta_p = wp.vec3(0.0) ang_delta_p = wp.vec3(0.0) lin_delta_c = wp.vec3(0.0) ang_delta_c = wp.vec3(0.0) rel_pose = wp.transform_inverse(X_wp) * X_wc rel_p = wp.transform_get_translation(rel_pose) # joint connection points # x_p = wp.transform_get_translation(X_wp) x_c = wp.transform_get_translation(X_wc) linear_compliance = joint_linear_compliance[tid] angular_compliance = joint_angular_compliance[tid] axis_start = joint_axis_start[tid] lin_axis_count = joint_axis_dim[tid, 0] ang_axis_count = joint_axis_dim[tid, 1] world_com_p = wp.transform_point(pose_p, com_p) world_com_c = wp.transform_point(pose_c, com_c) # handle positional constraints if type == wp.sim.JOINT_DISTANCE: r_p = wp.transform_get_translation(X_wp) - world_com_p r_c = wp.transform_get_translation(X_wc) - world_com_c lower = joint_limit_lower[axis_start] upper = joint_limit_upper[axis_start] if lower < 0.0 and upper < 0.0: # no limits return d = wp.length(rel_p) err = 0.0 if lower >= 0.0 and d < lower: err = d - lower # use a more descriptive direction vector for the constraint # in case the joint parent and child anchors are very close rel_p = err * wp.normalize(world_com_c - world_com_p) elif upper >= 0.0 and d > upper: err = d - upper if wp.abs(err) > 1e-9: # compute gradients linear_c = rel_p linear_p = -linear_c r_c = x_c - world_com_c angular_p = -wp.cross(r_p, linear_c) angular_c = wp.cross(r_c, linear_c) # constraint time derivative derr = ( wp.dot(linear_p, vel_p) + wp.dot(linear_c, vel_c) + wp.dot(angular_p, omega_p) + wp.dot(angular_c, omega_c) ) lambda_in = 0.0 compliance = linear_compliance ke = joint_target_ke[axis_start] if ke > 0.0: compliance = 1.0 / ke damping = joint_target_kd[axis_start] d_lambda = compute_positional_correction( err, derr, pose_p, pose_c, m_inv_p, m_inv_c, I_inv_p, I_inv_c, linear_p, linear_c, angular_p, angular_c, lambda_in, compliance, damping, dt, ) lin_delta_p += linear_p * (d_lambda * linear_relaxation) ang_delta_p += angular_p * (d_lambda * angular_relaxation) lin_delta_c += linear_c * (d_lambda * linear_relaxation) ang_delta_c += angular_c * (d_lambda * angular_relaxation) else: # compute joint target, stiffness, damping ke_sum = float(0.0) axis_limits = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) axis_mode = wp.vec3ub(wp.uint8(0), wp.uint8(0), wp.uint8(0)) axis_target_ke_kd = wp.mat33(0.0) # avoid a for loop here since local variables would need to be modified which is not yet differentiable if lin_axis_count > 0: axis = joint_axis[axis_start] lo_temp = axis * joint_limit_lower[axis_start] up_temp = axis * joint_limit_upper[axis_start] axis_limits = wp.spatial_vector(vec_min(lo_temp, up_temp), vec_max(lo_temp, up_temp)) mode = joint_axis_mode[axis_start] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_start] kd = joint_target_kd[axis_start] target = joint_target[axis_start] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke if lin_axis_count > 1: axis_idx = axis_start + 1 axis = joint_axis[axis_idx] lower = joint_limit_lower[axis_idx] upper = joint_limit_upper[axis_idx] axis_limits = update_joint_axis_limits(axis, lower, upper, axis_limits) mode = joint_axis_mode[axis_idx] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_idx] kd = joint_target_kd[axis_idx] target = joint_target[axis_idx] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke if lin_axis_count > 2: axis_idx = axis_start + 2 axis = joint_axis[axis_idx] lower = joint_limit_lower[axis_idx] upper = joint_limit_upper[axis_idx] axis_limits = update_joint_axis_limits(axis, lower, upper, axis_limits) mode = joint_axis_mode[axis_idx] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_idx] kd = joint_target_kd[axis_idx] target = joint_target[axis_idx] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke axis_target = axis_target_ke_kd[0] axis_stiffness = axis_target_ke_kd[1] axis_damping = axis_target_ke_kd[2] if ke_sum > 0.0: axis_target /= ke_sum axis_limits_lower = wp.spatial_top(axis_limits) axis_limits_upper = wp.spatial_bottom(axis_limits) frame_p = wp.quat_to_matrix(wp.transform_get_rotation(X_wp)) # note that x_c appearing in both is correct r_p = x_c - world_com_p r_c = x_c - wp.transform_point(pose_c, com_c) # for loop will be unrolled, so we can modify local variables for dim in range(3): e = rel_p[dim] mode = axis_mode[dim] # compute gradients linear_c = wp.vec3(frame_p[0, dim], frame_p[1, dim], frame_p[2, dim]) linear_p = -linear_c angular_p = -wp.cross(r_p, linear_c) angular_c = wp.cross(r_c, linear_c) # constraint time derivative derr = ( wp.dot(linear_p, vel_p) + wp.dot(linear_c, vel_c) + wp.dot(angular_p, omega_p) + wp.dot(angular_c, omega_c) ) err = 0.0 compliance = linear_compliance damping = 0.0 # consider joint limits irrespective of axis mode lower = axis_limits_lower[dim] upper = axis_limits_upper[dim] if e < lower: err = e - lower compliance = linear_compliance damping = 0.0 elif e > upper: err = e - upper compliance = linear_compliance damping = 0.0 else: target = axis_target[dim] if mode == JOINT_MODE_TARGET_POSITION: target = wp.clamp(target, lower, upper) if axis_stiffness[dim] > 0.0: err = e - target compliance = 1.0 / axis_stiffness[dim] damping = axis_damping[dim] elif mode == JOINT_MODE_TARGET_VELOCITY: if axis_stiffness[dim] > 0.0: err = (derr - target) * dt compliance = 1.0 / axis_stiffness[dim] damping = axis_damping[dim] if wp.abs(err) > 1e-9: lambda_in = 0.0 d_lambda = compute_positional_correction( err, derr, pose_p, pose_c, m_inv_p, m_inv_c, I_inv_p, I_inv_c, linear_p, linear_c, angular_p, angular_c, lambda_in, compliance, damping, dt, ) lin_delta_p += linear_p * (d_lambda * linear_relaxation) ang_delta_p += angular_p * (d_lambda * angular_relaxation) lin_delta_c += linear_c * (d_lambda * linear_relaxation) ang_delta_c += angular_c * (d_lambda * angular_relaxation) if ( type == wp.sim.JOINT_FIXED or type == wp.sim.JOINT_PRISMATIC or type == wp.sim.JOINT_REVOLUTE or type == wp.sim.JOINT_UNIVERSAL or type == wp.sim.JOINT_COMPOUND or type == wp.sim.JOINT_D6 ): # handle angular constraints # local joint rotations q_p = wp.transform_get_rotation(X_wp) q_c = wp.transform_get_rotation(X_wc) # make quats lie in same hemisphere if wp.dot(q_p, q_c) < 0.0: q_c *= -1.0 rel_q = wp.quat_inverse(q_p) * q_c qtwist = wp.normalize(wp.quat(rel_q[0], 0.0, 0.0, rel_q[3])) qswing = rel_q * wp.quat_inverse(qtwist) # decompose to a compound rotation each axis s = wp.sqrt(rel_q[0] * rel_q[0] + rel_q[3] * rel_q[3]) invs = 1.0 / s invscube = invs * invs * invs # handle axis-angle joints # rescale twist from quaternion space to angular err_0 = 2.0 * wp.asin(wp.clamp(qtwist[0], -1.0, 1.0)) err_1 = qswing[1] err_2 = qswing[2] # analytic gradients of swing-twist decomposition grad_0 = wp.quat(invs - rel_q[0] * rel_q[0] * invscube, 0.0, 0.0, -(rel_q[3] * rel_q[0]) * invscube) grad_1 = wp.quat( -rel_q[3] * (rel_q[3] * rel_q[2] + rel_q[0] * rel_q[1]) * invscube, rel_q[3] * invs, -rel_q[0] * invs, rel_q[0] * (rel_q[3] * rel_q[2] + rel_q[0] * rel_q[1]) * invscube, ) grad_2 = wp.quat( rel_q[3] * (rel_q[3] * rel_q[1] - rel_q[0] * rel_q[2]) * invscube, rel_q[0] * invs, rel_q[3] * invs, rel_q[0] * (rel_q[2] * rel_q[0] - rel_q[3] * rel_q[1]) * invscube, ) grad_0 *= 2.0 / wp.abs(qtwist[3]) # grad_0 *= 2.0 / wp.sqrt(1.0-qtwist[0]*qtwist[0]) # derivative of asin(x) = 1/sqrt(1-x^2) # rescale swing swing_sq = qswing[3] * qswing[3] # if swing axis magnitude close to zero vector, just treat in quaternion space angularEps = 1.0e-4 if swing_sq + angularEps < 1.0: d = wp.sqrt(1.0 - qswing[3] * qswing[3]) theta = 2.0 * wp.acos(wp.clamp(qswing[3], -1.0, 1.0)) scale = theta / d err_1 *= scale err_2 *= scale grad_1 *= scale grad_2 *= scale errs = wp.vec3(err_0, err_1, err_2) grad_x = wp.vec3(grad_0[0], grad_1[0], grad_2[0]) grad_y = wp.vec3(grad_0[1], grad_1[1], grad_2[1]) grad_z = wp.vec3(grad_0[2], grad_1[2], grad_2[2]) grad_w = wp.vec3(grad_0[3], grad_1[3], grad_2[3]) # compute joint target, stiffness, damping ke_sum = float(0.0) axis_limits = wp.spatial_vector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0) axis_mode = wp.vec3ub(wp.uint8(0), wp.uint8(0), wp.uint8(0)) axis_target_ke_kd = wp.mat33(0.0) # avoid a for loop here since local variables would need to be modified which is not yet differentiable if ang_axis_count > 0: axis_idx = axis_start + lin_axis_count axis = joint_axis[axis_idx] lo_temp = axis * joint_limit_lower[axis_idx] up_temp = axis * joint_limit_upper[axis_idx] axis_limits = wp.spatial_vector(vec_min(lo_temp, up_temp), vec_max(lo_temp, up_temp)) mode = joint_axis_mode[axis_idx] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_idx] kd = joint_target_kd[axis_idx] target = joint_target[axis_idx] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke if ang_axis_count > 1: axis_idx = axis_start + lin_axis_count + 1 axis = joint_axis[axis_idx] lower = joint_limit_lower[axis_idx] upper = joint_limit_upper[axis_idx] axis_limits = update_joint_axis_limits(axis, lower, upper, axis_limits) mode = joint_axis_mode[axis_idx] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_idx] kd = joint_target_kd[axis_idx] target = joint_target[axis_idx] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke if ang_axis_count > 2: axis_idx = axis_start + lin_axis_count + 2 axis = joint_axis[axis_idx] lower = joint_limit_lower[axis_idx] upper = joint_limit_upper[axis_idx] axis_limits = update_joint_axis_limits(axis, lower, upper, axis_limits) mode = joint_axis_mode[axis_idx] if mode != JOINT_MODE_LIMIT: # position or velocity target ke = joint_target_ke[axis_idx] kd = joint_target_kd[axis_idx] target = joint_target[axis_idx] axis_mode = update_joint_axis_mode(mode, axis, axis_mode) axis_target_ke_kd = update_joint_axis_target_ke_kd(axis, target, ke, kd, axis_target_ke_kd) ke_sum += ke axis_target = axis_target_ke_kd[0] axis_stiffness = axis_target_ke_kd[1] axis_damping = axis_target_ke_kd[2] if ke_sum > 0.0: axis_target /= ke_sum axis_limits_lower = wp.spatial_top(axis_limits) axis_limits_upper = wp.spatial_bottom(axis_limits) # if type == wp.sim.JOINT_D6: # wp.printf("axis_target: %f %f %f\t axis_stiffness: %f %f %f\t axis_damping: %f %f %f\t axis_limits_lower: %f %f %f \t axis_limits_upper: %f %f %f\n", # axis_target[0], axis_target[1], axis_target[2], # axis_stiffness[0], axis_stiffness[1], axis_stiffness[2], # axis_damping[0], axis_damping[1], axis_damping[2], # axis_limits_lower[0], axis_limits_lower[1], axis_limits_lower[2], # axis_limits_upper[0], axis_limits_upper[1], axis_limits_upper[2]) # # wp.printf("wp.sqrt(1.0-qtwist[0]*qtwist[0]) = %f\n", wp.sqrt(1.0-qtwist[0]*qtwist[0])) for dim in range(3): e = errs[dim] mode = axis_mode[dim] # analytic gradients of swing-twist decomposition grad = wp.quat(grad_x[dim], grad_y[dim], grad_z[dim], grad_w[dim]) quat_c = 0.5 * q_p * grad * wp.quat_inverse(q_c) angular_c = wp.vec3(quat_c[0], quat_c[1], quat_c[2]) angular_p = -angular_c # time derivative of the constraint derr = wp.dot(angular_p, omega_p) + wp.dot(angular_c, omega_c) err = 0.0 compliance = angular_compliance damping = 0.0 # consider joint limits irrespective of mode lower = axis_limits_lower[dim] upper = axis_limits_upper[dim] if e < lower: err = e - lower compliance = angular_compliance damping = 0.0 elif e > upper: err = e - upper compliance = angular_compliance damping = 0.0 else: target = axis_target[dim] if mode == JOINT_MODE_TARGET_POSITION: target = wp.clamp(target, lower, upper) if axis_stiffness[dim] > 0.0: err = e - target compliance = 1.0 / axis_stiffness[dim] damping = axis_damping[dim] elif mode == JOINT_MODE_TARGET_VELOCITY: if axis_stiffness[dim] > 0.0: err = (derr - target) * dt compliance = 1.0 / axis_stiffness[dim] damping = axis_damping[dim] d_lambda = ( compute_angular_correction( err, derr, pose_p, pose_c, I_inv_p, I_inv_c, angular_p, angular_c, 0.0, compliance, damping, dt ) * angular_relaxation ) # update deltas ang_delta_p += angular_p * d_lambda ang_delta_c += angular_c * d_lambda if id_p >= 0: wp.atomic_add(deltas, id_p, wp.spatial_vector(ang_delta_p, lin_delta_p)) if id_c >= 0: wp.atomic_add(deltas, id_c, wp.spatial_vector(ang_delta_c, lin_delta_c)) @wp.func def compute_contact_constraint_delta( err: float, tf_a: wp.transform, tf_b: wp.transform, m_inv_a: float, m_inv_b: float, I_inv_a: wp.mat33, I_inv_b: wp.mat33, linear_a: wp.vec3, linear_b: wp.vec3, angular_a: wp.vec3, angular_b: wp.vec3, relaxation: float, dt: float, ) -> float: denom = 0.0 denom += wp.length_sq(linear_a) * m_inv_a denom += wp.length_sq(linear_b) * m_inv_b q1 = wp.transform_get_rotation(tf_a) q2 = wp.transform_get_rotation(tf_b) # Eq. 2-3 (make sure to project into the frame of the body) rot_angular_a = wp.quat_rotate_inv(q1, angular_a) rot_angular_b = wp.quat_rotate_inv(q2, angular_b) denom += wp.dot(rot_angular_a, I_inv_a * rot_angular_a) denom += wp.dot(rot_angular_b, I_inv_b * rot_angular_b) deltaLambda = -err if denom > 0.0: deltaLambda /= dt * dt * denom return deltaLambda * relaxation @wp.func def compute_positional_correction( err: float, derr: float, tf_a: wp.transform, tf_b: wp.transform, m_inv_a: float, m_inv_b: float, I_inv_a: wp.mat33, I_inv_b: wp.mat33, linear_a: wp.vec3, linear_b: wp.vec3, angular_a: wp.vec3, angular_b: wp.vec3, lambda_in: float, compliance: float, damping: float, dt: float, ) -> float: denom = 0.0 denom += wp.length_sq(linear_a) * m_inv_a denom += wp.length_sq(linear_b) * m_inv_b q1 = wp.transform_get_rotation(tf_a) q2 = wp.transform_get_rotation(tf_b) # Eq. 2-3 (make sure to project into the frame of the body) rot_angular_a = wp.quat_rotate_inv(q1, angular_a) rot_angular_b = wp.quat_rotate_inv(q2, angular_b) denom += wp.dot(rot_angular_a, I_inv_a * rot_angular_a) denom += wp.dot(rot_angular_b, I_inv_b * rot_angular_b) alpha = compliance gamma = compliance * damping deltaLambda = -(err + alpha * lambda_in + gamma * derr) if denom + alpha > 0.0: deltaLambda /= dt * (dt + gamma) * denom + alpha return deltaLambda @wp.func def compute_angular_correction( err: float, derr: float, tf_a: wp.transform, tf_b: wp.transform, I_inv_a: wp.mat33, I_inv_b: wp.mat33, angular_a: wp.vec3, angular_b: wp.vec3, lambda_in: float, compliance: float, damping: float, dt: float, ) -> float: denom = 0.0 q1 = wp.transform_get_rotation(tf_a) q2 = wp.transform_get_rotation(tf_b) # Eq. 2-3 (make sure to project into the frame of the body) rot_angular_a = wp.quat_rotate_inv(q1, angular_a) rot_angular_b = wp.quat_rotate_inv(q2, angular_b) denom += wp.dot(rot_angular_a, I_inv_a * rot_angular_a) denom += wp.dot(rot_angular_b, I_inv_b * rot_angular_b) alpha = compliance gamma = compliance * damping deltaLambda = -(err + alpha * lambda_in + gamma * derr) if denom + alpha > 0.0: deltaLambda /= dt * (dt + gamma) * denom + alpha return deltaLambda @wp.kernel def solve_body_contact_positions( body_q: wp.array(dtype=wp.transform), body_qd: wp.array(dtype=wp.spatial_vector), body_com: wp.array(dtype=wp.vec3), body_m_inv: wp.array(dtype=float), body_I_inv: wp.array(dtype=wp.mat33), contact_count: wp.array(dtype=int), contact_body0: wp.array(dtype=int), contact_body1: wp.array(dtype=int), contact_point0: wp.array(dtype=wp.vec3), contact_point1: wp.array(dtype=wp.vec3), contact_offset0: wp.array(dtype=wp.vec3), contact_offset1: wp.array(dtype=wp.vec3), contact_normal: wp.array(dtype=wp.vec3), contact_thickness: wp.array(dtype=float), contact_shape0: wp.array(dtype=int), contact_shape1: wp.array(dtype=int), shape_materials: ModelShapeMaterials, relaxation: float, dt: float, contact_torsional_friction: float, contact_rolling_friction: float, # outputs deltas: wp.array(dtype=wp.spatial_vector), active_contact_point0: wp.array(dtype=wp.vec3), active_contact_point1: wp.array(dtype=wp.vec3), active_contact_distance: wp.array(dtype=float), contact_inv_weight: wp.array(dtype=float), ): tid = wp.tid() count = contact_count[0] if tid >= count: return body_a = contact_body0[tid] body_b = contact_body1[tid] if body_a == body_b: return if contact_shape0[tid] == contact_shape1[tid]: return # find body to world transform X_wb_a = wp.transform_identity() X_wb_b = wp.transform_identity() if body_a >= 0: X_wb_a = body_q[body_a] if body_b >= 0: X_wb_b = body_q[body_b] # compute body position in world space bx_a = wp.transform_point(X_wb_a, contact_point0[tid]) bx_b = wp.transform_point(X_wb_b, contact_point1[tid]) active_contact_point0[tid] = bx_a active_contact_point1[tid] = bx_b thickness = contact_thickness[tid] n = -contact_normal[tid] d = wp.dot(n, bx_b - bx_a) - thickness active_contact_distance[tid] = d if d >= 0.0: return m_inv_a = 0.0 m_inv_b = 0.0 I_inv_a = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) I_inv_b = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) # center of mass in body frame com_a = wp.vec3(0.0) com_b = wp.vec3(0.0) # body to world transform X_wb_a = wp.transform_identity() X_wb_b = wp.transform_identity() # angular velocities omega_a = wp.vec3(0.0) omega_b = wp.vec3(0.0) # contact offset in body frame offset_a = contact_offset0[tid] offset_b = contact_offset1[tid] if body_a >= 0: X_wb_a = body_q[body_a] com_a = body_com[body_a] m_inv_a = body_m_inv[body_a] I_inv_a = body_I_inv[body_a] omega_a = wp.spatial_top(body_qd[body_a]) if body_b >= 0: X_wb_b = body_q[body_b] com_b = body_com[body_b] m_inv_b = body_m_inv[body_b] I_inv_b = body_I_inv[body_b] omega_b = wp.spatial_top(body_qd[body_b]) # use average contact material properties mat_nonzero = 0 mu = 0.0 shape_a = contact_shape0[tid] shape_b = contact_shape1[tid] if shape_a >= 0: mat_nonzero += 1 mu += shape_materials.mu[shape_a] if shape_b >= 0: mat_nonzero += 1 mu += shape_materials.mu[shape_b] if mat_nonzero > 0: mu /= float(mat_nonzero) r_a = bx_a - wp.transform_point(X_wb_a, com_a) r_b = bx_b - wp.transform_point(X_wb_b, com_b) angular_a = -wp.cross(r_a, n) angular_b = wp.cross(r_b, n) if contact_inv_weight: if body_a >= 0: wp.atomic_add(contact_inv_weight, body_a, 1.0) if body_b >= 0: wp.atomic_add(contact_inv_weight, body_b, 1.0) lambda_n = compute_contact_constraint_delta( d, X_wb_a, X_wb_b, m_inv_a, m_inv_b, I_inv_a, I_inv_b, -n, n, angular_a, angular_b, relaxation, dt ) lin_delta_a = -n * lambda_n lin_delta_b = n * lambda_n ang_delta_a = angular_a * lambda_n ang_delta_b = angular_b * lambda_n # linear friction if mu > 0.0: # add on displacement from surface offsets, this ensures we include any rotational effects due to thickness from feature # need to use the current rotation to account for friction due to angular effects (e.g.: slipping contact) bx_a += wp.transform_vector(X_wb_a, offset_a) bx_b += wp.transform_vector(X_wb_b, offset_b) # update delta delta = bx_b - bx_a friction_delta = delta - wp.dot(n, delta) * n perp = wp.normalize(friction_delta) r_a = bx_a - wp.transform_point(X_wb_a, com_a) r_b = bx_b - wp.transform_point(X_wb_b, com_b) angular_a = -wp.cross(r_a, perp) angular_b = wp.cross(r_b, perp) err = wp.length(friction_delta) if err > 0.0: lambda_fr = compute_contact_constraint_delta( err, X_wb_a, X_wb_b, m_inv_a, m_inv_b, I_inv_a, I_inv_b, -perp, perp, angular_a, angular_b, 1.0, dt ) # limit friction based on incremental normal force, good approximation to limiting on total force lambda_fr = wp.max(lambda_fr, -lambda_n * mu) lin_delta_a -= perp * lambda_fr lin_delta_b += perp * lambda_fr ang_delta_a += angular_a * lambda_fr ang_delta_b += angular_b * lambda_fr torsional_friction = mu * contact_torsional_friction delta_omega = omega_b - omega_a if torsional_friction > 0.0: err = wp.dot(delta_omega, n) * dt if wp.abs(err) > 0.0: lin = wp.vec3(0.0) lambda_torsion = compute_contact_constraint_delta( err, X_wb_a, X_wb_b, m_inv_a, m_inv_b, I_inv_a, I_inv_b, lin, lin, -n, n, 1.0, dt ) lambda_torsion = wp.clamp(lambda_torsion, -lambda_n * torsional_friction, lambda_n * torsional_friction) ang_delta_a -= n * lambda_torsion ang_delta_b += n * lambda_torsion rolling_friction = mu * contact_rolling_friction if rolling_friction > 0.0: delta_omega -= wp.dot(n, delta_omega) * n err = wp.length(delta_omega) * dt if err > 0.0: lin = wp.vec3(0.0) roll_n = wp.normalize(delta_omega) lambda_roll = compute_contact_constraint_delta( err, X_wb_a, X_wb_b, m_inv_a, m_inv_b, I_inv_a, I_inv_b, lin, lin, -roll_n, roll_n, 1.0, dt ) lambda_roll = wp.max(lambda_roll, -lambda_n * rolling_friction) ang_delta_a -= roll_n * lambda_roll ang_delta_b += roll_n * lambda_roll if body_a >= 0: wp.atomic_add(deltas, body_a, wp.spatial_vector(ang_delta_a, lin_delta_a)) if body_b >= 0: wp.atomic_add(deltas, body_b, wp.spatial_vector(ang_delta_b, lin_delta_b)) @wp.kernel def update_body_velocities( poses: wp.array(dtype=wp.transform), poses_prev: wp.array(dtype=wp.transform), body_com: wp.array(dtype=wp.vec3), dt: float, qd_out: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() pose = poses[tid] pose_prev = poses_prev[tid] x = wp.transform_get_translation(pose) x_prev = wp.transform_get_translation(pose_prev) q = wp.transform_get_rotation(pose) q_prev = wp.transform_get_rotation(pose_prev) # Update body velocities according to Alg. 2 # XXX we consider the body COM as the origin of the body frame x_com = x + wp.quat_rotate(q, body_com[tid]) x_com_prev = x_prev + wp.quat_rotate(q_prev, body_com[tid]) # XXX consider the velocity of the COM v = (x_com - x_com_prev) / dt dq = q * wp.quat_inverse(q_prev) omega = 2.0 / dt * wp.vec3(dq[0], dq[1], dq[2]) if dq[3] < 0.0: omega = -omega qd_out[tid] = wp.spatial_vector(omega, v) @wp.kernel def apply_rigid_restitution( body_q: wp.array(dtype=wp.transform), body_qd: wp.array(dtype=wp.spatial_vector), body_q_prev: wp.array(dtype=wp.transform), body_qd_prev: wp.array(dtype=wp.spatial_vector), body_com: wp.array(dtype=wp.vec3), body_m_inv: wp.array(dtype=float), body_I_inv: wp.array(dtype=wp.mat33), contact_count: wp.array(dtype=int), contact_body0: wp.array(dtype=int), contact_body1: wp.array(dtype=int), contact_normal: wp.array(dtype=wp.vec3), contact_shape0: wp.array(dtype=int), contact_shape1: wp.array(dtype=int), shape_materials: ModelShapeMaterials, active_contact_distance: wp.array(dtype=float), active_contact_point0: wp.array(dtype=wp.vec3), active_contact_point1: wp.array(dtype=wp.vec3), contact_inv_weight: wp.array(dtype=float), gravity: wp.vec3, dt: float, # outputs deltas: wp.array(dtype=wp.spatial_vector), ): tid = wp.tid() count = contact_count[0] if tid >= count: return d = active_contact_distance[tid] if d >= 0.0: return # use average contact material properties mat_nonzero = 0 restitution = 0.0 shape_a = contact_shape0[tid] shape_b = contact_shape1[tid] if shape_a >= 0: mat_nonzero += 1 restitution += shape_materials.restitution[shape_a] if shape_b >= 0: mat_nonzero += 1 restitution += shape_materials.restitution[shape_b] if mat_nonzero > 0: restitution /= float(mat_nonzero) body_a = contact_body0[tid] body_b = contact_body1[tid] m_inv_a = 0.0 m_inv_b = 0.0 I_inv_a = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) I_inv_b = wp.mat33(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) # body to world transform X_wb_a_prev = wp.transform_identity() X_wb_b_prev = wp.transform_identity() # center of mass in body frame com_a = wp.vec3(0.0) com_b = wp.vec3(0.0) # previous velocity at contact points v_a = wp.vec3(0.0) v_b = wp.vec3(0.0) # new velocity at contact points v_a_new = wp.vec3(0.0) v_b_new = wp.vec3(0.0) # inverse mass used to compute the impulse inv_mass = 0.0 if body_a >= 0: X_wb_a_prev = body_q_prev[body_a] X_wb_a = body_q[body_a] m_inv_a = body_m_inv[body_a] I_inv_a = body_I_inv[body_a] com_a = body_com[body_a] if body_b >= 0: X_wb_b_prev = body_q_prev[body_b] X_wb_b = body_q[body_b] m_inv_b = body_m_inv[body_b] I_inv_b = body_I_inv[body_b] com_b = body_com[body_b] bx_a = active_contact_point0[tid] bx_b = active_contact_point1[tid] r_a = bx_a - wp.transform_point(X_wb_a, com_a) r_b = bx_b - wp.transform_point(X_wb_b, com_b) n = contact_normal[tid] if body_a >= 0: v_a = velocity_at_point(body_qd_prev[body_a], r_a) + gravity * dt v_a_new = velocity_at_point(body_qd[body_a], r_a) q_a = wp.transform_get_rotation(X_wb_a_prev) rxn = wp.quat_rotate_inv(q_a, wp.cross(r_a, n)) # Eq. 2 inv_mass_a = m_inv_a + wp.dot(rxn, I_inv_a * rxn) # if (contact_inv_weight): # if (contact_inv_weight[body_a] > 0.0): # inv_mass_a *= contact_inv_weight[body_a] inv_mass += inv_mass_a # inv_mass += m_inv_a + wp.dot(rxn, I_inv_a * rxn) if body_b >= 0: v_b = velocity_at_point(body_qd_prev[body_b], r_b) + gravity * dt v_b_new = velocity_at_point(body_qd[body_b], r_b) q_b = wp.transform_get_rotation(X_wb_b_prev) rxn = wp.quat_rotate_inv(q_b, wp.cross(r_b, n)) # Eq. 3 inv_mass_b = m_inv_b + wp.dot(rxn, I_inv_b * rxn) # if (contact_inv_weight): # if (contact_inv_weight[body_b] > 0.0): # inv_mass_b *= contact_inv_weight[body_b] inv_mass += inv_mass_b # inv_mass += m_inv_b + wp.dot(rxn, I_inv_b * rxn) if inv_mass == 0.0: return # Eq. 29 rel_vel_old = wp.dot(n, v_a - v_b) rel_vel_new = wp.dot(n, v_a_new - v_b_new) # Eq. 34 (Eq. 33 from the ACM paper, note the max operation) dv = n * (-rel_vel_new + wp.max(-restitution * rel_vel_old, 0.0)) # Eq. 33 p = dv / inv_mass if body_a >= 0: p_a = p if contact_inv_weight: if contact_inv_weight[body_a] > 0.0: p_a /= contact_inv_weight[body_a] q_a = wp.transform_get_rotation(X_wb_a) rxp = wp.quat_rotate_inv(q_a, wp.cross(r_a, p_a)) dq = wp.quat_rotate(q_a, I_inv_a * rxp) wp.atomic_add(deltas, body_a, wp.spatial_vector(dq, p_a * m_inv_a)) if body_b >= 0: p_b = p if contact_inv_weight: if contact_inv_weight[body_b] > 0.0: p_b /= contact_inv_weight[body_b] q_b = wp.transform_get_rotation(X_wb_b) rxp = wp.quat_rotate_inv(q_b, wp.cross(r_b, p_b)) dq = wp.quat_rotate(q_b, I_inv_b * rxp) wp.atomic_sub(deltas, body_b, wp.spatial_vector(dq, p_b * m_inv_b)) @wp.func def is_zero( value: wp.float32 ): return wp.abs(value) < 1e-5 @wp.kernel def attachment_constraint( particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), particle_indices: wp.array(dtype=int), constraint_points: wp.array(dtype=wp.vec3), constraint_norms: wp.array(dtype=wp.vec3), stiffness: wp.array(dtype=float), damping: wp.array(dtype=float), lambdas: wp.array(dtype=float), dt: float, # Output delta: wp.array(dtype=wp.vec3), ): tid = wp.tid() particle_index = particle_indices[tid] p = particle_x[particle_index] ke = stiffness[tid] kd = damping[tid] target_point = constraint_points[tid] norm = constraint_norms[tid] if is_zero(norm[0]) and is_zero(norm[1]) and is_zero(norm[2]): # Norm value is not set -- use an absolute constraint norm = wp.normalize(target_point - p) err = wp.dot(p - target_point, norm) if err < 0: # Apply constraint w = invmass[particle_index] denom = w # Note strict inequality for damping -- 0 damping is ok if denom <= 0.0 or ke <= 0.0 or kd < 0.0: return grad_p = norm v = particle_v[particle_index] alpha= 1.0 / (ke * dt * dt) gamma = kd / (ke * dt) grad_c_dot_v = dt * wp.dot(grad_p, v) # Note: dt because from the paper we want x_i - x^n, not v... dlambda = -1.0 * (err + alpha* lambdas[tid] + gamma * grad_c_dot_v) / ((1.0 + gamma) * denom + alpha) p_delta = w * dlambda * grad_p wp.atomic_add(delta, particle_index, p_delta) @wp.kernel def reference_shape_drag_constraint( particle_x: wp.array(dtype=wp.vec3), particle_v: wp.array(dtype=wp.vec3), invmass: wp.array(dtype=float), particle_radius: wp.array(dtype=float), particle_flags: wp.array(dtype=wp.uint32), particle_reference_label: wp.array(dtype=int), particle_normal: wp.array(dtype=wp.vec3), reference_shapes: wp.array(dtype=wp.uint64), reference_transform: wp.array(dtype=wp.transform), reference_scale: wp.array(dtype=wp.vec3), reference_marigin: float, reference_k: float, geo: ModelShapeGeometry, geo_transforms: wp.array(dtype=wp.transform), watertight_shape: int, cloth_reference_drag_particles: wp.array(dtype=int), # outputs deltas: wp.array(dtype=wp.vec3), ): """ Dim: num_particles inputs=[ state_in.particle_q, state_in.particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, model.particle_reference_label, model.particle_normal, model.cloth_reference_shape_ids, model.cloth_reference_transforms, model.cloth_reference_scale, model.cloth_reference_margin, model.cloth_reference_k, model.shape_geo.source, model.cloth_reference_watertight_whole_shape_index, model.cloth_reference_drag_particles ], outputs=[deltas], """ particle_index = wp.tid() ref_shape_index = particle_reference_label[particle_index] if cloth_reference_drag_particles[particle_index] == 0: return if (particle_flags[particle_index] & PARTICLE_FLAG_ACTIVE) == 0 or ref_shape_index < 0: return px = particle_x[particle_index] v = particle_v[particle_index] radius = particle_radius[particle_index] w = invmass[particle_index] # TODO Properly separate in-body and cloth-cloth drag into two separate constraints, with separate arrays of flags if cloth_reference_drag_particles[particle_index] == 1: # Self-collision resolution -- drag to the reference shape mesh = reference_shapes[ref_shape_index] X_ws = reference_transform[ref_shape_index] geo_scale = reference_scale[ref_shape_index] else: # Body collision resolution -- drag to the closest point on the whole body mesh = geo.source[watertight_shape] X_ws = geo_transforms[watertight_shape] geo_scale = geo.scale[watertight_shape] # transform particle position to shape local space X_sw = wp.transform_inverse(X_ws) x_local = wp.transform_point(X_sw, px) # geo description geo_thickness = geo.thickness[watertight_shape] # Assuming thickness is the same as in main collider # evaluate shape sdf err = 1.0e6 face_index = int(0) face_u = float(0.0) face_v = float(0.0) sign = float(0.0) point = wp.cw_div(x_local, geo_scale) if wp.mesh_query_point(mesh, point, 1000.0, sign, face_index, face_u, face_v): pass sign = mesh_query_inside_(geo.source[watertight_shape], wp.cw_div(x_local, geo_scale)) shape_p = wp.mesh_eval_position(mesh, face_index, face_u, face_v) # shape_v = wp.mesh_eval_velocity(mesh, face_index, face_u, face_v) shape_p = wp.cw_mul(shape_p, geo_scale) # shape_v = wp.cw_mul(shape_v, geo_scale) delta = wp.transform_vector(X_ws, shape_p - x_local) err = wp.length(delta) * sign # NOTE: Not doing the filtering by matching norms (particle-face) # Since it's a valid case: e.g. piece of trunk stuck in an arm # TODO Add back reference_k? # NOTE: One is the case for body-cloth collision, the other for cloth-cloth collision # NOTE: Not applying collision thickness as the gradient direction would be wrong for partickles in the thickness area if cloth_reference_drag_particles[particle_index] == 2 and err < 0. or cloth_reference_drag_particles[particle_index] == 1 and err > reference_marigin + radius: grad_p = wp.normalize(delta) denom = w * wp.dot(grad_p, grad_p) dlambda = abs(err) / denom p_delta = w * dlambda * grad_p wp.atomic_add(deltas, particle_index, p_delta) class XPBDIntegrator: """A implicit integrator using XPBD After constructing `Model` and `State` objects this time-integrator may be used to advance the simulation state forward in time. Example ------- .. code-block:: python integrator = wp.SemiImplicitIntegrator() # simulation loop for i in range(100): state = integrator.simulate(model, state_in, state_out, dt) """ def __init__( self, iterations=1, # NOTE: No need to have >1 if doing sim with substeps soft_body_relaxation=0.9, soft_contact_relaxation=0.9, joint_linear_relaxation=0.7, joint_angular_relaxation=0.4, rigid_contact_relaxation=0.8, rigid_contact_con_weighting=True, angular_damping=0.0, enable_restitution=False, ): self.iterations = iterations self.soft_body_relaxation = soft_body_relaxation self.soft_contact_relaxation = soft_contact_relaxation self.joint_linear_relaxation = joint_linear_relaxation self.joint_angular_relaxation = joint_angular_relaxation self.rigid_contact_relaxation = rigid_contact_relaxation self.rigid_contact_con_weighting = rigid_contact_con_weighting self.angular_damping = angular_damping self.enable_restitution = enable_restitution def simulate(self, model, state_in, state_out, dt, requires_grad=False): with wp.ScopedTimer("simulate", False): particle_q = None particle_qd = None if model.particle_count: if requires_grad: particle_q = wp.zeros_like(state_in.particle_q) particle_qd = wp.zeros_like(state_in.particle_qd) else: particle_q = state_out.particle_q particle_qd = state_out.particle_qd wp.launch( kernel=integrate_particles, dim=model.particle_count, inputs=[ state_in.particle_q, state_in.particle_qd, state_in.particle_f, model.particle_inv_mass, model.particle_flags, model.gravity, model.global_viscous_damping, dt ], outputs=[particle_q, particle_qd], device=model.device, ) if model.body_count: if model.joint_count: wp.launch( kernel=apply_joint_torques, dim=model.joint_count, inputs=[ state_in.body_q, model.body_com, model.joint_q_start, model.joint_qd_start, model.joint_type, model.joint_parent, model.joint_child, model.joint_X_p, model.joint_X_c, model.joint_axis_start, model.joint_axis_dim, model.joint_axis, model.joint_act, ], outputs=[state_in.body_f], device=model.device, ) wp.launch( kernel=integrate_bodies, dim=model.body_count, inputs=[ state_in.body_q, state_in.body_qd, state_in.body_f, model.body_com, model.body_mass, model.body_inertia, model.body_inv_mass, model.body_inv_inertia, model.gravity, self.angular_damping, dt, ], outputs=[state_out.body_q, state_out.body_qd], device=model.device, ) if model.spring_count: model.spring_constraint_lambdas.zero_() if model.edge_count: model.edge_constraint_lambdas.zero_() if model.attachment_constraint: model.vert_attach_constraint_lambdas.zero_() for i in range(self.iterations): if model.body_count: if requires_grad: out_body_q = wp.clone(state_out.body_q) out_body_qd = wp.clone(state_out.body_qd) state_out.body_deltas = wp.zeros_like(state_out.body_deltas) else: out_body_q = state_out.body_q out_body_qd = state_out.body_qd state_out.body_deltas.zero_() else: out_body_q = None out_body_qd = None # ---------------------------- # handle particles if model.particle_count: if requires_grad: deltas = wp.zeros_like(state_out.particle_f) else: deltas = state_out.particle_f deltas.zero_() # particle ground contact if model.ground: wp.launch( kernel=solve_particle_ground_contacts, dim=model.particle_count, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, model.soft_contact_ke, model.soft_contact_kd, model.soft_contact_kf, model.soft_contact_mu, model.ground_plane, dt, self.soft_contact_relaxation, ], outputs=[deltas], device=model.device, ) # particle - rigid body contacts (besides ground plane) if model.shape_count > 1: wp.launch( kernel=solve_particle_shape_contacts, dim=model.soft_contact_max, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, out_body_q, out_body_qd, model.body_com, model.body_inv_mass, model.body_inv_inertia, model.shape_body, model.shape_materials, model.soft_contact_mu, model.particle_adhesion, model.soft_contact_count, model.soft_contact_particle, model.soft_contact_shape, model.soft_contact_body_pos, model.soft_contact_body_vel, model.soft_contact_normal, model.soft_contact_max, dt, self.soft_contact_relaxation, ], # outputs outputs=[deltas, state_out.body_deltas], device=model.device, ) # Custom constraint for pushion points out of the body if model.cloth_reference_drag: wp.launch( kernel=reference_shape_drag_constraint, dim=model.particle_count, inputs=[ state_in.particle_q, state_in.particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, model.particle_reference_label, model.particle_normal, model.cloth_reference_shape_ids, model.cloth_reference_transforms, model.cloth_reference_scale, model.cloth_reference_margin, model.cloth_reference_k, model.shape_geo, model.shape_transform, model.cloth_reference_watertight_whole_shape_index, model.cloth_reference_drag_particles ], outputs=[deltas], device=model.device, ) # particle and mesh self-collisions if model.particle_max_radius > 0.0: if model.enable_particle_particle_collisions: wp.launch( kernel=solve_particle_particle_contacts, dim=model.particle_count, inputs=[ model.particle_grid.id, particle_q, particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, model.particle_mu, model.particle_cohesion, model.particle_max_radius, dt, self.soft_contact_relaxation, ], outputs=[deltas], device=model.device, ) if model.enable_triangle_particle_collisions: wp.launch( kernel=solve_particle_triangle_self_contacts, dim=model.point_tri_contact_max, inputs=[ particle_q, model.particle_inv_mass, model.particle_radius, model.particle_shape.id, model.point_tri_contact_count, model.point_tri_contact_pairs, model.point_tri_contact_filter, model.point_tri_contact_sidedness ], outputs=[deltas], device=model.device, ) # distance constraints and edge collision constraints if model.spring_count: if model.enable_edge_edge_collisions: wp.launch( kernel=solve_edge_edge_self_contact, dim=model.edge_contact_max, inputs=[ particle_q, model.particle_inv_mass, model.particle_radius, model.edge_contact_count, model.edge_contact_pairs, model.edge_contact_filter, model.edge_contact_normal ], outputs=[deltas], device=model.device, ) wp.launch( kernel=solve_springs, dim=model.spring_count, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.spring_indices, model.spring_rest_length, model.spring_stiffness, model.spring_damping, dt, model.spring_constraint_lambdas, ], outputs=[deltas], device=model.device, ) # bending constraints if model.edge_count: wp.launch( kernel=bending_constraint, dim=model.edge_count, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.edge_indices, model.edge_rest_angle, model.edge_bending_properties, dt, model.edge_constraint_lambdas, ], outputs=[deltas], device=model.device, ) # Attchement constraints if model.attachment_constraint: wp.launch( kernel=attachment_constraint, dim=model.attachment_count, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.attachment_indices, model.attachment_point, model.attachment_norm, model.attachment_stiffness, model.attachment_damping, model.vert_attach_constraint_lambdas, dt ], outputs=[deltas], device=model.device ) # tetrahedral FEM if model.tet_count: wp.launch( kernel=solve_tetrahedra, dim=model.tet_count, inputs=[ particle_q, particle_qd, model.particle_inv_mass, model.tet_indices, model.tet_poses, model.tet_activations, model.tet_materials, dt, self.soft_body_relaxation, ], outputs=[deltas], device=model.device, ) # apply particle deltas if requires_grad: new_particle_q = wp.clone(particle_q) new_particle_qd = wp.clone(particle_qd) else: new_particle_q = particle_q new_particle_qd = particle_qd wp.launch( kernel=apply_particle_deltas, dim=model.particle_count, inputs=[ state_in.particle_q, particle_q, model.particle_flags, deltas, dt, model.particle_max_velocity, ], outputs=[new_particle_q, new_particle_qd], device=model.device, ) if requires_grad: particle_q.assign(new_particle_q) particle_qd.assign(new_particle_qd) else: particle_q = new_particle_q particle_qd = new_particle_qd # handle rigid bodies # ---------------------------- if model.joint_count: wp.launch( kernel=solve_body_joints, dim=model.joint_count, inputs=[ state_out.body_q, state_out.body_qd, model.body_com, model.body_inv_mass, model.body_inv_inertia, model.joint_type, model.joint_enabled, model.joint_parent, model.joint_child, model.joint_X_p, model.joint_X_c, model.joint_limit_lower, model.joint_limit_upper, model.joint_axis_start, model.joint_axis_dim, model.joint_axis_mode, model.joint_axis, model.joint_target, model.joint_target_ke, model.joint_target_kd, model.joint_linear_compliance, model.joint_angular_compliance, self.joint_angular_relaxation, self.joint_linear_relaxation, dt, ], outputs=[state_out.body_deltas], device=model.device, ) # apply updates wp.launch( kernel=apply_body_deltas, dim=model.body_count, inputs=[ state_out.body_q, state_out.body_qd, model.body_com, model.body_inertia, model.body_inv_mass, model.body_inv_inertia, state_out.body_deltas, None, dt, ], outputs=[ out_body_q, out_body_qd, ], device=model.device, ) if model.body_count and requires_grad: # update state state_out.body_q.assign(out_body_q) state_out.body_qd.assign(out_body_qd) # Solve rigid contact constraints if model.rigid_contact_max and ( model.ground and model.shape_ground_contact_pair_count or model.shape_contact_pair_count ): rigid_contact_inv_weight = None if requires_grad: body_deltas = wp.zeros_like(state_out.body_deltas) rigid_active_contact_distance = wp.zeros_like(model.rigid_active_contact_distance) rigid_active_contact_point0 = wp.empty_like( model.rigid_active_contact_point0, requires_grad=True ) rigid_active_contact_point1 = wp.empty_like( model.rigid_active_contact_point1, requires_grad=True ) if self.rigid_contact_con_weighting: rigid_contact_inv_weight = wp.zeros_like(model.rigid_contact_inv_weight) else: body_deltas = state_out.body_deltas body_deltas.zero_() rigid_active_contact_distance = model.rigid_active_contact_distance rigid_active_contact_point0 = model.rigid_active_contact_point0 rigid_active_contact_point1 = model.rigid_active_contact_point1 rigid_active_contact_distance.zero_() if self.rigid_contact_con_weighting: rigid_contact_inv_weight = model.rigid_contact_inv_weight rigid_contact_inv_weight.zero_() wp.launch( kernel=solve_body_contact_positions, dim=model.rigid_contact_max, inputs=[ state_out.body_q, state_out.body_qd, model.body_com, model.body_inv_mass, model.body_inv_inertia, model.rigid_contact_count, model.rigid_contact_body0, model.rigid_contact_body1, model.rigid_contact_point0, model.rigid_contact_point1, model.rigid_contact_offset0, model.rigid_contact_offset1, model.rigid_contact_normal, model.rigid_contact_thickness, model.rigid_contact_shape0, model.rigid_contact_shape1, model.shape_materials, self.rigid_contact_relaxation, dt, model.rigid_contact_torsional_friction, model.rigid_contact_rolling_friction, ], outputs=[ body_deltas, rigid_active_contact_point0, rigid_active_contact_point1, rigid_active_contact_distance, rigid_contact_inv_weight, ], device=model.device, ) if self.enable_restitution and i == 0: # remember the contacts from the first iteration if requires_grad: model.rigid_active_contact_distance_prev = wp.clone(rigid_active_contact_distance) model.rigid_active_contact_point0_prev = wp.clone(rigid_active_contact_point0) model.rigid_active_contact_point1_prev = wp.clone(rigid_active_contact_point1) if self.rigid_contact_con_weighting: model.rigid_contact_inv_weight_prev = wp.clone(rigid_contact_inv_weight) else: model.rigid_contact_inv_weight_prev = None else: model.rigid_active_contact_distance_prev.assign(rigid_active_contact_distance) model.rigid_active_contact_point0_prev.assign(rigid_active_contact_point0) model.rigid_active_contact_point1_prev.assign(rigid_active_contact_point1) if self.rigid_contact_con_weighting: model.rigid_contact_inv_weight_prev.assign(rigid_contact_inv_weight) else: model.rigid_contact_inv_weight_prev = None if requires_grad: model.rigid_active_contact_distance = rigid_active_contact_distance model.rigid_active_contact_point0 = rigid_active_contact_point0 model.rigid_active_contact_point1 = rigid_active_contact_point1 body_q = wp.clone(state_out.body_q) body_qd = wp.clone(state_out.body_qd) else: body_q = state_out.body_q body_qd = state_out.body_qd # apply updates wp.launch( kernel=apply_body_deltas, dim=model.body_count, inputs=[ state_out.body_q, state_out.body_qd, model.body_com, model.body_inertia, model.body_inv_mass, model.body_inv_inertia, body_deltas, rigid_contact_inv_weight, dt, ], outputs=[ body_q, body_qd, ], device=model.device, ) if requires_grad: state_out.body_q = body_q state_out.body_qd = body_qd # update body velocities from position changes if model.body_count and not requires_grad: # causes gradient issues (probably due to numerical problems # when computing velocities from position changes) if requires_grad: out_body_qd = wp.clone(state_out.body_qd) else: out_body_qd = state_out.body_qd # update body velocities wp.launch( kernel=update_body_velocities, dim=model.body_count, inputs=[state_out.body_q, state_in.body_q, model.body_com, dt], outputs=[out_body_qd], device=model.device, ) if requires_grad: state_out.body_qd.assign(out_body_qd) if self.enable_restitution: if model.particle_count: if requires_grad: new_particle_qd = wp.clone(particle_qd) else: new_particle_qd = particle_qd wp.launch( kernel=apply_soft_restitution_ground, dim=model.particle_count, inputs=[ particle_q, particle_qd, state_in.particle_q, state_in.particle_qd, model.particle_inv_mass, model.particle_radius, model.particle_flags, model.soft_contact_restitution, model.ground_plane, dt, self.soft_contact_relaxation, ], outputs=[new_particle_qd], device=model.device, ) if requires_grad: particle_qd.assign(new_particle_qd) else: particle_qd = new_particle_qd if model.body_count: if requires_grad: state_out.body_deltas = wp.zeros_like(state_out.body_deltas) else: state_out.body_deltas.zero_() wp.launch( kernel=apply_rigid_restitution, dim=model.rigid_contact_max, inputs=[ state_out.body_q, state_out.body_qd, state_in.body_q, state_in.body_qd, model.body_com, model.body_inv_mass, model.body_inv_inertia, model.rigid_contact_count, model.rigid_contact_body0, model.rigid_contact_body1, model.rigid_contact_normal, model.rigid_contact_shape0, model.rigid_contact_shape1, model.shape_materials, model.rigid_active_contact_distance_prev, model.rigid_active_contact_point0_prev, model.rigid_active_contact_point1_prev, model.rigid_contact_inv_weight_prev, model.gravity, dt, ], outputs=[ state_out.body_deltas, ], device=model.device, ) wp.launch( kernel=apply_body_delta_velocities, dim=model.body_count, inputs=[ state_out.body_qd, state_out.body_deltas, ], outputs=[state_out.body_qd], device=model.device, ) # NOTE: Keep mesh points updated # TODO Only if model is indeed a mesh.. if model.particle_count: # update particle state wp.launch( kernel=replace_mesh_points, dim=model.particle_count, inputs=[ model.particle_shape.id, particle_q ], device=model.device, ) model.particle_shape.refit() return state_out